fix(remotecache): make cache-export loop detection linear and cancellable

cache/remotecache/v1/chains.go

@@ -60,13 +60,18 @@ func (c *CacheChains) normalize(ctx context.Context) error {
 	c.items = validated
 
 	for _, it := range c.items {
+		if err := ctx.Err(); err != nil {
+			return err
+		}
 		_, err := normalizeItem(it, st)
 		if err != nil {
 			return err
 		}
 	}
 
-	st.removeLoops(ctx)
+	if err := st.removeLoops(ctx); err != nil {
+		return err
+	}
 
 	items := make([]*item, 0, len(st.byKey))
 	for _, it := range st.byKey {

cache/remotecache/v1/chains_test.go

@@ -3,6 +3,7 @@ package cacheimport
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"testing"
 	"time"
 
@@ -95,3 +96,104 @@ func TestSimpleMarshal(t *testing.T) {
 func dgst(s string) digest.Digest {
 	return digest.FromBytes([]byte(s))
 }
+
+// addDiamondChain adds a chain of n fan-in "diamonds" (a_i -> {b_i, c_i} -> a_i+1)
+// to cc. The graph has 3n+1 items but 2^n distinct root-to-leaf paths.
+func addDiamondChain(cc *CacheChains, n int) {
+	prev := cc.Add(outputKey(dgst("a0"), 0))
+	prev.AddResult("", 0, time.Now(), &solver.Remote{})
+	for i := 0; i < n; i++ {
+		b := cc.Add(outputKey(dgst(fmt.Sprintf("b%d", i)), 0))
+		b.LinkFrom(prev, 0, "")
+		c := cc.Add(outputKey(dgst(fmt.Sprintf("c%d", i)), 0))
+		c.LinkFrom(prev, 0, "")
+		a := cc.Add(outputKey(dgst(fmt.Sprintf("a%d", i+1)), 0))
+		a.LinkFrom(b, 0, "")
+		a.LinkFrom(c, 1, "")
+		prev = a
+	}
+}
+
+// TestMarshalDiamondChain guards against loop detection enumerating every
+// path through the cache graph, which is exponential in the number of
+// diamonds and made cache export hang on large build graphs.
+func TestMarshalDiamondChain(t *testing.T) {
+	const n = 1000
+	cc := NewCacheChains()
+	addDiamondChain(cc, n)
+
+	done := make(chan error, 1)
+	var cfg *CacheConfig
+	go func() {
+		var err error
+		cfg, _, err = cc.Marshal(context.TODO())
+		done <- err
+	}()
+	select {
+	case err := <-done:
+		require.NoError(t, err)
+	case <-time.After(30 * time.Second):
+		t.Fatalf("marshaling a chain of %d diamonds did not finish in time", n)
+	}
+	require.Equal(t, 3*n+1, len(cfg.Records))
+}
+
+func TestMarshalCanceled(t *testing.T) {
+	cc := NewCacheChains()
+	addDiamondChain(cc, 10)
+
+	ctx, cancel := context.WithCancel(context.TODO())
+	cancel()
+	_, _, err := cc.Marshal(ctx)
+	require.ErrorIs(t, err, context.Canceled)
+}
+
+// TestMarshalLoop checks that loops created by deduplication during
+// normalization are still broken, including when the loop is reachable
+// through several paths (so its nodes are already fully explored when the
+// second path reaches them).
+func TestMarshalLoop(t *testing.T) {
+	cc := NewCacheChains()
+
+	root := cc.Add(outputKey(dgst("root"), 0))
+	root.AddResult("", 0, time.Now(), &solver.Remote{})
+	// diamond in front of the loop: root -> {p, q} -> x
+	p := cc.Add(outputKey(dgst("p"), 0))
+	p.LinkFrom(root, 0, "")
+	q := cc.Add(outputKey(dgst("q"), 0))
+	q.LinkFrom(root, 0, "")
+	x1 := cc.Add(outputKey(dgst("x"), 0))
+	x1.LinkFrom(p, 0, "")
+	x1.LinkFrom(q, 0, "")
+	y := cc.Add(outputKey(dgst("y"), 0))
+	y.LinkFrom(x1, 0, "")
+	// x2 has the same digest as x1 and one input alternative in common, so
+	// normalization merges it into x1, which creates the loop x -> y -> x.
+	x2 := cc.Add(outputKey(dgst("x"), 0))
+	x2.LinkFrom(p, 0, "")
+	x2.LinkFrom(y, 0, "")
+
+	cfg, _, err := cc.Marshal(context.TODO())
+	require.NoError(t, err)
+	require.Equal(t, 5, len(cfg.Records))
+
+	// the exported records must form a DAG
+	state := map[int]int{}
+	var visit func(i int)
+	visit = func(i int) {
+		require.NotEqual(t, 1, state[i], "loop in exported cache records")
+		if state[i] == 2 {
+			return
+		}
+		state[i] = 1
+		for _, inputs := range cfg.Records[i].Inputs {
+			for _, in := range inputs {
+				visit(in.LinkIndex)
+			}
+		}
+		state[i] = 2
+	}
+	for i := range cfg.Records {
+		visit(i)
+	}
+}

cache/remotecache/v1/utils.go

@@ -128,29 +128,43 @@ type normalizeState struct {
 	next  int
 }
 
-func (s *normalizeState) removeLoops(ctx context.Context) {
+func (s *normalizeState) removeLoops(ctx context.Context) error {
 	roots := []digest.Digest{}
 	for dgst, it := range s.byKey {
 		if len(it.links) == 0 {
 			roots = append(roots, dgst)
 		}
 	}
 
+	// visited holds the digests on the current DFS path; done holds the
+	// digests whose subtrees have been fully explored. Without done, every
+	// root-to-leaf path is enumerated, which is exponential in the number of
+	// fan-in "diamonds" (A->{B,C}->D) in the graph.
 	visited := map[digest.Digest]struct{}{}
+	done := map[digest.Digest]struct{}{}
 
 	for _, d := range roots {
-		s.checkLoops(ctx, d, visited)
+		if err := s.checkLoops(ctx, d, visited, done); err != nil {
+			return err
+		}
 	}
+	return nil
 }
 
-func (s *normalizeState) checkLoops(ctx context.Context, d digest.Digest, visited map[digest.Digest]struct{}) {
+func (s *normalizeState) checkLoops(ctx context.Context, d digest.Digest, visited, done map[digest.Digest]struct{}) error {
+	if _, ok := done[d]; ok {
+		return nil
+	}
+	if err := ctx.Err(); err != nil {
+		return err
+	}
 	it, ok := s.byKey[d]
 	if !ok {
-		return
+		return nil
 	}
 	links, ok := s.links[it]
 	if !ok {
-		return
+		return nil
 	}
 	visited[d] = struct{}{}
 	defer func() {
@@ -168,11 +182,13 @@ func (s *normalizeState) checkLoops(ctx context.Context, d digest.Digest, visite
 					bklog.G(ctx).Warnf("failed to remove looping cache key %s %s", d, id)
 				}
 				delete(links[l], id)
-			} else {
-				s.checkLoops(ctx, id, visited)
+			} else if err := s.checkLoops(ctx, id, visited, done); err != nil {
+				return err
 			}
 		}
 	}
+	done[d] = struct{}{}
+	return nil
 }
 
 func normalizeItem(it *item, state *normalizeState) (*item, error) {

solver/exporter.go

@@ -172,6 +172,9 @@ 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 {
+				if ctxErr := ctx.Err(); ctxErr != nil {
+					return nil, ctxErr
+				}
 				return nil, nil
 			}
 			for _, r := range recs {
@@ -184,6 +187,9 @@ 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 {
+				if ctxErr := ctx.Err(); ctxErr != nil {
+					return nil, ctxErr
+				}
 				return nil, nil
 			}
 			for _, r := range recs {