worker: remove unreferenced cache mount after release

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

cache/manager.go

@@ -36,6 +36,7 @@ type Accessor interface {
 	New(ctx context.Context, s ImmutableRef, opts ...RefOption) (MutableRef, error)
 	GetMutable(ctx context.Context, id string) (MutableRef, error) // Rebase?
 	IdentityMapping() *idtools.IdentityMapping
+	Metadata(string) *metadata.StorageItem
 }
 
 type Controller interface {
@@ -124,6 +125,16 @@ func (cm *cacheManager) GetFromSnapshotter(ctx context.Context, id string, opts
 	return cm.get(ctx, id, true, opts...)
 }
 
+func (cm *cacheManager) Metadata(id string) *metadata.StorageItem {
+	cm.mu.Lock()
+	defer cm.mu.Unlock()
+	r, ok := cm.records[id]
+	if !ok {
+		return nil
+	}
+	return r.Metadata()
+}
+
 // get requires manager lock to be taken
 func (cm *cacheManager) get(ctx context.Context, id string, fromSnapshotter bool, opts ...RefOption) (*immutableRef, error) {
 	rec, err := cm.getRecord(ctx, id, fromSnapshotter, opts...)

solver/llbsolver/bridge.go

@@ -105,7 +105,7 @@ func (b *llbBridge) Solve(ctx context.Context, req frontend.SolveRequest) (res *
 				ids = append(ids, id)
 			}
 			if err := b.eachWorker(func(w worker.Worker) error {
-				return w.PruneCacheMounts(ids)
+				return w.PruneCacheMounts(ctx, ids)
 			}); err != nil {
 				return nil, err
 			}

worker/base/worker.go

@@ -219,7 +219,7 @@ func (w *Worker) ResolveOp(v solver.Vertex, s frontend.FrontendLLBBridge, sm *se
 	return nil, errors.Errorf("could not resolve %v", v)
 }
 
-func (w *Worker) PruneCacheMounts(ids []string) error {
+func (w *Worker) PruneCacheMounts(ctx context.Context, ids []string) error {
 	mu := ops.CacheMountsLocker()
 	mu.Lock()
 	defer mu.Unlock()
@@ -233,6 +233,9 @@ func (w *Worker) PruneCacheMounts(ids []string) error {
 		for _, si := range sis {
 			for _, k := range si.Indexes() {
 				if k == id || strings.HasPrefix(k, id+":") {
+					if siCached := w.CacheManager.Metadata(si.ID()); siCached != nil {
+						si = siCached
+					}
 					if err := cache.CachePolicyDefault(si); err != nil {
 						return err
 					}
@@ -242,6 +245,10 @@ func (w *Worker) PruneCacheMounts(ids []string) error {
 					if err := si.Commit(); err != nil {
 						return err
 					}
+					// if ref is unused try to clean it up right away by releasing it
+					if mref, err := w.CacheManager.GetMutable(ctx, si.ID()); err == nil {
+						go mref.Release(context.TODO())
+					}
 					break
 				}
 			}

worker/worker.go

@@ -33,7 +33,7 @@ type Worker interface {
 	Prune(ctx context.Context, ch chan client.UsageInfo, opt ...client.PruneInfo) error
 	GetRemote(ctx context.Context, ref cache.ImmutableRef, createIfNeeded bool) (*solver.Remote, error)
 	FromRemote(ctx context.Context, remote *solver.Remote) (cache.ImmutableRef, error)
-	PruneCacheMounts(ids []string) error
+	PruneCacheMounts(ctx context.Context, ids []string) error
 }
 
 // Pre-defined label keys