Fix various lazy blob errors involving dedupe by chainID

client/client_test.go

@@ -222,6 +222,8 @@
 	testLayerLimitOnMounts,
 	testFrontendVerifyPlatforms,
 	testRunValidExitCodes,
+	testSameChainIDWithLazyBlobsCacheExport,
+	testSameChainIDWithLazyBlobsCacheMountBase,
 }
 
 func TestIntegration(t *testing.T) {
@@ -4816,7 +4818,7 @@
 			},
 		},
 	}, nil)
 	require.NoError(t, err)
 
 	dgst, ok := resp.ExporterResponse[exptypes.ExporterImageDigestKey]
 	require.Equal(t, true, ok)
@@ -5012,7 +5014,7 @@
 			},
 		},
 	}, nil)
 	require.NoError(t, err)
 
 	img, err := imageService.Get(ctx, target)
 	require.NoError(t, err)
@@ -7392,15 +7394,15 @@
 }
 
 func testMergeOpCacheInline(t *testing.T, sb integration.Sandbox) {
 	testMergeOpCache(t, sb, "inline")
 }
 
 func testMergeOpCacheMin(t *testing.T, sb integration.Sandbox) {
 	testMergeOpCache(t, sb, "min")
 }
 
 func testMergeOpCacheMax(t *testing.T, sb integration.Sandbox) {
 	testMergeOpCache(t, sb, "max")
 }
 
 func testMergeOpCache(t *testing.T, sb integration.Sandbox, mode string) {
@@ -10895,3 +10897,241 @@
 	require.Error(t, err)
 	require.ErrorContains(t, err, "exit code: 0")
 }
+
+func testSameChainIDWithLazyBlobsCacheExport(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/cache-load 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))
+}
+
+func testSameChainIDWithLazyBlobsCacheMountBase(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/cache-load 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)
+
+	// use the gzip image as a cache mount base, the cache ref will be deduped by
+	// chainID with the zstd layers made in the previous solve
+	def, err = llb.Image(busyboxZstdRef).Run(
+		llb.Shlex(`touch /mnt/bar`),
+		llb.AddMount("/mnt",
+			llb.Image(busyboxGzipRef),
+			llb.AsPersistentCacheDir("idc", llb.CacheMountShared),
+		),
+	).Root().Marshal(sb.Context())
+	require.NoError(t, err)
+	_, err = c.Solve(sb.Context(), def, SolveOpt{}, nil)
+	require.NoError(t, err)
+
+	// try to re-use the cache mount from before, ensure we successfully get
+	// the same one writen to in previous step
+	def, err = llb.Image(busyboxZstdRef).Run(
+		llb.Shlex(`stat /mnt/bar`),
+		llb.AddMount("/mnt",
+			llb.Image(busyboxGzipRef),
+			llb.AsPersistentCacheDir("idc", llb.CacheMountShared),
+		),
+	).Root().Marshal(sb.Context())
+	require.NoError(t, err)
+	_, err = c.Solve(sb.Context(), def, SolveOpt{}, nil)
+	require.NoError(t, err)
+}

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
@@ -172,7 +179,7 @@ func (e *exporter) ExportTo(ctx context.Context, t CacheExporterTarget, opt Cach
 		for _, dep := range deps {
 			recs, err := dep.CacheKey.Exporter.ExportTo(ctx, t, opt)
 			if err != nil {
-				return nil, nil
+				return nil, err
 			}
 			for _, r := range recs {
 				srcs[i] = append(srcs[i], expr{r: r, selector: dep.Selector})
@@ -184,7 +191,7 @@ func (e *exporter) ExportTo(ctx context.Context, t CacheExporterTarget, opt Cach
 		for _, de := range e.edge.secondaryExporters {
 			recs, err := de.cacheKey.CacheKey.Exporter.ExportTo(ctx, t, opt)
 			if err != nil {
-				return nil, nil
+				return nil, err
 			}
 			for _, r := range recs {
 				srcs[de.index] = append(srcs[de.index], expr{r: r, selector: de.cacheKey.Selector})

solver/llbsolver/mounts/mount.go

@@ -24,6 +24,7 @@ import (
 	"github.com/moby/buildkit/util/grpcerrors"
 	"github.com/moby/locker"
 	"github.com/moby/sys/userns"
+	digest "github.com/opencontainers/go-digest"
 	"github.com/pkg/errors"
 	"google.golang.org/grpc/codes"
 )
@@ -123,7 +124,18 @@ func (g *cacheRefGetter) getRefCacheDirNoCache(ctx context.Context, key string,
 		}
 		locked := false
 		for _, si := range sis {
-			if mRef, err := g.cm.GetMutable(ctx, si.ID()); err == nil {
+			mRef, err := g.cm.GetMutable(ctx, si.ID())
+			var needsRemoteProviders cache.NeedsRemoteProviderError
+			if errors.As(err, &needsRemoteProviders) && ref != nil {
+				descHandlers := cache.DescHandlers(make(map[digest.Digest]*cache.DescHandler))
+				for _, dgst := range needsRemoteProviders {
+					if handler := ref.DescHandler(dgst); handler != nil {
+						descHandlers[dgst] = handler
+					}
+				}
+				mRef, err = g.cm.GetMutable(ctx, si.ID(), descHandlers)
+			}
+			if err == nil {
 				bklog.G(ctx).Debugf("reusing ref for cache dir %q: %s", id, mRef.ID())
 				return mRef, nil
 			} else if errors.Is(err, cache.ErrLocked) {

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) {