feat(adapter): expose keyviz heatmap via AdminServer.GetKeyVizMatrix (#646)

## Summary

- Implements `AdminServer.GetKeyVizMatrix` against the proto declared in
`proto/admin.proto` (already generated; no proto edits).
- New narrow `KeyVizSampler` interface in the adapter package (just
`Snapshot(from, to time.Time) []keyviz.MatrixColumn`), so production
wires `*keyviz.MemSampler` while tests pass an in-memory fake.
- `AdminServer.RegisterSampler` mirrors `RegisterGroup`. Without it,
`GetKeyVizMatrix` returns `codes.Unavailable` so callers can distinguish
"keyviz disabled on this node" from "no data yet" (which is a successful
empty response).
- Pivots the column-major `MatrixColumn` slice into the row-major proto
layout: one `KeyVizRow` per `RouteID` with values aligned to a parallel
`column_unix_ms` slice. `KeyVizSeries` selection picks the matching
per-row counter; `UNSPECIFIED` defaults to `Reads`.
- `bucket_id` encodes `route:<id>` for individual slots and
`virtual:<syntheticID>` for aggregate buckets. Aggregate rows carry
`MemberRoutes` verbatim through `route_ids` and `route_count`.

Implements the read-side half of
`docs/admin_ui_key_visualizer_design.md` §5.2 / §6. The dispatch-side
`Observe` wiring is in #645.

## Test plan

- [x] `TestGetKeyVizMatrixReturnsUnavailableWhenSamplerNotRegistered` —
verifies `codes.Unavailable` when no sampler registered.
- [x] `TestGetKeyVizMatrixPivotsColumnsToRows` — two-column / two-route
fixture, verifies the missing-row-becomes-zero contract.
- [x] `TestGetKeyVizMatrixSeriesSelection` — table-driven across all
five enum values including `UNSPECIFIED` defaulting to `Reads`.
- [x] `TestGetKeyVizMatrixEncodesAggregateBucket` — virtual bucket
layout (`bucket_id` prefix, `aggregate=true`, `route_ids`,
`route_count`).
- [x] `go test -race -count=1 -run TestGetKeyVizMatrix ./adapter/...`
clean.
- [x] `golangci-lint run ./adapter/...` clean.

Author: bootjp

adapter/admin_grpc.go

@@ -535,7 +535,24 @@ func (s *AdminServer) GetKeyVizMatrix(
 	to := unixMsToTime(req.GetToUnixMs())
 	cols := sampler.Snapshot(from, to)
 	pickValue := matrixSeriesPicker(req.GetSeries())
-	return matrixToProto(cols, pickValue, int(req.GetRows())), nil
+	return matrixToProto(cols, pickValue, clampRowBudget(int(req.GetRows()))), nil
+}
+
+// keyVizRowBudgetCap is the upper bound on the per-request rows
+// budget — design doc §4.1 caps rows at 1024 to bound server work
+// (sort + payload) for adversarial / over-large requests.
+const keyVizRowBudgetCap = 1024
+
+// clampRowBudget enforces design §4.1's upper bound. A request of 0
+// (or negative) means "no cap" and is preserved; anything past the
+// cap is silently clamped — clients asking for more rows than the
+// server is willing to render get the most rows the server will
+// render, not an error.
+func clampRowBudget(requested int) int {
+	if requested > keyVizRowBudgetCap {
+		return keyVizRowBudgetCap
+	}
+	return requested
 }
 
 // unixMsToTime converts a Unix-millisecond timestamp into a time.Time,
@@ -549,21 +566,22 @@ func unixMsToTime(ms int64) time.Time {
 }
 
 // matrixSeriesPicker returns a callback that extracts the requested
-// counter from a MatrixRow. KEYVIZ_SERIES_UNSPECIFIED (and READS)
-// fall through to Reads so a default-valued request still returns
-// something useful.
+// counter from a MatrixRow. KEYVIZ_SERIES_UNSPECIFIED falls through
+// to Writes per design doc §4.1 — write traffic is the primary
+// signal the heatmap is built around, and the read path is wired in
+// a follow-up phase.
 func matrixSeriesPicker(series pb.KeyVizSeries) func(keyviz.MatrixRow) uint64 {
 	switch series {
-	case pb.KeyVizSeries_KEYVIZ_SERIES_WRITES:
-		return func(r keyviz.MatrixRow) uint64 { return r.Writes }
+	case pb.KeyVizSeries_KEYVIZ_SERIES_READS:
+		return func(r keyviz.MatrixRow) uint64 { return r.Reads }
 	case pb.KeyVizSeries_KEYVIZ_SERIES_READ_BYTES:
 		return func(r keyviz.MatrixRow) uint64 { return r.ReadBytes }
 	case pb.KeyVizSeries_KEYVIZ_SERIES_WRITE_BYTES:
 		return func(r keyviz.MatrixRow) uint64 { return r.WriteBytes }
-	case pb.KeyVizSeries_KEYVIZ_SERIES_UNSPECIFIED, pb.KeyVizSeries_KEYVIZ_SERIES_READS:
-		return func(r keyviz.MatrixRow) uint64 { return r.Reads }
+	case pb.KeyVizSeries_KEYVIZ_SERIES_UNSPECIFIED, pb.KeyVizSeries_KEYVIZ_SERIES_WRITES:
+		return func(r keyviz.MatrixRow) uint64 { return r.Writes }
 	default:
-		return func(r keyviz.MatrixRow) uint64 { return r.Reads }
+		return func(r keyviz.MatrixRow) uint64 { return r.Writes }
 	}
 }
 
@@ -608,6 +626,14 @@ func matrixToProto(cols []keyviz.MatrixColumn, pick func(keyviz.MatrixRow) uint6
 // applyKeyVizRowBudget caps rows to budget by total activity per row
 // (sum of per-column values), preserving the top-N rows. budget <= 0
 // means "no cap."
+//
+// NOTE: design doc §5.5 specifies a "lexicographic walk + greedy
+// merge of low-activity adjacent ranges" algorithm — we simplify to
+// activity-descending truncation for Phase 1 because it covers the
+// common UI need (highlight hotspots) without needing the synthetic
+// virtual-bucket plumbing the merge requires. Phase 2 should swap
+// this for the spec'd merge so low-activity ranges become coarse
+// aggregates instead of being silently dropped.
 func applyKeyVizRowBudget(rows []*pb.KeyVizRow, budget int) []*pb.KeyVizRow {
 	if budget <= 0 || len(rows) <= budget {
 		return rows

adapter/admin_grpc_keyviz_test.go

@@ -126,7 +126,7 @@ func TestGetKeyVizMatrixSeriesSelection(t *testing.T) {
 		series pb.KeyVizSeries
 		want   uint64
 	}{
-		{"unspecified defaults to reads", pb.KeyVizSeries_KEYVIZ_SERIES_UNSPECIFIED, 11},
+		{"unspecified defaults to writes", pb.KeyVizSeries_KEYVIZ_SERIES_UNSPECIFIED, 22},
 		{"reads", pb.KeyVizSeries_KEYVIZ_SERIES_READS, 11},
 		{"writes", pb.KeyVizSeries_KEYVIZ_SERIES_WRITES, 22},
 		{"read_bytes", pb.KeyVizSeries_KEYVIZ_SERIES_READ_BYTES, 333},
@@ -214,6 +214,35 @@ func TestGetKeyVizMatrixSurfacesRouteCountTruncation(t *testing.T) {
 	require.Equal(t, []uint64{2, 3}, r.RouteIds)
 }
 
+// TestGetKeyVizMatrixClampsRowsBudgetToCap pins design §4.1's
+// upper-bound: rows requests above the keyVizRowBudgetCap are
+// silently clamped down to the cap so a pathological client cannot
+// force the server to materialise an unbounded payload.
+func TestGetKeyVizMatrixClampsRowsBudgetToCap(t *testing.T) {
+	t.Parallel()
+	rows := make([]keyviz.MatrixRow, keyVizRowBudgetCap+5)
+	for i := range rows {
+		idx := uint64(i + 1) //nolint:gosec // i is bounded by keyVizRowBudgetCap+5
+		rows[i] = keyviz.MatrixRow{
+			RouteID: idx,
+			Start:   []byte{byte(i / 256), byte(i % 256)},
+			End:     []byte{byte((i + 1) / 256), byte((i + 1) % 256)},
+			Writes:  idx,
+		}
+	}
+	srv := newAdminServerWithFakeSampler(t, []keyviz.MatrixColumn{{
+		At:   time.Unix(1_700_000_000, 0),
+		Rows: rows,
+	}})
+
+	resp, err := srv.GetKeyVizMatrix(context.Background(), &pb.GetKeyVizMatrixRequest{
+		Series: pb.KeyVizSeries_KEYVIZ_SERIES_WRITES,
+		Rows:   uint32(keyVizRowBudgetCap + 1000),
+	})
+	require.NoError(t, err)
+	require.Len(t, resp.Rows, keyVizRowBudgetCap, "rows must be clamped to keyVizRowBudgetCap")
+}
+
 // TestGetKeyVizMatrixHonorsRowsBudget pins Codex round-1 P1 on
 // PR #646: a request with rows=N must return at most N rows. We
 // stage 4 routes with distinct activity totals and request rows=2;

keyviz/sampler.go

@@ -234,11 +234,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
@@ -463,16 +468,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)
 	}
@@ -486,6 +487,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
@@ -727,12 +753,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()
@@ -746,6 +773,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