tka: avoid quadratic filesystem scans

Since #17567, FS.ChildAUMs scans and decodes every active AUM file.
Authority reconstruction and compaction call it repeatedly, making
filesystem work quadratic in the number of AUMs.

Build per-FS indexes of active AUM hashes and parent-to-child
relationships in one scan. Use the indexes for graph queries while
continuing to decode returned AUM values from disk. Invalidate the
indexes after mutations, but retain them for empty purges.

A production-sized 825-AUM benchmark improves from 124 seconds to
4 seconds on macOS with Defender and from 5.9 seconds to 34 milliseconds
on Linux ext3. The checked-in benchmark exercises the 2,000-AUM maximum
linear history accepted by the traversal guardrail.

RELNOTE: Avoid Tailnet Lock startup failures with large authority histories.

Fixes #20735

Change-Id: I088136bc2e9d1b6bfdecb223c069d42400c9f63d
Signed-off-by: Michael Renner <terrorobe@github.com>

tka/tailchonk.go

@@ -20,6 +20,7 @@ import (
 	"github.com/fxamacker/cbor/v2"
 	"tailscale.com/atomicfile"
 	"tailscale.com/tstime"
+	"tailscale.com/util/set"
 	"tailscale.com/util/testenv"
 )
 
@@ -286,14 +287,30 @@ func (c *Mem) PurgeAUMs(hashes []AUMHash) error {
 
 // FS implements filesystem storage of TKA state.
 //
+// FS caches graph metadata derived from its storage directory. Mutations made
+// through this instance invalidate the cache. Graph changes made through
+// another FS or directly to the directory are not reflected in the indexes
+// until they are invalidated or this instance is reopened.
+//
 // FS implements the Chonk interface.
 type FS struct {
 	base string
 	mu   sync.RWMutex
+
+	// aumIndex and parentIndex are initialized together by buildIndexLocked.
+	// They avoid repeatedly scanning every AUM file while walking the authority
+	// graph. AUM contents are still decoded from disk on demand. A nil aumIndex
+	// means the indexes need to be rebuilt.
+	aumIndex    set.Set[AUMHash]
+	parentIndex map[AUMHash][]AUMHash
 }
 
-// ChonkDir returns an implementation of Chonk which uses the
-// given directory to store TKA state.
+// ChonkDir returns an implementation of Chonk which uses the given directory
+// to store TKA state.
+//
+// The returned FS expects exclusive write access to dir for its lifetime.
+// Graph changes made through another FS or directly to dir may not be reflected
+// in queries until the returned FS is reopened.
 func ChonkDir(dir string) (*FS, error) {
 	if err := os.MkdirAll(dir, 0755); err != nil && !os.IsExist(err) {
 		return nil, fmt.Errorf("creating chonk root dir: %v", err)
@@ -354,7 +371,12 @@ func (c *FS) aumDir(h AUMHash) (dir, base string) {
 func (c *FS) AUM(hash AUMHash) (AUM, error) {
 	c.mu.RLock()
 	defer c.mu.RUnlock()
+	return c.aumLocked(hash)
+}
 
+// aumLocked reads an AUM from disk. The caller must hold c.mu for reading or
+// writing.
+func (c *FS) aumLocked(hash AUMHash) (AUM, error) {
 	info, err := c.get(hash)
 	if err != nil {
 		if os.IsNotExist(err) {
@@ -401,20 +423,36 @@ func (c *FS) CommitTime(h AUMHash) (time.Time, error) {
 	return s.ModTime(), nil
 }
 
-// AUM returns any known AUMs with a specific parent hash.
+// ChildAUMs returns any known AUMs with a specific parent hash.
 func (c *FS) ChildAUMs(prevAUMHash AUMHash) ([]AUM, error) {
 	c.mu.RLock()
-	defer c.mu.RUnlock()
+	if c.aumIndex != nil {
+		defer c.mu.RUnlock()
+		return c.childAUMsFromIndexLocked(prevAUMHash)
+	}
+	c.mu.RUnlock()
 
-	var out []AUM
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if err := c.buildIndexLocked(); err != nil {
+		return nil, err
+	}
+	return c.childAUMsFromIndexLocked(prevAUMHash)
+}
 
-	err := c.scanHashes(func(info *fsHashInfo) {
-		if info.AUM != nil && bytes.Equal(info.AUM.PrevAUMHash, prevAUMHash[:]) {
-			out = append(out, *info.AUM)
+// childAUMsFromIndexLocked returns children from the in-memory index. The
+// caller must hold c.mu for reading or writing.
+func (c *FS) childAUMsFromIndexLocked(prevAUMHash AUMHash) ([]AUM, error) {
+	children := c.parentIndex[prevAUMHash]
+	out := make([]AUM, 0, len(children))
+	for i, h := range children {
+		aum, err := c.aumLocked(h)
+		if err != nil {
+			return nil, fmt.Errorf("reading child %d of %x: %w", i, prevAUMHash, err)
 		}
-	})
-
-	return out, err
+		out = append(out, aum)
+	}
+	return out, nil
 }
 
 func (c *FS) get(h AUMHash) (*fsHashInfo, error) {
@@ -442,72 +480,99 @@ func (c *FS) get(h AUMHash) (*fsHashInfo, error) {
 
 // Heads returns AUMs for which there are no children. In other
 // words, the latest AUM in all possible chains (the 'leaves').
-//
-// Heads is expected to be called infrequently compared to AUM() or
-// ChildAUMs(), so we haven't put any work into maintaining an index.
-// Instead, the full set of AUMs is scanned.
 func (c *FS) Heads() ([]AUM, error) {
 	c.mu.RLock()
-	defer c.mu.RUnlock()
+	if c.aumIndex != nil {
+		defer c.mu.RUnlock()
+		return c.headsFromIndexLocked()
+	}
+	c.mu.RUnlock()
 
-	// Scan the complete list of AUMs, and build a list of all parent hashes.
-	// This tells us which AUMs have children.
-	var parentHashes []AUMHash
-
-	allAUMs, err := c.AllAUMs()
-	if err != nil {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if err := c.buildIndexLocked(); err != nil {
 		return nil, err
 	}
+	return c.headsFromIndexLocked()
+}
 
-	for _, h := range allAUMs {
-		aum, err := c.AUM(h)
-		if err != nil {
-			return nil, err
-		}
-		parent, hasParent := aum.Parent()
-		if !hasParent {
-			continue
-		}
-		if !slices.Contains(parentHashes, parent) {
-			parentHashes = append(parentHashes, parent)
-		}
-	}
-
-	// Now scan a second time, and only include AUMs which weren't marked as
-	// the parent of any other AUM.
+// headsFromIndexLocked returns heads from the in-memory index. The caller
+// must hold c.mu for reading or writing.
+func (c *FS) headsFromIndexLocked() ([]AUM, error) {
 	out := make([]AUM, 0, 6) // 6 is arbitrary.
-
-	for _, h := range allAUMs {
-		if slices.Contains(parentHashes, h) {
+	for h := range c.aumIndex {
+		if len(c.parentIndex[h]) != 0 {
 			continue
 		}
-		aum, err := c.AUM(h)
+		aum, err := c.aumLocked(h)
 		if err != nil {
-			return nil, err
+			return nil, fmt.Errorf("reading head %x: %w", h, err)
 		}
 		out = append(out, aum)
 	}
-
 	return out, nil
 }
 
 // RemoveAll permanently and completely clears the TKA state.
 func (c *FS) RemoveAll() error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	c.invalidateIndexLocked()
 	return os.RemoveAll(c.base)
 }
 
 // AllAUMs returns all AUMs stored in the chonk.
 func (c *FS) AllAUMs() ([]AUMHash, error) {
 	c.mu.RLock()
-	defer c.mu.RUnlock()
+	if c.aumIndex != nil {
+		out := slices.Collect(maps.Keys(c.aumIndex))
+		c.mu.RUnlock()
+		return out, nil
+	}
+	c.mu.RUnlock()
 
-	out := make([]AUMHash, 0, 6) // 6 is arbitrary.
-	err := c.scanHashes(func(info *fsHashInfo) {
-		if info.AUM != nil {
-			out = append(out, info.AUM.Hash())
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if err := c.buildIndexLocked(); err != nil {
+		return nil, err
+	}
+	return slices.Collect(maps.Keys(c.aumIndex)), nil
+}
+
+// buildIndexLocked builds indexes of the active AUMs in a single filesystem
+// scan. If the index already exists, this is a no-op. The caller must hold
+// c.mu for writing.
+func (c *FS) buildIndexLocked() error {
+	if c.aumIndex != nil {
+		return nil
+	}
+
+	aumIndex := make(set.Set[AUMHash])
+	parentIndex := make(map[AUMHash][]AUMHash)
+	if err := c.scanHashes(func(info *fsHashInfo) {
+		if info.AUM == nil {
+			return
 		}
-	})
-	return out, err
+		h := info.AUM.Hash()
+		aumIndex.Add(h)
+		if parent, ok := info.AUM.Parent(); ok {
+			parentIndex[parent] = append(parentIndex[parent], h)
+		}
+	}); err != nil {
+		return err
+	}
+
+	c.aumIndex = aumIndex
+	c.parentIndex = parentIndex
+	return nil
+}
+
+// invalidateIndexLocked invalidates the in-memory indexes after storage is
+// changed. The caller must hold c.mu for writing.
+func (c *FS) invalidateIndexLocked() {
+	c.aumIndex = nil
+	c.parentIndex = nil
 }
 
 func (c *FS) scanHashes(eachHashInfo func(*fsHashInfo)) error {
@@ -594,6 +659,7 @@ func (c *FS) CommitVerifiedAUMs(updates []AUM) error {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
+	c.invalidateIndexLocked()
 	for i, aum := range updates {
 		h := aum.Hash()
 		err := c.commit(h, func(info *fsHashInfo) {
@@ -610,9 +676,14 @@ func (c *FS) CommitVerifiedAUMs(updates []AUM) error {
 
 // PurgeAUMs marks the specified AUMs for deletion from storage.
 func (c *FS) PurgeAUMs(hashes []AUMHash) error {
+	if len(hashes) == 0 {
+		return nil
+	}
+
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
+	c.invalidateIndexLocked()
 	now := time.Now()
 	for i, h := range hashes {
 		stored, err := c.get(h)

tka/tailchonk_test.go

@@ -104,22 +104,13 @@ func TestTailchonkFS_IgnoreTempFile(t *testing.T) {
 		}
 	}
 
-	// Check that calling AllAUMs() returns the single committed AUM
-	got, err := chonk.AllAUMs()
-	if err != nil {
-		t.Fatalf("AllAUMs() failed: %v", err)
-	}
-	want := []AUMHash{aum.Hash()}
-	if !slices.Equal(got, want) {
-		t.Fatalf("AllAUMs() is wrong: got %v, want %v", got, want)
-	}
-
 	// Write some temporary files which are named like partially-committed AUMs,
-	// then check that AllAUMs() only returns the single committed AUM.
+	// then check that the initial scan only returns the committed AUM.
 	writeAUMFile("AUM1234.tmp", "incomplete AUM\n")
 	writeAUMFile("AUM1234.tmp_123", "second incomplete AUM\n")
 
-	got, err = chonk.AllAUMs()
+	got, err := chonk.AllAUMs()
+	want := []AUMHash{aum.Hash()}
 	if err != nil {
 		t.Fatalf("AllAUMs() failed: %v", err)
 	}
@@ -165,6 +156,115 @@ func TestTailchonkFS_CannotUseFile(t *testing.T) {
 	}
 }
 
+// Indexed FS reads decode AUMs from disk, so mutable fields in returned values
+// must not alias subsequent results.
+func TestTailchonkFS_ReturnsIndependentAUMs(t *testing.T) {
+	chonk := must.Get(ChonkDir(t.TempDir()))
+	parent := AUM{MessageKind: AUMRemoveKey, KeyID: []byte{0}}
+	parentHash := parent.Hash()
+	child := AUM{MessageKind: AUMRemoveKey, KeyID: []byte{1}, PrevAUMHash: parentHash[:]}
+	must.Do(chonk.CommitVerifiedAUMs([]AUM{parent, child}))
+
+	got := must.Get(chonk.ChildAUMs(parentHash))
+	if len(got) != 1 {
+		t.Fatalf("ChildAUMs() returned %d children, want 1", len(got))
+	}
+	got[0].PrevAUMHash[0] ^= 0xff
+	got[0].KeyID[0] ^= 0xff
+
+	got = must.Get(chonk.ChildAUMs(parentHash))
+	if diff := cmp.Diff([]AUM{child}, got); diff != "" {
+		t.Fatalf("stored child changed through a returned AUM (-want, +got):\n%s", diff)
+	}
+}
+
+func TestTailchonkFS_ConcurrentIndexBuild(t *testing.T) {
+	chonk := must.Get(ChonkDir(t.TempDir()))
+	parent := AUM{MessageKind: AUMRemoveKey, KeyID: []byte{0}}
+	parentHash := parent.Hash()
+	child := AUM{MessageKind: AUMRemoveKey, KeyID: []byte{1}, PrevAUMHash: parentHash[:]}
+	must.Do(chonk.CommitVerifiedAUMs([]AUM{parent, child}))
+
+	start := make(chan struct{})
+	var wg sync.WaitGroup
+	for i := range 30 {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			<-start
+
+			switch i % 3 {
+			case 0:
+				if got, err := chonk.ChildAUMs(parentHash); err != nil || len(got) != 1 {
+					t.Errorf("ChildAUMs() = %v, %v; want one child", got, err)
+				}
+			case 1:
+				if got, err := chonk.Heads(); err != nil || len(got) != 1 {
+					t.Errorf("Heads() = %v, %v; want one head", got, err)
+				}
+			case 2:
+				if got, err := chonk.AllAUMs(); err != nil || len(got) != 2 {
+					t.Errorf("AllAUMs() = %v, %v; want two AUMs", got, err)
+				}
+			}
+		}()
+	}
+	close(start)
+	wg.Wait()
+}
+
+func BenchmarkTailchonkFSOpenLongChain(b *testing.B) {
+	// Exercise the largest linear AUM history accepted by maxScanIterations.
+	// Generate more history than needed, then retain a prefix so the benchmark
+	// size does not depend on checkpoint cadence.
+	const targetAUMs = maxScanIterations
+
+	signer := key.NewNLPrivate()
+	trustedKey := Key{Kind: Key25519, Public: signer.Public().Verifier(), Votes: 1}
+	storage := ChonkMem()
+	authority, _, err := Create(storage, CreateStateForTest(trustedKey), signer)
+	if err != nil {
+		b.Fatal(err)
+	}
+	SeedAUMs(b, SeedAUMConfig{
+		Count:  targetAUMs,
+		Signer: signer,
+		Nodes:  []SeedNode{CreateSeedNode(b, authority, storage)},
+	})
+
+	// SeedAUMs may insert checkpoints in addition to the requested updates.
+	// Walk from the final head to genesis, reverse the result, and retain
+	// exactly the supported maximum.
+	aums := make([]AUM, 0, targetAUMs)
+	for h := authority.Head(); ; {
+		aum := must.Get(storage.AUM(h))
+		aums = append(aums, aum)
+
+		parent, ok := aum.Parent()
+		if !ok {
+			break
+		}
+		h = parent
+	}
+	slices.Reverse(aums)
+
+	if len(aums) < targetAUMs {
+		b.Fatalf("seeded %d AUMs, want at least %d", len(aums), targetAUMs)
+	}
+	aums = aums[:targetAUMs]
+
+	dir := b.TempDir()
+	fs := must.Get(ChonkDir(dir))
+	must.Do(fs.CommitVerifiedAUMs(aums))
+
+	b.ResetTimer()
+	for b.Loop() {
+		fs := must.Get(ChonkDir(dir))
+		must.Get(Open(fs))
+	}
+	b.ReportMetric(float64(len(aums)), "AUMs")
+}
+
 func TestMarkActiveChain(t *testing.T) {
 	type aumTemplate struct {
 		AUM AUM

tstest/chonktest/chonktest.go

@@ -79,6 +79,25 @@ func RunChonkTests(t *testing.T, newChonk func(*testing.T) tka.Chonk) {
 		if diff := cmp.Diff(data, stored, cmpopts.SortSlices(aumHashesLess)); diff != "" {
 			t.Errorf("stored AUM differs (-want, +got):\n%s", diff)
 		}
+
+		// Commit another child after ChildAUMs has been called, then check that
+		// subsequent reads include it.
+		additional := tka.AUM{
+			MessageKind: tka.AUMRemoveKey,
+			KeyID:       []byte{5, 6},
+			PrevAUMHash: parentHash[:],
+		}
+		if err := chonk.CommitVerifiedAUMs([]tka.AUM{additional}); err != nil {
+			t.Fatalf("additional CommitVerifiedAUMs failed: %v", err)
+		}
+		stored, err = chonk.ChildAUMs(parentHash)
+		if err != nil {
+			t.Fatalf("ChildAUMs after additional commit failed: %v", err)
+		}
+		data = append(data, additional)
+		if diff := cmp.Diff(data, stored, cmpopts.SortSlices(aumHashesLess)); diff != "" {
+			t.Errorf("stored AUMs after additional commit differ (-want, +got):\n%s", diff)
+		}
 	})
 
 	t.Run("AUMMissing", func(t *testing.T) {