builder: fix blobs releasing via leases after pull

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
(cherry picked from commit 5c01d06f72)
Signed-off-by: Sebastiaan van Stijn <github@gone.nl>

builder/builder-next/adapters/containerimage/pull.go

@@ -13,7 +13,9 @@ import (
 
 	"github.com/containerd/containerd/content"
 	containerderrors "github.com/containerd/containerd/errdefs"
+	"github.com/containerd/containerd/gc"
 	"github.com/containerd/containerd/images"
+	"github.com/containerd/containerd/leases"
 	"github.com/containerd/containerd/platforms"
 	ctdreference "github.com/containerd/containerd/reference"
 	"github.com/containerd/containerd/remotes"
@@ -34,6 +36,7 @@ import (
 	"github.com/moby/buildkit/source"
 	"github.com/moby/buildkit/util/flightcontrol"
 	"github.com/moby/buildkit/util/imageutil"
+	"github.com/moby/buildkit/util/leaseutil"
 	"github.com/moby/buildkit/util/progress"
 	"github.com/moby/buildkit/util/resolver"
 	digest "github.com/opencontainers/go-digest"
@@ -54,6 +57,8 @@ type SourceOpt struct {
 	ImageStore      image.Store
 	RegistryHosts   docker.RegistryHosts
 	LayerStore      layer.Store
+	LeaseManager    leases.Manager
+	GarbageCollect  func(ctx context.Context) (gc.Stats, error)
 }
 
 // Source is the source implementation for accessing container images
@@ -101,7 +106,7 @@ func (is *Source) resolveRemote(ctx context.Context, ref string, platform *ocisp
 	key := "getconfig::" + ref + "::" + platforms.Format(p)
 	res, err := is.g.Do(ctx, key, func(ctx context.Context) (interface{}, error) {
 		res := resolver.DefaultPool.GetResolver(is.RegistryHosts, ref, "pull", sm, g)
-		dgst, dt, err := imageutil.Config(ctx, ref, res, is.ContentStore, nil, platform)
+		dgst, dt, err := imageutil.Config(ctx, ref, res, is.ContentStore, is.LeaseManager, platform)
 		if err != nil {
 			return nil, err
 		}
@@ -384,6 +389,17 @@ func (p *puller) Snapshot(ctx context.Context, g session.Group) (cache.Immutable
 
 	ongoing := newJobs(p.ref)
 
+	ctx, done, err := leaseutil.WithLease(ctx, p.is.LeaseManager, leases.WithExpiration(5*time.Minute), leaseutil.MakeTemporary)
+	if err != nil {
+		return nil, err
+	}
+	defer func() {
+		done(context.TODO())
+		if p.is.GarbageCollect != nil {
+			go p.is.GarbageCollect(context.TODO())
+		}
+	}()
+
 	pctx, stopProgress := context.WithCancel(ctx)
 
 	pw, _, ctx := progress.FromContext(ctx)
@@ -406,6 +422,8 @@ func (p *puller) Snapshot(ctx context.Context, g session.Group) (cache.Immutable
 
 	platform := platforms.Only(p.platform)
 
+	var nonLayers []digest.Digest
+
 	var (
 		schema1Converter *schema1.Converter
 		handlers         []images.Handler
@@ -426,6 +444,7 @@ func (p *puller) Snapshot(ctx context.Context, g session.Group) (cache.Immutable
 			case images.MediaTypeDockerSchema2Manifest, ocispec.MediaTypeImageManifest,
 				images.MediaTypeDockerSchema2ManifestList, ocispec.MediaTypeImageIndex,
 				images.MediaTypeDockerSchema2Config, ocispec.MediaTypeImageConfig:
+				nonLayers = append(nonLayers, desc.Digest)
 			default:
 				return nil, images.ErrSkipDesc
 			}
@@ -435,8 +454,6 @@ func (p *puller) Snapshot(ctx context.Context, g session.Group) (cache.Immutable
 
 		// Get all the children for a descriptor
 		childrenHandler := images.ChildrenHandler(p.is.ContentStore)
-		// Set any children labels for that content
-		childrenHandler = images.SetChildrenLabels(p.is.ContentStore, childrenHandler)
 		// Filter the children by the platform
 		childrenHandler = images.FilterPlatforms(childrenHandler, platform)
 		// Limit manifests pulled to the best match in an index
@@ -539,9 +556,6 @@ func (p *puller) Snapshot(ctx context.Context, g session.Group) (cache.Immutable
 
 	defer func() {
 		<-progressDone
-		for _, desc := range mfst.Layers {
-			p.is.ContentStore.Delete(context.TODO(), desc.Digest)
-		}
 	}()
 
 	r := image.NewRootFS()
@@ -557,6 +571,16 @@ func (p *puller) Snapshot(ctx context.Context, g session.Group) (cache.Immutable
 		return nil, err
 	}
 
+	// keep manifest blobs until ref is alive for cache
+	for _, nl := range nonLayers {
+		if err := p.is.LeaseManager.AddResource(ctx, leases.Lease{ID: ref.ID()}, leases.Resource{
+			ID:   nl.String(),
+			Type: "content",
+		}); err != nil {
+			return nil, err
+		}
+	}
+
 	// TODO: handle windows layers for cross platform builds
 
 	if p.src.RecordType != "" && cache.GetRecordType(ref) == "" {

builder/builder-next/controller.go

@@ -111,6 +111,7 @@ func newController(rt http.RoundTripper, opt Opt) (*control.Controller, error) {
 		PruneRefChecker: refChecker,
 		LeaseManager:    lm,
 		ContentStore:    store,
+		GarbageCollect:  mdb.GarbageCollect,
 	})
 	if err != nil {
 		return nil, err
@@ -125,6 +126,8 @@ func newController(rt http.RoundTripper, opt Opt) (*control.Controller, error) {
 		ReferenceStore:  dist.ReferenceStore,
 		RegistryHosts:   opt.RegistryHosts,
 		LayerStore:      dist.LayerStore,
+		LeaseManager:    lm,
+		GarbageCollect:  mdb.GarbageCollect,
 	})
 	if err != nil {
 		return nil, err