imageutil: keep config and manifest lease for time based cache

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

control/control.go

@@ -3,6 +3,7 @@ package control
 import (
 	"context"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	controlapi "github.com/moby/buildkit/api/services/control"
@@ -17,6 +18,7 @@ import (
 	"github.com/moby/buildkit/solver"
 	"github.com/moby/buildkit/solver/llbsolver"
 	"github.com/moby/buildkit/solver/pb"
+	"github.com/moby/buildkit/util/imageutil"
 	"github.com/moby/buildkit/util/throttle"
 	"github.com/moby/buildkit/worker"
 	"github.com/pkg/errors"
@@ -42,6 +44,7 @@ type Controller struct { // TODO: ControlService
 	gatewayForwarder *controlgateway.GatewayForwarder
 	throttledGC      func()
 	gcmu             sync.Mutex
+	buildCount       int64
 }
 
 func NewController(opt Opt) (*Controller, error) {
@@ -110,6 +113,10 @@ func (c *Controller) DiskUsage(ctx context.Context, r *controlapi.DiskUsageReque
 }
 
 func (c *Controller) Prune(req *controlapi.PruneRequest, stream controlapi.Control_PruneServer) error {
+	if atomic.LoadInt64(&c.buildCount) == 0 {
+		imageutil.CancelCacheLeases()
+	}
+
 	ch := make(chan client.UsageInfo)
 
 	eg, ctx := errgroup.WithContext(stream.Context())
@@ -207,6 +214,9 @@ func translateLegacySolveRequest(req *controlapi.SolveRequest) error {
 }
 
 func (c *Controller) Solve(ctx context.Context, req *controlapi.SolveRequest) (*controlapi.SolveResponse, error) {
+	atomic.AddInt64(&c.buildCount, 1)
+	defer atomic.AddInt64(&c.buildCount, -1)
+
 	if err := translateLegacySolveRequest(req); err != nil {
 		return nil, err
 	}

util/imageutil/config.go

@@ -3,6 +3,9 @@ package imageutil
 import (
 	"context"
 	"encoding/json"
+	"fmt"
+	"sync"
+	"time"
 
 	"github.com/containerd/containerd/content"
 	"github.com/containerd/containerd/images"
@@ -21,6 +24,18 @@ type ContentCache interface {
 	content.Provider
 }
 
+var leasesMu sync.Mutex
+var leasesF []func(context.Context) error
+
+func CancelCacheLeases() {
+	leasesMu.Lock()
+	for _, f := range leasesF {
+		f(context.TODO())
+	}
+	leasesF = nil
+	leasesMu.Unlock()
+}
+
 func Config(ctx context.Context, str string, resolver remotes.Resolver, cache ContentCache, leaseManager leases.Manager, p *specs.Platform) (digest.Digest, []byte, error) {
 	// TODO: fix buildkit to take interface instead of struct
 	var platform platforms.MatchComparer
@@ -35,12 +50,17 @@ func Config(ctx context.Context, str string, resolver remotes.Resolver, cache Co
 	}
 
 	if leaseManager != nil {
-		ctx2, done, err := leaseutil.WithLease(ctx, leaseManager)
+		ctx2, done, err := leaseutil.WithLease(ctx, leaseManager, leases.WithExpiration(5*time.Minute))
 		if err != nil {
 			return "", nil, errors.WithStack(err)
 		}
 		ctx = ctx2
-		defer done(ctx)
+		defer func() {
+			// this lease is not deleted to allow other components to access manifest/config from cache. It will be deleted after 5 min deadline or on pruning inactive builder
+			leasesMu.Lock()
+			leasesF = append(leasesF, done)
+			leasesMu.Unlock()
+		}()
 	}
 
 	desc := specs.Descriptor{
@@ -75,7 +95,7 @@ func Config(ctx context.Context, str string, resolver remotes.Resolver, cache Co
 
 	children := childrenConfigHandler(cache, platform)
 	if m, ok := cache.(content.Manager); ok {
-		children = images.SetChildrenLabels(m, children)
+		children = SetChildrenLabelsNonBlobs(m, children)
 	}
 
 	handlers := []images.Handler{
@@ -187,3 +207,39 @@ func DetectManifestBlobMediaType(dt []byte) (string, error) {
 	}
 	return images.MediaTypeDockerSchema2ManifestList, nil
 }
+
+func SetChildrenLabelsNonBlobs(manager content.Manager, f images.HandlerFunc) images.HandlerFunc {
+	return func(ctx context.Context, desc specs.Descriptor) ([]specs.Descriptor, error) {
+		children, err := f(ctx, desc)
+		if err != nil {
+			return children, err
+		}
+
+		if len(children) > 0 {
+			info := content.Info{
+				Digest: desc.Digest,
+				Labels: map[string]string{},
+			}
+			fields := []string{}
+			for i, ch := range children {
+				switch ch.MediaType {
+				case images.MediaTypeDockerSchema2Layer, images.MediaTypeDockerSchema2LayerGzip, specs.MediaTypeImageLayer, specs.MediaTypeImageLayerGzip:
+					continue
+				default:
+				}
+
+				info.Labels[fmt.Sprintf("containerd.io/gc.ref.content.%d", i)] = ch.Digest.String()
+				fields = append(fields, fmt.Sprintf("labels.containerd.io/gc.ref.content.%d", i))
+			}
+
+			if len(info.Labels) > 0 {
+				_, err := manager.Update(ctx, info, fields...)
+				if err != nil {
+					return nil, err
+				}
+			}
+		}
+
+		return children, err
+	}
+}

util/leaseutil/manager.go

@@ -10,15 +10,15 @@ import (
 	bolt "go.etcd.io/bbolt"
 )
 
-func WithLease(ctx context.Context, ls leases.Manager) (context.Context, func(context.Context) error, error) {
+func WithLease(ctx context.Context, ls leases.Manager, opts ...leases.Opt) (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))
+	l, err := ls.Create(ctx, append([]leases.Opt{leases.WithRandomID(), leases.WithExpiration(time.Hour)}, opts...)...)
 	if err != nil {
 		return nil, nil, err
 	}

util/pull/pull.go

@@ -132,7 +132,7 @@ func (p *Puller) Pull(ctx context.Context) (*Pulled, error) {
 		// Get all the children for a descriptor
 		childrenHandler := images.ChildrenHandler(p.ContentStore)
 		// Set any children labels for that content
-		childrenHandler = images.SetChildrenLabels(p.ContentStore, childrenHandler)
+		childrenHandler = imageutil.SetChildrenLabelsNonBlobs(p.ContentStore, childrenHandler)
 		// Filter the children by the platform
 		childrenHandler = images.FilterPlatforms(childrenHandler, platform)
 		// Limit manifests pulled to the best match in an index