build: move main solve request into main gateway call

Now, we always perform the full solve request in the main gateway call.
This ensures that progress works properly, and makes the lifetime
semantics much clearer.

NewResultContext abstracts the details of a successful/failed build, to
always return a single ResultContext, even though the details of how a
gateway is created is different:
- For a failed build, we can just keep the gateway open.
- For a successful build, we immediately open another gateway and
  re-evaluate the build definition in that gateway. This should give an
  instant cache hit (since the build was just successful).

Signed-off-by: Justin Chadwell <me@jedevc.com>

build/build.go

@@ -885,7 +885,7 @@ func BuildWithResultHandler(ctx context.Context, nodes []builder.Node, opt map[s
 
 					cc := c
 					var printRes map[string][]byte
-					rr, err := c.Build(ctx, so, "buildx", func(ctx context.Context, c gateway.Client) (*gateway.Result, error) {
+					buildFunc := func(ctx context.Context, c gateway.Client) (*gateway.Result, error) {
 						if opt.PrintFunc != nil {
 							if _, ok := req.FrontendOpt["frontend.caps"]; !ok {
 								req.FrontendOpt["frontend.caps"] = "moby.buildkit.frontend.subrequests+forward"
@@ -915,7 +915,6 @@ func BuildWithResultHandler(ctx context.Context, nodes []builder.Node, opt map[s
 							}
 
 							if fallback {
-								fmt.Println("falling back!")
 								req.FrontendOpt["build-arg:BUILDKIT_SYNTAX"] = printFallbackImage
 								res2, err2 := c.Solve(ctx, req)
 								if err2 != nil {
@@ -931,16 +930,16 @@ func BuildWithResultHandler(ctx context.Context, nodes []builder.Node, opt map[s
 						}
 
 						results.Set(resultKey(dp.driverIndex, k), res)
-						if resultHandleFunc != nil {
-							resultCtx, err := NewResultContext(ctx, cc, so, res)
-							if err == nil {
-								resultHandleFunc(dp.driverIndex, resultCtx)
-							} else {
-								logrus.Warnf("failed to record result: %s", err)
-							}
-						}
 						return res, nil
-					}, ch)
+					}
+					var rr *client.SolveResponse
+					if resultHandleFunc != nil {
+						var resultCtx *ResultContext
+						resultCtx, rr, err = NewResultContext(ctx, cc, so, "buildx", buildFunc, ch)
+						resultHandleFunc(dp.driverIndex, resultCtx)
+					} else {
+						rr, err = c.Build(ctx, so, "buildx", buildFunc, ch)
+					}
 					if err != nil {
 						return err
 					}

build/result.go

@@ -6,8 +6,6 @@ import (
 	"encoding/json"
 	"io"
 	"sync"
-	"sync/atomic"
-	"time"
 
 	controllerapi "github.com/docker/buildx/controller/pb"
 	"github.com/moby/buildkit/client"
@@ -22,140 +20,232 @@ import (
 	"golang.org/x/sync/errgroup"
 )
 
-func NewResultContext(ctx context.Context, c *client.Client, solveOpt client.SolveOpt, res *gateway.Result) (*ResultContext, error) {
+// NewResultContext wraps a call to client.Build, additionally returning a
-	def, err := getDefinition(ctx, res)
+// ResultContext alongside the standard response and error.
-	if err != nil {
+//
-		return nil, err
+// This ResultContext can be used to execute additional build steps in the same
-	}
+// context as the build occurred, which can allow easy debugging of build
-	return getResultAt(ctx, c, solveOpt, def, nil)
+// failures and successes.
-}
+//
-
+// If the returned ResultContext is not nil, the caller must call Done() on it.
-func getDefinition(ctx context.Context, res *gateway.Result) (*result.Result[*pb.Definition], error) {
+func NewResultContext(ctx context.Context, cc *client.Client, opt client.SolveOpt, product string, buildFunc gateway.BuildFunc, ch chan *client.SolveStatus) (*ResultContext, *client.SolveResponse, error) {
-	return result.ConvertResult(res, func(ref gateway.Reference) (*pb.Definition, error) {
+	// Create a new context to wrap the original, and cancel it when the
-		st, err := ref.ToState()
+	// caller-provided context is cancelled.
-		if err != nil {
+	//
-			return nil, err
+	// We derive the context from the background context so that we can forbid
-		}
+	// cancellation of the build request after <-done is closed (which we do
-		def, err := st.Marshal(ctx)
+	// before returning the ResultContext).
-		if err != nil {
+	baseCtx := ctx
-			return nil, err
+	ctx, cancel := context.WithCancelCause(context.Background())
+	done := make(chan struct{})
+	go func() {
+		select {
+		case <-baseCtx.Done():
+			cancel(baseCtx.Err())
+		case <-done:
+			// Once done is closed, we've recorded a ResultContext, so we
+			// shouldn't allow cancelling the underlying build request anymore.
 		}
-		return def.ToPB(), nil
+	}()
-	})
-}
-
-func getResultAt(ctx context.Context, c *client.Client, solveOpt client.SolveOpt, targets *result.Result[*pb.Definition], statusChan chan *client.SolveStatus) (*ResultContext, error) {
-	ctx, cancel := context.WithCancel(ctx)
-	defer cancel()
 
-	// forward SolveStatus
+	// Create a new channel to forward status messages to the original.
-	done := new(atomic.Bool)
+	//
-	defer done.Store(true)
+	// We do this so that we can discard status messages after the main portion
-	ch := make(chan *client.SolveStatus)
+	// of the build is complete. This is necessary for the solve error case,
+	// where the original gateway is kept open until the ResultContext is
+	// closed - we don't want progress messages from operations in that
+	// ResultContext to display after this function exits.
+	//
+	// Additionally, callers should wait for the progress channel to be closed.
+	// If we keep the session open and never close the progress channel, the
+	// caller will likely hang.
+	baseCh := ch
+	ch = make(chan *client.SolveStatus)
 	go func() {
 		for {
-			s := <-ch
+			s, ok := <-ch
-			if s == nil {
+			if !ok {
 				return
 			}
-			if done.Load() {
-				// Do not forward if the function returned because statusChan is possibly closed
-				continue
-			}
 			select {
-			case statusChan <- s:
+			case <-baseCh:
-			case <-ctx.Done():
+				// base channel is closed, discard status messages
+			default:
+				baseCh <- s
 			}
 		}
 	}()
+	defer close(baseCh)
+
+	var resCtx *ResultContext
+	var resp *client.SolveResponse
+	var respErr error
 
-	// get result
-	resultCtxCh := make(chan *ResultContext)
-	errCh := make(chan error)
 	go func() {
-		solveOpt := solveOpt
+		defer cancel(context.Canceled) // ensure no dangling processes
-		solveOpt.Ref = ""
+
-		buildDoneCh := make(chan struct{})
+		var res *gateway.Result
-		_, err := c.Build(context.Background(), solveOpt, "buildx", func(ctx context.Context, c gateway.Client) (*gateway.Result, error) {
+		var err error
-			doneErr := errors.Errorf("done")
+		resp, err = cc.Build(ctx, opt, product, func(ctx context.Context, c gateway.Client) (*gateway.Result, error) {
-			ctx, cancel := context.WithCancelCause(ctx)
+			var err error
-			defer cancel(doneErr)
+			res, err = buildFunc(ctx, c)
-
+
-			// force evaluation of all targets in parallel
+			if res != nil && err == nil {
-			results := make(map[*pb.Definition]*gateway.Result)
+				// Force evaluation of the build result (otherwise, we likely
-			resultsMu := sync.Mutex{}
+				// won't get a solve error)
-			eg, egCtx := errgroup.WithContext(ctx)
+				def, err2 := getDefinition(ctx, res)
-			targets.EachRef(func(def *pb.Definition) error {
+				if err2 != nil {
-				eg.Go(func() error {
+					return nil, err2
-					res2, err := c.Solve(egCtx, gateway.SolveRequest{
+				}
-						Evaluate:   true,
+				res, err = evalDefinition(ctx, c, def)
-						Definition: def,
+			}
-					})
+
-					if err != nil {
+			if err != nil {
-						return err
+				// Scenario 1: we failed to evaluate a node somewhere in the
-					}
+				// build graph.
-					resultsMu.Lock()
+				//
-					results[def] = res2
+				// In this case, we construct a ResultContext from this
-					resultsMu.Unlock()
+				// original Build session, and return it alongside the original
-					return nil
+				// build error. We then need to keep the gateway session open
-				})
+				// until the caller explicitly closes the ResultContext.
-				return nil
+
-			})
-			resultCtx := ResultContext{}
-			if err := eg.Wait(); err != nil {
 				var se *errdefs.SolveError
 				if errors.As(err, &se) {
-					resultCtx.solveErr = se
+					resCtx = &ResultContext{
-				} else {
+						done:     make(chan struct{}),
-					return nil, err
+						solveErr: se,
+						gwClient: c,
+						gwCtx:    ctx,
+					}
+					respErr = se
+					close(done)
+
+					// Block until the caller closes the ResultContext.
+					select {
+					case <-resCtx.done:
+					case <-ctx.Done():
+					}
 				}
 			}
-			res2, _ := result.ConvertResult(targets, func(def *pb.Definition) (gateway.Reference, error) {
+			return res, err
-				if res, ok := results[def]; ok {
+		}, ch)
-					return res.Ref, nil
+		if resCtx != nil {
-				}
+			return
-				return nil, nil
+		}
-			})
+		if err != nil {
+			// Something unexpected failed during the build, we didn't succeed,
+			// but we also didn't make it far enough to create a ResultContext.
+			respErr = err
+			close(done)
+			return
+		}
 
-			// Record the client and ctx as well so that containers can be created from the SolveError.
+		// Scenario 2: we successfully built the image with no errors.
-			resultCtx.res = res2
+		//
-			resultCtx.gwClient = c
+		// In this case, the original gateway session has now been closed
-			resultCtx.gwCtx = ctx
+		// since the Build has been completed. So, we need to create a new
-			resultCtx.gwDone = func() {
+		// gateway session to populate the ResultContext. To do this, we
-				cancel(doneErr)
+		// need to re-evaluate the target result, in this new session. This
-				// wait for Build() completion(or timeout) to ensure the Build's finalizing and avoiding an error "context canceled"
+		// should be instantaneous since the result should be cached.
-				select {
+
-				case <-buildDoneCh:
+		def, err := getDefinition(ctx, res)
-				case <-time.After(5 * time.Second):
+		if err != nil {
-				}
+			respErr = err
+			close(done)
+			return
+		}
+
+		// NOTE: ideally this second connection should be lazily opened
+		opt := opt
+		opt.Ref = ""
+		_, respErr = cc.Build(ctx, opt, "buildx", func(ctx context.Context, c gateway.Client) (*gateway.Result, error) {
+			res, err := evalDefinition(ctx, c, def)
+			if err != nil {
+				// This should probably not happen, since we've previously
+				// successfully evaluated the same result with no issues.
+				return nil, errors.Wrap(err, "inconsistent solve result")
 			}
-			select {
+			resCtx = &ResultContext{
-			case resultCtxCh <- &resultCtx:
+				done:     make(chan struct{}),
-			case <-ctx.Done():
+				res:      res,
-				return nil, ctx.Err()
+				gwClient: c,
+				gwCtx:    ctx,
 			}
+			close(done)
 
-			// wait for cleanup or cancel
+			// Block until the caller closes the ResultContext.
-			<-ctx.Done()
+			select {
-			if context.Cause(ctx) != doneErr { // doneErr is not an error.
+			case <-resCtx.done:
-				return nil, ctx.Err()
+			case <-ctx.Done():
 			}
-			return nil, nil
+			return nil, ctx.Err()
-		}, ch)
+		}, nil)
-		close(buildDoneCh)
+		if resCtx != nil {
-		if err != nil {
+			return
-			errCh <- err
 		}
+		close(done)
 	}()
 
+	// Block until the other thread signals that it's completed the build.
 	select {
-	case resultCtx := <-resultCtxCh:
+	case <-done:
-		return resultCtx, nil
+	case <-baseCtx.Done():
-	case err := <-errCh:
+		if respErr == nil {
+			respErr = baseCtx.Err()
+		}
+	}
+	return resCtx, resp, respErr
+}
+
+// getDefinition converts a gateway result into a collection of definitions for
+// each ref in the result.
+func getDefinition(ctx context.Context, res *gateway.Result) (*result.Result[*pb.Definition], error) {
+	return result.ConvertResult(res, func(ref gateway.Reference) (*pb.Definition, error) {
+		st, err := ref.ToState()
+		if err != nil {
+			return nil, err
+		}
+		def, err := st.Marshal(ctx)
+		if err != nil {
+			return nil, err
+		}
+		return def.ToPB(), nil
+	})
+}
+
+// evalDefinition performs the reverse of getDefinition, converting a
+// collection of definitions into a gateway result.
+func evalDefinition(ctx context.Context, c gateway.Client, defs *result.Result[*pb.Definition]) (*gateway.Result, error) {
+	// force evaluation of all targets in parallel
+	results := make(map[*pb.Definition]*gateway.Result)
+	resultsMu := sync.Mutex{}
+	eg, egCtx := errgroup.WithContext(ctx)
+	defs.EachRef(func(def *pb.Definition) error {
+		eg.Go(func() error {
+			res, err := c.Solve(egCtx, gateway.SolveRequest{
+				Evaluate:   true,
+				Definition: def,
+			})
+			if err != nil {
+				return err
+			}
+			resultsMu.Lock()
+			results[def] = res
+			resultsMu.Unlock()
+			return nil
+		})
+		return nil
+	})
+	if err := eg.Wait(); err != nil {
 		return nil, err
-	case <-ctx.Done():
-		return nil, ctx.Err()
 	}
+	res, _ := result.ConvertResult(defs, func(def *pb.Definition) (gateway.Reference, error) {
+		if res, ok := results[def]; ok {
+			return res.Ref, nil
+		}
+		return nil, nil
+	})
+	return res, nil
 }
 
 // ResultContext is a build result with the client that built it.
@@ -163,17 +253,18 @@ type ResultContext struct {
 	res      *gateway.Result
 	solveErr *errdefs.SolveError
 
-	gwClient   gateway.Client
+	done     chan struct{}
-	gwCtx      context.Context
+	doneOnce sync.Once
-	gwDone     func()
+
-	gwDoneOnce sync.Once
+	gwClient gateway.Client
+	gwCtx    context.Context
 
 	cleanups   []func()
 	cleanupsMu sync.Mutex
 }
 
 func (r *ResultContext) Done() {
-	r.gwDoneOnce.Do(func() {
+	r.doneOnce.Do(func() {
 		r.cleanupsMu.Lock()
 		cleanups := r.cleanups
 		r.cleanups = nil
@@ -181,7 +272,9 @@ func (r *ResultContext) Done() {
 		for _, f := range cleanups {
 			f()
 		}
-		r.gwDone()
+
+		close(r.done)
+		<-r.gwCtx.Done()
 	})
 }