replace context.WithCancel with WithCancelCause

Keep stack traces for cancellation errors where possible.

Signed-off-by: Tonis Tiigi <[email protected]>

Author: tonistiigi

.golangci.yml

@@ -41,6 +41,8 @@ linters-settings:
     forbid:
       - '^fmt\.Errorf(# use errors\.Errorf instead)?$'
       - '^logrus\.(Trace|Debug|Info|Warn|Warning|Error|Fatal)(f|ln)?(# use bklog\.G or bklog\.L instead of logrus directly)?$'
+      - '^context\.WithCancel(# use context\.WithCancelCause instead)?$'
+      - '^ctx\.Err(# use context\.Cause instead)?$'
   importas:
     alias:
       - pkg: "github.com/opencontainers/image-spec/specs-go/v1"

cache/manager.go

@@ -1258,7 +1258,7 @@ func (cm *cacheManager) prune(ctx context.Context, ch chan client.UsageInfo, opt
 
 	select {
 	case <-ctx.Done():
-		return ctx.Err()
+		return context.Cause(ctx)
 	default:
 		return cm.prune(ctx, ch, opt)
 	}

client/build_test.go

@@ -1266,8 +1266,8 @@ func testClientGatewayContainerCancelExecTty(t *testing.T, sb integration.Sandbo
 		defer pid1.Wait()
 		defer ctr.Release(ctx)
 
-		execCtx, cancel := context.WithCancel(ctx)
-		defer cancel()
+		execCtx, cancel := context.WithCancelCause(ctx)
+		defer cancel(errors.WithStack(context.Canceled))
 
 		prompt := newTestPrompt(execCtx, t, inputW, output)
 		pid2, err := ctr.Start(execCtx, client.StartRequest{
@@ -1281,7 +1281,7 @@ func testClientGatewayContainerCancelExecTty(t *testing.T, sb integration.Sandbo
 		require.NoError(t, err)
 
 		prompt.SendExpect("echo hi", "hi")
-		cancel()
+		cancel(errors.WithStack(context.Canceled))
 
 		err = pid2.Wait()
 		require.ErrorIs(t, err, context.Canceled)

client/client.go

@@ -205,7 +205,7 @@ func (c *Client) Wait(ctx context.Context) error {
 
 		select {
 		case <-ctx.Done():
-			return ctx.Err()
+			return context.Cause(ctx)
 		case <-time.After(time.Second):
 		}
 		c.conn.ResetConnectBackoff()

client/client_test.go

@@ -7407,8 +7407,8 @@ func testInvalidExporter(t *testing.T, sb integration.Sandbox) {
 
 // moby/buildkit#492
 func testParallelLocalBuilds(t *testing.T, sb integration.Sandbox) {
-	ctx, cancel := context.WithCancel(sb.Context())
-	defer cancel()
+	ctx, cancel := context.WithCancelCause(sb.Context())
+	defer cancel(errors.WithStack(context.Canceled))
 
 	c, err := New(ctx, sb.Address())
 	require.NoError(t, err)

client/llb/async.go

@@ -61,7 +61,7 @@ func (as *asyncState) Do(ctx context.Context, c *Constraints) error {
 		if err != nil {
 			select {
 			case <-ctx.Done():
-				if errors.Is(err, ctx.Err()) {
+				if errors.Is(err, context.Cause(ctx)) {
 					return res, err
 				}
 			default:

client/solve.go

@@ -106,8 +106,8 @@ func (c *Client) solve(ctx context.Context, def *llb.Definition, runGateway runG
 	}
 	eg, ctx := errgroup.WithContext(ctx)
 
-	statusContext, cancelStatus := context.WithCancel(context.Background())
-	defer cancelStatus()
+	statusContext, cancelStatus := context.WithCancelCause(context.Background())
+	defer cancelStatus(errors.WithStack(context.Canceled))
 
 	if span := trace.SpanFromContext(ctx); span.SpanContext().IsValid() {
 		statusContext = trace.ContextWithSpan(statusContext, span)
@@ -230,16 +230,16 @@ func (c *Client) solve(ctx context.Context, def *llb.Definition, runGateway runG
 		frontendAttrs[k] = v
 	}
 
-	solveCtx, cancelSolve := context.WithCancel(ctx)
+	solveCtx, cancelSolve := context.WithCancelCause(ctx)
 	var res *SolveResponse
 	eg.Go(func() error {
 		ctx := solveCtx
-		defer cancelSolve()
+		defer cancelSolve(errors.WithStack(context.Canceled))
 
 		defer func() { // make sure the Status ends cleanly on build errors
 			go func() {
 				<-time.After(3 * time.Second)
-				cancelStatus()
+				cancelStatus(errors.WithStack(context.Canceled))
 			}()
 			if !opt.SessionPreInitialized {
 				bklog.G(ctx).Debugf("stopping session")
@@ -298,7 +298,7 @@ func (c *Client) solve(ctx context.Context, def *llb.Definition, runGateway runG
 			select {
 			case <-solveCtx.Done():
 			case <-time.After(5 * time.Second):
-				cancelSolve()
+				cancelSolve(errors.WithStack(context.Canceled))
 			}
 
 			return err

cmd/buildkitd/main.go

@@ -223,8 +223,8 @@ func main() {
 		if os.Geteuid() > 0 {
 			return errors.New("rootless mode requires to be executed as the mapped root in a user namespace; you may use RootlessKit for setting up the namespace")
 		}
-		ctx, cancel := context.WithCancel(appcontext.Context())
-		defer cancel()
+		ctx, cancel := context.WithCancelCause(appcontext.Context())
+		defer cancel(errors.WithStack(context.Canceled))
 
 		cfg, err := config.LoadFile(c.GlobalString("config"))
 		if err != nil {
@@ -344,9 +344,9 @@ func main() {
 		select {
 		case serverErr := <-errCh:
 			err = serverErr
-			cancel()
+			cancel(err)
 		case <-ctx.Done():
-			err = ctx.Err()
+			err = context.Cause(ctx)
 		}
 
 		bklog.G(ctx).Infof("stopping server")
@@ -634,14 +634,14 @@ func unaryInterceptor(globalCtx context.Context, tp trace.TracerProvider) grpc.U
 	withTrace := otelgrpc.UnaryServerInterceptor(otelgrpc.WithTracerProvider(tp), otelgrpc.WithPropagators(propagators))
 
 	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
-		ctx, cancel := context.WithCancel(ctx)
-		defer cancel()
+		ctx, cancel := context.WithCancelCause(ctx)
+		defer cancel(errors.WithStack(context.Canceled))
 
 		go func() {
 			select {
 			case <-ctx.Done():
 			case <-globalCtx.Done():
-				cancel()
+				cancel(context.Cause(globalCtx))
 			}
 		}()
 

control/control.go

@@ -505,10 +505,10 @@ func (c *Controller) Session(stream controlapi.Control_SessionServer) error {
 	conn, closeCh, opts := grpchijack.Hijack(stream)
 	defer conn.Close()
 
-	ctx, cancel := context.WithCancel(stream.Context())
+	ctx, cancel := context.WithCancelCause(stream.Context())
 	go func() {
 		<-closeCh
-		cancel()
+		cancel(errors.WithStack(context.Canceled))
 	}()
 
 	err := c.opt.SessionManager.HandleConn(ctx, conn, opts)

executor/containerdexecutor/executor.go

@@ -243,7 +243,7 @@ func (w *containerdExecutor) Exec(ctx context.Context, id string, process execut
 		}
 		select {
 		case <-ctx.Done():
-			return ctx.Err()
+			return context.Cause(ctx)
 		case err, ok := <-details.done:
 			if !ok || err == nil {
 				return errors.Errorf("container %s has stopped", id)
@@ -336,8 +336,8 @@ func (w *containerdExecutor) runProcess(ctx context.Context, p containerd.Proces
 
 	// handle signals (and resize) in separate go loop so it does not
 	// potentially block the container cancel/exit status loop below.
-	eventCtx, eventCancel := context.WithCancel(ctx)
-	defer eventCancel()
+	eventCtx, eventCancel := context.WithCancelCause(ctx)
+	defer eventCancel(errors.WithStack(context.Canceled))
 	go func() {
 		for {
 			select {
@@ -403,7 +403,7 @@ func (w *containerdExecutor) runProcess(ctx context.Context, p containerd.Proces
 				}
 				select {
 				case <-ctx.Done():
-					exitErr.Err = errors.Wrap(ctx.Err(), exitErr.Error())
+					exitErr.Err = errors.Wrap(context.Cause(ctx), exitErr.Error())
 				default:
 				}
 				return exitErr