Pass in context to lease done function in client

Allows the client to choose the context to finish the lease.
This allows the client to switch contexts when the main context
used to the create the lease may have been cancelled.

Signed-off-by: Derek McGowan <derek@mcgstyle.net>

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 {

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"))

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

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)

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 {

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
 }

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)
 

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,

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

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