From 6748c8e79a765d5e9c6d03a13600715534864d62 Mon Sep 17 00:00:00 2001 From: Chang Su Date: Wed, 22 Apr 2026 07:50:52 -0700 Subject: [PATCH 1/4] feat(smg): add WorkerSyncAdapter for v2 worker: CRDT namespace MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bridges the `worker:` CRDT namespace (MeshKV) and the in-process WorkerRegistry. Outbound APIs (`on_worker_changed`, `on_worker_removed`) bincode-serialise WorkerState and write to the namespace. An inbound task spawned by `start` routes remote updates through WorkerRegistry::on_remote_worker_state — the same sink the existing WorkerStateSubscriber implementation feeds — so URL deduplication, health promotion, and the Registered event fan-out stay unchanged. No server.rs wiring yet. Step 5 (dual-mode deployment) will wire the adapter alongside teardown of the v1 MeshSyncManager worker sync path. Third in the Step 4 adapter sequence; builds on the auto-registered config: prefix merged in the previous PR. Signed-off-by: Chang Su --- model_gateway/src/lib.rs | 1 + model_gateway/src/mesh/mod.rs | 12 ++ model_gateway/src/mesh/worker_sync.rs | 252 ++++++++++++++++++++++++++ 3 files changed, 265 insertions(+) create mode 100644 model_gateway/src/mesh/mod.rs create mode 100644 model_gateway/src/mesh/worker_sync.rs diff --git a/model_gateway/src/lib.rs b/model_gateway/src/lib.rs index d227a5931c..9403807130 100644 --- a/model_gateway/src/lib.rs +++ b/model_gateway/src/lib.rs @@ -1,6 +1,7 @@ pub mod app_context; pub mod config; pub mod memory; +pub mod mesh; pub mod middleware; pub mod observability; pub mod policies; diff --git a/model_gateway/src/mesh/mod.rs b/model_gateway/src/mesh/mod.rs new file mode 100644 index 0000000000..836302da3f --- /dev/null +++ b/model_gateway/src/mesh/mod.rs @@ -0,0 +1,12 @@ +//! Gateway-side bridges between the v2 mesh KV and local registries. +//! +//! Each adapter scopes to one CRDT or stream prefix, serialises +//! state into the shared merge format, and routes remote updates +//! into the gateway-local component (worker registry, policy +//! registry, rate limiter, tree cache). Adapters are built against +//! the raw `smg_mesh` primitives so the gateway can migrate one +//! domain at a time without moving the v1 `MeshSyncManager` glue. + +pub mod worker_sync; + +pub use worker_sync::WorkerSyncAdapter; diff --git a/model_gateway/src/mesh/worker_sync.rs b/model_gateway/src/mesh/worker_sync.rs new file mode 100644 index 0000000000..44c6e35eb3 --- /dev/null +++ b/model_gateway/src/mesh/worker_sync.rs @@ -0,0 +1,252 @@ +//! `worker:` CRDT adapter: gateway ↔ mesh bridge for `WorkerState`. +//! +//! Outbound: `on_worker_changed` bincode-serialises a `WorkerState` +//! and writes it under `worker:{worker_id}`. `on_worker_removed` +//! writes a tombstone. Both map 1:1 to the v1 +//! `MeshSyncManager::sync_worker_state` / `remove_worker_state` +//! behaviour but target the typed CRDT namespace instead of the +//! untyped store. +//! +//! Inbound: `start` spawns a task that subscribes to the namespace +//! and routes each non-tombstone update through +//! `WorkerRegistry::on_remote_worker_state` — the same sink the v1 +//! `WorkerStateSubscriber` wires up, so registry behaviour +//! (URL-dedupe, health promotion, `Registered` event fan-out) is +//! unchanged. Tombstones are logged at debug; the registry does not +//! yet expose a remote-remove hook, and wiring one belongs in a +//! later PR together with the v1 mirror teardown. +//! +//! The adapter writes through `CrdtNamespace::put`, which fires +//! local subscribers in addition to gossiping. A local write +//! therefore echoes back through `start`'s loop and lands in +//! `on_remote_worker_state`. That path is idempotent — the URL +//! lookup short-circuits to a health refresh — so the loop is +//! self-limiting and matches the v1 behaviour where a local +//! `sync_worker_state` also re-appears in the same node's store +//! before gossip fans out. + +use std::sync::Arc; + +use bytes::Bytes; +use smg_mesh::{CrdtNamespace, WorkerState, WorkerStateSubscriber}; +use tracing::{debug, warn}; + +use crate::worker::WorkerRegistry; + +const PREFIX: &str = "worker:"; + +/// Bridge between the `worker:` CRDT namespace and the gateway's +/// in-process `WorkerRegistry`. +pub struct WorkerSyncAdapter { + workers: Arc, + worker_registry: Arc, +} + +impl std::fmt::Debug for WorkerSyncAdapter { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("WorkerSyncAdapter") + .field("prefix", &self.workers.prefix()) + .finish_non_exhaustive() + } +} + +impl WorkerSyncAdapter { + /// Build an adapter wrapping a `worker:`-scoped namespace and the + /// gateway's worker registry. Panics if the namespace is not + /// scoped to `worker:` so a mis-wired caller fails fast at + /// startup rather than silently routing updates to the wrong + /// prefix. + pub fn new(workers: Arc, worker_registry: Arc) -> Arc { + assert_eq!( + workers.prefix(), + PREFIX, + "WorkerSyncAdapter requires a namespace scoped to `{PREFIX}`", + ); + Arc::new(Self { + workers, + worker_registry, + }) + } + + /// Spawn the inbound task. It runs until the namespace's + /// subscriber channel closes (only on `MeshKV` drop). + pub fn start(self: &Arc) { + let this = Arc::clone(self); + let mut sub = self.workers.subscribe(""); + #[expect( + clippy::disallowed_methods, + reason = "subscription task ends automatically when the mesh KV drops and closes the channel; no handle needed" + )] + tokio::spawn(async move { + while let Some((key, value)) = sub.receiver.recv().await { + let Some(worker_id) = key.strip_prefix(PREFIX).filter(|s| !s.is_empty()) else { + warn!(key, "worker: subscription yielded unexpected key shape"); + continue; + }; + match value { + Some(fragments) => this.apply_incoming(worker_id, &fragments), + None => debug!( + worker_id, + "remote worker tombstone (no-op pending registry remove-by-mesh hook)" + ), + } + } + debug!("WorkerSyncAdapter subscription closed"); + }); + } + + fn apply_incoming(&self, worker_id: &str, fragments: &[Bytes]) { + // CRDT `put` publishes a single fragment today. The + // subscription API is shared with stream namespaces that can + // deliver chunked payloads, so concatenate defensively; the + // fast path is a single cheap Bytes copy. + let total = fragments.iter().map(Bytes::len).sum(); + let mut bytes = Vec::with_capacity(total); + for frag in fragments { + bytes.extend_from_slice(frag); + } + match bincode::deserialize::(&bytes) { + Ok(state) => self.worker_registry.on_remote_worker_state(&state), + Err(err) => warn!(worker_id, %err, "failed to decode WorkerState"), + } + } + + /// Publish a worker update to the cluster. Callers pass the + /// registry's current state; the adapter owns (de)serialisation + /// and key formatting. + pub fn on_worker_changed(&self, worker_id: &str, state: &WorkerState) { + match bincode::serialize(state) { + Ok(bytes) => self.workers.put(&format!("{PREFIX}{worker_id}"), bytes), + Err(err) => warn!(worker_id, %err, "failed to serialize WorkerState"), + } + } + + /// Publish a tombstone for a worker, removing it from the CRDT. + pub fn on_worker_removed(&self, worker_id: &str) { + self.workers.delete(&format!("{PREFIX}{worker_id}")); + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use smg_mesh::{MergeStrategy, MeshKV}; + use tokio::time::sleep; + + use super::*; + + fn worker_namespace(mesh: &MeshKV) -> Arc { + mesh.configure_crdt_prefix(PREFIX, MergeStrategy::LastWriterWins) + } + + fn sample_state(worker_id: &str, url: &str) -> WorkerState { + WorkerState { + worker_id: worker_id.into(), + model_id: "llama-3".into(), + url: url.into(), + health: true, + load: 0.25, + version: 1, + spec: vec![], + } + } + + #[tokio::test] + async fn on_worker_changed_writes_decodable_state() { + let mesh = MeshKV::new("node-a".into()); + let ns = worker_namespace(&mesh); + let registry = Arc::new(WorkerRegistry::new()); + let adapter = WorkerSyncAdapter::new(ns.clone(), registry); + + let state = sample_state("w1", "http://worker-a:8080"); + adapter.on_worker_changed("w1", &state); + + let raw = ns.get("worker:w1").expect("adapter wrote through to CRDT"); + let decoded: WorkerState = bincode::deserialize(&raw).unwrap(); + assert_eq!(decoded, state); + } + + #[tokio::test] + async fn on_worker_removed_tombstones_the_entry() { + let mesh = MeshKV::new("node-a".into()); + let ns = worker_namespace(&mesh); + let registry = Arc::new(WorkerRegistry::new()); + let adapter = WorkerSyncAdapter::new(ns.clone(), registry); + + let state = sample_state("w1", "http://worker-a:8080"); + adapter.on_worker_changed("w1", &state); + assert!(ns.get("worker:w1").is_some()); + + adapter.on_worker_removed("w1"); + assert!( + ns.get("worker:w1").is_none(), + "tombstone must hide the prior value from readers" + ); + } + + #[tokio::test] + async fn start_routes_remote_state_into_registry() { + // Two adapters over one store mimic a remote node's write + // (publisher) arriving at a local subscriber. The underlying + // store is shared so the publisher's put fires the + // subscriber adapter's subscription. + let mesh = MeshKV::new("node-a".into()); + let ns = worker_namespace(&mesh); + + let registry = Arc::new(WorkerRegistry::new()); + let subscriber = WorkerSyncAdapter::new(ns.clone(), registry.clone()); + subscriber.start(); + + let publisher_registry = Arc::new(WorkerRegistry::new()); + let publisher = WorkerSyncAdapter::new(ns, publisher_registry); + publisher.on_worker_changed("w1", &sample_state("w1", "http://remote:8080")); + + // Subscription fanout is async; poll briefly. + for _ in 0..20 { + if registry.get_by_url("http://remote:8080").is_some() { + return; + } + sleep(Duration::from_millis(10)).await; + } + panic!("registry did not see the remote worker"); + } + + #[tokio::test] + async fn start_ignores_malformed_payload() { + // A bad payload should not propagate into the registry and + // must not kill the spawned task — a subsequent valid write + // still lands. + let mesh = MeshKV::new("node-a".into()); + let ns = worker_namespace(&mesh); + let registry = Arc::new(WorkerRegistry::new()); + let adapter = WorkerSyncAdapter::new(ns.clone(), registry.clone()); + adapter.start(); + + ns.put("worker:bogus", b"not-bincode".to_vec()); + sleep(Duration::from_millis(20)).await; + assert!(registry.get_by_url("http://remote:8080").is_none()); + + let good = sample_state("w1", "http://remote:8080"); + ns.put( + "worker:w1", + bincode::serialize(&good).expect("state serializes"), + ); + for _ in 0..20 { + if registry.get_by_url("http://remote:8080").is_some() { + return; + } + sleep(Duration::from_millis(10)).await; + } + panic!("subscription task aborted after a bad payload"); + } + + #[tokio::test] + #[should_panic(expected = "WorkerSyncAdapter requires a namespace scoped to `worker:`")] + async fn new_rejects_wrong_prefix() { + let mesh = MeshKV::new("node-a".into()); + let ns = mesh.configure_crdt_prefix("policy:", MergeStrategy::LastWriterWins); + let registry = Arc::new(WorkerRegistry::new()); + let _ = WorkerSyncAdapter::new(ns, registry); + } +} From 4e7e665f1e385bc5452c8dc802f550410ac99030 Mon Sep 17 00:00:00 2001 From: Chang Su Date: Wed, 22 Apr 2026 07:58:31 -0700 Subject: [PATCH 2/4] refactor(smg): move mesh adapters under mesh/adapters/ subdir Makes room for non-adapter mesh glue (bootstrap, shutdown, config helpers) to land alongside without mixing with the per-namespace adapters. Pure rename + module re-wire; no behaviour change. Signed-off-by: Chang Su --- model_gateway/src/mesh/adapters/mod.rs | 8 ++++++++ .../src/mesh/{ => adapters}/worker_sync.rs | 0 model_gateway/src/mesh/mod.rs | 15 +++++---------- 3 files changed, 13 insertions(+), 10 deletions(-) create mode 100644 model_gateway/src/mesh/adapters/mod.rs rename model_gateway/src/mesh/{ => adapters}/worker_sync.rs (100%) diff --git a/model_gateway/src/mesh/adapters/mod.rs b/model_gateway/src/mesh/adapters/mod.rs new file mode 100644 index 0000000000..b999715c47 --- /dev/null +++ b/model_gateway/src/mesh/adapters/mod.rs @@ -0,0 +1,8 @@ +//! CRDT- and stream-namespace adapters that bridge `MeshKV` to +//! gateway-local state. Each adapter owns one prefix, serialises +//! domain types into the shared merge format, and routes remote +//! updates into the corresponding registry or cache. + +pub mod worker_sync; + +pub use worker_sync::WorkerSyncAdapter; diff --git a/model_gateway/src/mesh/worker_sync.rs b/model_gateway/src/mesh/adapters/worker_sync.rs similarity index 100% rename from model_gateway/src/mesh/worker_sync.rs rename to model_gateway/src/mesh/adapters/worker_sync.rs diff --git a/model_gateway/src/mesh/mod.rs b/model_gateway/src/mesh/mod.rs index 836302da3f..72696c8948 100644 --- a/model_gateway/src/mesh/mod.rs +++ b/model_gateway/src/mesh/mod.rs @@ -1,12 +1,7 @@ -//! Gateway-side bridges between the v2 mesh KV and local registries. -//! -//! Each adapter scopes to one CRDT or stream prefix, serialises -//! state into the shared merge format, and routes remote updates -//! into the gateway-local component (worker registry, policy -//! registry, rate limiter, tree cache). Adapters are built against -//! the raw `smg_mesh` primitives so the gateway can migrate one -//! domain at a time without moving the v1 `MeshSyncManager` glue. +//! Gateway-side glue for the v2 mesh: adapters that bridge the +//! typed `MeshKV` namespaces to local registries, plus bootstrap +//! and shutdown wiring added in later steps. -pub mod worker_sync; +pub mod adapters; -pub use worker_sync::WorkerSyncAdapter; +pub use adapters::WorkerSyncAdapter; From 5d300b02f0a076a453c317fcd97b42c5ff743b32 Mon Sep 17 00:00:00 2001 From: Chang Su Date: Wed, 22 Apr 2026 08:16:43 -0700 Subject: [PATCH 3/4] fix(smg): backfill WorkerRegistry from `worker:` namespace on start MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `start` previously only processed live subscription events, so a node that joined after workers had already gossiped would wait for the next unrelated write before its registry was populated — the v1 path explicitly replayed via `get_all_worker_states()` for exactly this reason. Subscribe first (so no replay/live race loses an event), then iterate the namespace and feed existing entries through the same `on_remote_worker_state` sink the live loop uses. Add a test that pre-seeds the namespace before `start` and asserts the registry sees the worker. Signed-off-by: Chang Su --- .../src/mesh/adapters/worker_sync.rs | 73 +++++++++++++++---- 1 file changed, 60 insertions(+), 13 deletions(-) diff --git a/model_gateway/src/mesh/adapters/worker_sync.rs b/model_gateway/src/mesh/adapters/worker_sync.rs index 44c6e35eb3..001cbfefea 100644 --- a/model_gateway/src/mesh/adapters/worker_sync.rs +++ b/model_gateway/src/mesh/adapters/worker_sync.rs @@ -68,8 +68,13 @@ impl WorkerSyncAdapter { }) } - /// Spawn the inbound task. It runs until the namespace's - /// subscriber channel closes (only on `MeshKV` drop). + /// Spawn the inbound task. Subscribes first (so no event is lost + /// during replay), then backfills the registry with any entries + /// already sitting in the CRDT — a node that joins after workers + /// have already gossiped would otherwise wait for the next + /// unrelated write before its routing table is complete. The task + /// runs until the namespace's subscriber channel closes (only on + /// `MeshKV` drop). pub fn start(self: &Arc) { let this = Arc::clone(self); let mut sub = self.workers.subscribe(""); @@ -78,13 +83,21 @@ impl WorkerSyncAdapter { reason = "subscription task ends automatically when the mesh KV drops and closes the channel; no handle needed" )] tokio::spawn(async move { + this.backfill_existing(); while let Some((key, value)) = sub.receiver.recv().await { let Some(worker_id) = key.strip_prefix(PREFIX).filter(|s| !s.is_empty()) else { warn!(key, "worker: subscription yielded unexpected key shape"); continue; }; match value { - Some(fragments) => this.apply_incoming(worker_id, &fragments), + Some(fragments) => { + let total = fragments.iter().map(Bytes::len).sum(); + let mut bytes = Vec::with_capacity(total); + for frag in fragments { + bytes.extend_from_slice(&frag); + } + this.apply_incoming(worker_id, &bytes); + } None => debug!( worker_id, "remote worker tombstone (no-op pending registry remove-by-mesh hook)" @@ -95,17 +108,23 @@ impl WorkerSyncAdapter { }); } - fn apply_incoming(&self, worker_id: &str, fragments: &[Bytes]) { - // CRDT `put` publishes a single fragment today. The - // subscription API is shared with stream namespaces that can - // deliver chunked payloads, so concatenate defensively; the - // fast path is a single cheap Bytes copy. - let total = fragments.iter().map(Bytes::len).sum(); - let mut bytes = Vec::with_capacity(total); - for frag in fragments { - bytes.extend_from_slice(frag); + /// Replay every entry currently in the `worker:` namespace into + /// the registry. Safe to run after `subscribe` — the subscription + /// loop is idempotent on repeated state (URL-dedupe short-circuits + /// to a health refresh), so overlap with a live event is fine. + fn backfill_existing(&self) { + for key in self.workers.keys("") { + let Some(worker_id) = key.strip_prefix(PREFIX).filter(|s| !s.is_empty()) else { + continue; + }; + if let Some(bytes) = self.workers.get(&key) { + self.apply_incoming(worker_id, &bytes); + } } - match bincode::deserialize::(&bytes) { + } + + fn apply_incoming(&self, worker_id: &str, bytes: &[u8]) { + match bincode::deserialize::(bytes) { Ok(state) => self.worker_registry.on_remote_worker_state(&state), Err(err) => warn!(worker_id, %err, "failed to decode WorkerState"), } @@ -241,6 +260,34 @@ mod tests { panic!("subscription task aborted after a bad payload"); } + #[tokio::test] + async fn start_backfills_preexisting_entries() { + // Rolling-restart scenario: the `worker:` namespace already + // contains gossiped state before `start` runs. The adapter + // must backfill the registry on spawn, not wait for the next + // live event. + let mesh = MeshKV::new("node-a".into()); + let ns = worker_namespace(&mesh); + + let seeded = sample_state("w-seeded", "http://seeded:8080"); + ns.put( + "worker:w-seeded", + bincode::serialize(&seeded).expect("state serializes"), + ); + + let registry = Arc::new(WorkerRegistry::new()); + let adapter = WorkerSyncAdapter::new(ns, registry.clone()); + adapter.start(); + + for _ in 0..20 { + if registry.get_by_url("http://seeded:8080").is_some() { + return; + } + sleep(Duration::from_millis(10)).await; + } + panic!("registry did not see the pre-existing worker"); + } + #[tokio::test] #[should_panic(expected = "WorkerSyncAdapter requires a namespace scoped to `worker:`")] async fn new_rejects_wrong_prefix() { From cbce753eb7fb899e93249ba738bcf38e6e6dfc87 Mon Sep 17 00:00:00 2001 From: Chang Su Date: Wed, 22 Apr 2026 08:54:41 -0700 Subject: [PATCH 4/4] fix(smg): drain subscription during backfill and warn on bad keys MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move `backfill_existing` out of the spawned recv task so the live loop starts draining immediately. `notify` uses `try_send` into a bounded mpsc (worker: capacity 1000), and a slow backfill inside the same task blocks the recv — on a busy startup that window could drop live updates the gossip path would not resend. Also mirror the live-path warning for malformed keys in `backfill_existing` so startup data-shape issues are observable. Signed-off-by: Chang Su --- .../src/mesh/adapters/worker_sync.rs | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/model_gateway/src/mesh/adapters/worker_sync.rs b/model_gateway/src/mesh/adapters/worker_sync.rs index 001cbfefea..57b43345f9 100644 --- a/model_gateway/src/mesh/adapters/worker_sync.rs +++ b/model_gateway/src/mesh/adapters/worker_sync.rs @@ -68,13 +68,16 @@ impl WorkerSyncAdapter { }) } - /// Spawn the inbound task. Subscribes first (so no event is lost - /// during replay), then backfills the registry with any entries - /// already sitting in the CRDT — a node that joins after workers - /// have already gossiped would otherwise wait for the next - /// unrelated write before its routing table is complete. The task - /// runs until the namespace's subscriber channel closes (only on - /// `MeshKV` drop). + /// Start the inbound path. Subscribes first so no live event is + /// lost, spawns the recv loop so it can start draining + /// immediately, and then backfills from the calling thread — any + /// entry already in the CRDT would otherwise have to wait for the + /// next unrelated write before the registry saw it. Running the + /// backfill outside the spawn keeps the live loop free to drain + /// concurrently: `notify` uses `try_send` into a bounded mpsc, so + /// a blocked recv while backfill is running could drop updates on + /// a busy startup. `on_remote_worker_state` is idempotent on URL, + /// so a key seen by both paths only refreshes health. pub fn start(self: &Arc) { let this = Arc::clone(self); let mut sub = self.workers.subscribe(""); @@ -83,7 +86,6 @@ impl WorkerSyncAdapter { reason = "subscription task ends automatically when the mesh KV drops and closes the channel; no handle needed" )] tokio::spawn(async move { - this.backfill_existing(); while let Some((key, value)) = sub.receiver.recv().await { let Some(worker_id) = key.strip_prefix(PREFIX).filter(|s| !s.is_empty()) else { warn!(key, "worker: subscription yielded unexpected key shape"); @@ -106,15 +108,17 @@ impl WorkerSyncAdapter { } debug!("WorkerSyncAdapter subscription closed"); }); + self.backfill_existing(); } /// Replay every entry currently in the `worker:` namespace into - /// the registry. Safe to run after `subscribe` — the subscription - /// loop is idempotent on repeated state (URL-dedupe short-circuits - /// to a health refresh), so overlap with a live event is fine. + /// the registry. Safe to run alongside the live subscription loop + /// — the sink is idempotent on URL (health refresh short-circuit), + /// so overlap with a concurrent live event is fine. fn backfill_existing(&self) { for key in self.workers.keys("") { let Some(worker_id) = key.strip_prefix(PREFIX).filter(|s| !s.is_empty()) else { + warn!(key, "worker: backfill yielded unexpected key shape"); continue; }; if let Some(bytes) = self.workers.get(&key) {