diff --git a/internal/admin/keyviz_fanout.go b/internal/admin/keyviz_fanout.go index 18fb369cd..c2b4a2b0d 100644 --- a/internal/admin/keyviz_fanout.go +++ b/internal/admin/keyviz_fanout.go @@ -449,7 +449,7 @@ func buildKeyVizPeerURL(peer string, params keyVizParams) (string, error) { if base.Host == "" { return "", fmt.Errorf("%w: peer base url %q has no host", errKeyVizPeer, peer) } - base.Path = "/admin/api/v1/keyviz/matrix" + base.Path = pathKeyVizMatrix q := base.Query() q.Set("series", string(params.series)) q.Set("rows", strconv.Itoa(params.rows)) diff --git a/internal/admin/keyviz_hotkeys_handler.go b/internal/admin/keyviz_hotkeys_handler.go new file mode 100644 index 000000000..a87ad6696 --- /dev/null +++ b/internal/admin/keyviz_hotkeys_handler.go @@ -0,0 +1,416 @@ +package admin + +import ( + "bytes" + "encoding/base64" + "errors" + "log/slog" + "net/http" + "sort" + "strconv" + "time" + + "github.com/bootjp/elastickv/keyviz" + "github.com/goccy/go-json" +) + +// KeyVizHotKeysSource is the narrow contract the hot-keys drill-down +// handler depends on. Production wires a *keyviz.MemSampler; tests use a +// stub that returns canned snapshots / bounds. +// +// HotKeysSnapshot must be safe to call concurrently with the +// aggregator's publish path (atomic.Pointer.Load — see keyviz/hot_keys.go). +// SubBucketBoundsFor returns ok=false when the route is unknown, an +// aggregate, or the index is out of range; nil hi means unbounded tail. +type KeyVizHotKeysSource interface { + HotKeysOptions() (enabled bool, capacity, sampleRate, maxKeyLen int) + HotKeysSnapshot(routeID uint64) *keyviz.KeyvizHotKeysSnapshot + SubBucketBoundsFor(routeID uint64, subBucket int) (lo, hi []byte, ok bool) +} + +// KeyVizHotKeysHandler serves GET /admin/api/v1/keyviz/hotkeys. +// +// The route is mounted on the admin auth chain (SessionAuth + audit) and +// further gated by the full-access role at the router — design §7 +// (privacy/auth) requires both because the response carries actual user +// key bytes. Returns 503 keyviz_disabled / hotkeys_disabled separately +// so the SPA can render a precise "feature off" message rather than a +// generic 404. +type KeyVizHotKeysHandler struct { + source KeyVizHotKeysSource + // snapshotWindow is the ±duration around SnapshotAt that the + // out_of_snapshot_window check accepts (Codex P2 round-5 L224). + // Defaults to keyviz.DefaultStep; tests can override via WithSnapshotWindow. + snapshotWindow time.Duration + now func() time.Time + logger *slog.Logger +} + +func NewKeyVizHotKeysHandler(source KeyVizHotKeysSource) *KeyVizHotKeysHandler { + return &KeyVizHotKeysHandler{ + source: source, + snapshotWindow: keyviz.DefaultStep, + now: func() time.Time { return time.Now().UTC() }, + logger: slog.Default(), + } +} + +// WithLogger overrides the slog destination so main.go can attach a +// component tag without changing the constructor signature. Nil is a no-op. +func (h *KeyVizHotKeysHandler) WithLogger(l *slog.Logger) *KeyVizHotKeysHandler { + if l == nil { + return h + } + h.logger = l + return h +} + +// WithClock lets tests inject a deterministic SnapshotAt for the +// out-of-window check. Nil is a no-op. +func (h *KeyVizHotKeysHandler) WithClock(now func() time.Time) *KeyVizHotKeysHandler { + if now == nil { + return h + } + h.now = now + return h +} + +// WithSnapshotWindow overrides the default keyvizStep used for the +// out_of_snapshot_window 400 (design §5). Non-positive is treated as a +// no-op so the default stays in force. +func (h *KeyVizHotKeysHandler) WithSnapshotWindow(d time.Duration) *KeyVizHotKeysHandler { + if d <= 0 { + return h + } + h.snapshotWindow = d + return h +} + +const ( + hotKeysDefaultTop = 20 + hotKeysMaxTopParam = 1024 +) + +// hotKeyResponse mirrors the design §5 wire shape. Counts are +// scaled-to-true-frequency estimates (sampled-counter × sample_rate), +// `degraded` is true iff drops > 0 OR skipped_long_keys > 0 (§6 fan-out +// merge OR), and `error_bound` is the per-node Space-Saving bound in +// scaled units (`sample_rate × sampled_n / m`). +type hotKeyResponse struct { + RouteID uint64 `json:"route_id"` + SubBucket *int `json:"sub_bucket,omitempty"` + Series string `json:"series"` + Keys []hotKeyResponseEntry `json:"keys"` + Approximate bool `json:"approximate"` + SampleRate int `json:"sample_rate"` + SampledN uint64 `json:"sampled_n"` + DroppedSamples uint64 `json:"dropped_samples"` + SkippedLongKeys uint64 `json:"skipped_long_keys"` + Degraded bool `json:"degraded"` + ErrorBound uint64 `json:"error_bound"` + SnapshotAt time.Time `json:"snapshot_at"` +} + +type hotKeyResponseEntry struct { + KeyB64 string `json:"key_b64"` + Count uint64 `json:"count"` +} + +func (h *KeyVizHotKeysHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + writeJSONError(w, http.StatusMethodNotAllowed, "method_not_allowed", "only GET") + return + } + if h.source == nil { + writeJSONError(w, http.StatusServiceUnavailable, "keyviz_disabled", + "key visualizer sampler is not configured on this node") + return + } + enabled, capacity, _, _ := h.source.HotKeysOptions() + if !enabled { + writeJSONError(w, http.StatusServiceUnavailable, "hotkeys_disabled", + "hot-keys drill-down is not enabled on this node (--keyvizHotKeysEnabled)") + return + } + + params, err := parseHotKeysParams(r, capacity) + if err != nil { + writeJSONError(w, http.StatusBadRequest, "invalid_query", err.Error()) + return + } + + snap := h.source.HotKeysSnapshot(params.routeID) + if snap == nil { + writeJSONError(w, http.StatusNotFound, "no_snapshot", + "no hot-keys snapshot is available for this route yet (it may be an aggregate, or no traffic has flushed since the feature was enabled)") + return + } + + if errCode := h.checkSnapshotWindow(params, snap); errCode != "" { + writeJSONError(w, http.StatusBadRequest, errCode, + "requested time window does not overlap the current snapshot; v1 returns only the current keyvizStep window") + return + } + + entries, errCode, errMsg := h.collectEntries(params, snap) + if errCode != "" { + writeJSONError(w, http.StatusBadRequest, errCode, errMsg) + return + } + + out := buildHotKeysResponse(params, snap, entries) + h.audit(r, params, len(out.Keys), out.Degraded) + w.Header().Set("Content-Type", "application/json; charset=utf-8") + w.Header().Set("Cache-Control", "no-store") + w.WriteHeader(http.StatusOK) + if err := json.NewEncoder(w).Encode(out); err != nil { + h.logger.LogAttrs(r.Context(), slog.LevelWarn, "admin keyviz hotkeys response encode failed", + slog.String("error", err.Error()), + ) + } +} + +// checkSnapshotWindow returns "" when the requested [from,to] window is +// acceptable, or the error_code string to send a 400 with otherwise. v1 +// rejects a window that does not overlap [SnapshotAt - snapshotWindow, +// SnapshotAt] (design §5; Codex P2 round-5 L224). Omitted bounds default +// to the snapshot window so a request with no time params always passes. +func (h *KeyVizHotKeysHandler) checkSnapshotWindow(p hotKeysParams, snap *keyviz.KeyvizHotKeysSnapshot) string { + if p.fromUnixMs == 0 && p.toUnixMs == 0 { + return "" + } + windowHigh := snap.SnapshotAt + windowLow := windowHigh.Add(-h.snapshotWindow) + reqLow := windowLow + reqHigh := windowHigh + if p.fromUnixMs != 0 { + reqLow = time.UnixMilli(p.fromUnixMs).UTC() + } + if p.toUnixMs != 0 { + reqHigh = time.UnixMilli(p.toUnixMs).UTC() + } + if reqHigh.Before(windowLow) || reqLow.After(windowHigh) { + return "out_of_snapshot_window" + } + return "" +} + +// collectEntries filters the snapshot entries to the requested sub-range +// (when sub_bucket is set), sorts descending by count, truncates to top. +// Returns (entries, "", "") on success or ("", errCode, errMsg) on a +// validation error (e.g. out-of-range sub_bucket). +func (h *KeyVizHotKeysHandler) collectEntries(p hotKeysParams, snap *keyviz.KeyvizHotKeysSnapshot) ([]keyviz.KeyvizHotKeyEntry, string, string) { + entries := snap.Entries + if p.subBucketSet { + lo, hi, ok := h.source.SubBucketBoundsFor(p.routeID, p.subBucket) + if !ok { + return nil, "invalid_query", "sub_bucket is out of range for this route" + } + filtered := make([]keyviz.KeyvizHotKeyEntry, 0, len(entries)) + for _, e := range entries { + if bytes.Compare(e.Key, lo) < 0 { + continue + } + if hi != nil && bytes.Compare(e.Key, hi) >= 0 { + continue + } + filtered = append(filtered, e) + } + entries = filtered + } else { + // Copy the snapshot's slice header so the sort below does not + // reorder the immutable published snapshot in place (a reader + // holding the snapshot reasonably expects it not to mutate + // under their feet). + entries = append([]keyviz.KeyvizHotKeyEntry(nil), entries...) + } + sort.Slice(entries, func(i, j int) bool { return entries[i].Count > entries[j].Count }) + if len(entries) > p.top { + entries = entries[:p.top] + } + return entries, "", "" +} + +// buildHotKeysResponse formats the wire body. count is scaled by the +// snapshot's sample_rate; error_bound is the same scaled +// (sample_rate × sampled_n / m) per-node bound from design §5. +func buildHotKeysResponse(p hotKeysParams, snap *keyviz.KeyvizHotKeysSnapshot, entries []keyviz.KeyvizHotKeyEntry) hotKeyResponse { + out := hotKeyResponse{ + RouteID: p.routeID, + Series: p.series, + Keys: make([]hotKeyResponseEntry, len(entries)), + Approximate: true, + SampleRate: snap.SampleRate, + SampledN: snap.SampledN, + DroppedSamples: snap.DroppedSamples, + SkippedLongKeys: snap.SkippedLongKeys, + Degraded: snap.DroppedSamples > 0 || snap.SkippedLongKeys > 0, + ErrorBound: scaledErrorBound(snap), + SnapshotAt: snap.SnapshotAt.UTC(), + } + if p.subBucketSet { + sb := p.subBucket + out.SubBucket = &sb + } + scale := positiveAsUint64(snap.SampleRate) + for i, e := range entries { + out.Keys[i] = hotKeyResponseEntry{ + KeyB64: base64.StdEncoding.EncodeToString(e.Key), + Count: e.Count * scale, + } + } + return out +} + +// scaledErrorBound is sample_rate × sampled_n / m, in scaled +// (true-frequency-estimate) units. Capacity is clamped at >=1 by +// NewMemSampler, so the division is safe; the guards below are +// belt-and-suspenders for snapshots produced by a future code path. +// (Local m, not `cap`, to avoid shadowing the builtin per gocritic.) +func scaledErrorBound(snap *keyviz.KeyvizHotKeysSnapshot) uint64 { + sr := positiveAsUint64(snap.SampleRate) + m := positiveAsUint64(snap.Capacity) + if m == 0 { + return 0 + } + return sr * snap.SampledN / m +} + +// positiveAsUint64 converts an int that is known to be non-negative +// (sampler-clamped to ≥1 in practice) to uint64 with a defensive guard +// so gosec G115 has no negative-wrap path to flag. +func positiveAsUint64(n int) uint64 { + if n < 0 { + return 0 + } + return uint64(n) +} + +func (h *KeyVizHotKeysHandler) audit(r *http.Request, p hotKeysParams, returned int, degraded bool) { + attrs := []slog.Attr{ + slog.String("component", "admin.keyviz.hotkeys"), + slog.String("method", r.Method), + slog.String("path", r.URL.Path), + slog.Uint64("route_id", p.routeID), + slog.Int("top", p.top), + slog.String("series", p.series), + slog.Int("returned", returned), + slog.Bool("degraded", degraded), + } + if p.subBucketSet { + attrs = append(attrs, slog.Int("sub_bucket", p.subBucket)) + } + if principal, ok := PrincipalFromContext(r.Context()); ok { + attrs = append(attrs, slog.String("principal", principal.AccessKey)) + } + h.logger.LogAttrs(r.Context(), slog.LevelInfo, "admin_audit", attrs...) +} + +type hotKeysParams struct { + routeID uint64 + subBucket int + subBucketSet bool + series string + top int + fromUnixMs int64 + toUnixMs int64 +} + +func parseHotKeysParams(r *http.Request, capacity int) (hotKeysParams, error) { + q := r.URL.Query() + routeID, err := parseRouteIDParam(q.Get("route_id")) + if err != nil { + return hotKeysParams{}, err + } + p := hotKeysParams{routeID: routeID, series: "writes", top: hotKeysDefaultTop} + if err := parseSubBucketParam(q.Get("sub_bucket"), &p); err != nil { + return hotKeysParams{}, err + } + if err := parseSeriesParam(q.Get("series"), &p); err != nil { + return hotKeysParams{}, err + } + if err := parseTopParam(q.Get("top"), capacity, &p); err != nil { + return hotKeysParams{}, err + } + if err := parseTimeWindowParams(q.Get("from_unix_ms"), q.Get("to_unix_ms"), &p); err != nil { + return hotKeysParams{}, err + } + return p, nil +} + +func parseRouteIDParam(raw string) (uint64, error) { + if raw == "" { + return 0, errors.New("route_id is required") + } + v, err := strconv.ParseUint(raw, 10, 64) + if err != nil { + return 0, errors.New("route_id must be an unsigned 64-bit integer") + } + return v, nil +} + +func parseSubBucketParam(raw string, p *hotKeysParams) error { + if raw == "" { + return nil + } + sb, err := strconv.Atoi(raw) + if err != nil || sb < 0 { + return errors.New("sub_bucket must be a non-negative integer") + } + p.subBucket = sb + p.subBucketSet = true + return nil +} + +func parseSeriesParam(raw string, p *hotKeysParams) error { + if raw == "" { + return nil + } + if raw != "writes" { + return errors.New("series: v1 supports only 'writes' (design §5)") + } + p.series = raw + return nil +} + +func parseTopParam(raw string, capacity int, p *hotKeysParams) error { + if raw != "" { + n, err := strconv.Atoi(raw) + if err != nil || n <= 0 { + return errors.New("top must be a positive integer") + } + if n > hotKeysMaxTopParam { + n = hotKeysMaxTopParam + } + p.top = n + } + if capacity > 0 && p.top > capacity { + p.top = capacity + } + return nil +} + +func parseTimeWindowParams(fromRaw, toRaw string, p *hotKeysParams) error { + if fromRaw != "" { + v, err := strconv.ParseInt(fromRaw, 10, 64) + if err != nil { + return errors.New("from_unix_ms must be an integer (unix milliseconds)") + } + p.fromUnixMs = v + } + if toRaw != "" { + v, err := strconv.ParseInt(toRaw, 10, 64) + if err != nil { + return errors.New("to_unix_ms must be an integer (unix milliseconds)") + } + p.toUnixMs = v + } + // Reject an inverted window before checkSnapshotWindow ever runs: + // a `to < from` interval makes the overlap check ambiguous and + // papers over what is almost certainly a client bug (gemini medium + // / claude 🔵 on PR #854). + if p.fromUnixMs != 0 && p.toUnixMs != 0 && p.fromUnixMs > p.toUnixMs { + return errors.New("from_unix_ms must be <= to_unix_ms") + } + return nil +} diff --git a/internal/admin/keyviz_hotkeys_handler_test.go b/internal/admin/keyviz_hotkeys_handler_test.go new file mode 100644 index 000000000..0b2ea2c4d --- /dev/null +++ b/internal/admin/keyviz_hotkeys_handler_test.go @@ -0,0 +1,358 @@ +package admin + +import ( + "bytes" + "context" + "encoding/base64" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/bootjp/elastickv/keyviz" + "github.com/goccy/go-json" + "github.com/stretchr/testify/require" +) + +// stubHotKeysSource mocks the sampler so handler tests don't have to +// drive a real *keyviz.MemSampler (clock + goroutine + tick). Setting +// disabled=true reproduces the --keyvizHotKeysEnabled=false path. +type stubHotKeysSource struct { + enabled bool + capacity int + sampleRate int + maxKeyLen int + snapshots map[uint64]*keyviz.KeyvizHotKeysSnapshot + bounds map[boundsKey]boundsValue +} + +type boundsKey struct { + routeID uint64 + subBucket int +} + +type boundsValue struct { + lo, hi []byte + ok bool +} + +func (s *stubHotKeysSource) HotKeysOptions() (bool, int, int, int) { + return s.enabled, s.capacity, s.sampleRate, s.maxKeyLen +} + +func (s *stubHotKeysSource) HotKeysSnapshot(routeID uint64) *keyviz.KeyvizHotKeysSnapshot { + return s.snapshots[routeID] +} + +func (s *stubHotKeysSource) SubBucketBoundsFor(routeID uint64, subBucket int) ([]byte, []byte, bool) { + v, ok := s.bounds[boundsKey{routeID, subBucket}] + if !ok { + return nil, nil, false + } + return v.lo, v.hi, v.ok +} + +func newStubSource() *stubHotKeysSource { + return &stubHotKeysSource{ + enabled: true, + capacity: 64, + sampleRate: 16, + maxKeyLen: 1024, + snapshots: map[uint64]*keyviz.KeyvizHotKeysSnapshot{}, + bounds: map[boundsKey]boundsValue{}, + } +} + +func newHandler(t *testing.T, src KeyVizHotKeysSource) http.Handler { + t.Helper() + h := NewKeyVizHotKeysHandler(src). + WithLogger(slog.New(slog.NewJSONHandler(&strings.Builder{}, nil))) + return h +} + +func TestHotKeysHandler_DisabledKeyVizReturns503(t *testing.T) { + t.Parallel() + h := newHandler(t, nil) + rec := serve(t, h, "GET", "/admin/api/v1/keyviz/hotkeys?route_id=1", nil) + require.Equal(t, http.StatusServiceUnavailable, rec.Code) + require.Contains(t, rec.Body.String(), "keyviz_disabled") +} + +func TestHotKeysHandler_DisabledHotKeysReturns503(t *testing.T) { + t.Parallel() + src := newStubSource() + src.enabled = false + h := newHandler(t, src) + rec := serve(t, h, "GET", "/admin/api/v1/keyviz/hotkeys?route_id=1", nil) + require.Equal(t, http.StatusServiceUnavailable, rec.Code) + require.Contains(t, rec.Body.String(), "hotkeys_disabled") +} + +func TestHotKeysHandler_MissingRouteIDReturns400(t *testing.T) { + t.Parallel() + h := newHandler(t, newStubSource()) + rec := serve(t, h, "GET", "/admin/api/v1/keyviz/hotkeys", nil) + require.Equal(t, http.StatusBadRequest, rec.Code) + require.Contains(t, rec.Body.String(), "route_id is required") +} + +func TestHotKeysHandler_InvalidRouteIDReturns400(t *testing.T) { + t.Parallel() + h := newHandler(t, newStubSource()) + rec := serve(t, h, "GET", "/admin/api/v1/keyviz/hotkeys?route_id=not-a-number", nil) + require.Equal(t, http.StatusBadRequest, rec.Code) +} + +func TestHotKeysHandler_NoSnapshotReturns404(t *testing.T) { + t.Parallel() + src := newStubSource() // no snapshots inserted + h := newHandler(t, src) + rec := serve(t, h, "GET", "/admin/api/v1/keyviz/hotkeys?route_id=42", nil) + require.Equal(t, http.StatusNotFound, rec.Code) + require.Contains(t, rec.Body.String(), "no_snapshot") +} + +func TestHotKeysHandler_WholeRouteReturnsScaledTopN(t *testing.T) { + t.Parallel() + src := newStubSource() + src.snapshots[1] = &keyviz.KeyvizHotKeysSnapshot{ + RouteID: 1, + SampledN: 100, + SampleRate: 16, + Capacity: 64, + SnapshotAt: time.Date(2026, 5, 28, 12, 0, 0, 0, time.UTC), + Entries: []keyviz.KeyvizHotKeyEntry{ + {Key: []byte("a"), Count: 1}, + {Key: []byte("hot"), Count: 50}, + {Key: []byte("middling"), Count: 5}, + }, + } + h := newHandler(t, src) + rec := serve(t, h, "GET", "/admin/api/v1/keyviz/hotkeys?route_id=1&top=2", nil) + require.Equal(t, http.StatusOK, rec.Code) + + var resp hotKeyResponse + require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp)) + require.Equal(t, uint64(1), resp.RouteID) + require.Nil(t, resp.SubBucket) + require.Equal(t, "writes", resp.Series) + require.True(t, resp.Approximate) + require.Equal(t, 16, resp.SampleRate) + require.Equal(t, uint64(100), resp.SampledN) + require.False(t, resp.Degraded) + // error_bound = 16 × 100 / 64 = 25 + require.Equal(t, uint64(25), resp.ErrorBound) + + require.Len(t, resp.Keys, 2, "top=2 truncates") + // Sorted descending by count. counts scaled by sample_rate (16): + require.Equal(t, base64.StdEncoding.EncodeToString([]byte("hot")), resp.Keys[0].KeyB64) + require.Equal(t, uint64(50*16), resp.Keys[0].Count, "count must be scaled by sample_rate") + require.Equal(t, base64.StdEncoding.EncodeToString([]byte("middling")), resp.Keys[1].KeyB64) + require.Equal(t, uint64(5*16), resp.Keys[1].Count) +} + +func TestHotKeysHandler_SubBucketFiltersByBytes(t *testing.T) { + t.Parallel() + src := newStubSource() + src.snapshots[1] = &keyviz.KeyvizHotKeysSnapshot{ + RouteID: 1, + SampledN: 10, + SampleRate: 1, + Capacity: 8, + SnapshotAt: time.Date(2026, 5, 28, 12, 0, 0, 0, time.UTC), + Entries: []keyviz.KeyvizHotKeyEntry{ + {Key: []byte("aaa"), Count: 1}, // outside (below lo) + {Key: []byte("mmm"), Count: 10}, // inside [m, t) + {Key: []byte("rrr"), Count: 3}, // inside + {Key: []byte("zzz"), Count: 7}, // outside (>= hi) + {Key: []byte("ttt"), Count: 99}, // outside (>= hi) + }, + } + // Sub-bucket 3 spans [m, t). + src.bounds[boundsKey{1, 3}] = boundsValue{lo: []byte("m"), hi: []byte("t"), ok: true} + h := newHandler(t, src) + rec := serve(t, h, "GET", "/admin/api/v1/keyviz/hotkeys?route_id=1&sub_bucket=3", nil) + require.Equal(t, http.StatusOK, rec.Code) + + var resp hotKeyResponse + require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp)) + require.NotNil(t, resp.SubBucket) + require.Equal(t, 3, *resp.SubBucket) + require.Len(t, resp.Keys, 2, "only mmm and rrr fall in [m, t)") + require.Equal(t, base64.StdEncoding.EncodeToString([]byte("mmm")), resp.Keys[0].KeyB64) + require.Equal(t, base64.StdEncoding.EncodeToString([]byte("rrr")), resp.Keys[1].KeyB64) +} + +func TestHotKeysHandler_SubBucketUnboundedTailFilter(t *testing.T) { + t.Parallel() + src := newStubSource() + src.snapshots[1] = &keyviz.KeyvizHotKeysSnapshot{ + RouteID: 1, + SampledN: 10, + SampleRate: 1, + Capacity: 8, + SnapshotAt: time.Date(2026, 5, 28, 12, 0, 0, 0, time.UTC), + Entries: []keyviz.KeyvizHotKeyEntry{ + {Key: []byte("aaa"), Count: 1}, // outside (below lo) + {Key: []byte("xxx"), Count: 10}, // inside (≥ lo, hi=nil) + {Key: []byte("zzz"), Count: 5}, + }, + } + src.bounds[boundsKey{1, 7}] = boundsValue{lo: []byte("m"), hi: nil, ok: true} + h := newHandler(t, src) + rec := serve(t, h, "GET", "/admin/api/v1/keyviz/hotkeys?route_id=1&sub_bucket=7", nil) + require.Equal(t, http.StatusOK, rec.Code) + var resp hotKeyResponse + require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp)) + require.Len(t, resp.Keys, 2) +} + +func TestHotKeysHandler_SubBucketOutOfRange(t *testing.T) { + t.Parallel() + src := newStubSource() + src.snapshots[1] = &keyviz.KeyvizHotKeysSnapshot{RouteID: 1, SnapshotAt: time.Now().UTC(), SampleRate: 1, Capacity: 8} + // bounds map has nothing for (1, 99) → ok=false + h := newHandler(t, src) + rec := serve(t, h, "GET", "/admin/api/v1/keyviz/hotkeys?route_id=1&sub_bucket=99", nil) + require.Equal(t, http.StatusBadRequest, rec.Code) + require.Contains(t, rec.Body.String(), "sub_bucket is out of range") +} + +func TestHotKeysHandler_DegradedFlagOR(t *testing.T) { + t.Parallel() + for _, tc := range []struct { + name string + dropped, skipped uint64 + wantDegraded bool + }{ + {"both zero", 0, 0, false}, + {"only drops", 7, 0, true}, + {"only long-key skips", 0, 11, true}, + {"both", 1, 1, true}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + src := newStubSource() + src.snapshots[1] = &keyviz.KeyvizHotKeysSnapshot{ + RouteID: 1, + SampledN: 10, + DroppedSamples: tc.dropped, + SkippedLongKeys: tc.skipped, + SampleRate: 1, + Capacity: 8, + SnapshotAt: time.Date(2026, 5, 28, 12, 0, 0, 0, time.UTC), + } + h := newHandler(t, src) + rec := serve(t, h, "GET", "/admin/api/v1/keyviz/hotkeys?route_id=1", nil) + require.Equal(t, http.StatusOK, rec.Code) + var resp hotKeyResponse + require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp)) + require.Equal(t, tc.wantDegraded, resp.Degraded) + require.Equal(t, tc.dropped, resp.DroppedSamples) + require.Equal(t, tc.skipped, resp.SkippedLongKeys) + }) + } +} + +func TestHotKeysHandler_OutOfSnapshotWindow(t *testing.T) { + t.Parallel() + src := newStubSource() + snapAt := time.Date(2026, 5, 28, 12, 0, 0, 0, time.UTC) + src.snapshots[1] = &keyviz.KeyvizHotKeysSnapshot{ + RouteID: 1, SampledN: 10, SampleRate: 1, Capacity: 8, SnapshotAt: snapAt, + } + h := NewKeyVizHotKeysHandler(src).WithSnapshotWindow(time.Minute) + + // Request a window 10 minutes in the past → outside [snapAt-1m, snapAt]. + past := snapAt.Add(-10 * time.Minute).UnixMilli() + url := "/admin/api/v1/keyviz/hotkeys?route_id=1&from_unix_ms=" + intStr(past) + + "&to_unix_ms=" + intStr(snapAt.Add(-9*time.Minute).UnixMilli()) + rec := serve(t, h, "GET", url, nil) + require.Equal(t, http.StatusBadRequest, rec.Code) + require.Contains(t, rec.Body.String(), "out_of_snapshot_window") + + // A window inside the snapshot window must pass. + url2 := "/admin/api/v1/keyviz/hotkeys?route_id=1&from_unix_ms=" + + intStr(snapAt.Add(-30*time.Second).UnixMilli()) + + "&to_unix_ms=" + intStr(snapAt.UnixMilli()) + rec2 := serve(t, h, "GET", url2, nil) + require.Equal(t, http.StatusOK, rec2.Code) +} + +// TestHotKeysHandler_InvertedTimeWindowReturns400 pins the +// parseTimeWindowParams from > to rejection (round-2 reviewer 🟡): an +// inverted [from, to] window must be rejected at the parameter layer +// with 400 invalid_query so the snapshot-window check never gets to +// evaluate an upside-down interval. +func TestHotKeysHandler_InvertedTimeWindowReturns400(t *testing.T) { + t.Parallel() + src := newStubSource() + src.snapshots[1] = &keyviz.KeyvizHotKeysSnapshot{ + RouteID: 1, SampledN: 10, SampleRate: 1, Capacity: 8, + SnapshotAt: time.Now().UTC(), + } + h := newHandler(t, src) + // from = 5000ms, to = 1000ms (to < from). + url := "/admin/api/v1/keyviz/hotkeys?route_id=1&from_unix_ms=5000&to_unix_ms=1000" + rec := serve(t, h, "GET", url, nil) + require.Equal(t, http.StatusBadRequest, rec.Code) + body := rec.Body.String() + require.Contains(t, body, "invalid_query") + require.Contains(t, body, "from_unix_ms") + require.Contains(t, body, "to_unix_ms") +} + +func TestHotKeysHandler_UnknownSeriesReturns400(t *testing.T) { + t.Parallel() + src := newStubSource() + src.snapshots[1] = &keyviz.KeyvizHotKeysSnapshot{RouteID: 1, SnapshotAt: time.Now().UTC(), SampleRate: 1, Capacity: 8} + h := newHandler(t, src) + rec := serve(t, h, "GET", "/admin/api/v1/keyviz/hotkeys?route_id=1&series=reads", nil) + require.Equal(t, http.StatusBadRequest, rec.Code) + require.Contains(t, rec.Body.String(), "writes") +} + +func TestHotKeysHandler_NonGetReturns405(t *testing.T) { + t.Parallel() + h := newHandler(t, newStubSource()) + rec := serve(t, h, "POST", "/admin/api/v1/keyviz/hotkeys?route_id=1", bytes.NewReader([]byte(`{}`))) + require.Equal(t, http.StatusMethodNotAllowed, rec.Code) +} + +// serve fires a single request through h and returns the recorded +// response. body may be nil for GETs. +func serve(t *testing.T, h http.Handler, method, urlStr string, body io.Reader) *httptest.ResponseRecorder { + t.Helper() + req := httptest.NewRequestWithContext(context.Background(), method, urlStr, body) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + return rec +} + +// intStr is strconv.FormatInt(x, 10) without importing strconv into every +// test that wants a single conversion. +func intStr(v int64) string { + const digits = "0123456789" + if v == 0 { + return "0" + } + neg := v < 0 + if neg { + v = -v + } + buf := [20]byte{} + i := len(buf) + for v > 0 { + i-- + buf[i] = digits[v%10] + v /= 10 + } + if neg { + i-- + buf[i] = '-' + } + return string(buf[i:]) +} diff --git a/internal/admin/server.go b/internal/admin/server.go index ee862e8c0..369becbf9 100644 --- a/internal/admin/server.go +++ b/internal/admin/server.go @@ -67,6 +67,15 @@ type ServerDeps struct { // preserves the legacy single-node behaviour. KeyVizFanout *KeyVizFanout + // KeyVizHotKeys backs the per-cell Top-K drill-down handler at + // /admin/api/v1/keyviz/hotkeys (design 2026_05_28_proposed_keyviz_hot_key_topk). + // Optional: nil makes the route 503 keyviz_disabled, the same shape + // as the matrix route's disabled response. A non-nil source whose + // HotKeysOptions() reports enabled=false serves 503 hotkeys_disabled + // instead, so the SPA can distinguish "keyviz off entirely" from + // "hot-keys drill-down off". + KeyVizHotKeys KeyVizHotKeysSource + // Queues is the SQS admin source — covers list, describe, and // delete via QueuesSource. Optional: a nil value disables // /admin/api/v1/sqs/queues{,/{name}} (the mux answers them @@ -136,8 +145,13 @@ func NewServer(deps ServerDeps) (*Server, error) { // a clearer "feature off" state than an unknown_endpoint 404. // Fan-out is opt-in: nil leaves the handler in single-node mode. keyviz := NewKeyVizHandler(deps.KeyViz).WithLogger(logger).WithFanout(deps.KeyVizFanout) + // Hot-keys drill-down: same "always registered, serves 503 when off" + // posture as the matrix handler. Auth + full-access role are layered + // at the route table (see buildAPIMux); the source's own enabled + // check produces hotkeys_disabled vs keyviz_disabled. + hotKeys := NewKeyVizHotKeysHandler(deps.KeyVizHotKeys).WithLogger(logger) sqs := buildSqsHandlerForDeps(deps, logger) - mux := buildAPIMux(auth, deps.Verifier, cluster, dynamo, s3, keyviz, sqs, logger) + mux := buildAPIMux(auth, deps.Verifier, cluster, dynamo, s3, keyviz, hotKeys, sqs, logger) router := NewRouterWithLeaderProbe(mux, deps.StaticFS, deps.LeaderProbe) return &Server{deps: deps, router: router, auth: auth, mux: mux}, nil } @@ -259,6 +273,7 @@ func (s *Server) APIHandler() http.Handler { // DELETE /admin/api/v1/s3/buckets/{name} (auth required, full role) // PUT /admin/api/v1/s3/buckets/{name}/acl (auth required, full role) // GET /admin/api/v1/keyviz/matrix (auth required) +// GET /admin/api/v1/keyviz/hotkeys (auth required, full role) // // Body limit applies uniformly. CSRF and Audit middleware apply to // write-capable protected endpoints; login and logout carry their own @@ -272,7 +287,7 @@ func (s *Server) APIHandler() http.Handler { // keyvizHandler is always non-nil even when the sampler is disabled — // it serves 503 keyviz_disabled itself so the SPA gets a clearer // signal than an unknown_endpoint 404 from the catch-all. -func buildAPIMux(auth *AuthService, verifier *Verifier, clusterHandler, dynamoHandler, s3Handler, keyvizHandler, sqsHandler http.Handler, logger *slog.Logger) http.Handler { +func buildAPIMux(auth *AuthService, verifier *Verifier, clusterHandler, dynamoHandler, s3Handler, keyvizHandler, keyvizHotKeysHandler, sqsHandler http.Handler, logger *slog.Logger) http.Handler { loginHandler := http.HandlerFunc(auth.HandleLogin) logoutHandler := http.HandlerFunc(auth.HandleLogout) @@ -323,6 +338,14 @@ func buildAPIMux(auth *AuthService, verifier *Verifier, clusterHandler, dynamoHa logoutChain := protectNoAudit(logoutHandler) clusterChain := protect(clusterHandler) keyvizChain := protect(keyvizHandler) + // Hot-keys drill-down requires the full-access role (design §7): + // the response carries actual user key bytes. SessionAuth (in the + // `protect` chain) authenticates the caller, RequireWriteRole + // enforces the role. CSRF wraps the whole stack because the + // response is information-disclosure-sensitive — a cross-site + // page should not be able to enumerate hot keys from an + // authenticated session. + keyvizHotKeysChain := protect(RequireWriteRole()(keyvizHotKeysHandler)) // Dynamo endpoints (reads and writes) share the protect chain // so a missing session or CSRF token 401s/403s the same way // regardless of method. The Audit middleware is a no-op for @@ -350,13 +373,14 @@ func buildAPIMux(auth *AuthService, verifier *Verifier, clusterHandler, dynamoHa } routes := apiRouteTable{ - login: loginChain, - logout: logoutChain, - cluster: clusterChain, - dynamo: dynamoChain, - s3: s3Chain, - keyviz: keyvizChain, - sqs: sqsChain, + login: loginChain, + logout: logoutChain, + cluster: clusterChain, + dynamo: dynamoChain, + s3: s3Chain, + keyviz: keyvizChain, + keyvizHotKeys: keyvizHotKeysChain, + sqs: sqsChain, } return http.HandlerFunc(routes.dispatch) } @@ -370,6 +394,7 @@ type apiRouteTable struct { login, logout, cluster http.Handler dynamo, s3, sqs http.Handler keyviz http.Handler + keyvizHotKeys http.Handler } // dispatch is the receiver method httpHandlerFunc adapts. Logic is @@ -407,8 +432,8 @@ func (t apiRouteTable) dispatch(w http.ResponseWriter, r *http.Request) { // equality and never against a nil receiver. func (t apiRouteTable) resourceHandlerFor(path string) http.Handler { switch { - case t.keyviz != nil && path == "/admin/api/v1/keyviz/matrix": - return t.keyviz + case isKeyVizPath(path): + return t.keyvizForPath(path) case t.dynamo != nil && isDynamoPath(path): return t.dynamo case t.s3 != nil && isS3Path(path): @@ -420,6 +445,39 @@ func (t apiRouteTable) resourceHandlerFor(path string) http.Handler { } } +// keyviz API paths. Constants here (not raw literals) because the +// strings appear in multiple dispatch helpers and in the fan-out URL +// builder (keyviz_fanout.go); goconst flags >2 raw occurrences. +const ( + pathKeyVizMatrix = "/admin/api/v1/keyviz/matrix" + pathKeyVizHotKeys = "/admin/api/v1/keyviz/hotkeys" +) + +// isKeyVizPath reports whether the URL belongs to the keyviz family +// (matrix / hotkeys). Used by resourceHandlerFor to dispatch with one +// case instead of one per endpoint, keeping the switch under cyclop. +func isKeyVizPath(path string) bool { + switch path { + case pathKeyVizMatrix, pathKeyVizHotKeys: + return true + } + return false +} + +// keyvizForPath returns the keyviz family handler for the given path, +// or nil when the path's specific handler isn't registered (e.g. +// hotkeys missing from a partial deploy). Centralised here so +// resourceHandlerFor stays small. +func (t apiRouteTable) keyvizForPath(path string) http.Handler { + switch path { + case pathKeyVizMatrix: + return t.keyviz + case pathKeyVizHotKeys: + return t.keyvizHotKeys + } + return nil +} + func isDynamoPath(p string) bool { return p == pathDynamoTables || strings.HasPrefix(p, pathPrefixDynamoTables) } diff --git a/keyviz/flusher.go b/keyviz/flusher.go index 0be239b35..2dfc460f7 100644 --- a/keyviz/flusher.go +++ b/keyviz/flusher.go @@ -35,3 +35,17 @@ func RunFlusher(ctx context.Context, s *MemSampler, step time.Duration) { } } } + +// RunHotKeysAggregator drives the per-route Top-K aggregator goroutine +// until ctx is cancelled, returning afterwards. Mirrors RunFlusher; the +// caller in main.go is expected to launch one of each in the same +// errgroup. No-op (blocks on ctx) when the sampler is nil or +// HotKeysEnabled is false — keeping the call site uniform so a future +// flip of the flag does not require restructuring startup wiring. +func RunHotKeysAggregator(ctx context.Context, s *MemSampler) { + if s == nil || s.hotKeys == nil { + <-ctx.Done() + return + } + s.hotKeys.run(ctx) +} diff --git a/keyviz/hot_keys.go b/keyviz/hot_keys.go new file mode 100644 index 000000000..b96fadbf0 --- /dev/null +++ b/keyviz/hot_keys.go @@ -0,0 +1,399 @@ +package keyviz + +import ( + "context" + "sync" + "sync/atomic" + "time" +) + +// Per-cell hot-key drill-down (design 2026_05_28_proposed_keyviz_hot_key_topk). +// Off by default; opt in via MemSamplerOptions.HotKeysEnabled. Disabled-case +// overhead is one early-return branch on Observe; enabled-case overhead is a +// length check, the splitmix64 `nextSampleRoll() % R == 0` sample gate (see +// nextSampleRoll for why we avoid math/rand/v2 entirely), a pool clone, and +// a non-blocking channel send. + +// hotKeyEvent ferries one sampled observation from the hot path to the +// aggregator. key is the pool's `*[]byte` (NOT the slice itself) so the +// aggregator can return it to the pool after copying the bytes into the +// Space-Saving sketch's own storage. The buffer is strictly transient +// between the hot path and the aggregator (design §4 pool-vs-snapshot +// ownership) — never aliased by an SS entry or a published snapshot. +type hotKeyEvent struct { + routeID uint64 + key *[]byte +} + +// KeyvizHotKeysSnapshot is the immutable per-route view a drill-down +// handler reads. The aggregator deep-copies every key bytes before +// publishing, so a reader can hold it across resets / evictions on the +// live sketch without races (design §4 published-snapshot ownership). +// +// Exported so the admin handler can read it; the type lives in the +// keyviz package because the aggregator owns and publishes it. +type KeyvizHotKeysSnapshot struct { + RouteID uint64 + SampledN uint64 // events the sketch actually saw this window + DroppedSamples uint64 // node-global: bounded-queue back-pressure drops + SkippedLongKeys uint64 // node-global: pre-sample length-cap rejects + SnapshotAt time.Time // when the aggregator deep-copied this snapshot + SampleRate int // R at the time of snapshot + Capacity int // m at the time of snapshot + Entries []KeyvizHotKeyEntry +} + +// KeyvizHotKeyEntry is one tracked (key, count) pair in a snapshot. +// Count is the sketch's raw count (sampled-stream observations); the +// admin layer scales it by SampleRate to estimate true frequency and +// computes error_bound = SampleRate * SampledN / Capacity. +type KeyvizHotKeyEntry struct { + Key []byte + Count uint64 +} + +// hotKeysAggregator runs as a single goroutine (single-writer to every +// per-route sketch). The hot path is the only writer to ch, the +// node-global drop/skip counters, and the per-route sampledN counter. +// Sketches and snapshots are owned exclusively by this goroutine. +type hotKeysAggregator struct { + s *MemSampler // back-reference, for table.Load() at publish time + ch chan hotKeyEvent + keyPool *sync.Pool // []byte buffers for hot-path key clones + capacity int // m + sampleRate int // R (used by Observe) + maxKeyLen int // hotKeysMaxKeyLen + step time.Duration + clock func() time.Time + // Node-global hot-path counters (atomic.Uint64 incremented by Observe, + // read once per publish by the aggregator). + dropped atomic.Uint64 + skipped atomic.Uint64 + // Per-route sampled-N counters (incremented by the aggregator only, + // reset on publish). Map mutation is guarded by the fact that the + // aggregator is the single writer. + perRouteN map[uint64]*uint64 + // Per-route Space-Saving sketches (single-writer-by-aggregator). + sketches map[uint64]*spaceSaving + // rngState backs the sample gate. We avoid math/rand/v2 entirely + // (gosec G404 flags the package as a weak crypto generator — but + // it's also the only built-in PRNG, and we'd otherwise need a + // //nolint). atomic.Uint64.Add gives a lock-free monotonic + // sequence; splitmix64 mixes it into a uniformly-distributed + // uint64, breaking phase-lock against periodic / ordered workloads + // (Codex P2 round-3 L135) without owning any shared Source state + // (Codex P1 round-4 L137). The state is process-local to this + // aggregator; concurrent Observe callers race on Add only. + rngState atomic.Uint64 +} + +// splitmix64 standard parameters — public-domain mixer used unchanged +// across implementations. The golden-ratio increment is coprime with +// every power of two so the atomic counter walks the whole uint64 space. +const ( + splitmix64GoldenInc uint64 = 0x9E3779B97F4A7C15 + splitmix64Mul1 uint64 = 0xBF58476D1CE4E5B9 + splitmix64Mul2 uint64 = 0x94D049BB133111EB + splitmix64Shift1 uint64 = 30 + splitmix64Shift2 uint64 = 27 + splitmix64Shift3 uint64 = 31 +) + +// nextSampleRoll returns a pseudo-random uint64 drawn from the +// aggregator's lock-free splitmix64 stream. The output is uniformly +// distributed over [0, 2^64), so `nextSampleRoll() % R == 0` is the +// design §4 probabilistic sample gate. +func (a *hotKeysAggregator) nextSampleRoll() uint64 { + z := a.rngState.Add(splitmix64GoldenInc) + z = (z ^ (z >> splitmix64Shift1)) * splitmix64Mul1 + z = (z ^ (z >> splitmix64Shift2)) * splitmix64Mul2 + return z ^ (z >> splitmix64Shift3) +} + +// applyHotKeysDefaults clamps every HotKeys* field on opts to its +// Default / Max range. Extracted out of NewMemSampler so the constructor +// stays below the cyclop ceiling as the hot-keys knobs grew the option +// surface (5 new fields). Caller takes a pointer because the defaults +// must mutate the struct seen by the rest of NewMemSampler. +func applyHotKeysDefaults(opts *MemSamplerOptions) { + if opts.HotKeysPerRoute <= 0 { + opts.HotKeysPerRoute = DefaultHotKeysPerRoute + } + if opts.HotKeysPerRoute > MaxHotKeysPerRoute { + opts.HotKeysPerRoute = MaxHotKeysPerRoute + } + if opts.HotKeysSampleRate <= 0 { + opts.HotKeysSampleRate = DefaultHotKeysSampleRate + } + if opts.HotKeysSampleRate > MaxHotKeysSampleRate { + opts.HotKeysSampleRate = MaxHotKeysSampleRate + } + if opts.HotKeysQueueSize <= 0 { + opts.HotKeysQueueSize = DefaultHotKeysQueueSize + } + if opts.HotKeysQueueSize > MaxHotKeysQueueSize { + opts.HotKeysQueueSize = MaxHotKeysQueueSize + } + if opts.HotKeysMaxKeyLen <= 0 { + opts.HotKeysMaxKeyLen = DefaultHotKeysMaxKeyLen + } + if opts.HotKeysMaxKeyLen > MaxHotKeysMaxKeyLen { + opts.HotKeysMaxKeyLen = MaxHotKeysMaxKeyLen + } +} + +func newHotKeysAggregator(s *MemSampler, opts MemSamplerOptions) *hotKeysAggregator { + now := opts.Now + if now == nil { + now = time.Now + } + maxKeyLen := opts.HotKeysMaxKeyLen + a := &hotKeysAggregator{ + s: s, + ch: make(chan hotKeyEvent, opts.HotKeysQueueSize), + capacity: opts.HotKeysPerRoute, + sampleRate: opts.HotKeysSampleRate, + maxKeyLen: maxKeyLen, + step: opts.Step, + clock: now, + perRouteN: map[uint64]*uint64{}, + sketches: map[uint64]*spaceSaving{}, + } + a.keyPool = &sync.Pool{ + New: func() any { + b := make([]byte, 0, maxKeyLen) + return &b + }, + } + return a +} + +// borrowKeyBuffer returns a pooled buffer cleared to zero length but +// with maxKeyLen capacity, so Observe can append without further +// allocation for any key within the configured limit. +func (a *hotKeysAggregator) borrowKeyBuffer() *[]byte { + // keyPool.New is set in newHotKeysAggregator to always return a + // *[]byte, so the assertion can never fail in practice. The + // comma-ok form keeps forcetypeassert happy and ensures a panic + // would carry a more useful message than the runtime's default if + // a future refactor ever changes the pool's New func. + bp, ok := a.keyPool.Get().(*[]byte) + if !ok || bp == nil { + b := make([]byte, 0, a.maxKeyLen) + bp = &b + } + *bp = (*bp)[:0] + return bp +} + +// releaseKeyBuffer returns the buffer to the pool. The aggregator +// calls this immediately after it has copied the key bytes into the +// SS entry's own storage (sketches.observe copies internally). +func (a *hotKeysAggregator) releaseKeyBuffer(bp *[]byte) { + *bp = (*bp)[:0] + a.keyPool.Put(bp) +} + +// observe is invoked from the sampler's hot path when hot-keys is +// enabled. It performs the design §4 ordering: +// 1. length check FIRST (skipped_long_keys counted deterministically, +// not subject to sampling — Codex P2 round-7 L116) +// 2. probabilistic sample gate (splitmix64 nextSampleRoll()%R == 0, +// lock-free — Codex P1 round-4 L137) +// 3. key clone via sync.Pool +// 4. non-blocking channel send (drop-on-full + drop counter — Codex +// P2 round-3 L138) +// +// Returns true iff the event was enqueued for the aggregator. The +// boolean is for tests; the production caller in Observe ignores it. +func (a *hotKeysAggregator) observe(routeID uint64, key []byte) bool { + if len(key) > a.maxKeyLen { + a.skipped.Add(1) + return false + } + // Probabilistic sample gate driven by the aggregator's lock-free + // splitmix64 stream (no shared Source ownership, no math/rand/v2 + // import — see nextSampleRoll). The u64NonNeg helper (sampler.go) + // gives gosec the explicit non-negative guard it needs for the + // int→uint64 cast; sampleRate is sampler-clamped to ≥1 anyway, so + // the path that would return 0 here is unreachable. + rate := u64NonNeg(a.sampleRate) + if rate == 0 { + rate = 1 + } + if a.nextSampleRoll()%rate != 0 { + return false + } + bp := a.borrowKeyBuffer() + *bp = append(*bp, key...) + select { + case a.ch <- hotKeyEvent{routeID: routeID, key: bp}: + // bp will be released by the aggregator after it copies the key + // into the SS entry's own storage. + return true + default: + a.dropped.Add(1) + a.releaseKeyBuffer(bp) + return false + } +} + +// run is the aggregator goroutine's main loop. Single select with two +// arms (drain events / tick), matching the design §4 pseudocode: +// +// for { +// select { +// case e := <-ch: updateSketch(e) +// case <-tick: +// for len(ch) > 0 { updateSketch(<-ch) } // best-effort pre-drain +// publishAndReset() +// } +// } +// +// Single-writer to every per-route SS sketch and to perRouteN, so no +// lock is needed on the update path. ctx cancellation drains a final +// publish so the last window's data isn't lost on shutdown. +func (a *hotKeysAggregator) run(ctx context.Context) { + tick := time.NewTicker(a.step) + defer tick.Stop() + for { + select { + case <-ctx.Done(): + // Best-effort final drain + publish so an operator hitting + // the drill-down right after a graceful shutdown gets the + // last window's data rather than an empty snapshot. + for len(a.ch) > 0 { + a.consume(<-a.ch) + } + a.publishAndReset() + return + case e := <-a.ch: + a.consume(e) + case <-tick.C: + // Pre-drain: process every event currently queued before we + // publish + reset, so the just-finished window's tail-end + // observations land in this snapshot rather than the next + // (Codex P2 round-7 L193 — boundary smear is now bounded by + // arrivals during publishAndReset, not by the channel depth + // at the tick). + for len(a.ch) > 0 { + a.consume(<-a.ch) + } + a.publishAndReset() + } + } +} + +// consume handles one event: drains the sample-N counter for the route, +// updates the Space-Saving sketch, and releases the pool buffer. +func (a *hotKeysAggregator) consume(e hotKeyEvent) { + n, ok := a.perRouteN[e.routeID] + if !ok { + n = new(uint64) + a.perRouteN[e.routeID] = n + } + *n++ + + sk, ok := a.sketches[e.routeID] + if !ok { + sk = newSpaceSaving(a.capacity) + a.sketches[e.routeID] = sk + } + // SS copies the bytes into its own entry storage; the pool buffer + // is then safe to return immediately. + sk.observe(*e.key) + a.releaseKeyBuffer(e.key) +} + +// publishAndReset is called from the tick branch after the pre-drain. +// It atomically (single-writer; no concurrent reader of the live +// sketch) deep-copies every per-route sketch into a fresh +// KeyvizHotKeysSnapshot, publishes it on the matching routeSlot via +// atomic.Pointer.Store, then resets the live sketches and the +// per-route / node-global counters so the next keyvizStep window +// starts empty (design §4 sketch reset; Codex P2 round-6 L186). +func (a *hotKeysAggregator) publishAndReset() { + // When no per-route sketch exists this window (e.g. every observe + // was filtered by the length cap, or no traffic at all on any + // route), we have nowhere to attach the node-global drop / skip + // counters — and Swap-resetting them here would silently discard + // the signal. Carry them forward into the next window instead; + // whichever publish tick first has a sketch will surface them. + if len(a.sketches) == 0 { + return + } + now := a.clock() + // Swap-and-capture (not Load + later Store(0)) — the hot path can + // still call a.dropped.Add(1) / a.skipped.Add(1) between a Load and + // a Store, and those increments would belong to neither the snapshot + // being published (already latched) nor the next window (the Store + // would clobber them), leaving `degraded=false` during a publish- + // straddling burst (claude bot 🔴 PR #854). + dropped := a.dropped.Swap(0) + skipped := a.skipped.Swap(0) + tbl := a.s.table.Load() + for rid, sk := range a.sketches { + slot := lookupSlotForRoute(tbl, rid) + if slot == nil { + // Route was removed mid-window: drop the sketch AND its + // per-route counter so we don't accumulate dead entries + // under route churn (gemini HIGH on PR #854 — without this, + // every removed route's m × key bytes plus a *uint64 stays + // in the aggregator forever). + delete(a.sketches, rid) + delete(a.perRouteN, rid) + continue + } + n := uint64(0) + if p, ok := a.perRouteN[rid]; ok { + n = *p + } + snap := &KeyvizHotKeysSnapshot{ + RouteID: rid, + SampledN: n, + DroppedSamples: dropped, + SkippedLongKeys: skipped, + SnapshotAt: now, + SampleRate: a.sampleRate, + Capacity: a.capacity, + Entries: toSnapshotEntries(sk.snapshot()), + } + slot.hotKeysSnap.Store(snap) + sk.reset() + } + // Per-route counters reset; the node-global ones were already + // captured-and-reset by the Swap(0) calls at the top of this + // function — issuing another Store(0) here would re-clobber any + // hot-path Add that arrived during the publish loop. + for k := range a.perRouteN { + *a.perRouteN[k] = 0 + } +} + +func toSnapshotEntries(es []ssEntrySnap) []KeyvizHotKeyEntry { + if len(es) == 0 { + return nil + } + out := make([]KeyvizHotKeyEntry, len(es)) + for i, e := range es { + // ssEntrySnap and KeyvizHotKeyEntry have identical field layout; + // a struct conversion is what staticcheck S1016 asks for and + // is cheaper than the field-by-field literal copy. + out[i] = KeyvizHotKeyEntry(e) + } + return out +} + +// lookupSlotForRoute resolves a real (non-aggregate) route's slot from +// the route table. Aggregate / virtual buckets are NOT eligible for +// hot-keys tracking (design §2.2): they coarsen multiple routes and +// "top-K within an aggregate" is not meaningful. +func lookupSlotForRoute(tbl *routeTable, routeID uint64) *routeSlot { + if tbl == nil { + return nil + } + if s, ok := tbl.slots[routeID]; ok && !s.Aggregate { + return s + } + return nil +} diff --git a/keyviz/hot_keys_test.go b/keyviz/hot_keys_test.go new file mode 100644 index 000000000..78cb87113 --- /dev/null +++ b/keyviz/hot_keys_test.go @@ -0,0 +1,360 @@ +package keyviz + +import ( + "bytes" + "context" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +// hotKeysTestCapacity is the per-route Space-Saving capacity used by +// every direct hotKeysTestSampler test. Inlined here (rather than a +// fn parameter) because every caller passes the same value — varying +// only along queue/key-length axes — and unparam catches the dead arg. +const hotKeysTestCapacity = 8 + +// hotKeysTestSampler constructs a MemSampler with hot-keys enabled and +// R=1 (every observe enqueues — deterministic for tests). The Step is +// set high so the aggregator's tick never fires during the test; +// publishAndReset is called explicitly when needed. +func hotKeysTestSampler(t *testing.T, queueSize, maxKeyLen int) *MemSampler { + t.Helper() + return NewMemSampler(MemSamplerOptions{ + Step: time.Hour, // never ticks during the test + HistoryColumns: 4, + HotKeysEnabled: true, + HotKeysPerRoute: hotKeysTestCapacity, + HotKeysSampleRate: 1, // every Observe is sampled + HotKeysQueueSize: queueSize, + HotKeysMaxKeyLen: maxKeyLen, + KeyBucketsPerRoute: 1, + }) +} + +// drainAndPublish synchronously drains the channel into the per-route +// SS sketches and triggers one publish+reset cycle. Test-only shortcut +// that exercises the same code path as the aggregator's tick arm. +func drainAndPublish(a *hotKeysAggregator) { + for len(a.ch) > 0 { + a.consume(<-a.ch) + } + a.publishAndReset() +} + +func TestHotKeysObservePublishesPerRouteSnapshot(t *testing.T) { + t.Parallel() + s := hotKeysTestSampler(t, 32, 64) + require.True(t, s.RegisterRoute(1, []byte("a"), []byte("z"), 0)) + + for i := 0; i < 5; i++ { + s.Observe(1, []byte("hot"), OpWrite, 0) + } + s.Observe(1, []byte("cold"), OpWrite, 0) + + require.Nil(t, s.HotKeysSnapshot(1), "no snapshot before publish") + drainAndPublish(s.hotKeys) + + snap := s.HotKeysSnapshot(1) + require.NotNil(t, snap) + require.Equal(t, uint64(1), snap.RouteID) + require.Equal(t, uint64(6), snap.SampledN, "every Observe enqueued at R=1") + require.Equal(t, 1, snap.SampleRate) + require.Equal(t, 8, snap.Capacity) + require.Zero(t, snap.DroppedSamples) + require.Zero(t, snap.SkippedLongKeys) + require.False(t, snapDegraded(snap)) + + got := map[string]uint64{} + for _, e := range snap.Entries { + got[string(e.Key)] = e.Count + } + require.Equal(t, uint64(5), got["hot"]) + require.Equal(t, uint64(1), got["cold"]) +} + +// snapDegraded mirrors the admin response's `degraded` rule (§5): +// drops > 0 OR skipped_long_keys > 0. +func snapDegraded(s *KeyvizHotKeysSnapshot) bool { + return s.DroppedSamples > 0 || s.SkippedLongKeys > 0 +} + +// TestHotKeysLongKeySkippedBeforeSampling pins Codex round-7 L116: a +// key longer than HotKeysMaxKeyLen increments `skipped_long_keys` on +// every Observe — UNCONDITIONALLY, before the sample gate. So a long +// hot key cannot escape the degraded signal just because the sampler +// happened to roll against it on every call. +func TestHotKeysLongKeySkippedBeforeSampling(t *testing.T) { + t.Parallel() + // Use R = 1000 so the sample gate practically never fires. The + // length check fires before the sample gate, so every long-key + // Observe must still bump the counter. + s := NewMemSampler(MemSamplerOptions{ + Step: time.Hour, + HistoryColumns: 4, + HotKeysEnabled: true, + HotKeysPerRoute: 8, + HotKeysSampleRate: 1000, + HotKeysQueueSize: 32, + HotKeysMaxKeyLen: 4, // very short cap + }) + require.True(t, s.RegisterRoute(1, []byte("a"), []byte("z"), 0)) + tooLong := bytes.Repeat([]byte("x"), 100) + for i := 0; i < 50; i++ { + s.Observe(1, tooLong, OpWrite, 0) + } + // Long-key check runs BEFORE the sample gate, so all 50 observes + // bump the node-global counter — independent of R. + require.Equal(t, uint64(50), s.hotKeys.skipped.Load(), + "long-key check must run BEFORE the sample gate (Codex P2 round-7 L116)") + + // No events enqueued (every observe was length-skipped), so + // publishAndReset has no sketch to attach the node-global signal + // to. publishAndReset must NOT discard the signal in that case — + // it should carry the skipped count forward to the next window. + drainAndPublish(s.hotKeys) + require.Nil(t, s.HotKeysSnapshot(1), "no enqueues -> no per-route snapshot") + require.Equal(t, uint64(50), s.hotKeys.skipped.Load(), + "node-global skipped must carry forward when no sketch exists to attach it to") + + // Second sampler at R=1 (every observe samples deterministically) + // to pin the surface-and-Swap behavior without flaking on the + // sample gate: seed the live `skipped` counter directly, observe + // one short key so a sketch exists, and verify the next publish + // captures the skipped count INTO that snapshot via Swap and + // resets the live counter atomically. + s2 := hotKeysTestSampler(t, 32, 4) + require.True(t, s2.RegisterRoute(1, []byte("a"), []byte("z"), 0)) + s2.hotKeys.skipped.Store(50) + s2.Observe(1, []byte("ok"), OpWrite, 0) + drainAndPublish(s2.hotKeys) + snap := s2.HotKeysSnapshot(1) + require.NotNil(t, snap) + require.Equal(t, uint64(50), snap.SkippedLongKeys, "skipped surfaced on the first sketch-bearing snapshot") + require.True(t, snapDegraded(snap), "non-zero skipped -> degraded") + require.Equal(t, uint64(0), s2.hotKeys.skipped.Load(), "Swap reset the live counter when published") +} + +// TestHotKeysChannelFullCountsDrops pins Codex round-3 L138: a sampled +// event that hits a full bounded channel does NOT silently disappear — +// the `dropped_samples` counter increments and degraded fires. +func TestHotKeysChannelFullCountsDrops(t *testing.T) { + t.Parallel() + // Queue size 4; no aggregator goroutine running, so nothing drains. + // At R=1, the 5th observe must drop. + s := hotKeysTestSampler(t, 4, 64) + require.True(t, s.RegisterRoute(1, []byte("a"), []byte("z"), 0)) + for i := 0; i < 10; i++ { + s.Observe(1, []byte("k"), OpWrite, 0) + } + require.Equal(t, 4, len(s.hotKeys.ch), "queue should be full (cap=4)") + require.Equal(t, uint64(6), s.hotKeys.dropped.Load(), "10 observes - 4 enqueued = 6 drops") +} + +// TestHotKeysAggregateRoutesNotTracked pins design §2.2: aggregates +// (RouteID's synthetic / Aggregate=true slots) are not eligible for +// hot-keys tracking. HotKeysSnapshot must return nil for them. +func TestHotKeysAggregateRoutesNotTracked(t *testing.T) { + t.Parallel() + s := NewMemSampler(MemSamplerOptions{ + Step: time.Hour, + HistoryColumns: 4, + MaxTrackedRoutes: 1, // force coarsening for the 2nd RegisterRoute + HotKeysEnabled: true, + HotKeysPerRoute: 8, + HotKeysSampleRate: 1, + HotKeysQueueSize: 32, + HotKeysMaxKeyLen: 64, + }) + require.True(t, s.RegisterRoute(1, []byte("a"), []byte("m"), 0), "1st route gets its own slot") + require.False(t, s.RegisterRoute(2, []byte("n"), []byte("z"), 0), "2nd route coarsens into a virtual bucket") + // Observe against route 2: hits the aggregate bucket, which is + // filtered out by the Observe-time !slot.Aggregate check. + for i := 0; i < 5; i++ { + s.Observe(2, []byte("would-be-hot"), OpWrite, 0) + } + drainAndPublish(s.hotKeys) + require.Nil(t, s.HotKeysSnapshot(2), "aggregate routes are not tracked") +} + +// TestHotKeysSnapshotResetClearsLiveSketch pins design §4 / Codex P2 +// round-6 L186: each publish RESETS the live sketch, so a key hot in +// one window cannot dominate the next window's drill-down. +func TestHotKeysSnapshotResetClearsLiveSketch(t *testing.T) { + t.Parallel() + s := hotKeysTestSampler(t, 32, 64) + require.True(t, s.RegisterRoute(1, []byte("a"), []byte("z"), 0)) + // Window 1: "old" key hot. + for i := 0; i < 10; i++ { + s.Observe(1, []byte("old"), OpWrite, 0) + } + drainAndPublish(s.hotKeys) + snap := s.HotKeysSnapshot(1) + require.NotNil(t, snap) + require.Equal(t, uint64(10), snap.SampledN) + + // Window 2: only one observe of a new key. The reset rule means + // the new snapshot reflects ONLY window-2 traffic — "old" must + // be absent, sampledN == 1. + s.Observe(1, []byte("new"), OpWrite, 0) + drainAndPublish(s.hotKeys) + snap2 := s.HotKeysSnapshot(1) + require.NotNil(t, snap2) + require.Equal(t, uint64(1), snap2.SampledN, "reset must clear sampledN") + require.Len(t, snap2.Entries, 1) + require.Equal(t, []byte("new"), snap2.Entries[0].Key) +} + +// TestHotKeysDisabledIsZeroCost pins the disabled-case contract: with +// HotKeysEnabled=false the sampler does not even allocate an +// aggregator, so the hot path's nil check short-circuits and no key +// bytes are retained anywhere. +func TestHotKeysDisabledIsZeroCost(t *testing.T) { + t.Parallel() + s := NewMemSampler(MemSamplerOptions{Step: time.Hour, HistoryColumns: 4, HotKeysEnabled: false}) + require.Nil(t, s.hotKeys, "disabled => no aggregator constructed") + require.True(t, s.RegisterRoute(1, []byte("a"), []byte("z"), 0)) + s.Observe(1, []byte("any"), OpWrite, 0) // must not panic; no-op for hot-keys + require.Nil(t, s.HotKeysSnapshot(1)) +} + +// TestHotKeysDropCounterAttachesToSnapshot pins the claude-bot 🔴 +// fix: dropped counts must reach the published snapshot via Swap (not +// Load + later Store, which would lose drops that arrive between). +func TestHotKeysDropCounterAttachesToSnapshot(t *testing.T) { + t.Parallel() + // Queue size 4; no aggregator goroutine — observes accumulate so + // the 5th onwards drops, and the dropped counter rises to 6. + s := hotKeysTestSampler(t, 4, 64) + require.True(t, s.RegisterRoute(1, []byte("a"), []byte("z"), 0)) + for i := 0; i < 10; i++ { + s.Observe(1, []byte("k"), OpWrite, 0) + } + require.Equal(t, uint64(6), s.hotKeys.dropped.Load(), "10 sends - 4 enqueued = 6 drops") + drainAndPublish(s.hotKeys) + snap := s.HotKeysSnapshot(1) + require.NotNil(t, snap) + require.Equal(t, uint64(6), snap.DroppedSamples, "Swap captures drops into the published snapshot") + require.Equal(t, uint64(0), s.hotKeys.dropped.Load(), "Swap reset the live counter") + require.True(t, snapDegraded(snap)) +} + +// TestHotKeysAggregatorReleasesRemovedRouteState pins the gemini-HIGH +// fix: when a route is removed mid-window, the aggregator must drop +// its sketch + per-route counter, otherwise removed routes leak m × +// key bytes (and a *uint64) indefinitely under route churn. +func TestHotKeysAggregatorReleasesRemovedRouteState(t *testing.T) { + t.Parallel() + s := hotKeysTestSampler(t, 32, 64) + require.True(t, s.RegisterRoute(1, []byte("a"), []byte("z"), 0)) + for i := 0; i < 4; i++ { + s.Observe(1, []byte("hot"), OpWrite, 0) + } + drainAndPublish(s.hotKeys) + require.Contains(t, s.hotKeys.sketches, uint64(1)) + require.Contains(t, s.hotKeys.perRouteN, uint64(1)) + + // Route removed → next publish must drop its sketch + counter. + s.RemoveRoute(1) + drainAndPublish(s.hotKeys) + require.NotContains(t, s.hotKeys.sketches, uint64(1), "removed route's sketch must be released") + require.NotContains(t, s.hotKeys.perRouteN, uint64(1), "removed route's counter must be released") +} + +// TestHotKeysAggregatorGracefulShutdownPublishes pins the ctx.Done arm +// of run(): a final drain + publish must fire so a last-second observe +// lands in a queryable snapshot rather than being silently discarded at +// shutdown. We use cancellation (not the tick) to drive the publish so +// the assertion is deterministic regardless of goroutine scheduling — +// the periodic-tick path is exercised by TestHotKeysAggregatorRaceFree. +func TestHotKeysAggregatorGracefulShutdownPublishes(t *testing.T) { + t.Parallel() + s := NewMemSampler(MemSamplerOptions{ + // Step is set very long so the periodic tick cannot fire + // before we cancel — the only publish must come from the + // ctx.Done arm. + Step: time.Hour, + HistoryColumns: 4, + HotKeysEnabled: true, + HotKeysPerRoute: 8, + HotKeysSampleRate: 1, + HotKeysQueueSize: 32, + HotKeysMaxKeyLen: 64, + }) + require.True(t, s.RegisterRoute(1, []byte("a"), []byte("z"), 0)) + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + defer close(done) + RunHotKeysAggregator(ctx, s) + }() + + for i := 0; i < 7; i++ { + s.Observe(1, []byte("hot"), OpWrite, 0) + } + // Cancel triggers the ctx.Done arm, which drains the channel + // and publishes one final snapshot before returning. + cancel() + <-done + + snap := s.HotKeysSnapshot(1) + require.NotNil(t, snap, "graceful shutdown must publish a final snapshot") + require.GreaterOrEqual(t, snap.SampledN, uint64(7), "all observed events present in the final snapshot") + require.Len(t, snap.Entries, 1) + require.Equal(t, []byte("hot"), snap.Entries[0].Key) +} + +// TestHotKeysAggregatorRaceFree exercises the hot path under +// concurrent observers + a live aggregator goroutine; -race must pass. +func TestHotKeysAggregatorRaceFree(t *testing.T) { + t.Parallel() + s := NewMemSampler(MemSamplerOptions{ + Step: 5 * time.Millisecond, + HistoryColumns: 4, + HotKeysEnabled: true, + HotKeysPerRoute: 16, + HotKeysSampleRate: 4, + HotKeysQueueSize: 256, + HotKeysMaxKeyLen: 64, + }) + require.True(t, s.RegisterRoute(1, []byte("a"), []byte("z"), 0)) + require.True(t, s.RegisterRoute(2, []byte("m"), []byte("z"), 0)) + + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + var wgRun sync.WaitGroup + wgRun.Add(1) + go func() { + defer wgRun.Done() + RunHotKeysAggregator(ctx, s) + }() + + var wg sync.WaitGroup + for w := 0; w < 8; w++ { + wg.Add(1) + go func(seed int) { + defer wg.Done() + key := []byte{byte('k' + seed%4)} + for i := 0; i < 500; i++ { + var rid uint64 = 1 + if i%2 == 1 { + rid = 2 + } + s.Observe(rid, key, OpWrite, 0) + } + }(w) + } + wg.Wait() + // Let the aggregator drain + publish a few ticks. + time.Sleep(50 * time.Millisecond) + // Snapshot reads must be lock-free and consistent. + s1 := s.HotKeysSnapshot(1) + s2 := s.HotKeysSnapshot(2) + require.NotNil(t, s1) + require.NotNil(t, s2) + cancel() + wgRun.Wait() +} diff --git a/keyviz/sampler.go b/keyviz/sampler.go index 1927e6e52..5a9c124ae 100644 --- a/keyviz/sampler.go +++ b/keyviz/sampler.go @@ -91,6 +91,21 @@ const ( DefaultKeyBucketsPerRoute = 1 ) +// Defaults / caps for the per-cell hot-key drill-down knobs (see +// docs/design/2026_05_28_proposed_keyviz_hot_key_topk.md §8). All +// off-by-default; HotKeysEnabled=false is the binary's existing +// behaviour — no extra hot-path cost, no real key bytes retained. +const ( + DefaultHotKeysPerRoute = 64 + MaxHotKeysPerRoute = 256 + DefaultHotKeysSampleRate = 16 + MaxHotKeysSampleRate = 1024 + DefaultHotKeysQueueSize = 8192 + MaxHotKeysQueueSize = 65536 + DefaultHotKeysMaxKeyLen = 1024 + MaxHotKeysMaxKeyLen = 4096 +) + // MaxKeyBucketsPerRoute caps KeyBucketsPerRoute. The cap is an // operationally-safe bound, not merely a finite one: per-route memory // is K × 4 × 8 bytes, so at the default MaxTrackedRoutes = 10_000 the @@ -144,6 +159,34 @@ type MemSamplerOptions struct { // [1, MaxKeyBucketsPerRoute] at construction. Virtual aggregate // buckets are never sub-divided regardless of this value. KeyBucketsPerRoute int + // HotKeysEnabled opts in to per-route Top-K hot-key tracking that + // backs the heatmap drill-down (Phase 2-A++; see + // docs/design/2026_05_28_proposed_keyviz_hot_key_topk.md). When + // false the hot path adds one early-return branch and nothing else + // — disabled-case behaviour is byte-identical to today. When true + // the sampler retains REAL key bytes in memory and exposes them on + // the admin drill-down API (gated behind admin auth + audit). + HotKeysEnabled bool + // HotKeysPerRoute is the Space-Saving sketch capacity m per route + // (default 64, cap 256 — design §8). Larger m tightens the noise + // floor (error_bound ≈ N_total / m) at the cost of memory. + HotKeysPerRoute int + // HotKeysSampleRate is R: the hot path enqueues 1 in R observes. + // Default 16, cap 1024. Higher R reduces hot-path cost but raises + // the probability that a heavy hitter is missed (Chernoff variance + // over the sampled stream). + HotKeysSampleRate int + // HotKeysQueueSize bounds the channel between the hot path and + // the aggregator goroutine. Defaults 8192, cap 65536. Drops past + // the queue are counted, not silent (`dropped_samples` → + // `degraded`). + HotKeysQueueSize int + // HotKeysMaxKeyLen caps the key length sampled into the hot-keys + // sketch. Default 1024 B, cap 4096 B. Longer keys bump the + // node-global `skipped_long_keys` counter and contribute to the + // snapshot's `degraded` flag — they are never truncated (truncation + // would alias different keys). + HotKeysMaxKeyLen int // Now overrides time.Now for tests; nil falls back to time.Now. Now func() time.Time } @@ -192,6 +235,12 @@ type MemSampler struct { // an individual row — even under register/remove churn. virtualIDCounter atomic.Uint64 + // hotKeys is the per-route Top-K aggregator (design 2026_05_28 + // _proposed_keyviz_hot_key_topk). nil when HotKeysEnabled is false + // — the Observe hot path then skips the feature with one branch. + // Read from the hot path; assigned exactly once at construction. + hotKeys *hotKeysAggregator + // 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 @@ -328,6 +377,15 @@ type routeSlot struct { subLo []byte subHi []byte subBuckets []subCounter + + // hotKeysSnap is the most-recently-published per-route Top-K + // snapshot (design 2026_05_28_proposed_keyviz_hot_key_topk §4). + // nil until the hot-keys aggregator's first publish for this route. + // Drill-down handlers load it lock-free; the aggregator is the + // single writer (Store) and deep-copies snapshot contents so a + // subsequent reset / eviction on the live sketch cannot mutate a + // snapshot a reader holds. + hotKeysSnap atomic.Pointer[KeyvizHotKeysSnapshot] } // subCounter holds one sub-range bucket's four counter families. A @@ -438,6 +496,7 @@ func NewMemSampler(opts MemSamplerOptions) *MemSampler { if opts.KeyBucketsPerRoute > MaxKeyBucketsPerRoute { opts.KeyBucketsPerRoute = MaxKeyBucketsPerRoute } + applyHotKeysDefaults(&opts) now := opts.Now if now == nil { now = time.Now @@ -449,9 +508,73 @@ func NewMemSampler(opts MemSamplerOptions) *MemSampler { groupTerms: map[uint64]uint64{}, } s.table.Store(newEmptyRouteTable()) + if opts.HotKeysEnabled { + s.hotKeys = newHotKeysAggregator(s, opts) + } return s } +// HotKeysSnapshot returns the most-recently-published per-route Top-K +// snapshot for the given individual route, or nil if hot-keys tracking +// is disabled, the route is unknown, the route is an aggregate, or no +// publish has fired yet. Lock-free: an atomic.Pointer.Load (Codex-P1 +// fix from the design's round-4 review — no shared mutable PRNG state, +// no per-route lock). +func (s *MemSampler) HotKeysSnapshot(routeID uint64) *KeyvizHotKeysSnapshot { + if s == nil { + return nil + } + slot := lookupSlotForRoute(s.table.Load(), routeID) + if slot == nil { + return nil + } + return slot.hotKeysSnap.Load() +} + +// HotKeysOptions returns the effective hot-keys configuration after the +// NewMemSampler clamps. Used by the admin handler so the response can +// echo the same `sample_rate` / `m` values the sketch was built with, +// regardless of what an operator passed on the CLI. nil-receiver-safe. +func (s *MemSampler) HotKeysOptions() (enabled bool, capacity, sampleRate, maxKeyLen int) { + if s == nil { + return false, 0, 0, 0 + } + return s.opts.HotKeysEnabled, s.opts.HotKeysPerRoute, s.opts.HotKeysSampleRate, s.opts.HotKeysMaxKeyLen +} + +// SubBucketBoundsFor returns the [lo, hi) key bounds of sub-bucket +// subBucket within the individual route routeID, mirroring what Flush +// emits as a MatrixRow's Start/End. The admin hot-keys handler uses +// this to filter the route's top-m snapshot to keys in the +// drill-down-clicked sub-range (design §5). +// +// ok is false when the route doesn't exist, is an aggregate (those are +// not eligible for hot-keys), or subBucket is out of range +// [0, len(subBuckets)). For a single-bucket slot (K=1 / unbounded / +// degenerate), bound calls with subBucket == 0 succeed and return the +// route's own Start/End. hi == nil means the sub-bucket extends to the +// route's unbounded tail. +// +// nil-receiver-safe. +func (s *MemSampler) SubBucketBoundsFor(routeID uint64, subBucket int) (lo, hi []byte, ok bool) { + if s == nil || subBucket < 0 { + return nil, nil, false + } + slot := lookupSlotForRoute(s.table.Load(), routeID) + if slot == nil { + return nil, nil, false + } + if subBucket >= len(slot.subBuckets) { + return nil, nil, false + } + start, end, _, _, _, _ := slot.snapshotMeta() + if len(slot.subBuckets) <= 1 { + return start, end, true + } + lo, hi = slot.subBucketBounds(subBucket, start, end) + return lo, hi, true +} + // 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+ @@ -545,7 +668,22 @@ func (s *MemSampler) Observe(routeID uint64, key []byte, op Op, valueLen int) { if byteCount > 0 { sc.writeBytes.Add(byteCount) } + s.observeHotKey(routeID, slot, key) + } +} + +// observeHotKey is the per-Observe hot-key sampler dispatch, kept out of +// Observe itself to stay under the cyclop ceiling. v1 tracks WRITES +// only (design §5 series-picker note); reads follow in v2 once the +// hot-path cost is measured. The nil check short-circuits in disabled +// deploys with a single branch — Observe's disabled-case allocation +// contract is preserved. Aggregates (Aggregate==true) are not tracked +// (design §2.2); lookupSlotForRoute filters them at publish too. +func (s *MemSampler) observeHotKey(routeID uint64, slot *routeSlot, key []byte) { + if s.hotKeys == nil || slot.Aggregate { + return } + s.hotKeys.observe(routeID, key) } // subBucketIndex maps key to a sub-range bucket index in diff --git a/keyviz/space_saving.go b/keyviz/space_saving.go new file mode 100644 index 000000000..e2e18bc2d --- /dev/null +++ b/keyviz/space_saving.go @@ -0,0 +1,115 @@ +package keyviz + +// Space-Saving sketch (Metwally et al.) — Misra–Gries-equivalent +// heavy-hitter detector backing the per-cell hot-key drill-down (Phase +// 2-A++; see docs/design/2026_05_28_proposed_keyviz_hot_key_topk.md +// §3). Maintains at most `capacity` (= `m`) entries; a key with +// frequency f over the observed stream of length N is reported with +// error ≤ N/m, and any key with f > N/m is guaranteed in the tracked +// set (over that stream). v1 keys the bookkeeping on the byte-string +// form of the key; an O(m) scan finds the min on eviction, which is +// cheap at the target m ≤ 256 and avoids a separate heap structure. +// +// NOT thread-safe: the hot-keys aggregator goroutine +// (`hot_keys.go::aggregator`) is the single writer and the only +// caller of observe / reset / snapshot. The published deep-copied +// `keyvizHotKeysSnapshot` is what drill-down readers see, lock-free +// via the slot's atomic.Pointer. + +type spaceSaving struct { + capacity int + entries map[string]*ssEntry +} + +type ssEntry struct { + // Key bytes owned by this sketch — independent of the pool buffer + // the hot path used to ferry the key into the aggregator, so the + // pool buffer can be released as soon as observe returns. + key []byte + count uint64 +} + +// ssEntrySnap is the deep-copied form returned by snapshot. The Key +// slice is freshly allocated, so a subsequent reset / eviction on the +// live sketch cannot mutate a published snapshot (Gemini medium on the +// design's round-1). +type ssEntrySnap struct { + Key []byte + Count uint64 +} + +func newSpaceSaving(capacity int) *spaceSaving { + if capacity < 1 { + capacity = 1 + } + return &spaceSaving{ + capacity: capacity, + entries: make(map[string]*ssEntry, capacity), + } +} + +// observe records one occurrence of key. The aggregator passes the +// pool-borrowed buffer; this method copies the bytes into an +// SS-owned slice on insert/replace so the caller can release the +// pool buffer immediately afterwards (design §4 pool-vs-snapshot). +func (s *spaceSaving) observe(key []byte) { + // map[string([]byte)] lookup is a compiler-special-cased zero-alloc + // access; only the insert path actually allocates a string key. + if e, ok := s.entries[string(key)]; ok { + e.count++ + return + } + kc := append([]byte(nil), key...) // SS-owned copy + if len(s.entries) < s.capacity { + s.entries[string(kc)] = &ssEntry{key: kc, count: 1} + return + } + // Capacity reached: evict the entry with the smallest counter and + // reuse it for the new key. The new counter is min_counter + 1 + // (Space-Saving's overestimate-on-insert rule, which bounds error + // at N/m). + var minE *ssEntry + var minK string + for k, e := range s.entries { + if minE == nil || e.count < minE.count { + minE = e + minK = k + } + } + delete(s.entries, minK) + minE.key = kc + minE.count++ + s.entries[string(kc)] = minE +} + +// snapshot returns a deep-copied list of every tracked (key, count) +// pair. The aggregator calls it under its tick before publishing; the +// returned slice never aliases the live sketch's storage, so an +// eviction or reset after publish is safe. +func (s *spaceSaving) snapshot() []ssEntrySnap { + if len(s.entries) == 0 { + return nil + } + out := make([]ssEntrySnap, 0, len(s.entries)) + for _, e := range s.entries { + out = append(out, ssEntrySnap{ + Key: append([]byte(nil), e.key...), + Count: e.count, + }) + } + return out +} + +// reset clears every entry. Called by the aggregator immediately after +// snapshot publish so the next keyvizStep window starts empty +// (design §4 sketch reset; Codex P2 L186). Capacity is preserved. +func (s *spaceSaving) reset() { + if len(s.entries) == 0 { + return + } + s.entries = make(map[string]*ssEntry, s.capacity) +} + +// len returns the current number of tracked entries (≤ capacity). +// Test helper; not used by production code. +func (s *spaceSaving) len() int { return len(s.entries) } diff --git a/keyviz/space_saving_test.go b/keyviz/space_saving_test.go new file mode 100644 index 000000000..5103496ba --- /dev/null +++ b/keyviz/space_saving_test.go @@ -0,0 +1,185 @@ +package keyviz + +import ( + "bytes" + "sort" + "testing" + + "github.com/stretchr/testify/require" + "pgregory.net/rapid" +) + +func TestSpaceSavingTracksUpToCapacity(t *testing.T) { + t.Parallel() + s := newSpaceSaving(4) + for _, k := range [][]byte{{0x01}, {0x02}, {0x03}, {0x04}} { + s.observe(k) + } + require.Equal(t, 4, s.len()) + snap := s.snapshot() + require.Len(t, snap, 4) + for _, e := range snap { + require.Equal(t, uint64(1), e.Count) + } +} + +func TestSpaceSavingIncrementsExisting(t *testing.T) { + t.Parallel() + s := newSpaceSaving(4) + for i := 0; i < 5; i++ { + s.observe([]byte{0xAA}) + } + require.Equal(t, 1, s.len()) + require.Equal(t, uint64(5), s.snapshot()[0].Count) +} + +func TestSpaceSavingEvictsMinAndIncrements(t *testing.T) { + t.Parallel() + s := newSpaceSaving(2) + s.observe([]byte("a")) // a:1 + s.observe([]byte("b")) // b:1 + for i := 0; i < 5; i++ { // b:6 + s.observe([]byte("b")) + } + // New key c evicts the minimum (a:1) and takes count = 1 + 1 = 2. + s.observe([]byte("c")) + require.Equal(t, 2, s.len()) + got := map[string]uint64{} + for _, e := range s.snapshot() { + got[string(e.Key)] = e.Count + } + require.Equal(t, uint64(6), got["b"]) + require.Equal(t, uint64(2), got["c"], "Space-Saving overestimate-on-insert: min_counter+1") + require.NotContains(t, got, "a", "evicted entry must be gone") +} + +func TestSpaceSavingResetClearsEntries(t *testing.T) { + t.Parallel() + s := newSpaceSaving(4) + for _, k := range [][]byte{{0x01}, {0x02}, {0x03}} { + s.observe(k) + } + s.reset() + require.Equal(t, 0, s.len()) + require.Nil(t, s.snapshot()) + // Capacity preserved after reset. + for _, k := range [][]byte{{0x10}, {0x11}, {0x12}, {0x13}, {0x14}} { + s.observe(k) + } + require.Equal(t, 4, s.len(), "capacity preserved across reset") +} + +// TestSpaceSavingSnapshotDeepCopy pins the Gemini-medium concern: a +// published snapshot must not alias the live sketch's storage, so a +// subsequent reset/eviction cannot mutate the snapshot a reader holds. +func TestSpaceSavingSnapshotDeepCopy(t *testing.T) { + t.Parallel() + s := newSpaceSaving(2) + s.observe([]byte("hello")) + snap := s.snapshot() + require.Len(t, snap, 1) + require.Equal(t, []byte("hello"), snap[0].Key) + // Mutating the live sketch must not touch the snapshot. + s.reset() + for i := 0; i < 100; i++ { + s.observe([]byte("world")) + } + require.Equal(t, []byte("hello"), snap[0].Key, "snapshot must be deep-copied") + require.Equal(t, uint64(1), snap[0].Count) +} + +// TestSpaceSavingObservedBufferOwnedIndependently pins the design's +// pool-vs-SS ownership rule: the caller's key buffer can be safely +// reused or mutated after observe returns, because SS makes its own +// copy on insert. +func TestSpaceSavingObservedBufferOwnedIndependently(t *testing.T) { + t.Parallel() + s := newSpaceSaving(4) + buf := []byte{'a', 'b', 'c'} + s.observe(buf) + // Mutate the caller buffer post-observe. + buf[0] = 'X' + buf[1] = 'Y' + buf[2] = 'Z' + snap := s.snapshot() + require.Len(t, snap, 1) + require.Equal(t, []byte("abc"), snap[0].Key, "SS must own a copy independent of the caller buffer") +} + +// TestSpaceSavingHeavyHitterGuarantee is the canonical Misra–Gries +// guarantee: a key with true frequency f > N/m is GUARANTEED in the +// tracked set after the stream is consumed. +func TestSpaceSavingHeavyHitterGuarantee(t *testing.T) { + t.Parallel() + rapid.Check(t, func(rt *rapid.T) { + m := rapid.IntRange(2, 32).Draw(rt, "m") + // Generate a stream of length N over a key alphabet of size + // noisyAlphabet. Pick one "hot" key with frequency strictly + // greater than N/m and assert it survives. + nNoisy := rapid.IntRange(0, 200).Draw(rt, "noisyCount") + hotFreq := rapid.IntRange(1, 200).Draw(rt, "hotFreq") + // Total stream length: + n := nNoisy + hotFreq + if hotFreq*m <= n { + // hotFreq must be > N/m for the guarantee. Skip cases that + // don't satisfy the precondition (Chernoff-style filter). + rt.SkipNow() + } + + s := newSpaceSaving(m) + // Interleave hot and noisy observations deterministically. + alphabet := []byte{0x10, 0x20, 0x30, 0x40, 0x50, 0x60, 0x70, 0x80} + hot := []byte("HOT") + hi, ni := 0, 0 + for hi < hotFreq || ni < nNoisy { + if hi < hotFreq && (ni >= nNoisy || (hi+ni)%2 == 0) { + s.observe(hot) + hi++ + } else { + k := []byte{alphabet[ni%len(alphabet)]} + s.observe(k) + ni++ + } + } + + // Heavy hitter MUST be present. + found := false + for _, e := range s.snapshot() { + if bytes.Equal(e.Key, hot) { + found = true + // hotFreq is from rapid.IntRange(1, 200) — strictly + // positive — so the conversion is wrap-safe; the + // guard keeps gosec G115 silent without a //nolint. + wantAtLeast := uint64(0) + if hotFreq > 0 { + wantAtLeast = uint64(hotFreq) + } + require.GreaterOrEqual(rt, e.Count, wantAtLeast, + "SS counter is an upper bound on true frequency, so >= true count") + break + } + } + require.True(rt, found, "key with f > N/m must be tracked (m=%d, N=%d, hotFreq=%d)", m, n, hotFreq) + }) +} + +// TestSpaceSavingSnapshotSortedByCountDescending is just a sanity for +// the drill-down handler: it relies on snapshot's contents but sorts +// them itself by descending count. This test pins the raw snapshot +// contract (unsorted, all entries returned). +func TestSpaceSavingSnapshotReturnsAllEntries(t *testing.T) { + t.Parallel() + s := newSpaceSaving(8) + counts := []int{10, 5, 1, 20, 3, 7} + for i, c := range counts { + k := []byte{byte(i)} + for j := 0; j < c; j++ { + s.observe(k) + } + } + snap := s.snapshot() + require.Len(t, snap, len(counts)) + sort.Slice(snap, func(i, j int) bool { return snap[i].Count > snap[j].Count }) + require.Equal(t, uint64(20), snap[0].Count) + require.Equal(t, uint64(1), snap[len(snap)-1].Count) +} diff --git a/main.go b/main.go index 9dcf87d08..f4fcab10c 100644 --- a/main.go +++ b/main.go @@ -194,6 +194,18 @@ var ( keyvizMaxMemberRoutesPerSlot = flag.Int("keyvizMaxMemberRoutesPerSlot", keyviz.DefaultMaxMemberRoutesPerSlot, "Maximum members listed on a virtual bucket; excess routes still drive the bucket counters") keyvizHistoryColumns = flag.Int("keyvizHistoryColumns", keyviz.DefaultHistoryColumns, "Maximum matrix columns retained in the keyviz ring buffer (each column = one Step)") keyvizKeyBucketsPerRoute = flag.Int("keyvizKeyBucketsPerRoute", keyviz.DefaultKeyBucketsPerRoute, "Order-preserving sub-range buckets per individual route for the hot-key heatmap; 1 disables sub-bucketing (route-granular, today's behaviour). Capped at 256; memory is ~K*32 bytes/route, so K_max ~= memBudget/(32*keyvizMaxTrackedRoutes)") + + // Hot-key drill-down (Phase 2-A++; design 2026_05_28_proposed_keyviz_hot_key_topk). + // Off by default — the disabled-case adds one early-return branch + // to Observe and retains zero real key bytes. When enabled, the + // sampler retains actual hot key bytes in memory and exposes them + // via the admin /keyviz/hotkeys drill-down (gated behind admin + // auth + audit). + keyvizHotKeysEnabled = flag.Bool("keyvizHotKeysEnabled", false, "Enable per-route Top-K hot-key drill-down (retains actual key bytes; admin auth + keyviz flag both required)") + keyvizHotKeysPerRoute = flag.Int("keyvizHotKeysPerRoute", keyviz.DefaultHotKeysPerRoute, "Space-Saving sketch capacity m per route (default 64, cap 256). Larger m tightens the error bound N_total/m at the cost of memory") + keyvizHotKeysSampleRate = flag.Int("keyvizHotKeysSampleRate", keyviz.DefaultHotKeysSampleRate, "Hot-keys hot-path sample rate R: 1-in-R observes are enqueued (default 16, cap 1024). Higher R reduces hot-path cost; raises Chernoff miss probability over the sampled stream") + keyvizHotKeysQueueSize = flag.Int("keyvizHotKeysQueueSize", keyviz.DefaultHotKeysQueueSize, "Bounded channel size between Observe and the hot-keys aggregator. Default 8192, cap 65536. Drops past this are counted (dropped_samples -> degraded)") + keyvizHotKeysMaxKeyLen = flag.Int("keyvizHotKeysMaxKeyLen", keyviz.DefaultHotKeysMaxKeyLen, "Maximum key length sampled into the hot-keys sketch (default 1024 B, cap 4096). Longer keys bump skipped_long_keys -> degraded, never truncated") // Phase 2-C cluster fan-out: comma-separated list of admin // HTTP endpoints (host:port or scheme://host:port). When set, // the admin keyviz handler aggregates the local matrix with @@ -1914,6 +1926,11 @@ func buildKeyVizSampler() *keyviz.MemSampler { MaxTrackedRoutes: *keyvizMaxTrackedRoutes, MaxMemberRoutesPerSlot: *keyvizMaxMemberRoutesPerSlot, KeyBucketsPerRoute: *keyvizKeyBucketsPerRoute, + HotKeysEnabled: *keyvizHotKeysEnabled, + HotKeysPerRoute: *keyvizHotKeysPerRoute, + HotKeysSampleRate: *keyvizHotKeysSampleRate, + HotKeysQueueSize: *keyvizHotKeysQueueSize, + HotKeysMaxKeyLen: *keyvizHotKeysMaxKeyLen, }) } @@ -1959,6 +1976,14 @@ func startKeyVizFlusher(ctx context.Context, eg *errgroup.Group, s *keyviz.MemSa s.Flush() return nil }) + // Hot-key drill-down aggregator: a separate goroutine on the same + // keyvizStep cadence. RunHotKeysAggregator is a no-op block when + // HotKeysEnabled is false, so calling it unconditionally is safe + // and keeps the startup wiring uniform regardless of the flag. + eg.Go(func() error { + keyviz.RunHotKeysAggregator(ctx, s) + return nil + }) } // startKeyVizLeaderTermPublisher polls each Raft group's current term diff --git a/main_admin.go b/main_admin.go index dd0a97d6d..e9290d6d9 100644 --- a/main_admin.go +++ b/main_admin.go @@ -1102,19 +1102,20 @@ func buildAdminHTTPServer(adminCfg *admin.Config, creds map[string]string, clust return nil, errors.Wrap(err, "open embedded admin SPA") } server, err := admin.NewServer(admin.ServerDeps{ - Signer: signer, - Verifier: verifier, - Credentials: admin.MapCredentialStore(creds), - Roles: adminCfg.RoleIndex(), - ClusterInfo: cluster, - Tables: tables, - Buckets: buckets, - Queues: queues, - Forwarder: forwarder, - KeyViz: keyvizSourceFromSampler(keyvizSampler), - KeyVizFanout: buildKeyVizFanout(adminCfg.Listen, keyvizFanoutCfg), - StaticFS: staticFS, - LeaderProbe: leaderProbe, + Signer: signer, + Verifier: verifier, + Credentials: admin.MapCredentialStore(creds), + Roles: adminCfg.RoleIndex(), + ClusterInfo: cluster, + Tables: tables, + Buckets: buckets, + Queues: queues, + Forwarder: forwarder, + KeyViz: keyvizSourceFromSampler(keyvizSampler), + KeyVizFanout: buildKeyVizFanout(adminCfg.Listen, keyvizFanoutCfg), + KeyVizHotKeys: keyvizHotKeysSourceFromSampler(keyvizSampler), + StaticFS: staticFS, + LeaderProbe: leaderProbe, AuthOpts: admin.AuthServiceOpts{ InsecureCookie: adminCfg.AllowInsecureDevCookie, }, @@ -1253,6 +1254,21 @@ func keyvizSourceFromSampler(s *keyviz.MemSampler) admin.KeyVizSource { return s } +// keyvizHotKeysSourceFromSampler is the keyvizSourceFromSampler twin +// for the per-cell hot-key drill-down handler. Returning an interface-nil +// (not a typed-nil) when the sampler is disabled is again load-bearing: +// the hot-keys handler's "keyviz disabled → 503" branch checks for +// interface-nil, and a typed-nil *MemSampler stored as a non-nil +// interface would route the request into MemSampler methods that +// already nil-check internally — correct but masking the explicit +// "feature off" signal we want the SPA to render. +func keyvizHotKeysSourceFromSampler(s *keyviz.MemSampler) admin.KeyVizHotKeysSource { + if s == nil { + return nil + } + return s +} + // buildKeyVizFanout assembles the Phase 2-C fan-out aggregator from // the operator-supplied flag values. selfListen is the local admin // listener address (used to filter the local node out of the peer