From 372df78cc8da912c12655ee92dfd033675d6028f Mon Sep 17 00:00:00 2001
From: Edgar Lee <edgarl@netflix.com>
Date: Wed, 25 Nov 2020 11:26:53 -0800
Subject: [PATCH] Guarantee err results are released when result proxy is
 released

Signed-off-by: Edgar Lee <edgarl@netflix.com>
---
 solver/llbsolver/bridge.go | 48 ++++++++++++++++++++++++++------------
 solver/llbsolver/solver.go | 15 ------------
 2 files changed, 33 insertions(+), 30 deletions(-)

diff --git a/solver/llbsolver/bridge.go b/solver/llbsolver/bridge.go
index 52a88a23..30d92320 100644
--- a/solver/llbsolver/bridge.go
+++ b/solver/llbsolver/bridge.go
@@ -16,6 +16,7 @@ import (
 	"github.com/moby/buildkit/session"
 	"github.com/moby/buildkit/solver"
 	"github.com/moby/buildkit/solver/errdefs"
+	llberrdefs "github.com/moby/buildkit/solver/llbsolver/errdefs"
 	"github.com/moby/buildkit/solver/pb"
 	"github.com/moby/buildkit/util/flightcontrol"
 	"github.com/moby/buildkit/worker"
@@ -144,41 +145,58 @@ func (b *llbBridge) Solve(ctx context.Context, req frontend.SolveRequest, sid st
 }
 
 type resultProxy struct {
-	cb       func(context.Context) (solver.CachedResult, error)
-	def      *pb.Definition
-	g        flightcontrol.Group
-	mu       sync.Mutex
-	released bool
-	v        solver.CachedResult
-	err      error
+	cb         func(context.Context) (solver.CachedResult, error)
+	def        *pb.Definition
+	g          flightcontrol.Group
+	mu         sync.Mutex
+	released   bool
+	v          solver.CachedResult
+	err        error
+	errResults []solver.Result
 }
 
 func newResultProxy(b *llbBridge, req frontend.SolveRequest) *resultProxy {
-	return &resultProxy{
+	rp := &resultProxy{
 		def: req.Definition,
-		cb: func(ctx context.Context) (solver.CachedResult, error) {
-			return b.loadResult(ctx, req.Definition, req.CacheImports)
-		},
 	}
+	rp.cb = func(ctx context.Context) (solver.CachedResult, error) {
+		res, err := b.loadResult(ctx, req.Definition, req.CacheImports)
+		var ee *llberrdefs.ExecError
+		if errors.As(err, &ee) {
+			ee.EachRef(func(res solver.Result) error {
+				rp.errResults = append(rp.errResults, res)
+				return nil
+			})
+		}
+		return res, err
+	}
+	return rp
 }
 
 func (rp *resultProxy) Definition() *pb.Definition {
 	return rp.def
 }
 
-func (rp *resultProxy) Release(ctx context.Context) error {
+func (rp *resultProxy) Release(ctx context.Context) (err error) {
 	rp.mu.Lock()
 	defer rp.mu.Unlock()
+	for _, res := range rp.errResults {
+		rerr := res.Release(ctx)
+		if rerr != nil {
+			err = rerr
+		}
+	}
 	if rp.v != nil {
 		if rp.released {
 			logrus.Warnf("release of already released result")
 		}
-		if err := rp.v.Release(ctx); err != nil {
-			return err
+		rerr := rp.v.Release(ctx)
+		if err != nil {
+			return rerr
 		}
 	}
 	rp.released = true
-	return nil
+	return
 }
 
 func (rp *resultProxy) wrapError(err error) error {
diff --git a/solver/llbsolver/solver.go b/solver/llbsolver/solver.go
index 1413bc8c..d3ef5852 100644
--- a/solver/llbsolver/solver.go
+++ b/solver/llbsolver/solver.go
@@ -16,7 +16,6 @@ import (
 	"github.com/moby/buildkit/frontend/gateway"
 	"github.com/moby/buildkit/session"
 	"github.com/moby/buildkit/solver"
-	"github.com/moby/buildkit/solver/llbsolver/errdefs"
 	"github.com/moby/buildkit/util/compression"
 	"github.com/moby/buildkit/util/entitlements"
 	"github.com/moby/buildkit/util/progress"
@@ -143,20 +142,6 @@ func (s *Solver) Solve(ctx context.Context, id string, sessionID string, req fro
 	res.EachRef(func(ref solver.ResultProxy) error {
 		eg.Go(func() error {
 			_, err := ref.Result(ctx2)
-			if err != nil {
-				// Also release any results referenced by exec errors.
-				var ee *errdefs.ExecError
-				if errors.As(err, &ee) {
-					ee.EachRef(func(res solver.Result) error {
-
-						workerRef, ok := res.Sys().(*worker.WorkerRef)
-						if !ok {
-							return nil
-						}
-						return workerRef.ImmutableRef.Release(ctx)
-					})
-				}
-			}
 			return err
 		})
 		return nil