keyviz: reuse retired slot on re-registration + test polish

Codex round-12 P2: Re-registering a removed RouteID inside the grace
window allocated a brand-new *routeSlot. Flush then drained both the
new live slot AND the old retired slot in the same column,
producing two rows with the same RouteID and splitting counts
between them. Add reclaimRetiredSlot — when RegisterRoute lands a
real (non-aggregate) routeID and a retired slot with that ID is
queued, pull it off retiredSlots and reuse it. The old counters and
any in-flight late-Observe writes converge with new traffic on the
same slot, and Flush emits a single row.
TestReRegisterIndividualReusesRetiredSlot asserts one row per
RouteID per column with correctly accumulated counts.

Claude bot nits:
- nextVirtualBucketID comment now spells out the atomic.Uint64
  wrap-around explicitly instead of the cryptic "MaxUint64+1" form.
- TestRejoinAsIndividualLetsBucketPruneFire was vacuously true —
  the bucket had no traffic so it never emitted a row, making the
  prune assertion unreachable. Restructure with two virtual members
  (so removing one leaves the bucket alive, not orphaned), drive
  traffic through the surviving member, and flush twice past grace
  so the post-prune MemberRoutes lands in a snapshot row.

Author: bootjp

keyviz/sampler.go

@@ -376,10 +376,25 @@ func (s *MemSampler) RegisterRoute(routeID uint64, start, end []byte) bool {
 
 	next := copyRouteTable(cur)
 	if len(next.slots) < s.opts.MaxTrackedRoutes {
-		slot := &routeSlot{
-			RouteID: routeID,
-			Start:   cloneBytes(start),
-			End:     cloneBytes(end),
+		slot := s.reclaimRetiredSlot(routeID)
+		if slot == nil {
+			slot = &routeSlot{
+				RouteID: routeID,
+				Start:   cloneBytes(start),
+				End:     cloneBytes(end),
+			}
+		} else {
+			// Re-registering the same routeID inside the grace window:
+			// reuse the retired slot so any in-flight Observe writers
+			// hitting the prior table snapshot land on the same slot
+			// the new traffic uses, and Flush emits a single row per
+			// RouteID instead of two (one live + one retired) with
+			// counts split across them. Refresh the metadata fields to
+			// match the new registration; counters are preserved.
+			slot.metaMu.Lock()
+			slot.Start = cloneBytes(start)
+			slot.End = cloneBytes(end)
+			slot.metaMu.Unlock()
 		}
 		next.slots[routeID] = slot
 		next.sortedSlots = appendSorted(next.sortedSlots, slot)
@@ -613,7 +628,36 @@ func (s *MemSampler) cancelPendingPruneFor(bucket *routeSlot, routeID uint64) {
 // aggregate, one individual — if the original route is later
 // re-registered as an individual slot.
 func (s *MemSampler) nextVirtualBucketID() uint64 {
-	return s.virtualIDCounter.Add(^uint64(0)) // subtract 1; counter starts at MaxUint64+1
+	// atomic.Uint64 starts at 0; adding ^uint64(0) wraps to MaxUint64
+	// on the first call, MaxUint64-1 on the second, etc.
+	return s.virtualIDCounter.Add(^uint64(0))
+}
+
+// reclaimRetiredSlot looks for a non-aggregate retired slot whose
+// RouteID matches the supplied routeID and, if found, removes it
+// from retiredSlots and returns it for reuse. This guarantees a
+// route removed and re-registered inside the grace window is
+// represented by a single *routeSlot — Flush would otherwise emit
+// two rows with the same RouteID (one from the new live slot, one
+// from the still-draining retired slot) and split counts across
+// them. Aggregate (orphaned virtual bucket) entries are left in
+// place because their RouteID lives in the synthetic namespace and
+// a real-ID match against an aggregate would be coincidental.
+func (s *MemSampler) reclaimRetiredSlot(routeID uint64) *routeSlot {
+	s.retiredMu.Lock()
+	defer s.retiredMu.Unlock()
+	var reclaimed *routeSlot
+	keep := s.retiredSlots[:0]
+	for _, r := range s.retiredSlots {
+		if reclaimed == nil && !r.slot.Aggregate && r.slot.RouteID == routeID {
+			reclaimed = r.slot
+			continue
+		}
+		keep = append(keep, r)
+	}
+	clearTail(s.retiredSlots, len(keep))
+	s.retiredSlots = keep
+	return reclaimed
 }
 
 // memberRoutesContains reports whether routeID is already listed in

keyviz/sampler_test.go

@@ -718,41 +718,88 @@ func TestVirtualBucketRouteIDIsSynthetic(t *testing.T) {
 }
 
 // TestRejoinAsIndividualLetsBucketPruneFire pins Codex round-11 P2:
-// when a removed virtual member rejoins as an individual slot (capacity
-// freed up), the deferred member-prune for the old bucket must still
-// execute — otherwise the bucket's MemberRoutes keeps a route that no
-// longer contributes traffic to it.
+// when a removed virtual member rejoins as an individual slot, the
+// deferred member-prune for the old bucket must still execute —
+// otherwise the bucket's MemberRoutes keeps a route that no longer
+// contributes traffic to it. Set up the bucket with two members so
+// removing one leaves the bucket alive (not orphaned), and traffic
+// continues to be drained across the grace boundary; the post-grace
+// row must show the prune actually happened.
 func TestRejoinAsIndividualLetsBucketPruneFire(t *testing.T) {
 	t.Parallel()
 	s, clk := newTestSampler(t, MemSamplerOptions{
 		Step:             time.Second,
-		HistoryColumns:   4,
-		MaxTrackedRoutes: 2,
+		HistoryColumns:   8,
+		MaxTrackedRoutes: 1,
 	})
 	mustRegister(t, s, 1, "a", "b")
-	mustRegister(t, s, 2, "m", "n")
+	if s.RegisterRoute(2, []byte("m"), []byte("n")) {
+		t.Fatal("route 2 should fold")
+	}
 	if s.RegisterRoute(3, []byte("y"), []byte("z")) {
-		t.Fatal("route 3 over budget should fold")
+		t.Fatal("route 3 should fold (same bucket)")
 	}
-	// Free capacity: remove route 2 (now slots have room) AND remove 3
-	// (queued in pendingPrunes). Re-register 3 — it now fits as an
-	// individual slot, so the old bucket should still get pruned.
+	s.Observe(2, OpRead, 0, 0)
+	s.Observe(3, OpRead, 0, 0)
+	// Free capacity (route 1) and remove the virtual member (route 3),
+	// which queues a prune against the bucket. Route 2 still keeps the
+	// bucket alive in virtualForRoute.
+	s.RemoveRoute(1)
 	s.RemoveRoute(3)
-	s.RemoveRoute(2)
 	if !s.RegisterRoute(3, []byte("y"), []byte("z")) {
-		t.Fatal("route 3 should now fit as individual slot")
+		t.Fatal("route 3 should fit individually now")
 	}
 
+	// Two flushes after grace: the first executes the prune, the
+	// second emits a row with the post-prune MemberRoutes.
 	clk.Advance(s.graceWindow() + time.Second)
-	s.Observe(3, OpRead, 0, 0)
+	s.Observe(2, OpRead, 0, 0)
+	s.Flush()
+	s.Observe(2, OpRead, 0, 0)
 	s.Flush()
 
 	rows := lastSnapshotColumn(t, s).Rows
+	agg := findAggregateRow(t, rows)
+	if memberRoutesContain(agg.MemberRoutes, 3) {
+		t.Fatalf("bucket still lists pruned route 3: %v", agg.MemberRoutes)
+	}
+	if !memberRoutesContain(agg.MemberRoutes, 2) {
+		t.Fatalf("bucket dropped still-active route 2: %v", agg.MemberRoutes)
+	}
+}
+
+// TestReRegisterIndividualReusesRetiredSlot pins Codex round-12 P2:
+// re-registering the same RouteID inside the grace window must reuse
+// the retired slot. Otherwise Flush emits two rows for the same
+// RouteID in one column (live + retired drain), splitting counts.
+func TestReRegisterIndividualReusesRetiredSlot(t *testing.T) {
+	t.Parallel()
+	s, _ := newTestSampler(t, MemSamplerOptions{
+		Step:           time.Second,
+		HistoryColumns: 4,
+	})
+	mustRegister(t, s, 1, "a", "b")
+	s.Observe(1, OpRead, 0, 0)
+	s.RemoveRoute(1)
+	mustRegister(t, s, 1, "a", "b")
+	s.Observe(1, OpRead, 0, 0)
+	s.Flush()
+
+	rows := lastSnapshotColumn(t, s).Rows
+	count := 0
+	var total uint64
 	for _, r := range rows {
-		if r.Aggregate && memberRoutesContain(r.MemberRoutes, 3) {
-			t.Fatalf("after rejoin-as-individual + grace, bucket still lists route 3: %+v", r)
+		if r.RouteID == 1 {
+			count++
+			total += r.Reads
 		}
 	}
+	if count != 1 {
+		t.Fatalf("got %d rows for RouteID 1 in one column; expected 1 — retired slot was not reclaimed", count)
+	}
+	if total != 2 {
+		t.Fatalf("total Reads for RouteID 1 = %d, want 2 (1 pre-remove + 1 post-rejoin)", total)
+	}
 }
 
 // TestReRegisterDuringPruneGraceCancelsPrune pins Codex round-9 P2: