feat: add topology-aware routing policy with prefix hashing and HRW - #1947
feat: add topology-aware routing policy with prefix hashing and HRW#1947Vrinda12-tech wants to merge 3 commits into
Conversation
Signed-off-by: Vrinda <vrindaml539@gmail.com>
github-actions
Bot
added
python-bindings
📝 WalkthroughWalkthroughAdds topology-aware routing for JSON conversations using extracted prefixes, load-aware rendezvous hashing, and an LRU prefix cache. It also exposes routing APIs, adds unit tests for each component, and introduces Criterion benchmarks. ChangesTopology routing
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Request
participant TopologyRouter
participant PrefixCache
participant RendezvousRouter
Request->>TopologyRouter: route(json_bytes)
TopologyRouter->>PrefixCache: lookup(prefix_hash)
alt cache hit
PrefixCache-->>TopologyRouter: worker_id
else cache miss or stale worker
TopologyRouter->>RendezvousRouter: route(prefix_hash)
RendezvousRouter-->>TopologyRouter: WorkerNode
TopologyRouter->>PrefixCache: insert(prefix_hash, worker_id)
end
TopologyRouter-->>Request: selected WorkerNode
Suggested labels: Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
![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 replaces the main gateway implementation with a stub and introduces a new topology-based routing module (TopologyRouter) under model_gateway/src/router/topology/. This module implements Rendezvous (HRW) routing, conversation prefix extraction from JSON payloads, and a prefix cache utilizing mini-moka and DashMap. Feedback on the changes highlights a memory leak in PrefixCache due to LRU evictions not cleaning up the reverse mapping, a correctness bug where overwriting a cache entry fails to update the reverse index, and a concurrency bottleneck in set_worker_load that can be resolved by taking &self instead of &mut self.
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.
| pub struct PrefixCache { | ||
| inner: Cache<u64, String>, | ||
| reverse: DashMap<String, HashSet<u64>>, | ||
| } | ||
|
|
||
| impl PrefixCache { | ||
| pub fn new(capacity: usize) -> Self { | ||
| Self { | ||
| inner: Cache::builder().max_capacity(capacity as u64).build(), | ||
| reverse: DashMap::new(), | ||
| } | ||
| } |
There was a problem hiding this comment.
Memory Leak due to LRU Eviction and Worker ID Cloning Optimization
When entries are evicted from the mini-moka cache (self.inner) due to capacity limits, they are never removed from the self.reverse map. Over time, this will cause self.reverse to grow indefinitely, leading to a memory leak.
To fix this, wrap reverse in an Arc and register an eviction_listener on the cache builder to clean up the reverse index when entries are evicted. Additionally, since worker IDs are frequently cloned on hot paths, use Arc<str> instead of String to make clones cheap.
pub struct PrefixCache {
inner: Cache<u64, std::sync::Arc<str>>,
reverse: std::sync::Arc<DashMap<std::sync::Arc<str>, HashSet<u64>>>,
}
impl PrefixCache {
pub fn new(capacity: usize) -> Self {
let reverse = std::sync::Arc::new(DashMap::new());
let reverse_clone = std::sync::Arc::clone(&reverse);
let inner = Cache::builder()
.max_capacity(capacity as u64)
.eviction_listener(move |key, value, _cause| {
if let Some(mut hashes) = reverse_clone.get_mut(&value) {
hashes.remove(&*key);
}
})
.build();
Self { inner, reverse }
}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).
| pub fn insert(&self, prefix_hash: u64, worker_id: String) { | ||
| self.reverse | ||
| .entry(worker_id.clone()) | ||
| .or_default() | ||
| .insert(prefix_hash); | ||
|
|
||
| self.inner.insert(prefix_hash, worker_id); | ||
| } |
There was a problem hiding this comment.
Correctness Bug on Cache Overwrite and Worker ID Cloning Optimization
When a prefix_hash is overwritten with a new worker_id, the old worker_id's entry in self.reverse is not cleaned up. If the old worker is later removed via remove_worker, it will invalidate the prefix_hash in self.inner, incorrectly removing the cache entry for the new worker.
To fix this, check if there is an existing worker for the prefix_hash and remove it from that worker's reverse set before inserting the new mapping. Also, use Arc<str> instead of String for worker_id to make cloning cheap on this hot path.
pub fn insert(&self, prefix_hash: u64, worker_id: std::sync::Arc<str>) {
if let Some(old_worker_id) = self.inner.get(&prefix_hash) {
if old_worker_id != worker_id {
if let Some(mut hashes) = self.reverse.get_mut(&old_worker_id) {
hashes.remove(&prefix_hash);
}
}
}
self.reverse
.entry(worker_id.clone())
.or_default()
.insert(prefix_hash);
self.inner.insert(prefix_hash, worker_id);
}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).
| pub fn set_worker_load(&mut self, worker_id: &str, load: f32) { | ||
| for worker in self.hrw_router.workers_mut() { | ||
| if worker.id == worker_id { | ||
| worker.set_load(load); | ||
| break; | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
Concurrency Bottleneck on Load Updates
set_worker_load unnecessarily takes &mut self, which prevents updating worker loads concurrently with routing requests (since route takes &self). Since WorkerNode::set_load uses atomic operations and only requires &self, set_worker_load can be changed to take &self to allow fully concurrent, lock-free load updates.
| pub fn set_worker_load(&mut self, worker_id: &str, load: f32) { | |
| for worker in self.hrw_router.workers_mut() { | |
| if worker.id == worker_id { | |
| worker.set_load(load); | |
| break; | |
| } | |
| } | |
| } | |
| pub fn set_worker_load(&self, worker_id: &str, load: f32) { | |
| for worker in self.hrw_router.workers() { | |
| if worker.id == worker_id { | |
| worker.set_load(load); | |
| break; | |
| } | |
| } | |
| } |
![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: fc33c0fef7
ℹ️ 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".
| assert_eq!(server_config.runtime_worker_threads, None); | ||
| } | ||
| fn main() { | ||
| println!("SMG gateway stub (binary not built)"); |
There was a problem hiding this comment.
In the launch contexts I checked, the README and BFCL/Tau2 scripts still invoke smg launch, but this binary now ignores every CLI argument, prints a stub message, and exits successfully without constructing the router or starting the HTTP server. Any packaged gateway built from this commit will look like it started successfully while no gateway is listening.
Useful? React with 👍 / 👎.
| self.reverse | ||
| .entry(worker_id.clone()) | ||
| .or_default() | ||
| .insert(prefix_hash); | ||
|
|
||
| self.inner.insert(prefix_hash, worker_id); |
There was a problem hiding this comment.
When a prefix is reassigned from one worker to another, this adds the hash to the new worker's reverse set but never removes it from the old worker's set. If remove_worker("old-worker") runs later, it invalidates the prefix even though inner currently maps it to the new worker, breaking cache affinity for remapped prefixes; the new overwrite_cleans_reverse test exercises exactly this case.
Useful? React with 👍 / 👎.
| pub fn remove_worker(&self, worker_id: &str) { | ||
| self.prefix_cache.remove_worker(worker_id); |
There was a problem hiding this comment.
When a worker is decommissioned or fails and callers invoke remove_worker, only cached prefixes are invalidated; the worker remains in hrw_router.workers(), so the next cache miss can route straight back to that removed worker. This leaves topology removal ineffective unless there is a separate, currently unexposed update path for TopologyRouter.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@model_gateway/src/router/topology/mod.rs`:
- Around line 52-60: Update Topology::set_worker_load to take &self instead of
&mut self, iterate through hrw_router.workers() rather than workers_mut(), and
continue calling WorkerNode::set_load for the matching worker so concurrent
atomic load updates do not require exclusive router access.
In `@model_gateway/src/router/topology/prefix_cache.rs`:
- Around line 6-17: Update PrefixCache::new to wrap reverse in an Arc and
configure Cache’s eviction_listener to remove each evicted prefix hash from its
worker’s reverse HashSet, handling both size evictions and replacements. Ensure
the listener removes empty worker entries as appropriate, and update PrefixCache
methods to use the shared Arc while preserving existing cache behavior.
🪄 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: db0ccfd4-74fd-4047-919b-fa06cc2374b4
📒 Files selected for processing (11)
bindings/golang/Cargo.tomlbindings/python/Cargo.tomlmodel_gateway/Cargo.tomlmodel_gateway/benches/topology_bench.rsmodel_gateway/src/lib.rsmodel_gateway/src/main.rsmodel_gateway/src/router/topology/consistent_hash.rsmodel_gateway/src/router/topology/error.rsmodel_gateway/src/router/topology/extractor.rsmodel_gateway/src/router/topology/mod.rsmodel_gateway/src/router/topology/prefix_cache.rs
💤 Files with no reviewable changes (2)
- bindings/python/Cargo.toml
- bindings/golang/Cargo.toml
| opencv-video = ["smg/opencv-video"] | ||
| vendored-openssl = ["smg/vendored-openssl"] | ||
|
|
||
| [profile.release] |
| vendored-openssl = ["smg/vendored-openssl"] | ||
|
|
||
| [profile.ci] | ||
| inherits = "release" |
| name = "smg" | ||
| path = "src/main.rs" | ||
|
|
||
| [[bin]] |
| smg-mcp.workspace = true | ||
| kv-index.workspace = true | ||
| smg-data-connector.workspace = true | ||
| llm-multimodal = { workspace = true, default-features = false } |
| @@ -1,15 +1,3 @@ | |||
| pub mod app_context; | |||
There was a problem hiding this comment.
please revert those lib changes
Signed-off-by: Vrinda <vrindaml539@gmail.com>
Signed-off-by: Vrinda <vrindaml539@gmail.com>
c8884dc to
bd7584a
Compare
github-actions
Bot
removed
python-bindings
|
@slin1237 I've reverted the out‑of‑scope files, added docstrings, and fixed the PR feedback. This is now ready for review. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: bd7584af70
ℹ️ 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".
| use smg::router::topology::{ | ||
| extract_conversation_prefix, PrefixCache, RendezvousRouter, TopologyRouter, WorkerNode, | ||
| }; |
There was a problem hiding this comment.
In the current crate root I checked, only pub mod routers is exposed and there is no src/router/mod.rs; Cargo metadata also auto-discovers this new bench target, so compiling benches via cargo bench --bench topology_bench or --all-targets cannot resolve smg::router::topology. Wire the new topology module into the library/module tree, or point the benchmark at an already exposed module, so the added benchmark target can build.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
model_gateway/src/router/topology/extractor.rs (1)
49-55: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSimplify redundant emptiness check.
The
if prefix_messages.is_empty()check is provably unreachable (dead code) becausemessagesis verified to be non-empty on line 30, and all match arms forprefix_endevaluate to>= 1. You can streamline the logic by removing the branch.♻️ Proposed refactor
- let prefix_messages = &messages[..prefix_end]; - let messages_to_hash = if prefix_messages.is_empty() { - messages - } else { - prefix_messages - }; + let messages_to_hash = &messages[..prefix_end];🤖 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 `@model_gateway/src/router/topology/extractor.rs` around lines 49 - 55, In the message selection logic surrounding prefix_end, remove the redundant prefix_messages.is_empty() conditional and use prefix_messages directly as messages_to_hash. Preserve the existing non-empty validation and prefix_end calculations.
🤖 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 `@model_gateway/src/router/topology/prefix_cache.rs`:
- Around line 68-77: Update remove_worker to pass the existing worker_id: &str
directly to self.reverse.remove, removing the temporary Arc<str> conversion
while preserving the existing hash invalidation behavior.
---
Outside diff comments:
In `@model_gateway/src/router/topology/extractor.rs`:
- Around line 49-55: In the message selection logic surrounding prefix_end,
remove the redundant prefix_messages.is_empty() conditional and use
prefix_messages directly as messages_to_hash. Preserve the existing non-empty
validation and prefix_end calculations.
🪄 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: 6f7c601b-1e25-4fbc-bf1f-74c9effcf5e3
📒 Files selected for processing (5)
model_gateway/src/router/topology/consistent_hash.rsmodel_gateway/src/router/topology/error.rsmodel_gateway/src/router/topology/extractor.rsmodel_gateway/src/router/topology/mod.rsmodel_gateway/src/router/topology/prefix_cache.rs
| /// Removes all entries belonging to a specific worker. | ||
| pub fn remove_worker(&self, worker_id: &str) { | ||
| let worker_id: Arc<str> = worker_id.into(); | ||
|
|
||
| if let Some((_, hashes)) = self.reverse.remove(&worker_id) { | ||
| for hash in hashes { | ||
| self.inner.invalidate(&hash); | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win
Avoid allocating an Arc<str> for lookups.
DashMap::remove accepts any borrowed form of the key. Since Arc<str> implements Borrow<str>, you can pass the &str directly to remove instead of unnecessarily allocating a new Arc<str> on the heap.
⚡ Proposed fix to avoid allocation
/// Removes all entries belonging to a specific worker.
pub fn remove_worker(&self, worker_id: &str) {
- let worker_id: Arc<str> = worker_id.into();
-
- if let Some((_, hashes)) = self.reverse.remove(&worker_id) {
+ if let Some((_, hashes)) = self.reverse.remove(worker_id) {
for hash in hashes {
self.inner.invalidate(&hash);
}
}
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| /// Removes all entries belonging to a specific worker. | |
| pub fn remove_worker(&self, worker_id: &str) { | |
| let worker_id: Arc<str> = worker_id.into(); | |
| if let Some((_, hashes)) = self.reverse.remove(&worker_id) { | |
| for hash in hashes { | |
| self.inner.invalidate(&hash); | |
| } | |
| } | |
| } | |
| /// Removes all entries belonging to a specific worker. | |
| pub fn remove_worker(&self, worker_id: &str) { | |
| if let Some((_, hashes)) = self.reverse.remove(worker_id) { | |
| for hash in hashes { | |
| self.inner.invalidate(&hash); | |
| } | |
| } | |
| } |
🤖 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 `@model_gateway/src/router/topology/prefix_cache.rs` around lines 68 - 77,
Update remove_worker to pass the existing worker_id: &str directly to
self.reverse.remove, removing the temporary Arc<str> conversion while preserving
the existing hash invalidation behavior.
slin1237
left a comment
slin1237
left a comment
There was a problem hiding this comment.
Thanks for the contribution — routing performance is something we care about, and the HRW + load-weighting idea is worth exploring. That said, this can't merge in its current shape, for one mechanical reason and one architectural one.
It doesn't compile in-tree. The new model_gateway/src/router/topology/ module is never declared (no lib.rs/mod.rs change), so cargo never builds the library files at all — while the bench is auto-discovered and fails CI with error[E0433]: cannot find 'router' in 'smg' (see the unit-tests job). mini_moka is also imported but not added to any Cargo.toml. The "tests pass / no warnings" results in the description could only have come from building these files as a standalone crate.
It re-implements subsystems that already exist, as a parallel island that no request can reach:
RendezvousRouter↔policies/consistent_hashing.rs(session affinity over a ring prebuilt inWorkerRegistry, with header routing keys)PrefixCache+ xxh64 fingerprinting ↔policies/cache_aware.rs,policies/prefix_hash.rs, and thekv_indexcrate's radix trees — and note the exact-hash lookup here is strictly weaker than what those provide: it only gives affinity for byte-identical prefixes, whilekv_indexdoes longest-prefix matchingextract_conversation_prefix(json_bytes)re-parses the raw body per request — requests are already parsed once at the boundary, and policies receive routing text/token ids viaSelectWorkerInfoWorkerNode { id, address, load }↔ theWorkertrait /WorkerRegistry— bypassing health checks, circuit breakers, and the existing load accounting- load-weighted selection ↔
least_load.rs/power_of_two.rs
This also means the sub-microsecond benchmarks measure the island in isolation rather than the policy interface the gateway actually routes through.
Two directions that would be very welcome:
- If load-weighted HRW beats the existing ring hashing for a workload you can show, contribute it as a
LoadBalancingPolicyimplementation insidemodel_gateway/src/policies/— a sibling ofconsistent_hashing.rs, operating onArc<dyn Worker>+SelectWorkerInfo(both the HTTP and gRPC arms), registered throughPolicyConfig+ the factory + CLI, using thexxhash-rustdependency already in the workspace. No new cache and no new worker type needed — that's roughly a 150-line policy plus tests. - If the goal is prefix/cache affinity, that machinery lives in
cache_aware/prefix_hash/kv_index— improvements there would be gladly reviewed.
Happy to give pointers on either path if you want to rework it.
Introducing Topology‑Aware Routing for SMG
Adds a high-performance, lock-free topology router featuring Highest Random Weight (HRW) hashing and bi‑directional LRU prefix caching.
Key Features
xxhash64for fast, deterministic prefix fingerprintsmini-mokawith automatic eviction and reverse index cleanupPerformance Benchmarks (Criterion)
All hot paths operate in sub-microsecond time with zero heap allocations during routing.
Files Added
model_gateway/src/router/topology/— Core topology modulemodel_gateway/benches/topology_bench.rs— Criterion benchmarksDependencies Added
xxhash-rustserde_jsonthiserrormini-mokadashmapcriterion(dev)Testing
cargo bench)Summary by CodeRabbit