Include DescHandler CacheOpts during export.

Before this, descriptor handlers were not included with calls to the
exporter, which then sometimes called LoadRef and failed to get a ref
because it was lazy. This change results in the DescHandlers of the
already loaded refs to get plugged into context so they can be re-used
by the exporter if it needs to load the ref again.

Signed-off-by: Erik Sipsma <erik@sipsma.dev>

cache/refs.go

@@ -34,6 +34,7 @@ type Ref interface {
 	RefMetadata
 	Release(context.Context) error
 	IdentityMapping() *idtools.IdentityMapping
+	DescHandler(digest.Digest) *DescHandler
 }
 
 type ImmutableRef interface {
@@ -313,12 +314,20 @@ type immutableRef struct {
 	descHandlers    DescHandlers
 }
 
+func (sr *immutableRef) DescHandler(dgst digest.Digest) *DescHandler {
+	return sr.descHandlers[dgst]
+}
+
 type mutableRef struct {
 	*cacheRecord
 	triggerLastUsed bool
 	descHandlers    DescHandlers
 }
 
+func (sr *mutableRef) DescHandler(dgst digest.Digest) *DescHandler {
+	return sr.descHandlers[dgst]
+}
+
 func (sr *immutableRef) Clone() ImmutableRef {
 	sr.mu.Lock()
 	ref := sr.ref(false, sr.descHandlers)

solver/cacheopts.go

@@ -21,8 +21,12 @@ func CacheOptGetterOf(ctx context.Context) func(keys ...interface{}) map[interfa
 	return nil
 }
 
+func WithCacheOptGetter(ctx context.Context, getter func(keys ...interface{}) map[interface{}]interface{}) context.Context {
+	return context.WithValue(ctx, cacheOptGetterKey{}, getter)
+}
+
 func withAncestorCacheOpts(ctx context.Context, start *state) context.Context {
-	return context.WithValue(ctx, cacheOptGetterKey{}, func(keys ...interface{}) map[interface{}]interface{} {
+	return WithCacheOptGetter(ctx, func(keys ...interface{}) map[interface{}]interface{} {
 		keySet := make(map[interface{}]struct{})
 		for _, k := range keys {
 			keySet[k] = struct{}{}

solver/llbsolver/solver.go

@@ -245,6 +245,13 @@ func (s *Solver) Solve(ctx context.Context, id string, sessionID string, req fro
 				if err != nil {
 					return err
 				}
+
+				workerRef, ok := r.Sys().(*worker.WorkerRef)
+				if !ok {
+					return errors.Errorf("invalid reference: %T", r.Sys())
+				}
+				ctx = withDescHandlerCacheOpts(ctx, workerRef.ImmutableRef)
+
 				// all keys have same export chain so exporting others is not needed
 				_, err = r.CacheKeys()[0].Exporter.ExportTo(ctx, e, solver.CacheExportOpt{
 					ResolveRemotes: workerRefResolver(solver.CompressionOpt{
@@ -308,6 +315,7 @@ func inlineCache(ctx context.Context, e remotecache.Exporter, res solver.CachedR
 			digests = append(digests, desc.Digest)
 		}
 
+		ctx = withDescHandlerCacheOpts(ctx, workerRef.ImmutableRef)
 		if _, err := res.CacheKeys()[0].Exporter.ExportTo(ctx, e, solver.CacheExportOpt{
 			ResolveRemotes: workerRefResolver(compressionopt, true, g), // load as many compression blobs as possible
 			Mode:           solver.CacheExportModeMin,
@@ -322,6 +330,20 @@ func inlineCache(ctx context.Context, e remotecache.Exporter, res solver.CachedR
 	return nil, nil
 }
 
+func withDescHandlerCacheOpts(ctx context.Context, ref cache.ImmutableRef) context.Context {
+	return solver.WithCacheOptGetter(ctx, func(keys ...interface{}) map[interface{}]interface{} {
+		vals := make(map[interface{}]interface{})
+		for _, k := range keys {
+			if key, ok := k.(cache.DescHandlerKey); ok {
+				if handler := ref.DescHandler(digest.Digest(key)); handler != nil {
+					vals[k] = handler
+				}
+			}
+		}
+		return vals
+	})
+}
+
 func (s *Solver) Status(ctx context.Context, id string, statusChan chan *client.SolveStatus) error {
 	j, err := s.solver.Get(id)
 	if err != nil {

worker/base/worker.go

@@ -232,15 +232,20 @@ func (w *Worker) LoadRef(ctx context.Context, id string, hidden bool) (cache.Imm
 			}
 			descHandlers := cache.DescHandlers(make(map[digest.Digest]*cache.DescHandler))
 			for k, v := range optGetter(keys...) {
-				if v, ok := v.(*cache.DescHandler); ok {
-					descHandlers[k.(digest.Digest)] = v
+				if key, ok := k.(cache.DescHandlerKey); ok {
+					if handler, ok := v.(*cache.DescHandler); ok {
+						descHandlers[digest.Digest(key)] = handler
+					}
 				}
 			}
 			opts = append(opts, descHandlers)
 			ref, err = w.CacheMgr.Get(ctx, id, opts...)
 		}
 	}
-	return ref, err
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to load ref")
+	}
+	return ref, nil
 }
 
 func (w *Worker) Executor() executor.Executor {