solver: release results of builds outliving Discard

A build still running when its job was discarded stored its results
in states that were already released, leaking them. Keep the states
until the last running build returns, and reject new builds on a
discarded job. Release refs of a gateway solve that returns after the
forwarder was discarded.

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

client/client_test.go

@@ -285,6 +285,7 @@ func TestClientGatewayIntegration(t *testing.T) {
 		testClientGatewayRepeatedExecError,
 		testClientBuildOpExecErrorReleased,
 		testClientGatewayExecErrorReleasedDuringLoad,
+		testClientGatewaySolveReleasedAfterDiscard,
 		testClientGatewayExecFileActionError,
 		testClientGatewaySlowCacheExecError,
 

client/gateway_shared_error_test.go

@@ -249,3 +249,70 @@ func testClientBuildOpExecErrorReleased(t *testing.T, sb integration.Sandbox) {
 	require.NoError(t, err)
 	checkAllReleasable(t, c, sb, true)
 }
+
+// testClientGatewaySolveReleasedAfterDiscard checks that the refs of a solve
+// that finishes after the gateway is discarded are released.
+func testClientGatewaySolveReleasedAfterDiscard(t *testing.T, sb integration.Sandbox) {
+	requiresLinux(t)
+	ctx, cancel := context.WithTimeoutCause(sb.Context(), time.Minute, errors.New("released late solve test timed out"))
+	defer cancel()
+
+	c, err := New(ctx, sb.Address())
+	require.NoError(t, err)
+	defer c.Close()
+
+	id := identity.NewID()
+	// The exec succeeds only after a marker appears in a shared cache mount.
+	// The marker is written once the first build has returned, so the solve
+	// always finishes after the gateway is discarded.
+	cache := llb.AsPersistentCacheDir(id, llb.CacheMountShared)
+	def, err := llb.Image("busybox:latest").Run(
+		llb.Shlexf("sh -c 'echo %s; i=0; while [ ! -e /cache/go ] && [ $i -lt 300 ]; do sleep 0.1; i=$((i+1)); done'", id),
+		llb.AddMount("/cache", llb.Scratch(), cache),
+	).Root().Marshal(ctx)
+	require.NoError(t, err)
+	markerDef, err := llb.Image("busybox:latest").Run(
+		llb.Shlexf("sh -c 'echo %s; touch /cache/go'", id),
+		llb.AddMount("/cache", llb.Scratch(), cache),
+	).Root().Marshal(ctx)
+	require.NoError(t, err)
+
+	started := make(chan struct{})
+	statusCh := make(chan *SolveStatus)
+	statusDone := make(chan struct{})
+	go func() {
+		defer close(statusDone)
+		var once sync.Once
+		for st := range statusCh {
+			for _, v := range st.Vertexes {
+				if v.Started != nil && strings.Contains(v.Name, id) {
+					once.Do(func() { close(started) })
+				}
+			}
+		}
+	}()
+
+	solveErr := make(chan error, 1)
+	_, err = c.Build(ctx, SolveOpt{}, "buildkit_test", func(ctx context.Context, gc gwclient.Client) (*gwclient.Result, error) {
+		go func() {
+			_, err := gc.Solve(ctx, gwclient.SolveRequest{Definition: def.ToPB(), Evaluate: true})
+			solveErr <- err
+		}()
+		// return while the exec is running so the gateway is discarded
+		select {
+		case <-started:
+		case <-ctx.Done():
+			return nil, errors.WithStack(context.Cause(ctx))
+		}
+		return gwclient.NewResult(), nil
+	}, statusCh)
+	require.NoError(t, err)
+	<-statusDone
+
+	// the gateway has been discarded, let the exec finish
+	_, err = c.Solve(ctx, markerDef, SolveOpt{}, nil)
+	require.NoError(t, err)
+	require.ErrorContains(t, <-solveErr, "gateway forwarder is closing")
+
+	checkAllReleasable(t, c, sb, true)
+}

frontend/gateway/gateway.go

@@ -838,6 +838,16 @@ func (lbf *llbBridgeForwarder) Solve(ctx context.Context, req *pb.SolveRequest)
 
 	lbf.mu.Lock()
 
+	if lbf.discarded {
+		// A solve that finished after Discard has no one to release its refs.
+		lbf.mu.Unlock()
+		res.EachRef(func(r solver.ResultProxy) error {
+			r.Release(context.TODO())
+			return nil
+		})
+		return nil, stack.Enable(status.Error(codes.Unavailable, "gateway forwarder is closing"))
+	}
+
 	if res.Refs != nil {
 		ids := make(map[string]string, len(res.Refs))
 		defs := make(map[string]*opspb.Definition, len(res.Refs))

solver/jobs.go

@@ -434,7 +434,9 @@ func (sb *subBuilder) EachValue(ctx context.Context, key string, fn func(any) er
 }
 
 type Job struct {
-	mu            sync.Mutex // protects completedTime, pw, span
+	mu            sync.Mutex // protects completedTime, pw, span, builds, discarded
+	builds        int
+	discarded     bool
 	list          *Solver
 	pr            *progress.MultiReader
 	pw            progress.Writer
@@ -846,11 +848,17 @@ func (jl *Solver) deleteIfUnreferenced(k digest.Digest, st *state) {
 }
 
 func (j *Job) Build(ctx context.Context, e Edge) (CachedResultWithProvenance, error) {
-	if span := trace.SpanFromContext(ctx); span.SpanContext().IsValid() {
-		j.mu.Lock()
-		j.span = span
+	j.mu.Lock()
+	if j.discarded {
 		j.mu.Unlock()
+		return nil, errors.Errorf("build on discarded job %s", j.id)
 	}
+	j.builds++
+	if span := trace.SpanFromContext(ctx); span.SpanContext().IsValid() {
+		j.span = span
+	}
+	j.mu.Unlock()
+	defer j.buildDone()
 
 	v, err := j.list.load(ctx, e.Vertex, nil, j)
 	if err != nil {
@@ -919,7 +927,32 @@ func (j *Job) CloseProgress() {
 	j.pw.Close()
 }
 
+// buildDone finishes a pending Discard when the last running build returns.
+func (j *Job) buildDone() {
+	j.mu.Lock()
+	j.builds--
+	last := j.discarded && j.builds == 0
+	j.mu.Unlock()
+	if last {
+		j.discard()
+	}
+}
+
+// Discard releases the job's states. While builds are still running the states
+// are released when the last one returns, so results that complete after
+// Discard are released with their state.
 func (j *Job) Discard() error {
+	j.mu.Lock()
+	j.discarded = true
+	running := j.builds > 0
+	j.mu.Unlock()
+	if running {
+		return nil
+	}
+	return j.discard()
+}
+
+func (j *Job) discard() error {
 	j.list.mu.Lock()
 	defer j.list.mu.Unlock()
 

solver/scheduler_test.go

@@ -1297,6 +1297,8 @@ func TestExecErrorRefsReleasedAfterJobDiscard(t *testing.T) {
 
 	// A failed exec can finish after its job has released the vertex.
 	require.NoError(t, j.Discard())
+	_, err = j.Build(ctx, g)
+	require.ErrorContains(t, err, "build on discarded job")
 	resumeExec()
 	select {
 	case err := <-done:
@@ -1608,15 +1610,15 @@ func TestMultipleCacheSources(t *testing.T) {
 		}),
 	}
 
-	res, err = j1.Build(ctx, g2)
+	res, err = j2.Build(ctx, g2)
 	require.NoError(t, err)
 	require.Equal(t, "result2", unwrap(res))
 
 	require.Equal(t, int64(2), cacheManager.loadCounter)
 	require.Equal(t, int64(0), cacheManager2.loadCounter)
 
-	require.NoError(t, j1.Discard())
-	j1 = nil
+	require.NoError(t, j2.Discard())
+	j2 = nil
 }
 
 func TestRepeatBuildWithIgnoreCache(t *testing.T) {