Repository navigation
refactor(worker): extract HashRing into its own module (Finding 5) - #1129
Merged
Merged
Conversation
… Worker trait
Finding 5 from .claude/plans/2026-04-14-worker-module-followup-cleanup.md.
Pulls the 118-line consistent-hash-ring data structure out of the
2134-line registry.rs into a focused worker/hash_ring.rs, and narrows
its public API so the type no longer depends on the Worker trait.
What changed
- Add model_gateway/src/worker/hash_ring.rs (new file, ~190 lines incl.
tests). Contains VIRTUAL_NODES_PER_WORKER, `HashRing`, `hash_position`,
and six unit tests that cover: empty ring, ring length with virtual
nodes, deterministic lookup, unhealthy-worker skip, all-unhealthy
fallback, and owned-String iterator input.
- model_gateway/src/worker/registry.rs:
* Delete the 118-line HashRing block (struct + impl + const).
* Update the file-header doc comment to reference the new module.
* Drop `use std::collections::HashSet` from the HashRing usage (kept
for other uses in the file).
* Import `hash_ring::HashRing` from the sibling module.
* Update `rebuild_hash_ring` to construct the ring from a URL
iterator: `HashRing::new(workers.value().iter().map(|w| w.url()))`
instead of `HashRing::new(&workers)`.
- model_gateway/src/worker/mod.rs:
* Register `pub mod hash_ring;`.
* Split the re-export — `pub use hash_ring::HashRing;` + `pub use
registry::WorkerRegistry;` — so downstream callers still write
`use crate::worker::HashRing` without touching imports.
- model_gateway/src/policies/prefix_hash.rs (5 test sites) and
model_gateway/src/policies/consistent_hashing.rs (2 test sites):
* Adapt callers to the new iterator-based signature:
`HashRing::new(workers.iter().map(|w| w.url()))`.
API change: `HashRing::new`
- Before: `pub fn new(workers: &[Arc<dyn Worker>]) -> Self`
- After: `pub fn new<I>(urls: I) -> Self where I: IntoIterator, I::Item: AsRef<str>`
Accepts `&[&str]`, `&[String]`, owned `Vec<String>`, or any iterator
of string-like items. The rebuild path in `registry.rs` feeds the
URLs directly from the model index without allocating an intermediate
Vec, matching the previous zero-alloc hot path.
Why
- HashRing's only dependency on the `Worker` trait was `worker.url()`
in its constructor. Threading the full trait object through the
signature coupled a pure consistent-hash data structure to the
gateway's worker abstraction. Narrowing the constructor to a URL
iterator removes that coupling so `HashRing` could, in a later PR,
move to a shared `common/` crate if the gossip/mesh layers ever
need the same ring.
- `registry.rs` had grown to 2134 lines mixing the ring implementation
with worker-registry wiring, model index management, and the state
subscriber. Extracting the ring drops registry.rs by ~115 net lines
and gives reviewers a focused 190-line file to touch when tuning
the ring.
- The inline test module on `hash_ring.rs` fills a real coverage gap
— the ring previously had zero direct tests (the policy tests
exercised it transitively, but a broken ring would surface in
policy-level assertions, not as a localized failure).
Verification
- cargo check -p smg — clean.
- cargo clippy -p smg --all-targets -- -D warnings — clean.
- cargo clippy -p smg-golang --all-targets -- -D warnings — clean.
- cargo check -p smg-python — clean.
- cargo test -p smg --lib — 544 passed; 0 failed; 4 ignored
(+6 new unit tests for HashRing).
Refs: Finding 5 in .claude/plans/2026-04-14-worker-module-followup-cleanup.md
Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughThe PR extracts the HashRing implementation into a new Changes
Sequence Diagram(s)(omitted) Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request refactors the HashRing implementation into a dedicated module and decouples it from the Worker trait by using worker URLs for construction. Feedback indicates that find_healthy_url performs a HashSet allocation on the hot path, which conflicts with the "allocation-free" claim in the documentation. Further improvements were suggested to optimize string formatting during ring initialization using a reusable buffer and to add an early exit in the lookup loop to avoid redundant iterations when all workers are unhealthy.
| //! | ||
| //! The ring maps a routing key to a worker URL using consistent hashing over | ||
| //! virtual nodes. The registry rebuilds one ring per model when workers are | ||
| //! added or removed, so lookups stay allocation-free on the request path. |
Contributor
There was a problem hiding this comment.
The documentation states that lookups stay "allocation-free on the request path", but find_healthy_url (line 109) allocates a HashSet on every call. This allocation occurs on the hot path of every request using consistent hashing. Consider updating the documentation to reflect this or exploring ways to avoid the allocation (e.g., using a stack-allocated structure for small worker counts).
References
- To prevent vulnerabilities from duplicate entries, use data structures that inherently enforce uniqueness, such as HashSet, instead of manually deduplicating collections like Vec.
Comment on lines
+52
to
+61
| for url in iter { | ||
| // Create Arc<str> once per worker, share across all virtual nodes. | ||
| let url: Arc<str> = Arc::from(url.as_ref()); | ||
|
|
||
| for vnode in 0..VIRTUAL_NODES_PER_WORKER { | ||
| let vnode_key = format!("{url}#{vnode}"); | ||
| let pos = Self::hash_position(&vnode_key); | ||
| entries.push((pos, Arc::clone(&url))); | ||
| } | ||
| } |
Contributor
There was a problem hiding this comment.
This loop performs VIRTUAL_NODES_PER_WORKER (150) string allocations per worker due to format!. While this happens during registry rebuilds rather than on the request path, it can be optimized by reusing a single String buffer and using write! to reduce the number of allocations to one per worker. Additionally, using Arc<str> for the URL improves performance by making clones cheap.
Suggested change
| for url in iter { | |
| // Create Arc<str> once per worker, share across all virtual nodes. | |
| let url: Arc<str> = Arc::from(url.as_ref()); | |
| for vnode in 0..VIRTUAL_NODES_PER_WORKER { | |
| let vnode_key = format!("{url}#{vnode}"); | |
| let pos = Self::hash_position(&vnode_key); | |
| entries.push((pos, Arc::clone(&url))); | |
| } | |
| } | |
| let mut vnode_key = String.new(); | |
| for url in iter { | |
| let url: Arc<str> = Arc::from(url.as_ref()); | |
| for vnode in 0..VIRTUAL_NODES_PER_WORKER { | |
| vnode_key.clear(); | |
| use std::fmt::Write as _; | |
| let _ = write!(&mut vnode_key, "{url}#{vnode}"); | |
| let pos = Self::hash_position(&vnode_key); | |
| entries.push((pos, Arc::clone(&url))); | |
| } | |
| } |
References
- For types that are frequently cloned on hot paths and represent a small, repeated set of values (e.g., worker IDs or tenant IDs), use an interned string type like Arc to improve performance by making clones cheap (atomic reference count increments).
Comment on lines
+111
to
+122
| for i in 0..self.entries.len() { | ||
| let (_, url) = &self.entries[(start + i) % self.entries.len()]; | ||
| let url_str: &str = url; | ||
|
|
||
| if !checked_urls.insert(url_str) { | ||
| continue; | ||
| } | ||
|
|
||
| if is_healthy(url_str) { | ||
| return Some(url_str); | ||
| } | ||
| } |
Contributor
There was a problem hiding this comment.
In the worst-case scenario where all workers are unhealthy, this loop iterates through all virtual nodes (worker_count * 150). Since checked_urls already tracks unique workers using a HashSet, the loop can be terminated early once all unique workers have been checked. This significantly improves performance when many workers are down.
for i in 0..self.entries.len() {
let (_, url) = &self.entries[(start + i) % self.entries.len()];
let url_str: &str = url;
if !checked_urls.insert(url_str) {
continue;
}
if is_healthy(url_str) {
return Some(url_str);
}
if checked_urls.len() >= self.worker_count() {
break;
}
}References
- To prevent vulnerabilities from duplicate entries, use data structures that inherently enforce uniqueness, such as HashSet, instead of manually deduplicating collections like Vec.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@model_gateway/src/worker/hash_ring.rs`:
- Around line 3-6: The doc claim about "allocation-free" lookups is incorrect
because find_healthy_url allocates a HashSet each call; either relax the comment
or eliminate the per-call allocation by changing find_healthy_url to avoid
building a HashSet (e.g., iterate ring entries directly and check health via the
existing healthy_urls map/structure, returning the first healthy URL found) or
accept a preallocated/borrowed set as an argument; update the file comment or
the function find_healthy_url accordingly and remove the HashSet allocation
around lines where the HashSet is created.
🪄 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: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 3d6f404c-6d28-473f-97d3-18b77d93e11a
📒 Files selected for processing (5)
model_gateway/src/policies/consistent_hashing.rsmodel_gateway/src/policies/prefix_hash.rsmodel_gateway/src/worker/hash_ring.rsmodel_gateway/src/worker/mod.rsmodel_gateway/src/worker/registry.rs
Comment on lines
+3
to
+6
| //! The ring maps a routing key to a worker URL using consistent hashing over | ||
| //! virtual nodes. The registry rebuilds one ring per model when workers are | ||
| //! added or removed, so lookups stay allocation-free on the request path. | ||
| //! |
There was a problem hiding this comment.
Correct the “allocation-free lookup” claim.
Line 5 says lookups are allocation-free, but Line 109 allocates a HashSet on every lookup. Please either relax the wording or remove per-call allocation in find_healthy_url.
🛠️ Proposed doc fix
-//! added or removed, so lookups stay allocation-free on the request path.
+//! added or removed, so lookups avoid per-request ring reconstruction on the request path.Also applies to: 109-110
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@model_gateway/src/worker/hash_ring.rs` around lines 3 - 6, The doc claim
about "allocation-free" lookups is incorrect because find_healthy_url allocates
a HashSet each call; either relax the comment or eliminate the per-call
allocation by changing find_healthy_url to avoid building a HashSet (e.g.,
iterate ring entries directly and check health via the existing healthy_urls
map/structure, returning the first healthy URL found) or accept a
preallocated/borrowed set as an argument; update the file comment or the
function find_healthy_url accordingly and remove the HashSet allocation around
lines where the HashSet is created.
Addresses an AI review comment on PR #1129. The module-level doc comment claimed "lookups stay allocation-free on the request path", but `HashRing::find_healthy_url` allocates a small bounded `HashSet` on every call to dedupe virtual-node hits while walking the ring. The `HashSet` allocation itself is *pre-existing*: it was copied verbatim from the original `registry.rs:117` implementation. What was new is the "allocation-free" claim — I fabricated it in the module header when moving the code to its own file. The old `registry.rs` docs only said "the ring is rebuilt when workers are added/removed, not per-request" (a true statement about *build* cadence). I over-rewrote that into a false *lookup* claim. Fix - model_gateway/src/worker/hash_ring.rs: * Rewrite the module doc to describe the actual cost model: amortized build, `O(log n)` binary search per lookup, plus a small bounded dedupe set. * Expand the `find_healthy_url` doc with an explicit "Cost per call" section so the allocation is discoverable from the public API. * Inline comment at the `HashSet::with_capacity` site explains why the dedupe set exists and that the capacity is bounded by `worker_count().min(16)`, typically a handful of slots. - model_gateway/src/worker/mod.rs: rustfmt-driven reorder of `pub use hash_ring::HashRing;` into alphabetical position between `error` and `http_client` (this is what was failing the format check on the parent PR). Why not also remove the allocation? - It pre-dates this PR and is not the issue the review flagged. The cost is bounded (~200 bytes for ≤16 slots) and dominates neither the `O(log n)` search nor the blake3 hashing. Optimizing it without profiling data would be premature, and mixing an algorithmic change into a doc-accuracy fix muddies the review. Tracked as a possible follow-up if profiling shows it matters. Verification - make fmt — clean. - cargo fmt --all --check — exit 0. - cargo clippy -p smg --all-targets -- -D warnings — clean. - cargo test -p smg --lib — 544 passed; 0 failed; 4 ignored. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
slin1237
added a commit
that referenced
this pull request
Apr 15, 2026
Audit found ~30 inline comments in worker/ that restated what the
next line of code already says. None of them encoded a non-obvious
invariant, workaround, or design rationale — they were just
narration of the form "// Increment counter" above `counter += 1`.
Following the project rule "default to writing no comments — only
add one when the WHY is non-obvious", this PR removes them so the
remaining comments in worker/ are uniformly load-bearing.
What was removed (per file)
- worker/circuit_breaker.rs (1)
* "// Update last failure time atomically" above an atomic store.
- worker/hash_ring.rs (4)
* "// Create Arc<str> once per worker, share across all virtual
nodes." (the extract-and-clone pattern speaks for itself)
* "// Sort by ring position for binary search."
(sort_unstable_by_key on the position is self-evident)
* "// Take first 8 bytes as u64." (u64::from_le_bytes(hash[..8]))
* "// Binary search to find first entry at or after key_pos."
(partition_point IS binary search)
Kept: the expanded HashSet rationale block at find_healthy_url —
that one explains WHY the dedupe set exists and that its capacity
is bounded by worker_count, which is load-bearing context for
reviewers per the recent allocation-claim review on #1129.
- worker/kv_event_monitor.rs (4)
* "// Remove subscription from handles under lock." in
on_worker_removed (the lock-scoped block already shows it)
* Three "// Store first" / "// Remove" narrations in the
apply_event dispatch tests (variable names already say it).
- worker/metrics_aggregator.rs (1, shortened)
* Tighten "openmetrics_parser doesn't handle colons in metric
names; replace with underscores" → "openmetrics_parser rejects
colons in metric names." (drops the "and replace with
underscores" tail since the .replace(":", "_") right below is
self-evident).
- worker/registry.rs (~16)
* Filter method narration in `WorkerRegistry::list` — four
"// Check X if specified" comments above their `if let
Some(...)` filters.
* "// Get total workers count efficiently from DashMap" above a
plain `.len()` call (the "efficiently" claim wasn't even
actionable — it just decorated the call).
* "// Update type/connection mode index if changed" above their
`if old != new` blocks.
* "// Remove from URL/type/connection mode mapping" narrations
in the worker removal path.
* "// Store worker" above `self.workers.insert(...)`. Kept the
nearby copy-on-write rationale + DashMap-key-ownership notes
because those encode WHY the clones are needed.
* Test setup narration in five tests: "// Create a worker with
labels", "// Remove worker", "// Create workers for different
models", "// Set config for a model", "// Set retry config for
the model", "// Create and register a worker", "// Remove the
worker". Kept the "Last write wins" / "Disable retries" /
"Remove first|last worker — config should still exist|be
cleaned up" comments because those state the test's expected
behavior, not its mechanics.
- worker/token_bucket.rs (3, one rewritten)
* "// Wait for notify signal from return_tokens()" above
`notify.notified().await` — the function-level doc above
already explains the refill_rate=0 path.
* "// Return a token - now we should be able to acquire" test
narration above `bucket.return_tokens(1.0)`.
* "// Wait - should NOT refill automatically" rewritten to
"// refill_rate=0 should NOT refill automatically even after
waiting" — keeps the assertion intent without narrating the
sleep call.
What was deliberately KEPT
- All `///` doc comments and `//!` module docs (out of scope).
- Section-header comments in `worker.rs` that group trait methods
(`// ── Health check counter accessors ──` etc.) — they add
semantic grouping to a 2000-line file.
- Comments explaining lock ordering, atomic ordering rationale,
borrow-checker workarounds, mesh-sync ordering invariants, and
test logic premises ("Last write wins", "Should receive Removed
event").
- The 7 `//` comments in `registry.rs` that document non-obvious
invariants (mesh sync ordering, replace lock lifecycle, model
retry config cleanup intent, etc.).
Net change: 6 files changed, 4 insertions(+), 42 deletions(-).
Verification
- cargo clippy -p smg --all-targets -- -D warnings — clean.
- cargo clippy -p smg-golang --all-targets -- -D warnings — clean.
- cargo check -p smg-python — clean.
- cargo test -p smg --lib — 546 passed; 0 failed; 4 ignored.
- make fmt — clean.
- rustup run nightly cargo fmt -- --check — EXIT=0 (matches CI).
Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
This was referenced Apr 15, 2026
3 of 4 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Finding 5 from
.claude/plans/2026-04-14-worker-module-followup-cleanup.md. Pulls the 118-line consistent-hash-ring data structure out of the 2134-lineregistry.rsinto a focusedworker/hash_ring.rs, and narrows its public API so the type no longer depends on theWorkertrait.This is a pure organizational refactor — zero behavior changes, no hot-path allocation regression, and the new file ships with six direct unit tests that fill a previously-empty test coverage gap for the ring.
What changed
New file
model_gateway/src/worker/hash_ring.rs(~190 lines incl. tests):VIRTUAL_NODES_PER_WORKERconstHashRingstruct + impl (new,find_healthy_url,is_empty,len,worker_count, privatehash_position)#[cfg(test)] mod tests— six new unit tests: empty ring, length-scales-with-virtual-nodes, deterministic lookup, unhealthy-worker skip, all-unhealthy fallback, owned-String iterator inputmodel_gateway/src/worker/registry.rs:HashRingblock (struct + impl + const)hash_ring::HashRingfrom the sibling modulerebuild_hash_ringto construct the ring from a URL iterator:HashRing::new(workers.value().iter().map(|w| w.url()))instead ofHashRing::new(&workers)model_gateway/src/worker/mod.rs:pub mod hash_ring;use crate::worker::HashRingwithout touching their imports:Caller updates (7 test sites total):
model_gateway/src/policies/prefix_hash.rs— 5 test sitesmodel_gateway/src/policies/consistent_hashing.rs— 2 test sitesAll changed from
HashRing::new(&workers)toHashRing::new(workers.iter().map(|w| w.url())).API change:
HashRing::newAccepts
&[&str],&[String], ownedVec<String>, or any iterator of string-like items. The rebuild path inregistry.rsfeeds URLs directly from the model index without allocating an intermediate Vec, matching the previous zero-alloc hot path.Why
HashRing's only dependency on theWorkertrait wasworker.url()in its constructor. Threading the full trait object through the signature coupled a pure consistent-hash data structure to the gateway's worker abstraction. Narrowing the constructor to a URL iterator removes that coupling soHashRingcould, in a later PR, move to a sharedcommon/crate if the gossip/mesh layers ever need the same ring.registry.rshad grown to 2134 lines mixing the ring implementation with worker-registry wiring, model index management, and the state subscriber. Extracting the ring dropsregistry.rsby ~115 net lines and gives reviewers a focused 190-line file to touch when tuning the ring.hash_ring.rsfills a real coverage gap — the ring previously had zero direct tests (the policy tests exercised it transitively, but a broken ring would surface in policy-level assertions, not as a localized failure).Test plan
cargo check -p smg— clean.cargo clippy -p smg --all-targets -- -D warnings— clean.cargo clippy -p smg-golang --all-targets -- -D warnings— clean.cargo check -p smg-python— clean.cargo test -p smg --lib— 544 passed; 0 failed; 4 ignored (+6 new HashRing unit tests compared to main's 538).Refs: Finding 5 in
.claude/plans/2026-04-14-worker-module-followup-cleanup.mdSummary by CodeRabbit