fix(solver): probe the keys a dependency reports after an edge merge

solver/edge.go

@@ -479,8 +479,7 @@ func (e *edge) processUpdate(upt pipe.Receiver) (depChanged bool) {
 			return
 		}
 
-		if len(dep.keys) < len(state.keys) {
-			newKeys := state.keys[len(dep.keys):]
+		if newKeys := unseenKeys(dep.keys, state.keys); len(newKeys) > 0 {
 			if e.cacheMap != nil {
 				e.probeCache(dep, withSelector(newKeys, e.cacheMap.Deps[dep.index].Selector))
 				dep.edgeState.keys = state.keys
@@ -989,6 +988,41 @@ func isEqualState(s1, s2 edgeState) bool {
 	return true
 }
 
+// unseenKeys returns the keys in next that are not in prev. The keys of a
+// dependency normally only grow, so next extends prev. But when the
+// dependency's edge has been merged into another edge, the update carries the
+// keys of that edge instead, which can differ from prev at every position. A
+// length comparison would skip those keys, and they would never be probed.
+func unseenKeys(prev, next []ExportableCacheKey) []ExportableCacheKey {
+	if extendsKeys(prev, next) {
+		return next[len(prev):]
+	}
+	seen := make(map[*CacheKey]struct{}, len(prev))
+	for _, k := range prev {
+		seen[k.CacheKey] = struct{}{}
+	}
+	var out []ExportableCacheKey
+	for _, k := range next {
+		if _, ok := seen[k.CacheKey]; !ok {
+			out = append(out, k)
+		}
+	}
+	return out
+}
+
+// extendsKeys reports whether next starts with the keys of prev.
+func extendsKeys(prev, next []ExportableCacheKey) bool {
+	if len(prev) > len(next) {
+		return false
+	}
+	for i, k := range prev {
+		if k.CacheKey != next[i].CacheKey {
+			return false
+		}
+	}
+	return true
+}
+
 func withSelector(keys []ExportableCacheKey, selector digest.Digest) []CacheKeyWithSelector {
 	out := make([]CacheKeyWithSelector, len(keys))
 	for i, k := range keys {

solver/scheduler_test.go

@@ -1415,6 +1415,115 @@ func TestMultipleCacheSources(t *testing.T) {
 	j1 = nil
 }
 
+// TestCacheSourceLookupAfterDepMerge covers a consumer whose dependency edge
+// is merged into an equivalent edge that has already completed. The consumer
+// first receives a key from its own dependency edge that only the local cache
+// can resolve, then the dependency merges into the completed edge, whose key
+// the imported cache can resolve. The consumer must probe that key too, and
+// load its result from the imported cache instead of executing.
+//
+// This is how a frontend that solves part of a graph and then the whole graph
+// in the same session sees a local source: the source key is session-scoped,
+// so the imported cache can only match its dependants through the content
+// based key that the first solve computed.
+func TestCacheSourceLookupAfterDepMerge(t *testing.T) {
+	t.Parallel()
+	ctx := context.TODO()
+
+	// Build the graph once in another session to fill the cache that is
+	// imported below.
+	importStorage := NewInMemoryCacheStorage()
+	importResults := NewInMemoryResultStorage()
+	exportCache := NewCacheManager(ctx, "imported", importStorage, importResults)
+
+	l0 := NewSolver(SolverOpt{
+		ResolveOpFunc: testOpResolver,
+		DefaultCache:  exportCache,
+	})
+	defer l0.Close()
+
+	// graph copies src onto base and runs a step on the result. sessionSeed
+	// stands in for the session-scoped key of a local source; the copy also
+	// has a content based key for src.
+	graph := func(suffix, sessionSeed, runValue string, cacheSource CacheManager) (copied, run Edge) {
+		base := Edge{Vertex: vtx(vtxOpt{
+			name:         "base-" + suffix,
+			cacheKeySeed: "base",
+			value:        "base",
+			cacheSource:  cacheSource,
+		})}
+		src := Edge{Vertex: vtx(vtxOpt{
+			name:         "src-" + suffix,
+			cacheKeySeed: sessionSeed,
+			value:        "content",
+			cacheSource:  cacheSource,
+		})}
+		copied = Edge{Vertex: vtx(vtxOpt{
+			name:             "copy-" + suffix,
+			cacheKeySeed:     "copy",
+			value:            "copied",
+			inputs:           []Edge{base, src},
+			slowCacheCompute: map[int]ResultBasedCacheFunc{1: digestFromResult},
+			cacheSource:      cacheSource,
+		})}
+		run = Edge{Vertex: vtx(vtxOpt{
+			name:         "run-" + suffix,
+			cacheKeySeed: "run",
+			value:        runValue,
+			inputs:       []Edge{copied},
+			cacheSource:  cacheSource,
+		})}
+		return copied, run
+	}
+
+	j0, err := l0.NewJob("j0")
+	require.NoError(t, err)
+	_, run := graph("export", "session-0", "result-run", nil)
+	res, err := j0.Build(ctx, run)
+	require.NoError(t, err)
+	require.Equal(t, "result-run", unwrap(res))
+	require.NoError(t, j0.Discard())
+
+	// The imported cache cannot record links, like a cache imported from a
+	// registry.
+	importCache := NewCacheManager(ctx, "imported", readOnlyCacheKeyStorage{importStorage}, importResults)
+
+	l := NewSolver(SolverOpt{
+		ResolveOpFunc: testOpResolver,
+		DefaultCache:  NewInMemoryCacheManager(),
+	})
+	defer l.Close()
+
+	j1, err := l.NewJob("j1")
+	require.NoError(t, err)
+	defer j1.Discard()
+
+	// The first solve computes the copy through its content based key and
+	// loads it from the imported cache. Storing that result links the
+	// session-scoped key of src to the copy in the local cache.
+	copied, _ := graph("sub", "session-1", "", importCache)
+	sub := Edge{Vertex: vtx(vtxOpt{
+		name:         "sub",
+		cacheKeySeed: "sub",
+		value:        "result-sub",
+		inputs:       []Edge{copied},
+		cacheSource:  importCache,
+	})}
+	res, err = j1.Build(ctx, sub)
+	require.NoError(t, err)
+	require.Equal(t, "result-sub", unwrap(res))
+
+	// The second solve uses new vertices for the same steps in the same
+	// session. Its copy first finds the local copy through the session-scoped
+	// key and then merges into the copy from the first solve.
+	_, run = graph("main", "session-1", "result-run-executed", importCache)
+	run.Vertex.(*vertex).setupCallCounters()
+	res, err = j1.Build(ctx, run)
+	require.NoError(t, err)
+	require.Equal(t, "result-run", unwrap(res))
+	require.Equal(t, int64(0), *run.Vertex.(*vertex).execCallCount)
+}
+
 func TestRepeatBuildWithIgnoreCache(t *testing.T) {
 	t.Parallel()
 	ctx := context.TODO()
@@ -3872,6 +3981,16 @@ func (cm *trackingCacheManager) Load(ctx context.Context, rec *CacheRecord) (Res
 	return cm.CacheManager.Load(ctx, rec)
 }
 
+// readOnlyCacheKeyStorage ignores new links, like the storage of a cache
+// imported from a registry.
+type readOnlyCacheKeyStorage struct {
+	CacheKeyStorage
+}
+
+func (readOnlyCacheKeyStorage) AddLink(string, CacheInfoLink, string) error {
+	return nil
+}
+
 func digestFromResult(ctx context.Context, res Result, _ session.Group) (digest.Digest, error) {
 	return digest.FromBytes([]byte(unwrap(res))), nil
 }