diff --git a/bindings/golang/src/policy.rs b/bindings/golang/src/policy.rs index b50b8ac557..4d1525f774 100644 --- a/bindings/golang/src/policy.rs +++ b/bindings/golang/src/policy.rs @@ -6,6 +6,7 @@ //! Rust gateway. use std::{ + any::Any, ffi::{CStr, CString}, os::raw::c_char, ptr, @@ -31,7 +32,7 @@ use smg::{ utils::{generate_tool_constraints, process_chat_messages}, }, worker::{ - circuit_breaker::CircuitBreaker, + circuit_breaker::{CircuitBreaker, CircuitState}, resilience::ResolvedResilience, worker::{RuntimeType, WorkerMetadata, WorkerRoutingKeyLoad}, ConnectionMode, Worker, WorkerResult, WorkerType, @@ -106,6 +107,10 @@ impl std::fmt::Debug for GrpcWorker { #[async_trait] impl Worker for GrpcWorker { + fn as_any(&self) -> &dyn Any { + self + } + fn url(&self) -> &str { &self.endpoint } @@ -130,11 +135,33 @@ impl Worker for GrpcWorker { self.status.store(status as u8, Ordering::Relaxed); } + fn revision(&self) -> u64 { + 0 + } + async fn check_health_async(&self) -> WorkerResult<()> { // FFI workers don't do their own health checks Ok(()) } + // FFI workers don't run the state machine — these counters are unused. + // The Go SDK manages worker health via direct set_healthy() calls. + fn consecutive_failures_increment(&self) -> usize { + 0 + } + fn consecutive_failures_reset(&self) {} + fn consecutive_successes_increment(&self) -> usize { + 0 + } + fn consecutive_successes_reset(&self) {} + fn total_pending_probes(&self) -> usize { + 0 + } + fn total_pending_probes_increment(&self) -> usize { + 0 + } + fn total_pending_probes_reset(&self) {} + fn load(&self) -> usize { self.load.load(Ordering::Relaxed) } @@ -147,8 +174,16 @@ impl Worker for GrpcWorker { self.load.fetch_sub(1, Ordering::Relaxed); } - fn worker_routing_key_load(&self) -> &WorkerRoutingKeyLoad { - &self.routing_key_load + fn routing_key_load(&self) -> usize { + self.routing_key_load.value() + } + + fn increment_routing_key_load(&self, routing_key: &str) { + self.routing_key_load.increment(routing_key); + } + + fn decrement_routing_key_load(&self, routing_key: &str) { + self.routing_key_load.decrement(routing_key); } fn processed_requests(&self) -> usize { @@ -163,8 +198,16 @@ impl Worker for GrpcWorker { &self.metadata } - fn circuit_breaker(&self) -> &CircuitBreaker { - &self.circuit_breaker + fn circuit_breaker_state(&self) -> CircuitState { + self.circuit_breaker.state() + } + + fn circuit_breaker_can_execute(&self) -> bool { + self.circuit_breaker.can_execute() + } + + fn record_circuit_breaker_outcome(&self, success: bool) { + self.circuit_breaker.record_outcome(success); } fn resilience(&self) -> &ResolvedResilience { diff --git a/model_gateway/src/observability/metrics_ws/collectors.rs b/model_gateway/src/observability/metrics_ws/collectors.rs index 3dc79309a2..e8b7b916b1 100644 --- a/model_gateway/src/observability/metrics_ws/collectors.rs +++ b/model_gateway/src/observability/metrics_ws/collectors.rs @@ -206,7 +206,7 @@ fn collect_workers(context: &AppContext) -> Value { "is_healthy": w.is_healthy(), "load": w.load(), "processed_requests": w.processed_requests(), - "circuit_breaker": w.circuit_breaker().state().to_string(), + "circuit_breaker": w.circuit_breaker_state().to_string(), }) }) .collect(); diff --git a/model_gateway/src/policies/manual.rs b/model_gateway/src/policies/manual.rs index a8186c43f3..c406fd452d 100644 --- a/model_gateway/src/policies/manual.rs +++ b/model_gateway/src/policies/manual.rs @@ -292,9 +292,7 @@ fn min_load_select(workers: &[Arc], healthy_indices: &[usize]) -> us } fn min_group_select(workers: &[Arc], healthy_indices: &[usize]) -> usize { - select_min_by(healthy_indices, |idx| { - workers[idx].worker_routing_key_load().value() - }) + select_min_by(healthy_indices, |idx| workers[idx].routing_key_load()) } #[cfg(test)] @@ -793,9 +791,7 @@ mod tests { assert_eq!(branch, ExecutionBranch::Vacant); let selected_idx = result.unwrap(); - workers[selected_idx] - .worker_routing_key_load() - .increment(&routing_key); + workers[selected_idx].increment_routing_key_load(&routing_key); } let distribution: HashMap<_, usize> = policy @@ -822,13 +818,13 @@ mod tests { let policy = ManualPolicy::with_config(config); let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]); - workers[0].worker_routing_key_load().increment("existing-1"); - workers[0].worker_routing_key_load().increment("existing-2"); - workers[1].worker_routing_key_load().increment("existing-3"); + workers[0].increment_routing_key_load("existing-1"); + workers[0].increment_routing_key_load("existing-2"); + workers[1].increment_routing_key_load("existing-3"); - assert_eq!(workers[0].worker_routing_key_load().value(), 2); - assert_eq!(workers[1].worker_routing_key_load().value(), 1); - assert_eq!(workers[2].worker_routing_key_load().value(), 0); + assert_eq!(workers[0].routing_key_load(), 2); + assert_eq!(workers[1].routing_key_load(), 1); + assert_eq!(workers[2].routing_key_load(), 0); let headers = headers_with_routing_key("new-key"); let info = SelectWorkerInfo { @@ -878,9 +874,9 @@ mod tests { let policy = ManualPolicy::with_config(config); let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); - workers[0].worker_routing_key_load().increment("key-0"); - workers[1].worker_routing_key_load().increment("key-1"); - workers[1].worker_routing_key_load().increment("key-2"); + workers[0].increment_routing_key_load("key-0"); + workers[1].increment_routing_key_load("key-1"); + workers[1].increment_routing_key_load("key-2"); let headers = headers_with_routing_key("new-key"); let info = SelectWorkerInfo { @@ -916,9 +912,9 @@ mod tests { let policy = ManualPolicy::with_config(config); let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); - workers[0].worker_routing_key_load().increment("key-1"); - workers[0].worker_routing_key_load().increment("key-2"); - workers[0].worker_routing_key_load().increment("key-3"); + workers[0].increment_routing_key_load("key-1"); + workers[0].increment_routing_key_load("key-2"); + workers[0].increment_routing_key_load("key-3"); let mut selected_worker_0 = false; for i in 0..50 { diff --git a/model_gateway/src/policies/mod.rs b/model_gateway/src/policies/mod.rs index 872c526a0b..d3cdbc485a 100644 --- a/model_gateway/src/policies/mod.rs +++ b/model_gateway/src/policies/mod.rs @@ -144,7 +144,7 @@ pub(crate) fn get_healthy_worker_indices(workers: &[Arc]) -> Vec Result<(), Box = router_manager.clone(); - // Health checker handle must outlive the server to keep the background task alive. - // HealthChecker aborts its task on Drop, so binding it here keeps it alive until - // the server shuts down. - let _health_checker = if config.router_config.health_check.disable_health_check { - info!("Global health checks disabled via CLI/config; skipping health checker"); + // WorkerManager owns the background health check loop. Its handle must + // outlive the server to keep the task alive — bind it here so its Drop + // (which aborts the task) runs at server shutdown. + let _worker_manager = if config.router_config.health_check.disable_health_check { + info!("Global health checks disabled via CLI/config; skipping WorkerManager"); None } else { - let hc = app_context.worker_registry.start_health_checker( - config.router_config.health_check.check_interval_secs, - config.router_config.health_check.remove_unhealthy_workers, + let manager = WorkerManager::start( + app_context.worker_registry.clone(), + WorkerManagerConfig { + default_check_interval_secs: config.router_config.health_check.check_interval_secs, + remove_unhealthy: config.router_config.health_check.remove_unhealthy_workers, + }, app_context.worker_job_queue.get().cloned(), ); debug!( - "Started health checker for workers with {}s interval", + "Started WorkerManager health check loop with {}s default interval", config.router_config.health_check.check_interval_secs ); - Some(hc) + Some(manager) }; // LoadMonitor groups are started dynamically when workers are registered. diff --git a/model_gateway/src/service_discovery.rs b/model_gateway/src/service_discovery.rs index ed0d5fd595..9be462e722 100644 --- a/model_gateway/src/service_discovery.rs +++ b/model_gateway/src/service_discovery.rs @@ -640,6 +640,7 @@ async fn handle_pod_event( ); let job = Job::RemoveWorker { url: old_url.clone(), + expected_revision: None, }; if let Some(job_queue) = app_context.worker_job_queue.get() { if let Err(e) = job_queue.submit(job).await { @@ -789,6 +790,7 @@ async fn handle_pod_deletion( let job = Job::RemoveWorker { url: worker_url.clone(), + expected_revision: None, }; if let Some(job_queue) = app_context.worker_job_queue.get() { @@ -930,6 +932,7 @@ async fn reconcile_pods( ); let job = Job::RemoveWorker { url: worker_url.clone(), + expected_revision: None, }; if let Some(job_queue) = app_context.worker_job_queue.get() { match job_queue.submit(job).await { diff --git a/model_gateway/src/worker/builder.rs b/model_gateway/src/worker/builder.rs index ec1241dee7..8bcd426843 100644 --- a/model_gateway/src/worker/builder.rs +++ b/model_gateway/src/worker/builder.rs @@ -10,7 +10,7 @@ use super::{ circuit_breaker::{CircuitBreaker, CircuitBreakerConfig}, resilience::ResolvedResilience, worker::{ - BasicWorker, ConnectionMode, RuntimeType, WorkerMetadata, WorkerRoutingKeyLoad, WorkerType, + BasicWorker, ConnectionMode, RuntimeType, WorkerMetadata, WorkerRuntime, WorkerType, DEFAULT_WORKER_HTTP_TIMEOUT_SECS, }, }; @@ -227,10 +227,7 @@ impl BasicWorkerBuilder { /// Build the BasicWorker instance pub fn build(mut self) -> BasicWorker { - use std::sync::{ - atomic::{AtomicU8, AtomicUsize}, - Arc, - }; + use std::sync::Arc; use tokio::sync::OnceCell; @@ -287,17 +284,11 @@ impl BasicWorkerBuilder { let resilience = self.resilience.unwrap_or_default(); BasicWorker { - load_counter: Arc::new(AtomicUsize::new(0)), - worker_routing_key_load: Arc::new(WorkerRoutingKeyLoad::new(&metadata.spec.url)), - processed_counter: Arc::new(AtomicUsize::new(0)), - status: Arc::new(AtomicU8::new(initial_status as u8)), - consecutive_failures: Arc::new(AtomicUsize::new(0)), - consecutive_successes: Arc::new(AtomicUsize::new(0)), - total_pending_probes: Arc::new(AtomicUsize::new(0)), - circuit_breaker: CircuitBreaker::with_config_and_label( + runtime: ArcSwap::from_pointee(WorkerRuntime::new(&metadata.spec.url, initial_status)), + circuit_breaker: ArcSwap::from_pointee(CircuitBreaker::with_config_and_label( self.circuit_breaker_config, metadata.spec.url.clone(), - ), + )), metadata, grpc_client, models_override: Arc::new(ArcSwap::from_pointee(WorkerModels::Wildcard)), diff --git a/model_gateway/src/worker/circuit_breaker.rs b/model_gateway/src/worker/circuit_breaker.rs index e92a179722..f17b5af5a6 100644 --- a/model_gateway/src/worker/circuit_breaker.rs +++ b/model_gateway/src/worker/circuit_breaker.rs @@ -8,7 +8,7 @@ use tracing::info; use crate::observability::metrics::Metrics; /// Circuit breaker configuration -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq)] pub struct CircuitBreakerConfig { /// Number of consecutive failures to open the circuit pub failure_threshold: u32, @@ -143,6 +143,11 @@ impl CircuitBreaker { &self.metric_label } + /// Get the configuration used to create this circuit breaker. + pub fn config(&self) -> &CircuitBreakerConfig { + &self.config + } + /// Check if a request can be executed (lock-free hot path) #[inline] pub fn can_execute(&self) -> bool { diff --git a/model_gateway/src/worker/manager.rs b/model_gateway/src/worker/manager.rs index 8c206fe44b..e638cba8a3 100644 --- a/model_gateway/src/worker/manager.rs +++ b/model_gateway/src/worker/manager.rs @@ -2,33 +2,46 @@ //! //! Provides worker lifecycle operations and fan-out request utilities. -use std::{collections::HashMap, sync::Arc, time::Duration}; +use std::{ + collections::{HashMap, HashSet}, + future::Future, + pin::Pin, + sync::Arc, + time::Duration, +}; use axum::response::{IntoResponse, Response}; use futures::{ future, - stream::{self, StreamExt}, + stream::{self, FuturesUnordered, StreamExt}, }; use http::StatusCode; use openai_protocol::worker::{ - FlushCacheResult, WorkerGroupKey, WorkerLoadInfo, WorkerLoadResponse, WorkerLoadsResult, + FlushCacheResult, HealthCheckConfig, WorkerGroupKey, WorkerLoadInfo, WorkerLoadResponse, + WorkerLoadsResult, WorkerStatus, }; use tokio::{ - sync::{watch, Mutex}, + sync::{broadcast, watch, Mutex, Notify}, task::JoinHandle, }; -use tracing::{debug, info}; +use tracing::{debug, error, info, warn}; use crate::{ + observability::metrics::{metrics_labels, Metrics}, policies::PolicyRegistry, worker::{ + event::WorkerEvent, metrics_aggregator::{self, MetricPack}, + registry::{WorkerDescriptor, WorkerId}, + worker::WorkerTypeExt, ConnectionMode, Worker, WorkerLoadManager, WorkerRegistry, WorkerType, }, + workflow::{Job, JobQueue}, }; const REQUEST_TIMEOUT: Duration = Duration::from_secs(5); const MAX_CONCURRENT: usize = 32; +const MAX_CONCURRENT_HEALTH_PROBES: usize = 128; /// Result of a fan-out request to a single worker struct WorkerResponse { @@ -85,7 +98,580 @@ impl IntoResponse for EngineMetricsResult { } } -pub struct WorkerManager; +/// Lifecycle coordinator for the worker fleet. +/// +/// Owns the background health check loop, applies the state machine to +/// probe outcomes, and triggers removal of `Failed` workers when +/// `--remove-unhealthy-workers` is set. Subscribes to `WorkerRegistry` +/// events to keep its internal schedule in sync with registrations, +/// removals, and replacements. +/// +/// The static fan-out helpers (`get_worker_urls`, `flush_cache_all`, +/// `get_all_worker_loads`, `get_engine_metrics`) are operational commands +/// that don't depend on lifecycle state and remain associated functions. +pub struct WorkerManager { + handle: Option>, + shutdown_notify: Arc, +} + +impl std::fmt::Debug for WorkerManager { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("WorkerManager").finish() + } +} + +/// Configuration for the WorkerManager health check loop. +#[derive(Debug, Clone)] +pub struct WorkerManagerConfig { + /// Default check interval used when a worker has no override. + pub default_check_interval_secs: u64, + /// If true, submit `Job::RemoveWorker` for workers that reach Failed. + pub remove_unhealthy: bool, +} + +impl WorkerManager { + /// Create and start the WorkerManager background loop. + /// + /// Spawns a single task that: + /// - Subscribes to `WorkerRegistry` events to maintain a per-worker + /// deadline schedule. + /// - Probes due workers via `Worker::check_health_async()`. + /// - Applies the state machine to probe outcomes and calls + /// `WorkerRegistry::transition_status()` to publish StatusChanged. + /// - Submits `Job::RemoveWorker` for Failed workers when removal is + /// enabled. + pub fn start( + registry: Arc, + config: WorkerManagerConfig, + job_queue: Option>, + ) -> Self { + let shutdown_notify = Arc::new(Notify::new()); + let shutdown_clone = shutdown_notify.clone(); + + // Subscribe BEFORE snapshotting the registry. Any registration that + // lands after this line either (a) is already in the snapshot below + // because it happened synchronously on this thread, or (b) arrives + // as a Registered event in the broadcast buffer and is idempotently + // applied by the event loop. The "a or b" dichotomy is what makes + // startup deterministic regardless of task scheduling. + let events_rx = registry.subscribe_events(); + + // Run the bootstrap reconcile synchronously on the caller's thread + // so the initial schedule is captured deterministically — not + // whenever the spawned task happens to be scheduled. A worker + // registered between WorkerManager::start() returning and the task + // running (e.g. the mesh replay loop in server.rs, which runs + // synchronously right after start()) would otherwise race the + // task's own reconcile call. + let mut next_check: HashMap = HashMap::new(); + reconcile_from_registry(®istry, &mut next_check, &config); + + let job_queue = if config.remove_unhealthy { + job_queue + } else { + None + }; + + #[expect( + clippy::disallowed_methods, + reason = "WorkerManager loop runs for the lifetime of the registry; handle is stored and abort() runs on drop" + )] + let handle = tokio::spawn(async move { + run_health_loop( + registry, + events_rx, + next_check, + config, + job_queue, + shutdown_clone, + ) + .await; + }); + + Self { + handle: Some(handle), + shutdown_notify, + } + } + + /// Gracefully shut down the WorkerManager loop, awaiting the task. + /// Prefer this over dropping when an async context is available — it + /// lets the in-flight probe iteration finish instead of aborting. + pub async fn shutdown(&mut self) { + self.shutdown_notify.notify_one(); + if let Some(handle) = self.handle.take() { + let _ = handle.await; + } + } +} + +impl Drop for WorkerManager { + fn drop(&mut self) { + if let Some(handle) = self.handle.take() { + handle.abort(); + } + } +} + +struct ProbeCompletion { + worker_id: WorkerId, + worker: Arc, + expected_revision: u64, + launched_status: WorkerStatus, + health_config: HealthCheckConfig, + probe_result: crate::worker::WorkerResult<()>, +} + +struct RemovalCandidate { + worker_id: WorkerId, + url: String, + expected_revision: u64, +} + +enum ProbeApplyResult { + Applied(Option<(WorkerStatus, WorkerStatus)>), + Stale, +} + +type ProbeFutures = FuturesUnordered + Send>>>; + +/// Background task body: deadline-driven probe loop + event subscription. +/// +/// The loop keeps deadline scheduling, in-flight probe tracking, and event +/// handling in one place. Probes run concurrently via `FuturesUnordered` so a +/// slow worker does not block unrelated health checks or registry events. +/// +/// `next_check` is seeded by `WorkerManager::start()` via a synchronous +/// `reconcile_from_registry` call, so this function never runs a bootstrap +/// reconcile of its own — by the time the task is scheduled, the caller's +/// thread has already captured a consistent registry snapshot. The only +/// in-loop reconcile is the lag-recovery rebuild triggered by +/// `RecvError::Lagged`. +async fn run_health_loop( + registry: Arc, + mut events_rx: broadcast::Receiver, + mut next_check: HashMap, + config: WorkerManagerConfig, + job_queue: Option>, + shutdown: Arc, +) { + let mut probes: ProbeFutures = FuturesUnordered::new(); + let mut in_flight: HashSet = HashSet::new(); + + loop { + let now = tokio::time::Instant::now(); + let removals = queue_due_probes( + ®istry, + &config, + &mut next_check, + &mut in_flight, + &mut probes, + now, + ); + for removal in removals { + if let Some(jq) = job_queue.as_ref() { + submit_removal_job( + &removal.worker_id, + &removal.url, + removal.expected_revision, + jq, + ) + .await; + } + } + + let sleep_until = next_check + .values() + .min() + .copied() + .unwrap_or_else(|| now + Duration::from_secs(config.default_check_interval_secs)); + + tokio::select! { + Some(completion) = probes.next(), if !probes.is_empty() => { + let worker_id = completion.worker_id.clone(); + in_flight.remove(&worker_id); + if matches!( + apply_probe_completion(®istry, completion, job_queue.as_ref()).await, + ProbeApplyResult::Applied(Some((_, WorkerStatus::Failed))) + ) { + next_check.remove(&worker_id); + } + } + () = tokio::time::sleep_until(sleep_until) => {} + event = events_rx.recv() => { + match event { + Ok(WorkerEvent::Registered { worker_id, worker }) => { + schedule_worker_at( + &mut next_check, + worker_id, + worker.status(), + &worker.metadata().health_config, + &config, + tokio::time::Instant::now(), + true, + ); + } + Ok(WorkerEvent::Removed { worker_id, .. }) => { + next_check.remove(&worker_id); + } + Ok(WorkerEvent::Replaced { worker_id, new, .. }) => { + schedule_worker_at( + &mut next_check, + worker_id, + new.status(), + &new.metadata().health_config, + &config, + tokio::time::Instant::now(), + true, + ); + } + Ok(WorkerEvent::StatusChanged { .. }) => { + // Self-published; nothing to do. + } + Err(broadcast::error::RecvError::Lagged(n)) => { + warn!( + "WorkerManager lagged {n} events; rebuilding schedule from registry" + ); + next_check.clear(); + reconcile_from_registry(®istry, &mut next_check, &config); + } + Err(broadcast::error::RecvError::Closed) => { + debug!("WorkerEvent channel closed; WorkerManager exiting"); + return; + } + } + } + () = shutdown.notified() => { + debug!("WorkerManager received shutdown signal"); + return; + } + } + } +} + +/// Side-effect-free startup reconcile: rebuild the schedule from the +/// current registry snapshot without publishing new events or removals. +fn reconcile_from_registry( + registry: &Arc, + next_check: &mut HashMap, + config: &WorkerManagerConfig, +) { + let now = tokio::time::Instant::now(); + for descriptor in registry.reconcile_snapshot() { + schedule_descriptor_at(next_check, descriptor, config, now, false); + } +} + +fn queue_due_probes( + registry: &Arc, + config: &WorkerManagerConfig, + next_check: &mut HashMap, + in_flight: &mut HashSet, + probes: &mut ProbeFutures, + now: tokio::time::Instant, +) -> Vec { + let capacity = MAX_CONCURRENT_HEALTH_PROBES.saturating_sub(in_flight.len()); + if capacity == 0 { + return Vec::new(); + } + + let due_ids: Vec = next_check + .iter() + .filter(|(worker_id, deadline)| now >= **deadline && !in_flight.contains(*worker_id)) + .map(|(worker_id, _)| worker_id.clone()) + .take(capacity) + .collect(); + + let mut removals = Vec::new(); + for worker_id in due_ids { + let Some(worker) = registry.get(&worker_id) else { + next_check.remove(&worker_id); + continue; + }; + + let health_config = worker.metadata().health_config.clone(); + if health_config.disable_health_check { + next_check.remove(&worker_id); + continue; + } + + let launched_status = worker.status(); + let expected_revision = worker.revision(); + if launched_status == WorkerStatus::Failed { + next_check.remove(&worker_id); + if config.remove_unhealthy { + removals.push(RemovalCandidate { + worker_id: worker_id.clone(), + url: worker.url().to_string(), + expected_revision, + }); + } + continue; + } + + let next_deadline = now + + Duration::from_secs(resolved_interval_secs( + &health_config, + config.default_check_interval_secs, + )); + next_check.insert(worker_id.clone(), next_deadline); + in_flight.insert(worker_id.clone()); + probes.push(Box::pin(async move { + let probe_result = worker.check_health_async().await; + ProbeCompletion { + worker_id, + worker, + expected_revision, + launched_status, + health_config, + probe_result, + } + })); + } + + removals +} + +async fn apply_probe_completion( + registry: &Arc, + completion: ProbeCompletion, + job_queue: Option<&Arc>, +) -> ProbeApplyResult { + let ProbeCompletion { + worker_id, + worker, + expected_revision, + launched_status, + health_config, + probe_result, + } = completion; + + let probe_ok = match probe_result { + Ok(()) => true, + Err(err) => { + warn!( + worker_url = %worker.url(), + error = %err, + "Health probe failed" + ); + false + } + }; + Metrics::record_worker_health_check( + worker.worker_type().as_metric_label(), + if probe_ok { + metrics_labels::CB_SUCCESS + } else { + metrics_labels::CB_FAILURE + }, + ); + + let Some(((), transition)) = + registry.apply_if_revision(&worker_id, expected_revision, |current_worker| { + if launched_status == WorkerStatus::Pending { + current_worker.total_pending_probes_increment(); + } + ( + (), + compute_next_status(current_worker, probe_ok, &health_config), + ) + }) + else { + debug!( + worker_url = %worker.url(), + expected_revision, + "Discarding stale probe outcome after worker replacement" + ); + return ProbeApplyResult::Stale; + }; + + if let Some((old, new)) = transition { + debug!( + worker_url = %worker.url(), + ?old, + ?new, + "Worker status transition" + ); + if new == WorkerStatus::Failed { + if let Some(jq) = job_queue { + submit_removal_job(&worker_id, worker.url(), expected_revision, jq).await; + } + } + } + + ProbeApplyResult::Applied(transition) +} + +fn resolved_interval_secs(health_config: &HealthCheckConfig, default_interval_secs: u64) -> u64 { + if health_config.check_interval_secs > 0 { + health_config.check_interval_secs + } else { + default_interval_secs + } +} + +fn schedule_descriptor_at( + next_check: &mut HashMap, + descriptor: WorkerDescriptor, + config: &WorkerManagerConfig, + now: tokio::time::Instant, + immediate: bool, +) { + if descriptor.disable_health_check { + next_check.remove(&descriptor.worker_id); + return; + } + if descriptor.status == WorkerStatus::Failed { + // Startup reconcile and lagged rebuild must be side-effect-free: + // do not reschedule already-failed workers for probing or removal. + next_check.remove(&descriptor.worker_id); + return; + } + + let delay = if immediate { + Duration::ZERO + } else { + Duration::from_secs(if descriptor.check_interval_secs > 0 { + descriptor.check_interval_secs + } else { + config.default_check_interval_secs + }) + }; + next_check.insert(descriptor.worker_id, now + delay); +} + +fn schedule_worker_at( + next_check: &mut HashMap, + worker_id: WorkerId, + status: WorkerStatus, + health_config: &HealthCheckConfig, + config: &WorkerManagerConfig, + now: tokio::time::Instant, + immediate: bool, +) { + schedule_descriptor_at( + next_check, + WorkerDescriptor { + worker_id, + status, + disable_health_check: health_config.disable_health_check, + check_interval_secs: health_config.check_interval_secs, + }, + config, + now, + immediate, + ); +} + +/// Apply the state machine to a probe outcome. Returns the next status if +/// a transition is needed, `None` if the worker stays in its current state. +/// +/// State machine rules: +/// - Pending → Ready on `success_threshold` consecutive successes +/// - Pending → Failed on `max_pending_probes` (10 × failure_threshold) total +/// - NotReady → Ready on `success_threshold` consecutive successes +/// - NotReady → Failed on `liveness_failure_threshold` (3 × failure_threshold) +/// - Ready → NotReady on `failure_threshold` consecutive failures +/// - Failed: terminal (handled outside this function — no transitions) +fn compute_next_status( + worker: &Arc, + probe_ok: bool, + health_config: &HealthCheckConfig, +) -> Option { + let current_status = worker.status(); + let success_threshold = health_config.success_threshold as usize; + let failure_threshold = health_config.failure_threshold as usize; + // Liveness threshold: tolerate longer outages before declaring Failed, + // analogous to K8s having separate readiness and liveness probes. + let liveness_threshold = failure_threshold * 3; + // Pending cap: prevent misconfigured URLs from sitting in Pending forever. + let max_pending_probes = failure_threshold * 10; + + if probe_ok { + worker.consecutive_failures_reset(); + let successes = worker.consecutive_successes_increment(); + + if matches!( + current_status, + WorkerStatus::Pending | WorkerStatus::NotReady + ) && successes >= success_threshold + { + worker.consecutive_successes_reset(); + worker.total_pending_probes_reset(); + return Some(WorkerStatus::Ready); + } + + // Even on a successful probe, enforce the Pending cap. A worker + // that flaps F,S,F,S,... never reaches success_threshold and would + // otherwise grow `total_pending_probes` without bound. + if current_status == WorkerStatus::Pending + && worker.total_pending_probes() >= max_pending_probes + { + worker.consecutive_successes_reset(); + worker.consecutive_failures_reset(); + return Some(WorkerStatus::Failed); + } + + None + } else { + worker.consecutive_successes_reset(); + let failures = worker.consecutive_failures_increment(); + + match current_status { + WorkerStatus::Ready => { + if failures >= failure_threshold { + worker.consecutive_failures_reset(); + return Some(WorkerStatus::NotReady); + } + } + WorkerStatus::NotReady => { + if failures >= liveness_threshold { + worker.consecutive_failures_reset(); + return Some(WorkerStatus::Failed); + } + } + WorkerStatus::Pending => { + if worker.total_pending_probes() >= max_pending_probes { + worker.consecutive_failures_reset(); + return Some(WorkerStatus::Failed); + } + } + WorkerStatus::Failed => { + // Terminal — handled outside. + } + } + + None + } +} + +async fn submit_removal_job( + worker_id: &WorkerId, + worker_url: &str, + expected_revision: u64, + job_queue: &Arc, +) { + let url = worker_url.to_string(); + warn!( + worker_id = %worker_id.as_str(), + worker_url = %url, + expected_revision, + "Removing failed worker from registry" + ); + if let Err(e) = job_queue + .submit(Job::RemoveWorker { + url: url.clone(), + expected_revision: Some(expected_revision), + }) + .await + { + error!( + worker_url = %url, + error = %e, + "Failed to submit worker removal job" + ); + } +} impl WorkerManager { pub fn get_worker_urls(registry: &Arc) -> Vec { @@ -548,3 +1134,299 @@ impl Drop for LoadMonitor { } } } + +#[cfg(test)] +mod tests { + use std::{collections::HashMap, sync::Arc}; + + use openai_protocol::worker::{HealthCheckConfig, WorkerStatus}; + + use super::*; + use crate::worker::{BasicWorkerBuilder, Worker, WorkerError, WorkerRegistry, WorkerType}; + + fn make_worker(url: &str, success_threshold: u32, failure_threshold: u32) -> Arc { + Arc::new( + BasicWorkerBuilder::new(url) + .worker_type(WorkerType::Regular) + .health_config(HealthCheckConfig { + success_threshold, + failure_threshold, + timeout_secs: 1, + check_interval_secs: 1, + disable_health_check: false, + }) + .build(), + ) + } + + fn cfg(success_threshold: u32, failure_threshold: u32) -> HealthCheckConfig { + HealthCheckConfig { + success_threshold, + failure_threshold, + timeout_secs: 1, + check_interval_secs: 1, + disable_health_check: false, + } + } + + #[test] + fn test_state_machine_pending_to_ready_after_success_threshold() { + let worker = make_worker("http://w:1", 2, 3); + assert_eq!(worker.status(), WorkerStatus::Pending); + + // First success: not yet promoted (1 < 2) + assert_eq!(compute_next_status(&worker, true, &cfg(2, 3)), None); + assert_eq!(worker.status(), WorkerStatus::Pending); + + // Second success: promoted Pending → Ready + let next = compute_next_status(&worker, true, &cfg(2, 3)); + assert_eq!(next, Some(WorkerStatus::Ready)); + } + + #[test] + fn test_state_machine_ready_to_notready_after_failure_threshold() { + let worker = make_worker("http://w:1", 2, 3); + worker.set_status(WorkerStatus::Ready); + + // 1 fail, 2 fail: still Ready + assert_eq!(compute_next_status(&worker, false, &cfg(2, 3)), None); + assert_eq!(compute_next_status(&worker, false, &cfg(2, 3)), None); + assert_eq!(worker.status(), WorkerStatus::Ready); + + // 3rd fail: Ready → NotReady + assert_eq!( + compute_next_status(&worker, false, &cfg(2, 3)), + Some(WorkerStatus::NotReady) + ); + } + + #[test] + fn test_state_machine_notready_to_failed_after_liveness_threshold() { + let worker = make_worker("http://w:1", 2, 3); + worker.set_status(WorkerStatus::NotReady); + + // liveness_threshold = 3 × failure_threshold = 9 + for i in 1..9 { + assert_eq!( + compute_next_status(&worker, false, &cfg(2, 3)), + None, + "iteration {i}" + ); + } + + // 9th consecutive failure → Failed + assert_eq!( + compute_next_status(&worker, false, &cfg(2, 3)), + Some(WorkerStatus::Failed) + ); + } + + #[test] + fn test_state_machine_pending_to_failed_after_max_pending_probes() { + let worker = make_worker("http://w:1", 2, 3); + // max_pending_probes = 10 × failure_threshold = 30 + + // Simulate 30 failed probes — increment counter manually since the + // loop usually does this before calling compute_next_status. + for _ in 0..29 { + worker.total_pending_probes_increment(); + assert_eq!(compute_next_status(&worker, false, &cfg(2, 3)), None); + } + worker.total_pending_probes_increment(); + // 30th: Pending → Failed + assert_eq!( + compute_next_status(&worker, false, &cfg(2, 3)), + Some(WorkerStatus::Failed) + ); + } + + #[test] + fn test_state_machine_pending_to_failed_on_success_when_cap_exceeded() { + // Flapping pattern: a Pending worker that flaps F,S,F,S,... and + // never reaches success_threshold should still hit max_pending_probes. + let worker = make_worker("http://w:1", 2, 3); + + // Simulate 30 attempts with the counter, then call compute on success. + for _ in 0..30 { + worker.total_pending_probes_increment(); + } + // Even on success, the cap fires. + assert_eq!( + compute_next_status(&worker, true, &cfg(2, 3)), + Some(WorkerStatus::Failed) + ); + } + + #[test] + fn test_state_machine_failed_is_terminal() { + let worker = make_worker("http://w:1", 2, 3); + worker.set_status(WorkerStatus::Failed); + + // Successful probes don't recover Failed. + assert_eq!(compute_next_status(&worker, true, &cfg(2, 3)), None); + assert_eq!(compute_next_status(&worker, true, &cfg(2, 3)), None); + assert_eq!(worker.status(), WorkerStatus::Failed); + + // Failed probes don't transition Failed anywhere either. + assert_eq!(compute_next_status(&worker, false, &cfg(2, 3)), None); + } + + #[test] + fn test_state_machine_notready_to_ready_on_success_threshold() { + let worker = make_worker("http://w:1", 2, 3); + worker.set_status(WorkerStatus::NotReady); + + assert_eq!(compute_next_status(&worker, true, &cfg(2, 3)), None); + assert_eq!( + compute_next_status(&worker, true, &cfg(2, 3)), + Some(WorkerStatus::Ready) + ); + } + + #[test] + fn test_state_machine_success_resets_failure_counter() { + let worker = make_worker("http://w:1", 2, 3); + worker.set_status(WorkerStatus::Ready); + + // 2 failures (not yet at threshold) + assert_eq!(compute_next_status(&worker, false, &cfg(2, 3)), None); + assert_eq!(compute_next_status(&worker, false, &cfg(2, 3)), None); + + // Single success resets the counter + assert_eq!(compute_next_status(&worker, true, &cfg(2, 3)), None); + + // Now 2 failures again — still no transition because counter was reset + assert_eq!(compute_next_status(&worker, false, &cfg(2, 3)), None); + assert_eq!(compute_next_status(&worker, false, &cfg(2, 3)), None); + assert_eq!(worker.status(), WorkerStatus::Ready); + + // 3rd failure now triggers transition + assert_eq!( + compute_next_status(&worker, false, &cfg(2, 3)), + Some(WorkerStatus::NotReady) + ); + } + + #[test] + fn test_reconcile_from_registry_skips_failed_workers_on_bootstrap() { + let registry = Arc::new(WorkerRegistry::new()); + let failed_worker = make_worker("http://failed:1", 2, 3); + failed_worker.set_status(WorkerStatus::Failed); + let failed_id = registry.register(failed_worker).unwrap(); + + let mut next_check = HashMap::new(); + reconcile_from_registry( + ®istry, + &mut next_check, + &WorkerManagerConfig { + default_check_interval_secs: 5, + remove_unhealthy: true, + }, + ); + + assert!( + !next_check.contains_key(&failed_id), + "bootstrap reconcile must not reschedule failed workers" + ); + } + + #[test] + fn test_reconcile_from_registry_captures_pending_and_ready_workers() { + // Positive complement to the "skips failed" test: the startup + // reconcile must pick up Pending (not-yet-probed) and Ready + // workers that existed at registry snapshot time. This is what + // makes WorkerManager::start() deterministic — the schedule is + // captured on the caller's thread, not whenever the spawned task + // happens to run. + let registry = Arc::new(WorkerRegistry::new()); + + let pending_worker = make_worker("http://pending:1", 2, 3); + assert_eq!(pending_worker.status(), WorkerStatus::Pending); + let pending_id = registry.register(pending_worker).unwrap(); + + let ready_worker = make_worker("http://ready:1", 2, 3); + ready_worker.set_status(WorkerStatus::Ready); + let ready_id = registry.register(ready_worker).unwrap(); + + let mut next_check = HashMap::new(); + reconcile_from_registry( + ®istry, + &mut next_check, + &WorkerManagerConfig { + default_check_interval_secs: 5, + remove_unhealthy: true, + }, + ); + + assert!( + next_check.contains_key(&pending_id), + "pending worker must be in the bootstrap schedule" + ); + assert!( + next_check.contains_key(&ready_id), + "ready worker must be in the bootstrap schedule" + ); + } + + #[tokio::test] + async fn test_worker_manager_start_is_deterministic_with_preexisting_workers() { + // End-to-end contract for fix #2: a worker that exists in the + // registry before WorkerManager::start() returns must be on the + // schedule the moment the spawned task begins running — no race + // with task scheduling. We can't observe `next_check` directly, + // so we run the full start/shutdown lifecycle with a very long + // probe interval (so no probe actually fires) and verify the + // happy path doesn't panic. Together with the reconcile unit + // test above, this covers both the "reconcile captures workers" + // and "start() calls reconcile synchronously" invariants. + let registry = Arc::new(WorkerRegistry::new()); + let worker = make_worker("http://pre-existing:1", 2, 3); + worker.set_status(WorkerStatus::Ready); + registry.register(worker).unwrap(); + + let mut manager = WorkerManager::start( + registry, + WorkerManagerConfig { + default_check_interval_secs: 3600, + remove_unhealthy: false, + }, + None, + ); + manager.shutdown().await; + } + + #[tokio::test] + async fn test_apply_probe_completion_discards_stale_probe_after_replace() { + let registry = Arc::new(WorkerRegistry::new()); + let worker = make_worker("http://w:1", 1, 1); + worker.set_status(WorkerStatus::Ready); + let worker_id = registry.register(worker.clone()).unwrap(); + let expected_revision = worker.revision(); + + let completion = ProbeCompletion { + worker_id: worker_id.clone(), + worker: worker.clone(), + expected_revision, + launched_status: WorkerStatus::Ready, + health_config: cfg(1, 1), + probe_result: Err(WorkerError::HealthCheckFailed { + url: worker.url().to_string(), + reason: "stale probe".to_string(), + }), + }; + + let replacement = make_worker("http://w:1", 1, 1); + assert!(registry.replace(&worker_id, replacement)); + + let current = registry.get(&worker_id).unwrap(); + assert_eq!(current.status(), WorkerStatus::Ready); + assert_eq!(current.revision(), expected_revision + 1); + + let result = apply_probe_completion(®istry, completion, None).await; + assert!(matches!(result, ProbeApplyResult::Stale)); + assert_eq!( + registry.get(&worker_id).unwrap().status(), + WorkerStatus::Ready + ); + } +} diff --git a/model_gateway/src/worker/registry.rs b/model_gateway/src/worker/registry.rs index 5233a61da2..914b330814 100644 --- a/model_gateway/src/worker/registry.rs +++ b/model_gateway/src/worker/registry.rs @@ -11,10 +11,7 @@ //! The ring is rebuilt only when workers are added/removed, not per-request. //! Uses virtual nodes (150 per worker) for even distribution and blake3 for stable hashing. -use std::{ - collections::{HashMap, HashSet}, - sync::Arc, -}; +use std::{collections::HashSet, sync::Arc}; use dashmap::{mapref::entry::Entry, DashMap}; use openai_protocol::worker::WorkerStatus; @@ -29,10 +26,9 @@ use crate::{ worker::{ circuit_breaker::CircuitState, event::WorkerEvent, - worker::{HealthChecker, RuntimeType, WorkerType}, + worker::{RuntimeType, WorkerType}, ConnectionMode, Worker, }, - workflow::{Job, JobQueue}, }; /// Number of virtual nodes per physical worker for even distribution. @@ -181,6 +177,15 @@ impl Default for WorkerId { } } +/// Side-effect-free worker snapshot for subscriber bootstrap or lag recovery. +#[derive(Debug, Clone)] +pub struct WorkerDescriptor { + pub worker_id: WorkerId, + pub status: WorkerStatus, + pub disable_health_check: bool, + pub check_interval_secs: u64, +} + /// Model index using immutable snapshots for lock-free reads. /// Each model maps to an Arc'd slice of workers that can be read without locking. /// Updates create new snapshots (copy-on-write semantics). @@ -465,6 +470,14 @@ impl WorkerRegistry { return false; } + if !new_worker.inherit_shared_state_from(&*old_worker) { + tracing::warn!( + worker_id = %worker_id.as_str(), + worker_url = old_worker.url(), + "replace() did not preserve shared mutable worker state" + ); + } + // Overwrite worker object atomically self.workers.insert(worker_id.clone(), new_worker.clone()); @@ -743,6 +756,22 @@ impl WorkerRegistry { .collect() } + /// Get a side-effect-free snapshot for startup reconcile or lag recovery. + pub fn reconcile_snapshot(&self) -> Vec { + self.workers + .iter() + .map(|entry| { + let worker = entry.value(); + WorkerDescriptor { + worker_id: entry.key().clone(), + status: worker.status(), + disable_health_check: worker.metadata().health_config.disable_health_check, + check_interval_secs: worker.metadata().health_config.check_interval_secs, + } + }) + .collect() + } + /// Get all worker URLs pub fn get_all_urls(&self) -> Vec { self.workers @@ -870,7 +899,7 @@ impl WorkerRegistry { ConnectionMode::Grpc => grpc_count += 1, } - match worker.circuit_breaker().state() { + match worker.circuit_breaker_state() { CircuitState::Open => cb_open_count += 1, CircuitState::HalfOpen => cb_half_open_count += 1, CircuitState::Closed => {} @@ -912,132 +941,117 @@ impl WorkerRegistry { (regular_count, pd_count) } - /// Start a deadline-driven health checker for all workers in the registry. + /// Atomically transition a worker's lifecycle status and emit a + /// `StatusChanged` event if it actually changed. + /// + /// This is a pure mutation primitive — the registry has no opinion on + /// when a worker should transition. The caller (typically `WorkerManager`) + /// owns the state machine logic. + /// + /// The per-worker mutation lock guarantees: + /// 1. The status read + write + event emit are atomic per worker. + /// 2. Two concurrent calls cannot interleave to publish events out of + /// order for the same worker. /// - /// Each worker is checked according to its own `health_config.check_interval_secs`. - /// The task sleeps until the next worker is due, so it only wakes when there is - /// actual work to do — zero CPU when idle, no polling. - pub(crate) fn start_health_checker( + /// Returns `Some((old, new))` if the status changed, `None` if the worker + /// is gone or the status was already `new_status`. + pub fn transition_status( &self, - default_interval_secs: u64, - remove_unhealthy: bool, - job_queue: Option>, - ) -> HealthChecker { - let shutdown_notify = Arc::new(tokio::sync::Notify::new()); - let shutdown_clone = shutdown_notify.clone(); - let workers_ref = self.workers.clone(); - let job_queue = if remove_unhealthy { job_queue } else { None }; - - #[expect( - clippy::disallowed_methods, - reason = "Health checker loop: runs for the lifetime of the registry, handle is stored in HealthChecker and abort() is called on drop" - )] - let handle = tokio::spawn(async move { - // next_check[url] = Instant when the worker is next due for a health check. - let mut next_check: HashMap = HashMap::new(); - - loop { - let now = tokio::time::Instant::now(); - - // Snapshot current workers from the registry - let workers: Vec> = workers_ref - .iter() - .map(|entry| entry.value().clone()) - .collect(); + worker_id: &WorkerId, + new_status: WorkerStatus, + ) -> Option<(WorkerStatus, WorkerStatus)> { + self.transition_status_inner(worker_id, None, new_status) + } - // Sync schedule with registry: add new workers, prune removed - // and disabled ones so stale deadlines don't cause wakeups. - let checkable_urls: HashSet = workers - .iter() - .filter(|w| !w.metadata().health_config.disable_health_check) - .map(|w| w.url().to_string()) - .collect(); - next_check.retain(|url, _| checkable_urls.contains(url)); - for url in &checkable_urls { - next_check.entry(url.clone()).or_insert(now); - } + /// Same as `transition_status()`, but becomes a no-op if the currently + /// installed worker revision no longer matches `expected_revision`. + pub fn transition_status_if_revision( + &self, + worker_id: &WorkerId, + expected_revision: u64, + new_status: WorkerStatus, + ) -> Option<(WorkerStatus, WorkerStatus)> { + self.transition_status_inner(worker_id, Some(expected_revision), new_status) + } - // Collect workers whose deadline has passed - let due_workers: Vec<_> = workers - .iter() - .filter(|w| !w.metadata().health_config.disable_health_check) - .filter(|w| { - next_check - .get(w.url()) - .is_some_and(|deadline| now >= *deadline) - }) - .cloned() - .collect(); + /// Apply a worker-local mutation while holding the per-worker lock and + /// optionally emit a `StatusChanged` event under the same lock. + /// + /// Used by `WorkerManager` so counter mutation and revision-checked status + /// transitions cannot race a same-URL `replace()`. + pub fn apply_if_revision( + &self, + worker_id: &WorkerId, + expected_revision: u64, + f: F, + ) -> Option<(T, Option<(WorkerStatus, WorkerStatus)>)> + where + F: FnOnce(&Arc) -> (T, Option), + { + let lock = self + .worker_mutation_locks + .entry(worker_id.clone()) + .or_insert_with(|| Arc::new(parking_lot::Mutex::new(()))) + .clone(); + let _guard = lock.lock(); - // Run due health checks in parallel and schedule the next deadline - if !due_workers.is_empty() { - for worker in &due_workers { - let secs = worker.metadata().health_config.check_interval_secs; - let secs = if secs > 0 { - secs - } else { - default_interval_secs - }; - next_check.insert( - worker.url().to_string(), - now + tokio::time::Duration::from_secs(secs), - ); - } - let futs: Vec<_> = due_workers - .into_iter() - .map(|w| async move { - let failed = w.check_health_async().await.is_err(); - (w, failed) - }) - .collect(); - let checked_workers = futures::future::join_all(futs).await; - - // Only remove Failed workers — not Pending (still starting) - // or NotReady (may recover). Failed is terminal, so we - // don't gate on *failed (a worker can reach Failed via - // max_pending_probes on a probe that returned Ok(false)). - if let Some(ref job_queue) = job_queue { - for (worker, _failed) in &checked_workers { - if worker.status() == WorkerStatus::Failed { - let url = worker.url().to_string(); - tracing::warn!( - worker_url = %url, - "Removing unhealthy worker from registry" - ); - next_check.remove(&url); - if let Err(e) = job_queue - .submit(Job::RemoveWorker { url: url.clone() }) - .await - { - tracing::error!( - worker_url = %url, - error = %e, - "Failed to submit worker removal job" - ); - } - } - } - } - } + let worker = self.workers.get(worker_id)?.clone(); + if worker.revision() != expected_revision { + return None; + } - // Sleep until the earliest deadline or until shutdown is signalled. - // If the registry is empty, sleep for the default interval then re-scan - // (new workers may have been added). - let sleep_until = next_check.values().min().copied().unwrap_or_else(|| { - now + tokio::time::Duration::from_secs(default_interval_secs) + let old_status = worker.status(); + let (result, candidate_status) = f(&worker); + let transition = match candidate_status { + Some(new_status) if new_status != old_status => { + worker.set_status(new_status); + let _ = self.event_tx.send(WorkerEvent::StatusChanged { + worker_id: worker_id.clone(), + worker: worker.clone(), + old_status, + new_status, }); - - tokio::select! { - () = tokio::time::sleep_until(sleep_until) => {} - () = shutdown_clone.notified() => { - tracing::debug!("Registry health checker shutting down"); - break; - } - } + Some((old_status, new_status)) } + _ => None, + }; + + Some((result, transition)) + } + + fn transition_status_inner( + &self, + worker_id: &WorkerId, + expected_revision: Option, + new_status: WorkerStatus, + ) -> Option<(WorkerStatus, WorkerStatus)> { + let lock = self + .worker_mutation_locks + .entry(worker_id.clone()) + .or_insert_with(|| Arc::new(parking_lot::Mutex::new(()))) + .clone(); + let _guard = lock.lock(); + + let worker = self.workers.get(worker_id)?.clone(); + if expected_revision.is_some_and(|revision| worker.revision() != revision) { + return None; + } + + let old_status = worker.status(); + if old_status == new_status { + return None; + } + + worker.set_status(new_status); + + let _ = self.event_tx.send(WorkerEvent::StatusChanged { + worker_id: worker_id.clone(), + worker: worker.clone(), + old_status, + new_status, }); - HealthChecker::new(handle, shutdown_notify) + Some((old_status, new_status)) } } @@ -1078,8 +1092,20 @@ impl smg_mesh::WorkerStateSubscriber for WorkerRegistry { worker.set_healthy(state.health); - // register_inner skips mesh sync to avoid version-bump loop. - if let Some(id) = self.register_inner(Arc::new(worker)) { + // register_inner skips OUTGOING mesh sync to avoid a version-bump + // loop on the CRDT. We still publish the local `Registered` event + // so in-process subscribers (WorkerManager's health scheduler) + // pick up mesh-imported workers via the same event path as any + // other registration. Without this, mesh-synced workers would + // never enter the health schedule and `--remove-unhealthy-workers` + // could not reach them. The event is a local broadcast only; it + // does not re-enter the mesh. + let worker: Arc = Arc::new(worker); + if let Some(id) = self.register_inner(worker.clone()) { + let _ = self.event_tx.send(WorkerEvent::Registered { + worker_id: id.clone(), + worker: worker.clone(), + }); tracing::info!( worker_id = %id.as_str(), url = %state.url, @@ -1128,7 +1154,17 @@ mod tests { use openai_protocol::model_card::ModelCard; use super::*; - use crate::worker::{circuit_breaker::CircuitBreakerConfig, BasicWorkerBuilder}; + use crate::worker::{ + circuit_breaker::{CircuitBreakerConfig, CircuitState}, + BasicWorkerBuilder, WorkerLoadGuard, + }; + + fn no_health_check() -> openai_protocol::worker::HealthCheckConfig { + openai_protocol::worker::HealthCheckConfig { + disable_health_check: true, + ..Default::default() + } + } #[test] fn test_worker_registry() { @@ -1229,131 +1265,117 @@ mod tests { assert_eq!(llama_workers_after[0].url(), "http://worker2:8080"); } - #[tokio::test] - async fn test_pending_worker_stays_pending_under_failed_probes() { - use openai_protocol::worker::HealthCheckConfig; + // Health-checker integration tests moved to worker/manager.rs along with + // the loop itself. The registry is now a pure collection — see + // `worker::manager::WorkerManager` tests. - // Pending workers do NOT transition to NotReady on failed probes — - // they stay Pending until either `success_threshold` consecutive - // successes (→Ready) or `max_pending_probes` total attempts (→Failed). - // This test verifies the Pending state is sticky during early failures. + #[test] + fn test_transition_status_emits_event_and_changes_status() { let registry = WorkerRegistry::new(); - - let mut labels = HashMap::new(); - labels.insert("model_id".to_string(), "test-model".to_string()); + let mut rx = registry.subscribe_events(); let worker: Arc = Arc::new( - BasicWorkerBuilder::new("http://127.0.0.1:1") + BasicWorkerBuilder::new("http://w1:8080") .worker_type(WorkerType::Regular) - .labels(labels) - .health_config(HealthCheckConfig { - failure_threshold: 1, - success_threshold: 1, - timeout_secs: 1, - check_interval_secs: 1, - disable_health_check: false, + .health_config(openai_protocol::worker::HealthCheckConfig { + disable_health_check: true, + ..Default::default() }) .build(), ); - assert_eq!(worker.status(), WorkerStatus::Pending); - - registry.register(worker.clone()).unwrap(); - let _hc = registry.start_health_checker(1, true, None); + let worker_id = registry.register(worker.clone()).unwrap(); + // Drain Registered event + let _ = rx.try_recv().unwrap(); - // Wait for a few probe attempts. With check_interval=1 and a - // non-existent URL, probes will fail. With failure_threshold=1, - // max_pending_probes is 10 — so after ~3s we expect Pending still. - tokio::time::sleep(tokio::time::Duration::from_secs(3)).await; + // Initial status is Ready (disable_health_check). Transition to NotReady. + let result = registry.transition_status(&worker_id, WorkerStatus::NotReady); + assert_eq!(result, Some((WorkerStatus::Ready, WorkerStatus::NotReady))); + assert_eq!(worker.status(), WorkerStatus::NotReady); - let stats = registry.stats(); - assert_eq!(stats.total_workers, 1); - assert_eq!(stats.healthy_workers, 0, "Pending workers are not healthy"); - // Worker should still be Pending — not yet Failed (max_pending_probes - // is 10, we've only run ~3 probes) and not NotReady (Pending workers - // don't transition to NotReady on failure). - let after = registry.get_by_url("http://127.0.0.1:1").unwrap(); - assert_eq!( - after.status(), - WorkerStatus::Pending, - "Pending should be sticky under early failures" - ); + let event = rx.try_recv().unwrap(); + match event { + WorkerEvent::StatusChanged { + old_status, + new_status, + .. + } => { + assert_eq!(old_status, WorkerStatus::Ready); + assert_eq!(new_status, WorkerStatus::NotReady); + } + other => panic!("Expected StatusChanged, got {other:?}"), + } } #[test] - fn test_failed_worker_can_be_removed_from_registry() { - // Drives the removal-on-Failed logic without needing the async - // health checker loop or a JobQueue. Verifies that the registry - // accepts removal of a Failed worker via remove() and that the - // worker is gone afterward. - use openai_protocol::worker::HealthCheckConfig; - + fn test_transition_status_no_op_when_status_unchanged() { let registry = WorkerRegistry::new(); + let mut rx = registry.subscribe_events(); let worker: Arc = Arc::new( - BasicWorkerBuilder::new("http://failed:8080") + BasicWorkerBuilder::new("http://w1:8080") .worker_type(WorkerType::Regular) - .health_config(HealthCheckConfig { + .health_config(openai_protocol::worker::HealthCheckConfig { disable_health_check: true, - ..HealthCheckConfig::default() + ..Default::default() }) .build(), ); - let worker_id = registry.register(worker.clone()).unwrap(); - assert_eq!(worker.status(), WorkerStatus::Ready); + let worker_id = registry.register(worker).unwrap(); + let _ = rx.try_recv().unwrap(); - // Simulate the state machine reaching Failed. - worker.set_status(WorkerStatus::Failed); + // Already Ready — transition to Ready is a no-op assert_eq!( - registry.get(&worker_id).unwrap().status(), - WorkerStatus::Failed + registry.transition_status(&worker_id, WorkerStatus::Ready), + None ); - - // The health checker's removal loop only acts on workers in Failed state. - // Verify we can remove a Failed worker through the standard remove() API - // (which is what Job::RemoveWorker eventually invokes). - assert!(registry.remove(&worker_id).is_some()); - assert!(registry.get(&worker_id).is_none()); - assert_eq!(registry.stats().total_workers, 0); + assert!(rx.try_recv().is_err(), "no event should be emitted"); } - #[tokio::test] - async fn test_health_checker_keeps_workers_when_remove_unhealthy_disabled() { - use openai_protocol::worker::HealthCheckConfig; - - // When --remove-unhealthy-workers is false, workers in any state - // (Pending, NotReady, Failed) should remain in the registry. + #[test] + fn test_transition_status_returns_none_for_missing_worker() { let registry = WorkerRegistry::new(); + let missing = WorkerId::from_string("nonexistent".to_string()); + assert_eq!( + registry.transition_status(&missing, WorkerStatus::Ready), + None + ); + } - let mut labels = HashMap::new(); - labels.insert("model_id".to_string(), "test-model".to_string()); + #[test] + fn test_transition_status_if_revision_rejects_stale_transition() { + let registry = WorkerRegistry::new(); - let worker: Box = Box::new( - BasicWorkerBuilder::new("http://127.0.0.1:1") + let worker: Arc = Arc::new( + BasicWorkerBuilder::new("http://w1:8080") .worker_type(WorkerType::Regular) - .labels(labels) - .health_config(HealthCheckConfig { - failure_threshold: 1, - success_threshold: 1, - timeout_secs: 1, - check_interval_secs: 1, - disable_health_check: false, - }) + .health_config(no_health_check()) .build(), ); + let worker_id = registry.register(worker.clone()).unwrap(); + let stale_revision = worker.revision(); - registry.register(Arc::from(worker)).unwrap(); - assert_eq!(registry.stats().total_workers, 1); - - // Start health checker with remove_unhealthy=false - let _hc = registry.start_health_checker(1, false, None); - - tokio::time::sleep(tokio::time::Duration::from_secs(3)).await; + let replacement: Arc = Arc::new( + BasicWorkerBuilder::new("http://w1:8080") + .worker_type(WorkerType::Regular) + .health_config(no_health_check()) + .priority(99) + .build(), + ); + assert!(registry.replace(&worker_id, replacement)); assert_eq!( - registry.stats().total_workers, - 1, - "Worker should remain in the registry when remove_unhealthy is false" + registry.transition_status_if_revision( + &worker_id, + stale_revision, + WorkerStatus::NotReady + ), + None ); + + let current = registry.get(&worker_id).unwrap(); + assert_eq!(current.status(), WorkerStatus::Ready); + assert_eq!(current.priority(), 99); + assert_eq!(current.revision(), stale_revision + 1); } #[test] @@ -1518,6 +1540,56 @@ mod tests { assert_eq!(after.priority(), 99, "new priority should be applied"); } + #[test] + fn test_replace_preserves_runtime_state_and_circuit_breaker() { + let registry = WorkerRegistry::new(); + let mut headers = http::HeaderMap::new(); + headers.insert("x-smg-routing-key", "sticky-key".parse().unwrap()); + + let first: Arc = Arc::new( + BasicWorkerBuilder::new("http://worker:8080") + .worker_type(WorkerType::Regular) + .health_config(no_health_check()) + .build(), + ); + let first_id = registry.register(first.clone()).unwrap(); + let initial_revision = first.revision(); + + first.set_status(WorkerStatus::NotReady); + first.increment_processed(); + let load_guard = WorkerLoadGuard::new(first.clone(), Some(&headers)); + + for _ in 0..5 { + first.record_outcome(503); + } + assert_eq!(first.circuit_breaker_state(), CircuitState::Open); + + let replacement: Arc = Arc::new( + BasicWorkerBuilder::new("http://worker:8080") + .worker_type(WorkerType::Regular) + .health_config(no_health_check()) + .priority(99) + .build(), + ); + assert!(registry.replace(&first_id, replacement)); + + let current = registry.get(&first_id).unwrap(); + assert_eq!(current.priority(), 99); + assert_eq!(current.status(), WorkerStatus::NotReady); + assert_eq!(current.load(), 1); + assert_eq!(current.routing_key_load(), 1); + assert_eq!(current.processed_requests(), 1); + assert_eq!(current.circuit_breaker_state(), CircuitState::Open); + assert_eq!(current.revision(), initial_revision + 1); + + first.increment_processed(); + assert_eq!(current.processed_requests(), 2); + + drop(load_guard); + assert_eq!(current.load(), 0); + assert_eq!(current.routing_key_load(), 0); + } + #[test] fn test_builder_default_status_is_pending() { // Without an explicit override, health-checked workers start Pending. @@ -1675,6 +1747,86 @@ mod tests { ); } + #[test] + fn test_mesh_imported_worker_emits_registered_event() { + use smg_mesh::{WorkerState, WorkerStateSubscriber}; + + // Mesh-imported workers must emit `WorkerEvent::Registered` so + // event-driven subscribers (WorkerManager's health scheduler) + // pick them up via the same path as any other registration. + // Without this event, a mesh worker would be routable but never + // probed locally, and `--remove-unhealthy-workers` could not + // reach it. + let registry = WorkerRegistry::new(); + let mut rx = registry.subscribe_events(); + + let state = WorkerState { + worker_id: "mesh-w1".into(), + model_id: "llama-3".into(), + url: "http://mesh-worker-event:8080".into(), + health: true, + load: 0.5, + version: 1, + spec: vec![], + }; + + registry.on_remote_worker_state(&state); + + let event = rx + .try_recv() + .expect("mesh import must broadcast a Registered event"); + match event { + WorkerEvent::Registered { worker, .. } => { + assert_eq!(worker.url(), "http://mesh-worker-event:8080"); + assert_eq!(worker.model_id(), "llama-3"); + } + other => panic!("Expected Registered event from mesh import, got: {other:?}"), + } + } + + #[test] + fn test_mesh_worker_state_update_is_silent_on_existing_worker() { + use smg_mesh::{WorkerState, WorkerStateSubscriber}; + + // Health updates for an already-registered worker go through the + // compat shim (set_healthy) without emitting an event — the + // existing worker is already on WorkerManager's schedule, and + // local probes will reconcile the state on the next tick. We + // verify the no-event behavior so a future regression (e.g. adding + // a StatusChanged emit here) is caught. + let registry = WorkerRegistry::new(); + + let worker: Arc = Arc::new( + BasicWorkerBuilder::new("http://mesh-existing:8080") + .model(ModelCard::new("llama-3")) + .health_config(openai_protocol::worker::HealthCheckConfig { + disable_health_check: true, + ..Default::default() + }) + .build(), + ); + registry.register(worker).unwrap(); + + let mut rx = registry.subscribe_events(); + + let state = WorkerState { + worker_id: "mesh-w1".into(), + model_id: "llama-3".into(), + url: "http://mesh-existing:8080".into(), + health: false, + load: 0.0, + version: 2, + spec: vec![], + }; + registry.on_remote_worker_state(&state); + + // No event is expected on the update path. + assert!( + rx.try_recv().is_err(), + "existing-worker update path should not broadcast an event" + ); + } + #[test] fn test_worker_event_broadcast() { let registry = WorkerRegistry::new(); diff --git a/model_gateway/src/worker/service.rs b/model_gateway/src/worker/service.rs index 57fcab63d7..5c6a0cc530 100644 --- a/model_gateway/src/worker/service.rs +++ b/model_gateway/src/worker/service.rs @@ -388,7 +388,10 @@ impl WorkerService { worker_id: worker_id_raw.to_string(), })?; - let job = Job::RemoveWorker { url: url.clone() }; + let job = Job::RemoveWorker { + url: url.clone(), + expected_revision: None, + }; let job_queue = self.get_job_queue()?; job_queue diff --git a/model_gateway/src/worker/worker.rs b/model_gateway/src/worker/worker.rs index d2409cfd79..e072cbf569 100644 --- a/model_gateway/src/worker/worker.rs +++ b/model_gateway/src/worker/worker.rs @@ -1,7 +1,8 @@ use std::{ + any::Any, fmt, sync::{ - atomic::{AtomicU8, AtomicUsize, Ordering}, + atomic::{AtomicU64, AtomicU8, AtomicUsize, Ordering}, Arc, LazyLock, }, time::Duration, @@ -116,7 +117,10 @@ impl fmt::Debug for WorkerRoutingKeyLoad { /// Core worker abstraction that represents a backend service #[async_trait] -pub trait Worker: Send + Sync + fmt::Debug { +pub trait Worker: Send + Sync + fmt::Debug + 'static { + /// Downcast support for same-URL replace state sharing. + fn as_any(&self) -> &dyn Any; + /// Get the worker's URL fn url(&self) -> &str; /// Get the worker's API key @@ -144,9 +148,25 @@ pub trait Worker: Send + Sync + fmt::Debug { /// Get the worker's lifecycle status. fn status(&self) -> WorkerStatus; + /// Get the current monotonic worker revision. + /// + /// Same-URL `replace()` increments the revision so stale probe outcomes + /// can be discarded without mutating the newly installed worker object. + fn revision(&self) -> u64 { + 0 + } + /// Set the worker's lifecycle status. fn set_status(&self, status: WorkerStatus); + /// Adopt shared mutable runtime state from a previous worker object. + /// + /// Used by same-URL `replace()` so in-flight traffic and counters remain + /// attached to a single shared runtime across the old and new objects. + fn inherit_shared_state_from(&self, _other: &dyn Worker) -> bool { + false + } + /// Check if the worker is currently healthy (status == Ready). /// /// This is a routing predicate — returns true only for `Ready` workers. @@ -171,9 +191,37 @@ pub trait Worker: Send + Sync + fmt::Debug { } } - /// Perform an async health check on the worker + /// Perform an async health check on the worker. + /// + /// Pure probe — does not mutate worker status, does not increment + /// counters beyond `consecutive_*` totals exposed via the accessors + /// below. The state machine lives in `WorkerManager`, which reads the + /// counters and applies transitions via `WorkerRegistry::transition_status()`. async fn check_health_async(&self) -> WorkerResult<()>; + // ── Health check counter accessors (used by WorkerManager state machine) ── + + /// Increment `consecutive_failures` and return the new value. + fn consecutive_failures_increment(&self) -> usize; + + /// Reset `consecutive_failures` to 0. + fn consecutive_failures_reset(&self); + + /// Increment `consecutive_successes` and return the new value. + fn consecutive_successes_increment(&self) -> usize; + + /// Reset `consecutive_successes` to 0. + fn consecutive_successes_reset(&self); + + /// Read `total_pending_probes` (lifetime probe attempts in Pending state). + fn total_pending_probes(&self) -> usize; + + /// Increment `total_pending_probes` and return the new value. + fn total_pending_probes_increment(&self) -> usize; + + /// Reset `total_pending_probes` to 0 (called when promoting Pending → Ready). + fn total_pending_probes_reset(&self); + /// Get the current load (number of active requests) fn load(&self) -> usize; @@ -186,8 +234,14 @@ pub trait Worker: Send + Sync + fmt::Debug { /// Reset the load counter to 0 (for sync/recovery) fn reset_load(&self) {} - /// Get the worker routing key load tracker - fn worker_routing_key_load(&self) -> &WorkerRoutingKeyLoad; + /// Get the current routing-key load cardinality. + fn routing_key_load(&self) -> usize; + + /// Increment the routing-key load tracker for an active key. + fn increment_routing_key_load(&self, routing_key: &str); + + /// Decrement the routing-key load tracker for a completed key. + fn decrement_routing_key_load(&self, routing_key: &str); /// Get the number of processed requests fn processed_requests(&self) -> usize; @@ -198,15 +252,18 @@ pub trait Worker: Send + Sync + fmt::Debug { /// Get worker-specific metadata fn metadata(&self) -> &WorkerMetadata; - /// Get the circuit breaker for this worker. - /// - /// **Do not call from routers.** Use `record_outcome(status_code)` to - /// record request outcomes and `is_available()` to check worker health. - fn circuit_breaker(&self) -> &CircuitBreaker; + /// Get the current circuit breaker state for observability/debugging. + fn circuit_breaker_state(&self) -> super::circuit_breaker::CircuitState; + + /// Check whether the current circuit breaker state allows execution. + fn circuit_breaker_can_execute(&self) -> bool; + + /// Record a request outcome against the circuit breaker. + fn record_circuit_breaker_outcome(&self, success: bool); /// Check if the worker is available (healthy + circuit closed/half-open) fn is_available(&self) -> bool { - self.is_healthy() && self.circuit_breaker().can_execute() + self.is_healthy() && self.circuit_breaker_can_execute() } /// Record the outcome of a request based on the HTTP status code. @@ -223,7 +280,7 @@ pub trait Worker: Send + Sync + fmt::Debug { .resilience() .retryable_status_codes .contains(&status_code); - self.circuit_breaker().record_outcome(!is_failure); + self.record_circuit_breaker_outcome(!is_failure); } /// Get the resolved resilience config for this worker. @@ -517,21 +574,55 @@ impl WorkerMetadata { } } +/// Shared mutable worker state preserved across same-URL replacements. +#[derive(Debug)] +pub struct WorkerRuntime { + status: AtomicU8, + consecutive_failures: AtomicUsize, + consecutive_successes: AtomicUsize, + total_pending_probes: AtomicUsize, + load_counter: AtomicUsize, + processed_counter: AtomicUsize, + worker_routing_key_load: WorkerRoutingKeyLoad, + revision: AtomicU64, +} + +impl WorkerRuntime { + pub fn new(url: &str, initial_status: WorkerStatus) -> Self { + Self { + status: AtomicU8::new(initial_status as u8), + consecutive_failures: AtomicUsize::new(0), + consecutive_successes: AtomicUsize::new(0), + total_pending_probes: AtomicUsize::new(0), + load_counter: AtomicUsize::new(0), + processed_counter: AtomicUsize::new(0), + worker_routing_key_load: WorkerRoutingKeyLoad::new(url), + revision: AtomicU64::new(0), + } + } + + fn status(&self) -> WorkerStatus { + WorkerStatus::from_u8(self.status.load(Ordering::Acquire)) + } + + fn set_status(&self, status: WorkerStatus) { + self.status.store(status as u8, Ordering::Release); + } + + fn revision(&self) -> u64 { + self.revision.load(Ordering::Acquire) + } + + fn bump_revision(&self) -> u64 { + self.revision.fetch_add(1, Ordering::AcqRel) + 1 + } +} + /// Basic worker implementation -#[derive(Clone)] pub struct BasicWorker { pub metadata: WorkerMetadata, - pub load_counter: Arc, - pub worker_routing_key_load: Arc, - pub processed_counter: Arc, - pub status: Arc, - pub consecutive_failures: Arc, - pub consecutive_successes: Arc, - /// Total health check probes attempted while in Pending state. - /// Unlike consecutive counters, this counts ALL attempts (not just consecutive). - /// Used to detect misconfigured workers stuck in Pending. - pub total_pending_probes: Arc, - pub circuit_breaker: CircuitBreaker, + pub runtime: ArcSwap, + pub circuit_breaker: ArcSwap, /// Lazily initialized gRPC client for gRPC workers. /// Uses OnceCell for lock-free reads after initialization. pub grpc_client: Arc>>, @@ -545,15 +636,28 @@ pub struct BasicWorker { pub resilience: ResolvedResilience, } +impl Clone for BasicWorker { + fn clone(&self) -> Self { + Self { + metadata: self.metadata.clone(), + runtime: ArcSwap::from(self.runtime.load_full()), + circuit_breaker: ArcSwap::from(self.circuit_breaker.load_full()), + grpc_client: Arc::clone(&self.grpc_client), + models_override: Arc::clone(&self.models_override), + http_client: self.http_client.clone(), + resilience: self.resilience.clone(), + } + } +} + impl fmt::Debug for BasicWorker { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let runtime = self.runtime.load(); f.debug_struct("BasicWorker") .field("metadata", &self.metadata) - .field( - "status", - &WorkerStatus::from_u8(self.status.load(Ordering::Relaxed)), - ) - .field("circuit_breaker", &self.circuit_breaker) + .field("status", &runtime.status()) + .field("revision", &runtime.revision()) + .field("circuit_breaker_state", &self.circuit_breaker_state()) .field("grpc_client", &"") .finish() } @@ -564,10 +668,32 @@ impl BasicWorker { let load = self.load(); Metrics::set_worker_requests_active(self.url(), load); } + + fn shared_runtime(&self) -> Arc { + self.runtime.load_full() + } + + fn install_shared_state_from_basic(&self, other: &BasicWorker) { + let shared_runtime = other.shared_runtime(); + shared_runtime.bump_revision(); + let shared_status = shared_runtime.status(); + self.runtime.store(shared_runtime); + Metrics::set_worker_health(self.url(), shared_status == WorkerStatus::Ready); + + let existing_cb = self.circuit_breaker.load(); + let other_cb = other.circuit_breaker.load_full(); + if other_cb.config() == existing_cb.config() { + self.circuit_breaker.store(other_cb); + } + } } #[async_trait] impl Worker for BasicWorker { + fn as_any(&self) -> &dyn Any { + self + } + fn url(&self) -> &str { &self.metadata.spec.url } @@ -585,148 +711,121 @@ impl Worker for BasicWorker { } fn status(&self) -> WorkerStatus { - WorkerStatus::from_u8(self.status.load(Ordering::Acquire)) + self.runtime.load().status() + } + + fn revision(&self) -> u64 { + self.runtime.load().revision() } fn set_status(&self, status: WorkerStatus) { - self.status.store(status as u8, Ordering::Release); + self.runtime.load().set_status(status); Metrics::set_worker_health(self.url(), status == WorkerStatus::Ready); } + fn inherit_shared_state_from(&self, other: &dyn Worker) -> bool { + let Some(other) = other.as_any().downcast_ref::() else { + return false; + }; + self.install_shared_state_from_basic(other); + true + } + async fn check_health_async(&self) -> WorkerResult<()> { + // Pure probe — no state machine, no counter mutation, no status + // changes. The HealthChecker (in WorkerManager) owns the state machine + // and calls registry.transition_status() to apply transitions. + // + // Returns Ok(()) if the underlying transport probe succeeded, + // Err(HealthCheckFailed) otherwise. Transport-level errors (gRPC + // connect failure, TLS handshake, DNS) are surfaced as Err so the + // caller can log them with worker URL context. if self.metadata.health_config.disable_health_check { - if self.status() != WorkerStatus::Ready { - self.set_status(WorkerStatus::Ready); - } return Ok(()); } - let current_status = self.status(); - let health_config = &self.metadata.health_config; - let worker_type_str = self.metadata.spec.worker_type.as_metric_label(); - - // Failed is terminal — skip the probe entirely. The health checker - // loop will eventually remove the worker (if --remove-unhealthy-workers - // is set) or it stays Failed until explicit re-registration. - if current_status == WorkerStatus::Failed { - return Err(WorkerError::HealthCheckFailed { - url: self.metadata.spec.url.clone(), - reason: "worker is in Failed state".to_string(), - }); - } - - // Track total probes in Pending for startup timeout detection - if current_status == WorkerStatus::Pending { - self.total_pending_probes.fetch_add(1, Ordering::Relaxed); - } - - // Transport-level errors (e.g. gRPC connect failure, TLS handshake - // failure, DNS resolution failure) are treated as failed probes - // rather than short-circuiting, so the Pending timeout and - // NotReady→Failed paths can observe them. The error is logged with - // the worker URL so operators can debug why a worker won't come Ready. - let health_result = match &self.metadata.spec.connection_mode { - ConnectionMode::Http => self.http_health_check().await, - ConnectionMode::Grpc => self.grpc_health_check().await, - } - .unwrap_or_else(|err| { - tracing::warn!( - worker_url = %self.metadata.spec.url, - error = %err, - "Health check probe transport error (treating as failed probe)" - ); - false - }); + let probe_ok = match &self.metadata.spec.connection_mode { + ConnectionMode::Http => self.http_health_check().await?, + ConnectionMode::Grpc => self.grpc_health_check().await?, + }; - if health_result { - self.consecutive_failures.store(0, Ordering::Release); - let successes = self.consecutive_successes.fetch_add(1, Ordering::AcqRel) + 1; - Metrics::record_worker_health_check(worker_type_str, metrics_labels::CB_SUCCESS); - - // Pending → Ready or NotReady → Ready on success_threshold. - // Failed is terminal — workers that reached Failed must be removed - // and re-registered. (replace() preserves the old status, so a - // metadata update on a Failed worker would not recover it.) - if matches!( - current_status, - WorkerStatus::Pending | WorkerStatus::NotReady - ) && successes >= health_config.success_threshold as usize - { - self.set_status(WorkerStatus::Ready); - self.consecutive_successes.store(0, Ordering::Release); - // Reset pending probe counter on successful promotion - self.total_pending_probes.store(0, Ordering::Relaxed); - } else if current_status == WorkerStatus::Pending { - // Even on a successful probe, enforce the Pending cap. A - // worker that flaps F,S,F,S,... never reaches success_threshold - // and would otherwise grow `total_pending_probes` without bound. - let max_pending = health_config.failure_threshold as usize * 10; - if self.total_pending_probes.load(Ordering::Relaxed) >= max_pending { - self.set_status(WorkerStatus::Failed); - self.consecutive_successes.store(0, Ordering::Release); - self.consecutive_failures.store(0, Ordering::Release); - } - } + if probe_ok { Ok(()) } else { - self.consecutive_successes.store(0, Ordering::Release); - let failures = self.consecutive_failures.fetch_add(1, Ordering::AcqRel) + 1; - Metrics::record_worker_health_check(worker_type_str, metrics_labels::CB_FAILURE); - - match current_status { - WorkerStatus::Ready => { - // Ready → NotReady on failure_threshold - if failures >= health_config.failure_threshold as usize { - self.set_status(WorkerStatus::NotReady); - self.consecutive_failures.store(0, Ordering::Release); - } - } - WorkerStatus::NotReady => { - // NotReady → Failed on liveness_failure_threshold (3× failure_threshold) - let liveness_threshold = health_config.failure_threshold as usize * 3; - if failures >= liveness_threshold { - self.set_status(WorkerStatus::Failed); - self.consecutive_failures.store(0, Ordering::Release); - } - } - WorkerStatus::Pending => { - // Pending → Failed on max_pending_probes (10× failure_threshold). - // Note: this uses `total_pending_probes` (total attempts), not - // `consecutive_failures`, because a worker flapping between - // success/failure in Pending without reaching success_threshold - // is also suspect. Reset consecutive_failures for consistency - // with other failure-path transitions (Failed is terminal, but - // this keeps counter state clean). - let max_pending = health_config.failure_threshold as usize * 10; - let total = self.total_pending_probes.load(Ordering::Relaxed); - if total >= max_pending { - self.set_status(WorkerStatus::Failed); - self.consecutive_failures.store(0, Ordering::Release); - } - } - WorkerStatus::Failed => { - // Already failed — no further transitions - } - } - Err(WorkerError::HealthCheckFailed { url: self.metadata.spec.url.clone(), - reason: format!("Health check failed (consecutive failures: {failures})"), + reason: "health probe returned non-success".to_string(), }) } } + fn consecutive_failures_increment(&self) -> usize { + self.runtime + .load() + .consecutive_failures + .fetch_add(1, Ordering::AcqRel) + + 1 + } + + fn consecutive_failures_reset(&self) { + self.runtime + .load() + .consecutive_failures + .store(0, Ordering::Release); + } + + fn consecutive_successes_increment(&self) -> usize { + self.runtime + .load() + .consecutive_successes + .fetch_add(1, Ordering::AcqRel) + + 1 + } + + fn consecutive_successes_reset(&self) { + self.runtime + .load() + .consecutive_successes + .store(0, Ordering::Release); + } + + fn total_pending_probes(&self) -> usize { + self.runtime + .load() + .total_pending_probes + .load(Ordering::Relaxed) + } + + fn total_pending_probes_increment(&self) -> usize { + self.runtime + .load() + .total_pending_probes + .fetch_add(1, Ordering::Relaxed) + + 1 + } + + fn total_pending_probes_reset(&self) { + self.runtime + .load() + .total_pending_probes + .store(0, Ordering::Relaxed); + } + fn load(&self) -> usize { - self.load_counter.load(Ordering::Relaxed) + self.runtime.load().load_counter.load(Ordering::Relaxed) } fn increment_load(&self) { - self.load_counter.fetch_add(1, Ordering::Relaxed); + self.runtime + .load() + .load_counter + .fetch_add(1, Ordering::Relaxed); self.update_running_requests_metrics(); } fn decrement_load(&self) { - if self + let runtime = self.runtime.load(); + if runtime .load_counter .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| { current.checked_sub(1) @@ -742,28 +841,56 @@ impl Worker for BasicWorker { } fn reset_load(&self) { - self.load_counter.store(0, Ordering::Relaxed); + self.runtime.load().load_counter.store(0, Ordering::Relaxed); self.update_running_requests_metrics(); } - fn worker_routing_key_load(&self) -> &WorkerRoutingKeyLoad { - &self.worker_routing_key_load + fn routing_key_load(&self) -> usize { + self.runtime.load().worker_routing_key_load.value() + } + + fn increment_routing_key_load(&self, routing_key: &str) { + self.runtime + .load() + .worker_routing_key_load + .increment(routing_key); + } + + fn decrement_routing_key_load(&self, routing_key: &str) { + self.runtime + .load() + .worker_routing_key_load + .decrement(routing_key); } fn processed_requests(&self) -> usize { - self.processed_counter.load(Ordering::Relaxed) + self.runtime + .load() + .processed_counter + .load(Ordering::Relaxed) } fn increment_processed(&self) { - self.processed_counter.fetch_add(1, Ordering::Relaxed); + self.runtime + .load() + .processed_counter + .fetch_add(1, Ordering::Relaxed); } fn metadata(&self) -> &WorkerMetadata { &self.metadata } - fn circuit_breaker(&self) -> &CircuitBreaker { - &self.circuit_breaker + fn circuit_breaker_state(&self) -> super::circuit_breaker::CircuitState { + self.circuit_breaker.load().state() + } + + fn circuit_breaker_can_execute(&self) -> bool { + self.circuit_breaker.load().can_execute() + } + + fn record_circuit_breaker_outcome(&self, success: bool) { + self.circuit_breaker.load().record_outcome(success); } fn resilience(&self) -> &ResolvedResilience { @@ -945,7 +1072,7 @@ impl WorkerLoadGuard { let routing_key = extract_routing_key(headers).map(String::from); if let Some(ref key) = routing_key { - worker.worker_routing_key_load().increment(key); + worker.increment_routing_key_load(key); } Self { @@ -959,7 +1086,7 @@ impl Drop for WorkerLoadGuard { fn drop(&mut self) { self.worker.decrement_load(); if let Some(ref key) = self.routing_key { - self.worker.worker_routing_key_load().decrement(key); + self.worker.decrement_routing_key_load(key); } } } @@ -1014,56 +1141,6 @@ impl http_body::Body for AttachedBody { } } -/// Health checker handle with graceful shutdown. -/// -/// The checker sleeps until the next worker is due for a health check, -/// so it wakes only when there is actual work to do. -pub(crate) struct HealthChecker { - handle: Option>, - shutdown_notify: Arc, -} - -impl fmt::Debug for HealthChecker { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("HealthChecker").finish() - } -} - -impl HealthChecker { - pub fn new( - handle: tokio::task::JoinHandle<()>, - shutdown_notify: Arc, - ) -> Self { - Self { - handle: Some(handle), - shutdown_notify, - } - } - - /// Shutdown the health checker gracefully. - /// Wakes the sleeping task immediately so it can exit cleanly. - /// Prefer this over dropping when you can `.await` — it lets the - /// current health-check iteration finish instead of aborting mid-flight. - #[expect( - dead_code, - reason = "Drop::drop handles abort; this exists for graceful shutdown when an async context is available" - )] - pub async fn shutdown(&mut self) { - self.shutdown_notify.notify_one(); - if let Some(handle) = self.handle.take() { - let _ = handle.await; - } - } -} - -impl Drop for HealthChecker { - fn drop(&mut self) { - if let Some(handle) = self.handle.take() { - handle.abort(); - } - } -} - /// Helper to convert Worker trait object to WorkerInfo struct. /// /// Both `is_healthy` and `status` are derived from the same atomic snapshot @@ -1622,7 +1699,7 @@ mod tests { .build(); assert!(worker.is_available()); - assert_eq!(worker.circuit_breaker().state(), CircuitState::Closed); + assert_eq!(worker.circuit_breaker_state(), CircuitState::Closed); worker.record_outcome(503); worker.record_outcome(503); @@ -1635,7 +1712,7 @@ mod tests { assert!(!worker.is_available()); assert!(worker.is_healthy()); - assert!(!worker.circuit_breaker().can_execute()); + assert!(!worker.circuit_breaker_can_execute()); } #[test] @@ -1661,10 +1738,10 @@ mod tests { thread::sleep(Duration::from_millis(150)); assert!(worker.is_available()); - assert_eq!(worker.circuit_breaker().state(), CircuitState::HalfOpen); + assert_eq!(worker.circuit_breaker_state(), CircuitState::HalfOpen); worker.record_outcome(200); - assert_eq!(worker.circuit_breaker().state(), CircuitState::Closed); + assert_eq!(worker.circuit_breaker_state(), CircuitState::Closed); } #[test] @@ -1682,7 +1759,7 @@ mod tests { } assert!(!dp_worker.is_available()); - assert_eq!(dp_worker.circuit_breaker().state(), CircuitState::Open); + assert_eq!(dp_worker.circuit_breaker_state(), CircuitState::Open); } #[tokio::test] @@ -1888,7 +1965,7 @@ mod tests { ); assert_eq!(worker.load(), 0); - assert_eq!(worker.worker_routing_key_load().value(), 0); + assert_eq!(worker.routing_key_load(), 0); let mut headers = http::HeaderMap::new(); headers.insert("x-smg-routing-key", "key-123".parse().unwrap()); @@ -1896,11 +1973,11 @@ mod tests { { let _guard = WorkerLoadGuard::new(worker.clone(), Some(&headers)); assert_eq!(worker.load(), 1); - assert_eq!(worker.worker_routing_key_load().value(), 1); + assert_eq!(worker.routing_key_load(), 1); } assert_eq!(worker.load(), 0); - assert_eq!(worker.worker_routing_key_load().value(), 0); + assert_eq!(worker.routing_key_load(), 0); } #[test] @@ -1914,16 +1991,16 @@ mod tests { ); assert_eq!(worker.load(), 0); - assert_eq!(worker.worker_routing_key_load().value(), 0); + assert_eq!(worker.routing_key_load(), 0); { let _guard = WorkerLoadGuard::new(worker.clone(), None); assert_eq!(worker.load(), 1); - assert_eq!(worker.worker_routing_key_load().value(), 0); + assert_eq!(worker.routing_key_load(), 0); } assert_eq!(worker.load(), 0); - assert_eq!(worker.worker_routing_key_load().value(), 0); + assert_eq!(worker.routing_key_load(), 0); } #[test] @@ -1941,19 +2018,19 @@ mod tests { let guard1 = WorkerLoadGuard::new(worker.clone(), Some(&headers)); assert_eq!(worker.load(), 1); - assert_eq!(worker.worker_routing_key_load().value(), 1); + assert_eq!(worker.routing_key_load(), 1); let guard2 = WorkerLoadGuard::new(worker.clone(), Some(&headers)); assert_eq!(worker.load(), 2); - assert_eq!(worker.worker_routing_key_load().value(), 1); + assert_eq!(worker.routing_key_load(), 1); drop(guard1); assert_eq!(worker.load(), 1); - assert_eq!(worker.worker_routing_key_load().value(), 1); + assert_eq!(worker.routing_key_load(), 1); drop(guard2); assert_eq!(worker.load(), 0); - assert_eq!(worker.worker_routing_key_load().value(), 0); + assert_eq!(worker.routing_key_load(), 0); } #[test] diff --git a/model_gateway/src/workflow/job_queue.rs b/model_gateway/src/workflow/job_queue.rs index 5cc493739c..da4dc1e78b 100644 --- a/model_gateway/src/workflow/job_queue.rs +++ b/model_gateway/src/workflow/job_queue.rs @@ -41,6 +41,7 @@ pub enum Job { }, RemoveWorker { url: String, + expected_revision: Option, }, InitializeWorkersFromConfig { router_config: Box, @@ -87,7 +88,7 @@ impl Job { match self { Job::AddWorker { config } => &config.url, Job::UpdateWorker { url, .. } => url, - Job::RemoveWorker { url } => url, + Job::RemoveWorker { url, .. } => url, Job::InitializeWorkersFromConfig { .. } => "startup", Job::InitializeMcpServers { .. } => "startup", Job::RegisterMcpServer { config } => &config.name, @@ -366,7 +367,10 @@ impl JobQueue { .wait_for_completion(instance_id, url, timeout_duration) .await } - Job::RemoveWorker { url } => { + Job::RemoveWorker { + url, + expected_revision, + } => { let engines = context .workflow_engines .get() @@ -375,6 +379,7 @@ impl JobQueue { let workflow_data = create_worker_removal_workflow_data( url.to_string(), context.router_config.dp_aware, + *expected_revision, Arc::clone(context), ); diff --git a/model_gateway/src/workflow/steps/local/find_workers_to_remove.rs b/model_gateway/src/workflow/steps/local/find_workers_to_remove.rs index 741c369a14..43faa5bb98 100644 --- a/model_gateway/src/workflow/steps/local/find_workers_to_remove.rs +++ b/model_gateway/src/workflow/steps/local/find_workers_to_remove.rs @@ -14,6 +14,7 @@ use crate::workflow::data::{WorkerList, WorkerRemovalWorkflowData}; pub struct WorkerRemovalRequest { pub url: String, pub dp_aware: bool, + pub expected_revision: Option, } /// Step to find workers to remove based on URL. @@ -35,10 +36,24 @@ impl StepExecutor for FindWorkersToRemoveStep { .as_ref() .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?; - let workers_to_remove = + let mut workers_to_remove = find_workers_by_url(&app_context.worker_registry, &request.url, request.dp_aware); - if workers_to_remove.is_empty() { + if let Some(expected_revision) = request.expected_revision { + workers_to_remove.retain(|worker| worker.revision() == expected_revision); + if workers_to_remove.is_empty() { + debug!( + worker_url = %request.url, + expected_revision, + "Skipping stale worker removal job after same-URL replacement" + ); + context.data.workers_to_remove = Some(WorkerList::new()); + context.data.actual_workers_to_remove = Some(Vec::new()); + context.data.worker_urls = Vec::new(); + context.data.affected_models = HashSet::new(); + return Ok(StepResult::Success); + } + } else if workers_to_remove.is_empty() { let error_msg = if request.dp_aware { format!("No workers found with prefix {}@", request.url) } else { diff --git a/model_gateway/src/workflow/steps/local/mod.rs b/model_gateway/src/workflow/steps/local/mod.rs index ab299a0dcf..446fac61c3 100644 --- a/model_gateway/src/workflow/steps/local/mod.rs +++ b/model_gateway/src/workflow/steps/local/mod.rs @@ -184,10 +184,15 @@ pub fn create_worker_update_workflow() -> WorkflowDefinition, app_context: Arc, ) -> WorkerRemovalWorkflowData { WorkerRemovalWorkflowData { - config: WorkerRemovalRequest { url, dp_aware }, + config: WorkerRemovalRequest { + url, + dp_aware, + expected_revision, + }, workers_to_remove: None, worker_urls: Vec::new(), affected_models: std::collections::HashSet::new(), diff --git a/model_gateway/src/workflow/steps/local/update_worker_properties.rs b/model_gateway/src/workflow/steps/local/update_worker_properties.rs index c89a9ac733..e5b1d4b432 100644 --- a/model_gateway/src/workflow/steps/local/update_worker_properties.rs +++ b/model_gateway/src/workflow/steps/local/update_worker_properties.rs @@ -122,10 +122,20 @@ impl StepExecutor for UpdateWorkerPropertiesStep { let new_worker: Arc = Arc::new(builder.build()); // Replace the worker in the registry (overwrite-then-diff) - app_context + let worker_id = app_context .worker_registry .register_or_replace(new_worker.clone()); + // Same-URL replace preserves the existing shared runtime. If the + // update disables health checks, force the lifecycle to Ready so + // the worker does not stay stuck in a non-routable pre-update + // status while the manager correctly skips future probes. + if updated_health_config.disable_health_check { + let _ = app_context + .worker_registry + .transition_status(&worker_id, WorkerStatus::Ready); + } + updated_workers.push(new_worker); }