Use cleanup.Background instead of context.Background for cleanup

Use the cleanup context to re-use values from the original context

Signed-off-by: Derek McGowan <derek@mcg.dev>

rootfs/diff.go

@@ -22,7 +22,7 @@ import (
 
 	"github.com/containerd/containerd/diff"
 	"github.com/containerd/containerd/mount"
-	"github.com/containerd/containerd/namespaces"
+	"github.com/containerd/containerd/pkg/cleanup"
 	"github.com/containerd/containerd/snapshots"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 )
@@ -32,13 +32,6 @@ import (
 // the content creation and the provided snapshotter and mount differ are used
 // for calculating the diff. The descriptor for the layer diff is returned.
 func CreateDiff(ctx context.Context, snapshotID string, sn snapshots.Snapshotter, d diff.Comparer, opts ...diff.Opt) (ocispec.Descriptor, error) {
-	// dctx is used to handle cleanup things just in case the param ctx
-	// has been canceled, which causes that the defer cleanup fails.
-	dctx := context.Background()
-	if ns, ok := namespaces.Namespace(ctx); ok {
-		dctx = namespaces.WithNamespace(dctx, ns)
-	}
-
 	info, err := sn.Stat(ctx, snapshotID)
 	if err != nil {
 		return ocispec.Descriptor{}, err
@@ -49,7 +42,9 @@ func CreateDiff(ctx context.Context, snapshotID string, sn snapshots.Snapshotter
 	if err != nil {
 		return ocispec.Descriptor{}, err
 	}
-	defer sn.Remove(dctx, lowerKey)
+	defer cleanup.Do(ctx, func(ctx context.Context) {
+		sn.Remove(ctx, lowerKey)
+	})
 
 	var upper []mount.Mount
 	if info.Kind == snapshots.KindActive {
@@ -63,7 +58,9 @@ func CreateDiff(ctx context.Context, snapshotID string, sn snapshots.Snapshotter
 		if err != nil {
 			return ocispec.Descriptor{}, err
 		}
-		defer sn.Remove(dctx, upperKey)
+		defer cleanup.Do(ctx, func(ctx context.Context) {
+			sn.Remove(ctx, upperKey)
+		})
 	}
 
 	return d.Compare(ctx, lower, upper, opts...)