solver: fix handling of lazy blobs in cache export

Before this change, lazy blobs were handled during cache export by
providing descHandlers from the ref being exported in llbsolver.

However, this didn't handle some max cache export cases that involve use
of read-write mounts. Specifically, if you exported cache for a ref from
a read-write mount in an ExecOp, the ref's descHandlers didn't include
handlers for any refs from the the rootfs of the ExecOp.

If any of those refs for the rootfs involved lazy blobs, any error would
get hit during cache export about lazy blobs. It's possible for the
rootfs to have lazy blobs in a few different ways, but the one tested in
the integ test added here involves two images with layers that get
deduped by chainID (i.e. uncompress to the same layer but have different
compressions). Image layer refs that find an existing ref w/ same
chainID will get a snapshot for free but stay lazy in terms of their
blobs, thus making it possible for an exec to run on top of them while
still considered lazy.

The fix here puts the CacheOptGetter logic in the cache export code
directly so that it can use the solver's information on dependencies to
find all possible descHandlers, including those for the rootfs in the
read-write mount case.

Signed-off-by: Erik Sipsma <[email protected]>

client/client_test.go

@@ -222,6 +222,7 @@ var allTests = []func(t *testing.T, sb integration.Sandbox){
 	testLayerLimitOnMounts,
 	testFrontendVerifyPlatforms,
 	testRunValidExitCodes,
+	testSameChainIDWithLazyBlobs,
 }
 
 func TestIntegration(t *testing.T) {
@@ -10895,3 +10896,144 @@ func testRunValidExitCodes(t *testing.T, sb integration.Sandbox) {
 	require.Error(t, err)
 	require.ErrorContains(t, err, "exit code: 0")
 }
+
+func testSameChainIDWithLazyBlobs(t *testing.T, sb integration.Sandbox) {
+	workers.CheckFeatureCompat(t, sb,
+		workers.FeatureCacheExport,
+		workers.FeatureCacheImport,
+		workers.FeatureCacheBackendRegistry,
+	)
+
+	c, err := New(sb.Context(), sb.Address())
+	require.NoError(t, err)
+	defer c.Close()
+
+	registry, err := sb.NewRegistry()
+	if errors.Is(err, integration.ErrRequirements) {
+		t.Skip(err.Error())
+	}
+	require.NoError(t, err)
+
+	// push the base busybox image, ensuring it uses gzip
+
+	def, err := llb.Image("busybox:latest").
+		Marshal(sb.Context())
+	require.NoError(t, err)
+	busyboxGzipRef := registry + "/buildkit/busyboxgzip:latest"
+	_, err = c.Solve(sb.Context(), def, SolveOpt{
+		Exports: []ExportEntry{
+			{
+				Type: ExporterImage,
+				Attrs: map[string]string{
+					"name":              busyboxGzipRef,
+					"push":              "true",
+					"compression":       "gzip",
+					"force-compression": "true",
+				},
+			},
+		},
+	}, nil)
+	require.NoError(t, err)
+
+	// push the base busybox image plus an extra layer, ensuring it uses zstd
+	// the extra layer allows us to avoid edge merge later
+	def, err = llb.Image("busybox:latest").
+		Run(llb.Shlex(`touch /foo`)).Root().
+		Marshal(sb.Context())
+	require.NoError(t, err)
+	busyboxZstdRef := registry + "/buildkit/busyboxzstd:latest"
+	_, err = c.Solve(sb.Context(), def, SolveOpt{
+		Exports: []ExportEntry{
+			{
+				Type: ExporterImage,
+				Attrs: map[string]string{
+					"name":              busyboxZstdRef,
+					"push":              "true",
+					"compression":       "zstd",
+					"force-compression": "true",
+				},
+			},
+		},
+	}, nil)
+	require.NoError(t, err)
+
+	ensurePruneAll(t, c, sb)
+
+	// create non-lazy cache refs for the zstd image
+	def, err = llb.Image(busyboxZstdRef).
+		Run(llb.Shlex(`true`)).Root().
+		Marshal(sb.Context())
+	require.NoError(t, err)
+	_, err = c.Solve(sb.Context(), def, SolveOpt{}, nil)
+	require.NoError(t, err)
+
+	// Create lazy cache refs for the gzip layers that will be deduped by chainID with
+	// the zstd layers made in the previous solve.
+	// Put a random file in the rootfs, run a cache invalidation step and then copy
+	// the random file to a r/w mnt.
+	def, err = llb.Image(busyboxGzipRef).
+		Run(llb.Shlex(`sh -c "cat /dev/urandom | head -c 100 > /rand"`)).Root().
+		Run(llb.Shlex(`echo `+identity.NewID())).Root().
+		Run(llb.Shlex(`cp /rand /mnt/rand`)).AddMount("/mnt", llb.Scratch()).
+		Marshal(sb.Context())
+	require.NoError(t, err)
+
+	outDir := t.TempDir()
+	_, err = c.Solve(sb.Context(), def, SolveOpt{
+		Exports: []ExportEntry{
+			{
+				Type:      ExporterLocal,
+				OutputDir: outDir,
+			},
+		},
+		CacheExports: []CacheOptionsEntry{
+			{
+				Type: "registry",
+				Attrs: map[string]string{
+					"ref":  registry + "/buildkit/idc:latest",
+					"mode": "max",
+				},
+			},
+		},
+	}, nil)
+	require.NoError(t, err)
+
+	rand1, err := os.ReadFile(filepath.Join(outDir, "rand"))
+	require.NoError(t, err)
+
+	ensurePruneAll(t, c, sb)
+
+	// Run the same steps as before but with a different cache invalidation step in the middle
+	// The random file should still be cached from earlier and thus the output should be the same
+	def, err = llb.Image(busyboxGzipRef).
+		Run(llb.Shlex(`sh -c "cat /dev/urandom | head -c 100 > /rand"`)).Root().
+		Run(llb.Shlex(`echo `+identity.NewID())).Root().
+		Run(llb.Shlex(`cp /rand /mnt/rand`)).AddMount("/mnt", llb.Scratch()).
+		Marshal(sb.Context())
+	require.NoError(t, err)
+
+	outDir = t.TempDir()
+	_, err = c.Solve(sb.Context(), def, SolveOpt{
+		Exports: []ExportEntry{
+			{
+				Type:      ExporterLocal,
+				OutputDir: outDir,
+			},
+		},
+		CacheImports: []CacheOptionsEntry{
+			{
+				Type: "registry",
+				Attrs: map[string]string{
+					"ref":  registry + "/buildkit/idc:latest",
+					"mode": "max",
+				},
+			},
+		},
+	}, nil)
+	require.NoError(t, err)
+
+	rand2, err := os.ReadFile(filepath.Join(outDir, "rand"))
+	require.NoError(t, err)
+
+	require.Equal(t, string(rand1), string(rand2))
+}

solver/exporter.go

@@ -117,6 +117,13 @@ func (e *exporter) ExportTo(ctx context.Context, t CacheExporterTarget, opt Cach
 			return nil, err
 		}
 
+		if e.edge != nil {
+			op, ok := e.edge.op.(*sharedOp)
+			if ok && op != nil && op.st != nil {
+				ctx = withAncestorCacheOpts(ctx, op.st)
+			}
+		}
+
 		remotes, err := cm.results.LoadRemotes(ctx, res, opt.CompressionOpt, opt.Session)
 		if err != nil {
 			return nil, err

solver/llbsolver/provenance.go

@@ -398,18 +398,9 @@ func NewProvenanceCreator(ctx context.Context, cp *provenance.Capture, res solve
 			return nil, err
 		}
 
-		wref, ok := r.Sys().(*worker.WorkerRef)
-		if !ok {
-			return nil, errors.Errorf("invalid worker ref %T", r.Sys())
-		}
-
 		addLayers = func() error {
 			e := newCacheExporter()
 
-			if wref.ImmutableRef != nil {
-				ctx = withDescHandlerCacheOpts(ctx, wref.ImmutableRef)
-			}
-
 			if _, err := r.CacheKeys()[0].Exporter.ExportTo(ctx, e, solver.CacheExportOpt{
 				ResolveRemotes:  resolveRemotes,
 				Mode:            solver.CacheExportModeRemoteOnly,

solver/llbsolver/solver.go

@@ -691,8 +691,6 @@ func runCacheExporters(ctx context.Context, exporters []RemoteCacheExporter, j *
 			err = inBuilderContext(ctx, j, exp.Exporter.Name(), id, func(ctx context.Context, _ session.Group) error {
 				prepareDone := progress.OneOff(ctx, "preparing build cache for export")
 				if err := result.EachRef(cached, inp, func(res solver.CachedResult, ref cache.ImmutableRef) error {
-					ctx = withDescHandlerCacheOpts(ctx, ref)
-
 					// Configure compression
 					compressionConfig := exp.Config().Compression
 
@@ -998,7 +996,6 @@ func inlineCache(ctx context.Context, ie inlineCacheExporter, res solver.CachedR
 		digests = append(digests, desc.Digest)
 	}
 
-	ctx = withDescHandlerCacheOpts(ctx, workerRef.ImmutableRef)
 	refCfg := cacheconfig.RefConfig{Compression: compressionopt}
 	if _, err := res.CacheKeys()[0].Exporter.ExportTo(ctx, ie, solver.CacheExportOpt{
 		ResolveRemotes: workerRefResolver(refCfg, true, g), // load as many compression blobs as possible
@@ -1011,20 +1008,6 @@ func inlineCache(ctx context.Context, ie inlineCacheExporter, res solver.CachedR
 	return ie.ExportForLayers(ctx, digests)
 }
 
-func withDescHandlerCacheOpts(ctx context.Context, ref cache.ImmutableRef) context.Context {
-	return solver.WithCacheOptGetter(ctx, func(includeAncestors bool, 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 {
 	if err := s.history.Status(ctx, id, statusChan); err != nil {
 		if !errors.Is(err, os.ErrNotExist) {