solver: fix failed exec ref ownership

ExecOp returned failed outputs together with ExecError. When the exec was
canceled, sharedOp dropped those outputs without releasing them, leaving
their cache records in use.

Make Op.Exec return nil outputs on error and move failed ExecOp outputs into
ExecError. sharedOp now retains clones of a cached error's refs until the
vertex is released, so builds sharing a failed vertex can still use its
snapshots after another build is discarded.

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

client/client_test.go

@@ -274,6 +274,7 @@ func TestClientGatewayIntegration(t *testing.T) {
 		testClientGatewayContainerPID1Tty,
 		testClientGatewayContainerSignal,
 		testClientGatewayExecError,
+		testClientGatewaySharedExecError,
 		testClientGatewayExecFileActionError,
 		testClientGatewaySlowCacheExecError,
 

client/gateway_shared_error_test.go

@@ -0,0 +1,106 @@
+package client
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/moby/buildkit/client/llb"
+	gwclient "github.com/moby/buildkit/frontend/gateway/client"
+	"github.com/moby/buildkit/identity"
+	"github.com/moby/buildkit/solver/errdefs"
+	"github.com/moby/buildkit/util/testutil/integration"
+	"github.com/pkg/errors"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/sync/errgroup"
+)
+
+func testClientGatewaySharedExecError(t *testing.T, sb integration.Sandbox) {
+	requiresLinux(t)
+	ctx, cancel := context.WithTimeoutCause(sb.Context(), time.Minute, errors.New("shared exec error test timed out"))
+	defer cancel()
+
+	c, err := New(ctx, sb.Address())
+	require.NoError(t, err)
+	defer c.Close()
+
+	def, err := llb.Image("busybox:latest").Run(
+		llb.Shlexf("sh -c 'echo %s > /data && exit 1'", identity.NewID()),
+	).Root().Marshal(ctx)
+	require.NoError(t, err)
+
+	solve := func(ctx context.Context, c gwclient.Client) (*errdefs.SolveError, error) {
+		_, err := c.Solve(ctx, gwclient.SolveRequest{
+			Definition: def.ToPB(),
+			Evaluate:   true,
+		})
+		var se *errdefs.SolveError
+		if !errors.As(err, &se) {
+			return nil, errors.Errorf("expected SolveError, got %v", err)
+		}
+		if len(se.MountIDs) != 1 || se.MountIDs[0] == "" {
+			return nil, errors.Errorf("expected failed rootfs mount, got %v", se.MountIDs)
+		}
+		return se, nil
+	}
+
+	firstReady := make(chan struct{})
+	secondReady := make(chan struct{})
+	firstDone := make(chan struct{})
+	eg, ctx := errgroup.WithContext(ctx)
+	eg.Go(func() error {
+		defer close(firstDone)
+		_, err := c.Build(ctx, SolveOpt{}, "buildkit_test", func(ctx context.Context, c gwclient.Client) (*gwclient.Result, error) {
+			if _, err := solve(ctx, c); err != nil {
+				return nil, err
+			}
+			close(firstReady)
+			// Keep the failed vertex alive until both builds hold its error.
+			select {
+			case <-secondReady:
+				return gwclient.NewResult(), nil
+			case <-ctx.Done():
+				return nil, errors.WithStack(context.Cause(ctx))
+			}
+		}, nil)
+		return err
+	})
+	eg.Go(func() error {
+		_, err := c.Build(ctx, SolveOpt{}, "buildkit_test", func(ctx context.Context, gc gwclient.Client) (*gwclient.Result, error) {
+			select {
+			case <-firstReady:
+			case <-ctx.Done():
+				return nil, errors.WithStack(context.Cause(ctx))
+			}
+			se, err := solve(ctx, gc)
+			if err != nil {
+				return nil, err
+			}
+			close(secondReady)
+			select {
+			case <-firstDone:
+			case <-ctx.Done():
+				return nil, errors.WithStack(context.Cause(ctx))
+			}
+
+			// Discarding the first build must not let GC remove the snapshot
+			// still needed by the second build's cached error.
+			if err := c.Prune(ctx, nil, PruneAll); err != nil {
+				return nil, err
+			}
+			ctr, err := gc.NewContainer(ctx, gwclient.NewContainerRequest{
+				Mounts: []gwclient.Mount{{Dest: "/", ResultID: se.MountIDs[0]}},
+			})
+			if err != nil {
+				return nil, err
+			}
+			if err := ctr.Release(ctx); err != nil {
+				return nil, err
+			}
+			return gwclient.NewResult(), nil
+		}, nil)
+		return err
+	})
+	require.NoError(t, eg.Wait())
+	checkAllReleasable(t, c, sb, true)
+}

frontend/dockerfile/dockerfile_cache_test.go

@@ -5,6 +5,7 @@ import (
 	"io"
 	"os"
 	"path/filepath"
+	"strconv"
 	"strings"
 	"testing"
 	"time"
@@ -23,6 +24,7 @@ import (
 	"github.com/moby/buildkit/util/testutil/workers"
 	ocispecs "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/pkg/errors"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/tonistiigi/fsutil"
 )
@@ -89,6 +91,61 @@ func testMissingCopySourceReleasesCache(t *testing.T, sb integration.Sandbox) {
 	checkAllReleasable(t, c, sb, false)
 }
 
+func testCanceledRunReleasesCache(t *testing.T, sb integration.Sandbox) {
+	integration.SkipOnPlatform(t, "windows")
+	f := getFrontend(t, sb)
+	dir := integration.Tmpdir(t, fstest.CreateFile("Dockerfile", []byte(`FROM busybox
+ARG BUST
+RUN echo "$BUST" > /parent && dd if=/dev/zero of=/a bs=1M count=1
+RUN dd if=/dev/zero of=/b bs=1M count=1 && echo CANCEL-READY && sleep 120
+`), 0600))
+	c, err := client.New(sb.Context(), sb.Address())
+	require.NoError(t, err)
+	defer c.Close()
+
+	for i := range 2 {
+		ctx, cancel := context.WithCancelCause(sb.Context())
+		ctx, timeoutCancel := context.WithTimeoutCause(ctx, time.Minute, errors.New("RUN did not reach cancellation point"))
+		status := make(chan *client.SolveStatus)
+		done := make(chan bool, 1)
+		go func() {
+			started := false
+			for st := range status {
+				for _, log := range st.Logs {
+					if strings.Contains(string(log.Data), "CANCEL-READY") {
+						started = true
+						cancel(errors.WithStack(context.Canceled))
+					}
+				}
+			}
+			done <- started
+		}()
+		_, err := f.Solve(ctx, c, client.SolveOpt{
+			FrontendAttrs: map[string]string{"build-arg:BUST": strconv.Itoa(i)},
+			LocalMounts: map[string]fsutil.FS{
+				dockerui.DefaultLocalNameDockerfile: dir,
+				dockerui.DefaultLocalNameContext:    dir,
+			},
+		}, status)
+		timeoutCancel()
+		cancel(errors.WithStack(context.Canceled))
+		require.True(t, <-done, "build must be cancelled during RUN")
+		require.Error(t, err)
+	}
+	// Cancellation must release the cache even while build history is retained.
+	require.EventuallyWithT(t, func(t *assert.CollectT) {
+		du, err := c.DiskUsage(sb.Context())
+		require.NoError(t, err)
+		for _, record := range du {
+			require.False(t, record.InUse, "cache record %s is still in use", record.ID)
+		}
+	}, 10*time.Second, 100*time.Millisecond)
+	require.NoError(t, c.Prune(sb.Context(), nil, client.PruneAll))
+	du, err := c.DiskUsage(sb.Context())
+	require.NoError(t, err)
+	require.Empty(t, du)
+}
+
 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

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

solver/jobs.go

@@ -1021,6 +1021,8 @@ type sharedOp struct {
 	execRes  *execRes
 	execDone bool
 	execErr  error
+	// execErrRefs retains error refs while execErr is cached.
+	execErrRefs []Result
 
 	cacheRes  []*CacheMap
 	cacheDone bool
@@ -1294,6 +1296,9 @@ func (s *sharedOp) Exec(ctx context.Context, inputs []Result) (outputs []Result,
 
 				s.execRes = &execRes{execRes: wrapShared(res), execExporters: subExporters}
 			}
+			if err != nil {
+				s.execErrRefs = cloneErrorRefs(err)
+			}
 			s.execErr = err
 		}
 		if s.execRes == nil || err != nil {
@@ -1344,6 +1349,9 @@ func (s *sharedOp) release() {
 			go r.Release(context.TODO())
 		}
 	}
+	for _, r := range s.execErrRefs {
+		go r.Release(context.TODO())
+	}
 }
 
 func initClientVertex(v Vertex) client.Vertex {
@@ -1439,6 +1447,20 @@ func WrapSlowCache(err error, index Index, res Result) error {
 	return &SlowCacheError{Index: index, Result: res, error: err}
 }
 
+func cloneErrorRefs(err error) (refs []Result) {
+	for ; err != nil; err = errors.Unwrap(err) {
+		if re, ok := err.(interface {
+			EachRef(func(Result) error) error
+		}); ok {
+			re.EachRef(func(r Result) error {
+				refs = append(refs, r.Clone())
+				return nil
+			})
+		}
+	}
+	return refs
+}
+
 func releaseError(err error) {
 	if err == nil {
 		return

solver/llbsolver/ops/exec.go

@@ -404,18 +404,11 @@ func (e *ExecOp) Exec(ctx context.Context, jobCtx solver.JobContext, inputs []so
 			}
 			execMounts := make([]solver.Result, len(e.op.Mounts))
 			copy(execMounts, execInputs)
+			// Failed outputs belong to ExecError, not to the caller.
 			for i, res := range results {
-				// res.Clone() (not res) is required: results[i] is owned by
-				// the caller (and ultimately released via the gateway / job
-				// path), while execMounts[i] is embedded in the ExecError
-				// returned below and released independently by the error
-				// owner. Sharing the same *workerRefResult here would mean
-				// two independent owners holding the same *WorkerRef, so a
-				// Release on one would also free the cache ref held by the
-				// other. See worker/result_test.go for the underlying
-				// ownership invariant.
-				execMounts[p.OutputRefs[i].MountIndex] = res.Clone()
+				execMounts[p.OutputRefs[i].MountIndex] = res
 			}
+			results = nil
 			for _, active := range p.Actives {
 				if active.NoCommit {
 					active.Ref.Release(context.TODO())

solver/types.go

@@ -176,6 +176,8 @@ type Op interface {
 	CacheMap(context.Context, JobContext, int) (*CacheMap, bool, error)
 
 	// Exec runs an operation given results from previous operations.
+	// On error, outputs must be nil; the operation must release any outputs
+	// it created or transfer their ownership to the error.
 	Exec(ctx context.Context, jobCtx JobContext, inputs []Result) (outputs []Result, err error)
 
 	// Acquire acquires the necessary resources to execute the `Op`.