[model-gateway] optimize worker registry and reduce lock contention in grpc client fetch (#15336)

sgl-model-gateway/src/core/circuit_breaker.rs

@@ -374,8 +374,14 @@ impl CircuitBreaker {
     }
 
     fn publish_gauge_metrics(&self) {
-        Metrics::set_worker_cb_consecutive_failures(&self.metric_label, self.failure_count());
+        Metrics::set_worker_cb_consecutive_failures(
-        Metrics::set_worker_cb_consecutive_successes(&self.metric_label, self.success_count());
+            &self.metric_label,
+            self.consecutive_failures(),
+        );
+        Metrics::set_worker_cb_consecutive_successes(
+            &self.metric_label,
+            self.consecutive_successes(),
+        );
     }
 }
 

sgl-model-gateway/src/core/worker.rs

@@ -10,7 +10,7 @@ use std::{
 use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
 use serde_json;
-use tokio::{sync::RwLock, time};
+use tokio::{sync::OnceCell, time};
 
 use super::{
     CircuitBreaker, Endpoint, ModelCard, ModelType, ProviderType, WorkerError, WorkerResult,
@@ -515,8 +515,9 @@ pub struct BasicWorker {
     pub consecutive_failures: Arc<AtomicUsize>,
     pub consecutive_successes: Arc<AtomicUsize>,
     pub circuit_breaker: CircuitBreaker,
-    /// Lazily initialized gRPC client for gRPC workers
+    /// Lazily initialized gRPC client for gRPC workers.
-    pub grpc_client: Arc<RwLock<Option<Arc<GrpcClient>>>>,
+    /// Uses OnceCell for lock-free reads after initialization.
+    pub grpc_client: Arc<OnceCell<Arc<GrpcClient>>>,
     /// Runtime-mutable models override (for lazy discovery)
     /// When set, overrides metadata.models for routing decisions.
     /// Uses std::sync::RwLock for synchronous access in supports_model().
@@ -715,19 +716,11 @@ impl Worker for BasicWorker {
         match self.metadata.connection_mode {
             ConnectionMode::Http => Ok(None),
             ConnectionMode::Grpc { .. } => {
-                {
+                // OnceCell provides lock-free reads after initialization.
-                    let client_guard = self.grpc_client.read().await;
+                // get_or_try_init only acquires internal lock on first call.
-                    if let Some(ref client) = *client_guard {
+                let client = self
-                        return Ok(Some(client.clone()));
+                    .grpc_client
-                    }
+                    .get_or_try_init(|| async {
-                }
-
-                let mut client_guard = self.grpc_client.write().await;
-
-                if let Some(ref client) = *client_guard {
-                    return Ok(Some(client.clone()));
-                }
-
                         let runtime_str = self.metadata.runtime_type.to_string();
                         tracing::info!(
                             "Lazily initializing gRPC client ({}) for worker: {}",
@@ -736,14 +729,12 @@ impl Worker for BasicWorker {
                         );
                         match GrpcClient::connect(&self.metadata.url, &runtime_str).await {
                             Ok(client) => {
-                        let client_arc = Arc::new(client);
-                        *client_guard = Some(client_arc.clone());
                                 tracing::info!(
                                     "Successfully connected gRPC client ({}) for worker: {}",
                                     runtime_str,
                                     self.metadata.url
                                 );
-                        Ok(Some(client_arc))
+                                Ok(Arc::new(client))
                             }
                             Err(e) => {
                                 tracing::error!(
@@ -757,23 +748,22 @@ impl Worker for BasicWorker {
                                 })
                             }
                         }
+                    })
+                    .await?;
+                Ok(Some(Arc::clone(client)))
             }
         }
     }
 
     async fn reset_grpc_client(&self) -> WorkerResult<()> {
-        match self.metadata.connection_mode {
+        // OnceCell doesn't support resetting. This is intentional for lock-free performance.
-            ConnectionMode::Http => Ok(()),
+        // If a connection fails, the worker should be removed and re-added.
-            ConnectionMode::Grpc { .. } => {
+        tracing::debug!(
-                let mut client_guard = self.grpc_client.write().await;
+            "reset_grpc_client called for {} (no-op with OnceCell)",
-                if client_guard.is_some() {
+            self.metadata.url
-                    tracing::info!("Resetting gRPC client for worker: {}", self.metadata.url);
+        );
-                    *client_guard = None;
-                }
         Ok(())
     }
-        }
-    }
 
     async fn grpc_health_check(&self) -> WorkerResult<bool> {
         let timeout = Duration::from_secs(self.metadata.health_config.timeout_secs);

sgl-model-gateway/src/core/worker_builder.rs

@@ -131,7 +131,7 @@ impl BasicWorkerBuilder {
             Arc, RwLock as StdRwLock,
         };
 
-        use tokio::sync::RwLock;
+        use tokio::sync::OnceCell;
 
         let bootstrap_host = match url::Url::parse(&self.url) {
             Ok(parsed) => parsed.host_str().unwrap_or("localhost").to_string(),
@@ -176,7 +176,16 @@ impl BasicWorkerBuilder {
             default_model_type: ModelType::LLM, // Standard LLM capabilities
         };
 
-        let grpc_client = Arc::new(RwLock::new(self.grpc_client.map(Arc::new)));
+        // Use OnceCell for lock-free gRPC client access after initialization
+        let grpc_client = Arc::new(match self.grpc_client {
+            Some(client) => {
+                let cell = OnceCell::new();
+                // Pre-set the client if provided (blocking set is fine during construction)
+                cell.set(Arc::new(client)).ok();
+                cell
+            }
+            None => OnceCell::new(),
+        });
 
         BasicWorker {
             metadata,

sgl-model-gateway/src/core/worker_registry.rs

@@ -1,8 +1,12 @@
 //! Worker Registry for multi-router support
 //!
 //! Provides centralized registry for workers with model-based indexing
+//!
+//! # Performance Optimizations
+//! The model index uses immutable Arc snapshots instead of RwLock for lock-free reads.
+//! This is critical for high-concurrency scenarios where many requests query the same model.
 
-use std::sync::{Arc, RwLock};
+use std::sync::Arc;
 
 use dashmap::DashMap;
 use uuid::Uuid;
@@ -36,8 +40,10 @@ impl Default for WorkerId {
     }
 }
 
-/// Model index type for O(1) lookups (stores Arc<dyn Worker> directly)
+/// Model index using immutable snapshots for lock-free reads.
-type ModelIndex = Arc<DashMap<String, Arc<RwLock<Vec<Arc<dyn Worker>>>>>>;
+/// Each model maps to an Arc'd slice of workers that can be read without locking.
+/// Updates create new snapshots (copy-on-write semantics).
+type ModelIndex = Arc<DashMap<String, Arc<[Arc<dyn Worker>]>>>;
 
 /// Worker registry with model-based indexing
 #[derive(Debug)]
@@ -45,8 +51,8 @@ pub struct WorkerRegistry {
     /// All workers indexed by ID
     workers: Arc<DashMap<WorkerId, Arc<dyn Worker>>>,
 
-    /// Model index for O(1) lookups (stores Arc<dyn Worker> directly)
+    /// Model index for O(1) lookups using immutable snapshots.
-    /// This replaces the previous dual-index approach for better memory efficiency
+    /// Uses Arc<[T]> instead of Arc<RwLock<Vec<T>>> for lock-free reads.
     model_index: ModelIndex,
 
     /// Workers indexed by worker type
@@ -87,14 +93,18 @@ impl WorkerRegistry {
         self.url_to_id
             .insert(worker.url().to_string(), worker_id.clone());
 
-        // Update model index for O(1) lookups
+        // Update model index for O(1) lookups using copy-on-write
+        // This creates a new immutable snapshot with the added worker
         let model_id = worker.model_id().to_string();
         self.model_index
             .entry(model_id)
-            .or_insert_with(|| Arc::new(RwLock::new(Vec::new())))
+            .and_modify(|existing| {
-            .write()
+                // Create new snapshot with the additional worker
-            .expect("RwLock for model_index is poisoned")
+                let mut new_workers: Vec<Arc<dyn Worker>> = existing.iter().cloned().collect();
-            .push(worker.clone());
+                new_workers.push(worker.clone());
+                *existing = Arc::from(new_workers.into_boxed_slice());
+            })
+            .or_insert_with(|| Arc::from(vec![worker.clone()].into_boxed_slice()));
 
         // Update type index (clone needed for DashMap key ownership)
         self.type_workers
@@ -117,13 +127,16 @@ impl WorkerRegistry {
             // Remove from URL mapping
             self.url_to_id.remove(worker.url());
 
-            // Remove from model index
+            // Remove from model index using copy-on-write
-            if let Some(model_index_entry) = self.model_index.get(worker.model_id()) {
+            // Create new snapshot without the removed worker
             let worker_url = worker.url();
-                model_index_entry
+            if let Some(mut entry) = self.model_index.get_mut(worker.model_id()) {
-                    .write()
+                let new_workers: Vec<Arc<dyn Worker>> = entry
-                    .expect("RwLock for model_index is poisoned")
+                    .iter()
-                    .retain(|w| w.url() != worker_url);
+                    .filter(|w| w.url() != worker_url)
+                    .cloned()
+                    .collect();
+                *entry = Arc::from(new_workers.into_boxed_slice());
             }
 
             // Remove from type index
@@ -165,23 +178,22 @@ impl WorkerRegistry {
         self.url_to_id.get(url).and_then(|id| self.get(&id))
     }
 
-    /// Get all workers for a model (O(1) optimized)
+    /// Empty worker slice constant for returning when no workers found
-    /// Uses the pre-indexed model_index for fast lookups
+    const EMPTY_WORKERS: &'static [Arc<dyn Worker>] = &[];
-    pub fn get_by_model(&self, model_id: &str) -> Vec<Arc<dyn Worker>> {
+
+    /// Get all workers for a model (O(1) optimized, lock-free)
+    /// Returns an Arc to the immutable worker slice - just an atomic refcount bump.
+    /// This is the fastest possible read path with zero contention.
+    pub fn get_by_model(&self, model_id: &str) -> Arc<[Arc<dyn Worker>]> {
         self.model_index
             .get(model_id)
-            .map(|workers| {
+            .map(|workers| Arc::clone(&workers))
-                workers
+            .unwrap_or_else(|| Arc::from(Self::EMPTY_WORKERS))
-                    .read()
-                    .expect("RwLock for model_index is poisoned")
-                    .clone()
-            })
-            .unwrap_or_default()
     }
 
     /// Alias for get_by_model for backwards compatibility
     #[inline]
-    pub fn get_by_model_fast(&self, model_id: &str) -> Vec<Arc<dyn Worker>> {
+    pub fn get_by_model_fast(&self, model_id: &str) -> Arc<[Arc<dyn Worker>]> {
         self.get_by_model(model_id)
     }
 
@@ -266,17 +278,11 @@ impl WorkerRegistry {
             .collect()
     }
 
-    /// Get all model IDs with workers
+    /// Get all model IDs with workers (lock-free)
     pub fn get_models(&self) -> Vec<String> {
         self.model_index
             .iter()
-            .filter(|entry| {
+            .filter(|entry| !entry.value().is_empty())
-                entry
-                    .value()
-                    .read()
-                    .map(|workers| !workers.is_empty())
-                    .unwrap_or(false)
-            })
             .map(|entry| entry.key().clone())
             .collect()
     }
@@ -299,8 +305,8 @@ impl WorkerRegistry {
     ) -> Vec<Arc<dyn Worker>> {
         // Start with the most efficient collection based on filters
         // Use model index when possible as it's O(1) lookup
-        let workers = if let Some(model) = model_id {
+        let workers: Vec<Arc<dyn Worker>> = if let Some(model) = model_id {
-            self.get_by_model_fast(model)
+            self.get_by_model_fast(model).to_vec()
         } else {
             self.get_all()
         };
@@ -340,20 +346,14 @@ impl WorkerRegistry {
             .collect()
     }
 
-    /// Get worker statistics
+    /// Get worker statistics (lock-free)
     pub fn stats(&self) -> WorkerRegistryStats {
         let total_workers = self.workers.len();
-        // Count models directly instead of allocating Vec via get_models()
+        // Count models directly instead of allocating Vec via get_models() (lock-free)
         let total_models = self
             .model_index
             .iter()
-            .filter(|entry| {
+            .filter(|entry| !entry.value().is_empty())
-                entry
-                    .value()
-                    .read()
-                    .map(|workers| !workers.is_empty())
-                    .unwrap_or(false)
-            })
             .count();
 
         let mut healthy_count = 0;

sgl-model-gateway/src/routers/http/pd_router.rs

@@ -705,8 +705,9 @@ impl PDRouter {
         let prefill_workers = if let Some(model) = effective_model_id {
             self.worker_registry
                 .get_by_model_fast(model)
-                .into_iter()
+                .iter()
                 .filter(|w| matches!(w.worker_type(), WorkerType::Prefill { .. }))
+                .cloned()
                 .collect()
         } else {
             self.worker_registry.get_prefill_workers()
@@ -715,8 +716,9 @@ impl PDRouter {
         let decode_workers = if let Some(model) = effective_model_id {
             self.worker_registry
                 .get_by_model_fast(model)
-                .into_iter()
+                .iter()
                 .filter(|w| matches!(w.worker_type(), WorkerType::Decode))
+                .cloned()
                 .collect()
         } else {
             self.worker_registry.get_decode_workers()