feat(token_rate_limit): sliding-window/token-bucket rate limiting (M1/M2/M6) - #796
feat(token_rate_limit): sliding-window/token-bucket rate limiting (M1/M2/M6)#796jordigilh wants to merge 16 commits into
Conversation
|
Unsigned commits: e156a1a, 47f6876, 82eab8e. Please sign your commits. |
c02d386 to
d3c7813
Compare
Walking skeleton for the reservation/reconcile/headers plumbing described in ai#658 (rate + burst config, admission-time reservation, post-response reconciliation against token_count's token.total, 429 + X-RateLimit-*-Tokens headers). Global bucket only; per-key buckets (M5), configurable estimation (M3), and token-type weighting (M4) are deliberately out of scope. This is explicitly a spike, not a candidate implementation: it uses a continuous token-bucket algorithm, while ai#658's current "How?" design section specifies sliding-window budgets (Rules/token_budgets/tiers) -- a different algorithm, not a scoped-down version of it (see review comment on discussion_r3758007992). Built to validate the reservation/reconcile/ headers pattern end-to-end and to ground the GA-readiness review of praxis-proxy#658 in working code; the admission/keying core here should not be adjusted in place once window semantics are confirmed -- expect a rebuild against whichever model is confirmed. Also surfaced a reliability gap in praxis-proxy itself: no hook reliably fires on response-body abort (upstream reset/timeout/disconnect), so reservations leak on those paths -- tracked in the praxis-proxy#658 review thread, not fixed here. Signed-off-by: Jordi Gil <jgil@redhat.com>
…#129) Add ai#129's bucket_key_header option to the token_rate_limit spike: one independent token bucket per unique value of a configured request header, falling back to a single shared bucket when the header is absent or the per-key map's hard cap is reached. Mirrors rate_limit's own PerIp DashMap pattern (soft/hard eviction caps), generalized from IpAddr keys to header values. Reservation and reconciliation are wired to operate on the same per-key bucket across request/response phases via a stashed token_rate_limit.bucket_key metadata entry, so multi-tenant admission decisions stay consistent end-to-end. Adds a token-rate-limit-per-app.yaml example and integration test demonstrating multiple apps sharing one budget rule, each keyed by its own request header value: each app draws down only its own bucket, with one app hitting 429s having zero effect on the others. Built with TDD (RED/GREEN per increment) as an extension of the existing walking-skeleton spike, not a candidate implementation for ai#658 -- see that commit's caveats on the token-bucket vs sliding-window design divergence still open on the PR. Signed-off-by: Jordi Gil <jgil@redhat.com>
…mples Replace odin/thor/loki naming and "customer scenario" framing with generic app-a/app-b/app-c placeholders across the per-app bucket key tests, integration test, and example config. The underlying scenario (independent per-app token budgets via bucket_key_header) is unchanged. Signed-off-by: Jordi Gil <jgil@redhat.com>
…backend
Replace the token-bucket state model with an exact sliding-window
reservation ledger (ported from nerdalert's poc/distributed-token-rate-limit-demo
spike branch, attributed in-file), behind a pluggable backend trait so
the same filter runs in-process (default) or against a shared Valkey
instance -- the "final scenario" of budgets shared across gateway
instances/replicas, layered on top of ai#129's per-app bucket_key_header
keying.
Config schema moves from rate/burst to window/capacity to match ai#658's
own proposal language, plus a reservation_timeout (how long an
unreconciled reservation is trusted before being conservatively charged)
and a backend block (memory | valkey, with ${ENV_VAR}-expandable url).
Adds unit coverage for gaps identified during review against the
project's TDD/pyramid conventions:
- a lost request (aborted before the response completes) is still
charged at its estimate once its reservation times out, so it can't
bypass the budget, but isn't locked out forever once the window
rolls over
- an empty or oversized bucket_key_header value falls back to the
shared budget rather than bypassing per-app isolation or minting its
own budget (a cardinality-exhaustion vector against the per-key cap)
- a malformed backend.url ${...} reference (multiple references,
embedded rather than whole-value, invalid variable name) fails
config load loudly instead of connecting to an unintended host
Also fixes the integration-tier example configs/tests, which still
referenced the removed rate/burst fields, and clears the clippy/fmt/doc
debt introduced by the port (missing docs, functions over the line
cap) via extraction rather than lint suppression.
Signed-off-by: Jordi Gil <jgil@redhat.com>
…oken_bucket) Replaces the single fixed-algorithm config with an ordered `rules:` list (ai#789 / praxis#551): each rule gets its own optional static header-value `match`, its own admission algorithm, and its own budget, mirroring GuardrailsFilter's `rules: Vec<RuleConfig>` shape. The first rule whose match is satisfied (or which has no match) applies; a request satisfying no rule is not rate limited by this filter instance. Two algorithms now live behind the same backend trait: - sliding_window: the existing exact trailing-window ledger, unchanged in behavior, now one algorithm choice among others instead of the only one. - token_bucket: new continuous-refill ledger (capacity, refill_rate tokens/sec), reusing the refill formula from Praxis's own lock-free traffic_management::token_bucket, extended with the reserve/reconcile split this filter needs (immediate-decrement reserve, credit-back reconcile). Both algorithms support in-process (default) and Valkey-backed shared state, so per-algorithm budgets are consistent across gateway replicas just like the existing sliding-window Valkey backend. Valkey-side duplication between the two algorithms' EVAL/reconcile-worker scaffolding was extracted into shared ValkeyEval/ReconcileWorker helpers rather than copy-pasted. The state-backend trait itself was widened with backend-agnostic reconcile_sync()/cleanup() methods, replacing on_response's direct coupling to the sliding-window Ledger type -- necessary groundwork for the filter to be indifferent to which algorithm a rule picked. Full pyramid coverage: unit tests for config parsing, both in-memory ledgers' business behavior (admission bounds, refill/window recovery, reconciliation, FedRAMP-guided DoS/bounds scenarios), and both Valkey backends (cross-instance isolation, fail-closed on backend unavailability); an integration test driving two Valkey-backed rules (one per algorithm) through the real HttpFilter to prove per-algorithm cross-replica isolation; and a manual end-to-end pass against the real gateway binary + docker-compose Valkey exercising admit/deny/refill- recovery for both algorithms before any demo work. Signed-off-by: Jordi Gil <jgil@redhat.com>
`refill_rate <= 0.0` never trips for NaN or +/-Infinity (IEEE-754 comparisons against NaN are always false, and Infinity > 0.0 is true), so both silently passed validation in both the in-process TokenBucketLedger and ValkeyTokenBucketBackend. NaN would poison every refill calculation for that bucket's state; Infinity would refill to capacity on any nonzero time delta, disabling the rate limit outright. For the Valkey-backed path this corrupts state shared across every gateway replica, not just one instance. `refill_rate: .nan`/`.inf` both parse cleanly from YAML via serde_yaml, so this was a real, config-reachable input-validation gap (FedRAMP SI-10), not theoretical. Found during this branch's own GA-readiness audit. Fixed both validation sites to `!refill_rate.is_finite() || refill_rate <= 0.0`, following RED (failing tests proving the bypass) then GREEN. Also adds examples/configs/token-rate-limit-mixed-algorithms.yaml: no runnable example previously exercised token_bucket or mixed per-rule algorithms (the only worked example lived in a rustdoc comment, not covered by the repo's own example-parsing test harness). Verified against `all_example_configs_parse` and synced into examples/README.md. Signed-off-by: Jordi Gil <jgil@redhat.com>
…ween backends ValkeyTokenBucketBackend::new duplicated TokenBucketConfig::validate's capacity/refill_rate checks (positivity, the 2^53 f64-precision bound, and the capacity/refill_rate ratio bound TOKEN_BUCKET_RESERVE_SCRIPT folds into a PEXPIRE TTL), which had grown past clippy's too-many-lines limit. Extracted validate_capacity_and_refill_rate(capacity, refill_rate) in token_bucket_ledger.rs and call it from both backends, so a config rejected on one backend is rejected identically on the other by construction instead of by keeping two copies in sync by hand. Adds the two capacity/refill_rate bound checks to the Valkey backend's own test module (previously only exercised through the in-memory ledger's tests), and makes TOKEN_BUCKET_RESERVE_SCRIPT pub(super) so the shared validator's doc comment can link to it across modules. Also: reworded a handful of code comments that referenced an internal, not-externally-published compliance-mapping shorthand with plain engineering language describing the same behavior; widened four separator comments in backend.rs back to the repo's 80-column convention; fixed a cargo fmt drift in tests.rs; and added an integration smoke test for examples/configs/token-rate-limit-mixed- algorithms.yaml, which lint-example-tests had flagged as uncovered. Signed-off-by: Jordi Gil <jgil@redhat.com>
…nted rationale Manual audit against this repo's clean-code standards (no automated comment-slop lint exists in Rust the way it does elsewhere): several comments re-derived rationale that already lives, in full, at a single authoritative source. - ValkeyTokenBucketBackend's construction-time tests re-explained why each bound exists (NaN/Infinity, the 2^53 ceiling, the PEXPIRE ratio) even though they all now go through the shared validate_capacity_and_refill_rate helper -- collapsed to a one-line pointer at the constant/test that already documents it. - CompiledRule::resolve_key's doc re-enumerated the exact fallback conditions already spelled out on FALLBACK_KEY's own doc comment. - MAX_CAPACITY_REFILL_RATE_RATIO_SECS/MAX_F64_SAFE_INTEGER's doc comments trimmed of restated phrasing while keeping the load-bearing "why". No behavior change; cargo test/clippy/doc/lint all still pass. Signed-off-by: Jordi Gil <jgil@redhat.com>
Add unit tests for branches that had no dedicated coverage: ledger and token-bucket key/reservation capacity denials, config validation error paths (zero bounds, zero reservation timeout, duplicate budget windows), expired-reservation reaping during reserve/reconcile, and the reconcile worker's bounded-retry-then-abandon path. Measured against a live Valkey instance (CI does not set TOKEN_RATE_LIMIT_VALKEY_URL, so these paths are otherwise never exercised): backend.rs 79.6% -> 94.4%, ledger.rs 94.0% -> 96.8%, token_bucket_ledger.rs 95.2% -> 97.2% line coverage. Signed-off-by: Jordi Gil <jgil@redhat.com>
ai#658 (the proposal PR this module's doc comments cited) was auto-closed today by the org-wide proposals migration to praxis-proxy/enhancements, and docs/proposals/00121_token-rate-limiting.md no longer exists in this repo. Repoint every reference to the proposal's new home instead of a closed PR, and regenerate docs/filters/token_rate_limit.md to match. Signed-off-by: Jordi Gil <jgil@redhat.com>
d3c7813 to
c644536
Compare
praxis-bot
left a comment
There was a problem hiding this comment.
Review: feat(token_rate_limit): per-header rate-limit bucket keys (ai#129)
Well-structured PR with thorough test coverage (112 unit tests, integration tests for all three example configs, both algorithms, both backends). The reserve/reconcile design, fail-closed Valkey behavior, and per-key budget isolation are all solid. Two issues worth addressing before merge.
| Severity | Count |
|---|---|
| Critical | 0 |
| Large | 2 |
| Medium | 1 |
All comments are inline.
| /// rather than hanging it indefinitely. | ||
| async fn connection(&self) -> Result<MultiplexedConnection, BackendError> { | ||
| tokio::time::timeout(VALKEY_TIMEOUT, self.client.get_multiplexed_async_connection()) | ||
| .await |
There was a problem hiding this comment.
Large: new TCP connection opened per Valkey operation
connection() calls self.client.get_multiplexed_async_connection() on every invocation, which opens a fresh TCP connection. Every reserve() call (on the request hot path) pays TCP handshake latency before it can run the EVAL. At production request rates this creates significant connection churn and risks FD exhaustion.
redis::aio::MultiplexedConnection is designed to be cloned and reused -- it multiplexes commands over a single TCP socket internally. Cache the connection in an Option<MultiplexedConnection> (or tokio::sync::OnceCell) and reconnect only on error, rather than opening a fresh socket per request.
The eval() method's doc comment already notes worst-case latency is 2x VALKEY_TIMEOUT (connect + command); caching the connection would cut that to 1x for the common case.
| | `rules[].r#match` | MatchConfig | no | Static header-value match condition. Every listed header must be present on the request with an exact value match (`ANDed`) for this rule to apply. Omit entirely for a catch-all rule. | | ||
| | `rules[].r#match.headers` | object<string, string> | yes | Every header must be present on the request with this exact value for the rule to match (`ANDed` across all entries). | | ||
| | `rules[].enabled` | bool | no | Whether this algorithm is enabled. | | ||
| | `rules[].level` | integer | no | Compression level for this algorithm. | |
There was a problem hiding this comment.
Large: generated docs contain phantom fields and are missing the actual algorithm fields
The configuration table lists rules[].enabled ("Whether this algorithm is enabled") and rules[].level ("Compression level for this algorithm") -- these don't exist in RuleConfig and appear to be copy-paste artifacts from a compression filter's doc template. Meanwhile, the actual algorithm-specific fields (algorithm, window, capacity, refill_rate) are absent from the table because they live inside the #[serde(flatten)]-ed AlgorithmConfig enum and the doc generator doesn't flatten them.
The reference table is the primary operator-facing config documentation. It currently documents fields that don't exist, silently omits fields that do, and would mislead anyone configuring this filter from the reference rather than copying the example verbatim.
Please regenerate after fixing the xtask generator, or manually correct this table to remove enabled/level and add the actual algorithm-discriminated fields.
| | Field | Type | Required | Description | | ||
| |-------|------|---------|-------------| | ||
| | `rules` | RuleConfig[] | yes | Evaluated in order; the first rule whose `match` is satisfied (or which has no `match` at all) applies to a given request. A request satisfying no rule's `match` is not rate limited by this filter instance -- add a trailing rule with no `match` to enforce a catch-all budget instead. | | ||
| | `rules[].name` | string | yes | Human-readable rule identifier, folded into Valkey key namespacing so distinct rules sharing one backend never collide. Renaming a live `valkey`-backed rule is therefore not a no-op for operators: it changes the Valkey key hash, so the old name's tracked budget is orphaned (left to expire on its own TTL) and the new name starts with a fresh budget. There's no migration/rename path today -- routine config hygiene (e.g. renaming `"gold"` to `"gold-tier"`) silently resets that rule's state. | |
There was a problem hiding this comment.
Medium: Rust raw identifier prefix leaks into operator-facing docs
rules[].r#match shows the Rust raw identifier syntax (r#). Operators configure this field as match: in YAML -- the r# prefix is a Rust-side concern that the doc generator shouldn't be surfacing. Same applies to rules[].r#match.headers on the next row.
Resolves conflicts in Cargo.lock (combine redis and quixotic-plecostomus-core dependency additions) and filters/src/register.rs (combine TokenRateLimitFilter and Sigv4SignFilter imports). Signed-off-by: Jordi Gil <jgil@redhat.com>
docs/filters/token_rate_limit.md was rendering two long-standing generator bugs: - `rules[].enabled`/`rules[].level` were phantom fields belonging to an unrelated `struct AlgorithmConfig` in the sibling praxis repo's compression filter. The doc generator's shared-item scan matched on bare type name only, so it picked up that struct instead of this module's `enum AlgorithmConfig` whenever both were in scope. Renamed the token_rate_limit enum to `RuleAlgorithm` to remove the collision. - `rules[].r#match` leaked its raw-identifier prefix into the rendered field name; serde serializes it as plain `match`. Added `strip_raw_ident_prefix` in filter_docs.rs and applied it in `serde_field_name`, with unit test coverage. Signed-off-by: Jordi Gil <jgil@redhat.com>
ValkeyEval::connection() opened a fresh multiplexed connection on every reserve/reconcile call instead of reusing one, defeating the purpose of a "multiplexed" connection and adding an avoidable TCP/TLS handshake to the request path on every call. Cache the connection behind a tokio::sync::Mutex, established lazily on first use and cloned (cheap: MultiplexedConnection is a handle onto one shared pipelined connection, not a socket per clone) on every subsequent call. Invalidate the cache on any command failure or timeout so a wedged/reset connection doesn't get reused indefinitely. Also updates two stale doc comments that still described the now-removed bucket_key_header as this filter's key source. Signed-off-by: Jordi Gil <jgil@redhat.com>
The merge with upstream/main pushed register_general_ai_filters over clippy's too_many_lines threshold. Split the token counting/usage/ rate-limiting registrations out into register_token_filters. Signed-off-by: Jordi Gil <jgil@redhat.com>
Drops bucket_key_header (ai#129), M5 of the original per-app budgets proposal. Every rule now resolves to a single shared budget (FALLBACK_KEY) regardless of request headers. M5's per-arbitrary-header keying conflicts with the quota-key design ai#790 already established for the Valkey-backed shared quota work (namespace + rule + principal + canonical model, not an arbitrary header value), and ai#790 itself hasn't been approved yet. Rather than carry two incompatible keying schemes, the team agreed to descope M5 out of this change; ai#129 tracks it separately pending the quota-key question being resolved. Removes the bucket_key_header config field, its resolve_key() call path, the per-app example config and its README entry, and the per-app unit/integration test coverage. This change keeps M1 (sliding-window), M2 (token-bucket), and M6 (Valkey backend) intact -- only the per-request keying layer is removed. Signed-off-by: Jordi Gil <jgil@redhat.com>
CI's `cargo +nightly fmt --all -- --check` flagged the RuleAlgorithm import in mod.rs as out of alphabetical order within its use-group. Signed-off-by: Jordi Gil <jgil@redhat.com>
nerdalert
left a comment
There was a problem hiding this comment.
Thanks for putting this together. The reservation/reconciliation split and shared backend direction are solid. I found four areas where tightening the trust boundary and hot-path behavior would make the implementation safer to operate.
| /// present on the request with an exact value match (`ANDed`) for | ||
| /// this rule to apply. Omit entirely for a catch-all rule. | ||
| #[serde(default)] | ||
| pub r#match: Option<MatchConfig>, |
There was a problem hiding this comment.
Thanks for documenting the matching semantics clearly. Could we make the trust boundary explicit here? These are ordinary client request headers, so a caller can omit or change a value such as x-app-id; with no catch-all, that makes the request unmatched and therefore unmetered. My preference would be to match trusted filter metadata established by authentication. If header matching must remain in this PR, please require/document that a preceding trusted filter overwrites the header and consider rejecting configurations without an explicit catch-all unless bypass is deliberately enabled.
|
|
||
| async fn on_request(&self, ctx: &mut HttpFilterContext<'_>) -> Result<FilterAction, FilterError> { | ||
| let now_ms = self.now_ms(); | ||
| let Some((rule_index, rule)) = self.matching_rule(&ctx.request.headers) else { |
There was a problem hiding this comment.
Could we avoid reserving quota for unrelated or invalid requests? With a catch-all rule, every request reaching this hook is charged, including readiness probes and malformed/non-inference traffic that never produces token usage; the current integration test even accounts for GET / consuming a reservation. That also gives unauthenticated traffic a cheap way to exhaust a shared budget. A focused fix would be to require authentication plus request/model validation before this filter and validate that ordering, or add an explicit request/path scope so only qualified inference requests reach reservation.
| ) | ||
| .into()); | ||
| } | ||
| if rule.estimate_tokens > capacity { |
There was a problem hiding this comment.
Could we apply a Lua-safe integer bound to sliding-window values here? The Valkey script converts capacity, estimate_tokens, and reconciled usage with Lua tonumber, so integers above 2^53 are not represented exactly and may produce incorrect admission decisions. The token-bucket implementation already defines a safe-integer ceiling; reusing that validation for sliding-window capacity/estimate and rejecting oversized actual usage before reconciliation would keep the memory and Valkey backends consistent.
| keys: &[String; N], | ||
| args: &[String], | ||
| ) -> Result<Vec<i64>, BackendError> { | ||
| let mut command = redis::cmd("EVAL"); |
There was a problem hiding this comment.
Connection reuse is a good improvement. One remaining hot-path optimization: this still sends the complete Lua source through EVAL for every reservation. Could we use redis::Script/EVALSHA, cache the script digest, and retry once with script loading on NOSCRIPT? That keeps the normal request-path command small while retaining bounded recovery after a Valkey restart or script-cache flush.
Summary
Implements the uncontested MVP core of the token rate limiting proposal (
00121_token-rate-limiting.mdin praxis-proxy/enhancements, epic ai#121):rules:, each an optional static header-value match condition bound to its own budget (catch-alldefaultrule when nomatch:is given)token_count'stoken.total) once the response completesX-RateLimit-*-Tokensheaders on hard denyTwo admission algorithms, chosen per rule via
algorithm:(sliding_window|token_bucket, per the maintainer framing on ai#789/praxis#551 that this is "a per-rule choice, similar to shadow/enforcement-action knobs elsewhere"), both behind a pluggable backend: in-process by default, or a shared Valkey backend (backend: {kind: valkey}) for state shared across gateway instances/replicas.Descoped from this PR: M5 (per-header/per-app bucket keys, ai#129) has been pulled out. It conflicted with the quota-key design ai#790 already established for the Valkey-backed shared quota work (
namespace + rule + principal + canonical model, not an arbitrary header value), and ai#790 itself isn't approved yet. Every rule now resolves to a single shared budget regardless of request headers; ai#129 tracks per-key budgets separately pending the quota-key question being resolved with ai#790's stakeholders.Also deliberately deferred (see
filters/src/token_rate_limit/mod.rsmodule doc for the full list and rationale): configurable estimation (M3), token-type-aware weighting (M4), multiple budgets/soft-limit tiers per rule, observability, and metering.Additional fixes included
reserve/reconcilecall.token_rate_limit'sAlgorithmConfigenum collided with an unrelated struct of the same name in a sibling filter, producing phantom fields in generated docs; renamed toRuleAlgorithm. Also fixed the generator leaking Rust'sr#raw-identifier prefix into thematchfield name.Testing
filters/src/token_rate_limit/, covering both algorithms, both backends, config validation, and reservation/reconciliation edge cases (expired-reservation reaping, lost-request handling, capacity denials on all three configured bounds).examples/configs/token-rate-limit.yaml,examples/configs/token-rate-limit-mixed-algorithms.yaml).make lint(clippy, fmt,cargo doc -D warnings, separator/filter-doc/example-test xtask lints) passes clean.Supersedes the closed draft #773, which called out the sliding-window-vs-token-bucket question as unresolved; this PR resolves it by supporting both, per rule.