feat(keyviz): sampler GroupID + LeaderTerm stamping (Phase 2-C+ PR-3a) (#709)

## Summary

- Add `GroupID` field to `routeSlot` and `RaftGroupID + LeaderTerm` to
`MatrixRow` (Phase 2-C+ §9.1 dedupe key)
- Add `MemSampler.SetLeaderTerm(groupID, term)` API and `groupTerms`
snapshot at Flush time
- Extend `RegisterRoute` with a `groupID uint64` parameter; update all
call sites (`main.go` reads `r.GroupID` from
`distribution.Engine.Stats()`)

## Why

Per `docs/design/2026_04_27_proposed_keyviz_cluster_fanout.md` §9.1, the
cluster fan-out aggregator dedupes write samples by `(routeID,
raftGroupID, leaderTerm, columnAt)` instead of the conservative
max-merge. Max-merge undercounts during a leadership flip when the new
leader observes a window the old leader was halfway through. Carrying
`RaftGroupID + LeaderTerm` on every row gives the merge enough
information to dedupe per-term and sum across terms.

## Scope (PR-3a — sampler API only)

This PR ships the sampler-side API. Wiring lands in follow-up PRs:

- **PR-3a (this)**: `routeSlot.GroupID`, `MatrixRow.RaftGroupID +
LeaderTerm`, `SetLeaderTerm` API, `RegisterRoute` signature, `Flush`
term snapshot.
- **PR-3b (follow-up)**: periodic ticker in `main.go` that polls engine
Status and calls `SetLeaderTerm`; proto + JSON wire-format additions.
- **PR-3c (follow-up)**: aggregator-side `(group, term)`-keyed dedupe in
`internal/admin/keyviz_fanout.go`.

With `SetLeaderTerm` never called (this PR alone), every row emits
`LeaderTerm=0` and the fan-out merge falls back to today's max-merge →
no behavior change for legacy deployments.

## Test plan

- [x] `go test -race -count=1 ./keyviz/...` — passes
- [x] `go test -race -count=1 ./internal/admin/...` — passes (uses
MemSampler)
- [x] `go test -race -count=1 .` (root, `main_keyviz_test.go`) — passes
- [x] `go build ./...` — clean
- [x] `golangci-lint run` — 0 issues
- [ ] Jepsen — N/A (sampler-internal change, no replication/MVCC impact)

## Five-lens self-review

1. **Data loss** — no on-disk format change; legacy max-merge fallback
preserves today's behavior when `SetLeaderTerm` is never called.
2. **Concurrency** — `groupTermsMu` is a fine-grained `RWMutex`;
`Observe` never touches it; `snapshotGroupTerms` clones at the top of
`Flush` so every row in a column observes a stable view, even if
`SetLeaderTerm` fires concurrently.
3. **Performance** — `Observe` hot path is unchanged (no new map lookup,
no new lock); `Flush` gains one `RLock` + map clone per column, bounded
by `len(groupTerms)` (typically ≤ a few groups).
4. **Data consistency** — `(routeID, raftGroupID, leaderTerm)` is a
strict superset of today's `routeID` dedupe key; rows with
`LeaderTerm=0` collapse to the legacy max-merge.
5. **Test coverage** — new `SetLeaderTerm` publishing + `Flush` snapshot
tests under `keyviz/sampler_test.go`; existing `RegisterRoute` test
signatures updated; race detector clean.

Author: bootjp

keyviz/flusher_test.go

@@ -9,7 +9,7 @@ import (
 func TestRunFlusherTicksUntilCancel(t *testing.T) {
 	t.Parallel()
 	s := NewMemSampler(MemSamplerOptions{Step: 5 * time.Millisecond, HistoryColumns: 16})
-	if !s.RegisterRoute(1, []byte("a"), []byte("b")) {
+	if !s.RegisterRoute(1, []byte("a"), []byte("b"), 0) {
 		t.Fatal("Register failed")
 	}
 	ctx, cancel := context.WithCancel(context.Background())

keyviz/sampler.go

@@ -160,6 +160,13 @@ type MemSampler struct {
 	// twice in the same column — once as an aggregate row, once as
 	// an individual row — even under register/remove churn.
 	virtualIDCounter atomic.Uint64
+
+	// groupTermsMu guards groupTerms. SetLeaderTerm and Flush both
+	// touch the map; Observe never does. The lock is fine-grained
+	// enough that contention is bounded by leader-term flips, which
+	// happen at most a few times per second across the whole cluster.
+	groupTermsMu sync.RWMutex
+	groupTerms   map[uint64]uint64
 }
 
 // retiredSlot tracks a removed slot through its post-removal grace
@@ -237,6 +244,14 @@ type routeTable struct {
 type routeSlot struct {
 	metaMu  sync.RWMutex
 	RouteID uint64
+	// GroupID is the Raft group this route belongs to. Phase 2-C+
+	// stamps this on every emitted MatrixRow so the cluster fan-out
+	// aggregator can dedupe write samples by (raftGroupID,
+	// leaderTerm) instead of the conservative max-merge that may
+	// undercount during a leadership flip. 0 means "no group
+	// attached" (legacy single-group deployments and the synthetic
+	// virtual-bucket slots both use 0).
+	GroupID uint64
 	Start   []byte
 	End     []byte
 	// Aggregate marks virtual buckets that fold multiple coarsened
@@ -266,7 +281,7 @@ type routeSlot struct {
 // Start/End/MemberRoutes with the live slot (which a later
 // RegisterRoute may extend, and which the snapshot API exports to
 // external consumers that may mutate the bounds).
-func (s *routeSlot) snapshotMeta() (start, end []byte, aggregate bool, members []uint64, membersTotal uint64) {
+func (s *routeSlot) snapshotMeta() (start, end []byte, aggregate bool, members []uint64, membersTotal, groupID uint64) {
 	s.metaMu.RLock()
 	defer s.metaMu.RUnlock()
 	start = cloneBytes(s.Start)
@@ -276,6 +291,7 @@ func (s *routeSlot) snapshotMeta() (start, end []byte, aggregate bool, members [
 		members = append([]uint64(nil), s.MemberRoutes...)
 	}
 	membersTotal = s.MemberRoutesTotal
+	groupID = s.GroupID
 	return
 }
 
@@ -299,6 +315,16 @@ type MatrixRow struct {
 	// prefix of this list when MemberRoutesTotal > len(MemberRoutes).
 	MemberRoutesTotal uint64
 
+	// RaftGroupID + LeaderTerm carry the route's Raft identity at the
+	// time the column was flushed. Stamped from the per-group term
+	// snapshot SetLeaderTerm publishes (Phase 2-C+ fan-out merge
+	// uses (RouteID/BucketID, RaftGroupID, LeaderTerm, columnAt) as
+	// the dedupe key). Zero values mean "term not tracked" — emitted
+	// when no SetLeaderTerm call has been made for the group yet, or
+	// for synthetic virtual-bucket slots that span groups.
+	RaftGroupID uint64
+	LeaderTerm  uint64
+
 	Reads      uint64
 	Writes     uint64
 	ReadBytes  uint64
@@ -331,14 +357,61 @@ func NewMemSampler(opts MemSamplerOptions) *MemSampler {
 		now = time.Now
 	}
 	s := &MemSampler{
-		opts:    opts,
-		now:     now,
-		history: newRingBuffer(opts.HistoryColumns),
+		opts:       opts,
+		now:        now,
+		history:    newRingBuffer(opts.HistoryColumns),
+		groupTerms: map[uint64]uint64{},
 	}
 	s.table.Store(newEmptyRouteTable())
 	return s
 }
 
+// SetLeaderTerm publishes the current Raft leader term for the given
+// group. Called by main.go on a periodic ticker that polls the engine
+// Status, and on every term-change observed via the engine. Phase 2-C+
+// stamps each MatrixRow's LeaderTerm from this snapshot at Flush time.
+//
+// Calling with term == 0 is allowed but treated as "term unknown"
+// during merge — the canonical (groupID, term) dedupe key collapses
+// to the legacy max-merge for cells whose LeaderTerm is 0. This lets
+// nodes that have not finished engine startup contribute partial data
+// without poisoning the merge.
+//
+// groupID == 0 is reserved for virtual aggregate buckets (which span
+// multiple real groups). Calls with groupID == 0 are silently ignored
+// so a future caller cannot accidentally stamp a non-zero term on
+// aggregate rows — they must remain LeaderTerm == 0 so the fan-out
+// merge falls back to max-merge for cross-group cells.
+//
+// nil-receiver-safe.
+func (s *MemSampler) SetLeaderTerm(groupID, term uint64) {
+	if s == nil || groupID == 0 {
+		return
+	}
+	s.groupTermsMu.Lock()
+	defer s.groupTermsMu.Unlock()
+	s.groupTerms[groupID] = term
+}
+
+// snapshotGroupTerms returns a copy of the per-group term map. Called
+// at the top of Flush so the column built below sees a stable view
+// of the term mapping even if SetLeaderTerm fires concurrently.
+// Returns nil when no terms have been published yet — `nil[k]` returns
+// the zero value in Go, so the row-builder hot path needs no nil guard
+// and avoids one allocation per Flush before SetLeaderTerm is wired.
+func (s *MemSampler) snapshotGroupTerms() map[uint64]uint64 {
+	s.groupTermsMu.RLock()
+	defer s.groupTermsMu.RUnlock()
+	if len(s.groupTerms) == 0 {
+		return nil
+	}
+	out := make(map[uint64]uint64, len(s.groupTerms))
+	for g, t := range s.groupTerms {
+		out[g] = t
+	}
+	return out
+}
+
 func newEmptyRouteTable() *routeTable {
 	return &routeTable{
 		slots:           map[uint64]*routeSlot{},
@@ -385,10 +458,18 @@ func (s *MemSampler) Observe(routeID uint64, op Op, keyLen, valueLen int) {
 }
 
 // RegisterRoute adds a (RouteID, [Start, End)) pair to the tracking
-// set. Returns true when the route gets its own slot, false when the
+// set. groupID is the Raft group the route belongs to; it is stamped
+// onto every MatrixRow this slot eventually emits so the Phase 2-C+
+// fan-out merge can dedupe write samples by (groupID, leaderTerm).
+// Pass 0 when the deployment doesn't use multiple groups (legacy /
+// single-group setups); the legacy max-merge fallback handles those.
+//
+// Returns true when the route gets its own slot, false when the
 // MaxTrackedRoutes cap was hit and the route was folded into a
 // virtual aggregate bucket. Idempotent: calling twice with the same
-// RouteID is a no-op (the original slot stays in place).
+// RouteID is a no-op (the original slot stays in place; a different
+// groupID on the second call is silently ignored — RegisterRoute is
+// not the right API to retag a slot, RemoveRoute + RegisterRoute is).
 //
 // If a previous RemoveRoute(routeID) queued a deferred member-prune
 // and the route is now re-registered into the SAME bucket inside the
@@ -397,7 +478,7 @@ func (s *MemSampler) Observe(routeID uint64, op Op, keyLen, valueLen int) {
 // attributing fresh traffic to it. Prunes for different buckets (or
 // when the route rejoins as an individual slot) are left alone so
 // the old bucket's MemberRoutes is correctly cleaned up.
-func (s *MemSampler) RegisterRoute(routeID uint64, start, end []byte) bool {
+func (s *MemSampler) RegisterRoute(routeID uint64, start, end []byte, groupID uint64) bool {
 	if s == nil {
 		return false
 	}
@@ -418,6 +499,7 @@ func (s *MemSampler) RegisterRoute(routeID uint64, start, end []byte) bool {
 		if slot == nil {
 			slot = &routeSlot{
 				RouteID:           routeID,
+				GroupID:           groupID,
 				Start:             cloneBytes(start),
 				End:               cloneBytes(end),
 				MemberRoutesTotal: 1,
@@ -431,6 +513,7 @@ func (s *MemSampler) RegisterRoute(routeID uint64, start, end []byte) bool {
 			// counts split across them. Refresh the metadata fields to
 			// match the new registration; counters are preserved.
 			slot.metaMu.Lock()
+			slot.GroupID = groupID
 			slot.Start = cloneBytes(start)
 			slot.End = cloneBytes(end)
 			slot.MemberRoutesTotal = 1
@@ -601,14 +684,21 @@ func (s *MemSampler) Flush() {
 		return
 	}
 	col := MatrixColumn{At: s.now()}
+	// Snapshot the per-group leader-term map once at the top of
+	// Flush so every row in this column sees a consistent view, even
+	// if SetLeaderTerm fires concurrently. Later rows can never
+	// observe a *newer* term than earlier rows in the same column,
+	// which preserves the merge-side invariant that all rows in a
+	// (groupID, columnAt) tuple share a single leaderTerm value.
+	terms := s.snapshotGroupTerms()
 	tbl := s.table.Load()
 	for _, slot := range tbl.sortedSlots {
-		col.Rows = appendDrainedRow(col.Rows, slot)
+		col.Rows = appendDrainedRow(col.Rows, slot, terms)
 	}
 
 	grace := s.graceWindow()
 	s.retiredMu.Lock()
-	s.retiredSlots = drainRetiredSlots(s.retiredSlots, &col.Rows, col.At, grace)
+	s.retiredSlots = drainRetiredSlots(s.retiredSlots, &col.Rows, col.At, grace, terms)
 	s.pendingPrunes = advancePendingPrunes(s.pendingPrunes, col.At, grace)
 	s.retiredMu.Unlock()
 
@@ -628,10 +718,10 @@ func (s *MemSampler) Flush() {
 // this final drain. The dropped tail of the backing array is zeroed
 // so released *routeSlot pointers do not stay GC-reachable through
 // the reused capacity.
-func drainRetiredSlots(retired []retiredSlot, rows *[]MatrixRow, now time.Time, grace time.Duration) []retiredSlot {
+func drainRetiredSlots(retired []retiredSlot, rows *[]MatrixRow, now time.Time, grace time.Duration, terms map[uint64]uint64) []retiredSlot {
 	keep := retired[:0]
 	for _, r := range retired {
-		*rows = appendDrainedRow(*rows, r.slot)
+		*rows = appendDrainedRow(*rows, r.slot, terms)
 		if now.Sub(r.retiredAt) < grace {
 			keep = append(keep, r)
 		}
@@ -825,17 +915,19 @@ func (s *MemSampler) HistoryColumns() int {
 // MatrixRow when any counter was non-zero. Idle slots are skipped.
 // Metadata is read under the slot's metaMu so a concurrent
 // RegisterRoute fold cannot race with the row materialisation.
-func appendDrainedRow(rows []MatrixRow, slot *routeSlot) []MatrixRow {
+func appendDrainedRow(rows []MatrixRow, slot *routeSlot, terms map[uint64]uint64) []MatrixRow {
 	reads := slot.reads.Swap(0)
 	writes := slot.writes.Swap(0)
 	readBytes := slot.readBytes.Swap(0)
 	writeBytes := slot.writeBytes.Swap(0)
 	if reads == 0 && writes == 0 && readBytes == 0 && writeBytes == 0 {
 		return rows
 	}
-	start, end, aggregate, members, membersTotal := slot.snapshotMeta()
+	start, end, aggregate, members, membersTotal, groupID := slot.snapshotMeta()
 	return append(rows, MatrixRow{
 		RouteID:           slot.RouteID,
+		RaftGroupID:       groupID,
+		LeaderTerm:        terms[groupID],
 		Start:             start,
 		End:               end,
 		Aggregate:         aggregate,