Make local cache non-lazy

Signed-off-by: Kohei Tokunaga <[email protected]>

Author: ktock

cache/manager.go

@@ -301,7 +301,14 @@ func (cm *cacheManager) GetByBlob(ctx context.Context, desc ocispecs.Descriptor,
 
 	cm.records[id] = rec
 
-	return rec.ref(true, descHandlers, nil), nil
+	ref := rec.ref(true, descHandlers, nil)
+	if s := unlazySessionOf(opts...); s != nil {
+		if err := ref.unlazy(ctx, ref.descHandlers, ref.progress, s, true); err != nil {
+			return nil, err
+		}
+	}
+
+	return ref, nil
 }
 
 // init loads all snapshots from metadata state and tries to load the records

cache/opts.go

@@ -35,3 +35,14 @@ type NeedsRemoteProviderError []digest.Digest //nolint:errname
 func (m NeedsRemoteProviderError) Error() string {
 	return fmt.Sprintf("missing descriptor handlers for lazy blobs %+v", []digest.Digest(m))
 }
+
+type Unlazy session.Group
+
+func unlazySessionOf(opts ...RefOption) session.Group {
+	for _, opt := range opts {
+		if opt, ok := opt.(session.Group); ok {
+			return opt
+		}
+	}
+	return nil
+}

cache/remotecache/local/local.go

@@ -98,7 +98,16 @@ func getContentStore(ctx context.Context, sm *session.Manager, g session.Group,
 	if err != nil {
 		return nil, err
 	}
-	return sessioncontent.NewCallerStore(caller, storeID), nil
+	return &unlazyProvider{sessioncontent.NewCallerStore(caller, storeID), g}, nil
+}
+
+type unlazyProvider struct {
+	content.Store
+	s session.Group
+}
+
+func (p *unlazyProvider) UnlazySession(desc ocispecs.Descriptor) session.Group {
+	return p.s
 }
 
 func attrsToCompression(attrs map[string]string) (*compression.Config, error) {

util/contentutil/multiprovider.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/containerd/containerd/content"
 	"github.com/containerd/containerd/errdefs"
+	"github.com/moby/buildkit/session"
 	digest "github.com/opencontainers/go-digest"
 	ocispecs "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/pkg/errors"
@@ -90,3 +91,23 @@ func (mp *MultiProvider) Add(dgst digest.Digest, p content.Provider) {
 	defer mp.mu.Unlock()
 	mp.sub[dgst] = p
 }
+
+func (mp *MultiProvider) UnlazySession(desc ocispecs.Descriptor) session.Group {
+	type unlazySession interface {
+		UnlazySession(ocispecs.Descriptor) session.Group
+	}
+
+	mp.mu.RLock()
+	if p, ok := mp.sub[desc.Digest]; ok {
+		mp.mu.RUnlock()
+		if cd, ok := p.(unlazySession); ok {
+			return cd.UnlazySession(desc)
+		}
+	} else {
+		mp.mu.RUnlock()
+	}
+	if cd, ok := mp.base.(unlazySession); ok {
+		return cd.UnlazySession(desc)
+	}
+	return nil
+}

worker/base/worker.go

@@ -477,6 +477,14 @@ func (w *Worker) FromRemote(ctx context.Context, remote *solver.Remote) (ref cac
 			cache.WithCreationTime(tm),
 			descHandlers,
 		}
+		if ul, ok := remote.Provider.(interface {
+			UnlazySession(ocispecs.Descriptor) session.Group
+		}); ok {
+			s := ul.UnlazySession(desc)
+			if s != nil {
+				opts = append(opts, cache.Unlazy(s))
+			}
+		}
 		if dh, ok := descHandlers[desc.Digest]; ok {
 			if ref, ok := dh.Annotations["containerd.io/distribution.source.ref"]; ok {
 				opts = append(opts, cache.WithImageRef(ref)) // can set by registry cache importer