Repository navigation
fix(router): buffer bodies when routing-key override is enabled - #2355
Conversation
Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
github-actions
Bot
added
tests
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (7)
Included review availability: 7 reviews are currently available. Your included PR review attempts over the past 7 days set your current allowance at 8 reviews per hour. 📝 WalkthroughSummary by CodeRabbit
WalkthroughRouting-key overrides now force automatically forwarded requests through buffered body routing. The body ChangesRouting-key override routing
Estimated code review effort: 3 (Moderate) | ~20 minutes Merge Risk: ⚪ Minimal · up to The change restores buffered request-body routing when routing-key override is enabled so body-derived routing data and precedence are preserved; no actionable merge-blocking risk remains after normal checks and review. Sequence Diagram(s)sequenceDiagram
participant HTTPRequest
participant request_body_path
participant PolicyRegistry
participant BodyPathPolicy
participant Worker
HTTPRequest->>request_body_path: evaluate request body path
request_body_path->>PolicyRegistry: routing_key_override_enabled()
PolicyRegistry-->>request_body_path: override enabled
request_body_path->>BodyPathPolicy: select buffered path
BodyPathPolicy-->>request_body_path: body RID remains available
request_body_path->>Worker: forward buffered request
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Warning Your free Security trial is over. An organization admin can activate Security or dismiss this notice. Comment |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| /// Whether sticky routing may derive its preferred key from the request | ||
| /// body's `rid`, requiring automatic body-path selection to keep the body | ||
| /// readable. | ||
| pub(crate) fn routing_key_override_enabled(&self) -> bool { |
There was a problem hiding this comment.
🟡 Nit: This accessor reports "the override may read body rid" purely from enabled, but select_worker only consults the sticky map when routing_key_override_applies(policy.name()) is true — i.e. never for manual / consistent_hashing (line 346). With --routing-key-override + --policy manual, the override is a no-op for selection, yet every request now takes the buffered path for a rid nothing will consume. any_policy_needs_request_text already handles this asymmetry by folding routing_key_override_applies into its per-policy scan; mirroring it here (override enabled and at least one registered policy the override can apply to, since the model — hence the policy — is inside the unread body) would keep those deployments streamable.
Related: because decide_body_path now checks routing_key_override first, the keyed_override waiver at line 735 can only be true in exactly the case where the result is discarded (the sole production caller is Router::request_body_path). Its doc comment — "a valid routing-key header under the sticky override supersedes the policy before it reads text" — describes a streamed outcome that can no longer happen, so it's worth updating or dropping alongside the branch.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
| .worker_registry | ||
| .get_routing_pool(crate::worker::UNKNOWN_MODEL_ID, RoutingPool::HttpRegular); | ||
| decide_body_path(&BodyPathInputs { | ||
| routing_key_override: self.policy_registry.routing_key_override_enabled(), |
There was a problem hiding this comment.
🟡 Nit: All BodyPathInputs fields are evaluated eagerly, so with the override on every request still pays for get_routing_pool (clones the whole HTTP-regular pool into a Vec), the mutates_request() scan, any_policy_needs_request_text (scans the default policy plus every per-model policy, parses two hint headers), and model_count() — only to be told Buffer("routing_key_override") unconditionally. That's pure per-request waste on the ingress hot path for the entire class of deployments this PR targets. An early return before gathering the pool avoids it:
fn request_body_path(&self, headers: &HeaderMap, wasm_request_hooks: bool) -> BodyPath {
if self.policy_registry.routing_key_override_enabled() {
return BodyPath::Buffer(REASON_ROUTING_KEY_OVERRIDE);
}
...(with routing_key_override: false left in the struct for the shared matrix, or the field dropped if the gate lives only here).
| let policy_registry = Arc::new(PolicyRegistry::with_override( | ||
| config.policy.clone(), | ||
| config.routing_key_override.clone(), | ||
| )); |
There was a problem hiding this comment.
🟡 Nit: Good fix, but the same gap remains in the sibling harness: tests/common/test_app.rs:51 (create_test_app, which does take a router_config) still builds PolicyRegistry::new(router_config.policy.clone()) and silently drops routing_key_override. Any future test that wires the override through that entry point will pass for the wrong reason. Worth converting it to with_override too so both harnesses match the production builder. (test_app.rs:245 takes no config, so it's unaffected in practice.)
Three merged router fixes changed behaviour a client can observe without leaving an end-to-end check behind: - #2417: an overload shed must answer 503 with its own worker_overload_protection_shed code and a Retry-After, not the generic no_available_workers. The engine is pinned to one running request so a burst piles up in its queue; a probe sent while the queue is deep must be shed, and the worker must serve again once the burst drains. - #2355: with --routing-key-override a body rid must outrank the routing key header. Two workers under the manual policy: a header key stays sticky, and a request whose header names one lineage and whose body rid names another lands on the rid's worker. The serving worker is observed through the gateway's in-flight load during a long generation. - #2427: a gRPC worker stopped under a live gateway must be reported unhealthy, requests must fail fast rather than hang, and after the worker restarts on the same port it must return to rotation. The test borrows the session pool's worker so it does not race a cached worker for the GPU. Signed-off-by: Alex McC <319643551+hello-alexmcc@users.noreply.github.com>
Three merged router fixes changed behaviour a client can observe without leaving an end-to-end check behind: - #2417: an overload shed must answer 503 with its own worker_overload_protection_shed code and a Retry-After, not the generic no_available_workers. The engine is pinned to one running request so a burst piles up in its queue; a probe sent while the queue is deep must be shed, and the worker must serve again once the burst drains. - #2355: with --routing-key-override a body rid must outrank the routing key header. Two workers under the manual policy: a header key stays sticky, and a request whose header names one lineage and whose body rid names another lands on the rid's worker. The serving worker is observed through the gateway's in-flight load during a long generation. - #2427: a gRPC worker stopped under a live gateway must be reported unhealthy, requests must fail fast rather than hang, and after the worker restarts on the same port it must return to rotation. The test borrows the session pool's worker so it does not race a cached worker for the GPU. Signed-off-by: Alex McC <319643551+hello-alexmcc@users.noreply.github.com>
1 participant
Description
Problem
#2286 made the stream-vs-buffer decision per request. When
--routing-key-overrideis enabled, a request carrying a valid routing-key header or tokens hint could qualify for the streamed pass-through, where the JSON body is never parsed. That silently changed the override's semantics:ridmust win over the routing-key headers, but streamed requests never see it — two requests from the sameridlineage with different header keys can land on different workersrid,input_ids) it would have received on the buffered pathBefore #2286, streaming was opt-in (
--stream-request-bodies-over, default off), so enabling the override guaranteed every body was parsed andridprecedence always held.Solution
Make an enabled routing-key override the first hard-buffer reason in the body-path decision. With the override on, every request takes the buffered typed path, restoring the pre-#2286 contract: body
ridwins over header keys and body tokens stay visible to selection. No new flags; the decision is observable through the existingsmg_router_request_body_path_totalcounter aspath="buffered", reason="routing_key_override".Changes
routers/common/body_policy.rs: addrouting_key_overridetoBodyPathInputsand returnBuffer("routing_key_override")ahead of every other reasonpolicies/registry.rs: add arouting_key_override_enabled()accessorrouters/http/router.rs: feed the flag into the decision and update the streamed-path commentconfig/types.rs,main.rs: document the buffering contract on the flagtests/common: build the harness policy registry with the config's routing-key override, matching the production builder (it previously dropped the override, so integration tests could not exercise it)Test Plan
cargo +nightly fmt --allcargo clippy --all-targets -- -D warningscargo test -p smg(unit + integration, all green)The new integration test
routing_key_override_buffers_and_prefers_body_ridreproduces the regression: retries disabled, a valid tokens hint, distinct per-request header keys, and a shared bodyridlineage. Against the pre-fix code both requests stream and split across workers by header key; with this fix they buffer, pin to one worker via the strippedridlineage, and the forwarded body keepsridandinput_idsintact.Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspasses