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/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/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-spec.md b/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md new file mode 100644 index 000000000..3c5581b97 --- /dev/null +++ b/docs/dev/2026-08-16_rate-limit-counter-persistence-spec.md @@ -0,0 +1,349 @@ +# Rate Limit Counter Persistence + +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 + +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. It does not implement Redis-Lua, and does not +build a seam for it in advance (§5). + +## 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. +- 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. +- 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 + +```text +Service.Acquire / RecordTokens / RouteAvailable / Status / Reset + │ + ▼ + limiter (in memory, authoritative) + │ snapshot() / restore() + ▼ + Store.LoadCounters / SaveCounters / DeleteCounter / DeleteAllCounters +``` + +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: + +- 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** — 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 + +One row per live window key. Concurrent rules (`period_seconds = 0`) +are never written. + +```text +scope TEXT -- user_path | provider | model +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 +requests_previous INT64 +tokens_window_start INT64 +tokens_current INT64 +tokens_previous INT64 +updated_at INT64 -- unix seconds; drives staleness collection +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, 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 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. + +Package type (exported only if tests in another package need it; +otherwise unexported is fine): + +```go +type windowSnapshot struct { + Scope, Subject, Partition 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.” +A shared rule’s `Partition` is `""`. + +## 7. Store methods + +Add to the existing `Store` interface. No second factory. + +- `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 (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. +- `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. + +## 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 **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 +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. + +A second mutex (`persistMu`) serializes snapshot writes against +reset/delete row deletes. `admit()` does not take it. + +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). + +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 + +```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. + +## 10. Service operations vs snapshot + +| Operation | Memory | Snapshot | +|---|---|---| +| `Admit` / `RecordTokens` | update windows | next flush / Close | +| `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` +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 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 +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 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 + +Persistence never fails a request. `admit()` does not see store errors. + +- **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 + past the existing shutdown budget. +- **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. + +## 12. Testing + +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 + +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 +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. +That work extracts whatever seam it needs from the concrete `limiter`; +this change deliberately leaves none behind (§5). 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..ec5263753 100644 --- a/internal/ratelimit/factory.go +++ b/internal/ratelimit/factory.go @@ -60,7 +60,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(flushIntervalFromSeconds(cfg.RateLimits.FlushInterval)), + ) if err != nil { return nil, err } diff --git a/internal/ratelimit/persist.go b/internal/ratelimit/persist.go new file mode 100644 index 000000000..70bd767d6 --- /dev/null +++ b/internal/ratelimit/persist.go @@ -0,0 +1,153 @@ +package ratelimit + +import ( + "context" + "log/slog" + "math" + "time" +) + +const persistTimeout = 5 * time.Second + +type persistState int + +const ( + persistIdle persistState = iota + persistStarting + 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. +func WithFlushInterval(interval time.Duration) ServiceOption { + return func(service *Service) { + if interval < 0 { + interval = 0 + } + service.flushInterval = interval + } +} + +// 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 + } + return time.Duration(min(int64(seconds), int64(math.MaxInt64/time.Second))) * 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) +} + +// 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 + } + + s.lifeMu.Lock() + if s.persistState != persistIdle { + s.lifeMu.Unlock() + return + } + s.persistState = persistStarting + s.lifeMu.Unlock() + + loadCtx, cancel := persistContext(ctx) + snapshots, err := s.store.LoadCounters(loadCtx) + cancel() + + s.lifeMu.Lock() + defer s.lifeMu.Unlock() + 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 + if len(snapshots) > 0 { + slog.Info("rate limit counters restored", "windows", len(snapshots)) + } +} + +func (s *Service) startFlushLoop() { + if s.flushInterval <= 0 { + return + } + s.flushStop = make(chan struct{}) + s.flushDone = make(chan struct{}) + interval := s.flushInterval + stop := s.flushStop + done := s.flushDone + go func() { + defer close(done) + ticker := time.NewTicker(interval) + defer ticker.Stop() + for { + select { + case <-ticker.C: + flushCtx, cancel := persistContext(context.Background()) + s.flush(flushCtx) + cancel() + case <-stop: + 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(ctx context.Context, scope RuleScope, subject string, periodSeconds int64) error { + if s == nil || s.store == nil { + return nil + } + // 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) +} + +func (s *Service) persistDeleteAll(ctx context.Context) error { + if s == nil || s.store == nil { + return nil + } + 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 new file mode 100644 index 000000000..2ed46a420 --- /dev/null +++ b/internal/ratelimit/persist_test.go @@ -0,0 +1,456 @@ +package ratelimit + +import ( + "context" + "errors" + "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") + } +} + +// 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 { + 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) + } + 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("Close without Start saved %d times, 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 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) + } +} + +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 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{{ + 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 +} + +func (s *recordingStore) SaveCounters(ctx context.Context, snapshots []WindowSnapshot) error { + 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 +} + +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()) +} + +// 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 11bf24a54..6fea58062 100644 --- a/internal/ratelimit/service.go +++ b/internal/ratelimit/service.go @@ -35,6 +35,13 @@ type Service struct { rules []Rule quotaTemplates bool + + flushInterval time.Duration + persistMu sync.Mutex + lifeMu sync.Mutex + persistState persistState + flushStop chan struct{} + flushDone chan struct{} } func NewService(ctx context.Context, store Store, options ...ServiceOption) (*Service, error) { @@ -57,9 +64,35 @@ 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.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() } } @@ -138,7 +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}) - return s.Refresh(ctx) + // 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 persistErr } func (s *Service) ReplaceConfigRules(ctx context.Context, rules []Rule) error { @@ -336,7 +377,7 @@ func (s *Service) ResetRule(scope RuleScope, subject string, periodSeconds int64 return err } s.limiter.reset(ruleKey{scope: scope, subject: subject, periodSeconds: periodSeconds}) - return nil + return s.persistDelete(context.Background(), scope, subject, periodSeconds) } // ResetAll clears every live window counter. @@ -345,7 +386,7 @@ func (s *Service) ResetAll() error { return ErrUnavailable } s.limiter.resetAll() - return nil + return s.persistDeleteAll(context.Background()) } func (s *Service) matchingRules(subjects Subjects) []Rule { diff --git a/internal/ratelimit/service_test.go b/internal/ratelimit/service_test.go index 7c95528ec..65f7d2e16 100644 --- a/internal/ratelimit/service_test.go +++ b/internal/ratelimit/service_test.go @@ -13,7 +13,10 @@ import ( // memStore is a minimal in-memory Store for service tests. type memStore struct { - rules []Rule + 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) { @@ -68,6 +71,62 @@ 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 +} + +// 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 { + now := time.Now().Unix() + if m.written == nil { + m.written = make(map[string]int64) + } + for _, snap := range snapshots { + m.written[snapshotIdentity(snap)] = 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 m.written[snapshotIdentity(snap)]+2*snap.PeriodSeconds >= now { + kept = append(kept, snap) + } + } + m.counters = kept + 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..cbcba56da --- /dev/null +++ b/internal/ratelimit/snapshot.go @@ -0,0 +1,133 @@ +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"` +} + +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..396bc9351 100644 --- a/internal/ratelimit/store.go +++ b/internal/ratelimit/store.go @@ -9,13 +9,21 @@ 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 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 Close() error } 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 0b4e8e7e4..e5629ac95 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,93 @@ 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 { + slog.Warn("rate limit counters: skipping malformed snapshot", "error", err) + continue + } + snapshots = append(snapshots, snap) + } + if err := cursor.Err(); err != nil { + return nil, fmt.Errorf("iterate rate limit counters: %w", err) + } + 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 { + writes = append(writes, mongo.NewUpdateOneModel(). + SetFilter(counterIdentityFilter(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 { + return fmt.Errorf("upsert 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 +} + +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}, + {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_mongodb_test.go b/internal/ratelimit/store_mongodb_test.go index a0892f911..26bfe94f5 100644 --- a/internal/ratelimit/store_mongodb_test.go +++ b/internal/ratelimit/store_mongodb_test.go @@ -224,3 +224,59 @@ func TestMongoDBStoreMigratesPreScopeDocuments(t *testing.T) { } }) } + +// 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) { + store, err := NewMongoDBStore(context.Background(), db) + if err != nil { + t.Fatalf("NewMongoDBStore() failed: %v", 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) + } + }) + }) +} + +// 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 594424ee7..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" @@ -25,6 +26,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 +81,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 +172,102 @@ 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 { + // 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) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("iterate rate limit counters: %w", err) + } + 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 { + for _, snap := range snapshots { + 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("upsert rate limit counter: %w", err) + } + } + if _, err := q.Exec(ctx, pruneCountersSQL, now); err != nil { + return fmt.Errorf("prune expired rate limit counters: %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 + 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..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" @@ -250,3 +251,52 @@ func TestSQLStoreMigratesPreScopeTable(t *testing.T) { } }) } + +func TestSQLStoreCounterRoundTrip(t *testing.T) { + runSQLStoreTest(t, func(t *testing.T, store *SQLStore) { + 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) + } + }) + }) +} + +// 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) + } +} 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..50b1ad873 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 --connect-timeout 1 --max-time 2 "$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 +```