Updates lease creation to respect existing leases

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

client.go

@@ -25,7 +25,6 @@ import (
 	"github.com/containerd/containerd/diff"
 	"github.com/containerd/containerd/errdefs"
 	"github.com/containerd/containerd/images"
-	"github.com/containerd/containerd/leases"
 	"github.com/containerd/containerd/namespaces"
 	"github.com/containerd/containerd/platforms"
 	"github.com/containerd/containerd/plugin"
@@ -138,13 +137,11 @@ func (c *Client) Containers(ctx context.Context, filters ...string) ([]Container
 // NewContainer will create a new container in container with the provided id
 // the id must be unique within the namespace
 func (c *Client) NewContainer(ctx context.Context, id string, opts ...NewContainerOpts) (Container, error) {
-	l, err := c.CreateLease(ctx)
+	ctx, done, err := c.withLease(ctx)
 	if err != nil {
 		return nil, err
 	}
-	defer l.Delete(ctx)
-
-	ctx = leases.WithLease(ctx, l.ID())
+	defer done()
 
 	container := containers.Container{
 		ID: id,
@@ -221,13 +218,11 @@ func (c *Client) Pull(ctx context.Context, ref string, opts ...RemoteOpt) (Image
 	}
 	store := c.ContentStore()
 
-	l, err := c.CreateLease(ctx)
+	ctx, done, err := c.withLease(ctx)
 	if err != nil {
 		return nil, err
 	}
-	defer l.Delete(ctx)
-
-	ctx = leases.WithLease(ctx, l.ID())
+	defer done()
 
 	name, desc, err := pullCtx.Resolver.Resolve(ctx, ref)
 	if err != nil {
@@ -596,13 +591,11 @@ func (c *Client) Import(ctx context.Context, ref string, reader io.Reader, opts
 		return nil, err
 	}
 
-	l, err := c.CreateLease(ctx)
+	ctx, done, err := c.withLease(ctx)
 	if err != nil {
 		return nil, err
 	}
-	defer l.Delete(ctx)
-
-	ctx = leases.WithLease(ctx, l.ID())
+	defer done()
 
 	switch iopts.format {
 	case ociImageFormat:

lease.go

@@ -5,6 +5,7 @@ import (
 	"time"
 
 	leasesapi "github.com/containerd/containerd/api/services/leases/v1"
+	"github.com/containerd/containerd/leases"
 )
 
 // Lease is used to hold a reference to active resources which have not been
@@ -50,6 +51,25 @@ func (c *Client) ListLeases(ctx context.Context) ([]Lease, error) {
 	return leases, nil
 }
 
+func (c *Client) withLease(ctx context.Context) (context.Context, func() error, error) {
+	_, ok := leases.Lease(ctx)
+	if ok {
+		return ctx, func() error {
+			return nil
+		}, nil
+	}
+
+	l, err := c.CreateLease(ctx)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	ctx = leases.WithLease(ctx, l.ID())
+	return ctx, func() error {
+		return l.Delete(ctx)
+	}, nil
+}
+
 // ID returns the lease ID
 func (l Lease) ID() string {
 	return l.id

task.go

@@ -18,7 +18,6 @@ import (
 	"github.com/containerd/containerd/diff"
 	"github.com/containerd/containerd/errdefs"
 	"github.com/containerd/containerd/images"
-	"github.com/containerd/containerd/leases"
 	"github.com/containerd/containerd/mount"
 	"github.com/containerd/containerd/plugin"
 	"github.com/containerd/containerd/rootfs"
@@ -358,13 +357,11 @@ func (t *task) Resize(ctx context.Context, w, h uint32) error {
 }
 
 func (t *task) Checkpoint(ctx context.Context, opts ...CheckpointTaskOpts) (Image, error) {
-	l, err := t.client.CreateLease(ctx)
+	ctx, done, err := t.client.withLease(ctx)
 	if err != nil {
 		return nil, err
 	}
-	defer l.Delete(ctx)
-
-	ctx = leases.WithLease(ctx, l.ID())
+	defer done()
 
 	request := &tasks.CheckpointTaskRequest{
 		ContainerID: t.id,