ops: fix deadlock on releasing shared mounts

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
(cherry picked from commit 6d907b6)

solver/llbsolver/ops/exec.go

@@ -56,7 +56,8 @@ type execOp struct {
 	platform  *pb.Platform
 	numInputs int
 
-	cacheMounts map[string]*cacheRefShare
+	cacheMounts   map[string]*cacheRefShare
+	cacheMountsMu sync.Mutex
 }
 
 func NewExecOp(v solver.Vertex, op *pb.Op_Exec, platform *pb.Platform, cm cache.Manager, sm *session.Manager, md *metadata.Store, exec executor.Executor, w worker.Worker) (solver.Op, error) {
@@ -222,21 +223,23 @@ func (e *execOp) getMountDeps() ([]dep, error) {
 
 func (e *execOp) getRefCacheDir(ctx context.Context, ref cache.ImmutableRef, id string, m *pb.Mount, sharing pb.CacheSharingOpt) (mref cache.MutableRef, err error) {
 	g := &cacheRefGetter{
-		locker:      CacheMountsLocker(),
-		cacheMounts: e.cacheMounts,
-		cm:          e.cm,
-		md:          e.md,
-		name:        fmt.Sprintf("cached mount %s from exec %s", m.Dest, strings.Join(e.op.Meta.Args, " ")),
+		locker:          &e.cacheMountsMu,
+		cacheMounts:     e.cacheMounts,
+		cm:              e.cm,
+		md:              e.md,
+		globalCacheRefs: sharedCacheRefs,
+		name:            fmt.Sprintf("cached mount %s from exec %s", m.Dest, strings.Join(e.op.Meta.Args, " ")),
 	}
 	return g.getRefCacheDir(ctx, ref, id, sharing)
 }
 
 type cacheRefGetter struct {
-	locker      sync.Locker
-	cacheMounts map[string]*cacheRefShare
-	cm          cache.Manager
-	md          *metadata.Store
-	name        string
+	locker          sync.Locker
+	cacheMounts     map[string]*cacheRefShare
+	cm              cache.Manager
+	md              *metadata.Store
+	globalCacheRefs *cacheRefs
+	name            string
 }
 
 func (g *cacheRefGetter) getRefCacheDir(ctx context.Context, ref cache.ImmutableRef, id string, sharing pb.CacheSharingOpt) (mref cache.MutableRef, err error) {
@@ -261,7 +264,7 @@ func (g *cacheRefGetter) getRefCacheDir(ctx context.Context, ref cache.Immutable
 
 	switch sharing {
 	case pb.CacheSharingOpt_SHARED:
-		return sharedCacheRefs.get(key, func() (cache.MutableRef, error) {
+		return g.globalCacheRefs.get(key, func() (cache.MutableRef, error) {
 			return g.getRefCacheDirNoCache(ctx, key, ref, id, false)
 		})
 	case pb.CacheSharingOpt_PRIVATE:
@@ -814,6 +817,9 @@ func CacheMountsLocker() sync.Locker {
 }
 
 func (r *cacheRefs) get(key string, fn func() (cache.MutableRef, error)) (cache.MutableRef, error) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
 	if r.shares == nil {
 		r.shares = map[string]*cacheRefShare{}
 	}
@@ -830,7 +836,6 @@ func (r *cacheRefs) get(key string, fn func() (cache.MutableRef, error)) (cache.
 
 	share = &cacheRefShare{MutableRef: mref, main: r, key: key, refs: map[*cacheRef]struct{}{}}
 	r.shares[key] = share
-
 	return share.clone(), nil
 }
 
@@ -844,6 +849,9 @@ type cacheRefShare struct {
 
 func (r *cacheRefShare) clone() cache.MutableRef {
 	cacheRef := &cacheRef{cacheRefShare: r}
+	if cacheRefCloneHijack != nil {
+		cacheRefCloneHijack()
+	}
 	r.mu.Lock()
 	r.refs[cacheRef] = struct{}{}
 	r.mu.Unlock()
@@ -852,22 +860,30 @@ func (r *cacheRefShare) clone() cache.MutableRef {
 
 func (r *cacheRefShare) release(ctx context.Context) error {
 	if r.main != nil {
-		r.main.mu.Lock()
-		defer r.main.mu.Unlock()
 		delete(r.main.shares, r.key)
 	}
 	return r.MutableRef.Release(ctx)
 }
 
+var cacheRefReleaseHijack func()
+var cacheRefCloneHijack func()
+
 type cacheRef struct {
 	*cacheRefShare
 }
 
 func (r *cacheRef) Release(ctx context.Context) error {
+	if r.main != nil {
+		r.main.mu.Lock()
+		defer r.main.mu.Unlock()
+	}
 	r.mu.Lock()
 	defer r.mu.Unlock()
 	delete(r.refs, r)
 	if len(r.refs) == 0 {
+		if cacheRefReleaseHijack != nil {
+			cacheRefReleaseHijack()
+		}
 		return r.release(ctx)
 	}
 	return nil