From aaa46f59944a4b649eeac0ad761a96413da26c63 Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Sun, 16 Aug 2026 15:00:36 +0200 Subject: [PATCH 01/12] docs(ratelimit): spec for persisting counters across restarts --- ...-16_rate-limit-counter-persistence-spec.md | 301 ++++++++++++++++++ 1 file changed, 301 insertions(+) create mode 100644 docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md diff --git a/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md b/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md new file mode 100644 index 000000000..5901fd3eb --- /dev/null +++ b/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md @@ -0,0 +1,301 @@ +# Rate Limit Counter Persistence + +Status: approved design +Date: 2026-08-16 +Extends: `docs/dev/2026-07-05_rate-limiting-spec.md` §8, §11.2, §11.7 + +## 1. Why + +Request and token windows live only in process memory. A restart or +`gomodel --reload` (SIGHUP) starts them at zero. Minute windows barely +notice; an hour or day cap — the thing operators use for daily provider +quotas and team allowances — silently resets. The original spec documented +this and listed two follow-ups: + +- §11.7 counter persistence across restarts (this work) +- §11.2 Redis live counters for exact multi-replica enforcement (later) + +This spec does §11.7 and extracts the seam §11.2 will need. It does not +implement Redis-Lua. + +## 2. Goals + +- Request and token sliding windows survive process restart, SIGTERM, and + `--reload`. +- `admit()` stays in memory. Persistence is a background snapshot, never on + the request path. +- A later Redis-Lua backend can replace the in-memory limiter without + changing `Acquire`, the usage tap, or `RouteAvailable`. +- Default on whenever rate limits are on. No extra dependency. Uses the + existing SQLite / Postgres / Mongo store that already holds the rules. +- Crash loses at most one flush interval. Graceful stop and reload lose + nothing. + +## 3. Non-goals + +- Shared counters across replicas. N instances still mean about N × the + configured limit. Redis-Lua is the follow-up. +- Persisting concurrency gauges. After a restart nothing is in flight. +- Pre-call token reservation, per-(path × model) rules, or any other item + from the original §11. +- A separate persist-on/off flag. `RATE_LIMITS_ENABLED=false` already + builds no limiter. `RATE_LIMITS_FLUSH_INTERVAL=0` is the escape hatch + that turns off the periodic loop only. + +## 4. Architecture + +Two seams, only the first used as a live store today. + +``` +Service.Acquire / RecordTokens / RouteAvailable / Status / Reset + │ + ▼ + counterBackend (unexported; package ratelimit) + │ + ├── memoryBackend current limiter + snapshot load/save + └── (later) redisLua live INCR; no snapshots +``` + +Snapshots are an implementation detail of `memoryBackend`. They persist +only the sliding windows. The backend talks to the existing `Store`, +which grows four methods implemented by the SQL and Mongo stores already +used for rules. + +Cardinality stays one row per rule, not per caller. + +## 5. `counterBackend` + +Unexported interface in `internal/ratelimit`, matching what `limiter` +already does: + +- `Admit(rules []Rule, now time.Time) (HeaderSnapshot, []ruleKey, *ExceededError)` +- `Available(rules []Rule, now time.Time) bool` +- `Release(held []ruleKey)` +- `RecordTokens(rules []Rule, tokens int64, now time.Time)` +- `Status(rule Rule, now time.Time) Status` +- `Reset(key ruleKey)` +- `ResetAll()` + +`Service` holds a `counterBackend` instead of a concrete `*limiter`. +`ruleKey` stays unexported. No public API change except `Service.Start` +and the new config field. + +The memory implementation is the current `limiter` plus snapshot +helpers. A future Redis-Lua type implements the same interface and +ignores the snapshot store. + +## 6. Snapshot data model + +One row per windowed rule. Concurrent rules (`period_seconds = 0`) are +never written. + +```text +scope TEXT -- user_path | provider | model +subject TEXT +period_seconds INT64 -- > 0 +requests_window_start INT64 -- unix seconds, 0 if unused +requests_current INT64 +requests_previous INT64 +tokens_window_start INT64 +tokens_current INT64 +tokens_previous INT64 +updated_at INT64 -- unix seconds, diagnostics only +PRIMARY KEY (scope, subject, period_seconds) +``` + +SQL table: `rate_limit_counters`. Mongo collection: `rate_limit_counters` +with a unique index on `(scope, subject, period_seconds)`. Created at +store init with `CREATE TABLE IF NOT EXISTS` / `CreateIndex`. Older +databases just gain the table; no backfill. + +Units match `windowCounter`: Unix seconds and `int64` counts. After +load, `advance()` already zeros a window more than one period old, so +rows do not need a TTL. + +Package type (exported only if tests in another package need it; +otherwise unexported is fine): + +```go +type windowSnapshot struct { + Scope, Subject string + PeriodSeconds int64 + RequestsWindowStart, RequestsCurrent, RequestsPrevious int64 + TokensWindowStart, TokensCurrent, TokensPrevious int64 +} +``` + +A rule that only limits requests leaves the token fields at zero, and +the reverse. Zero `windowStart` with zero counts means “no window yet.” + +## 7. Store methods + +Add to the existing `Store` interface. No second factory. + +- `LoadCounters(ctx) ([]windowSnapshot, error)` — all rows. +- `SaveCounters(ctx, []windowSnapshot) error` — one transaction that + **replaces** the table/collection with the provided set (delete all, + then insert). The memory backend always passes the full live window + set, so orphans disappear on the next flush. +- `DeleteCounter(ctx, scope, subject, periodSeconds) error` — one row. + Missing row is not an error. +- `DeleteAllCounters(ctx) error` — empty the table. + +`Close` is unchanged. + +## 8. Lifecycle + +Reload builds the next generation while the current one is still +serving, then shuts the old one down, then starts the new one +(`run.serveUntilShutdown` / `watchForReload`). That order is load-bearing. + +| Call | When | Persistence | +|---|---|---| +| `New` / `NewService` | `app.New`, tests | Empty limiter. No load, no flush, no loop. | +| `Start` | `App.startServer`, after the previous generation’s `Shutdown` | Load snapshot, apply rows whose key still matches a rule, then start the flush loop if `flush_interval > 0`. Mark the generation **active**. | +| `Close` (`Result.Close`) | `App.Shutdown` | Stop the loop. If **active**, write one last snapshot. Then close the store as today. | + +An idle replacement that is built and then discarded (failed listen, +process exiting during rebuild) is not active, so its `Close` must not +write an empty snapshot over the live one. + +`App.startServer` calls `rateLimits.Service.Start(ctx)` before the HTTP +server accepts connections. Tests that construct a `Service` and never +call `Start` keep today’s in-memory-only behavior. + +Flush copies the request and token maps under the limiter mutex, then +writes outside the lock. `admit()` never waits on storage. + +## 9. Configuration + +```yaml +rate_limits: + enabled: true + flush_interval: 1 # seconds; 0 = no periodic loop +``` + +```env +RATE_LIMITS_FLUSH_INTERVAL=1 +``` + +- Default `1` (set next to `Enabled: true` in `config.Load` defaults). +- `0` disables the periodic loop only. `Start` still loads. `Close` of + an active generation still writes once. `--reload` and SIGTERM stay + correct; SQLite’s single connection gets no extra writer on a timer. +- Negative values are rejected at config validation. +- No maximum. A large value just widens crash loss. +- Persistence is on whenever `RATE_LIMITS_ENABLED` is on. There is no + second flag. + +Document in `.env.template`, `config/config.example.yaml`, +`docs/features/rate-limits.mdx`, and `CLAUDE.md` / `Agents.md`. Remove +the wording that counters reset on restart / `--reload`. Keep the +wording that counters are per instance (N replicas ≈ N × limit) and +that concurrency is in-memory only. + +## 10. Service operations vs snapshot + +| Operation | Memory | Snapshot | +|---|---|---| +| `Admit` / `RecordTokens` | update windows | next flush / Close | +| `ResetRule` | `limiter.reset` | `DeleteCounter` immediately | +| `ResetAll` | `limiter.resetAll` | `DeleteAllCounters` immediately | +| `DeleteRule` | reset that key | `DeleteCounter` immediately | +| `ReplaceConfigRules` | refresh rules | after refresh, `DeleteCounter` for every snapshot key that is no longer a windowed rule | + +Reset and delete clear memory first (the operator-visible effect), then +the row. A failed row delete is logged; the in-memory reset still +stands. The next successful flush replace-all drops the stale row. + +## 11. Error handling + +Persistence never fails a request. `admit()` does not see store errors. + +- **Load failure:** log and start with empty windows. Do not block the + listener. Skip corrupt rows; restore the good ones. Still mark the + generation active so later flushes work. +- **Periodic flush failure:** log and retry on the next tick. Memory + stays authoritative. +- **Shutdown flush failure:** log and continue teardown. Do not hang + past the existing shutdown budget. +- **Reset/delete row failure:** log at error level. Memory is already + cleared. +- **Migrate:** creating the new table/collection fails store init the + same way a missing `rate_limits` table would — that is a hard start + error, not a soft persist error. + +## 12. Testing + +### Unit / store + +- Snapshot encode/restore: `estimate` after load matches the pre-flush + value for both request and token windows. +- A snapshot whose `windowStart` is more than one period old advances + to zero. +- Concurrent keys never appear in a snapshot. +- `New` does not read or write. `Start` loads. `Close` after `Start` + writes a final snapshot. `Close` without `Start` writes nothing. +- `flush_interval=0` still loads and still flushes on `Close`; the loop + never ticks. +- `ResetRule` / `ResetAll` / `DeleteRule` clear memory and the row; a + later `Start` on a new service does not resurrect them. +- `ReplaceConfigRules` dropping a config rule deletes that row. +- SQL store (and Mongo, same as rules) round-trip: replace, load, + delete one, delete all. Migration creates `rate_limit_counters` on an + existing DB. +- Config: default interval is 1; `RATE_LIMITS_FLUSH_INTERVAL=0` is + valid; a negative value is rejected. +- Existing admit/release/header tests stay in-memory. A recording + `Store` asserts `Admit` does not call `SaveCounters`. + +### Release E2E + +Add to `tests/e2e/release-e2e-scenarios.md` (after S204). Update the +file header count and the stateful-note list. + +Shared helper in the common environment block: + +- Export `RELEASE_STACK_DIR` (default `/tmp/gomodel-release-stack`). +- `reload_release_gateway ` sends `SIGHUP` to + `$RELEASE_STACK_DIR//server.pid`, then waits until that + gateway’s `logs/server.log` contains a **new** `configuration reloaded` + line. A request sent immediately after `kill -HUP` can still hit the + old generation, which still has in-memory counters — the log wait is + required. + +Use **hour** windows so the cap outlives the reload wait. Each scenario +creates a `$QA_SUFFIX`-scoped user-path rule and deletes it. + +- **S205 — Request-window counters survive `--reload` (SQLite).** + `max_requests=1` on `$BASE_URL`. First chat succeeds. Sleep ~2s so + the 1s flush lands. Reload `sqlite-main`. Second chat is `429` with + `code: rate_limit_exceeded`. Delete the rule. +- **S206 — `reset-one` stays cleared across `--reload`.** Same shape: + burn the hour window, `reset-one`, sleep, reload `sqlite-main`. Next + chat succeeds. Delete the rule. +- **S207 — Same request-window survival on PostgreSQL and MongoDB.** + S203-style loop over `$PG_BASE_URL` / `$MONGO_BASE_URL`, reloading + `pg-smoke` and `mongo-smoke`. Distinct paths, delete each rule. + +S205–S207 reload a shared gateway. That is safe in this sequential +runner (same class as S137). Token-window reload is covered by S157 +plus unit tests; concurrent is not persisted and is not an E2E case. + +## 13. Docs and comments + +- `docs/features/rate-limits.mdx` — windows survive restart and + `--reload`; concurrency does not; still per instance. +- `docs/dev/2026-07-05_rate-limiting-spec.md` §8 and §11.7 — mark + persistence done; §11.2 stays future work. +- `CLAUDE.md` / `Agents.md` — drop rate-limit counters from the + “in-memory state resets on reload” list; keep session affinity and + live log buffers. +- `.env.template` and `config/config.example.yaml` — document + `RATE_LIMITS_FLUSH_INTERVAL` / `flush_interval`. + +## 14. Follow-up (not this change) + +Redis-Lua `counterBackend`: every `Admit` is an atomic script on shared +keys. No snapshots. Chosen when `REDIS_URL` is set, or behind an +explicit config once someone is running HA and wants N replicas to +share one limit. The interface in §5 is the only preparation this +change makes for that work. From d87171a90b6510673df5491e8276113ea0f14cee Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Sun, 16 Aug 2026 15:15:11 +0200 Subject: [PATCH 02/12] docs(ratelimit): fix persistence spec races found in review --- ...-16_rate-limit-counter-persistence-spec.md | 77 +++++++++++++++---- 1 file changed, 61 insertions(+), 16 deletions(-) diff --git a/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md b/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md index 5901fd3eb..8a873d7e7 100644 --- a/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md +++ b/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md @@ -63,6 +63,12 @@ used for rules. Cardinality stays one row per rule, not per caller. +Several replicas sharing Postgres or Mongo last-write-wins on that +row. That is accepted: this work does not give shared live counters, and +the schema has no `instance_id`. A restart loads whoever flushed last. +Use one instance, or wait for the Redis-Lua backend, when the limit +must be exact across replicas. + ## 5. `counterBackend` Unexported interface in `internal/ratelimit`, matching what `limiter` @@ -134,8 +140,9 @@ Add to the existing `Store` interface. No second factory. - `LoadCounters(ctx) ([]windowSnapshot, error)` — all rows. - `SaveCounters(ctx, []windowSnapshot) error` — one transaction that **replaces** the table/collection with the provided set (delete all, - then insert). The memory backend always passes the full live window - set, so orphans disappear on the next flush. + then insert). The payload is only **current windowed rules** (period + > 0 still present in `Service.rules`), never leftover limiter-map + entries for a rule that was dropped. - `DeleteCounter(ctx, scope, subject, periodSeconds) error` — one row. Missing row is not an error. - `DeleteAllCounters(ctx) error` — empty the table. @@ -165,6 +172,22 @@ call `Start` keep today’s in-memory-only behavior. Flush copies the request and token maps under the limiter mutex, then writes outside the lock. `admit()` never waits on storage. +A second mutex (`persistMu`) serializes snapshot writes against +reset/delete row deletes. `admit()` does not take it. + +1. Flush: lock `persistMu`, copy maps (only keys that are still + windowed rules), `SaveCounters`, unlock. +2. `Reset` / `DeleteRule` / `ResetAll`: clear memory, lock `persistMu`, + delete the row(s), unlock. + +Without that, a flush that sampled before a reset can `SaveCounters` +after the row delete and resurrect the burned window (S206 would flake). + +Orphan snapshot rows (rule gone, row still there) are dropped only by +an **active** generation: `Start` applies matching keys and ignores the +rest; the next flush replace-all omits them. Construction must not +delete snapshot rows — see §10. + ## 9. Configuration ```yaml @@ -197,14 +220,27 @@ that concurrency is in-memory only. | Operation | Memory | Snapshot | |---|---|---| | `Admit` / `RecordTokens` | update windows | next flush / Close | -| `ResetRule` | `limiter.reset` | `DeleteCounter` immediately | -| `ResetAll` | `limiter.resetAll` | `DeleteAllCounters` immediately | -| `DeleteRule` | reset that key | `DeleteCounter` immediately | -| `ReplaceConfigRules` | refresh rules | after refresh, `DeleteCounter` for every snapshot key that is no longer a windowed rule | +| `ResetRule` | `limiter.reset` | `DeleteCounter` under `persistMu` | +| `ResetAll` | `limiter.resetAll` | `DeleteAllCounters` under `persistMu` | +| `DeleteRule` | reset that key **and drop it from the limiter maps** | `DeleteCounter` under `persistMu` | +| `ReplaceConfigRules` | refresh rules; **prune limiter maps** for keys that are no longer a windowed rule | **do not touch the store** | + +`ReplaceConfigRules` runs from `factory.New` → `seedConfiguredRules` +while the previous generation is still serving. Deleting snapshot rows +there would wipe a live hour/day window if the replacement is then +discarded (failed reload). Orphans are left for the **active** +generation: `Start` does not apply them; the next flush replace-all +drops them. + +When the rule set changes (`ReplaceConfigRules`, `DeleteRule`), prune +the request/token maps so a later flush cannot reinsert a dropped +rule’s window. `SaveCounters` is only ever given current windowed +rules. Reset and delete clear memory first (the operator-visible effect), then -the row. A failed row delete is logged; the in-memory reset still -stands. The next successful flush replace-all drops the stale row. +the row under `persistMu`. A failed row delete is logged; the in-memory +reset still stands. `persistMu` is what stops the next flush from +putting the row back. ## 11. Error handling @@ -238,7 +274,11 @@ Persistence never fails a request. `admit()` does not see store errors. never ticks. - `ResetRule` / `ResetAll` / `DeleteRule` clear memory and the row; a later `Start` on a new service does not resurrect them. -- `ReplaceConfigRules` dropping a config rule deletes that row. +- `ReplaceConfigRules` dropping a config rule prunes the limiter maps + and does not call `DeleteCounter`. After `Start`, the next flush + omit that key from `SaveCounters`. +- An in-flight flush cannot resurrect a completed `ResetRule` / + `DeleteRule` (`persistMu`). - SQL store (and Mongo, same as rules) round-trip: replace, load, delete one, delete all. Migration creates `rate_limit_counters` on an existing DB. @@ -266,12 +306,15 @@ Use **hour** windows so the cap outlives the reload wait. Each scenario creates a `$QA_SUFFIX`-scoped user-path rule and deletes it. - **S205 — Request-window counters survive `--reload` (SQLite).** - `max_requests=1` on `$BASE_URL`. First chat succeeds. Sleep ~2s so - the 1s flush lands. Reload `sqlite-main`. Second chat is `429` with - `code: rate_limit_exceeded`. Delete the rule. + `max_requests=1` on `$BASE_URL`. First chat succeeds. Reload + `sqlite-main` (old `Close` writes the snapshot; no sleep required + for the happy path). Second chat is `429` with + `code: rate_limit_exceeded`. Delete the rule. Crash-before-`Close` + (periodic flush only) is a unit test, not this scenario. - **S206 — `reset-one` stays cleared across `--reload`.** Same shape: - burn the hour window, `reset-one`, sleep, reload `sqlite-main`. Next - chat succeeds. Delete the rule. + burn the hour window, `reset-one`, reload `sqlite-main`. Next chat + succeeds. Delete the rule. `persistMu` is what makes this + deterministic. - **S207 — Same request-window survival on PostgreSQL and MongoDB.** S203-style loop over `$PG_BASE_URL` / `$MONGO_BASE_URL`, reloading `pg-smoke` and `mongo-smoke`. Distinct paths, delete each rule. @@ -282,8 +325,10 @@ plus unit tests; concurrent is not persisted and is not an E2E case. ## 13. Docs and comments -- `docs/features/rate-limits.mdx` — windows survive restart and - `--reload`; concurrency does not; still per instance. +- `docs/features/rate-limits.mdx` and `docs/advanced/cli.mdx` — + windows survive restart and `--reload`; concurrency does not; still + per instance. Drop the “counters start fresh” wording on the CLI + reload page. - `docs/dev/2026-07-05_rate-limiting-spec.md` §8 and §11.7 — mark persistence done; §11.2 stays future work. - `CLAUDE.md` / `Agents.md` — drop rate-limit counters from the From 37dfb1d36297282d04d0448890a0f10620562c10 Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Sun, 16 Aug 2026 15:21:49 +0200 Subject: [PATCH 03/12] docs(ratelimit): add review notes to persistence spec --- ...-16_rate-limit-counter-persistence-spec.md | 23 ++++++++++++------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md b/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md index 8a873d7e7..98d9f9b73 100644 --- a/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md +++ b/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md @@ -138,11 +138,14 @@ the reverse. Zero `windowStart` with zero counts means “no window yet.” Add to the existing `Store` interface. No second factory. - `LoadCounters(ctx) ([]windowSnapshot, error)` — all rows. -- `SaveCounters(ctx, []windowSnapshot) error` — one transaction that - **replaces** the table/collection with the provided set (delete all, - then insert). The payload is only **current windowed rules** (period - > 0 still present in `Service.rules`), never leftover limiter-map - entries for a rule that was dropped. +- `SaveCounters(ctx, []windowSnapshot) error` — **replaces** the + table/collection with the provided set (delete all, then insert). + SQL does this in one transaction. Mongo uses the same + transaction-plus-standalone-fallback as `ReplaceConfigRules` today + (replica-set e2e uses `rs0`; a hard transaction-only write would + fail open on standalone Mongo). The payload is only **current + windowed rules** (period > 0 still present in `Service.rules`), + never leftover limiter-map entries for a rule that was dropped. - `DeleteCounter(ctx, scope, subject, periodSeconds) error` — one row. Missing row is not an error. - `DeleteAllCounters(ctx) error` — empty the table. @@ -176,7 +179,8 @@ A second mutex (`persistMu`) serializes snapshot writes against reset/delete row deletes. `admit()` does not take it. 1. Flush: lock `persistMu`, copy maps (only keys that are still - windowed rules), `SaveCounters`, unlock. + windowed rules) **by value** (`windowCounter` structs, not the + pointers `admit()` still mutates), `SaveCounters`, unlock. 2. `Reset` / `DeleteRule` / `ResetAll`: clear memory, lock `persistMu`, delete the row(s), unlock. @@ -298,8 +302,11 @@ Shared helper in the common environment block: - `reload_release_gateway ` sends `SIGHUP` to `$RELEASE_STACK_DIR//server.pid`, then waits until that gateway’s `logs/server.log` contains a **new** `configuration reloaded` - line. A request sent immediately after `kill -HUP` can still hit the - old generation, which still has in-memory counters — the log wait is + line, then retries `/health` briefly. `configuration reloaded` is + logged after the old `Shutdown` and before the new + `StartWithListener`; the next request may sit in the held accept + queue until the new listener is up. A request sent immediately after + `kill -HUP` can still hit the old generation — the log wait is required. Use **hour** windows so the cap outlives the reload wait. Each scenario From 71031be2f60377e34493bec58813e879dccbf9da Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Sun, 16 Aug 2026 15:32:00 +0200 Subject: [PATCH 04/12] docs(ratelimit): persist per-child window partitions --- ...-16_rate-limit-counter-persistence-spec.md | 147 ++++++++++++------ 1 file changed, 102 insertions(+), 45 deletions(-) diff --git a/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md b/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md index 98d9f9b73..69c346430 100644 --- a/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md +++ b/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md @@ -1,8 +1,9 @@ # Rate Limit Counter Persistence -Status: approved design +Status: approved design (updated 2026-08-16 for `#670` per-child templates) Date: 2026-08-16 Extends: `docs/dev/2026-07-05_rate-limiting-spec.md` §8, §11.2, §11.7 +Depends on: `feat(quotas): add per-child user-path templates` (`#670`) ## 1. Why @@ -41,6 +42,10 @@ implement Redis-Lua. - A separate persist-on/off flag. `RATE_LIMITS_ENABLED=false` already builds no limiter. `RATE_LIMITS_FLUSH_INTERVAL=0` is the escape hatch that turns off the periodic loop only. +- Changing `#670` semantics or the `quota_templates` entitlement gate. + Persistence lives in OSS `internal/ratelimit`. Pro (`../gomodel-pro`) + only enables the capability; it gets snapshot restore of child + partitions for free. ## 4. Architecture @@ -61,13 +66,26 @@ only the sliding windows. The backend talks to the existing `Store`, which grows four methods implemented by the SQL and Mongo stores already used for rules. -Cardinality stays one row per rule, not per caller. +Cardinality is one row per **live window key**, not one per rule +definition: -Several replicas sharing Postgres or Mongo last-write-wins on that -row. That is accepted: this work does not give shared live counters, and -the schema has no `instance_id`. A restart loads whoever flushed last. -Use one instance, or wait for the Redis-Lua backend, when the limit -must be exact across replicas. +- A shared rule (`per_child=false`, and every provider/model rule) still + has one window. `partition` is empty. +- A per-child user-path template (`#670`) has one window per active + direct child. `partition` is `Rule.EffectiveSubject` (for example + `/customers/alice` under a template on `/customers`). Deeper paths + share that child’s counter, same as live enforcement. + +Idle child partitions are already dropped from memory after two periods +(`limiter_expiry.go`). The snapshot only writes keys that are still in +the limiter maps, so table size tracks live children, not historical +ones. + +Several replicas sharing Postgres or Mongo last-write-wins **per +window key**. That is accepted: this work does not give shared live +counters, and the schema has no `instance_id`. A restart loads whoever +flushed last. Use one instance, or wait for the Redis-Lua backend, when +the limit must be exact across replicas. ## 5. `counterBackend` @@ -83,8 +101,13 @@ already does: - `ResetAll()` `Service` holds a `counterBackend` instead of a concrete `*limiter`. -`ruleKey` stays unexported. No public API change except `Service.Start` -and the new config field. +`ruleKey` stays unexported and already includes `partition` (`#670`). +`Reset(key)` keeps today’s `sameDefinition` behavior: it clears every +partition of that `(scope, subject, period)` definition. + +Public API additions: `Service.Start` and the flush-interval config +field. `Service.Close` already exists (stops the child-expiry worker); +persistence `Close` must call it after the final snapshot. The memory implementation is the current `limiter` plus snapshot helpers. A future Redis-Lua type implements the same interface and @@ -92,12 +115,13 @@ ignores the snapshot store. ## 6. Snapshot data model -One row per windowed rule. Concurrent rules (`period_seconds = 0`) are -never written. +One row per live window key. Concurrent rules (`period_seconds = 0`) +are never written. ```text scope TEXT -- user_path | provider | model -subject TEXT +subject TEXT -- rule definition subject (template path) +partition TEXT -- "" for shared; child path for per-child period_seconds INT64 -- > 0 requests_window_start INT64 -- unix seconds, 0 if unused requests_current INT64 @@ -106,24 +130,32 @@ tokens_window_start INT64 tokens_current INT64 tokens_previous INT64 updated_at INT64 -- unix seconds, diagnostics only -PRIMARY KEY (scope, subject, period_seconds) +PRIMARY KEY (scope, subject, partition, period_seconds) ``` +`partition` is `ruleKey.partition` / `EffectiveSubject`. It is **not** +the request’s full user path. A template on `/customers` stores +`subject=/customers`, `partition=/customers/alice`, never +`/customers/alice/app`. + SQL table: `rate_limit_counters`. Mongo collection: `rate_limit_counters` -with a unique index on `(scope, subject, period_seconds)`. Created at -store init with `CREATE TABLE IF NOT EXISTS` / `CreateIndex`. Older -databases just gain the table; no backfill. +with a unique index on `(scope, subject, partition, period_seconds)`. +Created at store init with `CREATE TABLE IF NOT EXISTS` / `CreateIndex`. +Older databases just gain the table; no backfill. Units match `windowCounter`: Unix seconds and `int64` counts. After load, `advance()` already zeros a window more than one period old, so -rows do not need a TTL. +rows do not need a TTL. Load still skips a child row whose window is +already past the in-memory expiry horizon (`windowStart + 2*period`), +and must call `trackCounterExpiry` for every restored partition so the +existing worker keeps pruning. Package type (exported only if tests in another package need it; otherwise unexported is fine): ```go type windowSnapshot struct { - Scope, Subject string + Scope, Subject, Partition string PeriodSeconds int64 RequestsWindowStart, RequestsCurrent, RequestsPrevious int64 TokensWindowStart, TokensCurrent, TokensPrevious int64 @@ -132,6 +164,7 @@ type windowSnapshot struct { A rule that only limits requests leaves the token fields at zero, and the reverse. Zero `windowStart` with zero counts means “no window yet.” +A shared rule’s `Partition` is `""`. ## 7. Store methods @@ -143,11 +176,15 @@ Add to the existing `Store` interface. No second factory. SQL does this in one transaction. Mongo uses the same transaction-plus-standalone-fallback as `ReplaceConfigRules` today (replica-set e2e uses `rs0`; a hard transaction-only write would - fail open on standalone Mongo). The payload is only **current - windowed rules** (period > 0 still present in `Service.rules`), - never leftover limiter-map entries for a rule that was dropped. -- `DeleteCounter(ctx, scope, subject, periodSeconds) error` — one row. - Missing row is not an error. + fail open on standalone Mongo). The payload is every **live window + key** still in the limiter maps whose definition is a current + windowed rule (period > 0). That includes one row per active + per-child partition. It never includes leftover map entries for a + dropped definition, and never includes concurrent keys. +- `DeleteCounter(ctx, scope, subject, periodSeconds) error` — every + partition of that **definition** (matches `limiter.reset` / + `sameDefinition`). Missing rows are not an error. Resetting a + per-child template must not leave sibling children in the table. - `DeleteAllCounters(ctx) error` — empty the table. `Close` is unchanged. @@ -161,8 +198,8 @@ serving, then shuts the old one down, then starts the new one | Call | When | Persistence | |---|---|---| | `New` / `NewService` | `app.New`, tests | Empty limiter. No load, no flush, no loop. | -| `Start` | `App.startServer`, after the previous generation’s `Shutdown` | Load snapshot, apply rows whose key still matches a rule, then start the flush loop if `flush_interval > 0`. Mark the generation **active**. | -| `Close` (`Result.Close`) | `App.Shutdown` | Stop the loop. If **active**, write one last snapshot. Then close the store as today. | +| `Start` | `App.startServer`, after the previous generation’s `Shutdown` | Load snapshot, apply rows whose **definition** still matches a current rule and whose `partition` matches the rule’s mode (empty iff `!PerChild`), re-arm child expiry, then start the flush loop if `flush_interval > 0`. Mark the generation **active**. | +| `Close` (`Result.Close`) | `App.Shutdown` | Stop the flush loop. If **active**, write one last snapshot. Then `Service.Close()` (expiry worker, already wired today) and close the store. | An idle replacement that is built and then discarded (failed listen, process exiting during rebuild) is not active, so its `Close` must not @@ -178,11 +215,14 @@ writes outside the lock. `admit()` never waits on storage. A second mutex (`persistMu`) serializes snapshot writes against reset/delete row deletes. `admit()` does not take it. -1. Flush: lock `persistMu`, copy maps (only keys that are still - windowed rules) **by value** (`windowCounter` structs, not the - pointers `admit()` still mutates), `SaveCounters`, unlock. -2. `Reset` / `DeleteRule` / `ResetAll`: clear memory, lock `persistMu`, - delete the row(s), unlock. +1. Flush: lock `persistMu`, copy maps **by value** (`windowCounter` + structs, not the pointers `admit()` still mutates). Keep a key only + when its definition is still a windowed rule. Merge request+token + windows that share a `ruleKey` into one `windowSnapshot`. + `SaveCounters`, unlock. +2. `Reset` / `DeleteRule` / `ResetAll`: clear memory (all partitions of + the definition), lock `persistMu`, delete those snapshot rows, + unlock. Without that, a flush that sampled before a reset can `SaveCounters` after the row delete and resurrect the burned window (S206 would flake). @@ -224,10 +264,10 @@ that concurrency is in-memory only. | Operation | Memory | Snapshot | |---|---|---| | `Admit` / `RecordTokens` | update windows | next flush / Close | -| `ResetRule` | `limiter.reset` | `DeleteCounter` under `persistMu` | +| `ResetRule` | `limiter.reset` (all partitions of the definition) | `DeleteCounter` under `persistMu` (all partitions) | | `ResetAll` | `limiter.resetAll` | `DeleteAllCounters` under `persistMu` | -| `DeleteRule` | reset that key **and drop it from the limiter maps** | `DeleteCounter` under `persistMu` | -| `ReplaceConfigRules` | refresh rules; **prune limiter maps** for keys that are no longer a windowed rule | **do not touch the store** | +| `DeleteRule` | `limiter.reset` of that definition | `DeleteCounter` under `persistMu` | +| `ReplaceConfigRules` / `Refresh` | already prunes maps when a definition disappears or `PerChild` flips (`Service.Refresh` after `#670`) | **do not touch the store** | `ReplaceConfigRules` runs from `factory.New` → `seedConfiguredRules` while the previous generation is still serving. Deleting snapshot rows @@ -236,10 +276,15 @@ discarded (failed reload). Orphans are left for the **active** generation: `Start` does not apply them; the next flush replace-all drops them. -When the rule set changes (`ReplaceConfigRules`, `DeleteRule`), prune -the request/token maps so a later flush cannot reinsert a dropped -rule’s window. `SaveCounters` is only ever given current windowed -rules. +Do not add a second prune path. `Refresh` already drops limiter keys +when a definition is removed or shared/per-child mode changes; a later +flush therefore cannot reinsert those windows. `SaveCounters` is only +ever given live keys whose definition is still a windowed rule. + +A per-child row is applied on `Start` only if the matching definition +still has `PerChild=true`. A shared row is applied only if that +definition is not per-child. A mode flip therefore starts empty for +that definition, matching `Refresh`. Reset and delete clear memory first (the operator-visible effect), then the row under `persistMu`. A failed row delete is logged; the in-memory @@ -290,6 +335,14 @@ Persistence never fails a request. `admit()` does not see store errors. valid; a negative value is rejected. - Existing admit/release/header tests stay in-memory. A recording `Store` asserts `Admit` does not call `SaveCounters`. +- Per-child: two children of one template flush as two rows (same + `subject`, different `partition`); after `Start` each child is still + isolated. Reset of the template deletes both rows. A restored + partition is registered with the expiry worker. A shared-rule row is + not applied to a definition that is now `PerChild`, and the reverse. +- OSS without `quota_templates` still persists shared / provider / + model windows. Per-child config continues to abort startup / reject + admin writes; persistence does not change that gate. ### Release E2E @@ -310,7 +363,10 @@ Shared helper in the common environment block: required. Use **hour** windows so the cap outlives the reload wait. Each scenario -creates a `$QA_SUFFIX`-scoped user-path rule and deletes it. +creates a `$QA_SUFFIX`-scoped **shared** user-path rule (`per_child` +unset) and deletes it. The release stack is OSS and has no +`quota_templates` entitlement; a per-child admin write would 403. +Per-child persistence is the unit tests above, not this matrix. - **S205 — Request-window counters survive `--reload` (SQLite).** `max_requests=1` on `$BASE_URL`. First chat succeeds. Reload @@ -333,9 +389,9 @@ plus unit tests; concurrent is not persisted and is not an E2E case. ## 13. Docs and comments - `docs/features/rate-limits.mdx` and `docs/advanced/cli.mdx` — - windows survive restart and `--reload`; concurrency does not; still - per instance. Drop the “counters start fresh” wording on the CLI - reload page. + windows survive restart and `--reload`, including each active + per-child partition; concurrency does not; still per instance. Drop + the “counters start fresh” wording on the CLI reload page. - `docs/dev/2026-07-05_rate-limiting-spec.md` §8 and §11.7 — mark persistence done; §11.2 stays future work. - `CLAUDE.md` / `Agents.md` — drop rate-limit counters from the @@ -347,7 +403,8 @@ plus unit tests; concurrent is not persisted and is not an E2E case. ## 14. Follow-up (not this change) Redis-Lua `counterBackend`: every `Admit` is an atomic script on shared -keys. No snapshots. Chosen when `REDIS_URL` is set, or behind an -explicit config once someone is running HA and wants N replicas to -share one limit. The interface in §5 is the only preparation this -change makes for that work. +keys (the key must include `partition`, same as `ruleKey`). No +snapshots. Chosen when `REDIS_URL` is set, or behind an explicit config +once someone is running HA and wants N replicas to share one limit. +The interface in §5 is the only preparation this change makes for that +work. From 68fca872c2c387b9b55fdd28418a5ed1271e8410 Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Sun, 16 Aug 2026 15:46:32 +0200 Subject: [PATCH 05/12] feat(ratelimit): persist request and token windows across restart Snapshot sliding-window counters to the existing store so hour and day limits survive restart and --reload. Admission stays in memory. Per-child partitions are stored separately. Concurrency gauges are not persisted. --- .env.template | 10 +- CLAUDE.md | 4 +- config/config.example.yaml | 1 + config/config.go | 3 +- config/ratelimit.go | 8 + config/ratelimit_test.go | 26 ++ docs/advanced/cli.mdx | 8 +- docs/dev/2026-07-05_rate-limiting-spec.md | 12 +- ...-16_rate-limit-counter-persistence-plan.md | 60 +++++ docs/features/rate-limits.mdx | 11 +- internal/admin/handler_ratelimits_test.go | 11 + internal/app/app.go | 4 + internal/ratelimit/factory.go | 6 +- internal/ratelimit/persist.go | 103 +++++++ internal/ratelimit/persist_test.go | 255 ++++++++++++++++++ internal/ratelimit/service.go | 20 +- internal/ratelimit/service_test.go | 28 +- internal/ratelimit/snapshot.go | 134 +++++++++ internal/ratelimit/store.go | 7 +- internal/ratelimit/store_mongodb.go | 110 +++++++- internal/ratelimit/store_sql.go | 91 +++++++ internal/ratelimit/store_sql_test.go | 62 +++++ internal/ratelimit/types.go | 5 +- internal/server/ratelimit_support_test.go | 16 +- tests/e2e/release-e2e-scenarios.md | 153 ++++++++++- 25 files changed, 1116 insertions(+), 32 deletions(-) create mode 100644 docs/dev/2026-08-16_rate-limit-counter-persistence-plan.md create mode 100644 internal/ratelimit/persist.go create mode 100644 internal/ratelimit/persist_test.go create mode 100644 internal/ratelimit/snapshot.go diff --git a/.env.template b/.env.template index 340b66ca6..962e8cc67 100644 --- a/.env.template +++ b/.env.template @@ -377,11 +377,13 @@ # Rate Limits # ============================================================================= # Cap requests, tokens, and in-flight concurrency per user_path subtree. -# Breaches return 429 with Retry-After and x-ratelimit-* headers. Counters are -# in-memory per gateway instance and reset on restart; token limits (tpm/tph/ -# tpd) additionally require USAGE_ENABLED=true. Enabled by default; with no -# configured rules the check is a no-op. +# Breaches return 429 with Retry-After and x-ratelimit-* headers. Request and +# token windows are snapshotted to the configured store (default every 1s) so +# they survive restart and --reload. Concurrency gauges stay in-memory. +# Token limits (tpm/tph/tpd) additionally require USAGE_ENABLED=true. Enabled +# by default; with no configured rules the check is a no-op. # RATE_LIMITS_ENABLED=true +# RATE_LIMITS_FLUSH_INTERVAL=1 # Declare rules per user path with SET_RATE_LIMIT_ (double underscores # separate path segments, like SET_BUDGET_*). Names: rpm/tpm (per minute), diff --git a/CLAUDE.md b/CLAUDE.md index a4a9a6476..369b8a0d9 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -111,7 +111,7 @@ Full reference: `.env.template` and `config/config.yaml` - `GOMODEL_MASTER_KEY` (empty = unsafe mode). Managed API keys (dashboard API Keys page / `POST /admin/auth-keys`) carry a per-key `dashboard_access` flag (default false, changeable via `PUT /admin/auth-keys/{id}/dashboard-access`): only the master key and flagged keys can call the admin REST API endpoints under `/admin/*` (others get 403 `dashboard_access_denied`); the dashboard UI shell and static assets (`/admin/dashboard`, `/admin/static/*`) skip auth entirely — only the admin data they load is gated; model endpoints and `GET /v1/usage` stay open to every key, and the no-master-key lockout-recovery path (auth skipped on `/admin/*`) is unaffected. - `BODY_SIZE_LIMIT` ("10M") - `USER_PATH_HEADER` (`X-GoModel-User-Path`: Header used to read/write request `user_path` values) - - `PID_FILE` / `server.pid_file` (`data/gomodel.pid` next to a `./data` directory, otherwise the OS per-user data dir — same resolution as `SQLITE_PATH`): where the running gateway records its process id. `gomodel --reload` reads it and signals that process (SIGHUP; `kill -HUP` works too) to reload configuration without a restart, like `nginx -s reload`. The reload re-reads `.env` (exported variables still win over the file; variables removed from the file are unset) and the whole config, then rebuilds the application — so every setting reloads, not a curated subset. The replacement is built before the running one is stopped, so a broken config keeps the current one serving; the listening socket is held across generations, so no connection is refused mid-reload. `PORT` and `PID_FILE` changes still need a restart (warned about), and in-memory state — rate limit counters, session affinity pins, live log buffers — resets as it would on restart. `server.pid_file: ""` in `config.yaml` disables the pid file and `--reload` (an empty `PID_FILE` env var reads as unset and keeps the default). Not available on Windows (POSIX signals). + - `PID_FILE` / `server.pid_file` (`data/gomodel.pid` next to a `./data` directory, otherwise the OS per-user data dir — same resolution as `SQLITE_PATH`): where the running gateway records its process id. `gomodel --reload` reads it and signals that process (SIGHUP; `kill -HUP` works too) to reload configuration without a restart, like `nginx -s reload`. The reload re-reads `.env` (exported variables still win over the file; variables removed from the file are unset) and the whole config, then rebuilds the application — so every setting reloads, not a curated subset. The replacement is built before the running one is stopped, so a broken config keeps the current one serving; the listening socket is held across generations, so no connection is refused mid-reload. `PORT` and `PID_FILE` changes still need a restart (warned about), and in-memory state — session affinity pins, live log buffers — resets as it would on restart (request/token rate-limit windows are snapshotted to storage). `server.pid_file: ""` in `config.yaml` disables the pid file and `--reload` (an empty `PID_FILE` env var reads as unset and keeps the default). Not available on Windows (POSIX signals). - `ENABLE_PASSTHROUGH_ROUTES` (true: Enable provider-native passthrough routes under /p/{provider}/...) - `ALLOW_PASSTHROUGH_V1_ALIAS` (true: Allow /p/{provider}/v1/... aliases while keeping /p/{provider}/... canonical) - `ENABLED_PASSTHROUGH_PROVIDERS` (openai,anthropic,openrouter,kilo,zai,sglang,vllm,deepseek: Comma-separated list of enabled passthrough providers) @@ -125,7 +125,7 @@ Full reference: `.env.template` and `config/config.yaml` - **Audit logging:** `LOGGING_ENABLED` (true), `LOGGING_LOG_BODIES` (true), `LOGGING_LOG_AUDIO_BODIES` (false: refines `LOGGING_LOG_BODIES` for audio endpoints — base64 audio for both `/v1/audio/speech` output and `/v1/audio/transcriptions` upload (≤8 MB each, else `too_large`) + dashboard playback, plus transcription upload metadata; no effect unless `LOGGING_LOG_BODIES` is on, in which case audio-off records a placeholder), `LOGGING_LOG_HEADERS` (true), `LOGGING_RETENTION_DAYS` (30) - **Usage tracking:** `USAGE_ENABLED` (true), `ENFORCE_RETURNING_USAGE_DATA` (true), `USAGE_RETENTION_DAYS` (90). Callers can read their own status without admin access via `GET /v1/usage`: usage summary over a date window (`start_date`/`end_date`/`days`, default last 30 days UTC) plus budget and rate-limit statuses, all scoped to the caller's effective user path (managed key binding, else the user-path header). - **Budgets:** `BUDGETS_ENABLED` (true; no-op until budgets exist, and force-disabled with a warning when `USAGE_ENABLED=false` since spend is read from usage cost records). Every budget has a scope: `user_path` (subtree — a budget on `/team` covers `/team/app` but not `/team-alpha`) or `label` (matches a request label verbatim, no case folding; labels come from tagging headers and managed-key labels, and a request carrying several labels is charged against every matching label budget). Limits are an `amount` per period (`hourly`/`daily`/`weekly`/`monthly`, or a custom `period_seconds` — named periods reset on the calendar anchors configured under Settings → Budget Resets, only custom seconds are fixed windows). A breach returns 429 (`code: budget_exceeded`) with `Retry-After`; a budget with no recorded usage never blocks, and response cache hits return before enforcement. Managed in the dashboard (Budgets page: scope selector) / `/admin/budgets` (GET/PUT/DELETE + `POST .../reset-one`, `POST .../reset`; requests take `scope`+`subject`, with `user_path` as shorthand for user-path budgets), or as infrastructure-as-code under `budgets.{user_paths,labels}:` in `config.yaml` / `SET_BUDGET_` env vars (`period=amount` compact syntax or a JSON limit array; `__` separates path segments). Label budgets are YAML/admin-only — labels are matched verbatim and are not env-name safe. Config-sourced budgets are read-only in the dashboard and manual edits win over config seeds. A user-path budget can also carry `per_child: true` (GoModel Pro, `quota_templates` entitlement): the entry becomes a template giving every *direct* child of the subject its own independent budget rather than one shared subtree budget — `/customers/alice` and `/customers/alice/app` both spend from Alice's, `/customers/bob` from Bob's, and the template path itself matches nothing. Resetting the template starts a new period for every child. Without the entitlement `PER_CHILD_QUOTAS_ENABLED` is off, admin writes return 403 `quota_templates_not_entitled`, and `per_child` in `config.yaml`/`SET_BUDGET_*` aborts startup rather than silently degrading to a shared limit. Enforcement evaluates every matching budget in ONE batched store query (`Store.SumSpend`), so a wide match set costs one scan rather than one per budget. -- **Rate limits:** `RATE_LIMITS_ENABLED` (true; no-op until rules exist). Every rule has a scope: `user_path` (consumer control; subtree with ONE shared counter per rule — per-key limits = give each key its own path), `provider` (caps one configured provider instance across all consumers/models), or `model` (subject `openai/gpt-4o` pins one provider's model, bare `gpt-4o` covers it on any provider; matching case-insensitive). Limits: `max_requests`/`max_tokens` per period (`minute`/`hour`/`day`/custom `period_seconds`, sliding window) plus `concurrent` (period_seconds 0: `max_requests` = max in-flight; realtime sessions hold a slot for the session, batch submissions don't — and batch skips provider/model rules since batch files can mix models). Enforcement covers every model endpoint; user-path breaches return 429 (`code: rate_limit_exceeded`) with `Retry-After`, successes carry `x-ratelimit-{limit,remaining,reset}-{requests,tokens}` from the most-constrained matching rule; cache hits bypass. Saturated providers/models are instead routed around: virtual-model load balancing prefers targets with capacity (falling back to the first declared target when all are saturated, so the client gets an honest 429 rather than an unavailable-model error; saturation never affects catalog membership or /v1/models listing), a saturated primary route with configured failover rules skips the primary provider and is served by the sweep (which also skips saturated candidates), and only requests with no viable alternative get 429. Token windows are charged to the provider/model that actually executed (from the usage entry), so accounting stays correct under aliasing/failover. Managed in the dashboard (Rate Limits page: scope selector) / `/admin/rate-limits` (GET/PUT/DELETE + `POST .../reset-one`, `POST .../reset`; requests take `scope`+`subject`, with `user_path` as shorthand for user-path rules), or as infrastructure-as-code under `rate_limits.{user_paths,providers,models}:` in `config.yaml` / `SET_RATE_LIMIT_` env vars (`rpm/tpm/rph/tph/rpd/tpd/concurrent=N` compact syntax or a JSON rule array; `__` separates path segments) and `SET_PROVIDER_RATE_LIMIT_` (same syntax; suffix underscores become hyphens; model rules are YAML/admin-only). Env replaces the whole YAML entry for the same subject; config-sourced rules are read-only in the dashboard and manual edits win over config seeds, like budgets. A user-path rule can also carry `per_child: true` (GoModel Pro, `quota_templates` entitlement), turning the one shared subtree counter into a template with independent counters per *direct* child (descendants share their direct child's); provider and model rules cannot be per-child. Resetting or deleting the template clears every child's request and token window. The same entitlement gate as budgets applies: 403 `quota_templates_not_entitled` on admin writes, startup abort on config. Token limits are post-accounted from usage entries, so they require `USAGE_ENABLED=true` (startup warns otherwise) and one request can overshoot a token window. Counters are in-memory per instance (N replicas ≈ N× limit) and reset on restart — budgets remain the durable cross-instance control. +- **Rate limits:** `RATE_LIMITS_ENABLED` (true; no-op until rules exist). Every rule has a scope: `user_path` (consumer control; subtree with ONE shared counter per rule — per-key limits = give each key its own path), `provider` (caps one configured provider instance across all consumers/models), or `model` (subject `openai/gpt-4o` pins one provider's model, bare `gpt-4o` covers it on any provider; matching case-insensitive). Limits: `max_requests`/`max_tokens` per period (`minute`/`hour`/`day`/custom `period_seconds`, sliding window) plus `concurrent` (period_seconds 0: `max_requests` = max in-flight; realtime sessions hold a slot for the session, batch submissions don't — and batch skips provider/model rules since batch files can mix models). Enforcement covers every model endpoint; user-path breaches return 429 (`code: rate_limit_exceeded`) with `Retry-After`, successes carry `x-ratelimit-{limit,remaining,reset}-{requests,tokens}` from the most-constrained matching rule; cache hits bypass. Saturated providers/models are instead routed around: virtual-model load balancing prefers targets with capacity (falling back to the first declared target when all are saturated, so the client gets an honest 429 rather than an unavailable-model error; saturation never affects catalog membership or /v1/models listing), a saturated primary route with configured failover rules skips the primary provider and is served by the sweep (which also skips saturated candidates), and only requests with no viable alternative get 429. Token windows are charged to the provider/model that actually executed (from the usage entry), so accounting stays correct under aliasing/failover. Managed in the dashboard (Rate Limits page: scope selector) / `/admin/rate-limits` (GET/PUT/DELETE + `POST .../reset-one`, `POST .../reset`; requests take `scope`+`subject`, with `user_path` as shorthand for user-path rules), or as infrastructure-as-code under `rate_limits.{user_paths,providers,models}:` in `config.yaml` / `SET_RATE_LIMIT_` env vars (`rpm/tpm/rph/tph/rpd/tpd/concurrent=N` compact syntax or a JSON rule array; `__` separates path segments) and `SET_PROVIDER_RATE_LIMIT_` (same syntax; suffix underscores become hyphens; model rules are YAML/admin-only). Env replaces the whole YAML entry for the same subject; config-sourced rules are read-only in the dashboard and manual edits win over config seeds, like budgets. A user-path rule can also carry `per_child: true` (GoModel Pro, `quota_templates` entitlement), turning the one shared subtree counter into a template with independent counters per *direct* child (descendants share their direct child's); provider and model rules cannot be per-child. Resetting or deleting the template clears every child's request and token window. The same entitlement gate as budgets applies: 403 `quota_templates_not_entitled` on admin writes, startup abort on config. Token limits are post-accounted from usage entries, so they require `USAGE_ENABLED=true` (startup warns otherwise) and one request can overshoot a token window. Request/token windows are snapshotted to the store (`RATE_LIMITS_FLUSH_INTERVAL`, default 1s) so they survive restart and `--reload`; concurrency gauges stay in-memory. N replicas ≈ N× the configured limit — budgets remain the durable cross-instance control. - **Dashboard live logs:** - `DASHBOARD_LIVE_LOGS_ENABLED` (true): keep enabled for low-latency dashboard previews; set false only when live streams are not needed or memory/socket usage must be minimized. With `LOGGING_LOG_BODIES` also enabled, in-flight streamed responses render chunk-by-chunk in the request log and Interactions drawer (throttled `audit.stream` events, published only while a dashboard is connected; partial bodies are never buffered server-side). - `DASHBOARD_LIVE_LOGS_BUFFER_SIZE` (10000): effective size is capped at `DASHBOARD_LIVE_LOGS_REPLAY_LIMIT + 1` (older events can never be replayed); lower it below the replay limit only to shrink memory at the cost of more replay resets. Buffered events are compact previews — request/response bodies are never retained in the buffer (connected dashboards get them live; history hydrates from persisted audit entries). diff --git a/config/config.example.yaml b/config/config.example.yaml index c055d01af..785bee4b1 100644 --- a/config/config.example.yaml +++ b/config/config.example.yaml @@ -198,6 +198,7 @@ budgets: rate_limits: enabled: true # env: RATE_LIMITS_ENABLED; with no configured rules this has no effect + flush_interval: 1 # seconds; env: RATE_LIMITS_FLUSH_INTERVAL; 0 = no periodic snapshot (still load + shutdown flush) user_paths: # Env equivalent: # SET_RATE_LIMIT_USER__PATH__EXAMPLE="rpm=100,tpm=50000,rpd=10000,concurrent=10" diff --git a/config/config.go b/config/config.go index b4d0b6e45..19d03ad9e 100644 --- a/config/config.go +++ b/config/config.go @@ -154,7 +154,8 @@ func buildDefaultConfig() *Config { Enabled: true, }, RateLimits: RateLimitsConfig{ - Enabled: true, + Enabled: true, + FlushInterval: 1, }, Metrics: MetricsConfig{ Endpoint: "/metrics", diff --git a/config/ratelimit.go b/config/ratelimit.go index 38c0fd160..363b49a1a 100644 --- a/config/ratelimit.go +++ b/config/ratelimit.go @@ -15,6 +15,11 @@ type RateLimitsConfig struct { // Default: true. With no rules configured the check is a no-op. Enabled bool `yaml:"enabled" env:"RATE_LIMITS_ENABLED"` + // FlushInterval is how often live request/token windows are written to + // storage, in seconds. Default 1. 0 disables the periodic loop; Start + // still loads and Close of an active generation still writes once. + FlushInterval int `yaml:"flush_interval" env:"RATE_LIMITS_FLUSH_INTERVAL"` + // UserPaths declares rate limit rules by tracked user path. UserPaths []RateLimitUserPathConfig `yaml:"user_paths"` @@ -230,6 +235,9 @@ func validateRateLimitConfig(cfg *RateLimitsConfig) error { if cfg == nil { return nil } + if cfg.FlushInterval < 0 { + return fmt.Errorf("rate_limits.flush_interval must be >= 0") + } if !cfg.Enabled { return nil } diff --git a/config/ratelimit_test.go b/config/ratelimit_test.go index c2f37ddea..927cfd5f4 100644 --- a/config/ratelimit_test.go +++ b/config/ratelimit_test.go @@ -457,6 +457,9 @@ func TestRateLimitsEnabledByDefaultAndTogglable(t *testing.T) { if !result.Config.RateLimits.Enabled { t.Fatal("rate limits should be enabled by default") } + if result.Config.RateLimits.FlushInterval != 1 { + t.Fatalf("FlushInterval = %d, want 1", result.Config.RateLimits.FlushInterval) + } }) withTempDir(t, func(string) { @@ -471,6 +474,29 @@ func TestRateLimitsEnabledByDefaultAndTogglable(t *testing.T) { }) } +func TestRateLimitsFlushIntervalEnv(t *testing.T) { + clearAllConfigEnvVars(t) + + withTempDir(t, func(string) { + t.Setenv("RATE_LIMITS_FLUSH_INTERVAL", "0") + result, err := Load() + if err != nil { + t.Fatalf("Load() failed: %v", err) + } + if result.Config.RateLimits.FlushInterval != 0 { + t.Fatalf("FlushInterval = %d, want 0", result.Config.RateLimits.FlushInterval) + } + }) + + clearAllConfigEnvVars(t) + withTempDir(t, func(string) { + t.Setenv("RATE_LIMITS_FLUSH_INTERVAL", "-1") + if _, err := Load(); err == nil || !strings.Contains(err.Error(), "flush_interval") { + t.Fatalf("Load() error = %v, want flush_interval", err) + } + }) +} + func TestParseRateLimitEnvLimits_RejectsUnknownField(t *testing.T) { _, err := parseRateLimitEnvLimits(`[{"period":"minute","max_requsts":100}]`, true) if err == nil { diff --git a/docs/advanced/cli.mdx b/docs/advanced/cli.mdx index e3439cd67..9afc71aec 100644 --- a/docs/advanced/cli.mdx +++ b/docs/advanced/cli.mdx @@ -173,9 +173,11 @@ docker exec my-gateway /gomodel --reload needs a restart. The gateway logs a warning naming both ports. - **`PID_FILE`** — it names the process that is already running. - **`GOMODEL_DEMO_MODE`** — the demo warnings are wired up once at startup. -- **In-memory state** — rate limit counters, virtual-model session affinity, and - live log buffers start fresh, exactly as they would after a restart. Budgets - and usage are stored in the database and are unaffected. +- **In-memory state** — virtual-model session affinity and live log buffers + start fresh, exactly as they would after a restart. Request and token rate + limit windows are written on shutdown and restored by the next generation. + Concurrency gauges and budgets/usage behave as they do on restart (gauges + start empty; budgets and usage live in the database). For refreshing provider model catalogs and admin-managed data *without* re-reading configuration, the dashboard's runtime refresh (`POST /admin/runtime/refresh`) is diff --git a/docs/dev/2026-07-05_rate-limiting-spec.md b/docs/dev/2026-07-05_rate-limiting-spec.md index ac0daa512..539fdee11 100644 --- a/docs/dev/2026-07-05_rate-limiting-spec.md +++ b/docs/dev/2026-07-05_rate-limiting-spec.md @@ -261,11 +261,10 @@ three-backend layout (`store_sqlite.go`, `store_postgresql.go`, byte-for-byte the budget module pattern. No settings table: windows are epoch-aligned UTC; rate limits do not need budget-style calendar anchors. -**Counters are ephemeral.** A restart starts fresh windows. For minute windows -this is invisible; a `day` window can under-count after a restart. That -tradeoff is documented — durable long-horizon control is what budgets (DB-sum -based) are for. Bifrost makes the same tradeoff in its default mode (memory -counters, 10 s DB flush). +**Request and token windows are snapshotted** to the rule store (see +`docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md`). Admission stays +in memory. Concurrency gauges stay ephemeral. Multi-replica sharing is still +§11.2. ## 9. Multi-instance stance @@ -335,7 +334,8 @@ backends. instead of hard 429s (top LiteLLM enterprise upsell). 5. Upstream `x-ratelimit-*` passthrough when GoModel itself imposes no limit. 6. Workflow `features.rate_limit` gating, if a use case appears. -7. Counter persistence across restarts (periodic flush) for day windows. +7. ~~Counter persistence across restarts (periodic flush) for day windows.~~ + Done: `docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md`. 8. Attempt-level request counting for failover targets (today only the resolved primary route is charged a request; tokens are always correct). 9. Virtual-model (alias) subjects for model rules -- needs the requested diff --git a/docs/dev/2026-08-16_rate-limit-counter-persistence-plan.md b/docs/dev/2026-08-16_rate-limit-counter-persistence-plan.md new file mode 100644 index 000000000..eeaae8c7d --- /dev/null +++ b/docs/dev/2026-08-16_rate-limit-counter-persistence-plan.md @@ -0,0 +1,60 @@ +# Rate Limit Counter Persistence Implementation Plan + +> **For agentic workers:** Execute task-by-task. Spec: +> `docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md`. + +**Goal:** Persist request/token sliding windows across restart and `--reload` +using the existing SQL/Mongo store, without putting I/O on `admit()`. + +**Architecture:** Keep the in-memory limiter as the live store. Snapshot +windows (including per-child `partition` keys) to `rate_limit_counters`. +`New` does not load or flush. `Start` (from `App.startServer`) loads then +optionally ticks. `Close` of an **active** generation writes once. +`persistMu` serializes flush vs reset/delete. Construction never deletes +snapshot rows. + +**Tech Stack:** Go, existing `internal/ratelimit` limiter + SQL/Mongo stores, +`RATE_LIMITS_FLUSH_INTERVAL` (default 1s). + +--- + +## Files + +- Modify: `config/ratelimit.go`, `config/config.go`, `config/ratelimit_test.go` +- Modify: `internal/ratelimit/store.go`, `store_sql.go`, `store_mongodb.go`, + `service.go`, `factory.go`, `service_test.go`, `store_sql_test.go`, + `store_mongodb_test.go` +- Create: `internal/ratelimit/snapshot.go`, `persist.go`, `persist_test.go`, + `snapshot_test.go` +- Modify: `internal/app/app.go` +- Modify: docs, `.env.template`, `config/config.example.yaml`, + `tests/e2e/release-e2e-scenarios.md` + +### Task 1: Config + +`FlushInterval int` on `RateLimitsConfig` (`yaml:"flush_interval"` +`env:"RATE_LIMITS_FLUSH_INTERVAL"`). Default `1` in `config.Load`. Reject +`< 0`. `0` is valid (no periodic loop). + +### Task 2: Store + +Add `LoadCounters`, `SaveCounters` (replace-all), `DeleteCounter` (all +partitions of a definition), `DeleteAllCounters`. SQL table and Mongo +collection `rate_limit_counters`, PK +`(scope, subject, partition, period_seconds)`. Mongo uses +transaction-plus-standalone-fallback like `ReplaceConfigRules`. + +### Task 3: Snapshot + persist + +`windowSnapshot` with `Partition`. Limiter `snapshot`/`restore` copy +`windowCounter` by value. Restore skips mode-mismatch and expired child +windows (`windowStart + 2*period`); re-arms `trackCounterExpiry`. +`Service.Start` / flush loop / `Close` compose with existing expiry +`Close`. `Reset*`/`DeleteRule` delete rows under `persistMu`. +`Refresh`/`ReplaceConfigRules` do not touch the snapshot store. + +### Task 4: Wire + docs + E2E + +`factory` passes flush interval. `App.startServer` calls `Start` before +listen. Docs: windows survive bounce; concurrency does not; still per +instance. S205–S207 shared hour rules + `reload_release_gateway`. diff --git a/docs/features/rate-limits.mdx b/docs/features/rate-limits.mdx index 2add0576c..32de6f7eb 100644 --- a/docs/features/rate-limits.mdx +++ b/docs/features/rate-limits.mdx @@ -227,9 +227,14 @@ Details worth knowing: - A realtime session counts as one request and holds one concurrency slot for the whole session. A batch submission counts as one request and holds no concurrency slot. -- Counters live in memory, per gateway instance: with N replicas the - effective limit is about N times the configured value, and counters reset - on restart. Use budgets for durable, cross-instance control. +- Request and token windows live in memory for admission and are snapshotted + to the configured store about once a second (and on shutdown), so they + survive a restart and `gomodel --reload`. Concurrency gauges are not + persisted: after a bounce nothing is in flight. With N replicas the + effective limit is still about N times the configured value — use budgets + for durable, cross-instance control. Tune the snapshot interval with + `RATE_LIMITS_FLUSH_INTERVAL` (seconds; default `1`; `0` skips the periodic + write but still loads on start and flushes on a clean shutdown). ## Client behavior diff --git a/internal/admin/handler_ratelimits_test.go b/internal/admin/handler_ratelimits_test.go index 2f2d125cb..ffd5960ed 100644 --- a/internal/admin/handler_ratelimits_test.go +++ b/internal/admin/handler_ratelimits_test.go @@ -58,6 +58,17 @@ func (s *adminRateLimitStore) ReplaceConfigRules(ctx context.Context, rules []ra return s.UpsertRules(ctx, rules) } +func (s *adminRateLimitStore) LoadCounters(context.Context) ([]ratelimit.WindowSnapshot, error) { + return nil, nil +} +func (s *adminRateLimitStore) SaveCounters(context.Context, []ratelimit.WindowSnapshot) error { + return nil +} +func (s *adminRateLimitStore) DeleteCounter(context.Context, ratelimit.RuleScope, string, int64) error { + return nil +} +func (s *adminRateLimitStore) DeleteAllCounters(context.Context) error { return nil } + func (s *adminRateLimitStore) Close() error { return nil } func newRateLimitHandler(t *testing.T, store *adminRateLimitStore) (*Handler, *ratelimit.Service) { diff --git a/internal/app/app.go b/internal/app/app.go index 0b941819e..3a5193919 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -991,6 +991,10 @@ func (a *App) startServer(ctx context.Context, address string, start func(contex a.serverDone = done a.serverMu.Unlock() + if a.rateLimits != nil && a.rateLimits.Service != nil { + a.rateLimits.Service.Start(ctx) + } + slog.Info("starting server", "address", address) err := start(serverCtx) diff --git a/internal/ratelimit/factory.go b/internal/ratelimit/factory.go index 82eeb2bfe..12d886806 100644 --- a/internal/ratelimit/factory.go +++ b/internal/ratelimit/factory.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "sync" + "time" "go.mongodb.org/mongo-driver/v2/mongo" @@ -60,7 +61,10 @@ func newResult(ctx context.Context, cfg *config.Config, storeConn storage.Storag if err != nil { return nil, err } - service, err := NewService(ctx, store, WithQuotaTemplates(quotaTemplatesEnabled)) + service, err := NewService(ctx, store, + WithQuotaTemplates(quotaTemplatesEnabled), + WithFlushInterval(time.Duration(cfg.RateLimits.FlushInterval)*time.Second), + ) if err != nil { return nil, err } diff --git a/internal/ratelimit/persist.go b/internal/ratelimit/persist.go new file mode 100644 index 000000000..0cf331e13 --- /dev/null +++ b/internal/ratelimit/persist.go @@ -0,0 +1,103 @@ +package ratelimit + +import ( + "context" + "log/slog" + "time" +) + +// WithFlushInterval sets how often an active generation writes window +// snapshots. Zero disables the periodic loop; Start still loads and Close +// of an active generation still writes once. +func WithFlushInterval(interval time.Duration) ServiceOption { + return func(service *Service) { + if interval < 0 { + interval = 0 + } + service.flushInterval = interval + } +} + +func (s *Service) Start(ctx context.Context) { + if s == nil || s.store == nil { + return + } + s.loadCounters(ctx) + s.startFlushLoop() + s.active.Store(true) +} + +func (s *Service) loadCounters(ctx context.Context) { + snapshots, err := s.store.LoadCounters(ctx) + if err != nil { + slog.Warn("rate limit counters: load failed; starting empty", "error", err) + return + } + s.limiter.restore(snapshots, s.Rules(), time.Now().UTC()) +} + +func (s *Service) startFlushLoop() { + if s.flushInterval <= 0 { + return + } + s.flushStop = make(chan struct{}) + s.flushDone = make(chan struct{}) + go func() { + defer close(s.flushDone) + ticker := time.NewTicker(s.flushInterval) + defer ticker.Stop() + for { + select { + case <-ticker.C: + s.flush(context.Background()) + case <-s.flushStop: + return + } + } + }() +} + +func (s *Service) flush(ctx context.Context) { + if s == nil || s.store == nil { + return + } + s.persistMu.Lock() + defer s.persistMu.Unlock() + if err := s.store.SaveCounters(ctx, s.limiter.snapshot(s.Rules())); err != nil { + slog.Warn("rate limit counters: flush failed", "error", err) + } +} + +func (s *Service) persistDelete(scope RuleScope, subject string, periodSeconds int64) { + if s == nil || s.store == nil { + return + } + s.persistMu.Lock() + defer s.persistMu.Unlock() + if err := s.store.DeleteCounter(context.Background(), scope, subject, periodSeconds); err != nil { + slog.Error("rate limit counters: delete failed", "scope", scope, "subject", subject, "period_seconds", periodSeconds, "error", err) + } +} + +func (s *Service) persistDeleteAll() { + if s == nil || s.store == nil { + return + } + s.persistMu.Lock() + defer s.persistMu.Unlock() + if err := s.store.DeleteAllCounters(context.Background()); err != nil { + slog.Error("rate limit counters: delete-all failed", "error", err) + } +} + +func (s *Service) stopFlushAndSave() { + s.flushOnce.Do(func() { + if s.flushStop != nil { + close(s.flushStop) + <-s.flushDone + } + if s.active.Load() { + s.flush(context.Background()) + } + }) +} diff --git a/internal/ratelimit/persist_test.go b/internal/ratelimit/persist_test.go new file mode 100644 index 000000000..e9e98f1c7 --- /dev/null +++ b/internal/ratelimit/persist_test.go @@ -0,0 +1,255 @@ +package ratelimit + +import ( + "context" + "sync/atomic" + "testing" + "time" +) + +func TestSnapshotRoundTripPreservesEstimate(t *testing.T) { + now := time.Unix(1_700_000_000, 0).UTC() + period := PeriodHourSeconds + rule := Rule{ + Scope: ScopeUserPath, Subject: "/team", PeriodSeconds: period, + MaxRequests: new(int64(100)), MaxTokens: new(int64(500)), + } + + src := newLimiter() + if _, _, err := src.admit([]Rule{rule}, now); err != nil { + t.Fatalf("admit: %v", err) + } + src.recordTokens([]Rule{rule}, 40, now) + wantReq := src.status(rule, now).RequestsUsed + wantTok := src.status(rule, now).TokensUsed + + snaps := src.snapshot([]Rule{rule}) + if len(snaps) != 1 { + t.Fatalf("snapshots = %d, want 1", len(snaps)) + } + if snaps[0].Partition != "" { + t.Fatalf("partition = %q, want empty", snaps[0].Partition) + } + + dst := newLimiter() + dst.restore(snaps, []Rule{rule}, now) + if got := dst.status(rule, now).RequestsUsed; got != wantReq { + t.Fatalf("requests used = %d, want %d", got, wantReq) + } + if got := dst.status(rule, now).TokensUsed; got != wantTok { + t.Fatalf("tokens used = %d, want %d", got, wantTok) + } +} + +func TestSnapshotSkipsConcurrentAndExpiredChild(t *testing.T) { + now := time.Unix(1_700_000_000, 0).UTC() + shared := Rule{ + Scope: ScopeUserPath, Subject: "/team", PeriodSeconds: PeriodConcurrent, + MaxRequests: new(int64(3)), + } + template := Rule{ + Scope: ScopeUserPath, Subject: "/customers", PerChild: true, + PeriodSeconds: PeriodHourSeconds, MaxRequests: new(int64(10)), + } + child, ok := template.resolve(Subjects{UserPath: "/customers/alice"}) + if !ok { + t.Fatal("resolve child") + } + + src := newLimiter() + if _, _, err := src.admit([]Rule{shared, child}, now); err != nil { + t.Fatalf("admit: %v", err) + } + // Force the child window into the distant past so restore drops it. + src.mu.Lock() + for _, counter := range src.requests { + if counter != nil { + counter.windowStart = now.Unix() - 10*PeriodHourSeconds + } + } + src.mu.Unlock() + + snaps := src.snapshot([]Rule{shared, template}) + for _, snap := range snaps { + if snap.PeriodSeconds == PeriodConcurrent { + t.Fatal("concurrent snapshot written") + } + } + + dst := newLimiter() + dst.restore(snaps, []Rule{template}, now) + if got := dst.status(child, now).RequestsUsed; got != 0 { + t.Fatalf("expired child restored used = %d, want 0", got) + } +} + +func TestSnapshotIsolatesPerChildPartitions(t *testing.T) { + now := time.Unix(1_700_000_000, 0).UTC() + template := Rule{ + Scope: ScopeUserPath, Subject: "/customers", PerChild: true, + PeriodSeconds: PeriodHourSeconds, MaxRequests: new(int64(1)), + } + alice, _ := template.resolve(Subjects{UserPath: "/customers/alice/app"}) + bob, _ := template.resolve(Subjects{UserPath: "/customers/bob"}) + + src := newLimiter() + if _, _, err := src.admit([]Rule{alice}, now); err != nil { + t.Fatalf("alice admit: %v", err) + } + if _, _, err := src.admit([]Rule{bob}, now); err != nil { + t.Fatalf("bob admit: %v", err) + } + + snaps := src.snapshot([]Rule{template}) + if len(snaps) != 2 { + t.Fatalf("snapshots = %d, want 2", len(snaps)) + } + + dst := newLimiter() + dst.restore(snaps, []Rule{template}, now) + if _, _, err := dst.admit([]Rule{alice}, now); err == nil { + t.Fatal("alice should be exhausted after restore") + } + if _, _, err := dst.admit([]Rule{bob}, now); err == nil { + t.Fatal("bob should be exhausted after restore") + } +} + +func TestStartLoadsAndCloseFlushes(t *testing.T) { + now := time.Now().UTC() + rule := Rule{ + Scope: ScopeUserPath, Subject: "/team", PeriodSeconds: PeriodHourSeconds, + MaxRequests: new(int64(1)), Source: SourceManual, + } + store := &memStore{} + if err := store.UpsertRules(context.Background(), []Rule{rule}); err != nil { + t.Fatalf("seed: %v", err) + } + + first, err := NewService(context.Background(), store) + if err != nil { + t.Fatalf("NewService: %v", err) + } + if _, err := first.Acquire(onPath("/team"), now); err != nil { + t.Fatalf("Acquire: %v", err) + } + if len(store.counters) != 0 { + t.Fatal("New without Start wrote counters") + } + first.Start(context.Background()) + first.Close() + if len(store.counters) != 1 { + t.Fatalf("counters after Close = %d, want 1", len(store.counters)) + } + + second, err := NewService(context.Background(), store) + if err != nil { + t.Fatalf("second NewService: %v", err) + } + t.Cleanup(second.Close) + second.Start(context.Background()) + if _, err := second.Acquire(onPath("/team"), now); err == nil { + t.Fatal("restored window admitted a second request") + } +} + +func TestCloseWithoutStartDoesNotWrite(t *testing.T) { + store := &recordingStore{memStore: memStore{}} + if err := store.UpsertRules(context.Background(), []Rule{{ + Subject: "/", PeriodSeconds: PeriodHourSeconds, MaxRequests: new(int64(5)), Source: SourceManual, + }}); err != nil { + t.Fatalf("seed: %v", err) + } + service, err := NewService(context.Background(), store) + if err != nil { + t.Fatalf("NewService: %v", err) + } + if _, err := service.Acquire(onPath("/"), time.Now().UTC()); err != nil { + t.Fatalf("Acquire: %v", err) + } + service.Close() + if store.saves.Load() != 0 { + t.Fatalf("saves = %d, want 0", store.saves.Load()) + } +} + +func TestResetClearsPersistedWindow(t *testing.T) { + now := time.Now().UTC() + rule := Rule{ + Scope: ScopeUserPath, Subject: "/team", PeriodSeconds: PeriodHourSeconds, + MaxRequests: new(int64(1)), Source: SourceManual, + } + store := &memStore{} + if err := store.UpsertRules(context.Background(), []Rule{rule}); err != nil { + t.Fatalf("seed: %v", err) + } + service, err := NewService(context.Background(), store) + if err != nil { + t.Fatalf("NewService: %v", err) + } + service.Start(context.Background()) + if _, err := service.Acquire(onPath("/team"), now); err != nil { + t.Fatalf("Acquire: %v", err) + } + if err := service.ResetRule(ScopeUserPath, "/team", PeriodHourSeconds); err != nil { + t.Fatalf("ResetRule: %v", err) + } + service.Close() + + next, err := NewService(context.Background(), store) + if err != nil { + t.Fatalf("second NewService: %v", err) + } + t.Cleanup(next.Close) + next.Start(context.Background()) + if _, err := next.Acquire(onPath("/team"), now); err != nil { + t.Fatalf("Acquire after reset restore: %v", err) + } +} + +func TestAdmitDoesNotSave(t *testing.T) { + store := &recordingStore{memStore: memStore{}} + if err := store.UpsertRules(context.Background(), []Rule{{ + Subject: "/", PeriodSeconds: PeriodHourSeconds, MaxRequests: new(int64(5)), Source: SourceManual, + }}); err != nil { + t.Fatalf("seed: %v", err) + } + service, err := NewService(context.Background(), store) + if err != nil { + t.Fatalf("NewService: %v", err) + } + t.Cleanup(service.Close) + if _, err := service.Acquire(onPath("/"), time.Now().UTC()); err != nil { + t.Fatalf("Acquire: %v", err) + } + if store.saves.Load() != 0 { + t.Fatalf("Admit saved %d times", store.saves.Load()) + } +} + +func TestRestoreIgnoresSharedRowOnPerChildRule(t *testing.T) { + now := time.Unix(1_700_000_000, 0).UTC() + template := Rule{ + Scope: ScopeUserPath, Subject: "/customers", PerChild: true, + PeriodSeconds: PeriodHourSeconds, MaxRequests: new(int64(1)), + } + dst := newLimiter() + dst.restore([]WindowSnapshot{{ + Scope: string(ScopeUserPath), Subject: "/customers", Partition: "", + PeriodSeconds: PeriodHourSeconds, RequestsWindowStart: now.Unix(), RequestsCurrent: 1, + }}, []Rule{template}, now) + child, _ := template.resolve(Subjects{UserPath: "/customers/alice"}) + if _, _, err := dst.admit([]Rule{child}, now); err != nil { + t.Fatalf("shared row applied to per-child rule: %v", err) + } +} + +type recordingStore struct { + memStore + saves atomic.Int64 +} + +func (s *recordingStore) SaveCounters(ctx context.Context, snapshots []WindowSnapshot) error { + s.saves.Add(1) + return s.memStore.SaveCounters(ctx, snapshots) +} diff --git a/internal/ratelimit/service.go b/internal/ratelimit/service.go index 11bf24a54..1b09ec415 100644 --- a/internal/ratelimit/service.go +++ b/internal/ratelimit/service.go @@ -7,6 +7,7 @@ import ( "sort" "strings" "sync" + "sync/atomic" "time" ) @@ -35,6 +36,13 @@ type Service struct { rules []Rule quotaTemplates bool + + flushInterval time.Duration + persistMu sync.Mutex + flushStop chan struct{} + flushDone chan struct{} + flushOnce sync.Once + active atomic.Bool } func NewService(ctx context.Context, store Store, options ...ServiceOption) (*Service, error) { @@ -57,9 +65,14 @@ func NewService(ctx context.Context, store Store, options ...ServiceOption) (*Se return service, nil } -// Close stops the in-memory expiry cleanup worker. +// Close stops the flush loop, writes a final snapshot if this generation +// was started, and stops the in-memory expiry cleanup worker. func (s *Service) Close() { - if s != nil && s.limiter != nil { + if s == nil { + return + } + s.stopFlushAndSave() + if s.limiter != nil { s.limiter.close() } } @@ -138,6 +151,7 @@ func (s *Service) DeleteRule(ctx context.Context, scope RuleScope, subject strin return err } s.limiter.reset(ruleKey{scope: scope, subject: subject, periodSeconds: periodSeconds}) + s.persistDelete(scope, subject, periodSeconds) return s.Refresh(ctx) } @@ -336,6 +350,7 @@ func (s *Service) ResetRule(scope RuleScope, subject string, periodSeconds int64 return err } s.limiter.reset(ruleKey{scope: scope, subject: subject, periodSeconds: periodSeconds}) + s.persistDelete(scope, subject, periodSeconds) return nil } @@ -345,6 +360,7 @@ func (s *Service) ResetAll() error { return ErrUnavailable } s.limiter.resetAll() + s.persistDeleteAll() return nil } diff --git a/internal/ratelimit/service_test.go b/internal/ratelimit/service_test.go index 7c95528ec..a3591bc24 100644 --- a/internal/ratelimit/service_test.go +++ b/internal/ratelimit/service_test.go @@ -13,7 +13,8 @@ import ( // memStore is a minimal in-memory Store for service tests. type memStore struct { - rules []Rule + rules []Rule + counters []WindowSnapshot } func (m *memStore) ListRules(context.Context) ([]Rule, error) { @@ -68,6 +69,31 @@ func (m *memStore) ReplaceConfigRules(ctx context.Context, rules []Rule) error { func (m *memStore) Close() error { return nil } +func (m *memStore) LoadCounters(context.Context) ([]WindowSnapshot, error) { + return append([]WindowSnapshot(nil), m.counters...), nil +} + +func (m *memStore) SaveCounters(_ context.Context, snapshots []WindowSnapshot) error { + m.counters = append([]WindowSnapshot(nil), snapshots...) + return nil +} + +func (m *memStore) DeleteCounter(_ context.Context, scope RuleScope, subject string, periodSeconds int64) error { + kept := m.counters[:0] + for _, snap := range m.counters { + if snap.Scope != string(scope) || snap.Subject != subject || snap.PeriodSeconds != periodSeconds { + kept = append(kept, snap) + } + } + m.counters = kept + return nil +} + +func (m *memStore) DeleteAllCounters(context.Context) error { + m.counters = nil + return nil +} + // onPath builds request subjects for user-path-only tests. func onPath(path string) Subjects { return Subjects{UserPath: path} } diff --git a/internal/ratelimit/snapshot.go b/internal/ratelimit/snapshot.go new file mode 100644 index 000000000..869e0999c --- /dev/null +++ b/internal/ratelimit/snapshot.go @@ -0,0 +1,134 @@ +package ratelimit + +import "time" + +// WindowSnapshot is one persisted request/token sliding window. +// Partition is empty for a shared rule and the child path for a per-child +// template (Rule.EffectiveSubject). +type WindowSnapshot struct { + Scope string `bson:"scope"` + Subject string `bson:"subject"` + Partition string `bson:"partition"` + PeriodSeconds int64 `bson:"period_seconds"` + RequestsWindowStart int64 `bson:"requests_window_start"` + RequestsCurrent int64 `bson:"requests_current"` + RequestsPrevious int64 `bson:"requests_previous"` + TokensWindowStart int64 `bson:"tokens_window_start"` + TokensCurrent int64 `bson:"tokens_current"` + TokensPrevious int64 `bson:"tokens_previous"` + UpdatedAt int64 `bson:"updated_at,omitempty"` +} + +func definitionKey(scope RuleScope, subject string, periodSeconds int64) ruleKey { + return ruleKey{scope: scope, subject: subject, periodSeconds: periodSeconds} +} + +func (l *limiter) snapshot(rules []Rule) []WindowSnapshot { + l.mu.Lock() + defer l.mu.Unlock() + + perChild := make(map[ruleKey]bool, len(rules)) + for _, rule := range rules { + if rule.PeriodSeconds <= 0 { + continue + } + perChild[definitionKey(rule.Scope, rule.Subject, rule.PeriodSeconds)] = rule.PerChild + } + + byKey := make(map[ruleKey]*WindowSnapshot) + add := func(key ruleKey, kind counterKind, counter windowCounter) { + wantChild, ok := perChild[definitionKey(key.scope, key.subject, key.periodSeconds)] + if !ok { + return + } + if wantChild != (key.partition != "") { + return + } + snap := byKey[key] + if snap == nil { + snap = &WindowSnapshot{ + Scope: string(key.scope), + Subject: key.subject, + Partition: key.partition, + PeriodSeconds: key.periodSeconds, + } + byKey[key] = snap + } + switch kind { + case requestCounter: + snap.RequestsWindowStart = counter.windowStart + snap.RequestsCurrent = counter.current + snap.RequestsPrevious = counter.previous + case tokenCounter: + snap.TokensWindowStart = counter.windowStart + snap.TokensCurrent = counter.current + snap.TokensPrevious = counter.previous + } + } + + for key, counter := range l.requests { + add(key, requestCounter, *counter) + } + for key, counter := range l.tokens { + add(key, tokenCounter, *counter) + } + + out := make([]WindowSnapshot, 0, len(byKey)) + for _, snap := range byKey { + out = append(out, *snap) + } + return out +} + +func (l *limiter) restore(snapshots []WindowSnapshot, rules []Rule, now time.Time) { + l.mu.Lock() + defer l.mu.Unlock() + + byDef := make(map[ruleKey]Rule, len(rules)) + for _, rule := range rules { + byDef[definitionKey(rule.Scope, rule.Subject, rule.PeriodSeconds)] = rule + } + + nowUnix := now.Unix() + for _, snap := range snapshots { + if snap.PeriodSeconds <= 0 { + continue + } + rule, ok := byDef[definitionKey(RuleScope(snap.Scope), snap.Subject, snap.PeriodSeconds)] + if !ok { + continue + } + if rule.PerChild != (snap.Partition != "") { + continue + } + latest := max(snap.RequestsWindowStart, snap.TokensWindowStart) + if latest > 0 && latest+2*snap.PeriodSeconds < nowUnix { + continue + } + + key := ruleKey{ + scope: RuleScope(snap.Scope), + subject: snap.Subject, + partition: snap.Partition, + periodSeconds: snap.PeriodSeconds, + } + if snap.RequestsWindowStart != 0 || snap.RequestsCurrent != 0 || snap.RequestsPrevious != 0 { + counter := &windowCounter{ + windowStart: snap.RequestsWindowStart, + current: snap.RequestsCurrent, + previous: snap.RequestsPrevious, + } + l.requests[key] = counter + l.trackCounterExpiry(requestCounter, key, counter) + } + if snap.TokensWindowStart != 0 || snap.TokensCurrent != 0 || snap.TokensPrevious != 0 { + counter := &windowCounter{ + windowStart: snap.TokensWindowStart, + current: snap.TokensCurrent, + previous: snap.TokensPrevious, + } + l.tokens[key] = counter + l.trackCounterExpiry(tokenCounter, key, counter) + } + } +} diff --git a/internal/ratelimit/store.go b/internal/ratelimit/store.go index c7f56d082..fac424607 100644 --- a/internal/ratelimit/store.go +++ b/internal/ratelimit/store.go @@ -9,13 +9,16 @@ import ( var ErrNotFound = errors.New("rate limit rule not found") -// Store persists rate limit rule definitions. Live counters are in-memory -// and never stored. +// Store persists rate limit rule definitions and optional window snapshots. type Store interface { ListRules(ctx context.Context) ([]Rule, error) UpsertRules(ctx context.Context, rules []Rule) error DeleteRule(ctx context.Context, scope RuleScope, subject string, periodSeconds int64) error ReplaceConfigRules(ctx context.Context, rules []Rule) error + LoadCounters(ctx context.Context) ([]WindowSnapshot, error) + SaveCounters(ctx context.Context, snapshots []WindowSnapshot) error + DeleteCounter(ctx context.Context, scope RuleScope, subject string, periodSeconds int64) error + DeleteAllCounters(ctx context.Context) error Close() error } diff --git a/internal/ratelimit/store_mongodb.go b/internal/ratelimit/store_mongodb.go index 0b4e8e7e4..5667d0a96 100644 --- a/internal/ratelimit/store_mongodb.go +++ b/internal/ratelimit/store_mongodb.go @@ -6,6 +6,7 @@ import ( "fmt" "log/slog" "strings" + "time" "go.mongodb.org/mongo-driver/v2/bson" "go.mongodb.org/mongo-driver/v2/mongo" @@ -13,7 +14,8 @@ import ( ) type MongoDBStore struct { - rules *mongo.Collection + rules *mongo.Collection + counters *mongo.Collection } func NewMongoDBStore(ctx context.Context, database *mongo.Database) (*MongoDBStore, error) { @@ -21,7 +23,8 @@ func NewMongoDBStore(ctx context.Context, database *mongo.Database) (*MongoDBSto return nil, fmt.Errorf("database is required") } store := &MongoDBStore{ - rules: database.Collection("rate_limits"), + rules: database.Collection("rate_limits"), + counters: database.Collection("rate_limit_counters"), } if err := store.migratePreScopeDocuments(ctx); err != nil { return nil, err @@ -33,6 +36,18 @@ func NewMongoDBStore(ctx context.Context, database *mongo.Database) (*MongoDBSto if err != nil { return nil, fmt.Errorf("create rate limit indexes: %w", err) } + _, err = store.counters.Indexes().CreateOne(ctx, mongo.IndexModel{ + Keys: bson.D{ + {Key: "scope", Value: 1}, + {Key: "subject", Value: 1}, + {Key: "partition", Value: 1}, + {Key: "period_seconds", Value: 1}, + }, + Options: options.Index().SetUnique(true), + }) + if err != nil { + return nil, fmt.Errorf("create rate limit counter indexes: %w", err) + } return store, nil } @@ -340,6 +355,97 @@ func (s *MongoDBStore) configRulesWithoutManualCollisions(ctx context.Context, r return filtered, nil } +func (s *MongoDBStore) LoadCounters(ctx context.Context) ([]WindowSnapshot, error) { + cursor, err := s.counters.Find(ctx, bson.D{}) + if err != nil { + return nil, fmt.Errorf("list rate limit counters: %w", err) + } + defer cursor.Close(ctx) + + var snapshots []WindowSnapshot + for cursor.Next(ctx) { + var snap WindowSnapshot + if err := cursor.Decode(&snap); err != nil { + return nil, fmt.Errorf("decode rate limit counter: %w", err) + } + snapshots = append(snapshots, snap) + } + if err := cursor.Err(); err != nil { + return nil, fmt.Errorf("iterate rate limit counters: %w", err) + } + return snapshots, nil +} + +func (s *MongoDBStore) SaveCounters(ctx context.Context, snapshots []WindowSnapshot) error { + write := func(writeCtx context.Context) error { + if _, err := s.counters.DeleteMany(writeCtx, bson.D{}); err != nil { + return fmt.Errorf("clear rate limit counters: %w", err) + } + if len(snapshots) == 0 { + return nil + } + now := time.Now().Unix() + docs := make([]any, 0, len(snapshots)) + for _, snap := range snapshots { + snap.UpdatedAt = now + docs = append(docs, snap) + } + if _, err := s.counters.InsertMany(writeCtx, docs); err != nil { + return fmt.Errorf("insert rate limit counters: %w", err) + } + return nil + } + + session, err := s.counters.Database().Client().StartSession() + if err != nil { + return fmt.Errorf("start rate limit counter transaction: %w", err) + } + defer session.EndSession(ctx) + + _, err = session.WithTransaction(ctx, func(txCtx context.Context) (any, error) { + if err := write(txCtx); err != nil { + if isMongoTransactionCapabilityError(err) { + return nil, &mongoTransactionFallbackError{err: err} + } + return nil, err + } + return nil, nil + }) + if err != nil { + if fallbackErr := mongoTransactionFallbackCause(err); fallbackErr != nil || isMongoTransactionCapabilityError(err) { + if fallbackErr == nil { + fallbackErr = err + } + slog.Warn("MongoDB transactions unavailable for rate limit counters; falling back to non-transactional update", "error", fallbackErr) + if err := write(ctx); err != nil { + return fmt.Errorf("save rate limit counters without transaction: %w", errors.Join(fallbackErr, err)) + } + return nil + } + return fmt.Errorf("save rate limit counters transaction: %w", err) + } + return nil +} + +func (s *MongoDBStore) DeleteCounter(ctx context.Context, scope RuleScope, subject string, periodSeconds int64) error { + _, err := s.counters.DeleteMany(ctx, bson.D{ + {Key: "scope", Value: scope}, + {Key: "subject", Value: subject}, + {Key: "period_seconds", Value: periodSeconds}, + }) + if err != nil { + return fmt.Errorf("delete rate limit counters %s %s/%d: %w", scope, subject, periodSeconds, err) + } + return nil +} + +func (s *MongoDBStore) DeleteAllCounters(ctx context.Context) error { + if _, err := s.counters.DeleteMany(ctx, bson.D{}); err != nil { + return fmt.Errorf("delete all rate limit counters: %w", err) + } + return nil +} + func (s *MongoDBStore) Close() error { return nil } diff --git a/internal/ratelimit/store_sql.go b/internal/ratelimit/store_sql.go index 594424ee7..341683d6c 100644 --- a/internal/ratelimit/store_sql.go +++ b/internal/ratelimit/store_sql.go @@ -25,6 +25,22 @@ const sqlRateLimitsSchema = ` PRIMARY KEY (scope, subject, period_seconds) )` +const sqlRateLimitCountersSchema = ` + CREATE TABLE IF NOT EXISTS rate_limit_counters ( + scope TEXT NOT NULL, + subject TEXT NOT NULL, + partition TEXT NOT NULL DEFAULT '', + period_seconds ` + sqlx.TypeInt64 + ` NOT NULL, + requests_window_start ` + sqlx.TypeInt64 + ` NOT NULL DEFAULT 0, + requests_current ` + sqlx.TypeInt64 + ` NOT NULL DEFAULT 0, + requests_previous ` + sqlx.TypeInt64 + ` NOT NULL DEFAULT 0, + tokens_window_start ` + sqlx.TypeInt64 + ` NOT NULL DEFAULT 0, + tokens_current ` + sqlx.TypeInt64 + ` NOT NULL DEFAULT 0, + tokens_previous ` + sqlx.TypeInt64 + ` NOT NULL DEFAULT 0, + updated_at ` + sqlx.TypeInt64 + ` NOT NULL, + PRIMARY KEY (scope, subject, partition, period_seconds) + )` + // SQLStore stores rate limit rules in a SQL database. type SQLStore struct { db sqlx.DB @@ -64,6 +80,9 @@ func NewSQLStore(ctx context.Context, db sqlx.DB) (*SQLStore, error) { if err := db.Schema(ctx, `CREATE INDEX IF NOT EXISTS idx_rate_limits_subject ON rate_limits(scope, subject)`); err != nil { return nil, fmt.Errorf("failed to create rate limit index: %w", err) } + if err := db.Schema(ctx, sqlRateLimitCountersSchema); err != nil { + return nil, fmt.Errorf("failed to create rate_limit_counters table: %w", err) + } return &SQLStore{db: db}, nil } @@ -152,6 +171,78 @@ func (s *SQLStore) ReplaceConfigRules(ctx context.Context, rules []Rule) error { }) } +func (s *SQLStore) LoadCounters(ctx context.Context) ([]WindowSnapshot, error) { + rows, err := s.db.Query(ctx, ` + SELECT scope, subject, partition, period_seconds, + requests_window_start, requests_current, requests_previous, + tokens_window_start, tokens_current, tokens_previous + FROM rate_limit_counters + `) + if err != nil { + return nil, fmt.Errorf("list rate limit counters: %w", err) + } + defer rows.Close() + + var snapshots []WindowSnapshot + for rows.Next() { + var snap WindowSnapshot + if err := rows.Scan( + &snap.Scope, &snap.Subject, &snap.Partition, &snap.PeriodSeconds, + &snap.RequestsWindowStart, &snap.RequestsCurrent, &snap.RequestsPrevious, + &snap.TokensWindowStart, &snap.TokensCurrent, &snap.TokensPrevious, + ); err != nil { + return nil, fmt.Errorf("scan rate limit counter: %w", err) + } + snapshots = append(snapshots, snap) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("iterate rate limit counters: %w", err) + } + return snapshots, nil +} + +func (s *SQLStore) SaveCounters(ctx context.Context, snapshots []WindowSnapshot) error { + now := time.Now().Unix() + return s.db.InTx(ctx, func(q sqlx.Querier) error { + if _, err := q.Exec(ctx, `DELETE FROM rate_limit_counters`); err != nil { + return fmt.Errorf("clear rate limit counters: %w", err) + } + for _, snap := range snapshots { + if _, err := q.Exec(ctx, ` + INSERT INTO rate_limit_counters ( + scope, subject, partition, period_seconds, + requests_window_start, requests_current, requests_previous, + tokens_window_start, tokens_current, tokens_previous, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + `, snap.Scope, snap.Subject, snap.Partition, snap.PeriodSeconds, + snap.RequestsWindowStart, snap.RequestsCurrent, snap.RequestsPrevious, + snap.TokensWindowStart, snap.TokensCurrent, snap.TokensPrevious, now, + ); err != nil { + return fmt.Errorf("insert rate limit counter: %w", err) + } + } + return nil + }) +} + +func (s *SQLStore) DeleteCounter(ctx context.Context, scope RuleScope, subject string, periodSeconds int64) error { + _, err := s.db.Exec(ctx, ` + DELETE FROM rate_limit_counters + WHERE scope = ? AND subject = ? AND period_seconds = ? + `, scope, subject, periodSeconds) + if err != nil { + return fmt.Errorf("delete rate limit counters %s %s/%d: %w", scope, subject, periodSeconds, err) + } + return nil +} + +func (s *SQLStore) DeleteAllCounters(ctx context.Context) error { + if _, err := s.db.Exec(ctx, `DELETE FROM rate_limit_counters`); err != nil { + return fmt.Errorf("delete all rate limit counters: %w", err) + } + return nil +} + func (s *SQLStore) Close() error { return nil } diff --git a/internal/ratelimit/store_sql_test.go b/internal/ratelimit/store_sql_test.go index a02c2b1e0..129069a04 100644 --- a/internal/ratelimit/store_sql_test.go +++ b/internal/ratelimit/store_sql_test.go @@ -250,3 +250,65 @@ func TestSQLStoreMigratesPreScopeTable(t *testing.T) { } }) } + +func TestSQLStoreCounterRoundTrip(t *testing.T) { + runSQLStoreTest(t, func(t *testing.T, store *SQLStore) { + ctx := context.Background() + first := []WindowSnapshot{ + { + Scope: string(ScopeUserPath), Subject: "/customers", Partition: "/customers/alice", + PeriodSeconds: PeriodHourSeconds, RequestsWindowStart: 1700000000, RequestsCurrent: 3, + }, + { + Scope: string(ScopeUserPath), Subject: "/customers", Partition: "/customers/bob", + PeriodSeconds: PeriodHourSeconds, RequestsWindowStart: 1700000000, RequestsCurrent: 1, + }, + } + if err := store.SaveCounters(ctx, first); err != nil { + t.Fatalf("SaveCounters: %v", err) + } + got, err := store.LoadCounters(ctx) + if err != nil { + t.Fatalf("LoadCounters: %v", err) + } + if len(got) != 2 { + t.Fatalf("loaded = %d, want 2", len(got)) + } + + if err := store.DeleteCounter(ctx, ScopeUserPath, "/customers", PeriodHourSeconds); err != nil { + t.Fatalf("DeleteCounter: %v", err) + } + got, err = store.LoadCounters(ctx) + if err != nil { + t.Fatalf("LoadCounters after delete: %v", err) + } + if len(got) != 0 { + t.Fatalf("after delete = %d, want 0", len(got)) + } + + if err := store.SaveCounters(ctx, first); err != nil { + t.Fatalf("SaveCounters again: %v", err) + } + if err := store.SaveCounters(ctx, first[:1]); err != nil { + t.Fatalf("SaveCounters replace: %v", err) + } + got, err = store.LoadCounters(ctx) + if err != nil { + t.Fatalf("LoadCounters after replace: %v", err) + } + if len(got) != 1 || got[0].Partition != "/customers/alice" { + t.Fatalf("replaced = %+v, want alice only", got) + } + + if err := store.DeleteAllCounters(ctx); err != nil { + t.Fatalf("DeleteAllCounters: %v", err) + } + got, err = store.LoadCounters(ctx) + if err != nil { + t.Fatalf("LoadCounters after delete all: %v", err) + } + if len(got) != 0 { + t.Fatalf("after delete all = %d, want 0", len(got)) + } + }) +} diff --git a/internal/ratelimit/types.go b/internal/ratelimit/types.go index 267d9e19d..95c96ca4c 100644 --- a/internal/ratelimit/types.go +++ b/internal/ratelimit/types.go @@ -1,7 +1,8 @@ // Package ratelimit enforces request, token, and concurrency limits for the // AI gateway. Rules are scoped to a consumer user-path subtree, a provider, or -// a model. Rule definitions are persisted; live counters are in-memory and per -// instance. +// a model. Rule definitions are persisted. Request and token windows are +// snapshotted to the same store so they survive restart and reload; live +// admission stays in memory and per instance. package ratelimit import ( diff --git a/internal/server/ratelimit_support_test.go b/internal/server/ratelimit_support_test.go index 1e4d27949..4678819c5 100644 --- a/internal/server/ratelimit_support_test.go +++ b/internal/server/ratelimit_support_test.go @@ -28,8 +28,20 @@ func (s *staticRuleStore) UpsertRules(context.Context, []ratelimit.Rule) error { func (s *staticRuleStore) DeleteRule(context.Context, ratelimit.RuleScope, string, int64) error { return nil } -func (s *staticRuleStore) ReplaceConfigRules(context.Context, []ratelimit.Rule) error { return nil } -func (s *staticRuleStore) Close() error { return nil } +func (s *staticRuleStore) ReplaceConfigRules(context.Context, []ratelimit.Rule) error { + return nil +} +func (s *staticRuleStore) LoadCounters(context.Context) ([]ratelimit.WindowSnapshot, error) { + return nil, nil +} +func (s *staticRuleStore) SaveCounters(context.Context, []ratelimit.WindowSnapshot) error { + return nil +} +func (s *staticRuleStore) DeleteCounter(context.Context, ratelimit.RuleScope, string, int64) error { + return nil +} +func (s *staticRuleStore) DeleteAllCounters(context.Context) error { return nil } +func (s *staticRuleStore) Close() error { return nil } func newTestRateLimitService(t *testing.T, rules ...ratelimit.Rule) *ratelimit.Service { t.Helper() diff --git a/tests/e2e/release-e2e-scenarios.md b/tests/e2e/release-e2e-scenarios.md index cfff5e0ac..65e23139b 100644 --- a/tests/e2e/release-e2e-scenarios.md +++ b/tests/e2e/release-e2e-scenarios.md @@ -1,6 +1,6 @@ # Release E2E Curl Matrix -This file contains 204 end-to-end curl scenarios for release validation. +This file contains 207 end-to-end curl scenarios for release validation. These scenarios are prepared for execution across these local gateways: - `http://localhost:18080` - SQLite-backed main test gateway @@ -127,6 +127,11 @@ Stateful note: order. `S200` deliberately registers its managed key on the auth-enabled gateway rather than the no-master-key main SQLite gateway — see the note on that scenario for why +- `S205`-`S207` exercise request-window persistence across `SIGHUP` reload on + SQLite, reset-one across reload, and PostgreSQL/MongoDB parity. Each creates + a `$QA_SUFFIX`-scoped **shared** hour rule (not `per_child` — this stack has + no `quota_templates` entitlement) and deletes it. They reload a shared + gateway, which is safe in this sequential runner. - For stateful partial reruns, prefer a contiguous range that includes the prerequisite setup scenarios, or rerun with the same `--qa-suffix` and `--keep-artifacts` @@ -150,6 +155,34 @@ export BASE_URL=http://localhost:18080 export PG_BASE_URL=http://localhost:18081 export MONGO_BASE_URL=http://localhost:18082 export GR_BASE_URL=http://localhost:18083 +export RELEASE_STACK_DIR="${RELEASE_STACK_DIR:-/tmp/gomodel-release-stack}" + +reload_release_gateway() { + local gateway="$1" + local url="$2" + local pid_file="$RELEASE_STACK_DIR/$gateway/server.pid" + local log_file="$RELEASE_STACK_DIR/$gateway/logs/server.log" + local pid before + pid="$(cat "$pid_file")" + before="$(wc -l < "$log_file" | tr -d ' ')" + kill -HUP "$pid" + for _ in $(seq 1 50); do + if tail -n +$((before + 1)) "$log_file" 2>/dev/null | grep -Fq 'configuration reloaded'; then + for __ in $(seq 1 20); do + if curl -fsS "$url/health" >/dev/null 2>&1; then + return 0 + fi + sleep 0.1 + done + echo "error: $gateway did not become healthy after reload" >&2 + return 1 + fi + sleep 0.1 + done + echo "error: $gateway did not log configuration reloaded" >&2 + tail -n 40 "$log_file" >&2 || true + return 1 +} cat > "$QA_RUN_DIR/qa-openai-batch.jsonl" <<'EOF' {"custom_id":"qa-batch-1","method":"POST","url":"/v1/chat/completions","body":{"model":"gpt-4.1-nano","messages":[{"role":"user","content":"Reply with exactly QA_BATCH_FILE_OK"}],"max_tokens":20}} @@ -5230,3 +5263,121 @@ curl -fsS -H "$ADMIN_AUTH_HEADER" -X DELETE "$AUTH_BASE_URL/admin/budgets" \ -H 'Content-Type: application/json' \ -d "{\"scope\":\"label\",\"subject\":\"$QA_LBL\",\"budget_key\":{\"period\":\"daily\"}}" >/dev/null ``` + +## 25. Rate limit counters across reload + +These scenarios cover request-window persistence across `gomodel --reload` +(SIGHUP). Hour windows so the cap outlives the reload wait. Shared +user-path rules only — the OSS release stack has no `quota_templates` +entitlement. + +### S205 Request-window counters survive `--reload` (SQLite) + +Creates a one-request-per-hour rule, burns it, reloads `sqlite-main`, and +verifies the next request is still `429`. + +```bash +RL_PATH="/qa/ratelimit/persist/$QA_SUFFIX" +BODY_FILE="$QA_RUN_DIR/s205.body.json" + +curl -fsS -X PUT "$BASE_URL/admin/rate-limits" \ + -H 'Content-Type: application/json' \ + -d "{\"user_path\":\"$RL_PATH\",\"limit_key\":{\"period\":\"hour\"},\"max_requests\":1}" \ + | jq -e --arg p "$RL_PATH" 'any(.rate_limits[]?; .scope == "user_path" and .user_path == $p and .max_requests == 1)' >/dev/null + +curl -fsS -o "$BODY_FILE" "$BASE_URL/v1/chat/completions" \ + -H 'Content-Type: application/json' \ + -H "X-GoModel-User-Path: $RL_PATH/leaf" \ + -d '{"model":"gpt-4.1-nano","messages":[{"role":"user","content":"Reply with exactly QA_RL_PERSIST_OK"}],"max_tokens":20}' +assert_chat_response_contains "$BODY_FILE" "openai" "QA_RL_PERSIST_OK" + +reload_release_gateway sqlite-main "$BASE_URL" + +curl -sS -o "$BODY_FILE" -w '%{http_code}' "$BASE_URL/v1/chat/completions" \ + -H 'Content-Type: application/json' \ + -H "X-GoModel-User-Path: $RL_PATH/leaf" \ + -d '{"model":"gpt-4.1-nano","messages":[{"role":"user","content":"Reply with exactly QA_RL_PERSIST_BLOCKED"}],"max_tokens":20}' \ + | jq -R -e '. == "429"' >/dev/null +jq -e '.error.type == "rate_limit_error" and .error.code == "rate_limit_exceeded"' "$BODY_FILE" >/dev/null + +curl -fsS -X DELETE "$BASE_URL/admin/rate-limits" \ + -H 'Content-Type: application/json' \ + -d "{\"user_path\":\"$RL_PATH\",\"limit_key\":{\"period\":\"hour\"}}" \ + | jq -e --arg p "$RL_PATH" 'all(.rate_limits[]?; .user_path != $p)' >/dev/null +``` + +### S206 `reset-one` stays cleared across `--reload` + +Burns an hour window, resets it, reloads, and verifies the next request +succeeds. + +```bash +RL_PATH="/qa/ratelimit/persist-reset/$QA_SUFFIX" +BODY_FILE="$QA_RUN_DIR/s206.body.json" + +curl -fsS -X PUT "$BASE_URL/admin/rate-limits" \ + -H 'Content-Type: application/json' \ + -d "{\"user_path\":\"$RL_PATH\",\"limit_key\":{\"period\":\"hour\"},\"max_requests\":1}" \ + | jq -e --arg p "$RL_PATH" 'any(.rate_limits[]?; .user_path == $p and .max_requests == 1)' >/dev/null + +curl -fsS -o "$BODY_FILE" "$BASE_URL/v1/chat/completions" \ + -H 'Content-Type: application/json' \ + -H "X-GoModel-User-Path: $RL_PATH/leaf" \ + -d '{"model":"gpt-4.1-nano","messages":[{"role":"user","content":"Reply with exactly QA_RL_RESET_PERSIST_OK"}],"max_tokens":20}' +assert_chat_response_contains "$BODY_FILE" "openai" "QA_RL_RESET_PERSIST_OK" + +curl -fsS -X POST "$BASE_URL/admin/rate-limits/reset-one" \ + -H 'Content-Type: application/json' \ + -d "{\"user_path\":\"$RL_PATH\",\"period\":\"hour\"}" >/dev/null + +reload_release_gateway sqlite-main "$BASE_URL" + +curl -fsS -o "$BODY_FILE" "$BASE_URL/v1/chat/completions" \ + -H 'Content-Type: application/json' \ + -H "X-GoModel-User-Path: $RL_PATH/leaf" \ + -d '{"model":"gpt-4.1-nano","messages":[{"role":"user","content":"Reply with exactly QA_RL_RESET_PERSIST_AGAIN"}],"max_tokens":20}' +assert_chat_response_contains "$BODY_FILE" "openai" "QA_RL_RESET_PERSIST_AGAIN" + +curl -fsS -X DELETE "$BASE_URL/admin/rate-limits" \ + -H 'Content-Type: application/json' \ + -d "{\"user_path\":\"$RL_PATH\",\"limit_key\":{\"period\":\"hour\"}}" \ + | jq -e --arg p "$RL_PATH" 'all(.rate_limits[]?; .user_path != $p)' >/dev/null +``` + +### S207 Request-window counters survive `--reload` on PostgreSQL and MongoDB + +Same as S205 against the smoke gateways. + +```bash +for item in "pg-smoke $PG_BASE_URL" "mongo-smoke $MONGO_BASE_URL"; do + set -- $item + GW="$1" + URL="$2" + RL_PATH="/qa/ratelimit/persist-$GW/$QA_SUFFIX" + BODY_FILE="$QA_RUN_DIR/s207.$GW.body.json" + + curl -fsS -X PUT "$URL/admin/rate-limits" \ + -H 'Content-Type: application/json' \ + -d "{\"user_path\":\"$RL_PATH\",\"limit_key\":{\"period\":\"hour\"},\"max_requests\":1}" \ + | jq -e --arg p "$RL_PATH" 'any(.rate_limits[]?; .user_path == $p and .max_requests == 1)' >/dev/null + + curl -fsS -o "$BODY_FILE" "$URL/v1/chat/completions" \ + -H 'Content-Type: application/json' \ + -H "X-GoModel-User-Path: $RL_PATH/leaf" \ + -d '{"model":"gpt-4.1-nano","messages":[{"role":"user","content":"Reply with exactly QA_RL_PERSIST_BACKEND_OK"}],"max_tokens":20}' + assert_chat_response_contains "$BODY_FILE" "openai" "QA_RL_PERSIST_BACKEND_OK" + + reload_release_gateway "$GW" "$URL" + + curl -sS -o "$BODY_FILE" -w '%{http_code}' "$URL/v1/chat/completions" \ + -H 'Content-Type: application/json' \ + -H "X-GoModel-User-Path: $RL_PATH/leaf" \ + -d '{"model":"gpt-4.1-nano","messages":[{"role":"user","content":"Reply with exactly QA_RL_PERSIST_BACKEND_BLOCKED"}],"max_tokens":20}' \ + | jq -R -e '. == "429"' >/dev/null + + curl -fsS -X DELETE "$URL/admin/rate-limits" \ + -H 'Content-Type: application/json' \ + -d "{\"user_path\":\"$RL_PATH\",\"limit_key\":{\"period\":\"hour\"}}" \ + | jq -e --arg p "$RL_PATH" 'all(.rate_limits[]?; .user_path != $p)' >/dev/null +done +``` From c37f743515e0347b031d09f53487b0cc85846cef Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Sun, 16 Aug 2026 16:45:19 +0200 Subject: [PATCH 06/12] fix(ratelimit): make counter snapshots crash-safe Save upserts then prunes instead of delete-all-first. A failed load leaves the generation idle so it cannot flush an empty snapshot over durable windows. Start/Close are idempotent; reset returns store errors. --- config/config_test.go | 2 +- ...-16_rate-limit-counter-persistence-spec.md | 43 +++--- internal/ratelimit/factory.go | 3 +- internal/ratelimit/persist.go | 107 ++++++++++---- internal/ratelimit/persist_test.go | 135 ++++++++++++++++++ internal/ratelimit/service.go | 38 +++-- internal/ratelimit/store.go | 3 + internal/ratelimit/store_mongodb.go | 51 +++++-- internal/ratelimit/store_sql.go | 56 ++++++-- internal/ratelimit/store_sql_test.go | 22 ++- tests/e2e/release-e2e-scenarios.md | 2 +- 11 files changed, 371 insertions(+), 91 deletions(-) diff --git a/config/config_test.go b/config/config_test.go index efae480f5..128234f09 100644 --- a/config/config_test.go +++ b/config/config_test.go @@ -65,7 +65,7 @@ func clearAllConfigEnvVars(t *testing.T) { "USAGE_PRICING_RECALCULATION_ENABLED", "USAGE_BUFFER_SIZE", "USAGE_FLUSH_INTERVAL", "USAGE_RETENTION_DAYS", "BUDGETS_ENABLED", - "RATE_LIMITS_ENABLED", + "RATE_LIMITS_ENABLED", "RATE_LIMITS_FLUSH_INTERVAL", "DASHBOARD_LIVE_LOGS_ENABLED", "DASHBOARD_LIVE_LOGS_BUFFER_SIZE", "DASHBOARD_LIVE_LOGS_REPLAY_LIMIT", "DASHBOARD_LIVE_LOGS_HEARTBEAT_SECONDS", "GUARDRAILS_ENABLED", "ENABLE_GUARDRAILS_FOR_BATCH_PROCESSING", diff --git a/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md b/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md index 69c346430..e7e48cf24 100644 --- a/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md +++ b/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md @@ -51,7 +51,7 @@ implement Redis-Lua. Two seams, only the first used as a live store today. -``` +```text Service.Acquire / RecordTokens / RouteAvailable / Status / Reset │ ▼ @@ -170,17 +170,19 @@ A shared rule’s `Partition` is `""`. Add to the existing `Store` interface. No second factory. -- `LoadCounters(ctx) ([]windowSnapshot, error)` — all rows. -- `SaveCounters(ctx, []windowSnapshot) error` — **replaces** the - table/collection with the provided set (delete all, then insert). - SQL does this in one transaction. Mongo uses the same - transaction-plus-standalone-fallback as `ReplaceConfigRules` today - (replica-set e2e uses `rs0`; a hard transaction-only write would - fail open on standalone Mongo). The payload is every **live window - key** still in the limiter maps whose definition is a current - windowed rule (period > 0). That includes one row per active - per-child partition. It never includes leftover map entries for a - dropped definition, and never includes concurrent keys. +- `LoadCounters(ctx) ([]windowSnapshot, error)` — all rows. A + malformed row is skipped and logged; a query failure is a load + error. +- `SaveCounters(ctx, []windowSnapshot) error` — **upsert** each + snapshot, then delete rows that are no longer in the set. Never + delete-all first: a crash or a failed insert must leave the previous + generation intact. SQL does upsert+prune in one transaction. Mongo + uses the same algorithm, with the existing transaction-plus- + standalone-fallback. The payload is every **live window key** still + in the limiter maps whose definition is a current windowed rule + (period > 0). That includes one row per active per-child partition. + It never includes leftover map entries for a dropped definition, and + never includes concurrent keys. - `DeleteCounter(ctx, scope, subject, periodSeconds) error` — every partition of that **definition** (matches `limiter.reset` / `sameDefinition`). Missing rows are not an error. Resetting a @@ -198,8 +200,8 @@ serving, then shuts the old one down, then starts the new one | Call | When | Persistence | |---|---|---| | `New` / `NewService` | `app.New`, tests | Empty limiter. No load, no flush, no loop. | -| `Start` | `App.startServer`, after the previous generation’s `Shutdown` | Load snapshot, apply rows whose **definition** still matches a current rule and whose `partition` matches the rule’s mode (empty iff `!PerChild`), re-arm child expiry, then start the flush loop if `flush_interval > 0`. Mark the generation **active**. | -| `Close` (`Result.Close`) | `App.Shutdown` | Stop the flush loop. If **active**, write one last snapshot. Then `Service.Close()` (expiry worker, already wired today) and close the store. | +| `Start` | `App.startServer`, after the previous generation’s `Shutdown` | Load snapshot, apply rows whose **definition** still matches a current rule and whose `partition` matches the rule’s mode (empty iff `!PerChild`), re-arm child expiry, then start the flush loop if `flush_interval > 0`. Mark the generation **active** only if load succeeded. Start is idempotent. A failed load leaves the generation idle so it cannot flush an empty snapshot over durable windows. | +| `Close` (`Result.Close`) | `App.Shutdown` | Stop the flush loop. If **active**, write one last snapshot. Then stop the expiry worker and close the store. Close is idempotent. | An idle replacement that is built and then discarded (failed listen, process exiting during rebuild) is not active, so its `Close` must not @@ -264,9 +266,9 @@ that concurrency is in-memory only. | Operation | Memory | Snapshot | |---|---|---| | `Admit` / `RecordTokens` | update windows | next flush / Close | -| `ResetRule` | `limiter.reset` (all partitions of the definition) | `DeleteCounter` under `persistMu` (all partitions) | -| `ResetAll` | `limiter.resetAll` | `DeleteAllCounters` under `persistMu` | -| `DeleteRule` | `limiter.reset` of that definition | `DeleteCounter` under `persistMu` | +| `ResetRule` | `limiter.reset` (all partitions of the definition) | `DeleteCounter` under `persistMu` (all partitions); store errors are returned | +| `ResetAll` | `limiter.resetAll` | `DeleteAllCounters` under `persistMu`; store errors are returned | +| `DeleteRule` | `limiter.reset` of that definition | `DeleteCounter` under `persistMu`; store errors are returned | | `ReplaceConfigRules` / `Refresh` | already prunes maps when a definition disappears or `PerChild` flips (`Service.Refresh` after `#670`) | **do not touch the store** | `ReplaceConfigRules` runs from `factory.New` → `seedConfiguredRules` @@ -295,9 +297,10 @@ putting the row back. Persistence never fails a request. `admit()` does not see store errors. -- **Load failure:** log and start with empty windows. Do not block the - listener. Skip corrupt rows; restore the good ones. Still mark the - generation active so later flushes work. +- **Load failure:** log and leave the generation **idle**. Do not + block the listener. Do not flush. The previous snapshot stays on + disk. Skip corrupt rows; restore the good ones. Only a successful + load (including “no rows”) activates persistence. - **Periodic flush failure:** log and retry on the next tick. Memory stays authoritative. - **Shutdown flush failure:** log and continue teardown. Do not hang diff --git a/internal/ratelimit/factory.go b/internal/ratelimit/factory.go index 12d886806..ec5263753 100644 --- a/internal/ratelimit/factory.go +++ b/internal/ratelimit/factory.go @@ -5,7 +5,6 @@ import ( "errors" "fmt" "sync" - "time" "go.mongodb.org/mongo-driver/v2/mongo" @@ -63,7 +62,7 @@ func newResult(ctx context.Context, cfg *config.Config, storeConn storage.Storag } service, err := NewService(ctx, store, WithQuotaTemplates(quotaTemplatesEnabled), - WithFlushInterval(time.Duration(cfg.RateLimits.FlushInterval)*time.Second), + WithFlushInterval(flushIntervalFromSeconds(cfg.RateLimits.FlushInterval)), ) if err != nil { return nil, err diff --git a/internal/ratelimit/persist.go b/internal/ratelimit/persist.go index 0cf331e13..aaf56f036 100644 --- a/internal/ratelimit/persist.go +++ b/internal/ratelimit/persist.go @@ -3,9 +3,20 @@ package ratelimit import ( "context" "log/slog" + "math" "time" ) +const persistTimeout = 5 * time.Second + +type persistState int + +const ( + persistIdle persistState = iota + persistActive + persistClosed +) + // WithFlushInterval sets how often an active generation writes window // snapshots. Zero disables the periodic loop; Start still loads and Close // of an active generation still writes once. @@ -18,22 +29,63 @@ func WithFlushInterval(interval time.Duration) ServiceOption { } } +func flushIntervalFromSeconds(seconds int) time.Duration { + if seconds <= 0 { + return 0 + } + maxSeconds := int(math.MaxInt64 / int64(time.Second)) + if seconds > maxSeconds { + return time.Duration(math.MaxInt64) + } + return time.Duration(seconds) * time.Second +} + +func persistContext(parent context.Context) (context.Context, context.CancelFunc) { + if parent != nil { + if _, ok := parent.Deadline(); ok { + return context.WithCancel(parent) + } + return context.WithTimeout(parent, persistTimeout) + } + return context.WithTimeout(context.Background(), persistTimeout) +} + func (s *Service) Start(ctx context.Context) { if s == nil || s.store == nil { return } - s.loadCounters(ctx) + + s.lifeMu.Lock() + if s.persistState != persistIdle { + s.lifeMu.Unlock() + return + } + s.lifeMu.Unlock() + + loadCtx, cancel := persistContext(ctx) + err := s.loadCounters(loadCtx) + cancel() + + s.lifeMu.Lock() + defer s.lifeMu.Unlock() + if s.persistState != persistIdle { + return + } + if err != nil { + slog.Warn("rate limit counters: load failed; not persisting this generation", "error", err) + return + } s.startFlushLoop() - s.active.Store(true) + s.persistState = persistActive } -func (s *Service) loadCounters(ctx context.Context) { +func (s *Service) loadCounters(ctx context.Context) error { snapshots, err := s.store.LoadCounters(ctx) if err != nil { - slog.Warn("rate limit counters: load failed; starting empty", "error", err) - return + return err } s.limiter.restore(snapshots, s.Rules(), time.Now().UTC()) + return nil } func (s *Service) startFlushLoop() { @@ -42,15 +94,20 @@ func (s *Service) startFlushLoop() { } s.flushStop = make(chan struct{}) s.flushDone = make(chan struct{}) + interval := s.flushInterval + stop := s.flushStop + done := s.flushDone go func() { - defer close(s.flushDone) - ticker := time.NewTicker(s.flushInterval) + defer close(done) + ticker := time.NewTicker(interval) defer ticker.Stop() for { select { case <-ticker.C: - s.flush(context.Background()) - case <-s.flushStop: + flushCtx, cancel := persistContext(context.Background()) + s.flush(flushCtx) + cancel() + case <-stop: return } } @@ -68,36 +125,24 @@ func (s *Service) flush(ctx context.Context) { } } -func (s *Service) persistDelete(scope RuleScope, subject string, periodSeconds int64) { +func (s *Service) persistDelete(ctx context.Context, scope RuleScope, subject string, periodSeconds int64) error { if s == nil || s.store == nil { - return + return nil } + writeCtx, cancel := persistContext(ctx) + defer cancel() s.persistMu.Lock() defer s.persistMu.Unlock() - if err := s.store.DeleteCounter(context.Background(), scope, subject, periodSeconds); err != nil { - slog.Error("rate limit counters: delete failed", "scope", scope, "subject", subject, "period_seconds", periodSeconds, "error", err) - } + return s.store.DeleteCounter(writeCtx, scope, subject, periodSeconds) } -func (s *Service) persistDeleteAll() { +func (s *Service) persistDeleteAll(ctx context.Context) error { if s == nil || s.store == nil { - return + return nil } + writeCtx, cancel := persistContext(ctx) + defer cancel() s.persistMu.Lock() defer s.persistMu.Unlock() - if err := s.store.DeleteAllCounters(context.Background()); err != nil { - slog.Error("rate limit counters: delete-all failed", "error", err) - } -} - -func (s *Service) stopFlushAndSave() { - s.flushOnce.Do(func() { - if s.flushStop != nil { - close(s.flushStop) - <-s.flushDone - } - if s.active.Load() { - s.flush(context.Background()) - } - }) + return s.store.DeleteAllCounters(writeCtx) } diff --git a/internal/ratelimit/persist_test.go b/internal/ratelimit/persist_test.go index e9e98f1c7..37df7b356 100644 --- a/internal/ratelimit/persist_test.go +++ b/internal/ratelimit/persist_test.go @@ -2,6 +2,7 @@ package ratelimit import ( "context" + "errors" "sync/atomic" "testing" "time" @@ -244,6 +245,122 @@ func TestRestoreIgnoresSharedRowOnPerChildRule(t *testing.T) { } } +func TestFailedLoadDoesNotReplacePersistedWindows(t *testing.T) { + now := time.Now().Unix() + store := &failLoadStore{err: errors.New("store down")} + store.counters = []WindowSnapshot{{ + Scope: string(ScopeUserPath), Subject: "/team", PeriodSeconds: PeriodHourSeconds, + RequestsWindowStart: now, RequestsCurrent: 9, + }} + if err := store.UpsertRules(context.Background(), []Rule{{ + Scope: ScopeUserPath, Subject: "/team", PeriodSeconds: PeriodHourSeconds, + MaxRequests: new(int64(10)), Source: SourceManual, + }}); err != nil { + t.Fatalf("seed: %v", err) + } + service, err := NewService(context.Background(), store, WithFlushInterval(10*time.Millisecond)) + if err != nil { + t.Fatalf("NewService: %v", err) + } + service.Start(context.Background()) + if _, err := service.Acquire(onPath("/team"), time.Now().UTC()); err != nil { + t.Fatalf("Acquire: %v", err) + } + service.Close() + time.Sleep(30 * time.Millisecond) + if len(store.counters) != 1 || store.counters[0].RequestsCurrent != 9 { + t.Fatalf("persisted windows = %+v, want the pre-failure snapshot", store.counters) + } +} + +func TestStartIsIdempotentAndCloseStopsTheLoop(t *testing.T) { + store := &recordingStore{} + if err := store.UpsertRules(context.Background(), []Rule{{ + Subject: "/", PeriodSeconds: PeriodHourSeconds, MaxRequests: new(int64(50)), Source: SourceManual, + }}); err != nil { + t.Fatalf("seed: %v", err) + } + service, err := NewService(context.Background(), store, WithFlushInterval(15*time.Millisecond)) + if err != nil { + t.Fatalf("NewService: %v", err) + } + service.Start(context.Background()) + service.Start(context.Background()) + service.Close() + afterClose := store.saves.Load() + time.Sleep(50 * time.Millisecond) + if got := store.saves.Load(); got != afterClose { + t.Fatalf("saves after Close grew from %d to %d; an extra flush loop is still running", afterClose, got) + } +} + +func TestFlushIntervalWritesBeforeClose(t *testing.T) { + store := &recordingStore{} + if err := store.UpsertRules(context.Background(), []Rule{{ + Subject: "/", PeriodSeconds: PeriodHourSeconds, MaxRequests: new(int64(50)), Source: SourceManual, + }}); err != nil { + t.Fatalf("seed: %v", err) + } + service, err := NewService(context.Background(), store, WithFlushInterval(15*time.Millisecond)) + if err != nil { + t.Fatalf("NewService: %v", err) + } + t.Cleanup(service.Close) + if _, err := service.Acquire(onPath("/"), time.Now().UTC()); err != nil { + t.Fatalf("Acquire: %v", err) + } + service.Start(context.Background()) + deadline := time.Now().Add(200 * time.Millisecond) + for store.saves.Load() == 0 && time.Now().Before(deadline) { + time.Sleep(5 * time.Millisecond) + } + if store.saves.Load() == 0 { + t.Fatal("positive flush interval never wrote before Close") + } +} + +func TestFlushIntervalZeroOnlyWritesOnClose(t *testing.T) { + store := &recordingStore{} + if err := store.UpsertRules(context.Background(), []Rule{{ + Subject: "/", PeriodSeconds: PeriodHourSeconds, MaxRequests: new(int64(50)), Source: SourceManual, + }}); err != nil { + t.Fatalf("seed: %v", err) + } + service, err := NewService(context.Background(), store, WithFlushInterval(0)) + if err != nil { + t.Fatalf("NewService: %v", err) + } + if _, err := service.Acquire(onPath("/"), time.Now().UTC()); err != nil { + t.Fatalf("Acquire: %v", err) + } + service.Start(context.Background()) + time.Sleep(30 * time.Millisecond) + if store.saves.Load() != 0 { + t.Fatalf("interval 0 wrote %d times before Close", store.saves.Load()) + } + service.Close() + if store.saves.Load() != 1 { + t.Fatalf("saves after Close = %d, want 1", store.saves.Load()) + } +} + +func TestResetRuleReturnsPersistError(t *testing.T) { + store := &failDeleteStore{err: errors.New("delete failed")} + if err := store.UpsertRules(context.Background(), []Rule{{ + Subject: "/team", PeriodSeconds: PeriodHourSeconds, MaxRequests: new(int64(1)), Source: SourceManual, + }}); err != nil { + t.Fatalf("seed: %v", err) + } + service, err := NewService(context.Background(), store) + if err != nil { + t.Fatalf("NewService: %v", err) + } + t.Cleanup(service.Close) + if err := service.ResetRule(ScopeUserPath, "/team", PeriodHourSeconds); err == nil { + t.Fatal("ResetRule() error = nil, want persist error") + } +} + type recordingStore struct { memStore saves atomic.Int64 @@ -253,3 +370,21 @@ func (s *recordingStore) SaveCounters(ctx context.Context, snapshots []WindowSna s.saves.Add(1) return s.memStore.SaveCounters(ctx, snapshots) } + +type failLoadStore struct { + memStore + err error +} + +func (s *failLoadStore) LoadCounters(context.Context) ([]WindowSnapshot, error) { + return nil, s.err +} + +type failDeleteStore struct { + memStore + err error +} + +func (s *failDeleteStore) DeleteCounter(context.Context, RuleScope, string, int64) error { + return s.err +} diff --git a/internal/ratelimit/service.go b/internal/ratelimit/service.go index 1b09ec415..2c148ff4a 100644 --- a/internal/ratelimit/service.go +++ b/internal/ratelimit/service.go @@ -7,7 +7,6 @@ import ( "sort" "strings" "sync" - "sync/atomic" "time" ) @@ -39,10 +38,10 @@ type Service struct { flushInterval time.Duration persistMu sync.Mutex + lifeMu sync.Mutex + persistState persistState flushStop chan struct{} flushDone chan struct{} - flushOnce sync.Once - active atomic.Bool } func NewService(ctx context.Context, store Store, options ...ServiceOption) (*Service, error) { @@ -71,7 +70,28 @@ func (s *Service) Close() { if s == nil { return } - s.stopFlushAndSave() + s.lifeMu.Lock() + if s.persistState == persistClosed { + s.lifeMu.Unlock() + return + } + wasActive := s.persistState == persistActive + s.persistState = persistClosed + stop := s.flushStop + done := s.flushDone + s.flushStop = nil + s.flushDone = nil + s.lifeMu.Unlock() + + if stop != nil { + close(stop) + <-done + } + if wasActive { + flushCtx, cancel := persistContext(context.Background()) + s.flush(flushCtx) + cancel() + } if s.limiter != nil { s.limiter.close() } @@ -151,7 +171,9 @@ func (s *Service) DeleteRule(ctx context.Context, scope RuleScope, subject strin return err } s.limiter.reset(ruleKey{scope: scope, subject: subject, periodSeconds: periodSeconds}) - s.persistDelete(scope, subject, periodSeconds) + if err := s.persistDelete(ctx, scope, subject, periodSeconds); err != nil { + return err + } return s.Refresh(ctx) } @@ -350,8 +372,7 @@ func (s *Service) ResetRule(scope RuleScope, subject string, periodSeconds int64 return err } s.limiter.reset(ruleKey{scope: scope, subject: subject, periodSeconds: periodSeconds}) - s.persistDelete(scope, subject, periodSeconds) - return nil + return s.persistDelete(context.Background(), scope, subject, periodSeconds) } // ResetAll clears every live window counter. @@ -360,8 +381,7 @@ func (s *Service) ResetAll() error { return ErrUnavailable } s.limiter.resetAll() - s.persistDeleteAll() - return nil + return s.persistDeleteAll(context.Background()) } func (s *Service) matchingRules(subjects Subjects) []Rule { diff --git a/internal/ratelimit/store.go b/internal/ratelimit/store.go index fac424607..25c6fbec8 100644 --- a/internal/ratelimit/store.go +++ b/internal/ratelimit/store.go @@ -16,6 +16,9 @@ type Store interface { DeleteRule(ctx context.Context, scope RuleScope, subject string, periodSeconds int64) error ReplaceConfigRules(ctx context.Context, rules []Rule) error LoadCounters(ctx context.Context) ([]WindowSnapshot, error) + // SaveCounters upserts snapshots and deletes rows that are no longer + // present. It must not delete existing rows before the new values are + // durable, so a crash cannot wipe the previous generation. SaveCounters(ctx context.Context, snapshots []WindowSnapshot) error DeleteCounter(ctx context.Context, scope RuleScope, subject string, periodSeconds int64) error DeleteAllCounters(ctx context.Context) error diff --git a/internal/ratelimit/store_mongodb.go b/internal/ratelimit/store_mongodb.go index 5667d0a96..af9794d50 100644 --- a/internal/ratelimit/store_mongodb.go +++ b/internal/ratelimit/store_mongodb.go @@ -366,7 +366,8 @@ func (s *MongoDBStore) LoadCounters(ctx context.Context) ([]WindowSnapshot, erro for cursor.Next(ctx) { var snap WindowSnapshot if err := cursor.Decode(&snap); err != nil { - return nil, fmt.Errorf("decode rate limit counter: %w", err) + slog.Warn("rate limit counters: skipping malformed snapshot", "error", err) + continue } snapshots = append(snapshots, snap) } @@ -378,22 +379,20 @@ func (s *MongoDBStore) LoadCounters(ctx context.Context) ([]WindowSnapshot, erro func (s *MongoDBStore) SaveCounters(ctx context.Context, snapshots []WindowSnapshot) error { write := func(writeCtx context.Context) error { - if _, err := s.counters.DeleteMany(writeCtx, bson.D{}); err != nil { - return fmt.Errorf("clear rate limit counters: %w", err) - } - if len(snapshots) == 0 { - return nil - } now := time.Now().Unix() - docs := make([]any, 0, len(snapshots)) for _, snap := range snapshots { snap.UpdatedAt = now - docs = append(docs, snap) - } - if _, err := s.counters.InsertMany(writeCtx, docs); err != nil { - return fmt.Errorf("insert rate limit counters: %w", err) + _, err := s.counters.UpdateOne( + writeCtx, + counterIdentityFilter(snap), + bson.D{{Key: "$set", Value: snap}}, + options.UpdateOne().SetUpsert(true), + ) + if err != nil { + return fmt.Errorf("upsert rate limit counter: %w", err) + } } - return nil + return s.deleteOrphanCounters(writeCtx, snapshots) } session, err := s.counters.Database().Client().StartSession() @@ -427,6 +426,32 @@ func (s *MongoDBStore) SaveCounters(ctx context.Context, snapshots []WindowSnaps return nil } +func (s *MongoDBStore) deleteOrphanCounters(ctx context.Context, snapshots []WindowSnapshot) error { + if len(snapshots) == 0 { + if _, err := s.counters.DeleteMany(ctx, bson.D{}); err != nil { + return fmt.Errorf("clear rate limit counters: %w", err) + } + return nil + } + keep := make(bson.A, 0, len(snapshots)) + for _, snap := range snapshots { + keep = append(keep, counterIdentityFilter(snap)) + } + if _, err := s.counters.DeleteMany(ctx, bson.D{{Key: "$nor", Value: keep}}); err != nil { + return fmt.Errorf("prune rate limit counters: %w", err) + } + return nil +} + +func counterIdentityFilter(snap WindowSnapshot) bson.D { + return bson.D{ + {Key: "scope", Value: snap.Scope}, + {Key: "subject", Value: snap.Subject}, + {Key: "partition", Value: snap.Partition}, + {Key: "period_seconds", Value: snap.PeriodSeconds}, + } +} + func (s *MongoDBStore) DeleteCounter(ctx context.Context, scope RuleScope, subject string, periodSeconds int64) error { _, err := s.counters.DeleteMany(ctx, bson.D{ {Key: "scope", Value: scope}, diff --git a/internal/ratelimit/store_sql.go b/internal/ratelimit/store_sql.go index 341683d6c..c6c90b524 100644 --- a/internal/ratelimit/store_sql.go +++ b/internal/ratelimit/store_sql.go @@ -201,30 +201,62 @@ func (s *SQLStore) LoadCounters(ctx context.Context) ([]WindowSnapshot, error) { return snapshots, nil } +const upsertCounterSQL = ` + INSERT INTO rate_limit_counters ( + scope, subject, partition, period_seconds, + requests_window_start, requests_current, requests_previous, + tokens_window_start, tokens_current, tokens_previous, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(scope, subject, partition, period_seconds) DO UPDATE SET + requests_window_start = excluded.requests_window_start, + requests_current = excluded.requests_current, + requests_previous = excluded.requests_previous, + tokens_window_start = excluded.tokens_window_start, + tokens_current = excluded.tokens_current, + tokens_previous = excluded.tokens_previous, + updated_at = excluded.updated_at +` + func (s *SQLStore) SaveCounters(ctx context.Context, snapshots []WindowSnapshot) error { now := time.Now().Unix() return s.db.InTx(ctx, func(q sqlx.Querier) error { - if _, err := q.Exec(ctx, `DELETE FROM rate_limit_counters`); err != nil { - return fmt.Errorf("clear rate limit counters: %w", err) - } for _, snap := range snapshots { - if _, err := q.Exec(ctx, ` - INSERT INTO rate_limit_counters ( - scope, subject, partition, period_seconds, - requests_window_start, requests_current, requests_previous, - tokens_window_start, tokens_current, tokens_previous, updated_at - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - `, snap.Scope, snap.Subject, snap.Partition, snap.PeriodSeconds, + if _, err := q.Exec(ctx, upsertCounterSQL, + snap.Scope, snap.Subject, snap.Partition, snap.PeriodSeconds, snap.RequestsWindowStart, snap.RequestsCurrent, snap.RequestsPrevious, snap.TokensWindowStart, snap.TokensCurrent, snap.TokensPrevious, now, ); err != nil { - return fmt.Errorf("insert rate limit counter: %w", err) + return fmt.Errorf("upsert rate limit counter: %w", err) } } - return nil + return deleteOrphanCounters(ctx, q, snapshots) }) } +func deleteOrphanCounters(ctx context.Context, q sqlx.Querier, snapshots []WindowSnapshot) error { + if len(snapshots) == 0 { + if _, err := q.Exec(ctx, `DELETE FROM rate_limit_counters`); err != nil { + return fmt.Errorf("clear rate limit counters: %w", err) + } + return nil + } + var query strings.Builder + query.WriteString(`DELETE FROM rate_limit_counters WHERE NOT (`) + args := make([]any, 0, len(snapshots)*4) + for i, snap := range snapshots { + if i > 0 { + query.WriteString(` OR `) + } + query.WriteString(`(scope = ? AND subject = ? AND partition = ? AND period_seconds = ?)`) + args = append(args, snap.Scope, snap.Subject, snap.Partition, snap.PeriodSeconds) + } + query.WriteString(`)`) + if _, err := q.Exec(ctx, query.String(), args...); err != nil { + return fmt.Errorf("prune rate limit counters: %w", err) + } + return nil +} + func (s *SQLStore) DeleteCounter(ctx context.Context, scope RuleScope, subject string, periodSeconds int64) error { _, err := s.db.Exec(ctx, ` DELETE FROM rate_limit_counters diff --git a/internal/ratelimit/store_sql_test.go b/internal/ratelimit/store_sql_test.go index 129069a04..b62a7adbe 100644 --- a/internal/ratelimit/store_sql_test.go +++ b/internal/ratelimit/store_sql_test.go @@ -257,11 +257,15 @@ func TestSQLStoreCounterRoundTrip(t *testing.T) { first := []WindowSnapshot{ { Scope: string(ScopeUserPath), Subject: "/customers", Partition: "/customers/alice", - PeriodSeconds: PeriodHourSeconds, RequestsWindowStart: 1700000000, RequestsCurrent: 3, + PeriodSeconds: PeriodHourSeconds, + RequestsWindowStart: 1700000000, RequestsCurrent: 3, RequestsPrevious: 1, + TokensWindowStart: 1700000000, TokensCurrent: 40, TokensPrevious: 10, }, { Scope: string(ScopeUserPath), Subject: "/customers", Partition: "/customers/bob", - PeriodSeconds: PeriodHourSeconds, RequestsWindowStart: 1700000000, RequestsCurrent: 1, + PeriodSeconds: PeriodHourSeconds, + RequestsWindowStart: 1700000060, RequestsCurrent: 1, RequestsPrevious: 2, + TokensWindowStart: 1700000060, TokensCurrent: 7, TokensPrevious: 8, }, } if err := store.SaveCounters(ctx, first); err != nil { @@ -274,6 +278,20 @@ func TestSQLStoreCounterRoundTrip(t *testing.T) { if len(got) != 2 { t.Fatalf("loaded = %d, want 2", len(got)) } + byPart := map[string]WindowSnapshot{} + for _, snap := range got { + byPart[snap.Partition] = snap + } + alice := byPart["/customers/alice"] + if alice.RequestsWindowStart != 1700000000 || alice.RequestsCurrent != 3 || alice.RequestsPrevious != 1 || + alice.TokensWindowStart != 1700000000 || alice.TokensCurrent != 40 || alice.TokensPrevious != 10 { + t.Fatalf("alice = %+v", alice) + } + bob := byPart["/customers/bob"] + if bob.RequestsWindowStart != 1700000060 || bob.RequestsCurrent != 1 || bob.RequestsPrevious != 2 || + bob.TokensWindowStart != 1700000060 || bob.TokensCurrent != 7 || bob.TokensPrevious != 8 { + t.Fatalf("bob = %+v", bob) + } if err := store.DeleteCounter(ctx, ScopeUserPath, "/customers", PeriodHourSeconds); err != nil { t.Fatalf("DeleteCounter: %v", err) diff --git a/tests/e2e/release-e2e-scenarios.md b/tests/e2e/release-e2e-scenarios.md index 65e23139b..50b1ad873 100644 --- a/tests/e2e/release-e2e-scenarios.md +++ b/tests/e2e/release-e2e-scenarios.md @@ -169,7 +169,7 @@ reload_release_gateway() { for _ in $(seq 1 50); do if tail -n +$((before + 1)) "$log_file" 2>/dev/null | grep -Fq 'configuration reloaded'; then for __ in $(seq 1 20); do - if curl -fsS "$url/health" >/dev/null 2>&1; then + if curl -fsS --connect-timeout 1 --max-time 2 "$url/health" >/dev/null 2>&1; then return 0 fi sleep 0.1 From 2be766034a662ce626f8ef3cfbc77172bfb3021a Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Sun, 16 Aug 2026 17:10:51 +0200 Subject: [PATCH 07/12] fix(ratelimit): serialize persist start and bound SQL prune Restore snapshots only after Close cannot interleave, using a starting state. Prune orphans by existing keys instead of one OR query per live partition, so per-child cardinality cannot blow the SQL parameter limit. --- internal/ratelimit/persist.go | 17 ++++------ internal/ratelimit/persist_test.go | 48 ++++++++++++++++++++++++++++ internal/ratelimit/snapshot.go | 9 +++++- internal/ratelimit/store_sql.go | 42 +++++++++++++++++------- internal/ratelimit/store_sql_test.go | 23 +++++++++++++ 5 files changed, 116 insertions(+), 23 deletions(-) diff --git a/internal/ratelimit/persist.go b/internal/ratelimit/persist.go index aaf56f036..af6eac200 100644 --- a/internal/ratelimit/persist.go +++ b/internal/ratelimit/persist.go @@ -13,6 +13,7 @@ type persistState int const ( persistIdle persistState = iota + persistStarting persistActive persistClosed ) @@ -60,34 +61,28 @@ func (s *Service) Start(ctx context.Context) { s.lifeMu.Unlock() return } + s.persistState = persistStarting s.lifeMu.Unlock() loadCtx, cancel := persistContext(ctx) - err := s.loadCounters(loadCtx) + snapshots, err := s.store.LoadCounters(loadCtx) cancel() s.lifeMu.Lock() defer s.lifeMu.Unlock() - if s.persistState != persistIdle { + if s.persistState != persistStarting { return } if err != nil { + s.persistState = persistIdle slog.Warn("rate limit counters: load failed; not persisting this generation", "error", err) return } + s.limiter.restore(snapshots, s.Rules(), time.Now().UTC()) s.startFlushLoop() s.persistState = persistActive } -func (s *Service) loadCounters(ctx context.Context) error { - snapshots, err := s.store.LoadCounters(ctx) - if err != nil { - return err - } - s.limiter.restore(snapshots, s.Rules(), time.Now().UTC()) - return nil -} - func (s *Service) startFlushLoop() { if s.flushInterval <= 0 { return diff --git a/internal/ratelimit/persist_test.go b/internal/ratelimit/persist_test.go index 37df7b356..b52547700 100644 --- a/internal/ratelimit/persist_test.go +++ b/internal/ratelimit/persist_test.go @@ -344,6 +344,42 @@ func TestFlushIntervalZeroOnlyWritesOnClose(t *testing.T) { } } +func TestCloseDuringLoadDoesNotRestore(t *testing.T) { + now := time.Now().Unix() + store := &blockingLoadStore{ + started: make(chan struct{}), + release: make(chan struct{}), + } + store.counters = []WindowSnapshot{{ + Scope: string(ScopeUserPath), Subject: "/team", PeriodSeconds: PeriodHourSeconds, + RequestsWindowStart: now, RequestsCurrent: 1, + }} + if err := store.UpsertRules(context.Background(), []Rule{{ + Scope: ScopeUserPath, Subject: "/team", PeriodSeconds: PeriodHourSeconds, + MaxRequests: new(int64(1)), Source: SourceManual, + }}); err != nil { + t.Fatalf("seed: %v", err) + } + service, err := NewService(context.Background(), store) + if err != nil { + t.Fatalf("NewService: %v", err) + } + + done := make(chan struct{}) + go func() { + defer close(done) + service.Start(context.Background()) + }() + <-store.started + service.Close() + close(store.release) + <-done + + if _, err := service.Acquire(onPath("/team"), time.Now().UTC()); err != nil { + t.Fatalf("Acquire after Close-during-load: %v (window must not have been restored)", err) + } +} + func TestResetRuleReturnsPersistError(t *testing.T) { store := &failDeleteStore{err: errors.New("delete failed")} if err := store.UpsertRules(context.Background(), []Rule{{ @@ -388,3 +424,15 @@ type failDeleteStore struct { func (s *failDeleteStore) DeleteCounter(context.Context, RuleScope, string, int64) error { return s.err } + +type blockingLoadStore struct { + memStore + started chan struct{} + release chan struct{} +} + +func (s *blockingLoadStore) LoadCounters(context.Context) ([]WindowSnapshot, error) { + close(s.started) + <-s.release + return s.memStore.LoadCounters(context.Background()) +} diff --git a/internal/ratelimit/snapshot.go b/internal/ratelimit/snapshot.go index 869e0999c..f5011cd81 100644 --- a/internal/ratelimit/snapshot.go +++ b/internal/ratelimit/snapshot.go @@ -1,6 +1,9 @@ package ratelimit -import "time" +import ( + "strconv" + "time" +) // WindowSnapshot is one persisted request/token sliding window. // Partition is empty for a shared rule and the child path for a per-child @@ -23,6 +26,10 @@ func definitionKey(scope RuleScope, subject string, periodSeconds int64) ruleKey return ruleKey{scope: scope, subject: subject, periodSeconds: periodSeconds} } +func snapshotIdentity(s WindowSnapshot) string { + return s.Scope + "\x00" + s.Subject + "\x00" + s.Partition + "\x00" + strconv.FormatInt(s.PeriodSeconds, 10) +} + func (l *limiter) snapshot(rules []Rule) []WindowSnapshot { l.mu.Lock() defer l.mu.Unlock() diff --git a/internal/ratelimit/store_sql.go b/internal/ratelimit/store_sql.go index c6c90b524..7de964e8d 100644 --- a/internal/ratelimit/store_sql.go +++ b/internal/ratelimit/store_sql.go @@ -240,19 +240,39 @@ func deleteOrphanCounters(ctx context.Context, q sqlx.Querier, snapshots []Windo } return nil } - var query strings.Builder - query.WriteString(`DELETE FROM rate_limit_counters WHERE NOT (`) - args := make([]any, 0, len(snapshots)*4) - for i, snap := range snapshots { - if i > 0 { - query.WriteString(` OR `) + keep := make(map[string]struct{}, len(snapshots)) + for _, snap := range snapshots { + keep[snapshotIdentity(snap)] = struct{}{} + } + rows, err := q.Query(ctx, ` + SELECT scope, subject, partition, period_seconds + FROM rate_limit_counters + `) + if err != nil { + return fmt.Errorf("list rate limit counters for prune: %w", err) + } + defer rows.Close() + + var stale []WindowSnapshot + for rows.Next() { + var snap WindowSnapshot + if err := rows.Scan(&snap.Scope, &snap.Subject, &snap.Partition, &snap.PeriodSeconds); err != nil { + return fmt.Errorf("scan rate limit counter identity: %w", err) + } + if _, ok := keep[snapshotIdentity(snap)]; !ok { + stale = append(stale, snap) } - query.WriteString(`(scope = ? AND subject = ? AND partition = ? AND period_seconds = ?)`) - args = append(args, snap.Scope, snap.Subject, snap.Partition, snap.PeriodSeconds) } - query.WriteString(`)`) - if _, err := q.Exec(ctx, query.String(), args...); err != nil { - return fmt.Errorf("prune rate limit counters: %w", err) + if err := rows.Err(); err != nil { + return fmt.Errorf("iterate rate limit counters for prune: %w", err) + } + for _, snap := range stale { + if _, err := q.Exec(ctx, ` + DELETE FROM rate_limit_counters + WHERE scope = ? AND subject = ? AND partition = ? AND period_seconds = ? + `, snap.Scope, snap.Subject, snap.Partition, snap.PeriodSeconds); err != nil { + return fmt.Errorf("prune rate limit counter: %w", err) + } } return nil } diff --git a/internal/ratelimit/store_sql_test.go b/internal/ratelimit/store_sql_test.go index b62a7adbe..e2c203f83 100644 --- a/internal/ratelimit/store_sql_test.go +++ b/internal/ratelimit/store_sql_test.go @@ -3,6 +3,7 @@ package ratelimit import ( "context" "errors" + "strconv" "testing" "github.com/enterpilot/gomodel/internal/storage/sqlx" @@ -318,6 +319,28 @@ func TestSQLStoreCounterRoundTrip(t *testing.T) { t.Fatalf("replaced = %+v, want alice only", got) } + many := make([]WindowSnapshot, 0, 300) + for i := range 300 { + many = append(many, WindowSnapshot{ + Scope: string(ScopeUserPath), Subject: "/customers", + Partition: "/customers/" + strconv.Itoa(i), + PeriodSeconds: PeriodHourSeconds, RequestsCurrent: int64(i + 1), + }) + } + if err := store.SaveCounters(ctx, many); err != nil { + t.Fatalf("SaveCounters many: %v", err) + } + if err := store.SaveCounters(ctx, many[:1]); err != nil { + t.Fatalf("SaveCounters prune many: %v", err) + } + got, err = store.LoadCounters(ctx) + if err != nil { + t.Fatalf("LoadCounters after prune many: %v", err) + } + if len(got) != 1 || got[0].RequestsCurrent != 1 { + t.Fatalf("after pruning 299 orphans = %+v, want the first snapshot", got) + } + if err := store.DeleteAllCounters(ctx); err != nil { t.Fatalf("DeleteAllCounters: %v", err) } From 67d543cfb51f5ebbf85299763bc7bcdcc0249b53 Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Sun, 16 Aug 2026 17:38:21 +0200 Subject: [PATCH 08/12] refactor(ratelimit): make counter snapshots additive MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SaveCounters replaced the whole row set on every flush: it computed which persisted rows were missing from the payload and deleted them. That made a one-second timer destructive — a crash or a failed write between the delete and the insert could drop a live hour window, a Mongo standalone had to be wrapped in a transaction dance to narrow the gap, and two replicas sharing a store deleted each other's rows on every tick. Nothing needed the set semantics. A row that stops being written is garbage, not a row to erase on sight, and restore already ignores any window older than two of its periods. So a save now upserts what it has and deletes only rows that went two periods without a write — one bounded statement per backend instead of a read-modify-write, no transaction on the Mongo path, and no save that can destroy a window it did not write. Also: log restored window count on Start, drop the now-stale implementation plan doc, and align the spec with what shipped (no counterBackend seam — one implementation does not justify the interface). Verified live end to end on SQLite (hour window survives --reload, graceful restart and SIGKILL; reset-one stays cleared across reload), and the store suites now run green against real PostgreSQL and a standalone MongoDB. --- ...-16_rate-limit-counter-persistence-plan.md | 60 -------- ...-16_rate-limit-counter-persistence-spec.md | 133 +++++++++--------- internal/ratelimit/persist.go | 10 ++ internal/ratelimit/service_test.go | 30 +++- internal/ratelimit/snapshot.go | 13 +- internal/ratelimit/store.go | 8 +- internal/ratelimit/store_mongodb.go | 76 +++------- internal/ratelimit/store_mongodb_test.go | 85 +++++++++++ internal/ratelimit/store_sql.go | 56 ++------ internal/ratelimit/store_sql_test.go | 39 +++-- 10 files changed, 247 insertions(+), 263 deletions(-) delete mode 100644 docs/dev/2026-08-16_rate-limit-counter-persistence-plan.md diff --git a/docs/dev/2026-08-16_rate-limit-counter-persistence-plan.md b/docs/dev/2026-08-16_rate-limit-counter-persistence-plan.md deleted file mode 100644 index eeaae8c7d..000000000 --- a/docs/dev/2026-08-16_rate-limit-counter-persistence-plan.md +++ /dev/null @@ -1,60 +0,0 @@ -# Rate Limit Counter Persistence Implementation Plan - -> **For agentic workers:** Execute task-by-task. Spec: -> `docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md`. - -**Goal:** Persist request/token sliding windows across restart and `--reload` -using the existing SQL/Mongo store, without putting I/O on `admit()`. - -**Architecture:** Keep the in-memory limiter as the live store. Snapshot -windows (including per-child `partition` keys) to `rate_limit_counters`. -`New` does not load or flush. `Start` (from `App.startServer`) loads then -optionally ticks. `Close` of an **active** generation writes once. -`persistMu` serializes flush vs reset/delete. Construction never deletes -snapshot rows. - -**Tech Stack:** Go, existing `internal/ratelimit` limiter + SQL/Mongo stores, -`RATE_LIMITS_FLUSH_INTERVAL` (default 1s). - ---- - -## Files - -- Modify: `config/ratelimit.go`, `config/config.go`, `config/ratelimit_test.go` -- Modify: `internal/ratelimit/store.go`, `store_sql.go`, `store_mongodb.go`, - `service.go`, `factory.go`, `service_test.go`, `store_sql_test.go`, - `store_mongodb_test.go` -- Create: `internal/ratelimit/snapshot.go`, `persist.go`, `persist_test.go`, - `snapshot_test.go` -- Modify: `internal/app/app.go` -- Modify: docs, `.env.template`, `config/config.example.yaml`, - `tests/e2e/release-e2e-scenarios.md` - -### Task 1: Config - -`FlushInterval int` on `RateLimitsConfig` (`yaml:"flush_interval"` -`env:"RATE_LIMITS_FLUSH_INTERVAL"`). Default `1` in `config.Load`. Reject -`< 0`. `0` is valid (no periodic loop). - -### Task 2: Store - -Add `LoadCounters`, `SaveCounters` (replace-all), `DeleteCounter` (all -partitions of a definition), `DeleteAllCounters`. SQL table and Mongo -collection `rate_limit_counters`, PK -`(scope, subject, partition, period_seconds)`. Mongo uses -transaction-plus-standalone-fallback like `ReplaceConfigRules`. - -### Task 3: Snapshot + persist - -`windowSnapshot` with `Partition`. Limiter `snapshot`/`restore` copy -`windowCounter` by value. Restore skips mode-mismatch and expired child -windows (`windowStart + 2*period`); re-arms `trackCounterExpiry`. -`Service.Start` / flush loop / `Close` compose with existing expiry -`Close`. `Reset*`/`DeleteRule` delete rows under `persistMu`. -`Refresh`/`ReplaceConfigRules` do not touch the snapshot store. - -### Task 4: Wire + docs + E2E - -`factory` passes flush interval. `App.startServer` calls `Start` before -listen. Docs: windows survive bounce; concurrency does not; still per -instance. S205–S207 shared hour rules + `reload_release_gateway`. diff --git a/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md b/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md index e7e48cf24..2bf889db9 100644 --- a/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md +++ b/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md @@ -16,8 +16,8 @@ this and listed two follow-ups: - §11.7 counter persistence across restarts (this work) - §11.2 Redis live counters for exact multi-replica enforcement (later) -This spec does §11.7 and extracts the seam §11.2 will need. It does not -implement Redis-Lua. +This spec does §11.7. It does not implement Redis-Lua, and does not +build a seam for it in advance (§5). ## 2. Goals @@ -25,8 +25,6 @@ implement Redis-Lua. `--reload`. - `admit()` stays in memory. Persistence is a background snapshot, never on the request path. -- A later Redis-Lua backend can replace the in-memory limiter without - changing `Acquire`, the usage tap, or `RouteAvailable`. - Default on whenever rate limits are on. No extra dependency. Uses the existing SQLite / Postgres / Mongo store that already holds the rules. - Crash loses at most one flush interval. Graceful stop and reload lose @@ -49,22 +47,19 @@ implement Redis-Lua. ## 4. Architecture -Two seams, only the first used as a live store today. - ```text Service.Acquire / RecordTokens / RouteAvailable / Status / Reset │ ▼ - counterBackend (unexported; package ratelimit) - │ - ├── memoryBackend current limiter + snapshot load/save - └── (later) redisLua live INCR; no snapshots + limiter (in memory, authoritative) + │ snapshot() / restore() + ▼ + Store.LoadCounters / SaveCounters / DeleteCounter / DeleteAllCounters ``` -Snapshots are an implementation detail of `memoryBackend`. They persist -only the sliding windows. The backend talks to the existing `Store`, -which grows four methods implemented by the SQL and Mongo stores already -used for rules. +Snapshots persist only the sliding windows, through the existing `Store` +that already holds the rules — four new methods on the SQL and Mongo +implementations, no second factory and no new dependency. Cardinality is one row per **live window key**, not one per rule definition: @@ -82,36 +77,25 @@ the limiter maps, so table size tracks live children, not historical ones. Several replicas sharing Postgres or Mongo last-write-wins **per -window key**. That is accepted: this work does not give shared live -counters, and the schema has no `instance_id`. A restart loads whoever -flushed last. Use one instance, or wait for the Redis-Lua backend, when -the limit must be exact across replicas. - -## 5. `counterBackend` - -Unexported interface in `internal/ratelimit`, matching what `limiter` -already does: - -- `Admit(rules []Rule, now time.Time) (HeaderSnapshot, []ruleKey, *ExceededError)` -- `Available(rules []Rule, now time.Time) bool` -- `Release(held []ruleKey)` -- `RecordTokens(rules []Rule, tokens int64, now time.Time)` -- `Status(rule Rule, now time.Time) Status` -- `Reset(key ruleKey)` -- `ResetAll()` - -`Service` holds a `counterBackend` instead of a concrete `*limiter`. -`ruleKey` stays unexported and already includes `partition` (`#670`). -`Reset(key)` keeps today’s `sameDefinition` behavior: it clears every -partition of that `(scope, subject, period)` definition. - -Public API additions: `Service.Start` and the flush-interval config -field. `Service.Close` already exists (stops the child-expiry worker); -persistence `Close` must call it after the final snapshot. - -The memory implementation is the current `limiter` plus snapshot -helpers. A future Redis-Lua type implements the same interface and -ignores the snapshot store. +window key** — no replica's save deletes a key it did not write, so +replicas overwrite each other only where they overlap, and never wipe +each other's table. That is accepted: this work does not give shared +live counters, and the schema has no `instance_id`. A restart loads +whoever flushed that key last. Use one instance, or wait for the +Redis-Lua backend, when the limit must be exact across replicas. + +## 5. No backend interface (decided against) + +An earlier draft extracted a `counterBackend` interface so a Redis-Lua +implementation could drop in later. It is not built, on purpose: there +is one implementation, so the interface would be an abstraction with no +second caller to justify it. `Service` keeps its concrete `*limiter`, +and the Redis work — if it happens — extracts the seam then, against a +real second implementation rather than a guess at one. + +What this change actually adds to the package surface is `Service.Start` +and the flush-interval option. `Service.Close` already existed (it stops +the child-expiry worker) and now writes the final snapshot first. ## 6. Snapshot data model @@ -129,7 +113,7 @@ requests_previous INT64 tokens_window_start INT64 tokens_current INT64 tokens_previous INT64 -updated_at INT64 -- unix seconds, diagnostics only +updated_at INT64 -- unix seconds; drives staleness collection PRIMARY KEY (scope, subject, partition, period_seconds) ``` @@ -145,7 +129,8 @@ Older databases just gain the table; no backfill. Units match `windowCounter`: Unix seconds and `int64` counts. After load, `advance()` already zeros a window more than one period old, so -rows do not need a TTL. Load still skips a child row whose window is +rows need no TTL of their own; `SaveCounters` collects them once they +stop being written. Load still skips a child row whose window is already past the in-memory expiry horizon (`windowStart + 2*period`), and must call `trackCounterExpiry` for every restored partition so the existing worker keeps pruning. @@ -174,12 +159,18 @@ Add to the existing `Store` interface. No second factory. malformed row is skipped and logged; a query failure is a load error. - `SaveCounters(ctx, []windowSnapshot) error` — **upsert** each - snapshot, then delete rows that are no longer in the set. Never - delete-all first: a crash or a failed insert must leave the previous - generation intact. SQL does upsert+prune in one transaction. Mongo - uses the same algorithm, with the existing transaction-plus- - standalone-fallback. The payload is every **live window key** still - in the limiter maps whose definition is a current windowed rule + snapshot (stamping `updated_at`), then collect rows that went two of + their own periods without a write. Nothing is ever deleted to make + room for a write, so a crash mid-save costs at most the flush + interval, never a whole window. Two periods is the same staleness + bound `restore` applies, so collection can only remove rows a load + would have discarded anyway — which also means a save is not + destructive to rows this instance does not know about (see the + replica note in §4). SQL is one upsert loop plus one bounded + `DELETE` in a transaction; Mongo is one unordered `BulkWrite` plus + the equivalent `DeleteMany`, no transaction needed because no step + depends on another. The payload is every **live window key** still in + the limiter maps whose definition is a current windowed rule (period > 0). That includes one row per active per-child partition. It never includes leftover map entries for a dropped definition, and never includes concurrent keys. @@ -229,10 +220,10 @@ reset/delete row deletes. `admit()` does not take it. Without that, a flush that sampled before a reset can `SaveCounters` after the row delete and resurrect the burned window (S206 would flake). -Orphan snapshot rows (rule gone, row still there) are dropped only by -an **active** generation: `Start` applies matching keys and ignores the -rest; the next flush replace-all omits them. Construction must not -delete snapshot rows — see §10. +Orphan snapshot rows (rule gone, row still there) are inert: `Start` +applies matching keys and ignores the rest, and they are collected two +periods after the last write. Construction must not delete snapshot +rows — see §10. ## 9. Configuration @@ -274,9 +265,8 @@ that concurrency is in-memory only. `ReplaceConfigRules` runs from `factory.New` → `seedConfiguredRules` while the previous generation is still serving. Deleting snapshot rows there would wipe a live hour/day window if the replacement is then -discarded (failed reload). Orphans are left for the **active** -generation: `Start` does not apply them; the next flush replace-all -drops them. +discarded (failed reload). Orphans are harmless instead: `Start` does +not apply them, and staleness collection removes them. Do not add a second prune path. `Refresh` already drops limiter keys when a definition is removed or shared/per-child mode changes; a later @@ -289,9 +279,10 @@ definition is not per-child. A mode flip therefore starts empty for that definition, matching `Refresh`. Reset and delete clear memory first (the operator-visible effect), then -the row under `persistMu`. A failed row delete is logged; the in-memory -reset still stands. `persistMu` is what stops the next flush from -putting the row back. +the row under `persistMu`. A failed row delete is returned to the +caller: the in-memory reset stands, but the operator has to know the +window can come back on the next restart. `persistMu` is what stops the +next flush from putting the row back. ## 11. Error handling @@ -305,8 +296,9 @@ Persistence never fails a request. `admit()` does not see store errors. stays authoritative. - **Shutdown flush failure:** log and continue teardown. Do not hang past the existing shutdown budget. -- **Reset/delete row failure:** log at error level. Memory is already - cleared. +- **Reset/delete row failure:** returned to the admin caller. Memory is + already cleared, so the reset took effect for this process; the error + says the durable row may outlive it. - **Migrate:** creating the new table/collection fails store init the same way a missing `rate_limits` table would — that is a hard start error, not a soft persist error. @@ -331,9 +323,10 @@ Persistence never fails a request. `admit()` does not see store errors. omit that key from `SaveCounters`. - An in-flight flush cannot resurrect a completed `ResetRule` / `DeleteRule` (`persistMu`). -- SQL store (and Mongo, same as rules) round-trip: replace, load, - delete one, delete all. Migration creates `rate_limit_counters` on an - existing DB. +- SQL store (and Mongo, same as rules) round-trip: save, load, save a + subset (the omitted row survives — it is still restorable), collect a + row left unwritten for two periods, delete one, delete all. Migration + creates `rate_limit_counters` on an existing DB. - Config: default interval is 1; `RATE_LIMITS_FLUSH_INTERVAL=0` is valid; a negative value is rejected. - Existing admit/release/header tests stay in-memory. A recording @@ -405,9 +398,9 @@ plus unit tests; concurrent is not persisted and is not an E2E case. ## 14. Follow-up (not this change) -Redis-Lua `counterBackend`: every `Admit` is an atomic script on shared +Redis live counters: every `Admit` becomes an atomic script on shared keys (the key must include `partition`, same as `ruleKey`). No snapshots. Chosen when `REDIS_URL` is set, or behind an explicit config once someone is running HA and wants N replicas to share one limit. -The interface in §5 is the only preparation this change makes for that -work. +That work extracts whatever seam it needs from the concrete `limiter`; +this change deliberately leaves none behind (§5). diff --git a/internal/ratelimit/persist.go b/internal/ratelimit/persist.go index af6eac200..481b88b5a 100644 --- a/internal/ratelimit/persist.go +++ b/internal/ratelimit/persist.go @@ -51,6 +51,13 @@ func persistContext(parent context.Context) (context.Context, context.CancelFunc return context.WithTimeout(context.Background(), persistTimeout) } +// Start restores persisted windows and, from then on, keeps writing them. +// It is separate from NewService because a reload builds the next generation +// while the current one still serves: only the generation that gets to serve +// may touch the snapshot, or a replacement that is built and then discarded +// would flush its empty windows over the live ones. Start is idempotent, and +// a failed load leaves the generation idle rather than persisting from a +// blank slate. func (s *Service) Start(ctx context.Context) { if s == nil || s.store == nil { return @@ -81,6 +88,9 @@ func (s *Service) Start(ctx context.Context) { s.limiter.restore(snapshots, s.Rules(), time.Now().UTC()) s.startFlushLoop() s.persistState = persistActive + if len(snapshots) > 0 { + slog.Info("rate limit counters restored", "windows", len(snapshots)) + } } func (s *Service) startFlushLoop() { diff --git a/internal/ratelimit/service_test.go b/internal/ratelimit/service_test.go index a3591bc24..d49ae5d16 100644 --- a/internal/ratelimit/service_test.go +++ b/internal/ratelimit/service_test.go @@ -73,8 +73,36 @@ func (m *memStore) LoadCounters(context.Context) ([]WindowSnapshot, error) { return append([]WindowSnapshot(nil), m.counters...), nil } +// snapshotIdentity is the primary key every store keys a window row by. +func snapshotIdentity(s WindowSnapshot) string { + return s.Scope + "\x00" + s.Subject + "\x00" + s.Partition + "\x00" + strconv.FormatInt(s.PeriodSeconds, 10) +} + +// SaveCounters mirrors the real stores: upsert by identity, then collect rows +// that went two periods without a write. func (m *memStore) SaveCounters(_ context.Context, snapshots []WindowSnapshot) error { - m.counters = append([]WindowSnapshot(nil), snapshots...) + now := time.Now().Unix() + for _, snap := range snapshots { + snap.UpdatedAt = now + replaced := false + for i, existing := range m.counters { + if snapshotIdentity(existing) == snapshotIdentity(snap) { + m.counters[i] = snap + replaced = true + break + } + } + if !replaced { + m.counters = append(m.counters, snap) + } + } + kept := m.counters[:0] + for _, snap := range m.counters { + if snap.UpdatedAt+2*snap.PeriodSeconds >= now { + kept = append(kept, snap) + } + } + m.counters = kept return nil } diff --git a/internal/ratelimit/snapshot.go b/internal/ratelimit/snapshot.go index f5011cd81..b97d38aa1 100644 --- a/internal/ratelimit/snapshot.go +++ b/internal/ratelimit/snapshot.go @@ -1,9 +1,6 @@ package ratelimit -import ( - "strconv" - "time" -) +import "time" // WindowSnapshot is one persisted request/token sliding window. // Partition is empty for a shared rule and the child path for a per-child @@ -19,17 +16,15 @@ type WindowSnapshot struct { TokensWindowStart int64 `bson:"tokens_window_start"` TokensCurrent int64 `bson:"tokens_current"` TokensPrevious int64 `bson:"tokens_previous"` - UpdatedAt int64 `bson:"updated_at,omitempty"` + // UpdatedAt is stamped by the store on write and drives its staleness + // collection. Loads do not populate it; restore has no use for it. + UpdatedAt int64 `bson:"updated_at,omitempty"` } func definitionKey(scope RuleScope, subject string, periodSeconds int64) ruleKey { return ruleKey{scope: scope, subject: subject, periodSeconds: periodSeconds} } -func snapshotIdentity(s WindowSnapshot) string { - return s.Scope + "\x00" + s.Subject + "\x00" + s.Partition + "\x00" + strconv.FormatInt(s.PeriodSeconds, 10) -} - func (l *limiter) snapshot(rules []Rule) []WindowSnapshot { l.mu.Lock() defer l.mu.Unlock() diff --git a/internal/ratelimit/store.go b/internal/ratelimit/store.go index 25c6fbec8..396bc9351 100644 --- a/internal/ratelimit/store.go +++ b/internal/ratelimit/store.go @@ -16,9 +16,11 @@ type Store interface { DeleteRule(ctx context.Context, scope RuleScope, subject string, periodSeconds int64) error ReplaceConfigRules(ctx context.Context, rules []Rule) error LoadCounters(ctx context.Context) ([]WindowSnapshot, error) - // SaveCounters upserts snapshots and deletes rows that are no longer - // present. It must not delete existing rows before the new values are - // durable, so a crash cannot wipe the previous generation. + // SaveCounters upserts the given snapshots and collects rows no writer has + // refreshed for two of their own periods — the same staleness bound a load + // applies, so it only ever drops rows that could no longer be restored. + // Nothing is deleted to make room for a write, so a crash mid-save costs at + // most the flush interval rather than a whole window. SaveCounters(ctx context.Context, snapshots []WindowSnapshot) error DeleteCounter(ctx context.Context, scope RuleScope, subject string, periodSeconds int64) error DeleteAllCounters(ctx context.Context) error diff --git a/internal/ratelimit/store_mongodb.go b/internal/ratelimit/store_mongodb.go index af9794d50..b248d8c23 100644 --- a/internal/ratelimit/store_mongodb.go +++ b/internal/ratelimit/store_mongodb.go @@ -378,67 +378,31 @@ func (s *MongoDBStore) LoadCounters(ctx context.Context) ([]WindowSnapshot, erro } func (s *MongoDBStore) SaveCounters(ctx context.Context, snapshots []WindowSnapshot) error { - write := func(writeCtx context.Context) error { - now := time.Now().Unix() + now := time.Now().Unix() + if len(snapshots) > 0 { + writes := make([]mongo.WriteModel, 0, len(snapshots)) for _, snap := range snapshots { snap.UpdatedAt = now - _, err := s.counters.UpdateOne( - writeCtx, - counterIdentityFilter(snap), - bson.D{{Key: "$set", Value: snap}}, - options.UpdateOne().SetUpsert(true), - ) - if err != nil { - return fmt.Errorf("upsert rate limit counter: %w", err) - } + writes = append(writes, mongo.NewUpdateOneModel(). + SetFilter(counterIdentityFilter(snap)). + SetUpdate(bson.D{{Key: "$set", Value: snap}}). + SetUpsert(true)) } - return s.deleteOrphanCounters(writeCtx, snapshots) - } - - session, err := s.counters.Database().Client().StartSession() - if err != nil { - return fmt.Errorf("start rate limit counter transaction: %w", err) - } - defer session.EndSession(ctx) - - _, err = session.WithTransaction(ctx, func(txCtx context.Context) (any, error) { - if err := write(txCtx); err != nil { - if isMongoTransactionCapabilityError(err) { - return nil, &mongoTransactionFallbackError{err: err} - } - return nil, err + if _, err := s.counters.BulkWrite(ctx, writes, options.BulkWrite().SetOrdered(false)); err != nil { + return fmt.Errorf("upsert rate limit counters: %w", err) } - return nil, nil - }) - if err != nil { - if fallbackErr := mongoTransactionFallbackCause(err); fallbackErr != nil || isMongoTransactionCapabilityError(err) { - if fallbackErr == nil { - fallbackErr = err - } - slog.Warn("MongoDB transactions unavailable for rate limit counters; falling back to non-transactional update", "error", fallbackErr) - if err := write(ctx); err != nil { - return fmt.Errorf("save rate limit counters without transaction: %w", errors.Join(fallbackErr, err)) - } - return nil - } - return fmt.Errorf("save rate limit counters transaction: %w", err) - } - return nil -} - -func (s *MongoDBStore) deleteOrphanCounters(ctx context.Context, snapshots []WindowSnapshot) error { - if len(snapshots) == 0 { - if _, err := s.counters.DeleteMany(ctx, bson.D{}); err != nil { - return fmt.Errorf("clear rate limit counters: %w", err) - } - return nil - } - keep := make(bson.A, 0, len(snapshots)) - for _, snap := range snapshots { - keep = append(keep, counterIdentityFilter(snap)) } - if _, err := s.counters.DeleteMany(ctx, bson.D{{Key: "$nor", Value: keep}}); err != nil { - return fmt.Errorf("prune rate limit counters: %w", err) + // Collect the rows nobody writes any more, on the same staleness bound + // restore applies, so this only ever deletes rows a load would discard. + stale := bson.D{{Key: "$expr", Value: bson.D{{Key: "$lt", Value: bson.A{ + bson.D{{Key: "$add", Value: bson.A{ + "$updated_at", + bson.D{{Key: "$multiply", Value: bson.A{2, "$period_seconds"}}}, + }}}, + now, + }}}}} + if _, err := s.counters.DeleteMany(ctx, stale); err != nil { + return fmt.Errorf("prune expired rate limit counters: %w", err) } return nil } diff --git a/internal/ratelimit/store_mongodb_test.go b/internal/ratelimit/store_mongodb_test.go index a0892f911..cb394987b 100644 --- a/internal/ratelimit/store_mongodb_test.go +++ b/internal/ratelimit/store_mongodb_test.go @@ -224,3 +224,88 @@ func TestMongoDBStoreMigratesPreScopeDocuments(t *testing.T) { } }) } + +// TestMongoDBStoreCounterRoundTrip mirrors TestSQLStoreCounterRoundTrip: the +// same window survives a save, a partial save leaves an omitted row alone, and +// a row nobody writes for two periods is collected. +func TestMongoDBStoreCounterRoundTrip(t *testing.T) { + mongotest.Run(t, func(t *testing.T, db *mongo.Database) { + ctx := context.Background() + store, err := NewMongoDBStore(ctx, db) + if err != nil { + t.Fatalf("NewMongoDBStore() failed: %v", err) + } + + live := []WindowSnapshot{ + { + Scope: string(ScopeUserPath), Subject: "/customers", Partition: "/customers/alice", + PeriodSeconds: PeriodHourSeconds, + RequestsWindowStart: 1700000000, RequestsCurrent: 3, RequestsPrevious: 1, + TokensWindowStart: 1700000000, TokensCurrent: 40, TokensPrevious: 10, + }, + { + Scope: string(ScopeUserPath), Subject: "/customers", Partition: "/customers/bob", + PeriodSeconds: PeriodHourSeconds, + RequestsWindowStart: 1700000060, RequestsCurrent: 1, RequestsPrevious: 2, + }, + } + if err := store.SaveCounters(ctx, live); err != nil { + t.Fatalf("SaveCounters: %v", err) + } + got, err := store.LoadCounters(ctx) + if err != nil { + t.Fatalf("LoadCounters: %v", err) + } + if len(got) != 2 { + t.Fatalf("loaded = %+v, want both partitions", got) + } + byPartition := map[string]WindowSnapshot{} + for _, snap := range got { + byPartition[snap.Partition] = snap + } + alice := byPartition["/customers/alice"] + if alice.RequestsWindowStart != 1700000000 || alice.RequestsCurrent != 3 || alice.RequestsPrevious != 1 || + alice.TokensWindowStart != 1700000000 || alice.TokensCurrent != 40 || alice.TokensPrevious != 10 { + t.Fatalf("alice = %+v", alice) + } + + // Omitting bob does not delete bob: the row is still restorable. + if err := store.SaveCounters(ctx, live[:1]); err != nil { + t.Fatalf("SaveCounters partial: %v", err) + } + if got, err = store.LoadCounters(ctx); err != nil || len(got) != 2 { + t.Fatalf("after partial save = %+v (err %v), want both rows", got, err) + } + + // A row left unwritten for two of its periods is collected. + if _, err := db.Collection("rate_limit_counters").InsertOne(ctx, WindowSnapshot{ + Scope: string(ScopeUserPath), Subject: "/gone", PeriodSeconds: PeriodHourSeconds, + RequestsCurrent: 4, UpdatedAt: time.Now().Unix() - 3*PeriodHourSeconds, + }); err != nil { + t.Fatalf("seed stale counter: %v", err) + } + if err := store.SaveCounters(ctx, live); err != nil { + t.Fatalf("SaveCounters collecting: %v", err) + } + if got, err = store.LoadCounters(ctx); err != nil || len(got) != 2 { + t.Fatalf("after collection = %+v (err %v), want the two live rows", got, err) + } + + if err := store.DeleteCounter(ctx, ScopeUserPath, "/customers", PeriodHourSeconds); err != nil { + t.Fatalf("DeleteCounter: %v", err) + } + if got, err = store.LoadCounters(ctx); err != nil || len(got) != 0 { + t.Fatalf("after delete = %+v (err %v), want every partition gone", got, err) + } + + if err := store.SaveCounters(ctx, live); err != nil { + t.Fatalf("SaveCounters again: %v", err) + } + if err := store.DeleteAllCounters(ctx); err != nil { + t.Fatalf("DeleteAllCounters: %v", err) + } + if got, err = store.LoadCounters(ctx); err != nil || len(got) != 0 { + t.Fatalf("after delete all = %+v (err %v), want empty", got, err) + } + }) +} diff --git a/internal/ratelimit/store_sql.go b/internal/ratelimit/store_sql.go index 7de964e8d..387458acd 100644 --- a/internal/ratelimit/store_sql.go +++ b/internal/ratelimit/store_sql.go @@ -229,54 +229,22 @@ func (s *SQLStore) SaveCounters(ctx context.Context, snapshots []WindowSnapshot) return fmt.Errorf("upsert rate limit counter: %w", err) } } - return deleteOrphanCounters(ctx, q, snapshots) - }) -} - -func deleteOrphanCounters(ctx context.Context, q sqlx.Querier, snapshots []WindowSnapshot) error { - if len(snapshots) == 0 { - if _, err := q.Exec(ctx, `DELETE FROM rate_limit_counters`); err != nil { - return fmt.Errorf("clear rate limit counters: %w", err) + if _, err := q.Exec(ctx, pruneCountersSQL, now); err != nil { + return fmt.Errorf("prune expired rate limit counters: %w", err) } return nil - } - keep := make(map[string]struct{}, len(snapshots)) - for _, snap := range snapshots { - keep[snapshotIdentity(snap)] = struct{}{} - } - rows, err := q.Query(ctx, ` - SELECT scope, subject, partition, period_seconds - FROM rate_limit_counters - `) - if err != nil { - return fmt.Errorf("list rate limit counters for prune: %w", err) - } - defer rows.Close() - - var stale []WindowSnapshot - for rows.Next() { - var snap WindowSnapshot - if err := rows.Scan(&snap.Scope, &snap.Subject, &snap.Partition, &snap.PeriodSeconds); err != nil { - return fmt.Errorf("scan rate limit counter identity: %w", err) - } - if _, ok := keep[snapshotIdentity(snap)]; !ok { - stale = append(stale, snap) - } - } - if err := rows.Err(); err != nil { - return fmt.Errorf("iterate rate limit counters for prune: %w", err) - } - for _, snap := range stale { - if _, err := q.Exec(ctx, ` - DELETE FROM rate_limit_counters - WHERE scope = ? AND subject = ? AND partition = ? AND period_seconds = ? - `, snap.Scope, snap.Subject, snap.Partition, snap.PeriodSeconds); err != nil { - return fmt.Errorf("prune rate limit counter: %w", err) - } - } - return nil + }) } +// pruneCountersSQL collects the rows nobody writes any more: a rule that was +// deleted, or a window that fell out of memory. Two periods without a write is +// the same staleness bound restore applies, so this only ever deletes rows a +// load would have discarded. +const pruneCountersSQL = ` + DELETE FROM rate_limit_counters + WHERE updated_at + 2 * period_seconds < ? +` + func (s *SQLStore) DeleteCounter(ctx context.Context, scope RuleScope, subject string, periodSeconds int64) error { _, err := s.db.Exec(ctx, ` DELETE FROM rate_limit_counters diff --git a/internal/ratelimit/store_sql_test.go b/internal/ratelimit/store_sql_test.go index e2c203f83..690e68a8b 100644 --- a/internal/ratelimit/store_sql_test.go +++ b/internal/ratelimit/store_sql_test.go @@ -3,8 +3,8 @@ package ratelimit import ( "context" "errors" - "strconv" "testing" + "time" "github.com/enterpilot/gomodel/internal/storage/sqlx" "github.com/enterpilot/gomodel/internal/storage/sqlx/sqlxtest" @@ -305,40 +305,39 @@ func TestSQLStoreCounterRoundTrip(t *testing.T) { t.Fatalf("after delete = %d, want 0", len(got)) } + // A save that omits a row leaves it alone: it is still restorable + // until it goes two periods without a write. if err := store.SaveCounters(ctx, first); err != nil { t.Fatalf("SaveCounters again: %v", err) } if err := store.SaveCounters(ctx, first[:1]); err != nil { - t.Fatalf("SaveCounters replace: %v", err) + t.Fatalf("SaveCounters partial: %v", err) } got, err = store.LoadCounters(ctx) if err != nil { - t.Fatalf("LoadCounters after replace: %v", err) + t.Fatalf("LoadCounters after partial save: %v", err) } - if len(got) != 1 || got[0].Partition != "/customers/alice" { - t.Fatalf("replaced = %+v, want alice only", got) + if len(got) != 2 { + t.Fatalf("partial save dropped a live row: %+v", got) } - many := make([]WindowSnapshot, 0, 300) - for i := range 300 { - many = append(many, WindowSnapshot{ - Scope: string(ScopeUserPath), Subject: "/customers", - Partition: "/customers/" + strconv.Itoa(i), - PeriodSeconds: PeriodHourSeconds, RequestsCurrent: int64(i + 1), - }) - } - if err := store.SaveCounters(ctx, many); err != nil { - t.Fatalf("SaveCounters many: %v", err) + // A row nobody has written for two of its periods is collected by the + // next save. + if _, err := store.db.Exec(ctx, upsertCounterSQL, + string(ScopeUserPath), "/gone", "", PeriodHourSeconds, + 0, 4, 0, 0, 0, 0, time.Now().Unix()-3*PeriodHourSeconds, + ); err != nil { + t.Fatalf("seed stale counter: %v", err) } - if err := store.SaveCounters(ctx, many[:1]); err != nil { - t.Fatalf("SaveCounters prune many: %v", err) + if err := store.SaveCounters(ctx, first); err != nil { + t.Fatalf("SaveCounters collecting: %v", err) } got, err = store.LoadCounters(ctx) if err != nil { - t.Fatalf("LoadCounters after prune many: %v", err) + t.Fatalf("LoadCounters after collection: %v", err) } - if len(got) != 1 || got[0].RequestsCurrent != 1 { - t.Fatalf("after pruning 299 orphans = %+v, want the first snapshot", got) + if len(got) != 2 { + t.Fatalf("expired row not collected: %+v", got) } if err := store.DeleteAllCounters(ctx); err != nil { From fa8f7e5edf93fdd00aac741720a18c59cec71fcd Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Sun, 16 Aug 2026 17:48:10 +0200 Subject: [PATCH 09/12] test(ratelimit): share one counter suite across backends MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The SQLite and MongoDB counter tests asserted the same contract twice, in parallel prose. One suite now runs on every backend — the pattern mongotest was written for — with a seed hook for the one backend-specific step (aging a row past the collection horizon). Also drops updated_at from WindowSnapshot: it is the store's own write stamp, not part of the window, and only Mongo's loader ever populated it, so the two backends returned different values for the same row. Mongo now keeps it in a private document type, which is what let the shared suite compare loaded rows field by field. Folds TestAdmitDoesNotSave into TestCloseWithoutStartDoesNotWrite (the second already exercised the first) and tightens flushIntervalFromSeconds. --- internal/ratelimit/persist.go | 8 +- internal/ratelimit/persist_test.go | 33 ++----- internal/ratelimit/service_test.go | 9 +- internal/ratelimit/snapshot.go | 3 - internal/ratelimit/store_counters_test.go | 92 +++++++++++++++++++ internal/ratelimit/store_mongodb.go | 11 ++- internal/ratelimit/store_mongodb_test.go | 88 +++--------------- internal/ratelimit/store_sql_test.go | 107 ++-------------------- 8 files changed, 142 insertions(+), 209 deletions(-) create mode 100644 internal/ratelimit/store_counters_test.go diff --git a/internal/ratelimit/persist.go b/internal/ratelimit/persist.go index 481b88b5a..e821e4333 100644 --- a/internal/ratelimit/persist.go +++ b/internal/ratelimit/persist.go @@ -30,15 +30,13 @@ func WithFlushInterval(interval time.Duration) ServiceOption { } } +// flushIntervalFromSeconds converts the configured seconds, clamping an +// absurd value rather than letting it overflow into a negative duration. func flushIntervalFromSeconds(seconds int) time.Duration { if seconds <= 0 { return 0 } - maxSeconds := int(math.MaxInt64 / int64(time.Second)) - if seconds > maxSeconds { - return time.Duration(math.MaxInt64) - } - return time.Duration(seconds) * time.Second + return time.Duration(min(int64(seconds), int64(math.MaxInt64/time.Second))) * time.Second } func persistContext(parent context.Context) (context.Context, context.CancelFunc) { diff --git a/internal/ratelimit/persist_test.go b/internal/ratelimit/persist_test.go index b52547700..2f11abbd8 100644 --- a/internal/ratelimit/persist_test.go +++ b/internal/ratelimit/persist_test.go @@ -154,8 +154,12 @@ func TestStartLoadsAndCloseFlushes(t *testing.T) { } } -func TestCloseWithoutStartDoesNotWrite(t *testing.T) { - store := &recordingStore{memStore: memStore{}} +// TestServiceWithoutStartNeverWrites covers both halves of the "not this +// generation" rule: admission never touches storage, and a service that was +// built but never started writes nothing on the way out either — a discarded +// reload replacement must not flush its empty windows over the live ones. +func TestServiceWithoutStartNeverWrites(t *testing.T) { + store := &recordingStore{} if err := store.UpsertRules(context.Background(), []Rule{{ Subject: "/", PeriodSeconds: PeriodHourSeconds, MaxRequests: new(int64(5)), Source: SourceManual, }}); err != nil { @@ -168,9 +172,12 @@ func TestCloseWithoutStartDoesNotWrite(t *testing.T) { if _, err := service.Acquire(onPath("/"), time.Now().UTC()); err != nil { t.Fatalf("Acquire: %v", err) } + if store.saves.Load() != 0 { + t.Fatalf("Acquire saved %d times, want 0", store.saves.Load()) + } service.Close() if store.saves.Load() != 0 { - t.Fatalf("saves = %d, want 0", store.saves.Load()) + t.Fatalf("Close without Start saved %d times, want 0", store.saves.Load()) } } @@ -208,26 +215,6 @@ func TestResetClearsPersistedWindow(t *testing.T) { } } -func TestAdmitDoesNotSave(t *testing.T) { - store := &recordingStore{memStore: memStore{}} - if err := store.UpsertRules(context.Background(), []Rule{{ - Subject: "/", PeriodSeconds: PeriodHourSeconds, MaxRequests: new(int64(5)), Source: SourceManual, - }}); err != nil { - t.Fatalf("seed: %v", err) - } - service, err := NewService(context.Background(), store) - if err != nil { - t.Fatalf("NewService: %v", err) - } - t.Cleanup(service.Close) - if _, err := service.Acquire(onPath("/"), time.Now().UTC()); err != nil { - t.Fatalf("Acquire: %v", err) - } - if store.saves.Load() != 0 { - t.Fatalf("Admit saved %d times", store.saves.Load()) - } -} - func TestRestoreIgnoresSharedRowOnPerChildRule(t *testing.T) { now := time.Unix(1_700_000_000, 0).UTC() template := Rule{ diff --git a/internal/ratelimit/service_test.go b/internal/ratelimit/service_test.go index d49ae5d16..65f7d2e16 100644 --- a/internal/ratelimit/service_test.go +++ b/internal/ratelimit/service_test.go @@ -15,6 +15,8 @@ import ( type memStore struct { rules []Rule counters []WindowSnapshot + // written stamps each row the way a real store's updated_at column does. + written map[string]int64 } func (m *memStore) ListRules(context.Context) ([]Rule, error) { @@ -82,8 +84,11 @@ func snapshotIdentity(s WindowSnapshot) string { // that went two periods without a write. func (m *memStore) SaveCounters(_ context.Context, snapshots []WindowSnapshot) error { now := time.Now().Unix() + if m.written == nil { + m.written = make(map[string]int64) + } for _, snap := range snapshots { - snap.UpdatedAt = now + m.written[snapshotIdentity(snap)] = now replaced := false for i, existing := range m.counters { if snapshotIdentity(existing) == snapshotIdentity(snap) { @@ -98,7 +103,7 @@ func (m *memStore) SaveCounters(_ context.Context, snapshots []WindowSnapshot) e } kept := m.counters[:0] for _, snap := range m.counters { - if snap.UpdatedAt+2*snap.PeriodSeconds >= now { + if m.written[snapshotIdentity(snap)]+2*snap.PeriodSeconds >= now { kept = append(kept, snap) } } diff --git a/internal/ratelimit/snapshot.go b/internal/ratelimit/snapshot.go index b97d38aa1..cbcba56da 100644 --- a/internal/ratelimit/snapshot.go +++ b/internal/ratelimit/snapshot.go @@ -16,9 +16,6 @@ type WindowSnapshot struct { TokensWindowStart int64 `bson:"tokens_window_start"` TokensCurrent int64 `bson:"tokens_current"` TokensPrevious int64 `bson:"tokens_previous"` - // UpdatedAt is stamped by the store on write and drives its staleness - // collection. Loads do not populate it; restore has no use for it. - UpdatedAt int64 `bson:"updated_at,omitempty"` } func definitionKey(scope RuleScope, subject string, periodSeconds int64) ruleKey { diff --git a/internal/ratelimit/store_counters_test.go b/internal/ratelimit/store_counters_test.go new file mode 100644 index 000000000..61687c955 --- /dev/null +++ b/internal/ratelimit/store_counters_test.go @@ -0,0 +1,92 @@ +package ratelimit + +import ( + "context" + "testing" + "time" +) + +// runCounterStoreSuite asserts the counter half of the Store contract. Every +// backend runs the same body, so a window that round-trips on SQLite has to +// round-trip on PostgreSQL and MongoDB too. seedStale writes a row directly, +// bypassing SaveCounters, so the suite can age one without waiting an hour. +func runCounterStoreSuite(t *testing.T, store Store, seedStale func(t *testing.T, snap WindowSnapshot, updatedAt int64)) { + t.Helper() + ctx := context.Background() + live := []WindowSnapshot{ + { + Scope: string(ScopeUserPath), Subject: "/customers", Partition: "/customers/alice", + PeriodSeconds: PeriodHourSeconds, + RequestsWindowStart: 1700000000, RequestsCurrent: 3, RequestsPrevious: 1, + TokensWindowStart: 1700000000, TokensCurrent: 40, TokensPrevious: 10, + }, + { + Scope: string(ScopeUserPath), Subject: "/customers", Partition: "/customers/bob", + PeriodSeconds: PeriodHourSeconds, + RequestsWindowStart: 1700000060, RequestsCurrent: 1, RequestsPrevious: 2, + TokensWindowStart: 1700000060, TokensCurrent: 7, TokensPrevious: 8, + }, + } + loaded := func(t *testing.T, step string) map[string]WindowSnapshot { + t.Helper() + got, err := store.LoadCounters(ctx) + if err != nil { + t.Fatalf("LoadCounters %s: %v", step, err) + } + byPartition := make(map[string]WindowSnapshot, len(got)) + for _, snap := range got { + byPartition[snap.Partition] = snap + } + return byPartition + } + + // Every field survives the round trip, per partition. + if err := store.SaveCounters(ctx, live); err != nil { + t.Fatalf("SaveCounters: %v", err) + } + got := loaded(t, "after save") + for _, want := range live { + if got[want.Partition] != want { + t.Fatalf("%s = %+v, want %+v", want.Partition, got[want.Partition], want) + } + } + + // Omitting a row does not delete it: it stays restorable until it goes + // two periods without a write. + if err := store.SaveCounters(ctx, live[:1]); err != nil { + t.Fatalf("SaveCounters partial: %v", err) + } + if got = loaded(t, "after partial save"); len(got) != 2 { + t.Fatalf("partial save dropped a live row: %+v", got) + } + + // Two periods without a write makes it collectable by the next save. + seedStale(t, WindowSnapshot{ + Scope: string(ScopeUserPath), Subject: "/gone", PeriodSeconds: PeriodHourSeconds, + RequestsCurrent: 4, + }, time.Now().Unix()-3*PeriodHourSeconds) + if err := store.SaveCounters(ctx, live); err != nil { + t.Fatalf("SaveCounters collecting: %v", err) + } + if got = loaded(t, "after collection"); len(got) != 2 { + t.Fatalf("stale row not collected: %+v", got) + } + + // A reset clears every partition of the definition, not just one. + if err := store.DeleteCounter(ctx, ScopeUserPath, "/customers", PeriodHourSeconds); err != nil { + t.Fatalf("DeleteCounter: %v", err) + } + if got = loaded(t, "after delete"); len(got) != 0 { + t.Fatalf("after delete = %+v, want every partition gone", got) + } + + if err := store.SaveCounters(ctx, live); err != nil { + t.Fatalf("SaveCounters again: %v", err) + } + if err := store.DeleteAllCounters(ctx); err != nil { + t.Fatalf("DeleteAllCounters: %v", err) + } + if got = loaded(t, "after delete all"); len(got) != 0 { + t.Fatalf("after delete all = %+v, want empty", got) + } +} diff --git a/internal/ratelimit/store_mongodb.go b/internal/ratelimit/store_mongodb.go index b248d8c23..e5629ac95 100644 --- a/internal/ratelimit/store_mongodb.go +++ b/internal/ratelimit/store_mongodb.go @@ -377,15 +377,22 @@ func (s *MongoDBStore) LoadCounters(ctx context.Context) ([]WindowSnapshot, erro return snapshots, nil } +// counterDocument is a window row as stored: the snapshot plus the write +// stamp that drives staleness collection. The stamp is the store's own +// bookkeeping, so it stays out of WindowSnapshot. +type counterDocument struct { + WindowSnapshot `bson:",inline"` + UpdatedAt int64 `bson:"updated_at"` +} + func (s *MongoDBStore) SaveCounters(ctx context.Context, snapshots []WindowSnapshot) error { now := time.Now().Unix() if len(snapshots) > 0 { writes := make([]mongo.WriteModel, 0, len(snapshots)) for _, snap := range snapshots { - snap.UpdatedAt = now writes = append(writes, mongo.NewUpdateOneModel(). SetFilter(counterIdentityFilter(snap)). - SetUpdate(bson.D{{Key: "$set", Value: snap}}). + SetUpdate(bson.D{{Key: "$set", Value: counterDocument{WindowSnapshot: snap, UpdatedAt: now}}}). SetUpsert(true)) } if _, err := s.counters.BulkWrite(ctx, writes, options.BulkWrite().SetOrdered(false)); err != nil { diff --git a/internal/ratelimit/store_mongodb_test.go b/internal/ratelimit/store_mongodb_test.go index cb394987b..3ded7dab9 100644 --- a/internal/ratelimit/store_mongodb_test.go +++ b/internal/ratelimit/store_mongodb_test.go @@ -225,87 +225,21 @@ func TestMongoDBStoreMigratesPreScopeDocuments(t *testing.T) { }) } -// TestMongoDBStoreCounterRoundTrip mirrors TestSQLStoreCounterRoundTrip: the -// same window survives a save, a partial save leaves an omitted row alone, and -// a row nobody writes for two periods is collected. +// TestMongoDBStoreCounterRoundTrip runs the shared counter suite: MongoDB has +// its own upsert and staleness-collection queries, so it has to prove the same +// behaviour the SQL backends do. func TestMongoDBStoreCounterRoundTrip(t *testing.T) { mongotest.Run(t, func(t *testing.T, db *mongo.Database) { - ctx := context.Background() - store, err := NewMongoDBStore(ctx, db) + store, err := NewMongoDBStore(context.Background(), db) if err != nil { t.Fatalf("NewMongoDBStore() failed: %v", err) } - - live := []WindowSnapshot{ - { - Scope: string(ScopeUserPath), Subject: "/customers", Partition: "/customers/alice", - PeriodSeconds: PeriodHourSeconds, - RequestsWindowStart: 1700000000, RequestsCurrent: 3, RequestsPrevious: 1, - TokensWindowStart: 1700000000, TokensCurrent: 40, TokensPrevious: 10, - }, - { - Scope: string(ScopeUserPath), Subject: "/customers", Partition: "/customers/bob", - PeriodSeconds: PeriodHourSeconds, - RequestsWindowStart: 1700000060, RequestsCurrent: 1, RequestsPrevious: 2, - }, - } - if err := store.SaveCounters(ctx, live); err != nil { - t.Fatalf("SaveCounters: %v", err) - } - got, err := store.LoadCounters(ctx) - if err != nil { - t.Fatalf("LoadCounters: %v", err) - } - if len(got) != 2 { - t.Fatalf("loaded = %+v, want both partitions", got) - } - byPartition := map[string]WindowSnapshot{} - for _, snap := range got { - byPartition[snap.Partition] = snap - } - alice := byPartition["/customers/alice"] - if alice.RequestsWindowStart != 1700000000 || alice.RequestsCurrent != 3 || alice.RequestsPrevious != 1 || - alice.TokensWindowStart != 1700000000 || alice.TokensCurrent != 40 || alice.TokensPrevious != 10 { - t.Fatalf("alice = %+v", alice) - } - - // Omitting bob does not delete bob: the row is still restorable. - if err := store.SaveCounters(ctx, live[:1]); err != nil { - t.Fatalf("SaveCounters partial: %v", err) - } - if got, err = store.LoadCounters(ctx); err != nil || len(got) != 2 { - t.Fatalf("after partial save = %+v (err %v), want both rows", got, err) - } - - // A row left unwritten for two of its periods is collected. - if _, err := db.Collection("rate_limit_counters").InsertOne(ctx, WindowSnapshot{ - Scope: string(ScopeUserPath), Subject: "/gone", PeriodSeconds: PeriodHourSeconds, - RequestsCurrent: 4, UpdatedAt: time.Now().Unix() - 3*PeriodHourSeconds, - }); err != nil { - t.Fatalf("seed stale counter: %v", err) - } - if err := store.SaveCounters(ctx, live); err != nil { - t.Fatalf("SaveCounters collecting: %v", err) - } - if got, err = store.LoadCounters(ctx); err != nil || len(got) != 2 { - t.Fatalf("after collection = %+v (err %v), want the two live rows", got, err) - } - - if err := store.DeleteCounter(ctx, ScopeUserPath, "/customers", PeriodHourSeconds); err != nil { - t.Fatalf("DeleteCounter: %v", err) - } - if got, err = store.LoadCounters(ctx); err != nil || len(got) != 0 { - t.Fatalf("after delete = %+v (err %v), want every partition gone", got, err) - } - - if err := store.SaveCounters(ctx, live); err != nil { - t.Fatalf("SaveCounters again: %v", err) - } - if err := store.DeleteAllCounters(ctx); err != nil { - t.Fatalf("DeleteAllCounters: %v", err) - } - if got, err = store.LoadCounters(ctx); err != nil || len(got) != 0 { - t.Fatalf("after delete all = %+v (err %v), want empty", got, err) - } + runCounterStoreSuite(t, store, func(t *testing.T, snap WindowSnapshot, updatedAt int64) { + t.Helper() + doc := counterDocument{WindowSnapshot: snap, UpdatedAt: updatedAt} + if _, err := db.Collection("rate_limit_counters").InsertOne(context.Background(), doc); err != nil { + t.Fatalf("seed stale counter: %v", err) + } + }) }) } diff --git a/internal/ratelimit/store_sql_test.go b/internal/ratelimit/store_sql_test.go index 690e68a8b..adb7d0877 100644 --- a/internal/ratelimit/store_sql_test.go +++ b/internal/ratelimit/store_sql_test.go @@ -4,7 +4,6 @@ import ( "context" "errors" "testing" - "time" "github.com/enterpilot/gomodel/internal/storage/sqlx" "github.com/enterpilot/gomodel/internal/storage/sqlx/sqlxtest" @@ -254,101 +253,15 @@ func TestSQLStoreMigratesPreScopeTable(t *testing.T) { func TestSQLStoreCounterRoundTrip(t *testing.T) { runSQLStoreTest(t, func(t *testing.T, store *SQLStore) { - ctx := context.Background() - first := []WindowSnapshot{ - { - Scope: string(ScopeUserPath), Subject: "/customers", Partition: "/customers/alice", - PeriodSeconds: PeriodHourSeconds, - RequestsWindowStart: 1700000000, RequestsCurrent: 3, RequestsPrevious: 1, - TokensWindowStart: 1700000000, TokensCurrent: 40, TokensPrevious: 10, - }, - { - Scope: string(ScopeUserPath), Subject: "/customers", Partition: "/customers/bob", - PeriodSeconds: PeriodHourSeconds, - RequestsWindowStart: 1700000060, RequestsCurrent: 1, RequestsPrevious: 2, - TokensWindowStart: 1700000060, TokensCurrent: 7, TokensPrevious: 8, - }, - } - if err := store.SaveCounters(ctx, first); err != nil { - t.Fatalf("SaveCounters: %v", err) - } - got, err := store.LoadCounters(ctx) - if err != nil { - t.Fatalf("LoadCounters: %v", err) - } - if len(got) != 2 { - t.Fatalf("loaded = %d, want 2", len(got)) - } - byPart := map[string]WindowSnapshot{} - for _, snap := range got { - byPart[snap.Partition] = snap - } - alice := byPart["/customers/alice"] - if alice.RequestsWindowStart != 1700000000 || alice.RequestsCurrent != 3 || alice.RequestsPrevious != 1 || - alice.TokensWindowStart != 1700000000 || alice.TokensCurrent != 40 || alice.TokensPrevious != 10 { - t.Fatalf("alice = %+v", alice) - } - bob := byPart["/customers/bob"] - if bob.RequestsWindowStart != 1700000060 || bob.RequestsCurrent != 1 || bob.RequestsPrevious != 2 || - bob.TokensWindowStart != 1700000060 || bob.TokensCurrent != 7 || bob.TokensPrevious != 8 { - t.Fatalf("bob = %+v", bob) - } - - if err := store.DeleteCounter(ctx, ScopeUserPath, "/customers", PeriodHourSeconds); err != nil { - t.Fatalf("DeleteCounter: %v", err) - } - got, err = store.LoadCounters(ctx) - if err != nil { - t.Fatalf("LoadCounters after delete: %v", err) - } - if len(got) != 0 { - t.Fatalf("after delete = %d, want 0", len(got)) - } - - // A save that omits a row leaves it alone: it is still restorable - // until it goes two periods without a write. - if err := store.SaveCounters(ctx, first); err != nil { - t.Fatalf("SaveCounters again: %v", err) - } - if err := store.SaveCounters(ctx, first[:1]); err != nil { - t.Fatalf("SaveCounters partial: %v", err) - } - got, err = store.LoadCounters(ctx) - if err != nil { - t.Fatalf("LoadCounters after partial save: %v", err) - } - if len(got) != 2 { - t.Fatalf("partial save dropped a live row: %+v", got) - } - - // A row nobody has written for two of its periods is collected by the - // next save. - if _, err := store.db.Exec(ctx, upsertCounterSQL, - string(ScopeUserPath), "/gone", "", PeriodHourSeconds, - 0, 4, 0, 0, 0, 0, time.Now().Unix()-3*PeriodHourSeconds, - ); err != nil { - t.Fatalf("seed stale counter: %v", err) - } - if err := store.SaveCounters(ctx, first); err != nil { - t.Fatalf("SaveCounters collecting: %v", err) - } - got, err = store.LoadCounters(ctx) - if err != nil { - t.Fatalf("LoadCounters after collection: %v", err) - } - if len(got) != 2 { - t.Fatalf("expired row not collected: %+v", got) - } - - if err := store.DeleteAllCounters(ctx); err != nil { - t.Fatalf("DeleteAllCounters: %v", err) - } - got, err = store.LoadCounters(ctx) - if err != nil { - t.Fatalf("LoadCounters after delete all: %v", err) - } - if len(got) != 0 { - t.Fatalf("after delete all = %d, want 0", len(got)) - } + runCounterStoreSuite(t, store, func(t *testing.T, snap WindowSnapshot, updatedAt int64) { + t.Helper() + if _, err := store.db.Exec(context.Background(), upsertCounterSQL, + snap.Scope, snap.Subject, snap.Partition, snap.PeriodSeconds, + snap.RequestsWindowStart, snap.RequestsCurrent, snap.RequestsPrevious, + snap.TokensWindowStart, snap.TokensCurrent, snap.TokensPrevious, updatedAt, + ); err != nil { + t.Fatalf("seed stale counter: %v", err) + } + }) }) } From 204d2fc25c903e865417d4fb588b62c4e9006677 Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Sun, 16 Aug 2026 17:53:42 +0200 Subject: [PATCH 10/12] fix(ratelimit): skip malformed counter rows on SQL load The Mongo loader already skipped and logged an undecodable document, but the SQL loader failed the whole load on one bad row. Start treats a load error as "do not persist this generation", so a single unreadable row cost every other window both its restore and its persistence until someone deleted it. Both loaders now share the policy the spec states, with a regression test each (verified to fail without the skip): a query or iteration failure is still fatal, a row is not. --- internal/ratelimit/store_mongodb_test.go | 37 ++++++++++++++++++++++++ internal/ratelimit/store_sql.go | 7 ++++- internal/ratelimit/store_sql_test.go | 35 ++++++++++++++++++++++ 3 files changed, 78 insertions(+), 1 deletion(-) diff --git a/internal/ratelimit/store_mongodb_test.go b/internal/ratelimit/store_mongodb_test.go index 3ded7dab9..26bfe94f5 100644 --- a/internal/ratelimit/store_mongodb_test.go +++ b/internal/ratelimit/store_mongodb_test.go @@ -243,3 +243,40 @@ func TestMongoDBStoreCounterRoundTrip(t *testing.T) { }) }) } + +// TestMongoDBStoreLoadCountersSkipsMalformedDocument is the MongoDB half of +// TestSQLStoreLoadCountersSkipsMalformedRow: one undecodable document must not +// cost every other window its restore. +func TestMongoDBStoreLoadCountersSkipsMalformedDocument(t *testing.T) { + mongotest.Run(t, func(t *testing.T, db *mongo.Database) { + ctx := context.Background() + store, err := NewMongoDBStore(ctx, db) + if err != nil { + t.Fatalf("NewMongoDBStore() failed: %v", err) + } + good := WindowSnapshot{ + Scope: string(ScopeUserPath), Subject: "/team", PeriodSeconds: PeriodHourSeconds, + RequestsWindowStart: 1700000000, RequestsCurrent: 2, + } + if err := store.SaveCounters(ctx, []WindowSnapshot{good}); err != nil { + t.Fatalf("SaveCounters: %v", err) + } + if _, err := db.Collection("rate_limit_counters").InsertOne(ctx, bson.D{ + {Key: "scope", Value: string(ScopeUserPath)}, + {Key: "subject", Value: "/broken"}, + {Key: "partition", Value: ""}, + {Key: "period_seconds", Value: "not-a-number"}, + {Key: "updated_at", Value: time.Now().Unix()}, + }); err != nil { + t.Fatalf("seed malformed document: %v", err) + } + + got, err := store.LoadCounters(ctx) + if err != nil { + t.Fatalf("LoadCounters: %v", err) + } + if len(got) != 1 || got[0] != good { + t.Fatalf("loaded = %+v, want only the readable window %+v", got, good) + } + }) +} diff --git a/internal/ratelimit/store_sql.go b/internal/ratelimit/store_sql.go index 387458acd..51f26e68f 100644 --- a/internal/ratelimit/store_sql.go +++ b/internal/ratelimit/store_sql.go @@ -3,6 +3,7 @@ package ratelimit import ( "context" "fmt" + "log/slog" "strings" "time" @@ -191,7 +192,11 @@ func (s *SQLStore) LoadCounters(ctx context.Context) ([]WindowSnapshot, error) { &snap.RequestsWindowStart, &snap.RequestsCurrent, &snap.RequestsPrevious, &snap.TokensWindowStart, &snap.TokensCurrent, &snap.TokensPrevious, ); err != nil { - return nil, fmt.Errorf("scan rate limit counter: %w", err) + // One unreadable row must not cost every other window its + // restore: Start treats a load error as "do not persist this + // generation". Same policy as the Mongo loader. + slog.Warn("rate limit counters: skipping malformed snapshot", "error", err) + continue } snapshots = append(snapshots, snap) } diff --git a/internal/ratelimit/store_sql_test.go b/internal/ratelimit/store_sql_test.go index adb7d0877..1146add50 100644 --- a/internal/ratelimit/store_sql_test.go +++ b/internal/ratelimit/store_sql_test.go @@ -4,6 +4,7 @@ import ( "context" "errors" "testing" + "time" "github.com/enterpilot/gomodel/internal/storage/sqlx" "github.com/enterpilot/gomodel/internal/storage/sqlx/sqlxtest" @@ -265,3 +266,37 @@ func TestSQLStoreCounterRoundTrip(t *testing.T) { }) }) } + +// TestSQLStoreLoadCountersSkipsMalformedRow keeps one unreadable row from +// costing every other window its restore: Start treats a load error as "do not +// persist this generation". SQLite only — its dynamic typing is what lets a +// row hold a value the scan cannot read. +func TestSQLStoreLoadCountersSkipsMalformedRow(t *testing.T) { + ctx := context.Background() + db := sqlxtest.NewSQLite(t) + store, err := NewSQLStore(ctx, db) + if err != nil { + t.Fatalf("NewSQLStore: %v", err) + } + good := WindowSnapshot{ + Scope: string(ScopeUserPath), Subject: "/team", PeriodSeconds: PeriodHourSeconds, + RequestsWindowStart: 1700000000, RequestsCurrent: 2, + } + if err := store.SaveCounters(ctx, []WindowSnapshot{good}); err != nil { + t.Fatalf("SaveCounters: %v", err) + } + if _, err := db.Exec(ctx, upsertCounterSQL, + string(ScopeUserPath), "/broken", "", PeriodHourSeconds, + 0, "not-a-number", 0, 0, 0, 0, time.Now().Unix(), + ); err != nil { + t.Fatalf("seed malformed counter: %v", err) + } + + got, err := store.LoadCounters(ctx) + if err != nil { + t.Fatalf("LoadCounters: %v", err) + } + if len(got) != 1 || got[0] != good { + t.Fatalf("loaded = %+v, want only the readable window %+v", got, good) + } +} From 973488a0901ba3e4897f3d7de79736c3c1efb2e2 Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Sun, 16 Aug 2026 18:14:56 +0200 Subject: [PATCH 11/12] fix(ratelimit): refresh rules even when the snapshot row survives MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DeleteRule returned on a failed counter-row delete, skipping Refresh. The rule row was already gone from the store, so the service kept enforcing a deleted rule from its stale in-memory list — a worse outcome than the leftover snapshot row the error was about. Refresh now always runs and the row error is returned after it. Also takes persistMu before deriving the delete deadline: a reset queued behind a slow flush was spending its whole budget waiting for the lock. --- internal/ratelimit/persist.go | 10 ++++++---- internal/ratelimit/persist_test.go | 31 ++++++++++++++++++++++++++++++ internal/ratelimit/service.go | 9 +++++++-- 3 files changed, 44 insertions(+), 6 deletions(-) diff --git a/internal/ratelimit/persist.go b/internal/ratelimit/persist.go index e821e4333..70bd767d6 100644 --- a/internal/ratelimit/persist.go +++ b/internal/ratelimit/persist.go @@ -132,10 +132,12 @@ func (s *Service) persistDelete(ctx context.Context, scope RuleScope, subject st if s == nil || s.store == nil { return nil } - writeCtx, cancel := persistContext(ctx) - defer cancel() + // Take the lock before starting the clock: a delete that queued behind a + // slow flush would otherwise spend its whole budget waiting. s.persistMu.Lock() defer s.persistMu.Unlock() + writeCtx, cancel := persistContext(ctx) + defer cancel() return s.store.DeleteCounter(writeCtx, scope, subject, periodSeconds) } @@ -143,9 +145,9 @@ func (s *Service) persistDeleteAll(ctx context.Context) error { if s == nil || s.store == nil { return nil } - writeCtx, cancel := persistContext(ctx) - defer cancel() s.persistMu.Lock() defer s.persistMu.Unlock() + writeCtx, cancel := persistContext(ctx) + defer cancel() return s.store.DeleteAllCounters(writeCtx) } diff --git a/internal/ratelimit/persist_test.go b/internal/ratelimit/persist_test.go index 2f11abbd8..2ed46a420 100644 --- a/internal/ratelimit/persist_test.go +++ b/internal/ratelimit/persist_test.go @@ -423,3 +423,34 @@ func (s *blockingLoadStore) LoadCounters(context.Context) ([]WindowSnapshot, err <-s.release return s.memStore.LoadCounters(context.Background()) } + +// TestDeleteRuleStopsEnforcingWhenSnapshotDeleteFails: the rule row is already +// gone, so the in-memory refresh has to happen even when the snapshot row +// cannot be removed. The caller still hears about the row. +func TestDeleteRuleStopsEnforcingWhenSnapshotDeleteFails(t *testing.T) { + ctx := context.Background() + store := &failDeleteStore{err: errors.New("delete failed")} + if err := store.UpsertRules(ctx, []Rule{{ + Subject: "/team", PeriodSeconds: PeriodHourSeconds, MaxRequests: new(int64(1)), Source: SourceManual, + }}); err != nil { + t.Fatalf("seed: %v", err) + } + service, err := NewService(ctx, store) + if err != nil { + t.Fatalf("NewService: %v", err) + } + t.Cleanup(service.Close) + + if err := service.DeleteRule(ctx, ScopeUserPath, "/team", PeriodHourSeconds); err == nil { + t.Fatal("DeleteRule() error = nil, want the snapshot-row error") + } + if rules := service.Rules(); len(rules) != 0 { + t.Fatalf("rules = %+v, want the deleted rule gone from memory", rules) + } + // Two admissions: the deleted one-per-hour rule is no longer enforced. + for i := range 2 { + if _, err := service.Acquire(onPath("/team"), time.Now().UTC()); err != nil { + t.Fatalf("Acquire %d after delete: %v", i, err) + } + } +} diff --git a/internal/ratelimit/service.go b/internal/ratelimit/service.go index 2c148ff4a..6fea58062 100644 --- a/internal/ratelimit/service.go +++ b/internal/ratelimit/service.go @@ -171,10 +171,15 @@ func (s *Service) DeleteRule(ctx context.Context, scope RuleScope, subject strin return err } s.limiter.reset(ruleKey{scope: scope, subject: subject, periodSeconds: periodSeconds}) - if err := s.persistDelete(ctx, scope, subject, periodSeconds); err != nil { + // The rule is already gone from the store, so the refresh has to happen + // either way — dropping it on a snapshot-row failure would leave the + // deleted rule enforced in memory, which is worse than the stale row. + // The caller still hears about the row. + persistErr := s.persistDelete(ctx, scope, subject, periodSeconds) + if err := s.Refresh(ctx); err != nil { return err } - return s.Refresh(ctx) + return persistErr } func (s *Service) ReplaceConfigRules(ctx context.Context, rules []Rule) error { From f8372b820af1351021b6b6a559243b66127edd5c Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Sun, 16 Aug 2026 18:20:14 +0200 Subject: [PATCH 12/12] docs(ratelimit): trim the persistence spec to decisions The spec carried a per-test inventory, a restatement of each release scenario, and a "document it here, here and here" task list. The tests, the scenarios file and the docs themselves say all of that, and say it accurately as they change. Keeps what code cannot carry: why a generation persists only once it serves, why the reload helper has to wait for the log line, and what the suites are there to pin. 412 lines to 349. --- ...-16_rate-limit-counter-persistence-spec.md | 133 +++++------------- 1 file changed, 38 insertions(+), 95 deletions(-) diff --git a/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md b/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md index 2bf889db9..3c5581b97 100644 --- a/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md +++ b/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md @@ -246,12 +246,6 @@ RATE_LIMITS_FLUSH_INTERVAL=1 - Persistence is on whenever `RATE_LIMITS_ENABLED` is on. There is no second flag. -Document in `.env.template`, `config/config.example.yaml`, -`docs/features/rate-limits.mdx`, and `CLAUDE.md` / `Agents.md`. Remove -the wording that counters reset on restart / `--reload`. Keep the -wording that counters are per instance (N replicas ≈ N × limit) and -that concurrency is in-memory only. - ## 10. Service operations vs snapshot | Operation | Memory | Snapshot | @@ -305,98 +299,47 @@ Persistence never fails a request. `admit()` does not see store errors. ## 12. Testing -### Unit / store - -- Snapshot encode/restore: `estimate` after load matches the pre-flush - value for both request and token windows. -- A snapshot whose `windowStart` is more than one period old advances - to zero. -- Concurrent keys never appear in a snapshot. -- `New` does not read or write. `Start` loads. `Close` after `Start` - writes a final snapshot. `Close` without `Start` writes nothing. -- `flush_interval=0` still loads and still flushes on `Close`; the loop - never ticks. -- `ResetRule` / `ResetAll` / `DeleteRule` clear memory and the row; a - later `Start` on a new service does not resurrect them. -- `ReplaceConfigRules` dropping a config rule prunes the limiter maps - and does not call `DeleteCounter`. After `Start`, the next flush - omit that key from `SaveCounters`. -- An in-flight flush cannot resurrect a completed `ResetRule` / - `DeleteRule` (`persistMu`). -- SQL store (and Mongo, same as rules) round-trip: save, load, save a - subset (the omitted row survives — it is still restorable), collect a - row left unwritten for two periods, delete one, delete all. Migration - creates `rate_limit_counters` on an existing DB. -- Config: default interval is 1; `RATE_LIMITS_FLUSH_INTERVAL=0` is - valid; a negative value is rejected. -- Existing admit/release/header tests stay in-memory. A recording - `Store` asserts `Admit` does not call `SaveCounters`. -- Per-child: two children of one template flush as two rows (same - `subject`, different `partition`); after `Start` each child is still - isolated. Reset of the template deletes both rows. A restored - partition is registered with the expiry worker. A shared-rule row is - not applied to a definition that is now `PerChild`, and the reverse. -- OSS without `quota_templates` still persists shared / provider / - model windows. Per-child config continues to abort startup / reject - admin writes; persistence does not change that gate. +The suites live in `internal/ratelimit`; this section records only what +they are there to pin, not a per-test inventory. + +- `New` neither reads nor writes. `Start` loads. `Close` after `Start` + writes once. `Close` without `Start` writes nothing, and neither does + admission — the "only a serving generation persists" rule of §8. +- A failed load leaves the generation idle, so a store hiccup cannot + cost the persisted windows. +- `Close` during an in-flight load does not restore after shutdown. +- An in-flight flush cannot resurrect a completed reset or delete + (`persistMu`), and a failed row delete reaches the caller. +- Per-child partitions round-trip independently, re-arm the expiry + worker, and never cross-apply between shared and template modes. +- Concurrent keys are never written. +- One counter-store suite (`runCounterStoreSuite`) runs on SQLite, + PostgreSQL and MongoDB: field-exact round trip, an omitted row + surviving a save, a two-period-stale row being collected, and both + delete paths. A malformed row is skipped, not fatal — asserted per + backend, since each has its own decoder. ### Release E2E -Add to `tests/e2e/release-e2e-scenarios.md` (after S204). Update the -file header count and the stateful-note list. - -Shared helper in the common environment block: - -- Export `RELEASE_STACK_DIR` (default `/tmp/gomodel-release-stack`). -- `reload_release_gateway ` sends `SIGHUP` to - `$RELEASE_STACK_DIR//server.pid`, then waits until that - gateway’s `logs/server.log` contains a **new** `configuration reloaded` - line, then retries `/health` briefly. `configuration reloaded` is - logged after the old `Shutdown` and before the new - `StartWithListener`; the next request may sit in the held accept - queue until the new listener is up. A request sent immediately after - `kill -HUP` can still hit the old generation — the log wait is - required. - -Use **hour** windows so the cap outlives the reload wait. Each scenario -creates a `$QA_SUFFIX`-scoped **shared** user-path rule (`per_child` -unset) and deletes it. The release stack is OSS and has no -`quota_templates` entitlement; a per-child admin write would 403. -Per-child persistence is the unit tests above, not this matrix. - -- **S205 — Request-window counters survive `--reload` (SQLite).** - `max_requests=1` on `$BASE_URL`. First chat succeeds. Reload - `sqlite-main` (old `Close` writes the snapshot; no sleep required - for the happy path). Second chat is `429` with - `code: rate_limit_exceeded`. Delete the rule. Crash-before-`Close` - (periodic flush only) is a unit test, not this scenario. -- **S206 — `reset-one` stays cleared across `--reload`.** Same shape: - burn the hour window, `reset-one`, reload `sqlite-main`. Next chat - succeeds. Delete the rule. `persistMu` is what makes this - deterministic. -- **S207 — Same request-window survival on PostgreSQL and MongoDB.** - S203-style loop over `$PG_BASE_URL` / `$MONGO_BASE_URL`, reloading - `pg-smoke` and `mongo-smoke`. Distinct paths, delete each rule. - -S205–S207 reload a shared gateway. That is safe in this sequential -runner (same class as S137). Token-window reload is covered by S157 -plus unit tests; concurrent is not persisted and is not an E2E case. - -## 13. Docs and comments - -- `docs/features/rate-limits.mdx` and `docs/advanced/cli.mdx` — - windows survive restart and `--reload`, including each active - per-child partition; concurrency does not; still per instance. Drop - the “counters start fresh” wording on the CLI reload page. -- `docs/dev/2026-07-05_rate-limiting-spec.md` §8 and §11.7 — mark - persistence done; §11.2 stays future work. -- `CLAUDE.md` / `Agents.md` — drop rate-limit counters from the - “in-memory state resets on reload” list; keep session affinity and - live log buffers. -- `.env.template` and `config/config.example.yaml` — document - `RATE_LIMITS_FLUSH_INTERVAL` / `flush_interval`. - -## 14. Follow-up (not this change) +S205–S207 in `tests/e2e/release-e2e-scenarios.md` cover reload survival +on SQLite, `reset-one` staying cleared across a reload, and the same on +PostgreSQL and MongoDB. Hour windows, so the cap outlives the reload +wait. Shared user-path rules only: the release stack is OSS and a +per-child admin write would 403, so per-child persistence is covered by +the unit suites instead. + +The one non-obvious part is the shared `reload_release_gateway` helper. +It sends `SIGHUP`, then waits for a **new** `configuration reloaded` +line in that gateway's log before probing `/health`. That line is +logged after the old generation's `Shutdown` and before the new +`StartWithListener`, and a request sent straight after `kill -HUP` can +still be served by the old generation — without the log wait the +scenarios race. + +Crash-before-`Close` (periodic flush only) and token-window reload stay +unit tests; concurrency is not persisted, so it has no E2E case. + +## 13. Follow-up (not this change) Redis live counters: every `Admit` becomes an atomic script on shared keys (the key must include `partition`, same as `ruleKey`). No