builder: expand prune to buildkit

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

api/server/backend/build/backend.go

@@ -14,6 +14,7 @@ import (
 	"github.com/docker/docker/image"
 	"github.com/docker/docker/pkg/stringid"
 	"github.com/pkg/errors"
+	"golang.org/x/sync/errgroup"
 )
 
 // ImageComponent provides an interface for working with images
@@ -89,11 +90,33 @@ func (b *Backend) Build(ctx context.Context, config backend.BuildConfig) (string
 
 // PruneCache removes all cached build sources
 func (b *Backend) PruneCache(ctx context.Context) (*types.BuildCachePruneReport, error) {
-	size, err := b.fsCache.Prune(ctx)
+	eg, ctx := errgroup.WithContext(ctx)
+
+	var fsCacheSize uint64
+	eg.Go(func() error {
+		var err error
+		fsCacheSize, err = b.fsCache.Prune(ctx)
 		if err != nil {
-		return nil, errors.Wrap(err, "failed to prune build cache")
+			return errors.Wrap(err, "failed to prune fscache")
 		}
-	return &types.BuildCachePruneReport{SpaceReclaimed: size}, nil
+		return nil
+	})
+
+	var buildCacheSize int64
+	eg.Go(func() error {
+		var err error
+		buildCacheSize, err = b.buildkit.Prune(ctx)
+		if err != nil {
+			return errors.Wrap(err, "failed to prune build cache")
+		}
+		return nil
+	})
+
+	if err := eg.Wait(); err != nil {
+		return nil, err
+	}
+
+	return &types.BuildCachePruneReport{SpaceReclaimed: fsCacheSize + uint64(buildCacheSize)}, nil
 }
 
 func (b *Backend) Cancel(ctx context.Context, id string) error {

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

@@ -13,7 +13,7 @@ import (
 )
 
 func (s *snapshotter) EnsureLayer(ctx context.Context, key string) ([]layer.DiffID, error) {
-	if l, err := s.getLayer(key); err != nil {
+	if l, err := s.getLayer(key, true); err != nil {
 		return nil, err
 	} else if l != nil {
 		return getDiffChain(l), nil
@@ -57,7 +57,7 @@ func (s *snapshotter) EnsureLayer(ctx context.Context, key string) ([]layer.Diff
 	eg.Go(func() error {
 		parent := ""
 		if p := info.Parent; p != "" {
-			if l, err := s.getLayer(p); err != nil {
+			if l, err := s.getLayer(p, true); err != nil {
 				return err
 			} else if l != nil {
 				parent, err = getGraphID(l)

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

@@ -74,7 +74,7 @@ func NewSnapshotter(opt Opt) (snapshot.SnapshotterBase, error) {
 func (s *snapshotter) Prepare(ctx context.Context, key, parent string, opts ...snapshots.Opt) error {
 	origParent := parent
 	if parent != "" {
-		if l, err := s.getLayer(parent); err != nil {
+		if l, err := s.getLayer(parent, false); err != nil {
 			return err
 		} else if l != nil {
 			parent, err = getGraphID(l)
@@ -115,12 +115,16 @@ func (s *snapshotter) chainID(key string) (layer.ChainID, bool) {
 	return "", false
 }
 
-func (s *snapshotter) getLayer(key string) (layer.Layer, error) {
+func (s *snapshotter) getLayer(key string, withCommitted bool) (layer.Layer, error) {
 	s.mu.Lock()
 	l, ok := s.refs[key]
 	if !ok {
 		id, ok := s.chainID(key)
 		if !ok {
+			if !withCommitted {
+				s.mu.Unlock()
+				return nil, nil
+			}
 			if err := s.db.View(func(tx *bolt.Tx) error {
 				b := tx.Bucket([]byte(key))
 				if b == nil {
@@ -146,7 +150,7 @@ func (s *snapshotter) getLayer(key string) (layer.Layer, error) {
 			s.mu.Unlock()
 			return nil, err
 		}
-		s.refs[string(id)] = l
+		s.refs[key] = l
 		if err := s.db.Update(func(tx *bolt.Tx) error {
 			_, err := tx.CreateBucketIfNotExists([]byte(key))
 			return err
@@ -179,25 +183,28 @@ func (s *snapshotter) getGraphDriverID(key string) (string, bool) {
 }
 
 func (s *snapshotter) Stat(ctx context.Context, key string) (snapshots.Info, error) {
-	if l, err := s.getLayer(key); err != nil {
-		return snapshots.Info{}, err
-	} else if l != nil {
-		var parentID string
-		if p := l.Parent(); p != nil {
-			parentID = p.ChainID().String()
-		}
-		info := snapshots.Info{
-			Kind:   snapshots.KindCommitted,
-			Name:   key,
-			Parent: parentID,
-		}
-		return info, nil
-	}
-
 	inf := snapshots.Info{
 		Kind: snapshots.KindActive,
 	}
 
+	l, err := s.getLayer(key, false)
+	if err != nil {
+		return snapshots.Info{}, err
+	}
+	if l != nil {
+		if p := l.Parent(); p != nil {
+			inf.Parent = p.ChainID().String()
+		}
+		inf.Kind = snapshots.KindCommitted
+		inf.Name = string(key)
+		return inf, nil
+	}
+
+	l, err = s.getLayer(key, true)
+	if err != nil {
+		return snapshots.Info{}, err
+	}
+
 	id, committed := s.getGraphDriverID(key)
 	if committed {
 		inf.Kind = snapshots.KindCommitted
@@ -205,13 +212,22 @@ func (s *snapshotter) Stat(ctx context.Context, key string) (snapshots.Info, err
 
 	if err := s.db.View(func(tx *bolt.Tx) error {
 		b := tx.Bucket([]byte(id))
-		if b == nil {
+		if b == nil && l == nil {
-			return errors.Errorf("not found") // TODO: typed
+			return errors.Errorf("snapshot %s not found", id) // TODO: typed
 		}
 		inf.Name = string(key)
+		if b != nil {
 			v := b.Get(keyParent)
 			if v != nil {
 				inf.Parent = string(v)
+				return nil
+			}
+		}
+		if l != nil {
+			if p := l.Parent(); p != nil {
+				inf.Parent = p.ChainID().String()
+			}
+			inf.Kind = snapshots.KindCommitted
 		}
 		return nil
 	}); err != nil {
@@ -221,7 +237,7 @@ func (s *snapshotter) Stat(ctx context.Context, key string) (snapshots.Info, err
 }
 
 func (s *snapshotter) Mounts(ctx context.Context, key string) (snapshot.Mountable, error) {
-	l, err := s.getLayer(key)
+	l, err := s.getLayer(key, true)
 	if err != nil {
 		return nil, err
 	}
@@ -269,16 +285,17 @@ func (s *snapshotter) Mounts(ctx context.Context, key string) (snapshot.Mountabl
 }
 
 func (s *snapshotter) Remove(ctx context.Context, key string) error {
-	l, err := s.getLayer(key)
+	l, err := s.getLayer(key, true)
 	if err != nil {
 		return err
 	}
 
+	id, _ := s.getGraphDriverID(key)
+
 	var found bool
 	if err := s.db.Update(func(tx *bolt.Tx) error {
 		found = tx.Bucket([]byte(key)) != nil
 		if found {
-			id, _ := s.getGraphDriverID(key)
 			tx.DeleteBucket([]byte(key))
 			if id != key {
 				tx.DeleteBucket([]byte(id))
@@ -298,7 +315,6 @@ func (s *snapshotter) Remove(ctx context.Context, key string) error {
 		return nil
 	}
 
-	id, _ := s.getGraphDriverID(key)
 	return s.opt.GraphDriver.Remove(id)
 }
 
@@ -331,9 +347,9 @@ func (s *snapshotter) Update(ctx context.Context, info snapshots.Info, fieldpath
 	return s.Stat(ctx, info.Name)
 }
 
-func (s *snapshotter) Usage(ctx context.Context, key string) (snapshots.Usage, error) {
+func (s *snapshotter) Usage(ctx context.Context, key string) (us snapshots.Usage, retErr error) {
 	usage := snapshots.Usage{}
-	if l, err := s.getLayer(key); err != nil {
+	if l, err := s.getLayer(key, true); err != nil {
 		return usage, err
 	} else if l != nil {
 		s, err := l.DiffSize()
@@ -376,8 +392,17 @@ func (s *snapshotter) Usage(ctx context.Context, key string) (snapshots.Usage, e
 	}
 	var parent string
 	if info.Parent != "" {
+		if l, err := s.getLayer(info.Parent, false); err != nil {
+			return usage, err
+		} else if l != nil {
+			parent, err = getGraphID(l)
+			if err != nil {
+				return usage, err
+			}
+		} else {
 			parent, _ = s.getGraphDriverID(info.Parent)
 		}
+	}
 
 	diffSize, err := s.opt.GraphDriver.DiffSize(id, parent)
 	if err != nil {

builder/builder-next/builder.go

@@ -80,6 +80,34 @@ func (b *Builder) DiskUsage(ctx context.Context) ([]*types.BuildCache, error) {
 	return items, nil
 }
 
+func (b *Builder) Prune(ctx context.Context) (int64, error) {
+	ch := make(chan *controlapi.UsageRecord)
+
+	eg, ctx := errgroup.WithContext(ctx)
+
+	eg.Go(func() error {
+		defer close(ch)
+		return b.controller.Prune(&controlapi.PruneRequest{}, &pruneProxy{
+			streamProxy: streamProxy{ctx: ctx},
+			ch:          ch,
+		})
+	})
+
+	var size int64
+	eg.Go(func() error {
+		for r := range ch {
+			size += r.Size_
+		}
+		return nil
+	})
+
+	if err := eg.Wait(); err != nil {
+		return 0, err
+	}
+
+	return size, nil
+}
+
 func (b *Builder) Build(ctx context.Context, opt backend.BuildConfig) (*builder.Result, error) {
 	var out builder.Result
 
@@ -149,7 +177,7 @@ func (b *Builder) Build(ctx context.Context, opt backend.BuildConfig) (*builder.
 		defer close(ch)
 		return b.controller.Status(&controlapi.StatusRequest{
 			Ref: id,
-		}, &statusProxy{ctx: ctx, ch: ch})
+		}, &statusProxy{streamProxy: streamProxy{ctx: ctx}, ch: ch})
 	})
 
 	eg.Go(func() error {
@@ -188,37 +216,56 @@ func (b *Builder) Build(ctx context.Context, opt backend.BuildConfig) (*builder.
 	return &out, nil
 }
 
-type statusProxy struct {
+type streamProxy struct {
 	ctx context.Context
+}
+
+func (sp *streamProxy) SetHeader(_ grpcmetadata.MD) error {
+	return nil
+}
+
+func (sp *streamProxy) SendHeader(_ grpcmetadata.MD) error {
+	return nil
+}
+
+func (sp *streamProxy) SetTrailer(_ grpcmetadata.MD) {
+}
+
+func (sp *streamProxy) Context() context.Context {
+	return sp.ctx
+}
+func (sp *streamProxy) RecvMsg(m interface{}) error {
+	return io.EOF
+}
+
+type statusProxy struct {
+	streamProxy
 	ch chan *controlapi.StatusResponse
 }
 
-func (sp *statusProxy) SetHeader(_ grpcmetadata.MD) error {
-	return nil
-}
-
-func (sp *statusProxy) SendHeader(_ grpcmetadata.MD) error {
-	return nil
-}
-
-func (sp *statusProxy) SetTrailer(_ grpcmetadata.MD) {
-}
-
 func (sp *statusProxy) Send(resp *controlapi.StatusResponse) error {
 	return sp.SendMsg(resp)
 }
-
-func (sp *statusProxy) Context() context.Context {
-	return sp.ctx
-}
 func (sp *statusProxy) SendMsg(m interface{}) error {
 	if sr, ok := m.(*controlapi.StatusResponse); ok {
 		sp.ch <- sr
 	}
 	return nil
 }
-func (sp *statusProxy) RecvMsg(m interface{}) error {
+
-	return io.EOF
+type pruneProxy struct {
+	streamProxy
+	ch chan *controlapi.UsageRecord
+}
+
+func (sp *pruneProxy) Send(resp *controlapi.UsageRecord) error {
+	return sp.SendMsg(resp)
+}
+func (sp *pruneProxy) SendMsg(m interface{}) error {
+	if sr, ok := m.(*controlapi.UsageRecord); ok {
+		sp.ch <- sr
+	}
+	return nil
 }
 
 type contentStoreNoLabels struct {