From 98f9002389835a4d3c4f0bbd915af0f7956d4c38 Mon Sep 17 00:00:00 2001 From: Anna Tchernych Date: Sun, 25 Oct 2026 16:34:20 -0700 Subject: [PATCH 1/4] feat(global-router): let router replicas register with the discovery service Add the design's `RegisterRouter` call as `POST /routers`. A Global Router replica sends `{"replica": "", "url": ""}`, the shape of one `GET /routers` entry, on a timer. Each call is also its heartbeat, and the reply carries the lease, as for `POST /pools`. - Memory store: `MemoryReplicaStore` lists registered replicas until their lease ends, plus the `--router` replicas. A registration replaces a `--router` entry with the same ID while it lasts. - Kubernetes stores: `POST /routers` answers 501 and says why. MCS already lists every router replica from the router Service's EndpointSlices, so replicas need not register. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Anna Tchernych --- .../src/backends/kubernetes.rs | 2 + .../src/backends/memory.rs | 311 +++++++++++++++++- lib/global-discovery-service/src/main.rs | 21 +- .../src/pool_discovery.rs | 8 +- .../src/pool_discovery/replicas.rs | 18 +- lib/global-discovery-service/src/service.rs | 131 +++++++- 6 files changed, 466 insertions(+), 25 deletions(-) diff --git a/lib/global-discovery-service/src/backends/kubernetes.rs b/lib/global-discovery-service/src/backends/kubernetes.rs index b5312121401d..e7fd07b907e6 100644 --- a/lib/global-discovery-service/src/backends/kubernetes.rs +++ b/lib/global-discovery-service/src/backends/kubernetes.rs @@ -47,6 +47,8 @@ //! pools: Arc::new(McsPoolDirectory::new(client.clone(), mcs, suffix, true, locations)), //! routers: Arc::new(EndpointSliceReplicaDirectory::imported(client, &router, mcs)), //! registrar: None, +//! // `POST /routers` answers that router replicas need not register. +//! router_registrar: None, //! authenticator: Arc::new(TrustedNetworkAuthenticator), //! }; //! service::serve(listener, store, shutdown).await?; diff --git a/lib/global-discovery-service/src/backends/memory.rs b/lib/global-discovery-service/src/backends/memory.rs index 8c1f25d80484..0ec3716f4406 100644 --- a/lib/global-discovery-service/src/backends/memory.rs +++ b/lib/global-discovery-service/src/backends/memory.rs @@ -14,8 +14,8 @@ use std::time::Duration; use crate::pool_discovery::watch::{Directory, Snapshot, WatchStream}; use crate::pool_discovery::{ Completeness, LeasePolicy, PoolDiscoveryError, PoolEvent, PoolRecord, PoolRegistrar, - Registration, ReplicaEndpoint, ReplicaEvent, ReplicaId, Revision, VerifiedRelayIdentity, - WatchEvent, + Registration, ReplicaEndpoint, ReplicaEvent, ReplicaId, ReplicaRegistrar, Revision, + VerifiedRelayIdentity, WatchEvent, }; use async_trait::async_trait; use dynamo_global_router_shared::PoolKey; @@ -70,10 +70,7 @@ impl MemoryPoolStore { /// Fails unless `heartbeat_interval` is above zero and `lease_duration` /// is longer than it. pub fn new(lease: LeasePolicy) -> anyhow::Result { - anyhow::ensure!( - !lease.heartbeat_interval.is_zero() && lease.lease_duration > lease.heartbeat_interval, - "the lease duration must be longer than the heartbeat interval, which must be above zero" - ); + check_lease(&lease)?; Ok(Self { lease, pools: Arc::default(), @@ -106,6 +103,16 @@ impl Drop for MemoryPoolStore { } } +/// `heartbeat_interval` must be above zero and `lease_duration` longer than +/// it, so the sweep has a period and a registration outlives one interval. +fn check_lease(lease: &LeasePolicy) -> anyhow::Result<()> { + anyhow::ensure!( + !lease.heartbeat_interval.is_zero() && lease.lease_duration > lease.heartbeat_interval, + "the lease duration must be longer than the heartbeat interval, which must be above zero" + ); + Ok(()) +} + /// Removes expired records and publishes `Removed` for each. async fn sweep( pools: Arc>, @@ -287,6 +294,211 @@ impl Directory for StaticReplicaDirectory { } } +/// Router replicas in memory: the ones from configuration, and the ones +/// that register (`POST /routers`). A registration expires `lease_duration` +/// after its last renewal, and while it lasts it replaces a configured +/// replica with the same ID. A restarted process lists only the configured +/// replicas until the others register again, so lists are `Estimated`. +/// +/// Any authenticated caller may register any replica ID: like +/// [`MemoryPoolStore`], V1 assumes a trusted network. +pub struct MemoryReplicaStore { + lease: LeasePolicy, + replicas: Arc>, + /// Sent while holding the `replicas` lock, as in [`MemoryPoolStore`]. + events: broadcast::Sender, + sweeping: OnceLock<()>, + stop: CancellationToken, +} + +struct Replicas { + configured: HashMap, + registered: HashMap, +} + +struct RegisteredReplica { + endpoint: ReplicaEndpoint, + expires_at: Instant, +} + +impl Replicas { + /// A live registration, else the configured replica. + fn served(&self, id: &ReplicaId, now: Instant) -> Option<&ReplicaEndpoint> { + self.registered + .get(id) + .filter(|registered| registered.expires_at > now) + .map(|registered| ®istered.endpoint) + .or_else(|| self.configured.get(id)) + } + + fn snapshot(&self, now: Instant) -> Snapshot { + let registered = |id: &ReplicaId| { + self.registered + .get(id) + .is_some_and(|registered| registered.expires_at > now) + }; + Snapshot { + items: self + .registered + .values() + .filter(|registered| registered.expires_at > now) + .map(|registered| registered.endpoint.clone()) + .chain( + self.configured + .values() + .filter(|endpoint| !registered(&endpoint.replica)) + .cloned(), + ) + .collect(), + completeness: Completeness::Estimated, + } + } +} + +impl MemoryReplicaStore { + /// `configured` replicas are listed until shutdown. Fails, like + /// [`MemoryPoolStore::new`], for a lease without room for heartbeats. + pub fn new(configured: Vec, lease: LeasePolicy) -> anyhow::Result { + check_lease(&lease)?; + let configured = configured + .into_iter() + .map(|endpoint| (endpoint.replica.clone(), endpoint)) + .collect(); + Ok(Self { + lease, + replicas: Arc::new(Mutex::new(Replicas { + configured, + registered: HashMap::new(), + })), + events: broadcast::channel(EVENT_BUFFER).0, + sweeping: OnceLock::new(), + stop: CancellationToken::new(), + }) + } + + fn lock(&self) -> MutexGuard<'_, Replicas> { + self.replicas.lock().unwrap_or_else(PoisonError::into_inner) + } + + /// Every trait method runs in a Tokio runtime, so the spawn is safe. + fn ensure_sweeping(&self) { + self.sweeping.get_or_init(|| { + tokio::spawn(sweep_replicas( + Arc::clone(&self.replicas), + self.events.clone(), + self.lease.heartbeat_interval, + self.stop.clone(), + )); + }); + } +} + +impl Drop for MemoryReplicaStore { + fn drop(&mut self) { + self.stop.cancel(); + } +} + +/// Removes expired registrations. A configured replica goes back to its +/// configured address; any other replica is `Removed`. +async fn sweep_replicas( + replicas: Arc>, + events: broadcast::Sender, + every: Duration, + stop: CancellationToken, +) { + let mut ticks = tokio::time::interval(every); + ticks.set_missed_tick_behavior(MissedTickBehavior::Delay); + loop { + tokio::select! { + biased; + () = stop.cancelled() => return, + _ = ticks.tick() => {} + } + let now = Instant::now(); + let mut replicas = replicas.lock().unwrap_or_else(PoisonError::into_inner); + let Replicas { + configured, + registered, + } = &mut *replicas; + registered.retain(|id, registered| { + let live = registered.expires_at > now; + if !live { + // No subscribers is not an error. + match configured.get(id) { + None => { + let _ = events.send(WatchEvent::Removed(id.clone())); + } + Some(endpoint) if *endpoint != registered.endpoint => { + let _ = events.send(WatchEvent::Added(endpoint.clone())); + } + Some(_) => {} + } + } + live + }); + } +} + +#[async_trait] +impl ReplicaRegistrar for MemoryReplicaStore { + /// An unchanged registration only renews its lease. A new or changed one + /// is an `Added` event. + async fn register( + &self, + endpoint: ReplicaEndpoint, + ) -> Result { + self.ensure_sweeping(); + let now = Instant::now(); + let mut replicas = self.lock(); + let changed = replicas.served(&endpoint.replica, now) != Some(&endpoint); + replicas.registered.insert( + endpoint.replica.clone(), + RegisteredReplica { + endpoint: endpoint.clone(), + expires_at: now + self.lease.lease_duration, + }, + ); + if changed { + // No subscribers is not an error. + let _ = self.events.send(WatchEvent::Added(endpoint)); + } + Ok(Registration { lease: self.lease }) + } +} + +#[async_trait] +impl Directory for MemoryReplicaStore { + async fn check_connection(&self) -> Result<(), PoolDiscoveryError> { + Ok(()) + } + + async fn list(&self) -> Result, PoolDiscoveryError> { + Ok(self.lock().snapshot(Instant::now())) + } + + async fn get(&self, key: &ReplicaId) -> Result, PoolDiscoveryError> { + Ok(self.lock().served(key, Instant::now()).cloned()) + } + + async fn list_and_watch( + &self, + cancel_token: Option, + ) -> Result, PoolDiscoveryError> { + self.ensure_sweeping(); + let replicas = Arc::clone(&self.replicas); + let events = self.events.clone(); + Ok(watch_stream( + move || { + let replicas = replicas.lock().unwrap_or_else(PoisonError::into_inner); + (events.subscribe(), Some(replicas.snapshot(Instant::now()))) + }, + cancel_token.unwrap_or_default(), + self.stop.clone(), + )) + } +} + #[cfg(test)] mod tests { use crate::pool_discovery::{ @@ -449,6 +661,10 @@ mod tests { MemoryPoolStore::new(lease).is_err(), "{heartbeat_interval}s / {lease_duration}s" ); + assert!( + MemoryReplicaStore::new(Vec::new(), lease).is_err(), + "{heartbeat_interval}s / {lease_duration}s" + ); } } @@ -480,4 +696,87 @@ mod tests { cancel.cancel(); assert!(watch.next().await.is_none()); } + + fn router(id: &str, url: &str) -> ReplicaEndpoint { + ReplicaEndpoint { + replica: ReplicaId(id.to_owned()), + url: Url::parse(url).unwrap(), + } + } + + #[tokio::test(start_paused = true)] + async fn a_registered_router_is_listed_until_its_lease_ends() { + let store = MemoryReplicaStore::new(Vec::new(), LEASE).unwrap(); + let replica = router("ohio/global-router-0", "http://10.42.0.5:50051"); + let mut watch = store.list_and_watch(None).await.unwrap(); + assert!( + matches!(watch.next().await, Some(Ok(WatchEvent::Resync(s))) if s.items.is_empty()) + ); + + let registration = store.register(replica.clone()).await.unwrap(); + assert_eq!(registration.lease, LEASE); + assert!(matches!(watch.next().await, Some(Ok(WatchEvent::Added(r))) if r == replica)); + + // Heartbeats every 10 s keep the 15 s lease alive; renewing an + // unchanged registration is not a change. + for _ in 0..3 { + tokio::time::advance(Duration::from_secs(10)).await; + store.register(replica.clone()).await.unwrap(); + } + assert_eq!( + store.list().await.unwrap().items, + std::slice::from_ref(&replica) + ); + + // Without heartbeats it ends: reads skip it, then the sweep removes it. + tokio::time::advance(Duration::from_secs(16)).await; + assert!(store.get(&replica.replica).await.unwrap().is_none()); + assert!( + matches!(watch.next().await, Some(Ok(WatchEvent::Removed(id))) if id == replica.replica) + ); + } + + #[tokio::test(start_paused = true)] + async fn a_registration_replaces_a_configured_router_while_it_lasts() { + let configured = router("ohio/global-router-0", "http://10.42.0.5:50051"); + let moved = router("ohio/global-router-0", "http://10.42.0.9:50051"); + let other = router("west/global-router-0", "http://10.43.0.5:50051"); + let store = MemoryReplicaStore::new(vec![configured.clone()], LEASE).unwrap(); + let mut watch = store.list_and_watch(None).await.unwrap(); + assert!(matches!(watch.next().await, Some(Ok(WatchEvent::Added(r))) if r == configured)); + assert!(matches!( + watch.next().await, + Some(Ok(WatchEvent::Resync(_))) + )); + + // Registering the configured address is not a change. + store.register(configured.clone()).await.unwrap(); + store.register(moved.clone()).await.unwrap(); + store.register(other.clone()).await.unwrap(); + assert!(matches!(watch.next().await, Some(Ok(WatchEvent::Added(r))) if r == moved)); + assert!(matches!(watch.next().await, Some(Ok(WatchEvent::Added(r))) if r == other)); + assert_eq!( + store.get(&configured.replica).await.unwrap(), + Some(moved.clone()) + ); + + // When the registrations end, the configured address comes back and + // the other replica goes. + tokio::time::advance(Duration::from_secs(16)).await; + let mut changes = Vec::new(); + for _ in 0..2 { + changes.push(watch.next().await.unwrap().unwrap()); + } + assert!( + changes + .iter() + .any(|event| matches!(event, WatchEvent::Added(r) if *r == configured)) + ); + assert!( + changes + .iter() + .any(|event| matches!(event, WatchEvent::Removed(id) if *id == other.replica)) + ); + assert_eq!(store.list().await.unwrap().items, [configured]); + } } diff --git a/lib/global-discovery-service/src/main.rs b/lib/global-discovery-service/src/main.rs index 0674a4c091c5..c709418a9e29 100644 --- a/lib/global-discovery-service/src/main.rs +++ b/lib/global-discovery-service/src/main.rs @@ -18,9 +18,10 @@ use clap::{Parser, ValueEnum}; use dynamo_global_discovery_service::backends::kubernetes::{ EndpointSliceReplicaDirectory, McsPoolDirectory, RouterService, TrustedNetworkAuthenticator, }; -use dynamo_global_discovery_service::backends::memory::{MemoryPoolStore, StaticReplicaDirectory}; +use dynamo_global_discovery_service::backends::memory::{MemoryPoolStore, MemoryReplicaStore}; use dynamo_global_discovery_service::pool_discovery::{ LeasePolicy, McsImplementation, PoolDirectory, PoolRegistrar, ReplicaEndpoint, ReplicaId, + ReplicaRegistrar, }; use dynamo_global_discovery_service::service::{self, DiscoveryStore}; use kube::Client; @@ -37,14 +38,15 @@ struct Args { /// Where pools and router replicas come from. #[arg(long, value_enum, default_value_t = Store::Memory)] store: Store, - /// Memory store: how often a PoolRelay registers again. + /// Memory store: how often a PoolRelay or router replica registers again. #[arg(long, default_value_t = 5)] heartbeat_interval_secs: u64, /// Memory store: how long a registration lasts. #[arg(long, default_value_t = 15)] lease_duration_secs: u64, - /// Memory store: a Global Router replica as `=`. Repeat it for - /// each replica. + /// Memory store: a Global Router replica as `=`, listed until + /// shutdown. Repeat it for each replica. Replicas can also register + /// themselves with `POST /routers`. #[arg(long = "router", value_parser = parse_router)] routers: Vec, /// Kubernetes stores: each PoolRelay Service is named @@ -64,7 +66,8 @@ struct Args { #[derive(Clone, Copy, ValueEnum)] enum Store { - /// Registrations in this process; router replicas from `--router`. + /// Registrations in this process, from pools and router replicas; router + /// replicas also from `--router`. Memory, /// Imported EndpointSlices, named the Karmada way. No registrations. Karmada, @@ -131,12 +134,15 @@ fn memory_store(args: &Args) -> Result { { bail!("--router {} is given twice", repeated.replica.0); } + let routers = Arc::new(MemoryReplicaStore::new(args.routers.clone(), lease)?); let pools: Arc = store.clone(); let registrar: Arc = store; + let router_registrar: Arc = routers.clone(); Ok(DiscoveryStore { pools, - routers: Arc::new(StaticReplicaDirectory::new(args.routers.clone())), + routers, registrar: Some(registrar), + router_registrar: Some(router_registrar), authenticator: Arc::new(TrustedNetworkAuthenticator), }) } @@ -173,6 +179,9 @@ async fn kubernetes_store(args: &Args, mcs: McsImplementation) -> Result; pub type ReplicaEvent = WatchEvent; + +/// Write side behind `RegisterRouter`. A router replica calls `register` on +/// a timer, and each call is also its heartbeat, as with +/// [`super::PoolRegistrar`]. Only stores that cannot find router replicas +/// themselves take registrations: in Kubernetes mode, the router Service's +/// EndpointSlices that MCS imports already list every replica. +#[async_trait] +pub trait ReplicaRegistrar: Send + Sync { + async fn register(&self, endpoint: ReplicaEndpoint) + -> Result; +} diff --git a/lib/global-discovery-service/src/service.rs b/lib/global-discovery-service/src/service.rs index b05291ac9f6c..4d0467c2fa4a 100644 --- a/lib/global-discovery-service/src/service.rs +++ b/lib/global-discovery-service/src/service.rs @@ -6,6 +6,9 @@ //! - `POST /pools`: a PoolRelay registers its pool. It calls this on a timer, //! and each call is also the pool's heartbeat. //! - `GET /routers`: the Global Router replicas a PoolRelay connects to. +//! - `POST /routers`: a Global Router replica registers itself, on a timer +//! like `POST /pools`. Only a store that cannot list router replicas +//! itself takes it; the MCS store answers `501`, since MCS lists them. //! - `GET /pools`: the pools a Global Router routes to. //! - `GET /healthz`: OK while the service can answer from its store. Use it //! for readiness; liveness needs only the TCP port. @@ -26,8 +29,8 @@ use std::sync::Arc; use crate::pool_discovery::{ Capacity, Completeness, LeasePolicy, PeerCredentials, PoolDirectory, PoolDiscoveryError, - PoolRecord, PoolRegistrar, RelayAuthenticator, ReplicaDirectory, ReplicaEndpoint, Revision, - verify, + PoolRecord, PoolRegistrar, RelayAuthenticator, ReplicaDirectory, ReplicaEndpoint, ReplicaId, + ReplicaRegistrar, Revision, verify, }; use axum::extract::{ConnectInfo, State}; use axum::http::StatusCode; @@ -47,14 +50,17 @@ pub struct DiscoveryStore { pub routers: Arc, /// `None` when pools are announced another way, as with MCS. pub registrar: Option>, - /// Names the caller of `POST /pools` from its connection. + /// `None` when the store lists router replicas itself, as with MCS. + pub router_registrar: Option>, + /// Names the caller of `POST /pools` and `POST /routers` from its + /// connection. pub authenticator: Arc, } pub fn router(store: DiscoveryStore) -> Router { Router::new() .route("/pools", get(list_pools).post(register_pool)) - .route("/routers", get(list_routers)) + .route("/routers", get(list_routers).post(register_router)) .route("/healthz", get(healthz)) .with_state(Arc::new(store)) } @@ -125,6 +131,39 @@ async fn list_routers(State(store): State>) -> Response { } } +async fn register_router( + State(store): State>, + ConnectInfo(peer): ConnectInfo, + Json(registration): Json, +) -> Response { + let Some(registrar) = &store.router_registrar else { + return ( + StatusCode::NOT_IMPLEMENTED, + "router replicas need not register with this store: it lists them from the router \ + Service's EndpointSlices, which MCS imports", + ) + .into_response(); + }; + // The same check as `POST /pools`. A router replica needs no identity + // beyond passing it. + if let Err(error) = verify( + &*store.authenticator, + &PeerCredentials::PeerAddress(peer.ip()), + ) + .await + { + return error_response(error); + } + let endpoint = match registration.into_endpoint() { + Ok(endpoint) => endpoint, + Err(error) => return (StatusCode::BAD_REQUEST, error).into_response(), + }; + match registrar.register(endpoint).await { + Ok(registration) => Json(LeaseView::from(registration.lease)).into_response(), + Err(error) => error_response(error), + } +} + /// Lists what `GET /pools` and `GET /routers` would serve. A Kubernetes store /// answers from its cache once that has synced, so an API server outage does /// not turn every replica unready at once. The first probe also starts the @@ -266,8 +305,27 @@ impl From for WireCapacity { } } -/// `POST /pools` reply: when to register again, and how long a registration -/// lasts. +/// `POST /routers` body: one entry of `GET /routers`. +#[derive(Debug, Deserialize)] +struct RouterRegistration { + replica: String, + url: Url, +} + +impl RouterRegistration { + fn into_endpoint(self) -> Result { + if self.replica.trim().is_empty() { + return Err("replica must not be blank"); + } + Ok(ReplicaEndpoint { + replica: ReplicaId(self.replica), + url: self.url, + }) + } +} + +/// `POST /pools` and `POST /routers` reply: when to register again, and how +/// long a registration lasts. #[derive(Serialize)] struct LeaseView { heartbeat_interval_s: u64, @@ -356,7 +414,7 @@ mod tests { use super::*; use crate::backends::kubernetes::TrustedNetworkAuthenticator; - use crate::backends::memory::{MemoryPoolStore, StaticReplicaDirectory}; + use crate::backends::memory::{MemoryPoolStore, MemoryReplicaStore}; const LEASE: LeasePolicy = LeasePolicy { heartbeat_interval: Duration::from_secs(5), @@ -389,10 +447,12 @@ mod tests { fn memory_store() -> DiscoveryStore { let pools = Arc::new(MemoryPoolStore::new(LEASE).unwrap()); + let routers = Arc::new(MemoryReplicaStore::new(vec![router_replica()], LEASE).unwrap()); DiscoveryStore { pools: pools.clone(), - routers: Arc::new(StaticReplicaDirectory::new(vec![router_replica()])), + routers: routers.clone(), registrar: Some(pools), + router_registrar: Some(routers), authenticator: Arc::new(TrustedNetworkAuthenticator), } } @@ -414,8 +474,12 @@ mod tests { } async fn post_pool(base: &str, body: &Value) -> (reqwest::StatusCode, Value) { + post_json(base, "/pools", body).await + } + + async fn post_json(base: &str, path: &str, body: &Value) -> (reqwest::StatusCode, Value) { let response = reqwest::Client::new() - .post(format!("{base}/pools")) + .post(format!("{base}{path}")) .header("content-type", "application/json") .body(body.to_string()) .send() @@ -473,6 +537,54 @@ mod tests { shutdown.cancel(); } + #[tokio::test] + async fn register_router_then_list_routers() { + let (base, shutdown) = start(memory_store()).await; + let west = json!({"replica": "aws-west/global-router-0", "url": "http://10.43.0.5:50051"}); + let (status, lease) = post_json(&base, "/routers", &west).await; + assert_eq!(status, reqwest::StatusCode::OK); + assert_eq!( + lease, + json!({"heartbeat_interval_s": 5, "lease_duration_s": 15}) + ); + + let (_, routers) = get_json(&base, "/routers").await; + let mut listed = routers["routers"].as_array().unwrap().clone(); + listed.sort_by(|a, b| a["replica"].as_str().cmp(&b["replica"].as_str())); + assert_eq!( + listed, + [ + json!({"replica": "aws-ohio/global-router-0", "url": "http://10.42.0.5:50051/"}), + json!({"replica": "aws-west/global-router-0", "url": "http://10.43.0.5:50051/"}), + ] + ); + + let blank = json!({"replica": " ", "url": "http://10.43.0.6:50051"}); + let (status, _) = post_json(&base, "/routers", &blank).await; + assert_eq!(status, reqwest::StatusCode::BAD_REQUEST); + shutdown.cancel(); + } + + #[tokio::test] + async fn store_that_lists_routers_says_they_need_not_register() { + let mut store = memory_store(); + store.router_registrar = None; + let (base, shutdown) = start(store).await; + let response = reqwest::Client::new() + .post(format!("{base}/routers")) + .header("content-type", "application/json") + .body( + json!({"replica": "aws-west/global-router-0", "url": "http://10.43.0.5:50051"}) + .to_string(), + ) + .send() + .await + .unwrap(); + assert_eq!(response.status(), reqwest::StatusCode::NOT_IMPLEMENTED); + assert!(response.text().await.unwrap().contains("MCS")); + shutdown.cancel(); + } + #[tokio::test] async fn invalid_key_is_a_bad_request() { let (base, shutdown) = start(memory_store()).await; @@ -526,6 +638,7 @@ mod tests { pools: Arc::new(Down), routers: Arc::new(Down), registrar: None, + router_registrar: None, authenticator: Arc::new(TrustedNetworkAuthenticator), }) .await; From 871b21c3154871217bd5d3783f9b8c534a950c7d Mon Sep 17 00:00:00 2001 From: Anna Tchernych Date: Sun, 25 Oct 2026 17:26:16 -0700 Subject: [PATCH 2/4] fix(global-router): publish a router re-registration after an unswept expiry `MemoryReplicaStore::register` decided whether to publish `Added` from `served()`, which hides a registration that has expired but that the sweep has not removed yet. If such a registration overrode a `--router` entry and the router then registered the configured address again, no event went out, and the new live registration kept the sweep from publishing the fallback, so watchers kept the old URL while `list` and `get` returned the new one. Replacing any expired registration is now a change. Only a live registration of the same endpoint, or no registration with the same configured endpoint, skips the event. Also set the new `router_registrar` field in the global-ingress round-robin test, which no longer compiled. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Anna Tchernych --- .../src/backends/memory.rs | 41 ++++++++++++++++++- lib/global-ingress/tests/round_robin.rs | 1 + 2 files changed, 40 insertions(+), 2 deletions(-) diff --git a/lib/global-discovery-service/src/backends/memory.rs b/lib/global-discovery-service/src/backends/memory.rs index 0ec3716f4406..c98f5778d753 100644 --- a/lib/global-discovery-service/src/backends/memory.rs +++ b/lib/global-discovery-service/src/backends/memory.rs @@ -451,7 +451,14 @@ impl ReplicaRegistrar for MemoryReplicaStore { self.ensure_sweeping(); let now = Instant::now(); let mut replicas = self.lock(); - let changed = replicas.served(&endpoint.replica, now) != Some(&endpoint); + let unchanged = match replicas.registered.get(&endpoint.replica) { + // Only a live registration of the same endpoint makes this a + // renewal. Watchers may still hold an expired registration whose + // fallback the sweep has not published yet, so replacing one is + // always a change. + Some(previous) => previous.expires_at > now && previous.endpoint == endpoint, + None => replicas.configured.get(&endpoint.replica) == Some(&endpoint), + }; replicas.registered.insert( endpoint.replica.clone(), RegisteredReplica { @@ -459,7 +466,7 @@ impl ReplicaRegistrar for MemoryReplicaStore { expires_at: now + self.lease.lease_duration, }, ); - if changed { + if !unchanged { // No subscribers is not an error. let _ = self.events.send(WatchEvent::Added(endpoint)); } @@ -736,6 +743,36 @@ mod tests { ); } + #[tokio::test(start_paused = true)] + async fn replacing_a_registration_the_sweep_has_not_removed_is_a_change() { + let configured = router("ohio/global-router-0", "http://10.42.0.5:50051"); + let moved = router("ohio/global-router-0", "http://10.42.0.9:50051"); + let store = MemoryReplicaStore::new(vec![configured.clone()], LEASE).unwrap(); + // Sweeps run every 5 s from here. + let mut watch = store.list_and_watch(None).await.unwrap(); + assert!(matches!(watch.next().await, Some(Ok(WatchEvent::Added(r))) if r == configured)); + assert!(matches!( + watch.next().await, + Some(Ok(WatchEvent::Resync(_))) + )); + + tokio::time::advance(Duration::from_secs(2)).await; + store.register(moved.clone()).await.unwrap(); + assert!(matches!(watch.next().await, Some(Ok(WatchEvent::Added(r))) if r == moved)); + + // At 18 s the registration has expired (at 17 s), but the last sweep + // ran at 15 s, so watchers still hold `moved`. Registering the + // configured address again must tell them. + for _ in 0..16 { + tokio::time::advance(Duration::from_secs(1)).await; + } + store.register(configured.clone()).await.unwrap(); + let event = tokio::time::timeout(Duration::from_secs(1), watch.next()) + .await + .expect("watchers must learn the configured address again"); + assert!(matches!(event, Some(Ok(WatchEvent::Added(r))) if r == configured)); + } + #[tokio::test(start_paused = true)] async fn a_registration_replaces_a_configured_router_while_it_lasts() { let configured = router("ohio/global-router-0", "http://10.42.0.5:50051"); diff --git a/lib/global-ingress/tests/round_robin.rs b/lib/global-ingress/tests/round_robin.rs index a2da58024c25..0aa1ad291e77 100644 --- a/lib/global-ingress/tests/round_robin.rs +++ b/lib/global-ingress/tests/round_robin.rs @@ -34,6 +34,7 @@ async fn discovery_service() -> (Server, GlobalDiscoveryClient) { pools: pools.clone(), routers: Arc::new(StaticReplicaDirectory::new(Vec::new())), registrar: Some(pools), + router_registrar: None, authenticator: Arc::new(TrustedNetworkAuthenticator), }; let (listener, url) = common::bind().await; From 7391493fc50d64c3d425f5bed7eba301675e7f06 Mon Sep 17 00:00:00 2001 From: Anna Tchernych Date: Sun, 25 Oct 2026 17:45:31 -0700 Subject: [PATCH 3/4] refactor(global-router): answer POST /routers with 501 until its design settles Router registration is not designed yet: how a router proves it is a router, and what it sends, are open. The service trusts every caller, a V1 choice made for pools on a trusted network, but a fake router entry is worse than a fake pool: every PoolRelay would connect to it and stream its pool's state there. The memory store, which took these registrations, is also where a trusted network is least likely. So no store takes router registrations, and `POST /routers` answers 501 with the reason: - MCS stores: router replicas need not register; MCS lists them. - Memory store: not implemented yet; list router replicas with --router. `DiscoveryStore::router_source` names where router replicas come from and replaces `router_registrar`. `MemoryReplicaStore`, the `ReplicaRegistrar` trait and the re-registration fix go away with it. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Anna Tchernych --- .../src/backends/kubernetes.rs | 3 +- .../src/backends/memory.rs | 348 +----------------- lib/global-discovery-service/src/main.rs | 25 +- .../src/pool_discovery.rs | 8 +- .../src/pool_discovery/replicas.rs | 18 +- lib/global-discovery-service/src/service.rs | 184 ++++----- lib/global-ingress/tests/round_robin.rs | 4 +- 7 files changed, 94 insertions(+), 496 deletions(-) diff --git a/lib/global-discovery-service/src/backends/kubernetes.rs b/lib/global-discovery-service/src/backends/kubernetes.rs index e7fd07b907e6..32cc255f4fe5 100644 --- a/lib/global-discovery-service/src/backends/kubernetes.rs +++ b/lib/global-discovery-service/src/backends/kubernetes.rs @@ -46,9 +46,8 @@ //! let store = DiscoveryStore { //! pools: Arc::new(McsPoolDirectory::new(client.clone(), mcs, suffix, true, locations)), //! routers: Arc::new(EndpointSliceReplicaDirectory::imported(client, &router, mcs)), +//! router_source: RouterSource::Mcs, //! registrar: None, -//! // `POST /routers` answers that router replicas need not register. -//! router_registrar: None, //! authenticator: Arc::new(TrustedNetworkAuthenticator), //! }; //! service::serve(listener, store, shutdown).await?; diff --git a/lib/global-discovery-service/src/backends/memory.rs b/lib/global-discovery-service/src/backends/memory.rs index c98f5778d753..8c1f25d80484 100644 --- a/lib/global-discovery-service/src/backends/memory.rs +++ b/lib/global-discovery-service/src/backends/memory.rs @@ -14,8 +14,8 @@ use std::time::Duration; use crate::pool_discovery::watch::{Directory, Snapshot, WatchStream}; use crate::pool_discovery::{ Completeness, LeasePolicy, PoolDiscoveryError, PoolEvent, PoolRecord, PoolRegistrar, - Registration, ReplicaEndpoint, ReplicaEvent, ReplicaId, ReplicaRegistrar, Revision, - VerifiedRelayIdentity, WatchEvent, + Registration, ReplicaEndpoint, ReplicaEvent, ReplicaId, Revision, VerifiedRelayIdentity, + WatchEvent, }; use async_trait::async_trait; use dynamo_global_router_shared::PoolKey; @@ -70,7 +70,10 @@ impl MemoryPoolStore { /// Fails unless `heartbeat_interval` is above zero and `lease_duration` /// is longer than it. pub fn new(lease: LeasePolicy) -> anyhow::Result { - check_lease(&lease)?; + anyhow::ensure!( + !lease.heartbeat_interval.is_zero() && lease.lease_duration > lease.heartbeat_interval, + "the lease duration must be longer than the heartbeat interval, which must be above zero" + ); Ok(Self { lease, pools: Arc::default(), @@ -103,16 +106,6 @@ impl Drop for MemoryPoolStore { } } -/// `heartbeat_interval` must be above zero and `lease_duration` longer than -/// it, so the sweep has a period and a registration outlives one interval. -fn check_lease(lease: &LeasePolicy) -> anyhow::Result<()> { - anyhow::ensure!( - !lease.heartbeat_interval.is_zero() && lease.lease_duration > lease.heartbeat_interval, - "the lease duration must be longer than the heartbeat interval, which must be above zero" - ); - Ok(()) -} - /// Removes expired records and publishes `Removed` for each. async fn sweep( pools: Arc>, @@ -294,218 +287,6 @@ impl Directory for StaticReplicaDirectory { } } -/// Router replicas in memory: the ones from configuration, and the ones -/// that register (`POST /routers`). A registration expires `lease_duration` -/// after its last renewal, and while it lasts it replaces a configured -/// replica with the same ID. A restarted process lists only the configured -/// replicas until the others register again, so lists are `Estimated`. -/// -/// Any authenticated caller may register any replica ID: like -/// [`MemoryPoolStore`], V1 assumes a trusted network. -pub struct MemoryReplicaStore { - lease: LeasePolicy, - replicas: Arc>, - /// Sent while holding the `replicas` lock, as in [`MemoryPoolStore`]. - events: broadcast::Sender, - sweeping: OnceLock<()>, - stop: CancellationToken, -} - -struct Replicas { - configured: HashMap, - registered: HashMap, -} - -struct RegisteredReplica { - endpoint: ReplicaEndpoint, - expires_at: Instant, -} - -impl Replicas { - /// A live registration, else the configured replica. - fn served(&self, id: &ReplicaId, now: Instant) -> Option<&ReplicaEndpoint> { - self.registered - .get(id) - .filter(|registered| registered.expires_at > now) - .map(|registered| ®istered.endpoint) - .or_else(|| self.configured.get(id)) - } - - fn snapshot(&self, now: Instant) -> Snapshot { - let registered = |id: &ReplicaId| { - self.registered - .get(id) - .is_some_and(|registered| registered.expires_at > now) - }; - Snapshot { - items: self - .registered - .values() - .filter(|registered| registered.expires_at > now) - .map(|registered| registered.endpoint.clone()) - .chain( - self.configured - .values() - .filter(|endpoint| !registered(&endpoint.replica)) - .cloned(), - ) - .collect(), - completeness: Completeness::Estimated, - } - } -} - -impl MemoryReplicaStore { - /// `configured` replicas are listed until shutdown. Fails, like - /// [`MemoryPoolStore::new`], for a lease without room for heartbeats. - pub fn new(configured: Vec, lease: LeasePolicy) -> anyhow::Result { - check_lease(&lease)?; - let configured = configured - .into_iter() - .map(|endpoint| (endpoint.replica.clone(), endpoint)) - .collect(); - Ok(Self { - lease, - replicas: Arc::new(Mutex::new(Replicas { - configured, - registered: HashMap::new(), - })), - events: broadcast::channel(EVENT_BUFFER).0, - sweeping: OnceLock::new(), - stop: CancellationToken::new(), - }) - } - - fn lock(&self) -> MutexGuard<'_, Replicas> { - self.replicas.lock().unwrap_or_else(PoisonError::into_inner) - } - - /// Every trait method runs in a Tokio runtime, so the spawn is safe. - fn ensure_sweeping(&self) { - self.sweeping.get_or_init(|| { - tokio::spawn(sweep_replicas( - Arc::clone(&self.replicas), - self.events.clone(), - self.lease.heartbeat_interval, - self.stop.clone(), - )); - }); - } -} - -impl Drop for MemoryReplicaStore { - fn drop(&mut self) { - self.stop.cancel(); - } -} - -/// Removes expired registrations. A configured replica goes back to its -/// configured address; any other replica is `Removed`. -async fn sweep_replicas( - replicas: Arc>, - events: broadcast::Sender, - every: Duration, - stop: CancellationToken, -) { - let mut ticks = tokio::time::interval(every); - ticks.set_missed_tick_behavior(MissedTickBehavior::Delay); - loop { - tokio::select! { - biased; - () = stop.cancelled() => return, - _ = ticks.tick() => {} - } - let now = Instant::now(); - let mut replicas = replicas.lock().unwrap_or_else(PoisonError::into_inner); - let Replicas { - configured, - registered, - } = &mut *replicas; - registered.retain(|id, registered| { - let live = registered.expires_at > now; - if !live { - // No subscribers is not an error. - match configured.get(id) { - None => { - let _ = events.send(WatchEvent::Removed(id.clone())); - } - Some(endpoint) if *endpoint != registered.endpoint => { - let _ = events.send(WatchEvent::Added(endpoint.clone())); - } - Some(_) => {} - } - } - live - }); - } -} - -#[async_trait] -impl ReplicaRegistrar for MemoryReplicaStore { - /// An unchanged registration only renews its lease. A new or changed one - /// is an `Added` event. - async fn register( - &self, - endpoint: ReplicaEndpoint, - ) -> Result { - self.ensure_sweeping(); - let now = Instant::now(); - let mut replicas = self.lock(); - let unchanged = match replicas.registered.get(&endpoint.replica) { - // Only a live registration of the same endpoint makes this a - // renewal. Watchers may still hold an expired registration whose - // fallback the sweep has not published yet, so replacing one is - // always a change. - Some(previous) => previous.expires_at > now && previous.endpoint == endpoint, - None => replicas.configured.get(&endpoint.replica) == Some(&endpoint), - }; - replicas.registered.insert( - endpoint.replica.clone(), - RegisteredReplica { - endpoint: endpoint.clone(), - expires_at: now + self.lease.lease_duration, - }, - ); - if !unchanged { - // No subscribers is not an error. - let _ = self.events.send(WatchEvent::Added(endpoint)); - } - Ok(Registration { lease: self.lease }) - } -} - -#[async_trait] -impl Directory for MemoryReplicaStore { - async fn check_connection(&self) -> Result<(), PoolDiscoveryError> { - Ok(()) - } - - async fn list(&self) -> Result, PoolDiscoveryError> { - Ok(self.lock().snapshot(Instant::now())) - } - - async fn get(&self, key: &ReplicaId) -> Result, PoolDiscoveryError> { - Ok(self.lock().served(key, Instant::now()).cloned()) - } - - async fn list_and_watch( - &self, - cancel_token: Option, - ) -> Result, PoolDiscoveryError> { - self.ensure_sweeping(); - let replicas = Arc::clone(&self.replicas); - let events = self.events.clone(); - Ok(watch_stream( - move || { - let replicas = replicas.lock().unwrap_or_else(PoisonError::into_inner); - (events.subscribe(), Some(replicas.snapshot(Instant::now()))) - }, - cancel_token.unwrap_or_default(), - self.stop.clone(), - )) - } -} - #[cfg(test)] mod tests { use crate::pool_discovery::{ @@ -668,10 +449,6 @@ mod tests { MemoryPoolStore::new(lease).is_err(), "{heartbeat_interval}s / {lease_duration}s" ); - assert!( - MemoryReplicaStore::new(Vec::new(), lease).is_err(), - "{heartbeat_interval}s / {lease_duration}s" - ); } } @@ -703,117 +480,4 @@ mod tests { cancel.cancel(); assert!(watch.next().await.is_none()); } - - fn router(id: &str, url: &str) -> ReplicaEndpoint { - ReplicaEndpoint { - replica: ReplicaId(id.to_owned()), - url: Url::parse(url).unwrap(), - } - } - - #[tokio::test(start_paused = true)] - async fn a_registered_router_is_listed_until_its_lease_ends() { - let store = MemoryReplicaStore::new(Vec::new(), LEASE).unwrap(); - let replica = router("ohio/global-router-0", "http://10.42.0.5:50051"); - let mut watch = store.list_and_watch(None).await.unwrap(); - assert!( - matches!(watch.next().await, Some(Ok(WatchEvent::Resync(s))) if s.items.is_empty()) - ); - - let registration = store.register(replica.clone()).await.unwrap(); - assert_eq!(registration.lease, LEASE); - assert!(matches!(watch.next().await, Some(Ok(WatchEvent::Added(r))) if r == replica)); - - // Heartbeats every 10 s keep the 15 s lease alive; renewing an - // unchanged registration is not a change. - for _ in 0..3 { - tokio::time::advance(Duration::from_secs(10)).await; - store.register(replica.clone()).await.unwrap(); - } - assert_eq!( - store.list().await.unwrap().items, - std::slice::from_ref(&replica) - ); - - // Without heartbeats it ends: reads skip it, then the sweep removes it. - tokio::time::advance(Duration::from_secs(16)).await; - assert!(store.get(&replica.replica).await.unwrap().is_none()); - assert!( - matches!(watch.next().await, Some(Ok(WatchEvent::Removed(id))) if id == replica.replica) - ); - } - - #[tokio::test(start_paused = true)] - async fn replacing_a_registration_the_sweep_has_not_removed_is_a_change() { - let configured = router("ohio/global-router-0", "http://10.42.0.5:50051"); - let moved = router("ohio/global-router-0", "http://10.42.0.9:50051"); - let store = MemoryReplicaStore::new(vec![configured.clone()], LEASE).unwrap(); - // Sweeps run every 5 s from here. - let mut watch = store.list_and_watch(None).await.unwrap(); - assert!(matches!(watch.next().await, Some(Ok(WatchEvent::Added(r))) if r == configured)); - assert!(matches!( - watch.next().await, - Some(Ok(WatchEvent::Resync(_))) - )); - - tokio::time::advance(Duration::from_secs(2)).await; - store.register(moved.clone()).await.unwrap(); - assert!(matches!(watch.next().await, Some(Ok(WatchEvent::Added(r))) if r == moved)); - - // At 18 s the registration has expired (at 17 s), but the last sweep - // ran at 15 s, so watchers still hold `moved`. Registering the - // configured address again must tell them. - for _ in 0..16 { - tokio::time::advance(Duration::from_secs(1)).await; - } - store.register(configured.clone()).await.unwrap(); - let event = tokio::time::timeout(Duration::from_secs(1), watch.next()) - .await - .expect("watchers must learn the configured address again"); - assert!(matches!(event, Some(Ok(WatchEvent::Added(r))) if r == configured)); - } - - #[tokio::test(start_paused = true)] - async fn a_registration_replaces_a_configured_router_while_it_lasts() { - let configured = router("ohio/global-router-0", "http://10.42.0.5:50051"); - let moved = router("ohio/global-router-0", "http://10.42.0.9:50051"); - let other = router("west/global-router-0", "http://10.43.0.5:50051"); - let store = MemoryReplicaStore::new(vec![configured.clone()], LEASE).unwrap(); - let mut watch = store.list_and_watch(None).await.unwrap(); - assert!(matches!(watch.next().await, Some(Ok(WatchEvent::Added(r))) if r == configured)); - assert!(matches!( - watch.next().await, - Some(Ok(WatchEvent::Resync(_))) - )); - - // Registering the configured address is not a change. - store.register(configured.clone()).await.unwrap(); - store.register(moved.clone()).await.unwrap(); - store.register(other.clone()).await.unwrap(); - assert!(matches!(watch.next().await, Some(Ok(WatchEvent::Added(r))) if r == moved)); - assert!(matches!(watch.next().await, Some(Ok(WatchEvent::Added(r))) if r == other)); - assert_eq!( - store.get(&configured.replica).await.unwrap(), - Some(moved.clone()) - ); - - // When the registrations end, the configured address comes back and - // the other replica goes. - tokio::time::advance(Duration::from_secs(16)).await; - let mut changes = Vec::new(); - for _ in 0..2 { - changes.push(watch.next().await.unwrap().unwrap()); - } - assert!( - changes - .iter() - .any(|event| matches!(event, WatchEvent::Added(r) if *r == configured)) - ); - assert!( - changes - .iter() - .any(|event| matches!(event, WatchEvent::Removed(id) if *id == other.replica)) - ); - assert_eq!(store.list().await.unwrap().items, [configured]); - } } diff --git a/lib/global-discovery-service/src/main.rs b/lib/global-discovery-service/src/main.rs index c709418a9e29..c3699a127e70 100644 --- a/lib/global-discovery-service/src/main.rs +++ b/lib/global-discovery-service/src/main.rs @@ -18,12 +18,11 @@ use clap::{Parser, ValueEnum}; use dynamo_global_discovery_service::backends::kubernetes::{ EndpointSliceReplicaDirectory, McsPoolDirectory, RouterService, TrustedNetworkAuthenticator, }; -use dynamo_global_discovery_service::backends::memory::{MemoryPoolStore, MemoryReplicaStore}; +use dynamo_global_discovery_service::backends::memory::{MemoryPoolStore, StaticReplicaDirectory}; use dynamo_global_discovery_service::pool_discovery::{ LeasePolicy, McsImplementation, PoolDirectory, PoolRegistrar, ReplicaEndpoint, ReplicaId, - ReplicaRegistrar, }; -use dynamo_global_discovery_service::service::{self, DiscoveryStore}; +use dynamo_global_discovery_service::service::{self, DiscoveryStore, RouterSource}; use kube::Client; use tokio::net::TcpListener; use tokio_util::sync::CancellationToken; @@ -38,15 +37,14 @@ struct Args { /// Where pools and router replicas come from. #[arg(long, value_enum, default_value_t = Store::Memory)] store: Store, - /// Memory store: how often a PoolRelay or router replica registers again. + /// Memory store: how often a PoolRelay registers again. #[arg(long, default_value_t = 5)] heartbeat_interval_secs: u64, /// Memory store: how long a registration lasts. #[arg(long, default_value_t = 15)] lease_duration_secs: u64, - /// Memory store: a Global Router replica as `=`, listed until - /// shutdown. Repeat it for each replica. Replicas can also register - /// themselves with `POST /routers`. + /// Memory store: a Global Router replica as `=`. Repeat it for + /// each replica. #[arg(long = "router", value_parser = parse_router)] routers: Vec, /// Kubernetes stores: each PoolRelay Service is named @@ -66,8 +64,7 @@ struct Args { #[derive(Clone, Copy, ValueEnum)] enum Store { - /// Registrations in this process, from pools and router replicas; router - /// replicas also from `--router`. + /// Registrations in this process; router replicas from `--router`. Memory, /// Imported EndpointSlices, named the Karmada way. No registrations. Karmada, @@ -134,15 +131,13 @@ fn memory_store(args: &Args) -> Result { { bail!("--router {} is given twice", repeated.replica.0); } - let routers = Arc::new(MemoryReplicaStore::new(args.routers.clone(), lease)?); let pools: Arc = store.clone(); let registrar: Arc = store; - let router_registrar: Arc = routers.clone(); Ok(DiscoveryStore { pools, - routers, + routers: Arc::new(StaticReplicaDirectory::new(args.routers.clone())), + router_source: RouterSource::Configured, registrar: Some(registrar), - router_registrar: Some(router_registrar), authenticator: Arc::new(TrustedNetworkAuthenticator), }) } @@ -178,10 +173,8 @@ async fn kubernetes_store(args: &Args, mcs: McsImplementation) -> Result; pub type ReplicaEvent = WatchEvent; - -/// Write side behind `RegisterRouter`. A router replica calls `register` on -/// a timer, and each call is also its heartbeat, as with -/// [`super::PoolRegistrar`]. Only stores that cannot find router replicas -/// themselves take registrations: in Kubernetes mode, the router Service's -/// EndpointSlices that MCS imports already list every replica. -#[async_trait] -pub trait ReplicaRegistrar: Send + Sync { - async fn register(&self, endpoint: ReplicaEndpoint) - -> Result; -} diff --git a/lib/global-discovery-service/src/service.rs b/lib/global-discovery-service/src/service.rs index 4d0467c2fa4a..4935fdfbb205 100644 --- a/lib/global-discovery-service/src/service.rs +++ b/lib/global-discovery-service/src/service.rs @@ -6,9 +6,8 @@ //! - `POST /pools`: a PoolRelay registers its pool. It calls this on a timer, //! and each call is also the pool's heartbeat. //! - `GET /routers`: the Global Router replicas a PoolRelay connects to. -//! - `POST /routers`: a Global Router replica registers itself, on a timer -//! like `POST /pools`. Only a store that cannot list router replicas -//! itself takes it; the MCS store answers `501`, since MCS lists them. +//! - `POST /routers`: `RegisterRouter` from the design, not available yet. +//! It answers `501` with the reason (see [`RouterSource`]). //! - `GET /pools`: the pools a Global Router routes to. //! - `GET /healthz`: OK while the service can answer from its store. Use it //! for readiness; liveness needs only the TCP port. @@ -29,8 +28,8 @@ use std::sync::Arc; use crate::pool_discovery::{ Capacity, Completeness, LeasePolicy, PeerCredentials, PoolDirectory, PoolDiscoveryError, - PoolRecord, PoolRegistrar, RelayAuthenticator, ReplicaDirectory, ReplicaEndpoint, ReplicaId, - ReplicaRegistrar, Revision, verify, + PoolRecord, PoolRegistrar, RelayAuthenticator, ReplicaDirectory, ReplicaEndpoint, Revision, + verify, }; use axum::extract::{ConnectInfo, State}; use axum::http::StatusCode; @@ -48,15 +47,27 @@ use url::Url; pub struct DiscoveryStore { pub pools: Arc, pub routers: Arc, + /// Where `routers` comes from. `POST /routers` gives it as the reason. + pub router_source: RouterSource, /// `None` when pools are announced another way, as with MCS. pub registrar: Option>, - /// `None` when the store lists router replicas itself, as with MCS. - pub router_registrar: Option>, - /// Names the caller of `POST /pools` and `POST /routers` from its - /// connection. + /// Names the caller of `POST /pools` from its connection. pub authenticator: Arc, } +/// Where the service finds router replicas. Router replicas cannot register +/// yet: `POST /routers` answers `501`. With MCS they need not. Without it, +/// how a router proves it is a router, and what it sends, are not decided. +/// Until then every caller would be trusted, and a fake router entry would +/// make every PoolRelay stream its pool's state to it. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RouterSource { + /// `--router` configuration, in the memory store. + Configured, + /// The router Service's EndpointSlices that MCS imports. + Mcs, +} + pub fn router(store: DiscoveryStore) -> Router { Router::new() .route("/pools", get(list_pools).post(register_pool)) @@ -131,37 +142,20 @@ async fn list_routers(State(store): State>) -> Response { } } -async fn register_router( - State(store): State>, - ConnectInfo(peer): ConnectInfo, - Json(registration): Json, -) -> Response { - let Some(registrar) = &store.router_registrar else { - return ( - StatusCode::NOT_IMPLEMENTED, +/// `RegisterRouter` from the design. It reads no body: the answer is the +/// same for any request. +async fn register_router(State(store): State>) -> Response { + let reason = match store.router_source { + RouterSource::Mcs => { "router replicas need not register with this store: it lists them from the router \ - Service's EndpointSlices, which MCS imports", - ) - .into_response(); - }; - // The same check as `POST /pools`. A router replica needs no identity - // beyond passing it. - if let Err(error) = verify( - &*store.authenticator, - &PeerCredentials::PeerAddress(peer.ip()), - ) - .await - { - return error_response(error); - } - let endpoint = match registration.into_endpoint() { - Ok(endpoint) => endpoint, - Err(error) => return (StatusCode::BAD_REQUEST, error).into_response(), + Service's EndpointSlices, which MCS imports" + } + RouterSource::Configured => { + "router registration is not implemented yet: how a router proves it is a router, \ + and what it sends, are not decided. List router replicas with --router" + } }; - match registrar.register(endpoint).await { - Ok(registration) => Json(LeaseView::from(registration.lease)).into_response(), - Err(error) => error_response(error), - } + (StatusCode::NOT_IMPLEMENTED, reason).into_response() } /// Lists what `GET /pools` and `GET /routers` would serve. A Kubernetes store @@ -305,27 +299,8 @@ impl From for WireCapacity { } } -/// `POST /routers` body: one entry of `GET /routers`. -#[derive(Debug, Deserialize)] -struct RouterRegistration { - replica: String, - url: Url, -} - -impl RouterRegistration { - fn into_endpoint(self) -> Result { - if self.replica.trim().is_empty() { - return Err("replica must not be blank"); - } - Ok(ReplicaEndpoint { - replica: ReplicaId(self.replica), - url: self.url, - }) - } -} - -/// `POST /pools` and `POST /routers` reply: when to register again, and how -/// long a registration lasts. +/// `POST /pools` reply: when to register again, and how long a registration +/// lasts. #[derive(Serialize)] struct LeaseView { heartbeat_interval_s: u64, @@ -414,7 +389,7 @@ mod tests { use super::*; use crate::backends::kubernetes::TrustedNetworkAuthenticator; - use crate::backends::memory::{MemoryPoolStore, MemoryReplicaStore}; + use crate::backends::memory::{MemoryPoolStore, StaticReplicaDirectory}; const LEASE: LeasePolicy = LeasePolicy { heartbeat_interval: Duration::from_secs(5), @@ -447,12 +422,11 @@ mod tests { fn memory_store() -> DiscoveryStore { let pools = Arc::new(MemoryPoolStore::new(LEASE).unwrap()); - let routers = Arc::new(MemoryReplicaStore::new(vec![router_replica()], LEASE).unwrap()); DiscoveryStore { pools: pools.clone(), - routers: routers.clone(), + routers: Arc::new(StaticReplicaDirectory::new(vec![router_replica()])), + router_source: RouterSource::Configured, registrar: Some(pools), - router_registrar: Some(routers), authenticator: Arc::new(TrustedNetworkAuthenticator), } } @@ -474,12 +448,8 @@ mod tests { } async fn post_pool(base: &str, body: &Value) -> (reqwest::StatusCode, Value) { - post_json(base, "/pools", body).await - } - - async fn post_json(base: &str, path: &str, body: &Value) -> (reqwest::StatusCode, Value) { let response = reqwest::Client::new() - .post(format!("{base}{path}")) + .post(format!("{base}/pools")) .header("content-type", "application/json") .body(body.to_string()) .send() @@ -538,51 +508,41 @@ mod tests { } #[tokio::test] - async fn register_router_then_list_routers() { - let (base, shutdown) = start(memory_store()).await; - let west = json!({"replica": "aws-west/global-router-0", "url": "http://10.43.0.5:50051"}); - let (status, lease) = post_json(&base, "/routers", &west).await; - assert_eq!(status, reqwest::StatusCode::OK); - assert_eq!( - lease, - json!({"heartbeat_interval_s": 5, "lease_duration_s": 15}) - ); - - let (_, routers) = get_json(&base, "/routers").await; - let mut listed = routers["routers"].as_array().unwrap().clone(); - listed.sort_by(|a, b| a["replica"].as_str().cmp(&b["replica"].as_str())); - assert_eq!( - listed, - [ - json!({"replica": "aws-ohio/global-router-0", "url": "http://10.42.0.5:50051/"}), - json!({"replica": "aws-west/global-router-0", "url": "http://10.43.0.5:50051/"}), - ] - ); - - let blank = json!({"replica": " ", "url": "http://10.43.0.6:50051"}); - let (status, _) = post_json(&base, "/routers", &blank).await; - assert_eq!(status, reqwest::StatusCode::BAD_REQUEST); - shutdown.cancel(); - } - - #[tokio::test] - async fn store_that_lists_routers_says_they_need_not_register() { - let mut store = memory_store(); - store.router_registrar = None; - let (base, shutdown) = start(store).await; - let response = reqwest::Client::new() - .post(format!("{base}/routers")) - .header("content-type", "application/json") - .body( + async fn routers_cannot_register_and_the_answer_says_why() { + for (source, reason) in [ + (RouterSource::Configured, "not decided"), + (RouterSource::Mcs, "need not register"), + ] { + let mut store = memory_store(); + store.router_source = source; + let (base, shutdown) = start(store).await; + // Any body gets the same answer. + for body in [ json!({"replica": "aws-west/global-router-0", "url": "http://10.43.0.5:50051"}) .to_string(), - ) - .send() - .await - .unwrap(); - assert_eq!(response.status(), reqwest::StatusCode::NOT_IMPLEMENTED); - assert!(response.text().await.unwrap().contains("MCS")); - shutdown.cancel(); + String::new(), + ] { + let response = reqwest::Client::new() + .post(format!("{base}/routers")) + .body(body) + .send() + .await + .unwrap(); + assert_eq!( + response.status(), + reqwest::StatusCode::NOT_IMPLEMENTED, + "{source:?}" + ); + assert!( + response.text().await.unwrap().contains(reason), + "{source:?}" + ); + } + // Listing still works. + let (status, _) = get_json(&base, "/routers").await; + assert_eq!(status, reqwest::StatusCode::OK); + shutdown.cancel(); + } } #[tokio::test] @@ -637,8 +597,8 @@ mod tests { let (base, shutdown) = start(DiscoveryStore { pools: Arc::new(Down), routers: Arc::new(Down), + router_source: RouterSource::Mcs, registrar: None, - router_registrar: None, authenticator: Arc::new(TrustedNetworkAuthenticator), }) .await; diff --git a/lib/global-ingress/tests/round_robin.rs b/lib/global-ingress/tests/round_robin.rs index 0aa1ad291e77..669b2fbf0586 100644 --- a/lib/global-ingress/tests/round_robin.rs +++ b/lib/global-ingress/tests/round_robin.rs @@ -18,7 +18,7 @@ use dynamo_global_discovery_client::{ use dynamo_global_discovery_service::backends::kubernetes::TrustedNetworkAuthenticator; use dynamo_global_discovery_service::backends::memory::{MemoryPoolStore, StaticReplicaDirectory}; use dynamo_global_discovery_service::pool_discovery::LeasePolicy; -use dynamo_global_discovery_service::service::{self, DiscoveryStore}; +use dynamo_global_discovery_service::service::{self, DiscoveryStore, RouterSource}; use dynamo_global_router::{DiscoveryPoolCatalog, PoolCatalog}; use reqwest::StatusCode; use tokio_util::sync::CancellationToken; @@ -33,8 +33,8 @@ async fn discovery_service() -> (Server, GlobalDiscoveryClient) { let store = DiscoveryStore { pools: pools.clone(), routers: Arc::new(StaticReplicaDirectory::new(Vec::new())), + router_source: RouterSource::Configured, registrar: Some(pools), - router_registrar: None, authenticator: Arc::new(TrustedNetworkAuthenticator), }; let (listener, url) = common::bind().await; From a1b7d5c12f5db091a8cbaed0461e7e45b407a40d Mon Sep 17 00:00:00 2001 From: Anna Tchernych Date: Sun, 25 Oct 2026 18:02:15 -0700 Subject: [PATCH 4/4] test(global-router): keep one POST /routers case per reason Address review comments. The MCS answer is checked with an empty body, since it must not depend on the body, and the memory-store answer with a registration body. Drop the GET /routers check, which `routers_use_the_client_shape` already covers, and a field comment that repeated `RouterSource`'s docs. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Anna Tchernych --- lib/global-discovery-service/src/service.rs | 54 ++++++++++----------- 1 file changed, 25 insertions(+), 29 deletions(-) diff --git a/lib/global-discovery-service/src/service.rs b/lib/global-discovery-service/src/service.rs index 4935fdfbb205..c1d1d8e18423 100644 --- a/lib/global-discovery-service/src/service.rs +++ b/lib/global-discovery-service/src/service.rs @@ -47,7 +47,6 @@ use url::Url; pub struct DiscoveryStore { pub pools: Arc, pub routers: Arc, - /// Where `routers` comes from. `POST /routers` gives it as the reason. pub router_source: RouterSource, /// `None` when pools are announced another way, as with MCS. pub registrar: Option>, @@ -509,38 +508,35 @@ mod tests { #[tokio::test] async fn routers_cannot_register_and_the_answer_says_why() { - for (source, reason) in [ - (RouterSource::Configured, "not decided"), - (RouterSource::Mcs, "need not register"), + let registration = + json!({"replica": "aws-west/global-router-0", "url": "http://10.43.0.5:50051"}); + // The MCS answer must not depend on the body, so that case sends none. + for (source, body, reason) in [ + (RouterSource::Mcs, String::new(), "need not register"), + ( + RouterSource::Configured, + registration.to_string(), + "not decided", + ), ] { let mut store = memory_store(); store.router_source = source; let (base, shutdown) = start(store).await; - // Any body gets the same answer. - for body in [ - json!({"replica": "aws-west/global-router-0", "url": "http://10.43.0.5:50051"}) - .to_string(), - String::new(), - ] { - let response = reqwest::Client::new() - .post(format!("{base}/routers")) - .body(body) - .send() - .await - .unwrap(); - assert_eq!( - response.status(), - reqwest::StatusCode::NOT_IMPLEMENTED, - "{source:?}" - ); - assert!( - response.text().await.unwrap().contains(reason), - "{source:?}" - ); - } - // Listing still works. - let (status, _) = get_json(&base, "/routers").await; - assert_eq!(status, reqwest::StatusCode::OK); + let response = reqwest::Client::new() + .post(format!("{base}/routers")) + .body(body) + .send() + .await + .unwrap(); + assert_eq!( + response.status(), + reqwest::StatusCode::NOT_IMPLEMENTED, + "{source:?}" + ); + assert!( + response.text().await.unwrap().contains(reason), + "{source:?}" + ); shutdown.cancel(); } }