controller: avoid "context canceled" error on cleanup

Signed-off-by: Kohei Tokunaga <ktokunaga.mail@gmail.com>

Author: ktock

build/result.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	controllerapi "github.com/docker/buildx/controller/pb"
 	"github.com/moby/buildkit/client"
@@ -73,9 +74,8 @@ func getResultAt(ctx context.Context, c *client.Client, solveOpt client.SolveOpt
 	resultCtxCh := make(chan *ResultContext)
 	errCh := make(chan error)
 	go func() {
+		buildDoneCh := make(chan struct{})
 		_, err := c.Build(context.Background(), solveOpt, "buildx", func(ctx context.Context, c gateway.Client) (*gateway.Result, error) {
-			ctx, cancel := context.WithCancel(ctx)
-			defer cancel()
 			resultCtx := ResultContext{}
 			res2, err := c.Solve(ctx, gateway.SolveRequest{
 				Evaluate:   true,
@@ -90,18 +90,35 @@ func getResultAt(ctx context.Context, c *client.Client, solveOpt client.SolveOpt
 				}
 			}
 			// Record the client and ctx as well so that containers can be created from the SolveError.
+			doneCh := make(chan struct{})
 			resultCtx.res = res2
 			resultCtx.gwClient = c
 			resultCtx.gwCtx = ctx
-			resultCtx.gwDone = cancel
+			resultCtx.gwDone = func() {
+				close(doneCh)
+				// wait for Build() completion(or timeout) to ensure the Build's finalizing and avoiding an error "context canceled"
+				select {
+				case <-buildDoneCh:
+				case <-time.After(5 * time.Second):
+				}
+			}
 			select {
 			case resultCtxCh <- &resultCtx:
+			case <-doneCh:
+				return nil, nil
+			case <-ctx.Done():
+				return nil, ctx.Err()
+			}
+
+			// wait for cleanup or cancel
+			select {
+			case <-doneCh:
 			case <-ctx.Done():
 				return nil, ctx.Err()
 			}
-			<-ctx.Done()
 			return nil, nil
 		}, ch)
+		close(buildDoneCh)
 		if err != nil {
 			errCh <- err
 		}