builder: fixes after rebase

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>

api/server/router/system/system_routes.go

@@ -97,7 +97,7 @@ func (s *systemRouter) getDiskUsage(ctx context.Context, w http.ResponseWriter,
 		var err error
 		buildCache, err = s.builder.DiskUsage(ctx)
 		if err != nil {
-			return pkgerrors.Wrap(err, "error getting fscache build cache usage")
+			return pkgerrors.Wrap(err, "error getting build cache usage")
 		}
 		return nil
 	})

builder/builder-next/adapters/containerimage/pull.go

@@ -291,15 +291,14 @@ func (p *puller) Snapshot(ctx context.Context) (cache.ImmutableRef, error) {
 	}
 
 	if p.config != nil {
-		img, err := p.is.ImageStore.Get(image.ID(digest.Digest(p.config)))
-		if err != nil {
-			return nil, err
+		img, err := p.is.ImageStore.Get(image.ID(digest.FromBytes(p.config)))
+		if err == nil {
+			ref, err := p.is.CacheAccessor.GetFromSnapshotter(ctx, string(img.RootFS.ChainID()), cache.WithDescription(fmt.Sprintf("from local %s", p.ref)))
+			if err != nil {
+				return nil, err
+			}
+			return ref, nil
 		}
-		ref, err := p.is.CacheAccessor.GetFromSnapshotter(ctx, string(img.RootFS.ChainID()), cache.WithDescription(fmt.Sprintf("from local %s", p.ref)))
-		if err != nil {
-			return nil, err
-		}
-		return ref, nil
 	}
 
 	ongoing := newJobs(p.ref)

builder/builder-next/adapters/snapshot/snapshot.go

@@ -307,6 +307,9 @@ func (s *snapshotter) Remove(ctx context.Context, key string) error {
 	}
 
 	if l != nil {
+		s.mu.Lock()
+		delete(s.refs, key)
+		s.mu.Unlock()
 		_, err := s.opt.LayerStore.Release(l)
 		return err
 	}

builder/builder-next/builder.go

@@ -109,6 +109,15 @@ func (b *Builder) Prune(ctx context.Context) (int64, error) {
 }
 
 func (b *Builder) Build(ctx context.Context, opt backend.BuildConfig) (*builder.Result, error) {
+	if buildID := opt.Options.BuildID; buildID != "" {
+		b.mu.Lock()
+		ctx, b.jobs[buildID] = context.WithCancel(ctx)
+		b.mu.Unlock()
+		defer func() {
+			delete(b.jobs, buildID)
+		}()
+	}
+
 	var out builder.Result
 
 	id := identity.NewID()

builder/builder-next/exporter/export.go

@@ -73,24 +73,27 @@ func (e *imageExporterInstance) Export(ctx context.Context, ref cache.ImmutableR
 	}
 	config := e.config
 
-	layersDone := oneOffProgress(ctx, "exporting layers")
+	var diffs []digest.Digest
+	if ref != nil {
+		layersDone := oneOffProgress(ctx, "exporting layers")
 
-	if err := ref.Finalize(ctx); err != nil {
-		return nil, err
+		if err := ref.Finalize(ctx); err != nil {
+			return nil, err
+		}
+
+		diffIDs, err := e.opt.Differ.EnsureLayer(ctx, ref.ID())
+		if err != nil {
+			return nil, err
+		}
+
+		diffs = make([]digest.Digest, len(diffIDs))
+		for i := range diffIDs {
+			diffs[i] = digest.Digest(diffIDs[i])
+		}
+
+		layersDone(nil)
 	}
 
-	diffIDs, err := e.opt.Differ.EnsureLayer(ctx, ref.ID())
-	if err != nil {
-		return nil, err
-	}
-
-	diffs := make([]digest.Digest, len(diffIDs))
-	for i := range diffIDs {
-		diffs[i] = digest.Digest(diffIDs[i])
-	}
-
-	layersDone(nil)
-
 	if len(config) == 0 {
 		var err error
 		config, err = emptyImageConfig()