ensure context.WithoutCancel in defer funcs

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

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 {

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)

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

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 {

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.

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()
 		}()
 

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)
 		}
 	}()

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 {

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))
 		}
 	}()
 

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))
 		}
 	}()
 

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
 		}
 	}()
 

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 {

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)

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)
 		}
 	}()
 

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)
 			}
 		}()
 

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 {

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

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))
 		}
 	}()