pull: add cache for gcr bug

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

source/containerimage/pull.go

@@ -182,6 +182,9 @@ func (p *puller) Snapshot(ctx context.Context) (cache.ImmutableRef, error) {
 		}
 	}
 
+	// workaround for gcr, authentication not supported on blob endpoints
+	pull.EnsureManifestRequested(ctx, p.Puller.Resolver, p.Puller.Src.String())
+
 	pulled, err := p.Puller.Pull(ctx)
 	if err != nil {
 		return nil, err

util/pull/resolver.go

@@ -3,11 +3,14 @@ package pull
 import (
 	"context"
 	"net/http"
+	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/containerd/containerd/images"
 	"github.com/containerd/containerd/remotes"
 	"github.com/containerd/containerd/remotes/docker"
+	distreference "github.com/docker/distribution/reference"
 	"github.com/moby/buildkit/session"
 	"github.com/moby/buildkit/session/auth"
 	"github.com/moby/buildkit/source"
@@ -16,7 +19,17 @@ import (
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 )
 
+var cache *resolverCache
+
+func init() {
+	cache = newResolverCache()
+}
+
 func NewResolver(ctx context.Context, rfn resolver.ResolveOptionsFunc, sm *session.Manager, imageStore images.Store, mode source.ResolveMode, ref string) remotes.Resolver {
+	if res := cache.Get(ctx, ref); res != nil {
+		return withLocal(res, imageStore, mode)
+	}
+
 	opt := docker.ResolverOptions{
 		Client: http.DefaultClient,
 	}
@@ -26,12 +39,35 @@ func NewResolver(ctx context.Context, rfn resolver.ResolveOptionsFunc, sm *sessi
 	opt.Credentials = getCredentialsFromSession(ctx, sm)
 
 	r := docker.NewResolver(opt)
+	r = cache.Add(ctx, ref, r)
 
+	return withLocal(r, imageStore, mode)
+}
+
+func EnsureManifestRequested(ctx context.Context, res remotes.Resolver, ref string) {
+	rr := res
+	lr, ok := res.(withLocalResolver)
+	if ok {
+		if atomic.LoadInt64(&lr.counter) > 0 {
+			return
+		}
+		rr = lr.Resolver
+	}
+	cr, ok := rr.(*cachedResolver)
+	if !ok {
+		return
+	}
+	if atomic.LoadInt64(&cr.counter) == 0 {
+		res.Resolve(ctx, ref)
+	}
+}
+
+func withLocal(r remotes.Resolver, imageStore images.Store, mode source.ResolveMode) remotes.Resolver {
 	if imageStore == nil || mode == source.ResolveModeForcePull {
 		return r
 	}
 
-	return withLocalResolver{r, imageStore, mode}
+	return withLocalResolver{r, imageStore, mode, 0}
 }
 
 func getCredentialsFromSession(ctx context.Context, sm *session.Manager) func(string) (string, string, error) {
@@ -66,13 +102,15 @@ func getCredentialsFromSession(ctx context.Context, sm *session.Manager) func(st
 
 type withLocalResolver struct {
 	remotes.Resolver
-	is   images.Store
-	mode source.ResolveMode
+	is      images.Store
+	mode    source.ResolveMode
+	counter int64
 }
 
 func (r withLocalResolver) Resolve(ctx context.Context, ref string) (string, ocispec.Descriptor, error) {
 	if r.mode == source.ResolveModePreferLocal {
 		if img, err := r.is.Get(ctx, ref); err == nil {
+			atomic.AddInt64(&r.counter, 1)
 			return ref, img.Target, nil
 		}
 	}
@@ -90,3 +128,80 @@ func (r withLocalResolver) Resolve(ctx context.Context, ref string) (string, oci
 
 	return "", ocispec.Descriptor{}, err
 }
+
+type resolverCache struct {
+	mu sync.Mutex
+	m  map[string]cachedResolver
+}
+
+type cachedResolver struct {
+	timeout time.Time
+	remotes.Resolver
+	counter int64
+}
+
+func (cr *cachedResolver) Resolve(ctx context.Context, ref string) (name string, desc ocispec.Descriptor, err error) {
+	atomic.AddInt64(&cr.counter, 1)
+	return cr.Resolver.Resolve(ctx, ref)
+}
+
+func (r *resolverCache) Add(ctx context.Context, ref string, resolver remotes.Resolver) remotes.Resolver {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	ref = r.repo(ref) + "-" + session.FromContext(ctx)
+
+	cr, ok := r.m[ref]
+	cr.timeout = time.Now().Add(time.Minute)
+	if ok {
+		return &cr
+	}
+
+	cr.Resolver = resolver
+	r.m[ref] = cr
+	return &cr
+}
+
+func (r *resolverCache) repo(refStr string) string {
+	ref, err := distreference.ParseNormalizedNamed(refStr)
+	if err != nil {
+		return refStr
+	}
+	return ref.Name()
+}
+
+func (r *resolverCache) Get(ctx context.Context, ref string) remotes.Resolver {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	ref = r.repo(ref) + "-" + session.FromContext(ctx)
+
+	cr, ok := r.m[ref]
+	if !ok {
+		return nil
+	}
+	return &cr
+}
+
+func (r *resolverCache) clean(now time.Time) {
+	r.mu.Lock()
+	for k, cr := range r.m {
+		if now.After(cr.timeout) {
+			delete(r.m, k)
+		}
+	}
+	r.mu.Unlock()
+}
+
+func newResolverCache() *resolverCache {
+	rc := &resolverCache{
+		m: map[string]cachedResolver{},
+	}
+	t := time.NewTicker(time.Minute)
+	go func() {
+		for {
+			rc.clean(<-t.C)
+		}
+	}()
+	return rc
+}