solver: fix reference leak on slow cache errors

Release error input refs with solver state teardown. Give gateways
their own clones and preserve result wrappers during cleanup.

Add regression coverage for missing COPY sources and job cleanup.

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

frontend/dockerfile/dockerfile_cache_test.go

@@ -67,6 +67,28 @@ FROM nanoserver
 	checkAllReleasable(t, c, sb, true)
 }
 
+func testMissingCopySourceReleasesCache(t *testing.T, sb integration.Sandbox) {
+	f := getFrontend(t, sb)
+	dir := integration.Tmpdir(t,
+		fstest.CreateFile("Dockerfile", []byte("FROM scratch\nCOPY missing.txt /\n"), 0600),
+	)
+
+	c, err := client.New(sb.Context(), sb.Address())
+	require.NoError(t, err)
+	defer c.Close()
+
+	for range 2 {
+		_, err := f.Solve(sb.Context(), c, client.SolveOpt{
+			LocalMounts: map[string]fsutil.FS{
+				dockerui.DefaultLocalNameDockerfile: dir,
+				dockerui.DefaultLocalNameContext:    dir,
+			},
+		}, nil)
+		require.ErrorContains(t, err, "missing.txt")
+	}
+	checkAllReleasable(t, c, sb, false)
+}
+
 func testExportCacheLoop(t *testing.T, sb integration.Sandbox) {
 	workers.CheckFeatureCompat(t, sb, workers.FeatureCacheExport, workers.FeatureCacheImport, workers.FeatureCacheBackendLocal)
 	f := getFrontend(t, sb)

frontend/dockerfile/dockerfile_test.go

@@ -59,6 +59,7 @@ var allTests = integration.TestFuncs(
 
 	// dockerfile_cache_test.go
 	testCacheReleased,
+	testMissingCopySourceReleasesCache,
 	testExportCacheLoop,
 	testCacheMultiPlatformImportExport,
 	testImageManifestCacheImportExport,

frontend/gateway/forwarder/forward.go

@@ -36,7 +36,7 @@ func LLBBridgeToGatewayClient(ctx context.Context, llbBridge frontend.FrontendLL
 		sid:               sid,
 		sm:                sm,
 		workers:           w,
-		workerRefByID:     make(map[string]*worker.WorkerRef),
+		resultByID:        make(map[string]solver.Result),
 		executor:          exec,
 		mounts:            make(map[string]snapshot.Mounter),
 	}
@@ -46,17 +46,18 @@ func LLBBridgeToGatewayClient(ctx context.Context, llbBridge frontend.FrontendLL
 
 type BridgeClient struct {
 	frontend.FrontendLLBBridge
-	mu            sync.Mutex
-	opts          map[string]string
-	inputs        map[string]*opspb.Definition
-	sid           string
-	sm            *session.Manager
-	refs          []*ref
-	workers       worker.Infos
-	workerRefByID map[string]*worker.WorkerRef
-	buildOpts     client.BuildOpts
-	ctrs          []client.Container
-	executor      executor.Executor
+	mu         sync.Mutex
+	opts       map[string]string
+	inputs     map[string]*opspb.Definition
+	sid        string
+	sm         *session.Manager
+	refs       []*ref
+	workers    worker.Infos
+	resultByID map[string]solver.Result
+	discarded  bool
+	buildOpts  client.BuildOpts
+	ctrs       []client.Container
+	executor   executor.Executor
 
 	mounts   map[string]snapshot.Mounter
 	mountsMu sync.Mutex
@@ -159,7 +160,7 @@ func (c *BridgeClient) wrapSolveError(solveErr error) error {
 	}
 	if errors.As(solveErr, &sce) {
 		var err error
-		inputIDs, err = c.registerResultIDs(sce.Result)
+		inputIDs, err = c.registerSlowCacheResult(sce.Result)
 		if err != nil {
 			return err
 		}
@@ -172,6 +173,28 @@ func (c *BridgeClient) registerResultIDs(results ...solver.Result) (ids []string
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
+	return c.registerResultIDsLocked(results...)
+}
+
+// registerSlowCacheResult registers a clone of a result borrowed from the
+// solver job. The clone is only taken before the client is discarded, as the
+// job may be released after that.
+func (c *BridgeClient) registerSlowCacheResult(res solver.Result) ([]string, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if c.discarded {
+		return nil, nil
+	}
+	// Validate before cloning so a type error can't leak the clone. Once
+	// cloned, registerResultIDsLocked either keeps or releases it.
+	if _, ok := res.Sys().(*worker.WorkerRef); !ok {
+		return nil, errors.Errorf("unexpected type for result, got %T", res.Sys())
+	}
+	return c.registerResultIDsLocked(res.Clone())
+}
+
+func (c *BridgeClient) registerResultIDsLocked(results ...solver.Result) (ids []string, err error) {
 	ids = make([]string, len(results))
 	for i, res := range results {
 		if res == nil {
@@ -183,15 +206,16 @@ func (c *BridgeClient) registerResultIDs(results ...solver.Result) (ids []string
 		}
 		id := workerRef.ID()
 		ids[i] = id
-		if existing, ok := c.workerRefByID[id]; ok {
-			if existing != workerRef {
-				if err := workerRef.Release(context.TODO()); err != nil {
+		if existing, ok := c.resultByID[id]; ok {
+			if existing != res {
+				if err := res.Release(context.TODO()); err != nil {
 					return ids, errors.WithStack(err)
 				}
 			}
 			continue
 		}
-		c.workerRefByID[id] = workerRef
+		// Keep the result wrapper so Release balances split-result ownership.
+		c.resultByID[id] = res
 	}
 	return ids, nil
 }
@@ -228,10 +252,13 @@ func (c *BridgeClient) discard(err error) {
 
 	c.discardMounts()
 
-	for id, workerRef := range c.workerRefByID {
-		workerRef.Release(context.TODO())
-		delete(c.workerRefByID, id)
+	c.mu.Lock()
+	c.discarded = true
+	for id, res := range c.resultByID {
+		res.Release(context.TODO())
+		delete(c.resultByID, id)
 	}
+	c.mu.Unlock()
 	for _, r := range c.refs {
 		if r != nil {
 			r.resultProxy.Release(context.TODO())
@@ -288,7 +315,10 @@ func (c *BridgeClient) NewContainer(ctx context.Context, req client.NewContainer
 				}
 			} else if m.ResultID != "" {
 				var ok bool
-				workerRef, ok = c.workerRefByID[m.ResultID]
+				res, found := c.resultByID[m.ResultID]
+				if found {
+					workerRef, ok = res.Sys().(*worker.WorkerRef)
+				}
 				if !ok {
 					return errors.Errorf("failed to find ref %s for %q mount", m.ResultID, m.Dest)
 				}

frontend/gateway/gateway.go

@@ -374,9 +374,10 @@ func (lbf *llbBridgeForwarder) Discard() {
 		mount.Unmount()
 	}
 
-	for id, workerRef := range lbf.workerRefByID {
-		workerRef.Release(context.TODO())
-		delete(lbf.workerRefByID, id)
+	lbf.discarded = true
+	for id, res := range lbf.resultByID {
+		res.Release(context.TODO())
+		delete(lbf.resultByID, id)
 	}
 	if lbf.err != nil && lbf.result != nil {
 		lbf.result.EachRef(func(r solver.ResultProxy) error {
@@ -433,19 +434,19 @@ func NewBridgeForwarder(ctx context.Context, llbBridge frontend.FrontendLLBBridg
 
 func newBridgeForwarder(ctx context.Context, llbBridge frontend.FrontendLLBBridge, exec executor.Executor, workers worker.Infos, inputs map[string]*opspb.Definition, sid string, sm *session.Manager) *llbBridgeForwarder {
 	lbf := &llbBridgeForwarder{
-		callCtx:       ctx,
-		llbBridge:     llbBridge,
-		refs:          map[string]solver.ResultProxy{},
-		workerRefByID: map[string]*worker.WorkerRef{},
-		doneCh:        make(chan struct{}),
-		pipe:          newPipe(),
-		workers:       workers,
-		inputs:        inputs,
-		sid:           sid,
-		sm:            sm,
-		ctrs:          map[string]gwclient.Container{},
-		mounts:        map[string]snapshot.Mounter{},
-		executor:      exec,
+		callCtx:    ctx,
+		llbBridge:  llbBridge,
+		refs:       map[string]solver.ResultProxy{},
+		resultByID: map[string]solver.Result{},
+		doneCh:     make(chan struct{}),
+		pipe:       newPipe(),
+		workers:    workers,
+		inputs:     inputs,
+		sid:        sid,
+		sm:         sm,
+		ctrs:       map[string]gwclient.Container{},
+		mounts:     map[string]snapshot.Mounter{},
+		executor:   exec,
 	}
 	return lbf
 }
@@ -540,11 +541,12 @@ type LLBBridgeForwarder interface {
 }
 
 type llbBridgeForwarder struct {
-	mu            sync.Mutex
-	callCtx       context.Context
-	llbBridge     frontend.FrontendLLBBridge
-	refs          map[string]solver.ResultProxy
-	workerRefByID map[string]*worker.WorkerRef
+	mu         sync.Mutex
+	callCtx    context.Context
+	llbBridge  frontend.FrontendLLBBridge
+	refs       map[string]solver.ResultProxy
+	resultByID map[string]solver.Result
+	discarded  bool
 	// lastRef      solver.CachedResult
 	// lastRefs     map[string]solver.CachedResult
 	// err          error
@@ -685,7 +687,7 @@ func (lbf *llbBridgeForwarder) wrapSolveError(solveErr error) error {
 	}
 	if errors.As(solveErr, &sce) {
 		var err error
-		inputIDs, err = lbf.registerResultIDs(sce.Result)
+		inputIDs, err = lbf.registerSlowCacheResult(sce.Result)
 		if err != nil {
 			return err
 		}
@@ -698,6 +700,28 @@ func (lbf *llbBridgeForwarder) registerResultIDs(results ...solver.Result) (ids
 	lbf.mu.Lock()
 	defer lbf.mu.Unlock()
 
+	return lbf.registerResultIDsLocked(results...)
+}
+
+// registerSlowCacheResult registers a clone of a result borrowed from the
+// solver job. The clone is only taken before the forwarder is discarded, as the
+// job may be released after that.
+func (lbf *llbBridgeForwarder) registerSlowCacheResult(res solver.Result) ([]string, error) {
+	lbf.mu.Lock()
+	defer lbf.mu.Unlock()
+
+	if lbf.discarded {
+		return nil, nil
+	}
+	// Validate before cloning so a type error can't leak the clone. Once
+	// cloned, registerResultIDsLocked either keeps or releases it.
+	if _, ok := res.Sys().(*worker.WorkerRef); !ok {
+		return nil, errors.Errorf("unexpected type for result, got %T", res.Sys())
+	}
+	return lbf.registerResultIDsLocked(res.Clone())
+}
+
+func (lbf *llbBridgeForwarder) registerResultIDsLocked(results ...solver.Result) (ids []string, err error) {
 	ids = make([]string, len(results))
 	for i, res := range results {
 		if res == nil {
@@ -709,15 +733,16 @@ func (lbf *llbBridgeForwarder) registerResultIDs(results ...solver.Result) (ids
 		}
 		id := workerRef.ID()
 		ids[i] = id
-		if existing, ok := lbf.workerRefByID[id]; ok {
-			if existing != workerRef {
-				if err := workerRef.Release(context.TODO()); err != nil {
+		if existing, ok := lbf.resultByID[id]; ok {
+			if existing != res {
+				if err := res.Release(context.TODO()); err != nil {
 					return ids, errors.WithStack(err)
 				}
 			}
 			continue
 		}
-		lbf.workerRefByID[id] = workerRef
+		// Keep the result wrapper so Release balances split-result ownership.
+		lbf.resultByID[id] = res
 	}
 	return ids, nil
 }
@@ -1113,7 +1138,10 @@ func (lbf *llbBridgeForwarder) NewContainer(ctx context.Context, in *pb.NewConta
 		var workerRef *worker.WorkerRef
 		if m.ResultID != "" {
 			var ok bool
-			workerRef, ok = lbf.workerRefByID[m.ResultID]
+			res, found := lbf.resultByID[m.ResultID]
+			if found {
+				workerRef, ok = res.Sys().(*worker.WorkerRef)
+			}
 			if !ok {
 				refProxy, err := lbf.convertRef(m.ResultID)
 				if err != nil {

solver/jobs.go

@@ -88,6 +88,21 @@ func (s *state) Cleanup(fn func() error) error {
 	return nil
 }
 
+// cloneUntilRelease borrows the solver's lifetime lock so the state and its
+// dependencies stay alive until the clone has been registered for cleanup.
+func (s *state) cloneUntilRelease(res Result) Result {
+	s.solver.mu.RLock()
+	defer s.solver.mu.RUnlock()
+	if s.solver.actives[s.vtx.Digest()] != s {
+		return nil
+	}
+	clone := res.Clone()
+	s.Cleanup(func() error {
+		return clone.Release(context.TODO())
+	})
+	return clone
+}
+
 func (s *state) ResolverCache() ResolverCache {
 	return s
 }
@@ -238,7 +253,9 @@ func (s *state) setEdge(index Index, targetEdge *edge, targetState *state) {
 
 	if targetState != nil {
 		targetState.addJobs(s, map[*state]struct{}{})
+		targetState.mu.Lock()
 		targetState.releasers = append(targetState.releasers, s.releasers...)
+		targetState.mu.Unlock()
 		s.releasers = nil
 
 		targetState.allPwMu.Lock()
@@ -1054,9 +1071,6 @@ func (s *sharedOp) LoadCache(ctx context.Context, rec *CacheRecord) (Result, fun
 // evaluated, hence "slow" cache.
 func (s *sharedOp) CalcSlowCache(ctx context.Context, index Index, p PreprocessFunc, f ResultBasedCacheFunc, res Result) (dgst digest.Digest, err error) {
 	defer func() {
-		if err != nil {
-			err = WrapSlowCache(err, index, res.Clone())
-		}
 		err = errdefs.WithOp(err, s.st.vtx.Sys(), s.st.vtx.Options().Description)
 		err = errdefs.WrapVertex(err, s.st.origDigest)
 	}()
@@ -1117,6 +1131,10 @@ func (s *sharedOp) CalcSlowCache(ctx context.Context, index Index, p PreprocessF
 		if complete {
 			if err == nil {
 				s.slowCacheRes[index] = key
+			} else if errResult := s.st.cloneUntilRelease(res); errResult != nil {
+				// The error is memoized, so the input is cloned only once and
+				// kept available until the solver state is discarded.
+				err = WrapSlowCache(err, index, errResult)
 			}
 			s.slowCacheErr[index] = err
 		}
@@ -1389,7 +1407,10 @@ func notifyStarted(ctx context.Context, v *client.Vertex, cached bool) func(err
 
 type SlowCacheError struct {
 	error
-	Index  Index
+	Index Index
+	// Result is owned by the solver state that produced the error. Callers
+	// must not release it and need to Clone it while the job is still active
+	// if they want to keep a reference.
 	Result Result
 }
 

solver/scheduler_test.go

@@ -1108,7 +1108,7 @@ func TestSlowCache(t *testing.T) {
 	j1 = nil
 }
 
-func TestSlowCacheErrorResultCloneRelease(t *testing.T) {
+func TestSlowCacheErrorResultReleasedWithJob(t *testing.T) {
 	t.Parallel()
 	ctx := t.Context()
 
@@ -1160,14 +1160,13 @@ func TestSlowCacheErrorResultCloneRelease(t *testing.T) {
 	require.ErrorAs(t, err, &sce)
 	require.NotNil(t, sce.Result)
 
-	require.NoError(t, j.Discard())
-	j = nil
-
 	require.Never(t, func() bool {
 		return releaseCount.Load() != 0
 	}, 100*time.Millisecond, 10*time.Millisecond)
 
-	require.NoError(t, sce.Result.Release(ctx))
+	require.NoError(t, j.Discard())
+	j = nil
+
 	require.Eventually(t, func() bool {
 		return releaseCount.Load() == 1
 	}, time.Second, 10*time.Millisecond)