pull: fix content blobs deletion on pull race

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>

snapshot/blobmapping/snapshotter.go

@@ -2,6 +2,7 @@ package blobmapping
 
 import (
 	"context"
+	"time"
 
 	"github.com/containerd/containerd/content"
 	"github.com/containerd/containerd/snapshots"
@@ -122,6 +123,16 @@ func (s *Snapshotter) SetBlob(ctx context.Context, key string, diffID, blobsum d
 			return err
 		}
 	}
+	// update gc.root cause blob might be held by lease only
+	if _, err := s.opt.Content.Update(ctx, content.Info{
+		Digest: blobsum,
+		Labels: map[string]string{
+			"containerd.io/gc.root": time.Now().UTC().Format(time.RFC3339Nano),
+		},
+	}, "labels.containerd.io/gc.root"); err != nil {
+		return err
+	}
+
 	md, _ := s.opt.MetadataStore.Get(key)
 
 	v, err := metadata.NewValue(DiffPair{DiffID: diffID, Blobsum: blobsum})

snapshot/containerd/content.go

@@ -100,7 +100,17 @@ func (c *nsContent) writer(ctx context.Context, retries int, opts ...content.Wri
 			}
 		}
 	}
-	return w, err
+	return &nsWriter{Writer: w, ns: c.ns}, err
+}
+
+type nsWriter struct {
+	content.Writer
+	ns string
+}
+
+func (w *nsWriter) Commit(ctx context.Context, size int64, expected digest.Digest, opts ...content.Opt) error {
+	ctx = namespaces.WithNamespace(ctx, w.ns)
+	return w.Writer.Commit(ctx, size, expected, opts...)
 }
 
 type noGCContentStore struct {

source/containerimage/pull.go

@@ -8,6 +8,7 @@ import (
 	"github.com/containerd/containerd/content"
 	"github.com/containerd/containerd/diff"
 	"github.com/containerd/containerd/images"
+	"github.com/containerd/containerd/leases"
 	"github.com/containerd/containerd/platforms"
 	"github.com/docker/distribution/reference"
 	"github.com/moby/buildkit/cache"
@@ -36,6 +37,7 @@ type SourceOpt struct {
 	CacheAccessor cache.Accessor
 	ImageStore    images.Store // optional
 	ResolverOpt   resolver.ResolveOptionsFunc
+	LeaseManager  leases.Manager
 }
 
 type imageSource struct {
@@ -102,6 +104,7 @@ func (is *imageSource) Resolve(ctx context.Context, id source.Identifier, sm *se
 		Src:          imageIdentifier.Reference,
 		Resolver:     pull.NewResolver(ctx, is.ResolverOpt, sm, is.ImageStore, imageIdentifier.ResolveMode, imageIdentifier.Reference.String()),
 		Platform:     &platform,
+		LeaseManager: is.LeaseManager,
 	}
 	p := &puller{
 		CacheAccessor: is.CacheAccessor,

util/leaseutil/manager.go

@@ -0,0 +1,104 @@
+package leaseutil
+
+import (
+	"context"
+	"time"
+
+	"github.com/containerd/containerd/leases"
+	"github.com/containerd/containerd/metadata"
+	"github.com/containerd/containerd/namespaces"
+	bolt "go.etcd.io/bbolt"
+)
+
+func WithLease(ctx context.Context, ls leases.Manager) (context.Context, func(context.Context) error, error) {
+	_, ok := leases.FromContext(ctx)
+	if ok {
+		return ctx, func(context.Context) error {
+			return nil
+		}, nil
+	}
+
+	l, err := ls.Create(ctx, leases.WithRandomID(), leases.WithExpiration(time.Hour))
+	if err != nil {
+		return nil, nil, err
+	}
+
+	ctx = leases.WithLease(ctx, l.ID)
+	return ctx, func(ctx context.Context) error {
+		return ls.Delete(ctx, l)
+	}, nil
+}
+
+func NewManager(mdb *metadata.DB) leases.Manager {
+	return &local{db: mdb}
+}
+
+type local struct {
+	db *metadata.DB
+}
+
+func (l *local) Create(ctx context.Context, opts ...leases.Opt) (leases.Lease, error) {
+	var lease leases.Lease
+	if err := l.db.Update(func(tx *bolt.Tx) error {
+		var err error
+		lease, err = metadata.NewLeaseManager(tx).Create(ctx, opts...)
+		return err
+	}); err != nil {
+		return leases.Lease{}, err
+	}
+	return lease, nil
+}
+
+func (l *local) Delete(ctx context.Context, lease leases.Lease, opts ...leases.DeleteOpt) error {
+	var do leases.DeleteOptions
+	for _, opt := range opts {
+		if err := opt(ctx, &do); err != nil {
+			return err
+		}
+	}
+
+	if err := l.db.Update(func(tx *bolt.Tx) error {
+		return metadata.NewLeaseManager(tx).Delete(ctx, lease)
+	}); err != nil {
+		return err
+	}
+
+	return nil
+
+}
+
+func (l *local) List(ctx context.Context, filters ...string) ([]leases.Lease, error) {
+	var ll []leases.Lease
+	if err := l.db.View(func(tx *bolt.Tx) error {
+		var err error
+		ll, err = metadata.NewLeaseManager(tx).List(ctx, filters...)
+		return err
+	}); err != nil {
+		return nil, err
+	}
+	return ll, nil
+}
+
+func WithNamespace(lm leases.Manager, ns string) leases.Manager {
+	return &nsLM{Manager: lm, ns: ns}
+}
+
+type nsLM struct {
+	leases.Manager
+	ns string
+}
+
+func (l *nsLM) Create(ctx context.Context, opts ...leases.Opt) (leases.Lease, error) {
+	ctx = namespaces.WithNamespace(ctx, l.ns)
+	return l.Manager.Create(ctx, opts...)
+}
+
+func (l *nsLM) Delete(ctx context.Context, lease leases.Lease, opts ...leases.DeleteOpt) error {
+	ctx = namespaces.WithNamespace(ctx, l.ns)
+	return l.Manager.Delete(ctx, lease, opts...)
+}
+
+func (l *nsLM) List(ctx context.Context, filters ...string) ([]leases.Lease, error) {
+	ctx = namespaces.WithNamespace(ctx, l.ns)
+	return l.Manager.List(ctx, filters...)
+}

util/pull/pull.go

@@ -9,6 +9,7 @@ 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/platforms"
 	"github.com/containerd/containerd/reference"
 	"github.com/containerd/containerd/remotes"
@@ -17,6 +18,7 @@ import (
 	ctdsnapshot "github.com/containerd/containerd/snapshots"
 	"github.com/moby/buildkit/snapshot"
 	"github.com/moby/buildkit/util/imageutil"
+	"github.com/moby/buildkit/util/leaseutil"
 	"github.com/moby/buildkit/util/progress"
 	digest "github.com/opencontainers/go-digest"
 	"github.com/opencontainers/image-spec/identity"
@@ -30,6 +32,7 @@ type Puller struct {
 	Applier      diff.Applier
 	Src          reference.Spec
 	Platform     *ocispec.Platform
+	LeaseManager leases.Manager
 	// See NewResolver()
 	Resolver    remotes.Resolver
 	resolveOnce sync.Once
@@ -94,6 +97,12 @@ func (p *Puller) Pull(ctx context.Context) (*Pulled, error) {
 		platform = platforms.Default()
 	}
 
+	ctx, done, err := leaseutil.WithLease(ctx, p.LeaseManager)
+	if err != nil {
+		return nil, err
+	}
+	defer done(ctx)
+
 	ongoing := newJobs(p.ref)
 
 	pctx, stopProgress := context.WithCancel(ctx)

vendor/modules.txt

@@ -52,6 +52,8 @@ github.com/containerd/containerd/api/services/content/v1
 github.com/containerd/containerd/content/proxy
 github.com/containerd/containerd/services/content/contentserver
 github.com/containerd/containerd/reference
+github.com/containerd/containerd/leases
+github.com/containerd/containerd/metadata
 github.com/containerd/containerd/remotes/docker/schema1
 github.com/containerd/containerd/images/archive
 github.com/containerd/containerd/archive
@@ -60,7 +62,6 @@ github.com/containerd/containerd/log
 github.com/containerd/containerd/api/services/introspection/v1
 github.com/containerd/containerd/diff/apply
 github.com/containerd/containerd/diff/walking
-github.com/containerd/containerd/metadata
 github.com/containerd/containerd/labels
 github.com/containerd/containerd/version
 github.com/containerd/containerd/api/services/containers/v1
@@ -75,7 +76,6 @@ github.com/containerd/containerd/api/services/version/v1
 github.com/containerd/containerd/api/types
 github.com/containerd/containerd/defaults
 github.com/containerd/containerd/events
-github.com/containerd/containerd/leases
 github.com/containerd/containerd/leases/proxy
 github.com/containerd/containerd/pkg/dialer
 github.com/containerd/containerd/plugin

worker/base/worker.go

@@ -12,6 +12,7 @@ import (
 	"github.com/containerd/containerd/content"
 	"github.com/containerd/containerd/diff"
 	"github.com/containerd/containerd/images"
+	"github.com/containerd/containerd/leases"
 	"github.com/containerd/containerd/rootfs"
 	cdsnapshot "github.com/containerd/containerd/snapshots"
 	"github.com/docker/docker/pkg/idtools"
@@ -72,6 +73,7 @@ type WorkerOpt struct {
 	ImageStore         images.Store // optional
 	ResolveOptionsFunc resolver.ResolveOptionsFunc
 	IdentityMapping    *idtools.IdentityMapping
+	LeaseManager       leases.Manager
 }
 
 // Worker is a local worker instance with dedicated snapshotter, cache, and so on.
@@ -113,6 +115,7 @@ func NewWorker(opt WorkerOpt) (*Worker, error) {
 		ImageStore:    opt.ImageStore,
 		CacheAccessor: cm,
 		ResolverOpt:   opt.ResolveOptionsFunc,
+		LeaseManager:  opt.LeaseManager,
 	})
 	if err != nil {
 		return nil, err

worker/containerd/containerd.go

@@ -15,6 +15,7 @@ import (
 	"github.com/moby/buildkit/executor/containerdexecutor"
 	"github.com/moby/buildkit/identity"
 	containerdsnapshot "github.com/moby/buildkit/snapshot/containerd"
+	"github.com/moby/buildkit/util/leaseutil"
 	"github.com/moby/buildkit/util/network"
 	"github.com/moby/buildkit/util/throttle"
 	"github.com/moby/buildkit/util/winlayers"
@@ -112,6 +113,7 @@ func newContainerd(root string, client *containerd.Client, snapshotterName, ns s
 		Differ:        winlayers.NewWalkingDiffWithWindows(cs, df),
 		ImageStore:    client.ImageService(),
 		Platforms:     platforms,
+		LeaseManager:  leaseutil.WithNamespace(client.LeasesService(), ns),
 	}
 	return opt, nil
 }

worker/runc/runc.go

@@ -17,6 +17,7 @@ import (
 	"github.com/moby/buildkit/executor/oci"
 	"github.com/moby/buildkit/executor/runcexecutor"
 	containerdsnapshot "github.com/moby/buildkit/snapshot/containerd"
+	"github.com/moby/buildkit/util/leaseutil"
 	"github.com/moby/buildkit/util/network"
 	"github.com/moby/buildkit/util/throttle"
 	"github.com/moby/buildkit/util/winlayers"
@@ -110,6 +111,7 @@ func NewWorkerOpt(root string, snFactory SnapshotterFactory, rootless bool, proc
 		ImageStore:      nil, // explicitly
 		Platforms:       []specs.Platform{platforms.Normalize(platforms.DefaultSpec())},
 		IdentityMapping: idmap,
+		LeaseManager:    leaseutil.WithNamespace(leaseutil.NewManager(mdb), "buildkit"),
 	}
 	return opt, nil
 }