ensure context.WithoutCancel in defer funcs

tonistiigi · tonistiigi · commit 4103099d94ed · 2024-06-12T19:18:32.000-07:00
Signed-off-by: Tonis Tiigi &lt;tonistiigi@gmail.com&gt;
diff --git a/cache/blobs_linux.go b/cache/blobs_linux.go
@@ -46,12 +46,13 @@ func (sr *immutableRef) tryComputeOverlayBlob(ctx context.Context, lower, upper
 
 	defer func() {
 		if cw != nil {
+			ctx = context.WithoutCancel(ctx)
 			// after commit success cw will be set to nil, if cw isn't nil, error
 			// happened before commit, we should abort this ingest, and because the
 			// error may incured by ctx cancel, use a new context here. And since
 			// cm.Close will unlock this ref in the content store, we invoke abort
 			// to remove the ingest root in advance.
-			if aerr := sr.cm.ContentStore.Abort(context.Background(), ref); aerr != nil {
+			if aerr := sr.cm.ContentStore.Abort(ctx, ref); aerr != nil {
 				bklog.G(ctx).WithError(aerr).Warnf("failed to abort writer %q", ref)
 			}
 			if cerr := cw.Close(); cerr != nil {
diff --git a/cache/manager.go b/cache/manager.go
@@ -168,7 +168,7 @@ func (cm *cacheManager) GetByBlob(ctx context.Context, desc ocispecs.Descriptor,
 	releaseParent := false
 	defer func() {
 		if releaseParent || rerr != nil && p != nil {
-			p.Release(context.TODO())
+			p.Release(context.WithoutCancel(ctx))
 		}
 	}()
 
@@ -242,7 +242,8 @@ func (cm *cacheManager) GetByBlob(ctx context.Context, desc ocispecs.Descriptor,
 
 	defer func() {
 		if rerr != nil {
-			if err := cm.LeaseManager.Delete(context.TODO(), leases.Lease{
+			ctx := context.WithoutCancel(ctx)
+			if err := cm.LeaseManager.Delete(ctx, leases.Lease{
 				ID: l.ID,
 			}); err != nil {
 				bklog.G(ctx).Errorf("failed to remove lease: %+v", err)
@@ -423,7 +424,7 @@ func (cm *cacheManager) getRecord(ctx context.Context, id string, opts ...RefOpt
 	}
 	defer func() {
 		if retErr != nil {
-			parents.release(context.TODO())
+			parents.release(context.WithoutCancel(ctx))
 		}
 	}()
 
@@ -513,7 +514,7 @@ func (cm *cacheManager) getRecord(ctx context.Context, id string, opts ...RefOpt
 func (cm *cacheManager) parentsOf(ctx context.Context, md *cacheMetadata, opts ...RefOption) (ps parentRefs, rerr error) {
 	defer func() {
 		if rerr != nil {
-			ps.release(context.TODO())
+			ps.release(context.WithoutCancel(ctx))
 		}
 	}()
 	if parentID := md.getParent(); parentID != "" {
@@ -580,7 +581,7 @@ func (cm *cacheManager) New(ctx context.Context, s ImmutableRef, sess session.Gr
 
 	defer func() {
 		if err != nil && parent != nil {
-			parent.Release(context.TODO())
+			parent.Release(context.WithoutCancel(ctx))
 		}
 	}()
 
@@ -597,7 +598,8 @@ func (cm *cacheManager) New(ctx context.Context, s ImmutableRef, sess session.Gr
 
 	defer func() {
 		if err != nil {
-			if err := cm.LeaseManager.Delete(context.TODO(), leases.Lease{
+			ctx := context.WithoutCancel(ctx)
+			if err := cm.LeaseManager.Delete(ctx, leases.Lease{
 				ID: l.ID,
 			}); err != nil {
 				bklog.G(ctx).Errorf("failed to remove lease: %+v", err)
@@ -705,7 +707,7 @@ func (cm *cacheManager) Merge(ctx context.Context, inputParents []ImmutableRef,
 	dhs := make(map[digest.Digest]*DescHandler)
 	defer func() {
 		if rerr != nil {
-			parents.release(context.TODO())
+			parents.release(context.WithoutCancel(ctx))
 		}
 	}()
 	for _, inputParent := range inputParents {
@@ -798,7 +800,8 @@ func (cm *cacheManager) createMergeRef(ctx context.Context, parents parentRefs,
 	}
 	defer func() {
 		if rerr != nil {
-			if err := cm.LeaseManager.Delete(context.TODO(), leases.Lease{
+			ctx := context.WithoutCancel(ctx)
+			if err := cm.LeaseManager.Delete(ctx, leases.Lease{
 				ID: l.ID,
 			}); err != nil {
 				bklog.G(ctx).Errorf("failed to remove lease: %+v", err)
@@ -834,7 +837,7 @@ func (cm *cacheManager) Diff(ctx context.Context, lower, upper ImmutableRef, pg
 	dhs := make(map[digest.Digest]*DescHandler)
 	defer func() {
 		if rerr != nil {
-			parents.release(context.TODO())
+			parents.release(context.WithoutCancel(ctx))
 		}
 	}()
 	for i, inputParent := range []ImmutableRef{lower, upper} {
@@ -889,7 +892,7 @@ func (cm *cacheManager) Diff(ctx context.Context, lower, upper ImmutableRef, pg
 			mergeParents := parentRefs{mergeParents: make([]*immutableRef, len(upperLayers)-len(lowerLayers))}
 			defer func() {
 				if rerr != nil {
-					mergeParents.release(context.TODO())
+					mergeParents.release(context.WithoutCancel(ctx))
 				}
 			}()
 			for i := len(lowerLayers); i < len(upperLayers); i++ {
@@ -954,7 +957,8 @@ func (cm *cacheManager) createDiffRef(ctx context.Context, parents parentRefs, d
 	}
 	defer func() {
 		if rerr != nil {
-			if err := cm.LeaseManager.Delete(context.TODO(), leases.Lease{
+			ctx := context.WithoutCancel(ctx)
+			if err := cm.LeaseManager.Delete(ctx, leases.Lease{
 				ID: l.ID,
 			}); err != nil {
 				bklog.G(ctx).Errorf("failed to remove lease: %+v", err)
diff --git a/cache/manager_test.go b/cache/manager_test.go
@@ -1168,7 +1168,7 @@ func TestLoopLeaseContent(t *testing.T) {
 	allRefs := []ImmutableRef{ref}
 	defer func() {
 		for _, ref := range allRefs {
-			ref.Release(ctx)
+			ref.Release(context.WithoutCancel(ctx))
 		}
 	}()
 	var chain []ocispecs.Descriptor
diff --git a/cache/refs.go b/cache/refs.go
@@ -1045,6 +1045,7 @@ func (sr *immutableRef) withRemoteSnapshotLabelsStargzMode(ctx context.Context,
 			return errors.Wrapf(err, "failed to add tmp remote labels for remote snapshot")
 		}
 		defer func() {
+			ctx := context.WithoutCancel(ctx)
 			for k := range info.Labels {
 				info.Labels[k] = "" // Remove labels appended in this call
 			}
@@ -1106,6 +1107,7 @@ func (sr *immutableRef) prepareRemoteSnapshotsStargzMode(ctx context.Context, s
 					info, err := r.cm.Snapshotter.Stat(ctx, snapshotID)
 					if err == nil { // usable as remote snapshot without unlazying.
 						defer func() {
+							ctx := context.WithoutCancel(ctx)
 							// Remove tmp labels appended in this func
 							if info.Labels != nil {
 								for k := range tmpLabels {
diff --git a/cache/remote.go b/cache/remote.go
@@ -39,7 +39,7 @@ func (sr *immutableRef) GetRemotes(ctx context.Context, createIfNeeded bool, ref
 	if err != nil {
 		return nil, err
 	}
-	defer done(ctx)
+	defer done(context.WithoutCancel(ctx))
 
 	// fast path if compression variants aren't required
 	// NOTE: compressionopt is applied only to *newly created layers* if Force != true.
diff --git a/client/build_test.go b/client/build_test.go
@@ -432,13 +432,13 @@ func testClientGatewayContainerExecPipe(t *testing.T, sb integration.Sandbox) {
 			Args: []string{"sleep", "10"},
 		})
 		if err != nil {
-			ctr.Release(ctx)
+			ctr.Release(context.WithoutCancel(ctx))
 			return nil, err
 		}
 
 		defer func() {
 			// cancel pid1
-			ctr.Release(ctx)
+			ctr.Release(context.WithoutCancel(ctx))
 			pid1.Wait()
 		}()
 
diff --git a/executor/containerdexecutor/executor.go b/executor/containerdexecutor/executor.go
@@ -166,7 +166,7 @@ func (w *containerdExecutor) Run(ctx context.Context, id string, root executor.M
 	}
 
 	defer func() {
-		if err1 := container.Delete(context.TODO()); err == nil && err1 != nil {
+		if err1 := container.Delete(context.WithoutCancel(ctx)); err == nil && err1 != nil {
 			err = errors.Wrapf(err1, "failed to delete container %s", id)
 		}
 	}()
@@ -190,7 +190,7 @@ func (w *containerdExecutor) Run(ctx context.Context, id string, root executor.M
 	}
 
 	defer func() {
-		if _, err1 := task.Delete(context.TODO(), containerd.WithProcessKill); err == nil && err1 != nil {
+		if _, err1 := task.Delete(context.WithoutCancel(ctx), containerd.WithProcessKill); err == nil && err1 != nil {
 			err = errors.Wrapf(err1, "failed to delete task %s", id)
 		}
 	}()
diff --git a/executor/runcexecutor/executor.go b/executor/runcexecutor/executor.go
@@ -237,7 +237,7 @@ func (w *runcExecutor) Run(ctx context.Context, id string, root executor.Mount,
 	}
 	defer mount.Unmount(rootFSPath, 0)
 
-	defer executor.MountStubsCleaner(ctx, rootFSPath, mounts, meta.RemoveMountStubsRecursive)()
+	defer executor.MountStubsCleaner(context.WithoutCancel(ctx), rootFSPath, mounts, meta.RemoveMountStubsRecursive)()
 
 	uid, gid, sgids, err := oci.GetUser(rootFSPath, meta.User)
 	if err != nil {
diff --git a/exporter/containerimage/export.go b/exporter/containerimage/export.go
@@ -227,7 +227,7 @@ func (e *imageExporterInstance) Export(ctx context.Context, src *exporter.Source
 	}
 	defer func() {
 		if descref == nil {
-			done(context.TODO())
+			done(context.WithoutCancel(ctx))
 		}
 	}()
 
diff --git a/exporter/oci/export.go b/exporter/oci/export.go
@@ -156,7 +156,7 @@ func (e *imageExporterInstance) Export(ctx context.Context, src *exporter.Source
 	}
 	defer func() {
 		if descref == nil {
-			done(context.TODO())
+			done(context.WithoutCancel(ctx))
 		}
 	}()
 
diff --git a/frontend/gateway/gateway.go b/frontend/gateway/gateway.go
@@ -154,8 +154,9 @@ func (gf *gatewayFrontend) Solve(ctx context.Context, llbBridge frontend.Fronten
 			return nil, err
 		}
 		defer func() {
+			ctx := context.WithoutCancel(ctx)
 			devRes.EachRef(func(ref solver.ResultProxy) error {
-				return ref.Release(context.TODO())
+				return ref.Release(ctx)
 			})
 		}()
 		if devRes.Ref == nil {
@@ -247,8 +248,9 @@ func (gf *gatewayFrontend) Solve(ctx context.Context, llbBridge frontend.Fronten
 			return nil, err
 		}
 		defer func() {
+			ctx := context.WithoutCancel(ctx)
 			res.EachRef(func(ref solver.ResultProxy) error {
-				return ref.Release(context.TODO())
+				return ref.Release(ctx)
 			})
 		}()
 		if res.Ref == nil {
@@ -1155,7 +1157,7 @@ func (lbf *llbBridgeForwarder) NewContainer(ctx context.Context, in *pb.NewConta
 	}
 	defer func() {
 		if err != nil {
-			ctr.Release(ctx) // ensure release on error
+			ctr.Release(context.WithoutCancel(ctx)) // ensure release on error
 		}
 	}()
 
diff --git a/solver/combinedcache.go b/solver/combinedcache.go
@@ -73,11 +73,12 @@ func (cm *combinedCacheManager) Load(ctx context.Context, rec *CacheRecord) (res
 		return nil, err
 	}
 	defer func() {
+		ctx := context.WithoutCancel(ctx)
 		for i, res := range results {
 			if err == nil && i == 0 {
 				continue
 			}
-			res.Result.Release(context.TODO())
+			res.Result.Release(ctx)
 		}
 	}()
 	if rec.cacheManager != cm.main && cm.main != nil {
diff --git a/solver/llbsolver/file/refmanager.go b/solver/llbsolver/file/refmanager.go
@@ -46,10 +46,11 @@ func (rm *RefManager) Prepare(ctx context.Context, ref fileoptypes.Ref, readonly
 	}
 	defer func() {
 		if rerr != nil {
+			ctx := context.WithoutCancel(ctx)
 			if err := mr.SetCachePolicyDefault(); err != nil {
 				bklog.G(ctx).Errorf("failed to reset FileOp mutable ref cachepolicy: %v", err)
 			}
-			mr.Release(context.TODO())
+			mr.Release(ctx)
 		}
 	}()
 	m, err := mr.Mount(ctx, readonly, g)
diff --git a/solver/llbsolver/history.go b/solver/llbsolver/history.go
@@ -133,7 +133,7 @@ func (h *HistoryQueue) migrateV2() error {
 		if err != nil {
 			return err
 		}
-		defer release(ctx)
+		defer release(context.WithoutCancel(ctx))
 		return b.ForEach(func(key, dt []byte) error {
 			recs, err := h.opt.LeaseManager.ListResources(ctx, leases.Lease{ID: h.leaseID(string(key))})
 			if err != nil {
@@ -518,7 +518,7 @@ func (h *HistoryQueue) update(ctx context.Context, rec controlapi.BuildHistoryRe
 
 		defer func() {
 			if err != nil && created {
-				h.hLeaseManager.Delete(ctx, l)
+				h.hLeaseManager.Delete(context.WithoutCancel(ctx), l)
 			}
 		}()
 
@@ -598,7 +598,7 @@ func (h *HistoryQueue) OpenBlobWriter(ctx context.Context, mt string) (_ *Writer
 
 	defer func() {
 		if err != nil {
-			h.hLeaseManager.Delete(ctx, l)
+			h.hLeaseManager.Delete(context.WithoutCancel(ctx), l)
 		}
 	}()
 
diff --git a/solver/llbsolver/ops/file.go b/solver/llbsolver/ops/file.go
@@ -354,7 +354,7 @@ func (s *FileOpSolver) Solve(ctx context.Context, inputs []fileoptypes.Ref, acti
 	defer func() {
 		for _, in := range s.ins {
 			if in.ref == nil && in.mount != nil {
-				in.mount.Release(context.TODO())
+				in.mount.Release(context.WithoutCancel(ctx))
 			}
 		}
 	}()
@@ -432,6 +432,7 @@ func (s *FileOpSolver) getInput(ctx context.Context, idx int, inputs []fileoptyp
 
 		defer func() {
 			if err != nil && inpMount != nil {
+				ctx := context.WithoutCancel(ctx)
 				inputRes := make([]solver.Result, len(inputs))
 				for i, input := range inputs {
 					inputRes[i] = worker.NewWorkerRefResult(input.(cache.ImmutableRef), s.w)
@@ -458,8 +459,9 @@ func (s *FileOpSolver) getInput(ctx context.Context, idx int, inputs []fileoptyp
 
 				err = errdefs.WithExecErrorWithContext(ctx, err, inputRes, outputRes)
 			}
+			ctx := context.WithoutCancel(ctx)
 			for _, m := range toRelease {
-				m.Release(context.TODO())
+				m.Release(ctx)
 			}
 		}()
 
diff --git a/solver/llbsolver/solver.go b/solver/llbsolver/solver.go
@@ -506,7 +506,7 @@ func (s *Solver) Solve(ctx context.Context, id string, sessionID string, req fro
 		if err := s.gatewayForwarder.RegisterBuild(ctx, id, fwd); err != nil {
 			return nil, err
 		}
-		defer s.gatewayForwarder.UnregisterBuild(ctx, id)
+		defer s.gatewayForwarder.UnregisterBuild(context.WithoutCancel(ctx), id)
 	}
 
 	if !internal {
diff --git a/source/containerimage/pull.go b/source/containerimage/pull.go
@@ -218,14 +218,14 @@ func (p *puller) Snapshot(ctx context.Context, g session.Group) (ir cache.Immuta
 	}
 	defer func() {
 		if p.releaseTmpLeases != nil {
-			p.releaseTmpLeases(context.TODO())
+			p.releaseTmpLeases(context.WithoutCancel(ctx))
 		}
 	}()
 
 	var current cache.ImmutableRef
 	defer func() {
 		if err != nil && current != nil {
-			current.Release(context.TODO())
+			current.Release(context.WithoutCancel(ctx))
 		}
 	}()
 
@@ -255,7 +255,7 @@ func (p *puller) Snapshot(ctx context.Context, g session.Group) (ir cache.Immuta
 			if err != nil {
 				return nil, err
 			}
-			defer done(ctx)
+			defer done(context.WithoutCancel(ctx))
 
 			if _, err := p.PullManifests(ctx, getResolver); err != nil {
 				return nil, err
diff --git a/source/git/source.go b/source/git/source.go
@@ -497,7 +497,7 @@ func (gs *gitSourceHandler) Snapshot(ctx context.Context, g session.Group) (out
 
 	defer func() {
 		if retErr != nil && checkoutRef != nil {
-			checkoutRef.Release(context.TODO())
+			checkoutRef.Release(context.WithoutCancel(ctx))
 		}
 	}()
 
@@ -646,7 +646,7 @@ func (gs *gitSourceHandler) Snapshot(ctx context.Context, g session.Group) (out
 
 	defer func() {
 		if retErr != nil {
-			snap.Release(context.TODO())
+			snap.Release(context.WithoutCancel(ctx))
 		}
 	}()
 
diff --git a/source/http/source_test.go b/source/http/source_test.go
diff --git a/source/local/source.go b/source/local/source.go

Original file line number	Diff line number	Diff line change
`@@ -1168,7 +1168,7 @@ func TestLoopLeaseContent(t *testing.T) {`
`1168`	`1168`	`allRefs := []ImmutableRef{ref}`
`1169`	`1169`	`defer func() {`
`1170`	`1170`	`for _, ref := range allRefs {`
`1171`		`- ref.Release(ctx)`
	`1171`	`+ ref.Release(context.WithoutCancel(ctx))`
`1172`	`1172`	`}`
`1173`	`1173`	`}()`
`1174`	`1174`	`var chain []ocispecs.Descriptor`
Original file line number	Diff line number	Diff line change
`@@ -39,7 +39,7 @@ func (sr *immutableRef) GetRemotes(ctx context.Context, createIfNeeded bool, ref`
`39`	`39`	`if err != nil {`
`40`	`40`	`return nil, err`
`41`	`41`	`}`
`42`		`- defer done(ctx)`
	`42`	`+ defer done(context.WithoutCancel(ctx))`
`43`	`43`
`44`	`44`	`// fast path if compression variants aren't required`
`45`	`45`	`// NOTE: compressionopt is applied only to newly created layers if Force != true.`
Original file line number	Diff line number	Diff line change
`@@ -166,7 +166,7 @@ func (w *containerdExecutor) Run(ctx context.Context, id string, root executor.M`
`166`	`166`	`}`
`167`	`167`
`168`	`168`	`defer func() {`
`169`		`- if err1 := container.Delete(context.TODO()); err == nil && err1 != nil {`
	`169`	`+ if err1 := container.Delete(context.WithoutCancel(ctx)); err == nil && err1 != nil {`
`170`	`170`	`err = errors.Wrapf(err1, "failed to delete container %s", id)`
`171`	`171`	`}`
`172`	`172`	`}()`
`@@ -190,7 +190,7 @@ func (w *containerdExecutor) Run(ctx context.Context, id string, root executor.M`
`190`	`190`	`}`
`191`	`191`
`192`	`192`	`defer func() {`
`193`		`- if _, err1 := task.Delete(context.TODO(), containerd.WithProcessKill); err == nil && err1 != nil {`
	`193`	`+ if _, err1 := task.Delete(context.WithoutCancel(ctx), containerd.WithProcessKill); err == nil && err1 != nil {`
`194`	`194`	`err = errors.Wrapf(err1, "failed to delete task %s", id)`
`195`	`195`	`}`
`196`	`196`	`}()`
Original file line number	Diff line number	Diff line change
`@@ -237,7 +237,7 @@ func (w *runcExecutor) Run(ctx context.Context, id string, root executor.Mount,`
`237`	`237`	`}`
`238`	`238`	`defer mount.Unmount(rootFSPath, 0)`
`239`	`239`
`240`		`- defer executor.MountStubsCleaner(ctx, rootFSPath, mounts, meta.RemoveMountStubsRecursive)()`
	`240`	`+ defer executor.MountStubsCleaner(context.WithoutCancel(ctx), rootFSPath, mounts, meta.RemoveMountStubsRecursive)()`
`241`	`241`
`242`	`242`	`uid, gid, sgids, err := oci.GetUser(rootFSPath, meta.User)`
`243`	`243`	`if err != nil {`
Original file line number	Diff line number	Diff line change
`@@ -227,7 +227,7 @@ func (e imageExporterInstance) Export(ctx context.Context, src exporter.Source`
`227`	`227`	`}`
`228`	`228`	`defer func() {`
`229`	`229`	`if descref == nil {`
`230`		`- done(context.TODO())`
	`230`	`+ done(context.WithoutCancel(ctx))`
`231`	`231`	`}`
`232`	`232`	`}()`
`233`	`233`
Original file line number	Diff line number	Diff line change
`@@ -156,7 +156,7 @@ func (e imageExporterInstance) Export(ctx context.Context, src exporter.Source`
`156`	`156`	`}`
`157`	`157`	`defer func() {`
`158`	`158`	`if descref == nil {`
`159`		`- done(context.TODO())`
	`159`	`+ done(context.WithoutCancel(ctx))`
`160`	`160`	`}`
`161`	`161`	`}()`
`162`	`162`