@@ -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)

-	if err != nil {

-		return nil, errors.Wrap(err, "failed to prune build cache")

+	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 errors.Wrap(err, "failed to prune fscache")

+		}

+		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: size}, nil

+

+	return &types.BuildCachePruneReport{SpaceReclaimed: fsCacheSize + uint64(buildCacheSize)}, nil

 }

 func (b *Backend) Cancel(ctx context.Context, id string) error {

@@ -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)

@@ -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,23 +183,26 @@ 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 {

+	inf := snapshots.Info{

+		Kind: snapshots.KindActive,

+	}

+

+	l, err := s.getLayer(key, false)

+	if err != nil {

 		return snapshots.Info{}, err

-	} else if l != nil {

-		var parentID string

+	}

+	if l != nil {

 		if p := l.Parent(); p != nil {

-			parentID = p.ChainID().String()

-		}

-		info := snapshots.Info{

-			Kind:   snapshots.KindCommitted,

-			Name:   key,

-			Parent: parentID,

+			inf.Parent = p.ChainID().String()

 		}

-		return info, nil

+		inf.Kind = snapshots.KindCommitted

+		inf.Name = string(key)

+		return inf, nil

 	}

-	inf := snapshots.Info{

-		Kind: snapshots.KindActive,

+	l, err = s.getLayer(key, true)

+	if err != nil {

+		return snapshots.Info{}, err

 	}

 	id, committed := s.getGraphDriverID(key)

@@ -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 {

-			return errors.Errorf("not found") // TODO: typed

+		if b == nil && l == nil {

+			return errors.Errorf("snapshot %s not found", id) // TODO: typed

 		}

 		inf.Name = string(key)

-		v := b.Get(keyParent)

-		if v != nil {

-			inf.Parent = string(v)

+		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,7 +392,16 @@ func (s *snapshotter) Usage(ctx context.Context, key string) (snapshots.Usage, e

 	}

 	var parent string

 	if info.Parent != "" {

-		parent, _ = s.getGraphDriverID(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)

@@ -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,28 +216,35 @@ func (b *Builder) Build(ctx context.Context, opt backend.BuildConfig) (*builder.

 	return &out, nil

 }

-type statusProxy struct {

+type streamProxy struct {

 	ctx context.Context

-	ch  chan *controlapi.StatusResponse

 }

-func (sp *statusProxy) SetHeader(_ grpcmetadata.MD) error {

+func (sp *streamProxy) SetHeader(_ grpcmetadata.MD) error {

 	return nil

 }

-func (sp *statusProxy) SendHeader(_ grpcmetadata.MD) error {

+func (sp *streamProxy) SendHeader(_ grpcmetadata.MD) error {

 	return nil

 }

-func (sp *statusProxy) SetTrailer(_ grpcmetadata.MD) {

+func (sp *streamProxy) SetTrailer(_ grpcmetadata.MD) {

 }

-func (sp *statusProxy) Send(resp *controlapi.StatusResponse) error {

-	return sp.SendMsg(resp)

+func (sp *streamProxy) Context() context.Context {

+	return sp.ctx

+}

+func (sp *streamProxy) RecvMsg(m interface{}) error {

+	return io.EOF

 }

-func (sp *statusProxy) Context() context.Context {

-	return sp.ctx

+type statusProxy struct {

+	streamProxy

+	ch chan *controlapi.StatusResponse

+}

+

+func (sp *statusProxy) Send(resp *controlapi.StatusResponse) error {

+	return sp.SendMsg(resp)

 }

 func (sp *statusProxy) SendMsg(m interface{}) error {

 	if sr, ok := m.(*controlapi.StatusResponse); ok {

@@ -217,8 +252,20 @@ func (sp *statusProxy) SendMsg(m interface{}) error {

 	}

 	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 {