cache: Fix walkBlob to call callback only for unlazied contents

cache/manager.go

@@ -143,7 +143,7 @@ func (cm *cacheManager) GetByBlob(ctx context.Context, desc ocispecs.Descriptor,
 	descHandlers := descHandlersOf(opts...)
 	if desc.Digest != "" && (descHandlers == nil || descHandlers[desc.Digest] == nil) {
 		if _, err := cm.ContentStore.Info(ctx, desc.Digest); errors.Is(err, cerrdefs.ErrNotFound) {
-			return nil, NeedsRemoteProviderError([]digest.Digest{desc.Digest})
+			return nil, NeedsRemoteProviderError([]DigestDescriptionPair{{Digest: desc.Digest}})
 		} else if err != nil {
 			return nil, err
 		}
@@ -396,7 +396,7 @@ func (cm *cacheManager) getRecord(ctx context.Context, id string, opts ...RefOpt
 			if isLazy, err := cr.isLazy(ctx); err != nil {
 				return err
 			} else if isLazy && dhs[blob] == nil {
-				missing = append(missing, blob)
+				missing = append(missing, DigestDescriptionPair{Digest: blob, Description: cr.GetDescription()})
 			}
 			return nil
 		}); err != nil {

cache/opts.go

@@ -30,10 +30,19 @@ func descHandlersOf(opts ...RefOption) DescHandlers {
 
 type DescHandlerKey digest.Digest
 
-type NeedsRemoteProviderError []digest.Digest //nolint:errname
+type NeedsRemoteProviderError []DigestDescriptionPair //nolint:errname
+
+type DigestDescriptionPair struct {
+	Digest      digest.Digest
+	Description string
+}
+
+func (d DigestDescriptionPair) String() string {
+	return fmt.Sprintf("%s: %s", d.Digest, d.Description)
+}
 
 func (m NeedsRemoteProviderError) Error() string {
-	return fmt.Sprintf("missing descriptor handlers for lazy blobs %+v", []digest.Digest(m))
+	return fmt.Sprintf("missing descriptor handlers for lazy blobs %+v", []DigestDescriptionPair(m))
 }
 
 type Unlazy session.Group

cache/refs.go

@@ -49,6 +49,7 @@ type Ref interface {
 	RefMetadata
 	Release(context.Context) error
 	IdentityMapping() *idtools.IdentityMapping
+	DescHandlers() DescHandlers
 	DescHandler(digest.Digest) *DescHandler
 }
 
@@ -612,6 +613,13 @@ func (sr *immutableRef) LayerChain() RefList {
 	return l
 }
 
+func (sr *immutableRef) DescHandlers() DescHandlers {
+	// clone to prevent mutation of internal state
+	dhs := make(DescHandlers)
+	maps.Copy(dhs, sr.descHandlers)
+	return dhs
+}
+
 func (sr *immutableRef) DescHandler(dgst digest.Digest) *DescHandler {
 	return sr.descHandlers[dgst]
 }
@@ -640,6 +648,13 @@ func (sr *mutableRef) traceLogFields() logrus.Fields {
 	return m
 }
 
+func (sr *mutableRef) DescHandlers() DescHandlers {
+	// clone to prevent mutation of internal state
+	dhs := make(DescHandlers)
+	maps.Copy(dhs, sr.descHandlers)
+	return dhs
+}
+
 func (sr *mutableRef) DescHandler(dgst digest.Digest) *DescHandler {
 	return sr.descHandlers[dgst]
 }
@@ -833,6 +848,11 @@ func getBlobWithCompression(ctx context.Context, cs content.Store, desc ocispecs
 }
 
 func walkBlob(ctx context.Context, cs content.Store, desc ocispecs.Descriptor, f func(ocispecs.Descriptor) bool) error {
+	if _, err := cs.Info(ctx, desc.Digest); errors.Is(err, cerrdefs.ErrNotFound) {
+		return nil // this blob doesn't exist in the content store. Don't call the callback.
+	} else if err != nil {
+		return err
+	}
 	if !f(desc) {
 		return nil
 	}

client/client_test.go

@@ -222,6 +222,8 @@ var allTests = []func(t *testing.T, sb integration.Sandbox){
 	testLayerLimitOnMounts,
 	testFrontendVerifyPlatforms,
 	testRunValidExitCodes,
+	testSameChainIDWithLazyBlobsCacheExport,
+	testSameChainIDWithLazyBlobsCacheMountBase,
 }
 
 func TestIntegration(t *testing.T) {
@@ -10895,3 +10897,243 @@ func testRunValidExitCodes(t *testing.T, sb integration.Sandbox) {
 	require.Error(t, err)
 	require.ErrorContains(t, err, "exit code: 0")
 }
+
+func testSameChainIDWithLazyBlobsCacheExport(t *testing.T, sb integration.Sandbox) {
+	requiresLinux(t)
+	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) {
+	requiresLinux(t)
+	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 written 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)
+}