diff --git a/client.go b/client.go
index 7aebd1ced..77a54f02d 100644
--- a/client.go
+++ b/client.go
@@ -199,7 +199,7 @@ func (c *Client) NewContainer(ctx context.Context, id string, opts ...NewContain
 	if err != nil {
 		return nil, err
 	}
-	defer done()
+	defer done(ctx)
 
 	container := containers.Container{
 		ID: id,
@@ -284,7 +284,7 @@ func (c *Client) Pull(ctx context.Context, ref string, opts ...RemoteOpt) (Image
 	if err != nil {
 		return nil, err
 	}
-	defer done()
+	defer done(ctx)
 
 	name, desc, err := pullCtx.Resolver.Resolve(ctx, ref)
 	if err != nil {
@@ -561,7 +561,7 @@ func (c *Client) Import(ctx context.Context, importer images.Importer, reader io
 	if err != nil {
 		return nil, err
 	}
-	defer done()
+	defer done(ctx)
 
 	imgrecs, err := importer.Import(ctx, c.ContentStore(), reader)
 	if err != nil {
diff --git a/cmd/ctr/commands/snapshots/snapshots.go b/cmd/ctr/commands/snapshots/snapshots.go
index ec3bc7566..9fd79f264 100644
--- a/cmd/ctr/commands/snapshots/snapshots.go
+++ b/cmd/ctr/commands/snapshots/snapshots.go
@@ -127,7 +127,7 @@ var diffCommand = cli.Command{
 		if err != nil {
 			return err
 		}
-		defer done()
+		defer done(ctx)
 
 		var desc ocispec.Descriptor
 		labels := commands.LabelArgs(context.StringSlice("label"))
diff --git a/content/testsuite/testsuite.go b/content/testsuite/testsuite.go
index 0f77f59b2..29c203238 100644
--- a/content/testsuite/testsuite.go
+++ b/content/testsuite/testsuite.go
@@ -55,7 +55,7 @@ func ContentSuite(t *testing.T, name string, storeFn func(ctx context.Context, r
 // ContextWrapper is used to decorate new context used inside the test
 // before using the context on the content store.
 // This can be used to support leasing and multiple namespaces tests.
-type ContextWrapper func(ctx context.Context) (context.Context, func() error, error)
+type ContextWrapper func(ctx context.Context) (context.Context, func(context.Context) error, error)
 
 type wrapperKey struct{}
 
@@ -98,13 +98,13 @@ func makeTest(t *testing.T, name string, storeFn func(ctx context.Context, root
 
 		w, ok := ctx.Value(wrapperKey{}).(ContextWrapper)
 		if ok {
-			var done func() error
+			var done func(context.Context) error
 			ctx, done, err = w(ctx)
 			if err != nil {
 				t.Fatalf("Error wrapping context: %+v", err)
 			}
 			defer func() {
-				if err := done(); err != nil && !t.Failed() {
+				if err := done(ctx); err != nil && !t.Failed() {
 					t.Fatalf("Wrapper release failed: %+v", err)
 				}
 			}()
@@ -542,7 +542,7 @@ func checkCrossNSShare(ctx context.Context, t *testing.T, cs content.Store) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	defer done()
+	defer done(ctx2)
 
 	w, err := cs.Writer(ctx2, ref, size, d)
 	if err != nil {
@@ -593,7 +593,7 @@ func checkCrossNSAppend(ctx context.Context, t *testing.T, cs content.Store) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	defer done()
+	defer done(ctx2)
 
 	extra := []byte("appended bytes")
 	size2 := size + int64(len(extra))
diff --git a/content_test.go b/content_test.go
index c7112e683..e72bd9d70 100644
--- a/content_test.go
+++ b/content_test.go
@@ -41,7 +41,7 @@ func newContentStore(ctx context.Context, root string) (context.Context, content
 		name  = testsuite.Name(ctx)
 	)
 
-	wrap := func(ctx context.Context) (context.Context, func() error, error) {
+	wrap := func(ctx context.Context) (context.Context, func(context.Context) error, error) {
 		n := atomic.AddUint64(&count, 1)
 		ctx = namespaces.WithNamespace(ctx, fmt.Sprintf("%s-n%d", name, n))
 		return client.WithLease(ctx)
diff --git a/image.go b/image.go
index 5ae6db232..1af706c7f 100644
--- a/image.go
+++ b/image.go
@@ -108,7 +108,7 @@ func (i *image) Unpack(ctx context.Context, snapshotterName string) error {
 	if err != nil {
 		return err
 	}
-	defer done()
+	defer done(ctx)
 
 	layers, err := i.getLayers(ctx, platforms.Default())
 	if err != nil {
diff --git a/lease.go b/lease.go
index 5fc7833f8..8cf3e5879 100644
--- a/lease.go
+++ b/lease.go
@@ -68,10 +68,10 @@ func (c *Client) ListLeases(ctx context.Context) ([]Lease, error) {
 }
 
 // WithLease attaches a lease on the context
-func (c *Client) WithLease(ctx context.Context) (context.Context, func() error, error) {
+func (c *Client) WithLease(ctx context.Context) (context.Context, func(context.Context) error, error) {
 	_, ok := leases.Lease(ctx)
 	if ok {
-		return ctx, func() error {
+		return ctx, func(context.Context) error {
 			return nil
 		}, nil
 	}
@@ -82,7 +82,7 @@ func (c *Client) WithLease(ctx context.Context) (context.Context, func() error,
 	}
 
 	ctx = leases.WithLease(ctx, l.ID())
-	return ctx, func() error {
+	return ctx, func(ctx context.Context) error {
 		return l.Delete(ctx)
 	}, nil
 }
diff --git a/metadata/content_test.go b/metadata/content_test.go
index 1ba4a027d..272f08040 100644
--- a/metadata/content_test.go
+++ b/metadata/content_test.go
@@ -51,9 +51,11 @@ func createContentStore(ctx context.Context, root string) (context.Context, cont
 		count uint64
 		name  = testsuite.Name(ctx)
 	)
-	wrap := func(ctx context.Context) (context.Context, func() error, error) {
+	wrap := func(ctx context.Context) (context.Context, func(context.Context) error, error) {
 		n := atomic.AddUint64(&count, 1)
-		return namespaces.WithNamespace(ctx, fmt.Sprintf("%s-n%d", name, n)), func() error { return nil }, nil
+		return namespaces.WithNamespace(ctx, fmt.Sprintf("%s-n%d", name, n)), func(context.Context) error {
+			return nil
+		}, nil
 	}
 	ctx = testsuite.SetContextWrapper(ctx, wrap)
 
diff --git a/task.go b/task.go
index 42740b518..a8f0e1f73 100644
--- a/task.go
+++ b/task.go
@@ -391,7 +391,7 @@ func (t *task) Checkpoint(ctx context.Context, opts ...CheckpointTaskOpts) (Imag
 	if err != nil {
 		return nil, err
 	}
-	defer done()
+	defer done(ctx)
 
 	request := &tasks.CheckpointTaskRequest{
 		ContainerID: t.id,
diff --git a/vendor.conf b/vendor.conf
index e5947ba30..69db3d76b 100644
--- a/vendor.conf
+++ b/vendor.conf
@@ -43,7 +43,7 @@ github.com/gotestyourself/gotestyourself 44dbf532bbf5767611f6f2a61bded572e337010
 github.com/google/go-cmp v0.1.0
 
 # cri dependencies
-github.com/containerd/cri 0c876040681ebe8a291fa2cebefdcc2796fa3fc8
+github.com/containerd/cri fd18145c4b01fffff53cbf350012abe7ff83ebe9 https://github.com/dmcgowan/cri-containerd
 github.com/blang/semver v3.1.0
 github.com/containernetworking/cni v0.6.0
 github.com/containernetworking/plugins v0.6.0
diff --git a/vendor/github.com/containerd/cri/pkg/containerd/importer/importer.go b/vendor/github.com/containerd/cri/pkg/containerd/importer/importer.go
index b97b593c1..6fc7a1991 100644
--- a/vendor/github.com/containerd/cri/pkg/containerd/importer/importer.go
+++ b/vendor/github.com/containerd/cri/pkg/containerd/importer/importer.go
@@ -82,7 +82,7 @@ func Import(ctx context.Context, client *containerd.Client, reader io.Reader) (_
 		return nil, err
 	}
 	// TODO(random-liu): Fix this after containerd client is fixed (containerd/containerd#2193)
-	defer done() // nolint: errcheck
+	defer done(ctx) // nolint: errcheck
 
 	cs := client.ContentStore()
 	is := client.ImageService()