c8d/builder: Lease layer snapshots

Create a lease for the snapshot and hold it until the layer is released.

Signed-off-by: Paweł Gronowski <pawel.gronowski@docker.com>

daemon/containerd/image_builder.go

@@ -2,6 +2,7 @@ package containerd
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"os"
@@ -183,7 +184,6 @@ func newROLayerForImage(ctx context.Context, imgDesc *ocispec.Descriptor, i *Ima
 		platMatcher = platforms.Only(*platform)
 	}
 
-	// this needs it's own context + lease so that it doesn't get cleaned before we're ready
 	confDesc, err := containerdimages.Config(ctx, i.client.ContentStore(), *imgDesc, platMatcher)
 	if err != nil {
 		return nil, err
@@ -194,23 +194,46 @@ func newROLayerForImage(ctx context.Context, imgDesc *ocispec.Descriptor, i *Ima
 		return nil, err
 	}
 
+	// TODO(vvoland): Check if image is unpacked, and unpack it if it's not.
 	imageSnapshotID := identity.ChainID(diffIDs).String()
 
+	snapshotter := i.StorageDriver()
+	_, lease, err := createLease(ctx, i.client.LeasesService())
+	if err != nil {
+		return nil, errdefs.System(fmt.Errorf("failed to lease image snapshot %s: %w", imageSnapshotID, err))
+	}
+
 	return &rolayer{
 		key:                imageSnapshotID,
 		c:                  i.client,
-		snapshotter:        i.snapshotter,
+		snapshotter:        snapshotter,
 		diffID:             "", // Image RO layer doesn't have a diff.
 		contentStoreDigest: "",
+		lease:              &lease,
 	}, nil
 }
 
+func createLease(ctx context.Context, lm leases.Manager) (context.Context, leases.Lease, error) {
+	lease, err := lm.Create(ctx,
+		leases.WithExpiration(time.Hour*24),
+		leases.WithLabels(map[string]string{
+			"org.mobyproject.lease.classicbuilder": "true",
+		}),
+	)
+	if err != nil {
+		return nil, leases.Lease{}, fmt.Errorf("failed to create a lease for snapshot: %w", err)
+	}
+
+	return leases.WithLease(ctx, lease.ID), lease, nil
+}
+
 type rolayer struct {
 	key                string
 	c                  *containerd.Client
 	snapshotter        string
 	diffID             digest.Digest
 	contentStoreDigest digest.Digest
+	lease              *leases.Lease
 }
 
 func (rl *rolayer) ContentStoreDigest() digest.Digest {
@@ -225,22 +248,35 @@ func (rl *rolayer) DiffID() layer.DiffID {
 }
 
 func (rl *rolayer) Release() error {
+	if rl.lease != nil {
+		lm := rl.c.LeasesService()
+		err := lm.Delete(context.TODO(), *rl.lease)
+		if err != nil {
+			return err
+		}
+		rl.lease = nil
+	}
 	return nil
 }
 
 // NewRWLayer creates a new read-write layer for the builder
-func (rl *rolayer) NewRWLayer() (builder.RWLayer, error) {
+func (rl *rolayer) NewRWLayer() (_ builder.RWLayer, outErr error) {
 	snapshotter := rl.c.SnapshotService(rl.snapshotter)
 
-	// we need this here for the prepared snapshots or
-	// we'll have racy behaviour where sometimes they
-	// will get GC'd before we commit/use them
-	ctx, _, err := rl.c.WithLease(context.TODO(), leases.WithRandomID(), leases.WithExpiration(1*time.Hour))
+	key := stringid.GenerateRandomID()
+
+	ctx, lease, err := createLease(context.TODO(), rl.c.LeasesService())
 	if err != nil {
-		return nil, fmt.Errorf("failed to create lease for commit: %w", err)
+		return nil, err
 	}
+	defer func() {
+		if outErr != nil {
+			if err := rl.c.LeasesService().Delete(ctx, lease); err != nil {
+				log.G(ctx).WithError(err).Warn("failed to remove lease after NewRWLayer error")
+			}
+		}
+	}()
 
-	key := stringid.GenerateRandomID()
 	mounts, err := snapshotter.Prepare(ctx, key, rl.key)
 	if err != nil {
 		return nil, err
@@ -260,6 +296,7 @@ func (rl *rolayer) NewRWLayer() (builder.RWLayer, error) {
 		c:           rl.c,
 		snapshotter: rl.snapshotter,
 		root:        root,
+		lease:       &lease,
 	}, nil
 }
 
@@ -269,23 +306,31 @@ type rwlayer struct {
 	c           *containerd.Client
 	snapshotter string
 	root        string
+	lease       *leases.Lease
 }
 
 func (rw *rwlayer) Root() string {
 	return rw.root
 }
 
-func (rw *rwlayer) Commit() (builder.ROLayer, error) {
-	// we need this here for the prepared snapshots or
-	// we'll have racy behaviour where sometimes they
-	// will get GC'd before we commit/use them
-	ctx, _, err := rw.c.WithLease(context.TODO(), leases.WithRandomID(), leases.WithExpiration(1*time.Hour))
-	if err != nil {
-		return nil, fmt.Errorf("failed to create lease for commit: %w", err)
-	}
+func (rw *rwlayer) Commit() (_ builder.ROLayer, outErr error) {
 	snapshotter := rw.c.SnapshotService(rw.snapshotter)
 
 	key := stringid.GenerateRandomID()
+
+	lm := rw.c.LeasesService()
+	ctx, lease, err := createLease(context.TODO(), lm)
+	if err != nil {
+		return nil, err
+	}
+	defer func() {
+		if outErr != nil {
+			if err := lm.Delete(ctx, lease); err != nil {
+				log.G(ctx).WithError(err).Warn("failed to remove lease after NewRWLayer error")
+			}
+		}
+	}()
+
 	err = snapshotter.Commit(ctx, key, rw.key)
 	if err != nil && !cerrdefs.IsAlreadyExists(err) {
 		return nil, err
@@ -315,28 +360,35 @@ func (rw *rwlayer) Commit() (builder.ROLayer, error) {
 		snapshotter:        rw.snapshotter,
 		diffID:             diffID,
 		contentStoreDigest: desc.Digest,
+		lease:              &lease,
 	}, nil
 }
 
-func (rw *rwlayer) Release() error {
-	snapshotter := rw.c.SnapshotService(rw.snapshotter)
-	err := snapshotter.Remove(context.TODO(), rw.key)
-	if err != nil && !cerrdefs.IsNotFound(err) {
-		return err
-	}
-
+func (rw *rwlayer) Release() (outErr error) {
 	if rw.root == "" { // nothing to release
 		return nil
 	}
-	if err := mount.UnmountAll(rw.root, 0); err != nil {
+
+	if err := mount.UnmountAll(rw.root, 0); err != nil && !errors.Is(err, os.ErrNotExist) {
 		log.G(context.TODO()).WithError(err).WithField("root", rw.root).Error("failed to unmount ROLayer")
 		return err
 	}
-	if err := os.Remove(rw.root); err != nil {
+	if err := os.Remove(rw.root); err != nil && !errors.Is(err, os.ErrNotExist) {
 		log.G(context.TODO()).WithError(err).WithField("dir", rw.root).Error("failed to remove mount temp dir")
 		return err
 	}
 	rw.root = ""
+
+	if rw.lease != nil {
+		lm := rw.c.LeasesService()
+		err := lm.Delete(context.TODO(), *rw.lease)
+		if err != nil {
+			log.G(context.TODO()).WithError(err).Warn("failed to delete lease when releasing RWLayer")
+		} else {
+			rw.lease = nil
+		}
+	}
+
 	return nil
 }