Fix cache cannot reuse lazy layers

cache/remotecache/v1/cachestorage.go

@@ -291,7 +291,7 @@ func (cs *cacheResultStorage) LoadRemotes(ctx context.Context, res solver.CacheR
 	return nil, errors.WithStack(solver.ErrNotFound)
 }
 
-func (cs *cacheResultStorage) Exists(id string) bool {
+func (cs *cacheResultStorage) Exists(ctx context.Context, id string) bool {
 	return cs.byResultID(id) != nil
 }
 

client/client_test.go

@@ -3796,18 +3796,19 @@ func testStargzLazyRegistryCacheImportExport(t *testing.T, sb integration.Sandbo
 
 	// stargz layers should be lazy even for executing something on them
 	def, err = baseDef.
-		Run(llb.Args([]string{"/bin/touch", "/bar"})).
+		Run(llb.Args([]string{"sh", "-c", "cat /dev/urandom | head -c 100 | sha256sum > unique"})).
 		Marshal(sb.Context())
 	require.NoError(t, err)
 	target := registry + "/buildkit/testlazyimage:" + identity.NewID()
-	_, err = c.Solve(sb.Context(), def, SolveOpt{
+	resp, err := c.Solve(sb.Context(), def, SolveOpt{
 		Exports: []ExportEntry{
 			{
 				Type: ExporterImage,
 				Attrs: map[string]string{
 					"name":                                   target,
 					"push":                                   "true",
 					"store":                                  "true",
+					"oci-mediatypes":                         "true",
 					"unsafe-internal-store-allow-incomplete": "true",
 				},
 			},
@@ -3820,9 +3821,25 @@ func testStargzLazyRegistryCacheImportExport(t *testing.T, sb integration.Sandbo
 				},
 			},
 		},
+		CacheExports: []CacheOptionsEntry{
+			{
+				Type: "registry",
+				Attrs: map[string]string{
+					"ref":            sgzCache,
+					"compression":    "estargz",
+					"oci-mediatypes": "true",
+				},
+			},
+		},
 	}, nil)
 	require.NoError(t, err)
 
+	dgst, ok := resp.ExporterResponse[exptypes.ExporterImageDigestKey]
+	require.Equal(t, ok, true)
+
+	unique, err := readFileInImage(sb.Context(), t, c, target+"@"+dgst, "/unique")
+	require.NoError(t, err)
+
 	img, err := imageService.Get(ctx, target)
 	require.NoError(t, err)
 
@@ -3843,6 +3860,40 @@ func testStargzLazyRegistryCacheImportExport(t *testing.T, sb integration.Sandbo
 	_, err = contentStore.Info(ctx, manifest.Layers[len(manifest.Layers)-1].Digest)
 	require.NoError(t, err)
 
+	// Run build again and check if cache is reused
+	resp, err = c.Solve(sb.Context(), def, SolveOpt{
+		Exports: []ExportEntry{
+			{
+				Type: ExporterImage,
+				Attrs: map[string]string{
+					"name":                                   target,
+					"push":                                   "true",
+					"store":                                  "true",
+					"oci-mediatypes":                         "true",
+					"unsafe-internal-store-allow-incomplete": "true",
+				},
+			},
+		},
+		CacheImports: []CacheOptionsEntry{
+			{
+				Type: "registry",
+				Attrs: map[string]string{
+					"ref": sgzCache,
+				},
+			},
+		},
+	}, nil)
+	require.NoError(t, err)
+
+	dgst2, ok := resp.ExporterResponse[exptypes.ExporterImageDigestKey]
+	require.Equal(t, ok, true)
+
+	unique2, err := readFileInImage(sb.Context(), t, c, target+"@"+dgst2, "/unique")
+	require.NoError(t, err)
+
+	require.Equal(t, dgst, dgst2)
+	require.EqualValues(t, unique, unique2)
+
 	// clear all local state out
 	err = imageService.Delete(ctx, img.Name, images.SynchronousDelete())
 	require.NoError(t, err)
@@ -4086,11 +4137,11 @@ func testStargzLazyPull(t *testing.T, sb integration.Sandbox) {
 
 	// stargz layers should be lazy even for executing something on them
 	def, err = llb.Image(sgzImage).
-		Run(llb.Args([]string{"/bin/touch", "/foo"})).
+		Run(llb.Args([]string{"sh", "-c", "cat /dev/urandom | head -c 100 | sha256sum > unique"})).
 		Marshal(sb.Context())
 	require.NoError(t, err)
 	target := registry + "/buildkit/testlazyimage:" + identity.NewID()
-	_, err = c.Solve(sb.Context(), def, SolveOpt{
+	resp, err := c.Solve(sb.Context(), def, SolveOpt{
 		Exports: []ExportEntry{
 			{
 				Type: ExporterImage,
@@ -4106,6 +4157,12 @@ func testStargzLazyPull(t *testing.T, sb integration.Sandbox) {
 	}, nil)
 	require.NoError(t, err)
 
+	dgst, ok := resp.ExporterResponse[exptypes.ExporterImageDigestKey]
+	require.Equal(t, ok, true)
+
+	unique, err := readFileInImage(sb.Context(), t, c, target+"@"+dgst, "/unique")
+	require.NoError(t, err)
+
 	img, err := imageService.Get(ctx, target)
 	require.NoError(t, err)
 
@@ -4126,6 +4183,32 @@ func testStargzLazyPull(t *testing.T, sb integration.Sandbox) {
 	_, err = contentStore.Info(ctx, manifest.Layers[len(manifest.Layers)-1].Digest)
 	require.NoError(t, err)
 
+	// Run build again and check if cache is reused
+	resp, err = c.Solve(sb.Context(), def, SolveOpt{
+		Exports: []ExportEntry{
+			{
+				Type: ExporterImage,
+				Attrs: map[string]string{
+					"name":                                   target,
+					"push":                                   "true",
+					"store":                                  "true",
+					"oci-mediatypes":                         "true",
+					"unsafe-internal-store-allow-incomplete": "true",
+				},
+			},
+		},
+	}, nil)
+	require.NoError(t, err)
+
+	dgst2, ok := resp.ExporterResponse[exptypes.ExporterImageDigestKey]
+	require.Equal(t, ok, true)
+
+	unique2, err := readFileInImage(sb.Context(), t, c, target+"@"+dgst2, "/unique")
+	require.NoError(t, err)
+
+	require.Equal(t, dgst, dgst2)
+	require.EqualValues(t, unique, unique2)
+
 	// clear all local state out
 	err = imageService.Delete(ctx, img.Name, images.SynchronousDelete())
 	require.NoError(t, err)

control/control.go

@@ -181,9 +181,9 @@ func (c *Controller) Prune(req *controlapi.PruneRequest, stream controlapi.Contr
 	defer func() {
 		if didPrune {
 			if c, ok := c.cache.(interface {
-				ReleaseUnreferenced() error
+				ReleaseUnreferenced(context.Context) error
 			}); ok {
-				if err := c.ReleaseUnreferenced(); err != nil {
+				if err := c.ReleaseUnreferenced(ctx); err != nil {
 					bklog.G(ctx).Errorf("failed to release cache metadata: %+v", err)
 				}
 			}

solver/cache_test.go

@@ -49,7 +49,7 @@ func TestInMemoryCache(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, len(keys), 1)
 
-	matches, err := m.Records(keys[0])
+	matches, err := m.Records(context.TODO(), keys[0])
 	require.NoError(t, err)
 	require.Equal(t, len(matches), 1)
 
@@ -65,7 +65,7 @@ func TestInMemoryCache(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, len(keys), 1)
 
-	matches, err = m.Records(keys[0])
+	matches, err = m.Records(context.TODO(), keys[0])
 	require.NoError(t, err)
 	require.Equal(t, len(matches), 1)
 
@@ -99,7 +99,7 @@ func TestInMemoryCache(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, len(keys), 1)
 
-	matches, err = m.Records(keys[0])
+	matches, err = m.Records(context.TODO(), keys[0])
 	require.NoError(t, err)
 	require.Equal(t, len(matches), 1)
 
@@ -121,7 +121,7 @@ func TestInMemoryCache(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, len(keys), 1)
 
-	matches, err = m.Records(keys[0])
+	matches, err = m.Records(context.TODO(), keys[0])
 	require.NoError(t, err)
 	require.Equal(t, len(matches), 2)
 
@@ -175,7 +175,7 @@ func TestInMemoryCacheSelector(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, len(keys), 1)
 
-	matches, err := m.Records(keys[0])
+	matches, err := m.Records(context.TODO(), keys[0])
 	require.NoError(t, err)
 	require.Equal(t, len(matches), 1)
 
@@ -203,7 +203,7 @@ func TestInMemoryCacheSelectorNested(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, len(keys), 1)
 
-	matches, err := m.Records(keys[0])
+	matches, err := m.Records(context.TODO(), keys[0])
 	require.NoError(t, err)
 	require.Equal(t, len(matches), 1)
 
@@ -223,7 +223,7 @@ func TestInMemoryCacheSelectorNested(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, len(keys), 1)
 
-	matches, err = m.Records(keys[0])
+	matches, err = m.Records(context.TODO(), keys[0])
 	require.NoError(t, err)
 	require.Equal(t, len(matches), 1)
 
@@ -253,7 +253,7 @@ func TestInMemoryCacheReleaseParent(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, len(keys), 1)
 
-	matches, err := m.Records(keys[0])
+	matches, err := m.Records(context.TODO(), keys[0])
 	require.NoError(t, err)
 	require.Equal(t, len(matches), 1)
 
@@ -265,15 +265,15 @@ func TestInMemoryCacheReleaseParent(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, len(keys), 1)
 
-	matches, err = m.Records(keys[0])
+	matches, err = m.Records(context.TODO(), keys[0])
 	require.NoError(t, err)
 	require.Equal(t, len(matches), 0)
 
 	keys, err = m.Query(depKeys(expKey(keys[0])), 0, dgst("bar"), 0)
 	require.NoError(t, err)
 	require.Equal(t, len(keys), 1)
 
-	matches, err = m.Records(keys[0])
+	matches, err = m.Records(context.TODO(), keys[0])
 	require.NoError(t, err)
 	require.Equal(t, len(matches), 1)
 
@@ -311,15 +311,15 @@ func TestInMemoryCacheRestoreOfflineDeletion(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, len(keys), 1)
 
-	matches, err := m.Records(keys[0])
+	matches, err := m.Records(context.TODO(), keys[0])
 	require.NoError(t, err)
 	require.Equal(t, len(matches), 0)
 
 	keys, err = m.Query(depKeys(expKey(keys[0])), 0, dgst("bar"), 0)
 	require.NoError(t, err)
 	require.Equal(t, len(keys), 1)
 
-	matches, err = m.Records(keys[0])
+	matches, err = m.Records(context.TODO(), keys[0])
 	require.NoError(t, err)
 	require.Equal(t, len(matches), 1)
 }

solver/cachemanager.go

@@ -25,7 +25,7 @@ func NewCacheManager(ctx context.Context, id string, storage CacheKeyStorage, re
 		results: results,
 	}
 
-	if err := cm.ReleaseUnreferenced(); err != nil {
+	if err := cm.ReleaseUnreferenced(ctx); err != nil {
 		bklog.G(ctx).Errorf("failed to release unreferenced cache metadata: %+v", err)
 	}
 
@@ -40,10 +40,10 @@ type cacheManager struct {
 	results CacheResultStorage
 }
 
-func (c *cacheManager) ReleaseUnreferenced() error {
+func (c *cacheManager) ReleaseUnreferenced(ctx context.Context) error {
 	return c.backend.Walk(func(id string) error {
 		return c.backend.WalkResults(id, func(cr CacheResult) error {
-			if !c.results.Exists(cr.ID) {
+			if !c.results.Exists(ctx, cr.ID) {
 				c.backend.Release(cr.ID)
 			}
 			return nil
@@ -112,10 +112,10 @@ func (c *cacheManager) Query(deps []CacheKeyWithSelector, input Index, dgst dige
 	return keys, nil
 }
 
-func (c *cacheManager) Records(ck *CacheKey) ([]*CacheRecord, error) {
+func (c *cacheManager) Records(ctx context.Context, ck *CacheKey) ([]*CacheRecord, error) {
 	outs := make([]*CacheRecord, 0)
 	if err := c.backend.WalkResults(c.getID(ck), func(r CacheResult) error {
-		if c.results.Exists(r.ID) {
+		if c.results.Exists(ctx, r.ID) {
 			outs = append(outs, &CacheRecord{
 				ID:           r.ID,
 				cacheManager: c,

solver/cachestorage.go

@@ -49,5 +49,5 @@ type CacheResultStorage interface {
 	Save(Result, time.Time) (CacheResult, error)
 	Load(ctx context.Context, res CacheResult) (Result, error)
 	LoadRemotes(ctx context.Context, res CacheResult, compression *compression.Config, s session.Group) ([]*Remote, error)
-	Exists(id string) bool
+	Exists(ctx context.Context, id string) bool
 }

solver/combinedcache.go

@@ -100,7 +100,7 @@ func (cm *combinedCacheManager) Save(key *CacheKey, s Result, createdAt time.Tim
 	return cm.main.Save(key, s, createdAt)
 }
 
-func (cm *combinedCacheManager) Records(ck *CacheKey) ([]*CacheRecord, error) {
+func (cm *combinedCacheManager) Records(ctx context.Context, ck *CacheKey) ([]*CacheRecord, error) {
 	if len(ck.ids) == 0 {
 		return nil, errors.Errorf("no results")
 	}
@@ -112,7 +112,7 @@ func (cm *combinedCacheManager) Records(ck *CacheKey) ([]*CacheRecord, error) {
 	for c := range ck.ids {
 		func(c *cacheManager) {
 			eg.Go(func() error {
-				recs, err := c.Records(ck)
+				recs, err := c.Records(ctx, ck)
 				if err != nil {
 					return err
 				}

solver/edge.go

@@ -405,7 +405,7 @@ func (e *edge) processUpdate(upt pipe.Receiver) (depChanged bool) {
 					} else {
 						for _, k := range keys {
 							k.vtx = e.edge.Vertex.Digest()
-							records, err := e.op.Cache().Records(k)
+							records, err := e.op.Cache().Records(context.Background(), k)
 							if err != nil {
 								bklog.G(context.TODO()).Errorf("error receiving cache records: %v", err)
 								continue
@@ -583,7 +583,7 @@ func (e *edge) recalcCurrentState() {
 			}
 		}
 
-		records, err := e.op.Cache().Records(mergedKey)
+		records, err := e.op.Cache().Records(context.Background(), mergedKey)
 		if err != nil {
 			bklog.G(context.TODO()).Errorf("error receiving cache records: %v", err)
 			continue

solver/jobs.go

@@ -680,7 +680,18 @@ func (s *sharedOp) IgnoreCache() bool {
 }
 
 func (s *sharedOp) Cache() CacheManager {
-	return s.st.combinedCacheManager()
+	return &cacheWithCacheOpts{s.st.combinedCacheManager(), s.st}
+}
+
+type cacheWithCacheOpts struct {
+	CacheManager
+	st *state
+}
+
+func (c cacheWithCacheOpts) Records(ctx context.Context, ck *CacheKey) ([]*CacheRecord, error) {
+	// Allow Records accessing to cache opts through ctx. This enable to use remote provider
+	// during checking the cache existence.
+	return c.CacheManager.Records(withAncestorCacheOpts(ctx, c.st), ck)
 }
 
 func (s *sharedOp) LoadCache(ctx context.Context, rec *CacheRecord) (Result, error) {

solver/llbsolver/bridge.go

@@ -329,12 +329,12 @@ func (lcm *lazyCacheManager) Query(inp []solver.CacheKeyWithSelector, inputIndex
 	}
 	return lcm.main.Query(inp, inputIndex, dgst, outputIndex)
 }
-func (lcm *lazyCacheManager) Records(ck *solver.CacheKey) ([]*solver.CacheRecord, error) {
+func (lcm *lazyCacheManager) Records(ctx context.Context, ck *solver.CacheKey) ([]*solver.CacheRecord, error) {
 	lcm.wait()
 	if lcm.main == nil {
 		return nil, nil
 	}
-	return lcm.main.Records(ck)
+	return lcm.main.Records(ctx, ck)
 }
 func (lcm *lazyCacheManager) Load(ctx context.Context, rec *solver.CacheRecord) (solver.Result, error) {
 	if err := lcm.wait(); err != nil {

solver/memorycachestorage.go

@@ -303,7 +303,7 @@ func (s *inMemoryResultStore) LoadRemotes(_ context.Context, _ CacheResult, _ *c
 	return nil, nil
 }
 
-func (s *inMemoryResultStore) Exists(id string) bool {
+func (s *inMemoryResultStore) Exists(ctx context.Context, id string) bool {
 	_, ok := s.m.Load(id)
 	return ok
 }