keyviz: track hidden bucket members for accurate route_count

Round-2 review fixes for PR #646 (Codex P2 ×2):

P2 #1: foldIntoBucket dedup only consulted the visible MemberRoutes
slice, so re-registering a past-cap (hidden) routeID inside the
prune grace incremented MemberRoutesTotal a second time —
route_count drifted upward on remove/re-register churn even though
the bucket membership didn't grow.

P2 #2: pruneMemberRoute only decremented MemberRoutesTotal for
routes visible in MemberRoutes. Past-cap members were absent from
that list, so removals never reduced the total — route_count stayed
permanently inflated and route_ids_truncated stayed permanently
true even after the route was fully gone.

Add a `hiddenMembers map[uint64]struct{}` set on routeSlot. Past-cap
folds land there instead of being silently dropped from accounting:
- foldIntoBucket dedups against both MemberRoutes and hiddenMembers
  before incrementing.
- pruneMemberRoute checks hiddenMembers when the routeID isn't in
  the visible list, deletes from the set (releasing the map when it
  empties), and decrements the total.

Refactor: extract addMemberToBucket helper so foldIntoBucket stays
under the cyclop budget.

Tests:
- TestPastCapMemberRejoinDoesNotInflateTotal — Remove + Register a
  hidden member → MemberRoutesTotal stays the same, no drift.
- TestPastCapMemberPruneDecrementsTotal — Remove + grace expiry of a
  hidden member → MemberRoutesTotal decrements after the prune fires
  on the second post-grace flush.
- Updated TestMemberRoutesCappedAtConfiguredCap to assert
  MemberRoutesTotal=8 alongside the existing visible-cap=3 check.

Author: bootjp

keyviz/sampler.go

@@ -226,11 +226,16 @@ type routeSlot struct {
 	// routes together (Snapshot surfaces this in MatrixRow).
 	Aggregate    bool
 	MemberRoutes []uint64
-	// MemberRoutesTotal counts every distinct routeID that has folded
-	// into this bucket, including ones beyond MaxMemberRoutesPerSlot
-	// (which still contribute to the counters but are not appended to
-	// MemberRoutes). Always equals len(MemberRoutes) for individual
-	// (non-Aggregate) slots.
+	// hiddenMembers stores routeIDs folded past
+	// MaxMemberRoutesPerSlot. Used to dedup re-folds (so
+	// MemberRoutesTotal doesn't drift on remove+re-register churn for
+	// past-cap routes) and to drive accurate decrements in
+	// pruneMemberRoute. nil for individual slots.
+	hiddenMembers map[uint64]struct{}
+	// MemberRoutesTotal is len(MemberRoutes) + len(hiddenMembers) — the
+	// authoritative count of distinct routes folded into this bucket.
+	// Always equals len(MemberRoutes) for individual (non-Aggregate)
+	// slots.
 	MemberRoutesTotal uint64
 
 	reads      atomic.Uint64
@@ -455,16 +460,12 @@ func (s *MemSampler) RegisterRoute(routeID uint64, start, end []byte) bool {
 // to preserve Flush's key-order contract.
 //
 // MemberRoutes growth is capped by MaxMemberRoutesPerSlot — beyond
-// that cap the bucket counters still absorb the route's traffic, but
-// the routeID is not added to the visible member list.
+// that cap the bucket counters still absorb the route's traffic and
+// the routeID is recorded in hiddenMembers (so dedup is correct on
+// rejoin) but not appended to the visible MemberRoutes list.
 func (s *MemSampler) foldIntoBucket(next *routeTable, bucket *routeSlot, routeID uint64, start, end []byte) {
 	bucket.metaMu.Lock()
-	if !memberRoutesContains(bucket.MemberRoutes, routeID) {
-		bucket.MemberRoutesTotal++
-		if len(bucket.MemberRoutes) < s.opts.MaxMemberRoutesPerSlot {
-			bucket.MemberRoutes = append(bucket.MemberRoutes, routeID)
-		}
-	}
+	addMemberToBucket(bucket, routeID, s.opts.MaxMemberRoutesPerSlot)
 	if len(end) == 0 || (len(bucket.End) != 0 && bytesGT(end, bucket.End)) {
 		bucket.End = cloneBytes(end)
 	}
@@ -478,6 +479,31 @@ func (s *MemSampler) foldIntoBucket(next *routeTable, bucket *routeSlot, routeID
 	}
 }
 
+// addMemberToBucket records routeID against bucket, choosing the
+// visible MemberRoutes list when there's room and falling back to
+// the hiddenMembers set past the cap. Both lists are deduped so a
+// rejoin during the prune grace doesn't inflate MemberRoutesTotal.
+// Caller holds bucket.metaMu.
+func addMemberToBucket(bucket *routeSlot, routeID uint64, visibleCap int) {
+	if memberRoutesContains(bucket.MemberRoutes, routeID) {
+		return
+	}
+	if bucket.hiddenMembers != nil {
+		if _, ok := bucket.hiddenMembers[routeID]; ok {
+			return
+		}
+	}
+	if len(bucket.MemberRoutes) < visibleCap {
+		bucket.MemberRoutes = append(bucket.MemberRoutes, routeID)
+	} else {
+		if bucket.hiddenMembers == nil {
+			bucket.hiddenMembers = make(map[uint64]struct{})
+		}
+		bucket.hiddenMembers[routeID] = struct{}{}
+	}
+	bucket.MemberRoutesTotal++
+}
+
 // RemoveRoute drops a RouteID from tracking. Counts accumulated since
 // the last flush are NOT lost: the retired slot (or, for virtual-bucket
 // members, just the membership entry) is queued for one final drain by
@@ -719,12 +745,13 @@ func bucketStillReferenced(virtualForRoute map[uint64]*routeSlot, bucket *routeS
 	return false
 }
 
-// pruneMemberRoute removes routeID from bucket.MemberRoutes under the
-// bucket's metaMu so a concurrent snapshotMeta reader sees a
-// consistent view. MemberRoutesTotal is decremented when the routeID
-// was visible in MemberRoutes (the only case we can confidently
-// account for) — routes pruned past the visible cap stay in the
-// total because we don't track individual past-cap members.
+// pruneMemberRoute removes routeID from bucket.MemberRoutes (or
+// hiddenMembers) under the bucket's metaMu so a concurrent
+// snapshotMeta reader sees a consistent view. MemberRoutesTotal is
+// decremented whenever the routeID was actually present in either
+// list — including past-cap members in hiddenMembers — so the
+// reported route_count stays truthful across remove/re-register
+// churn.
 func pruneMemberRoute(bucket *routeSlot, routeID uint64) {
 	bucket.metaMu.Lock()
 	defer bucket.metaMu.Unlock()
@@ -738,6 +765,15 @@ func pruneMemberRoute(bucket *routeSlot, routeID uint64) {
 		filtered = append(filtered, m)
 	}
 	bucket.MemberRoutes = filtered
+	if !removed && bucket.hiddenMembers != nil {
+		if _, ok := bucket.hiddenMembers[routeID]; ok {
+			delete(bucket.hiddenMembers, routeID)
+			if len(bucket.hiddenMembers) == 0 {
+				bucket.hiddenMembers = nil
+			}
+			removed = true
+		}
+	}
 	if removed && bucket.MemberRoutesTotal > 0 {
 		bucket.MemberRoutesTotal--
 	}

keyviz/sampler_test.go

@@ -916,6 +916,94 @@ func TestMemberRoutesCappedAtConfiguredCap(t *testing.T) {
 	if agg.Reads != 8 {
 		t.Fatalf("bucket Reads = %d, want 8 (counters must absorb traffic past cap)", agg.Reads)
 	}
+	if agg.MemberRoutesTotal != 8 {
+		t.Fatalf("MemberRoutesTotal = %d, want 8 (true count of contributors past cap)", agg.MemberRoutesTotal)
+	}
+}
+
+// TestPastCapMemberRejoinDoesNotInflateTotal pins Codex round-2 P2
+// on PR #646: re-registering a hidden (past-cap) routeID inside the
+// prune grace must not double-count MemberRoutesTotal. The hidden
+// member set serves as dedup so route_count stays truthful across
+// remove/re-register churn for routes the visible MemberRoutes list
+// never carried.
+func TestPastCapMemberRejoinDoesNotInflateTotal(t *testing.T) {
+	t.Parallel()
+	s, _ := newTestSampler(t, MemSamplerOptions{
+		Step:                   time.Second,
+		HistoryColumns:         4,
+		MaxTrackedRoutes:       1,
+		MaxMemberRoutesPerSlot: 1, // visible cap=1; routes 3+4 land in hidden set
+	})
+	mustRegister(t, s, 1, "a", "b")
+	if s.RegisterRoute(2, []byte("c"), []byte("d")) {
+		t.Fatal("route 2 should fold (visible)")
+	}
+	if s.RegisterRoute(3, []byte("e"), []byte("f")) {
+		t.Fatal("route 3 should fold (hidden — past cap)")
+	}
+	if s.RegisterRoute(4, []byte("g"), []byte("h")) {
+		t.Fatal("route 4 should fold (hidden — past cap)")
+	}
+	s.Observe(2, OpRead, 0, 0)
+	s.Flush()
+	beforeTotal := findAggregateRow(t, lastSnapshotColumn(t, s).Rows).MemberRoutesTotal
+	if beforeTotal != 3 {
+		t.Fatalf("baseline MemberRoutesTotal = %d, want 3", beforeTotal)
+	}
+	// Remove the past-cap route 3 and re-register it inside the prune
+	// grace. Without the hidden-member dedup foldIntoBucket would
+	// increment MemberRoutesTotal again, drifting route_count up.
+	s.RemoveRoute(3)
+	if s.RegisterRoute(3, []byte("e"), []byte("f")) {
+		t.Fatal("route 3 should fold again")
+	}
+	s.Observe(2, OpRead, 0, 0)
+	s.Flush()
+	afterTotal := findAggregateRow(t, lastSnapshotColumn(t, s).Rows).MemberRoutesTotal
+	if afterTotal != 3 {
+		t.Fatalf("after remove+re-register, MemberRoutesTotal = %d, want 3 (no drift)", afterTotal)
+	}
+}
+
+// TestPastCapMemberPruneDecrementsTotal pins the other half of Codex
+// round-2 P2: when grace expires for a hidden (past-cap) member, the
+// prune must decrement MemberRoutesTotal — otherwise route_count
+// stays inflated even after the route is fully retired.
+func TestPastCapMemberPruneDecrementsTotal(t *testing.T) {
+	t.Parallel()
+	s, clk := newTestSampler(t, MemSamplerOptions{
+		Step:                   time.Second,
+		HistoryColumns:         8,
+		MaxTrackedRoutes:       1,
+		MaxMemberRoutesPerSlot: 1,
+	})
+	mustRegister(t, s, 1, "a", "b")
+	if s.RegisterRoute(2, []byte("c"), []byte("d")) {
+		t.Fatal("route 2 should fold (visible)")
+	}
+	if s.RegisterRoute(3, []byte("e"), []byte("f")) {
+		t.Fatal("route 3 should fold (hidden — past cap)")
+	}
+	s.Observe(2, OpRead, 0, 0)
+	s.Flush()
+	require := func(want uint64, ctx string) {
+		t.Helper()
+		s.Observe(2, OpRead, 0, 0)
+		s.Flush()
+		got := findAggregateRow(t, lastSnapshotColumn(t, s).Rows).MemberRoutesTotal
+		if got != want {
+			t.Fatalf("%s: MemberRoutesTotal = %d, want %d", ctx, got, want)
+		}
+	}
+	require(2, "baseline")
+	s.RemoveRoute(3) // hidden member — pendingPrune queued
+	clk.Advance(s.graceWindow() + time.Second)
+	// First flush past grace: drains the bucket (snapshot still
+	// shows Total=2) THEN fires the pending prune. The next flush is
+	// the first one whose drain reflects the post-prune state.
+	require(2, "drain runs before prune fires")
+	require(1, "after past-cap prune fires")
 }
 
 // TestRetiredTailClearedAfterDrop pins Codex round-7 P2: after a