Add throttling middleware for rate and in-flight limit.#4171
Add throttling middleware for rate and in-flight limit.#4171genericaccount-de wants to merge 1 commit into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAPI gateway throttling now uses named rate-limit and in-flight limit zones instead of global ChangesZone-based throttling migration
Estimated code review effort: 5 (Critical) | ~120 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
4646634 to
48abd86
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (2)
gears/system/api-gateway/tests/throttling_tests.rs (1)
252-268: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winFixed-delay synchronization risks flaky CI failures.
Using
sleep(20ms)to assume the first spawned request has acquired the in-flight permit before firing the second request is a timing assumption, not a guarantee. Under CI load this can intermittently let the second request land before the permit is taken, causing a spuriousOKinstead ofTOO_MANY_REQUESTS.Consider adding explicit synchronization — e.g., a
tokio::sync::Notify/oneshot signaled fromslow_handlerright after it starts (before the internalsleep), and await that signal instead of a fixed delay.♻️ Example fix sketch
-async fn slow_handler() -> Json<TestResponse> { - sleep(Duration::from_millis(200)).await; +async fn slow_handler() -> Json<TestResponse> { + STARTED.notify_one(); // signal readiness before the slow work + sleep(Duration::from_millis(200)).await; Json(TestResponse { message: "slow".to_owned(), }) }Then
await STARTED.notified()in the test instead of a fixedsleep.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@gears/system/api-gateway/tests/throttling_tests.rs` around lines 252 - 268, The throttling test currently relies on a fixed `tokio::time::sleep` to wait for the first spawned request to acquire the in-flight permit, which is flaky under load. Update the `throttling_tests.rs` flow around `first` and the `slow_handler` setup to use explicit synchronization instead of timing assumptions. Add a `tokio::sync::Notify` or oneshot signal from `slow_handler` as soon as it starts holding the permit, then have the test await that signal before sending the second request. This keeps the `TOO_MANY_REQUESTS` assertion deterministic.gears/system/api-gateway/src/middleware/throttling.rs (1)
91-107: 🚀 Performance & Scalability | 🔵 Trivial | ⚖️ Poor tradeoffSoft
max_keyseviction can't shrink below concurrently-active keys.
gate()only evicts entries withArc::strong_count(v) <= 1; if all tracked keys are currently in-flight (strong_count > 1), the map can grow pastmax_keysindefinitely under sustained high-cardinality traffic. This is a reasonable "soft cap" tradeoff for a first implementation givenKeyGates must stay alive while held, but worth a follow-up note/metric if unbounded key cardinality (e.g., unauthenticated IP keys) is a realistic threat model.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@gears/system/api-gateway/src/middleware/throttling.rs` around lines 91 - 107, The soft-cap behavior in InFlightZone::gate can leave the keys map growing beyond max_keys when all entries are still in flight, so add follow-up observability or a clear note around this tradeoff. Update gate() to surface when retain() fails to reclaim entries (for example via a metric/log on the key eviction path) and document that the cap is only best-effort while KeyGate references are active.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@gears/system/api-gateway/src/config.rs`:
- Around line 421-424: The validate_status helper currently accepts any valid
HTTP status, including 2xx/3xx values, for throttle responses. Update
validate_status in config.rs to keep using StatusCode::from_u16 for parsing,
then add an explicit check that only 4xx/5xx statuses are accepted for
limit-exceeded/throttle configuration and return an error otherwise. Make sure
the error message stays tied to zone and response_status_code so callers can
identify invalid non-error statuses.
In `@gears/system/api-gateway/src/middleware/throttling.rs`:
- Around line 12-14: Update the module-level docs in the throttling middleware
to match the current behavior: the comment near the top of throttling.rs is
stale because enforce() now performs live per-request rate-limit and in-flight
checks. Replace the reference to throttling_rules with the actual config
concepts used by this code path, such as rate_limit_zones and
in_flight_limit_zones, so the docs accurately describe what enforce() and the
surrounding middleware do.
- Around line 310-391: Rate-limit headers are being attached to the request in
enforce, so successful responses never expose them to clients. Update enforce
and the downstream success path around next.run(req) to add RateLimit-Policy,
RateLimit-Limit, RateLimit-Remaining, and legacy X-RateLimit-* onto the Response
before returning it, including the in-flight permit branch. Use the existing
symbols enforce, throttle_response, and the rate_zone/inflight_zone handling to
locate where the response headers should be applied.
- Around line 164-238: The throttling builder currently creates separate zone
state per auth partition, so the same named zone gets different limiter
instances before and after auth. Update the shared build path in build() so rate
zones and in-flight zones are constructed once and reused across both
partitions, instead of having independent HashMap/Arc state per require_ctx.
Reuse the same zone caches when populating ThrottlingInner routes so any
operation referencing the same zone name shares the same Arc<RateZone> or
Arc<InFlightZone> regardless of require_security_context.
In `@libs/toolkit/src/api/operation_builder.rs`:
- Around line 272-275: The ThrottlingSpec contract currently carries a key_type
value that is not enforced by runtime throttling, which can drift from the
actual gateway zone configuration. Update the operation builder/types around
ThrottlingSpec and the throttling map construction logic so key strategy is
sourced only from the referenced zone config, or validate any provided key_type
against the zone before accepting it. Use the ThrottlingSpec, key_type,
identity_extractor, and the throttling map builder code paths to remove the
stale field or enforce consistency.
---
Nitpick comments:
In `@gears/system/api-gateway/src/middleware/throttling.rs`:
- Around line 91-107: The soft-cap behavior in InFlightZone::gate can leave the
keys map growing beyond max_keys when all entries are still in flight, so add
follow-up observability or a clear note around this tradeoff. Update gate() to
surface when retain() fails to reclaim entries (for example via a metric/log on
the key eviction path) and document that the cap is only best-effort while
KeyGate references are active.
In `@gears/system/api-gateway/tests/throttling_tests.rs`:
- Around line 252-268: The throttling test currently relies on a fixed
`tokio::time::sleep` to wait for the first spawned request to acquire the
in-flight permit, which is flaky under load. Update the `throttling_tests.rs`
flow around `first` and the `slow_handler` setup to use explicit synchronization
instead of timing assumptions. Add a `tokio::sync::Notify` or oneshot signal
from `slow_handler` as soon as it starts holding the permit, then have the test
await that signal before sending the second request. This keeps the
`TOO_MANY_REQUESTS` assertion deterministic.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 974a0bd4-001e-42ea-a095-f1fb35faf48f
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (29)
Cargo.tomlconfig/e2e-local.yamlconfig/e2e-scope-enforcement.yamlconfig/e2e-tr-authz.yamlconfig/mini-chat.yamlconfig/no-auth.yamlconfig/quickstart.yamlgears/mini-chat/deploy/helm/mini-chat/templates/configmap.yamlgears/mini-chat/mini-chat/src/api/rest/routes/chats.rsgears/system/api-gateway/Cargo.tomlgears/system/api-gateway/src/config.rsgears/system/api-gateway/src/gear.rsgears/system/api-gateway/src/gear.rs.origgears/system/api-gateway/src/middleware/mime_validation.rsgears/system/api-gateway/src/middleware/mod.rsgears/system/api-gateway/src/middleware/rate_limit.rsgears/system/api-gateway/src/middleware/throttling.rsgears/system/api-gateway/tests/access_log_tests.rsgears/system/api-gateway/tests/body_limit_tests.rsgears/system/api-gateway/tests/http_metrics_tests.rsgears/system/api-gateway/tests/middleware_order.rsgears/system/api-gateway/tests/middleware_order.rs.origgears/system/api-gateway/tests/mime_validation_integration.rsgears/system/api-gateway/tests/rate_limit_tests.rsgears/system/api-gateway/tests/throttling_tests.rslibs/toolkit/src/api/mod.rslibs/toolkit/src/api/openapi_registry.rslibs/toolkit/src/api/operation_builder.rstesting/e2e/gears/mini_chat/config/base.yaml
💤 Files with no reviewable changes (12)
- gears/mini-chat/deploy/helm/mini-chat/templates/configmap.yaml
- config/no-auth.yaml
- gears/system/api-gateway/tests/access_log_tests.rs
- config/e2e-scope-enforcement.yaml
- config/mini-chat.yaml
- config/e2e-local.yaml
- gears/system/api-gateway/tests/middleware_order.rs.orig
- gears/system/api-gateway/tests/rate_limit_tests.rs
- config/e2e-tr-authz.yaml
- testing/e2e/gears/mini_chat/config/base.yaml
- gears/system/api-gateway/src/gear.rs.orig
- gears/system/api-gateway/src/middleware/rate_limit.rs
adc2cfd to
d7ff18b
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (1)
gears/mini-chat/deploy/helm/mini-chat/templates/configmap.yaml (1)
65-85: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThrottling limits hardcoded in production Helm template; consider exposing via
.Values.Unlike
bind_addr(line 54) which is parameterized through.Values.service.port, the rate/in-flight limits here (10/s, burst10, in-flight100,backlog_limit: 0) are hardcoded directly in the production chart and identical to the local dev config. Ops teams cannot tune throttling per-environment without editing the template, and these values (esp. zero backlog) may be too restrictive for production chat-creation traffic.Consider surfacing
rate_limit,burst_limit,in_flight_limit, andbacklog_limitas chart values with these as defaults.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@gears/mini-chat/deploy/helm/mini-chat/templates/configmap.yaml` around lines 65 - 85, The throttling settings in the mini-chat Helm ConfigMap are hardcoded instead of being configurable per environment. Update the configmap template to read the rate-limit, burst-limit, in-flight-limit, and backlog-limit values from chart values with the current numbers as defaults, using the existing config sections for rate_limit_zones and in_flight_limit_zones as the places to wire them in. This will let ops tune the mini-chat chat-create throttling without modifying the template while keeping the current behavior as the default.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@gears/mini-chat/deploy/helm/mini-chat/templates/configmap.yaml`:
- Around line 65-85: The throttling settings in the mini-chat Helm ConfigMap are
hardcoded instead of being configurable per environment. Update the configmap
template to read the rate-limit, burst-limit, in-flight-limit, and backlog-limit
values from chart values with the current numbers as defaults, using the
existing config sections for rate_limit_zones and in_flight_limit_zones as the
places to wire them in. This will let ops tune the mini-chat chat-create
throttling without modifying the template while keeping the current behavior as
the default.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 697e6620-d8b4-419f-8f52-39742791e61d
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (29)
Cargo.tomlconfig/e2e-local.yamlconfig/e2e-scope-enforcement.yamlconfig/e2e-tr-authz.yamlconfig/mini-chat.yamlconfig/no-auth.yamlconfig/quickstart.yamlgears/mini-chat/deploy/helm/mini-chat/templates/configmap.yamlgears/mini-chat/mini-chat/src/api/rest/routes/chats.rsgears/system/api-gateway/Cargo.tomlgears/system/api-gateway/src/config.rsgears/system/api-gateway/src/gear.rsgears/system/api-gateway/src/gear.rs.origgears/system/api-gateway/src/middleware/mime_validation.rsgears/system/api-gateway/src/middleware/mod.rsgears/system/api-gateway/src/middleware/rate_limit.rsgears/system/api-gateway/src/middleware/throttling.rsgears/system/api-gateway/tests/access_log_tests.rsgears/system/api-gateway/tests/body_limit_tests.rsgears/system/api-gateway/tests/http_metrics_tests.rsgears/system/api-gateway/tests/middleware_order.rsgears/system/api-gateway/tests/middleware_order.rs.origgears/system/api-gateway/tests/mime_validation_integration.rsgears/system/api-gateway/tests/rate_limit_tests.rsgears/system/api-gateway/tests/throttling_tests.rslibs/toolkit/src/api/mod.rslibs/toolkit/src/api/openapi_registry.rslibs/toolkit/src/api/operation_builder.rstesting/e2e/gears/mini_chat/config/base.yaml
💤 Files with no reviewable changes (9)
- config/no-auth.yaml
- gears/system/api-gateway/tests/middleware_order.rs.orig
- config/e2e-scope-enforcement.yaml
- config/e2e-local.yaml
- gears/system/api-gateway/src/middleware/rate_limit.rs
- config/e2e-tr-authz.yaml
- gears/system/api-gateway/tests/access_log_tests.rs
- gears/system/api-gateway/src/gear.rs.orig
- gears/system/api-gateway/tests/rate_limit_tests.rs
✅ Files skipped from review due to trivial changes (2)
- gears/system/api-gateway/src/middleware/mime_validation.rs
- Cargo.toml
🚧 Files skipped from review as they are similar to previous changes (15)
- gears/system/api-gateway/tests/body_limit_tests.rs
- gears/mini-chat/mini-chat/src/api/rest/routes/chats.rs
- gears/system/api-gateway/tests/mime_validation_integration.rs
- gears/system/api-gateway/src/middleware/mod.rs
- libs/toolkit/src/api/mod.rs
- gears/system/api-gateway/tests/middleware_order.rs
- config/quickstart.yaml
- gears/system/api-gateway/Cargo.toml
- gears/system/api-gateway/tests/http_metrics_tests.rs
- gears/system/api-gateway/src/gear.rs
- libs/toolkit/src/api/openapi_registry.rs
- gears/system/api-gateway/tests/throttling_tests.rs
- gears/system/api-gateway/src/middleware/throttling.rs
- libs/toolkit/src/api/operation_builder.rs
- gears/system/api-gateway/src/config.rs
ecfbf77 to
f080d81
Compare
There was a problem hiding this comment.
I think generic part of throttling should not be in api gateway crate
There was a problem hiding this comment.
Done, moved generic parts to the api-gateway.
cb8a5db to
a310889
Compare
…ute. Signed-off-by: ga <7kozlik+github@gmail.com>
a310889 to
ece9145
Compare
| /// peer address from `ConnectInfo`, else `"unknown"`. | ||
| fn client_ip(req: &Request) -> String { | ||
| let headers = req.headers(); | ||
| if let Some(ip) = headers |
There was a problem hiding this comment.
CRITICAL
client_ip trusts the leftmost X-Forwarded-For (and X-Real-IP) value, which is fully client-controlled, as the throttling bucket key with no trusted-proxy validation.
Pre-auth IP-keyed zones are the only limit an unauthenticated caller hits; a spoofed or per-request-unique XFF both bypasses the limit entirely and lets an attacker attribute load to victim IPs.
Derive the client IP from ConnectInfo or from a configured trusted-proxy hop count (rightmost untrusted XFF entry), and ignore client-supplied X-Forwarded-For/X-Real-IP unless the immediate peer is a trusted proxy.
| .ok_or_else(|| anyhow!("throttling: rate_limit zone '{name}' has rps = 0"))?; | ||
| let burst = NonZeroU32::new(cfg.burst_limit) | ||
| .ok_or_else(|| anyhow!("throttling: rate_limit zone '{name}' has burst_limit = 0"))?; | ||
| let limiter = RateLimiter::keyed(Quota::per_second(rps).allow_burst(burst)) |
There was a problem hiding this comment.
HIGH
RateLimitZone.max_keys is validated and documented as an LRU-eviction bound, but get_or_build_rate_zone never applies it to the governor keyed limiter, so its internal state store grows one entry per distinct key with no eviction.
For IP-keyed pre-auth zones the key is attacker-controlled via X-Forwarded-For, so unbounded distinct keys exhaust process memory and take down the gateway despite the configured bound.
Enforce max_keys on the rate limiter (periodic retain_recent/shrink prune bounded by a lifecycle CancellationToken, or a bounded keyed store) so the tracked-key count cannot exceed the configured limit.
| } | ||
| // Soft `max_keys` cap: drop gates no longer referenced by in-flight requests. | ||
| if self.keys.len() as u64 >= self.cfg.max_keys { | ||
| self.keys.retain(|_, v| Arc::strong_count(v) > 1); |
There was a problem hiding this comment.
HIGH
gate() runs DashMap::retain on the request hot path once len() >= max_keys; it write-locks and scans every shard, and when keys are actively in-flight (strong_count > 1) it evicts nothing, so the map keeps growing and the full O(n) scan repeats on each subsequent new-key request.
Under a flood of distinct keys (attacker-controlled IPs pre-auth) this becomes a per-request O(n) all-shard write-locking scan that serializes gate lookups and fails to bound memory — the cap becomes a DoS amplifier.
Replace the inline retain with a bounded eviction that does not write-lock all shards per request (segmented LRU with a size cap, or a periodic background sweep gated by a CancellationToken), and skip the scan when it cannot free entries.
| #[derive(Clone, Default)] | ||
| pub struct ThrottlingSpec { | ||
| /// Name of the rate-limit zone this operation participates in. | ||
| pub rate_limit_zone: String, |
There was a problem hiding this comment.
MEDIUM
ThrottlingSpec models zone references as bare String, with the empty string overloaded as a sentinel for "no zone in this category" (see if thr.rate_limit_zone.is_empty() in the gateway).
The absent-zone state is representable only by convention, so every consumer must remember the sentinel rule and an accidental empty string is a silently-valid "no throttling" — a type the compiler could have rejected.
Change both fields to Option<String> (or a ZoneName newtype) so "no zone" is expressed as None and callers cannot pass an ambiguous empty string.
| } | ||
|
|
||
| fn log_throttled(key: &ThrottleKey, kind: &str, id: &str) { | ||
| tracing::debug!( |
There was a problem hiding this comment.
MEDIUM
Enforced throttling rejections (real 429s) are logged at tracing::debug! and emit no metric, while non-enforced dry-run would-be-rejections are logged at tracing::warn!.
At the default info level operators get zero visibility into actual client throttling, and with no rejection counter, capacity/tuning problems on this critical path are undiagnosable in production.
Emit a rejection metric (counter keyed by zone/kind) and raise the enforced-rejection log to at least info, keeping the high-cardinality key value out of the message body.
| }); | ||
|
|
||
| // Let the first request acquire the permit before firing the second. | ||
| tokio::time::sleep(Duration::from_millis(20)).await; |
There was a problem hiding this comment.
LOW
The in-flight concurrency test coordinates the two requests with a fixed 20ms sleep, so the required ordering (first request holds the sole permit before the second arrives) is not guaranteed.
On a loaded runner the spawned first request may not acquire the permit within 20ms, causing spurious failures or masking a real regression instead of deterministically exercising the gate.
Replace the sleep with an explicit synchronization signal (a tokio::sync::Notify or Barrier the slow handler fires on entry) so the second request is issued only after the permit is confirmed held.
| }); | ||
|
|
||
| let throttling = builder.spec.throttling.as_ref().expect("throttling set"); | ||
| assert_eq!(throttling.rate_limit_zone, "rl_identity"); |
There was a problem hiding this comment.
LOW
builder_with_throttling_sets_spec_and_extractor asserts the zone-name fields equal the exact strings just passed to with_throttling, which only confirms the setter stored its input.
with_throttling merely assigns the spec, so echoing the same names back exercises no logic and catches no regression; only the extractor-invocation and Clone/Debug assertions carry value.
Drop the two zone-name echo assertions and keep/extend the ones that verify behavior — that the stored closure is invokable and the Arc<dyn Fn> field survives Clone/Debug.
|
Design-level comments: This is separate from the code-level findings. It concerns the throttling model itself.
IdentityKeyFn = Arc<dyn Fn(&axum::extract::Request) -> String> is a Rust closure — not serializable, not callable by an external gateway. Identity keying does not translate at all. IdentityKeyFn pulls axum::extract::Request into the toolkit contract type (operation_builder.rs:261). That leaks the runtime into the contract layer — a direct problem on feature/toolkit_contracts, where contracts are meant to be portable. Zone definitions live in ApiGatewayConfig, not in the contract. require_security_context / dry_run are concepts of the local middleware pipeline only. For an external target the key must be declarative (enum Key { Ip, Header(String), JwtClaim(String), … }) and the zone model must live in the contract, not the gateway gear.
Empty string as a "no zone" sentinel (if thr.rate_limit_zone.is_empty()) instead of Option.
Numeric limits are not duplicated (defined once per zone), only the binding boilerplate is.
Suggestion:
|
2 participants
Adds throttling middleware for rate and in-flight limit. Throttling is defined by common zones, which are reused by separate API endpoints by means of OperationsBuilder.
Summary by CodeRabbit
Retry-Afterand rate-limit metadata headers.