feat(main): wire keyviz sampler end-to-end (flags, RunFlusher, coordinator + admin)

[bootjp]

Summary

End-to-end keyviz wiring on each node binary. Stacked on top of #645 (coordinator WithSampler) and #646 (admin server RegisterSampler) — those two PRs are merged into this branch via merge commits, so this PR's incremental diff is main.go + main_keyviz_test.go.

• Three new flags: --keyvizEnabled (off by default — opt-in), --keyvizStep, --keyvizMaxTrackedRoutes, --keyvizMaxMemberRoutesPerSlot. Defaults pull from keyviz.Default* so flag help and code stay in sync.
• buildKeyVizSampler() returns *keyviz.MemSampler when enabled, nil otherwise — the single decision point for "keyviz on/off."
• seedKeyVizRoutes() copies the engine's startup route catalogue into the sampler. Route-watch propagation post-startup is a follow-up (Phase 3 in the design doc).
• startKeyVizFlusher() runs RunFlusher in the existing errgroup and calls one final Flush() after ctx fires so the in-progress step is harvested at graceful shutdown.
• Coordinator: chained .WithSampler(...) onto the existing .WithLeaseReadObserver(...) call so the dispatch hot path observes mutations.
• AdminServer: setupAdminService now takes the *MemSampler and only calls RegisterSampler when it's non-nil — operators with keyviz disabled get codes.Unavailable on GetKeyVizMatrix instead of a spurious empty success.

Implements docs/admin_ui_key_visualizer_design.md §5 / §10. After this lands the heatmap is fully usable end-to-end.

Test plan

• TestBuildKeyVizSamplerHonorsEnabledFlag — flag on/off contract.
• TestSeedKeyVizRoutesCopiesEngineCatalogue — engine.Stats() seed path.
• TestSeedKeyVizRoutesNoOpOnNilSampler — disabled-sampler safety.
• TestStartKeyVizFlusherReturnsAfterCancel — graceful shutdown drains the pre-cancel counters via the final Flush.
• go build ., go vet ., golangci-lint run ./... clean.
• go test -race -count=1 -run 'TestBuildKeyVizSampler|TestSeedKeyVizRoutes|TestStartKeyVizFlusher' . passes.

Merge order

1. Merge feat(kv): wire keyviz.Sampler into ShardedCoordinator dispatch path #645 (coordinator wiring).
2. Merge feat(adapter): expose keyviz heatmap via AdminServer.GetKeyVizMatrix #646 (admin server wiring).
3. Rebase this branch on main (the merge commits in this PR collapse to no-ops) and merge.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[coderabbitai]

Warning

Rate limit exceeded

@bootjp has exceeded the limit for the number of commits that can be reviewed per hour. Please wait 52 minutes and 31 seconds before requesting another review.

Your organization is not enrolled in usage-based pricing. Contact your admin to enable usage-based pricing to continue reviews beyond the rate limit, or try again in 52 minutes and 31 seconds.

⌛ How to resolve this issue?

After the wait time has elapsed, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

We recommend that you space out your commits to avoid hitting the rate limit.

🚦 How do rate limits work?

CodeRabbit enforces hourly rate limits for each developer per organization.

Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout.

Please see our FAQ for further information.

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: edb0b4a1-0eda-45c8-a606-7c66d2da0ab1

📥 Commits

Reviewing files that changed from the base of the PR and between 29011bb and 7e14751.

📒 Files selected for processing (7)

• adapter/admin_grpc.go
• adapter/admin_grpc_keyviz_test.go
• keyviz/sampler.go
• kv/sharded_coordinator.go
• kv/sharded_coordinator_sampler_test.go
• main.go
• main_keyviz_test.go

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feat/keyviz-main-wiring

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 7bc737ae79

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Seed keyviz routes after catalog snapshot is applied

seedKeyVizRoutes runs before setupDistributionCatalog, but setupDistributionCatalog calls distribution.EnsureCatalogSnapshot, which applies a catalog snapshot back into the engine (including durable non-zero RouteIDs). That means the sampler is seeded with pre-snapshot IDs (typically 0 from buildEngine), while later dispatches observe post-snapshot IDs, so Observe misses and drops samples for normal traffic. In practice this can leave keyviz mostly empty right after startup even when writes are flowing.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Apply the request row budget in GetKeyVizMatrix

This handler never reads req.GetRows(), so it always returns every row from the snapshot regardless of caller-requested resolution. On large route counts, this can blow up response size and processing time for admin queries that are supposed to request a bounded Y-axis (for example 256 rows), making the endpoint much more expensive than requested.

Useful? React with 👍 / 👎.

[bootjp]

@claude review

[claude]

Claude finished @bootjp's task in 2m 34s —— View job

---

Code Review — PR #647

Tasks

• Gather context and understand request
• Read changed files (main.go, main_keyviz_test.go, adapter/admin_grpc.go, kv/sharded_coordinator.go, test files)
• Review keyviz design doc for alignment
• Analyze correctness, concurrency, error handling, test coverage
• Post comprehensive review

---

Summary

The overall architecture is clean: single construction point (buildKeyVizSampler), nil-safe threading through coordinator and admin server, and the KeyVizSampler interface in adapter/ giving good test isolation. The comment/doc quality is high. However there is one functional correctness bug and one incomplete implementation that need addressing before merge.

---

Critical — Bug

Seed fires before EnsureCatalogSnapshot assigns real RouteIDs (main.go:292)

// main.go (current order)
seedKeyVizRoutes(sampler, cfg.engine)          // line 292 — RouteIDs are all 0 here
distCatalog, err := setupDistributionCatalog(…) // line 296 — EnsureCatalogSnapshot runs here

UpdateRoute (called during engine construction) deliberately leaves RouteID = 0 — see distribution/engine.go:123-135. EnsureCatalogSnapshot inside setupDistributionCatalog is what assigns non-zero RouteIDs (1, 2, …) and writes them back into the engine via engine.ApplySnapshot. So when seedKeyVizRoutes calls engine.Stats() at line 292, every route has RouteID = 0.

The sampler registers one slot under RouteID = 0. All coordinator-dispatched mutations arrive with RouteIDs 1, 2, … — which miss the slot and are silently dropped. The heatmap will be empty for the lifetime of the process until route-watch propagation lands (Phase 3).

The fix is one-line: move seedKeyVizRoutes to after setupDistributionCatalog returns successfully.

distCatalog, err := setupDistributionCatalog(ctx, runtimes, cfg.engine)
if err != nil {
    return err
}
seedKeyVizRoutes(sampler, cfg.engine)   // ← move here; engine now has real RouteIDs

Fix this →

This also invalidates the existing test TestSeedKeyVizRoutesCopiesEngineCatalogue: it calls engine.UpdateRoute (which leaves RouteID=0) and then observes the same RouteID=0, so it passes today but does not guard against this regression. After fixing the ordering, add a test that calls engine.ApplySnapshot with non-zero RouteIDs before seeding, confirming the sampler registers those IDs correctly.

---

Moderate — Incomplete

req.GetRows() budget never applied (adapter/admin_grpc.go:509-537)

The proto defines uint32 rows = 4 on GetKeyVizMatrixRequest as a Y-axis row budget, but GetKeyVizMatrix and matrixToProto never read it. On a cluster with, say, 50 000 routes, a caller requesting rows=256 for a dashboard gets back all 50 000 rows.

The missing enforcement point is in matrixToProto:

// after sortKeyVizRowsByStart(resp.Rows), add:
if req.GetRows() > 0 && len(resp.Rows) > int(req.GetRows()) {
    resp.Rows = resp.Rows[:req.GetRows()]
}

Since matrixToProto already sorts by Start, truncating at the end is deterministic. Alternatively, the cap could be applied before the row-major pivot loop to avoid building rows that will be discarded. Either way, the current behaviour violates the proto contract.

Fix this →

---

Minor

--keyvizHistoryColumns flag missing

buildKeyVizSampler hardcodes HistoryColumns at zero (falling back to DefaultHistoryColumns = 1440, 24 h at 60 s). Operators cannot tune the retention window without recompiling. Given that the other three options are all exposed as flags, this asymmetry seems unintentional. Not a blocker, but worth a follow-up flag or at least a // TODO comment.

Flusher docstring describes the wrong shutdown signal

The startKeyVizFlusher docstring says "calls one final Flush() after ctx fires so the in-progress step is harvested at graceful shutdown". The actual code does RunFlusher(ctx, s, s.Step()); s.Flush() — the Flush() runs after RunFlusher returns, not after the eg drains. That's fine semantically (it's what the comment means), but the phrasing "after ctx fires" slightly understates that the goroutine exits first, then Flush() is called synchronously before the goroutine return. The behaviour is correct; only the wording is slightly imprecise.

sortKeyVizRowsByStart tie-breaking on equal Start keys

The secondary sort key is BucketId (string order: "route:N" vs "virtual:N"). During a split, two routes can briefly share the same Start. The tie-break by BucketId string-sorts virtual buckets before real routes ("r" < "v"), which may or may not match what the UI expects. Worth calling out in the design doc comment.

---

What's good

• The KeyVizSampler interface in adapter/ (rather than a concrete dep on *keyviz.MemSampler) is the right call for testability, and the fakeKeyVizSampler in the test is a clean demonstration.
• keyVizSamplerForCoordinator correctly returns a nil interface (not a typed nil) so the hot-path if c.sampler == nil guard works as expected. The comment explaining this is valuable.
• RegisterSampler guarded by the existing groupsMu RWMutex is correct; concurrent RPC reads take RLock so there's no data race.
• The observeMutation placement (after route resolution, before group fan-out) is exactly right — it fires once per resolved (RouteID, key) pair regardless of cross-shard dispatch count.
• Test coverage for series selection, aggregate encoding, and the nil-sampler Unavailable signal is thorough.

---

[bootjp]

@claude review

[claude]

Claude finished @bootjp's task in 4m 17s —— View job

---

Code Review — PR #647

Tasks

• Gather context and understand request
• Read changed files (main.go, main_keyviz_test.go, adapter/admin_grpc.go, keyviz/sampler.go, kv/sharded_coordinator.go + test files)
• Analyze diff vs origin/main
• Five-lens review (data loss, concurrency, performance, consistency, test coverage)
• Post comprehensive review

---

Round-1 Fixes — Both Addressed ✅

Seed ordering (main.go:299–305): seedKeyVizRoutes is now correctly placed after setupDistributionCatalog, with a clear explanatory comment. The comment is excellent — it cites exactly why seeding earlier would register RouteID=0 and cause observations to miss.

Row budget (adapter/admin_grpc.go:603): applyKeyVizRowBudget is implemented, called in matrixToProto, and has a dedicated test (TestGetKeyVizMatrixHonorsRowsBudget). The activity-descending sort + truncate + final Start-key sort is the correct order for this operation.

---

Moderate — Read path not wired

observeMutation fires only from groupMutations (the write dispatch path). Reads — LeaseReadForKey, Coordinate.ReadForKey, and any follower-local read paths — never call Observe. The design doc §5.1 explicitly calls for read and write sampling:

"Writes are sampled exactly once by the current Raft leader before proposal. Reads are sampled by the node that actually serves the read."

And the Phase 2 success criterion (§10) includes "leader/follower reads." Since keyviz.OpRead is defined, MatrixRow.Reads is populated by the sampler, and the proto response includes KEYVIZ_SERIES_READS — operators enabling --keyvizEnabled today will get a write-only heatmap with the Reads series permanently zero. That mismatch with the design doc contract and the flag's description ("feeds AdminServer.GetKeyVizMatrix") should be called out explicitly.

If this is intentional scoping (writes-first in this phase, reads in the next), add a comment or a TODO in observeMutation and call it out in the PR description so reviewers and future contributors know Phase 2 is still open on this point.

Fix or document this →

---

Minor — Test gap in TestSeedKeyVizRoutesCopiesEngineCatalogue

The test uses engine.UpdateRoute(...) which leaves RouteID=0, then calls s.Observe(r.RouteID, keyviz.OpRead, 1, 1) against the same RouteID=0. This passes because seedKeyVizRoutes registered RouteID=0, and Observe(0, ...) finds it.

In production, EnsureCatalogSnapshot assigns RouteID=1 (first route) via routeDescriptorsFromEngine → routeID := uint64(1). The test doesn't guard the regression where a future change seeds before EnsureCatalogSnapshot (all IDs 0 again) — the test would still pass because it exercises the same 0 → 0 path throughout. Adding a subtest that calls engine.ApplySnapshot with a non-zero RouteID snapshot before seeding would close this gap. This is not blocking but was flagged in round-1 and remains open.

---

Minor — Missing --keyvizHistoryColumns flag

HistoryColumns is hardcoded to 0 (→ DefaultHistoryColumns=1440, 24 h). The other three tunables (--keyvizStep, --keyvizMaxTrackedRoutes, --keyvizMaxMemberRoutesPerSlot) are all flags. An operator who wants a shorter retention window (e.g., 2 h for a high-cardinality cluster) has no knob. Not a blocker for this PR, but add a // TODO: expose via --keyvizHistoryColumns comment in buildKeyVizSampler or open a follow-up issue.

---

Minor — startKeyVizFlusher always spawns a goroutine when sampler is nil

func startKeyVizFlusher(ctx context.Context, eg *errgroup.Group, s *keyviz.MemSampler) {
    eg.Go(func() error {
        keyviz.RunFlusher(ctx, s, s.Step()) // parks on <-ctx.Done() when s==nil
        s.Flush()                            // no-op when s==nil
        return nil
    })
}

When --keyvizEnabled=false, this goroutine just waits for context cancellation. The overhead is negligible (one parked goroutine), but a nil guard would make the "disabled" path explicit:

if s == nil {
    return
}

Not a blocker — the current code is correct and the docstring says "Nil-safe: a disabled sampler reduces the goroutine to ctx-wait + a no-op Flush."

---

What's good

• Both round-1 issues fixed and tested. TestGetKeyVizMatrixHonorsRowsBudget and the seed-ordering comment are well-written.
• keyVizSamplerForCoordinator returns an interface-nil (not a typed nil) so the if c.sampler == nil guard in observeMutation fires correctly. The comment explaining this invariant is valuable.
• applyKeyVizRowBudget ordering: activity-descending sort → truncate → re-sort by Start is correct and matches the test expectation.
• newKeyVizRowFrom correctly uses MemberRoutesTotal for RouteCount and sets RouteIdsTruncated when total > len(MemberRoutes). The TestGetKeyVizMatrixSurfacesRouteCountTruncation test is exactly right.
• RegisterSampler guarded by groupsMu — consistent with RegisterGroup and the read side in GetKeyVizMatrix. No data race.
• Seed ordering comment at main.go:299–305 is precise: it explains the exact mechanism (EnsureCatalogSnapshot → ApplySnapshot → non-zero IDs), names the failure mode (RouteID=0 → Observe misses), and references the Phase 3 follow-up. This is the kind of comment future contributors will thank you for.
• fakeKeyVizSampler in the test is a clean, deep-copying implementation that mirrors the real contract without standing up a full sampler.
• 9 new tests are well-structured with clear doc comments that pin both positive contracts and failure modes.

---

Summary

The PR is in mergeable shape. Both round-1 findings are fixed. The read-path gap is the most substantive remaining item — it's either a deliberate phase boundary (in which case document it) or an oversight (in which case wire the read paths before merge). The other items are low-priority follow-ups.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 7e1475104a

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Decrement total for pruned routes beyond visible cap

When a virtual bucket exceeds MaxMemberRoutesPerSlot, removing one of the hidden members will enqueue a prune but pruneMemberRoute only decrements MemberRoutesTotal if the route was present in MemberRoutes. In that common churn path removed is false, so the total never drops, and GetKeyVizMatrix later reports inflated route_count (and stale truncation) for that bucket. This makes the admin heatmap metadata drift over time even though the route has been removed.

Useful? React with 👍 / 👎.

[bootjp]

@claude review

[claude]

Claude finished @bootjp's task in 4m 32s —— View job

---

Code Review — PR #647 (Round 3)

Tasks

• Gather context from PR and prior review rounds
• Read all 7 changed files + supporting keyviz helpers (flusher.go, ring_buffer.go)
• Five-lens review pass (data loss, concurrency, perf, consistency, coverage)
• Post comprehensive review

---

Round 1+2 Fixes — Confirmed ✅

Seed ordering (main.go:299–306): seedKeyVizRoutes correctly placed after setupDistributionCatalog with a precise comment explaining the RouteID=0 failure mode. ✅

Row budget (adapter/admin_grpc.go:603): applyKeyVizRowBudget implemented, called in matrixToProto, covered by TestGetKeyVizMatrixHonorsRowsBudget. ✅

---

Moderate — MemberRoutesTotal drift on past-cap route churn

This is the Codex P2 finding (keyviz/sampler.go:751). The code comment at line 735 partially acknowledges it:

"routes pruned past the visible cap stay in the total because we don't track individual past-cap members"

But the drift is worse than the comment implies — re-registration also double-counts:

1. Route A registered above MaxTrackedRoutes cap → MemberRoutesTotal++, not appended to MemberRoutes (capped)
2. Route A removed → deferred pruneMemberRoute fires after grace window → removed=false (route never was in MemberRoutes) → MemberRoutesTotal unchanged (should decrement)
3. Route A re-registered → memberRoutesContains returns false (not in MemberRoutes) → MemberRoutesTotal++ again

Each remove+re-register cycle for a past-cap route increments MemberRoutesTotal by 1 permanently. On a cluster where routes above the cap churn (splits, migrations), MemberRoutesTotal grows without bound:

• route_count in the heatmap becomes permanently inflated
• route_ids_truncated stays true forever even after all past-cap routes have departed
• Operators drilling into hot buckets see wildly wrong attribution counts

A conservative fix for the prune side: maintain a pastCapMembers map[uint64]struct{} alongside MemberRoutes inside routeSlot, so pruneMemberRoute can detect and decrement past-cap removals. This adds memory per virtual bucket but bounds the drift. A simpler band-aid: cap MemberRoutesTotal at len(MemberRoutes) + some_sentinel and document the approximation clearly.

Fix or track this →

This was flagged by Codex in round 2 but not addressed before merge. Follow-up PR recommended.

---

Minor — applyKeyVizRowBudget uses unstable sort

// adapter/admin_grpc.go:615
sort.Slice(rows, func(i, j int) bool {
    return rowActivityTotal(rows[i]) > rowActivityTotal(rows[j])
})

sort.Slice is not guaranteed to be stable. When two routes have identical activity totals, the selection is non-deterministic across calls. On a dashboard polling GetKeyVizMatrix at regular intervals, this can cause "flickering" — different route sets returned for the same time window depending on internal sort order.

One-word fix: sort.SliceStable.

Fix this →

---

Still Open from Round 2

Read path not wired: observeMutation at kv/sharded_coordinator.go:978 includes a comment "reads never reach this path" — this is correct as a writes-only Phase 2 MVP. But the design doc §5.1 includes reads in the contract. The observeMutation comment should explicitly note this is deferred to Phase 3 (or equivalent) rather than just saying reads don't reach here, so the limitation is visible to future contributors.

Missing --keyvizHistoryColumns flag: HistoryColumns is hardcoded to DefaultHistoryColumns=1440 (24 h). All other tunables are flags. Add a // TODO: expose via --keyvizHistoryColumns in buildKeyVizSampler if not exposing in this PR.

---

What's confirmed correct

• keyVizSamplerForCoordinator returns an interface-nil (not a typed-nil) so c.sampler == nil fires correctly on the hot path. Comment explaining this is exactly right.
• Lock ordering: routesMu → retiredMu in RegisterRoute/RemoveRoute; Flush takes only retiredMu then historyMu. No cycle, no deadlock.
• COW routeTable with atomic.Pointer.Store: Observe is a single Load + map lookup with no mutex on the hot path. Mutations build a fresh copy under routesMu, then publish atomically.
• RunFlusher nil-safe (if s == nil { <-ctx.Done(); return }). startKeyVizFlusher with a nil sampler reduces to a parked goroutine that exits on cancel — correct and documented.
• cancelPendingPruneFor acquires retiredMu while holding routesMu, which is the consistent order everywhere. Flush never touches routesMu. No deadlock.
• rebuildSorted uses a seen map to deduplicate virtual-bucket pointers so slots don't appear twice in sortedSlots when both slots and virtualForRoute reference the same *routeSlot. Correct.
• ringBuffer.Range calls cloneColumn so the returned slice is independent of the stored history — concurrent flushing cannot corrupt a Snapshot response in flight.
• 9 tests well-structured with clear doc comments pinning both positive contracts and failure modes. fakeKeyVizSampler correctly deep-copies rows so test assertions don't share memory with stored state.

---

Summary

Both round-1 critical issues are fixed and well-tested. The main open item is the MemberRoutesTotal drift for past-cap virtual bucket members — the code acknowledges the prune side but the re-registration double-count makes it compound on every churn cycle. This needs a follow-up PR before keyviz is used in any high-cardinality production deployment. The sort.SliceStable fix in applyKeyVizRowBudget is a trivial improvement worth applying at any time.