[model-gateway] improve lock contention and allocation in middleware (#16405)

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

@@ -3,15 +3,19 @@ use std::{
     time::{Duration, Instant},
 };
 
-use tokio::sync::{Mutex, Notify};
+use parking_lot::Mutex;
+use tokio::sync::Notify;
 use tracing::{debug, trace};
 
-/// Token bucket for rate limiting
+/// Token bucket for rate limiting.
 ///
 /// This implementation provides:
 /// - Smooth rate limiting with configurable refill rate
 /// - Burst capacity handling
-/// - Fair queuing for waiting requests
+/// - Fair queuing for waiting requests via Notify
+/// - Sync token return for Drop handlers (via `return_tokens_sync`)
+///
+/// Uses `parking_lot::Mutex` for sync-compatible locking (no async required).
 #[derive(Clone)]
 pub struct TokenBucket {
     inner: Arc<Mutex<TokenBucketInner>>,
@@ -30,7 +34,7 @@ impl TokenBucket {
     ///
     /// # Arguments
     /// * `capacity` - Maximum number of tokens (burst capacity)
-    /// * `refill_rate` - Tokens added per second
+    /// * `refill_rate` - Tokens added per second (0 for pure concurrency limiting)
     pub fn new(capacity: usize, refill_rate: usize) -> Self {
         let capacity = capacity as f64;
         // Allow refill_rate=0 for pure concurrency limiting (semaphore behavior)
@@ -48,9 +52,16 @@ impl TokenBucket {
         }
     }
 
-    /// Try to acquire tokens immediately
+    /// Try to acquire tokens immediately.
+    ///
+    /// Returns `Ok(())` if tokens were acquired, `Err(())` if insufficient tokens.
     pub async fn try_acquire(&self, tokens: f64) -> Result<(), ()> {
-        let mut inner = self.inner.lock().await;
+        self.try_acquire_sync(tokens)
+    }
+
+    /// Sync version of try_acquire (for internal use).
+    fn try_acquire_sync(&self, tokens: f64) -> Result<(), ()> {
+        let mut inner = self.inner.lock();
 
         let now = Instant::now();
         let elapsed = now.duration_since(inner.last_refill).as_secs_f64();
@@ -77,15 +88,17 @@ impl TokenBucket {
         }
     }
 
-    /// Acquire tokens, waiting if necessary
+    /// Acquire tokens, waiting if necessary.
+    ///
+    /// When `refill_rate=0`, waits indefinitely for tokens to be returned via `return_tokens()`.
+    /// Use `acquire_timeout()` to set an appropriate timeout.
     pub async fn acquire(&self, tokens: f64) -> Result<(), tokio::time::error::Elapsed> {
         if self.try_acquire(tokens).await.is_ok() {
             return Ok(());
         }
 
         // When refill_rate=0 (pure concurrency limiting), tokens only come back
-        // via return_tokens(), so we wait indefinitely on notify signal.
+        // via return_tokens(), so we wait on notify signal only.
-        // The caller should use acquire_timeout() to set an appropriate timeout.
         if self.refill_rate == 0.0 {
             debug!(
                 "Token bucket: waiting indefinitely for {} tokens (refill_rate=0)",
@@ -93,17 +106,17 @@ impl TokenBucket {
             );
 
             loop {
+                // Wait for notify signal from return_tokens()
+                self.notify.notified().await;
+
                 if self.try_acquire(tokens).await.is_ok() {
                     return Ok(());
                 }
-
-                // Wait for notify signal from return_tokens()
-                self.notify.notified().await;
             }
         }
 
         let wait_time = {
-            let inner = self.inner.lock().await;
+            let inner = self.inner.lock();
             let tokens_needed = tokens - inner.tokens;
             let wait_secs = (tokens_needed / self.refill_rate).max(0.0);
             Duration::from_secs_f64(wait_secs)
@@ -119,7 +132,6 @@ impl TokenBucket {
                 if self.try_acquire(tokens).await.is_ok() {
                     return;
                 }
-
                 tokio::select! {
                     _ = self.notify.notified() => {},
                     _ = tokio::time::sleep(Duration::from_millis(10)) => {},
@@ -131,7 +143,7 @@ impl TokenBucket {
         Ok(())
     }
 
-    /// Acquire tokens with custom timeout
+    /// Acquire tokens with custom timeout.
     pub async fn acquire_timeout(
         &self,
         tokens: f64,
@@ -140,20 +152,30 @@ impl TokenBucket {
         tokio::time::timeout(timeout, self.acquire(tokens)).await?
     }
 
-    /// Return tokens to the bucket (for cancelled requests)
+    /// Return tokens to the bucket (sync version).
-    pub async fn return_tokens(&self, tokens: f64) {
+    ///
-        let mut inner = self.inner.lock().await;
+    /// This is safe to call from sync contexts (e.g., Drop handlers).
-        inner.tokens = (inner.tokens + tokens).min(self.capacity);
+    /// Uses `parking_lot::Mutex` which never blocks indefinitely.
+    pub fn return_tokens_sync(&self, tokens: f64) {
+        {
+            let mut inner = self.inner.lock();
+            inner.tokens = (inner.tokens + tokens).min(self.capacity);
+            debug!(
+                "Token bucket: returned {} tokens, {} available",
+                tokens, inner.tokens
+            );
+        } // Release lock before notify
         self.notify.notify_waiters();
-        debug!(
-            "Token bucket: returned {} tokens, {} available",
-            tokens, inner.tokens
-        );
     }
 
-    /// Get current available tokens (for monitoring)
+    /// Return tokens to the bucket (async version for API compatibility).
+    pub async fn return_tokens(&self, tokens: f64) {
+        self.return_tokens_sync(tokens);
+    }
+
+    /// Get current available tokens (for monitoring).
     pub async fn available_tokens(&self) -> f64 {
-        let mut inner = self.inner.lock().await;
+        let mut inner = self.inner.lock();
 
         let now = Instant::now();
         let elapsed = now.duration_since(inner.last_refill).as_secs_f64();
@@ -240,4 +262,18 @@ mod tests {
         let result = bucket.acquire_timeout(1.0, Duration::from_secs(1)).await;
         assert!(result.is_ok());
     }
+
+    #[tokio::test]
+    async fn test_return_tokens_sync() {
+        // Test that sync return works correctly
+        let bucket = TokenBucket::new(2, 0);
+
+        assert!(bucket.try_acquire(1.0).await.is_ok());
+        assert!(bucket.try_acquire(1.0).await.is_ok());
+        assert!(bucket.try_acquire(1.0).await.is_err());
+
+        // Use sync return
+        bucket.return_tokens_sync(1.0);
+        assert!(bucket.try_acquire(1.0).await.is_ok());
+    }
 }

sgl-model-gateway/src/middleware.rs

@@ -69,23 +69,12 @@ impl TokenGuardBody {
 impl Drop for TokenGuardBody {
     fn drop(&mut self) {
         if let Some(bucket) = self.token_bucket.take() {
-            let tokens = self.tokens;
             debug!(
                 "TokenGuardBody: stream ended, returning {} tokens to bucket",
-                tokens
+                self.tokens
             );
-            if let Ok(handle) = tokio::runtime::Handle::try_current() {
+            // Use lock-free sync return - no runtime needed, guaranteed token return
-                handle.spawn(async move {
+            bucket.return_tokens_sync(self.tokens);
-                    bucket.return_tokens(tokens).await;
-                });
-            } else {
-                // Runtime not available (e.g., during shutdown)
-                // Tokens will be lost, but this is acceptable during shutdown
-                warn!(
-                    "TokenGuardBody: Cannot return {} tokens - no Tokio runtime available",
-                    tokens
-                );
-            }
         }
     }
 }
@@ -159,7 +148,7 @@ pub async fn auth_middleware(
 /// Alphanumeric characters for request ID generation (as bytes for O(1) indexing)
 const REQUEST_ID_CHARS: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
 
-/// Generate OpenAI-compatible request ID based on endpoint
+/// Generate OpenAI-compatible request ID based on endpoint.
 fn generate_request_id(path: &str) -> String {
     let prefix = if path.contains("/chat/completions") {
         "chatcmpl-"

sgl-model-gateway/src/observability/inflight_tracker.rs

@@ -60,11 +60,9 @@ impl InFlightRequestTracker {
         let now = Instant::now();
         let inf_idx = AGE_BUCKET_LABELS.len() - 1;
 
-        let instants: Vec<Instant> = self.requests.iter().map(|entry| *entry.value()).collect();
-
         let mut non_cumulative_counts = [0usize; AGE_BUCKET_LABELS.len()];
-        for inst in instants {
+        for entry in self.requests.iter() {
-            let age_secs = now.duration_since(inst).as_secs();
+            let age_secs = now.duration_since(*entry.value()).as_secs();
             let bucket_idx = AGE_BUCKET_BOUNDS
                 .iter()
                 .position(|&bound| age_secs <= bound)

sgl-model-gateway/src/observability/metrics.rs

@@ -31,6 +31,7 @@ static STRING_INTERNER: Lazy<DashMap<String, Arc<str>>> = Lazy::new(DashMap::new
 ///
 /// This function is designed for high-throughput scenarios where the same
 /// strings (model IDs, worker URLs) appear repeatedly. The first call allocates,
+/// subsequent calls just clone the Arc (very cheap - just a ref count increment).
 pub fn intern_string(s: &str) -> Arc<str> {
     // Fast path: check if already interned
     if let Some(entry) = STRING_INTERNER.get(s) {

sgl-model-gateway/src/observability/otel_trace.rs

@@ -29,6 +29,11 @@ use tracing_subscriber::{
 
 use super::events::get_module_path as events_module_path;
 
+/// Whether OpenTelemetry tracing is enabled.
+///
+/// This flag guards access to TRACER and PROVIDER. We use Release/Acquire
+/// ordering to ensure proper synchronization: writes to TRACER/PROVIDER
+/// happen-before the Release store, and Acquire loads happen-before reads.
 static ENABLED: AtomicBool = AtomicBool::new(false);
 static TRACER: OnceLock<SdkTracer> = OnceLock::new();
 static PROVIDER: OnceLock<TracerProvider> = OnceLock::new();
@@ -84,7 +89,8 @@ where
 
 pub fn otel_tracing_init(enable: bool, otlp_endpoint: Option<&str>) -> Result<()> {
     if !enable {
-        ENABLED.store(false, Ordering::Relaxed);
+        // Use Release to ensure any prior OTEL state changes are visible
+        ENABLED.store(false, Ordering::Release);
         return Ok(());
     }
 
@@ -136,7 +142,9 @@ pub fn otel_tracing_init(enable: bool, otlp_endpoint: Option<&str>) -> Result<()
 
     let _ = global::set_tracer_provider(provider);
 
-    ENABLED.store(true, Ordering::Relaxed);
+    // Use Release ordering: all writes to TRACER/PROVIDER happen-before this store,
+    // so any thread that loads ENABLED with Acquire will see the initialized state.
+    ENABLED.store(true, Ordering::Release);
 
     eprintln!("[tracing] OpenTelemetry initialized successfully");
     Ok(())
@@ -163,9 +171,13 @@ where
     Ok(Box::new(layer))
 }
 
+/// Check if OpenTelemetry tracing is enabled.
+///
+/// Uses Acquire ordering to synchronize with the Release store in `otel_tracing_init`,
+/// ensuring that if this returns true, TRACER and PROVIDER are fully initialized.
 #[inline]
 pub fn is_otel_enabled() -> bool {
-    ENABLED.load(Ordering::Relaxed)
+    ENABLED.load(Ordering::Acquire)
 }
 
 pub async fn flush_spans_async() -> Result<()> {
@@ -186,9 +198,11 @@ pub async fn flush_spans_async() -> Result<()> {
 }
 
 pub fn shutdown_otel() {
-    if ENABLED.load(Ordering::Relaxed) {
+    // Use Acquire to ensure we see any prior OTEL operations
+    if ENABLED.load(Ordering::Acquire) {
         global::shutdown_tracer_provider();
-        ENABLED.store(false, Ordering::Relaxed);
+        // Use Release to ensure shutdown completes before flag is cleared
+        ENABLED.store(false, Ordering::Release);
         eprintln!("[tracing] OpenTelemetry shut down");
     }
 }