From 3756668adbbb48cf027a2229ba560d3777ed871e Mon Sep 17 00:00:00 2001
From: Tibor Vass <tibor@docker.com>
Date: Fri, 13 Nov 2020 02:23:01 +0000
Subject: [PATCH] builder-next: Refactor using buildkit's resolver pool

Signed-off-by: Tibor Vass <tibor@docker.com>
---
 .../adapters/containerimage/pull.go           | 129 +-----------------
 1 file changed, 6 insertions(+), 123 deletions(-)

diff --git a/builder/builder-next/adapters/containerimage/pull.go b/builder/builder-next/adapters/containerimage/pull.go
index a979d1b4d6..65f7f58d5c 100644
--- a/builder/builder-next/adapters/containerimage/pull.go
+++ b/builder/builder-next/adapters/containerimage/pull.go
@@ -9,7 +9,6 @@ import (
 	"path"
 	"runtime"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/containerd/containerd/content"
@@ -36,7 +35,7 @@ import (
 	"github.com/moby/buildkit/util/flightcontrol"
 	"github.com/moby/buildkit/util/imageutil"
 	"github.com/moby/buildkit/util/progress"
-	_ "github.com/moby/buildkit/util/resolver"
+	"github.com/moby/buildkit/util/resolver"
 	digest "github.com/opencontainers/go-digest"
 	"github.com/opencontainers/image-spec/identity"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
@@ -60,18 +59,12 @@ type SourceOpt struct {
 // Source is the source implementation for accessing container images
 type Source struct {
 	SourceOpt
-	g             flightcontrol.Group
-	resolverCache *resolverCache
+	g flightcontrol.Group
 }
 
 // NewSource creates a new image source
 func NewSource(opt SourceOpt) (*Source, error) {
-	is := &Source{
-		SourceOpt:     opt,
-		resolverCache: nil, //newResolverCache(),
-	}
-
-	return is, nil
+	return &Source{SourceOpt: opt}, nil
 }
 
 // ID returns image scheme identifier
@@ -79,17 +72,6 @@ func (is *Source) ID() string {
 	return source.DockerImageScheme
 }
 
-func (is *Source) getResolver(hosts docker.RegistryHosts, ref string, sm *session.Manager, g session.Group) remotes.Resolver {
-	if res := is.resolverCache.Get(ref, g); res != nil {
-		return res
-	}
-	//auth := resolver.NewSessionAuthenticator(sm, g)
-	//r := resolver.New(hosts, auth)
-	//r = is.resolverCache.Add(ref, auth, r, g)
-	//return r
-	return nil
-}
-
 func (is *Source) resolveLocal(refStr string) (*image.Image, error) {
 	ref, err := distreference.ParseNormalizedNamed(refStr)
 	if err != nil {
@@ -112,7 +94,8 @@ func (is *Source) resolveRemote(ctx context.Context, ref string, platform *ocisp
 		dt   []byte
 	}
 	res, err := is.g.Do(ctx, ref, func(ctx context.Context) (interface{}, error) {
-		dgst, dt, err := imageutil.Config(ctx, ref, is.getResolver(is.RegistryHosts, ref, sm, g), is.ContentStore, nil, platform)
+		res := resolver.DefaultPool.GetResolver(is.RegistryHosts, ref, "pull", sm, g)
+		dgst, dt, err := imageutil.Config(ctx, ref, res, is.ContentStore, nil, platform)
 		if err != nil {
 			return nil, err
 		}
@@ -199,20 +182,13 @@ type puller struct {
 	desc             ocispec.Descriptor
 	ref              string
 	resolveErr       error
-	resolverInstance remotes.Resolver
-	resolverOnce     sync.Once
 	config           []byte
 	platform         ocispec.Platform
 	sm               *session.Manager
 }
 
 func (p *puller) resolver(g session.Group) remotes.Resolver {
-	p.resolverOnce.Do(func() {
-		if p.resolverInstance == nil {
-			p.resolverInstance = p.is.getResolver(p.is.RegistryHosts, p.src.Reference.String(), p.sm, g)
-		}
-	})
-	return p.resolverInstance
+	return resolver.DefaultPool.GetResolver(p.is.RegistryHosts, p.src.Reference.String(), "pull", p.sm, g)
 }
 
 func (p *puller) mainManifestKey(dgst digest.Digest, platform ocispec.Platform) (digest.Digest, error) {
@@ -428,10 +404,6 @@ func (p *puller) Snapshot(ctx context.Context, g session.Group) (cache.Immutable
 	}
 
 	platform := platforms.Only(p.platform)
-	// workaround for GCR bug that requires a request to manifest endpoint for authentication to work.
-	// if current resolver has not used manifests do a dummy request.
-	// in most cases resolver should be cached and extra request is not needed.
-	ensureManifestRequested(ctx, p.resolver(g), p.ref)
 
 	var (
 		schema1Converter *schema1.Converter
@@ -847,95 +819,6 @@ func resolveModeToString(rm source.ResolveMode) string {
 	return ""
 }
 
-type resolverCache struct {
-	mu sync.Mutex
-	m  map[string]cachedResolver
-}
-
-type cachedResolver struct {
-	counter int64 // needs to be 64bit aligned for 32bit systems
-	timeout time.Time
-	remotes.Resolver
-	auth *struct{}
-}
-
-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(ref string, auth *struct{}, resolver remotes.Resolver, g session.Group) *cachedResolver {
-	r.mu.Lock()
-	defer r.mu.Unlock()
-
-	ref = r.repo(ref)
-
-	cr, ok := r.m[ref]
-	cr.timeout = time.Now().Add(time.Minute)
-	if ok {
-		return &cr
-	}
-
-	cr.Resolver = resolver
-	cr.auth = auth
-	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(ref string, g session.Group) *cachedResolver {
-	r.mu.Lock()
-	defer r.mu.Unlock()
-
-	ref = r.repo(ref)
-
-	cr, ok := r.m[ref]
-	if ok {
-		return &cr
-	}
-	return nil
-}
-
-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
-}
-
-func ensureManifestRequested(ctx context.Context, res remotes.Resolver, ref string) {
-	cr, ok := res.(*cachedResolver)
-	if !ok {
-		return
-	}
-	if atomic.LoadInt64(&cr.counter) == 0 {
-		res.Resolve(ctx, ref)
-	}
-}
-
 func platformMatches(img *image.Image, p *ocispec.Platform) bool {
 	if img.Architecture != p.Architecture {
 		return false