Repository navigation
perf(cache-aware): gather worker routing state in one pass - #1758
Conversation
|
Warning Review limit reached
More reviews will be available in 3 hours, 16 minutes, and 33 seconds. Learn how PR review limits work. Your organization has used up its prepaid credits, and credit purchases are no longer available. Enable the review add-on in the billing tab to keep reviews running — you're only billed for reviews past your plan's rate limits ($0.25/file). ⌛ How to resolve this issue?After more reviews become available, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans include higher PR review limits than trial, open-source, and free plans. In all cases, reviews become available again over time. During sustained high-volume PR review activity, CodeRabbit may temporarily slow when the next review becomes available. Please see our Fair Usage Limits Policy for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughIntroduces a ChangesRoutingState contract and cache-aware single-pass routing
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request optimizes worker selection in the cache-aware load balancing policy by introducing a single-pass routing_state() snapshot to retrieve worker status, load, and processed request counts, reducing ArcSwap guard overhead. The reviewer feedback identifies a performance concern regarding the creation of a HashMap (url_to_idx) on every request, which introduces heap allocation and hashing overhead on the hot path. To resolve this, the reviewer suggests replacing the HashMap lookup with an allocation-free linear scan over healthy_indices and using a stack-allocated SmallVec for healthy_indices to eliminate heap allocations entirely.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| workers: &[Arc<dyn Worker>], | ||
| text: &str, | ||
| healthy_indices: &[usize], | ||
| url_to_idx: &HashMap<&str, usize>, | ||
| min_load_idx: Option<usize>, | ||
| model_id: &str, | ||
| ) -> Option<usize> { |
There was a problem hiding this comment.
Remove the url_to_idx parameter from select_worker_with_text.
workers: &[Arc<dyn Worker>],
text: &str,
healthy_indices: &[usize],
min_load_idx: Option<usize>,
model_id: &str,
) -> Option<usize> {![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Clean perf optimization — the single-pass gather is well-structured, the tie-break logic from #1714 is correctly preserved, and the test helper mirrors the production gather faithfully. One 🟡 nit on a subtle (beneficial) behavioral tightening in the cache-hit tenant resolution path where the CB check is now included.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c1ea4012a4
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
c1ea401 to
40775d3
Compare
|
Addressed the review feedback (gemini-code-assist, chatgpt-codex, claude): dropped the per-request |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 40775d3615
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| let tenant_url: &str = &result.tenant; | ||
| workers | ||
| .iter() | ||
| .position(|w| w.url() == tenant_url) | ||
| .filter(|&idx| workers[idx].is_healthy()) | ||
| } else { | ||
| healthy_indices | ||
| .iter() | ||
| .min_by_key(|&&idx| { | ||
| (workers[idx].load(), workers[idx].processed_requests(), idx) | ||
| }) | ||
| .copied() | ||
| .find(|&idx| workers[idx].url() == tenant_url) |
There was a problem hiding this comment.
When a cache hit points at a tenant that is no longer in healthy_indices (for example, the newly handled case where its circuit breaker is open), this branch leaves selected_idx as None; the post-closure fallback below still returns healthy_indices.first() without incrementing the processed counter, so all such requests are redirected to the first registry entry rather than the precomputed min_load_idx the comment says should be used. The same cache-hit fallback logic appears in the string path, so this should fall back to min_load_idx there as well.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
40775d3 to
4057120
Compare
|
Re-measured the shipped (scan) version vs the earlier (map) version back-to-back at 2048 workers, same conditions:
Dropping the map cut the routing marginal +0.181 → +0.024 — cache_aware routing is now within noise of round_robin at 2048 workers (was ~+0.40 unoptimized). Test Plan updated. |
select_worker made several O(workers) passes per request (the healthy filter, the is_imbalanced load fold, and the cache-hit/miss worker scans), and each per-worker access — status, circuit breaker, load — took its own arc_swap guard. At high worker counts that per-worker guard traffic dominated routing CPU. Read each worker once via a new Worker::routing_state() that shares the runtime guard for status+load+processed, gathering the healthy set, load min/max and the min-load index in a single pass. is_imbalanced and the min-load fallback consume the gathered bounds; the cache-hit tenant lookup is a hash-free scan over the gathered healthy indices (url() is a cheap field read). The (load, processed_requests, idx) min-load tie-break from #1714 rides the same guard, so it costs nothing extra. Selection is unchanged except that a cache hit no longer routes onto a Ready-but-circuit-broken worker (it falls through to min-load, like the rest of selection). No-GPU sim (4 threads, 2048 HTTP workers, shared-prefix load), A/B back-to-back: cache_aware routing is now within noise of round_robin at 2048 workers (+0.024 cpu_ms/req marginal, vs +0.18 with a per-request url map and ~+0.40 unoptimized). 26 cache_aware + 107 policy + 207 worker tests pass. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
4057120 to
c3f0f09
Compare
1 participant
Description
Problem
cache_awareworker selection made several O(workers) passes per request — the healthy filter (get_healthy_worker_indices), theis_imbalancedrequest-count fold, and the cache-hit tenant→worker / cache-miss min-load scans. Each per-worker access went through anArcSwapguard:is_healthy()→status()andload()both load the worker'sruntimeguard, andcircuit_breaker_can_execute()loads a second (circuit-breaker) guard. So a healthy worker was read ~3 times per request via separate guarded virtual calls.In a no-GPU benchmark (4 reactor threads, saturating shared-prefix load, mock fleet) this per-worker guard traffic is what makes cache_aware's per-request CPU grow with worker count. gdb sampling at 2048 workers attributed the cost to
arc_swapguard load/drop + per-worker dynamic dispatch in the selection loops — not the radix-tree insert, which is cheap and worker-count-independent.Solution
Gather everything the balanced and imbalanced paths need in a single pass over the workers, reading each worker once via a new
Worker::routing_state()that shares theruntimeguard forstatus+load+processed_requests(the circuit breaker is a separateArcSwap, so it keeps its own guard). The pass collects the healthy indices, the load min/max for the imbalance count-spread, and the min-load index — with the(load, processed_requests, idx)tie-break from #1714, free here sinceprocessedrides the same guard.is_imbalancedand the min-load fallback consume the gathered values instead of re-scanning; a cache hit resolves its tenant→worker with a hash-free linear scan over the gathered healthy indices (url()is a cheap field read — no per-request map).Net: ~3 guarded virtual calls/worker across ~3 passes → 2 guards + 1 virtual call/worker in a single pass. Worker selection is unchanged except that a cache hit no longer routes onto a
Ready-but-circuit-broken worker — it now falls through to min-load, consistent with the rest of selection.Changes
worker/worker.rs: addRoutingState { healthy, can_execute, load, processed }andWorker::routing_state(). The default impl composes the existing accessors;BasicWorkeroverrides it to readstatus/load/processedunder oneruntimeguard.policies/cache_aware.rs:select_workerperforms the single-pass gather;is_imbalancedtakes the precomputed load bounds; the token-tree, string-tree, event-driven and min-load paths consume the gathered min-load index; cache-hit tenant→worker resolution is a hash-free scan over the gathered healthy indices.Test Plan
No-GPU simulation: 4 reactor threads (saturating), mock fleet (
--engine realistic), 16 shared prefixes at 0.8 fraction, 8k offered req/s,cache_awarewith the imbalance valve held open so the balanced (tree-route) path is exercised. Metric is gatewaycpu_ms/req(CPU-time per completed request; lower is better), reported as the marginal overround_robinat the same worker count so tokenization/proxy overhead cancels and the policy cost is isolated. Single-run, so figures carry measurement noise (~±0.07 at these windows); the marginal is the robust signal.Routing marginal over round_robin, HTTP path (string tree — not masked by tokenization, so routing cost is visible). The shipped (hash-free scan) and the earlier (per-request map) revisions were A/B'd at 2048 workers back-to-back under identical conditions:
url→indexmap (earlier revision)The shipped version's cache_aware routing is within noise of round_robin at 2048 workers (+0.024 cpu_ms/req), down from +0.181 with the per-request map and ~+0.40 unoptimized — the per-worker
ArcSwapguard traffic that scaled with worker count is gone. At 64 workers the marginal is within measurement noise for all variants. gRPC path (token tree) is unchanged — per-request tokenization dominates there either way. (Simulator numbers; absolute magnitudes need a live A/B.)Gates on this branch (rebased on
main):Full
cargo test -p smg --libshows one unrelated failure,middleware::metrics::tests::distinct_ids_on_matched_route_do_not_grow_interner, which passes in isolation (a shared-interner test sensitive to parallel execution) and is untouched by this change.--all-featuresis not run locally (it pulls in theopencv-videofeature absent in this environment, per CONTRIBUTING); CI runs the full matrix.Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit