improve stacks of cancels from defers

In this case the current stack trace points to the line
where the context was created. Instead the stack should be
captured when the defer is running so the return path to
the defer call is also part of the stack.

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

cache/remotecache/local/local.go

@@ -107,7 +107,7 @@ func getContentStore(ctx context.Context, sm *session.Manager, g session.Group,
 	}
 	timeoutCtx, cancel := context.WithCancelCause(ctx)
 	timeoutCtx, _ = context.WithTimeoutCause(timeoutCtx, 5*time.Second, errors.WithStack(context.DeadlineExceeded))
-	defer cancel(errors.WithStack(context.Canceled))
+	defer func() { cancel(errors.WithStack(context.Canceled)) }()
 
 	caller, err := sm.Get(timeoutCtx, sessionID, false)
 	if err != nil {

client/build_test.go

@@ -1236,7 +1236,7 @@ func testClientGatewayContainerCancelExecTty(t *testing.T, sb integration.Sandbo
 		defer ctr.Release(ctx)
 
 		execCtx, cancel := context.WithCancelCause(ctx)
-		defer cancel(errors.WithStack(context.Canceled))
+		defer func() { cancel(errors.WithStack(context.Canceled)) }()
 
 		prompt := newTestPrompt(execCtx, t, inputW, output)
 		pid2, err := ctr.Start(execCtx, client.StartRequest{

client/client_test.go

@@ -8018,7 +8018,7 @@ func testInvalidExporter(t *testing.T, sb integration.Sandbox) {
 func testParallelLocalBuilds(t *testing.T, sb integration.Sandbox) {
 	integration.SkipOnPlatform(t, "windows")
 	ctx, cancel := context.WithCancelCause(sb.Context())
-	defer cancel(errors.WithStack(context.Canceled))
+	defer func() { cancel(errors.WithStack(context.Canceled)) }()
 
 	c, err := New(ctx, sb.Address())
 	require.NoError(t, err)

cmd/buildctl/common/common.go

@@ -86,7 +86,7 @@ func ResolveClient(c *cli.Context) (*client.Client, error) {
 		ctx2, cancel := context.WithCancelCause(ctx)
 		ctx2, _ = context.WithTimeoutCause(ctx2, timeout*time.Second, errors.WithStack(context.DeadlineExceeded))
 		ctx = ctx2
-		defer cancel(errors.WithStack(context.Canceled))
+		defer func() { cancel(errors.WithStack(context.Canceled)) }()
 	}
 
 	cl, err := client.New(ctx, c.GlobalString("addr"), opts...)

cmd/buildkitd/main.go

@@ -233,7 +233,7 @@ func main() {
 			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.WithCancelCause(appcontext.Context())
-		defer cancel(errors.WithStack(context.Canceled))
+		defer func() { cancel(errors.WithStack(context.Canceled)) }()
 
 		cfg, err := config.LoadFile(c.GlobalString("config"))
 		if err != nil {

control/gateway/gateway.go

@@ -61,7 +61,7 @@ func (gwf *GatewayForwarder) lookupForwarder(ctx context.Context) (gateway.LLBBr
 
 	ctx, cancel := context.WithCancelCause(ctx)
 	ctx, _ = context.WithTimeoutCause(ctx, 3*time.Second, errors.WithStack(context.DeadlineExceeded))
-	defer cancel(errors.WithStack(context.Canceled))
+	defer func() { cancel(errors.WithStack(context.Canceled)) }()
 
 	go func() {
 		<-ctx.Done()

executor/runcexecutor/executor.go

@@ -547,7 +547,7 @@ func (k procKiller) Kill(ctx context.Context) (err error) {
 	// shorter timeout but here as a fail-safe for future refactoring.
 	ctx, cancel := context.WithCancelCause(ctx)
 	ctx, _ = context.WithTimeoutCause(ctx, 10*time.Second, errors.WithStack(context.DeadlineExceeded))
-	defer cancel(errors.WithStack(context.Canceled))
+	defer func() { cancel(errors.WithStack(context.Canceled)) }()
 
 	if k.pidfile == "" {
 		// for `runc run` process we use `runc kill` to terminate the process
@@ -694,7 +694,7 @@ func (p *procHandle) WaitForReady(ctx context.Context) error {
 func (p *procHandle) WaitForStart(ctx context.Context, startedCh <-chan int, started func()) error {
 	ctx, cancel := context.WithCancelCause(ctx)
 	ctx, _ = context.WithTimeoutCause(ctx, 10*time.Second, errors.WithStack(context.DeadlineExceeded))
-	defer cancel(errors.WithStack(context.Canceled))
+	defer func() { cancel(errors.WithStack(context.Canceled)) }()
 	select {
 	case <-ctx.Done():
 		return errors.New("go-runc started message never received")

exporter/local/export.go

@@ -81,7 +81,7 @@ func (e *localExporter) Config() *exporter.Config {
 func (e *localExporterInstance) Export(ctx context.Context, inp *exporter.Source, _ exptypes.InlineCache, sessionID string) (map[string]string, exporter.DescriptorReference, error) {
 	timeoutCtx, cancel := context.WithCancelCause(ctx)
 	timeoutCtx, _ = context.WithTimeoutCause(timeoutCtx, 5*time.Second, errors.WithStack(context.DeadlineExceeded))
-	defer cancel(errors.WithStack(context.Canceled))
+	defer func() { cancel(errors.WithStack(context.Canceled)) }()
 
 	if e.opts.Epoch == nil {
 		if tm, ok, err := epoch.ParseSource(inp); err != nil {

exporter/oci/export.go

@@ -218,7 +218,7 @@ func (e *imageExporterInstance) Export(ctx context.Context, src *exporter.Source
 
 	timeoutCtx, cancel := context.WithCancelCause(ctx)
 	timeoutCtx, _ = context.WithTimeoutCause(timeoutCtx, 5*time.Second, errors.WithStack(context.DeadlineExceeded))
-	defer cancel(errors.WithStack(context.Canceled))
+	defer func() { cancel(errors.WithStack(context.Canceled)) }()
 
 	caller, err := e.opt.SessionManager.Get(timeoutCtx, sessionID, false)
 	if err != nil {

exporter/tar/export.go

@@ -165,7 +165,7 @@ func (e *localExporterInstance) Export(ctx context.Context, inp *exporter.Source
 
 	timeoutCtx, cancel := context.WithCancelCause(ctx)
 	timeoutCtx, _ = context.WithTimeoutCause(timeoutCtx, 5*time.Second, errors.WithStack(context.DeadlineExceeded))
-	defer cancel(errors.WithStack(context.Canceled))
+	defer func() { cancel(errors.WithStack(context.Canceled)) }()
 
 	caller, err := e.opt.SessionManager.Get(timeoutCtx, sessionID, false)
 	if err != nil {

frontend/dockerfile/dockerfile_test.go

@@ -3493,7 +3493,7 @@ COPY . .
 
 	ctx, cancel := context.WithCancelCause(sb.Context())
 	ctx, _ = context.WithTimeoutCause(ctx, 15*time.Second, errors.WithStack(context.DeadlineExceeded))
-	defer cancel(errors.WithStack(context.Canceled))
+	defer func() { cancel(errors.WithStack(context.Canceled)) }()
 
 	c, err := client.New(ctx, sb.Address())
 	require.NoError(t, err)

session/filesync/filesync.go

@@ -206,7 +206,7 @@ func FSSync(ctx context.Context, c session.Caller, opt FSSendRequestOpt) error {
 	opts[keyDirName] = []string{opt.Name}
 
 	ctx, cancel := context.WithCancelCause(ctx)
-	defer cancel(errors.WithStack(context.Canceled))
+	defer func() { cancel(errors.WithStack(context.Canceled)) }()
 
 	client := NewFileSyncClient(c.Conn())
 

session/group.go

@@ -74,7 +74,7 @@ func (sm *Manager) Any(ctx context.Context, g Group, f func(context.Context, str
 
 		timeoutCtx, cancel := context.WithCancelCause(ctx)
 		timeoutCtx, _ = context.WithTimeoutCause(timeoutCtx, 5*time.Second, errors.WithStack(context.DeadlineExceeded))
-		defer cancel(errors.WithStack(context.Canceled))
+		defer func() { cancel(errors.WithStack(context.Canceled)) }()
 		c, err := sm.Get(timeoutCtx, id, false)
 		if err != nil {
 			lastErr = err

session/manager.go

@@ -99,7 +99,7 @@ func (sm *Manager) HandleConn(ctx context.Context, conn net.Conn, opts map[strin
 // caller needs to take lock, this function will release it
 func (sm *Manager) handleConn(ctx context.Context, conn net.Conn, opts map[string][]string) error {
 	ctx, cancel := context.WithCancelCause(ctx)
-	defer cancel(errors.WithStack(context.Canceled))
+	defer func() { cancel(errors.WithStack(context.Canceled)) }()
 
 	opts = canonicalHeaders(opts)
 
@@ -154,7 +154,7 @@ func (sm *Manager) Get(ctx context.Context, id string, noWait bool) (Caller, err
 	}
 
 	ctx, cancel := context.WithCancelCause(ctx)
-	defer cancel(errors.WithStack(context.Canceled))
+	defer func() { cancel(errors.WithStack(context.Canceled)) }()
 
 	go func() {
 		<-ctx.Done()

session/session.go

@@ -96,7 +96,7 @@ func (s *Session) Run(ctx context.Context, dialer Dialer) error {
 	s.cancelCtx = cancel
 	s.done = make(chan struct{})
 
-	defer cancel(errors.WithStack(context.Canceled))
+	defer func() { cancel(errors.WithStack(context.Canceled)) }()
 	defer close(s.done)
 
 	meta := make(map[string][]string)

solver/jobs.go

@@ -626,7 +626,7 @@ func (jl *Solver) NewJob(id string) (*Job, error) {
 func (jl *Solver) Get(id string) (*Job, error) {
 	ctx, cancel := context.WithCancelCause(context.Background())
 	ctx, _ = context.WithTimeoutCause(ctx, 6*time.Second, errors.WithStack(context.DeadlineExceeded))
-	defer cancel(errors.WithStack(context.Canceled))
+	defer func() { cancel(errors.WithStack(context.Canceled)) }()
 
 	go func() {
 		<-ctx.Done()

solver/llbsolver/solver.go

@@ -204,7 +204,7 @@ func (s *Solver) recordBuildHistory(ctx context.Context, id string, req frontend
 
 		ctx, cancel := context.WithCancelCause(ctx)
 		ctx, _ = context.WithTimeoutCause(ctx, 300*time.Second, errors.WithStack(context.DeadlineExceeded))
-		defer cancel(errors.WithStack(context.Canceled))
+		defer func() { cancel(errors.WithStack(context.Canceled)) }()
 
 		var mu sync.Mutex
 		ch := make(chan *client.SolveStatus)

solver/scheduler.go

@@ -227,7 +227,7 @@ func (s *scheduler) build(ctx context.Context, edge Edge) (CachedResult, error)
 	s.mu.Unlock()
 
 	ctx, cancel := context.WithCancelCause(ctx)
-	defer cancel(errors.WithStack(context.Canceled))
+	defer func() { cancel(errors.WithStack(context.Canceled)) }()
 
 	go func() {
 		<-ctx.Done()

solver/scheduler_test.go

@@ -580,7 +580,7 @@ func TestSingleCancelParallel(t *testing.T) {
 		}()
 
 		ctx, cancel := context.WithCancelCause(ctx)
-		defer cancel(errors.WithStack(context.Canceled))
+		defer func() { cancel(errors.WithStack(context.Canceled)) }()
 
 		g := Edge{
 			Vertex: vtx(vtxOpt{

source/containerimage/ocilayout.go

@@ -126,7 +126,7 @@ func (r *ociLayoutResolver) withCaller(ctx context.Context, f func(context.Conte
 	if r.store.SessionID != "" {
 		timeoutCtx, cancel := context.WithCancelCause(ctx)
 		timeoutCtx, _ = context.WithTimeoutCause(timeoutCtx, 5*time.Second, errors.WithStack(context.DeadlineExceeded))
-		defer cancel(errors.WithStack(context.Canceled))
+		defer func() { cancel(errors.WithStack(context.Canceled)) }()
 
 		caller, err := r.sm.Get(timeoutCtx, r.store.SessionID, false)
 		if err != nil {

source/local/source.go

@@ -143,7 +143,7 @@ func (ls *localSourceHandler) Snapshot(ctx context.Context, g session.Group) (ca
 
 	timeoutCtx, cancel := context.WithCancelCause(ctx)
 	timeoutCtx, _ = context.WithTimeoutCause(timeoutCtx, 5*time.Second, errors.WithStack(context.DeadlineExceeded))
-	defer cancel(errors.WithStack(context.Canceled))
+	defer func() { cancel(errors.WithStack(context.Canceled)) }()
 
 	caller, err := ls.sm.Get(timeoutCtx, sessionID, false)
 	if err != nil {

util/flightcontrol/flightcontrol.go

@@ -118,7 +118,7 @@ func newCall[T any](fn func(ctx context.Context) (T, error)) *call[T] {
 func (c *call[T]) run() {
 	defer c.closeProgressWriter(errors.WithStack(context.Canceled))
 	ctx, cancel := context.WithCancelCause(c.ctx)
-	defer cancel(errors.WithStack(context.Canceled))
+	defer func() { cancel(errors.WithStack(context.Canceled)) }()
 	v, err := c.fn(ctx)
 	c.mu.Lock()
 	c.result = v
@@ -157,7 +157,7 @@ func (c *call[T]) wait(ctx context.Context) (v T, err error) {
 	}
 
 	ctx, cancel := context.WithCancelCause(ctx)
-	defer cancel(errors.WithStack(context.Canceled))
+	defer func() { cancel(errors.WithStack(context.Canceled)) }()
 
 	c.ctxs = append(c.ctxs, ctx)
 

util/testutil/integration/registry.go

@@ -69,7 +69,7 @@ http:
 
 	ctx, cancel := context.WithCancelCause(context.Background())
 	ctx, _ = context.WithTimeoutCause(ctx, 5*time.Second, errors.WithStack(context.DeadlineExceeded))
-	defer cancel(errors.WithStack(context.Canceled))
+	defer func() { cancel(errors.WithStack(context.Canceled)) }()
 	url, err = detectPort(ctx, rc)
 	if err != nil {
 		return "", nil, err

util/testutil/integration/run.go

@@ -198,7 +198,7 @@ func Run(t *testing.T, testCases []Test, opt ...TestOpt) {
 						defer sandboxLimiter.Release(1)
 
 						ctx, cancel := context.WithCancelCause(ctx)
-						defer cancel(errors.WithStack(context.Canceled))
+						defer func() { cancel(errors.WithStack(context.Canceled)) }()
 
 						sb, closer, err := newSandbox(ctx, br, getMirror(), mv)
 						require.NoError(t, err)

util/tracing/otlptracegrpc/client.go

@@ -70,7 +70,7 @@ func (c *client) UploadTraces(ctx context.Context, protoSpans []*tracepb.Resourc
 	}
 
 	ctx, cancel := c.connection.ContextWithStop(ctx)
-	defer cancel(errors.WithStack(context.Canceled))
+	defer func() { cancel(errors.WithStack(context.Canceled)) }()
 	ctx, tCancel := context.WithCancelCause(ctx)
 	ctx, _ = context.WithTimeoutCause(ctx, 30*time.Second, errors.WithStack(context.DeadlineExceeded))
 	defer tCancel(errors.WithStack(context.Canceled))