diff --git a/sgl-model-gateway/src/core/token_bucket.rs b/sgl-model-gateway/src/core/token_bucket.rs index d3efa120e..d0e291800 100644 --- a/sgl-model-gateway/src/core/token_bucket.rs +++ b/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>, @@ -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. - // The caller should use acquire_timeout() to set an appropriate timeout. + // via return_tokens(), so we wait on notify signal only. 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) - pub async fn return_tokens(&self, tokens: f64) { - let mut inner = self.inner.lock().await; - inner.tokens = (inner.tokens + tokens).min(self.capacity); + /// Return tokens to the bucket (sync version). + /// + /// This is safe to call from sync contexts (e.g., Drop handlers). + /// 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()); + } } diff --git a/sgl-model-gateway/src/middleware.rs b/sgl-model-gateway/src/middleware.rs index 84d51b321..cc4ee546c 100644 --- a/sgl-model-gateway/src/middleware.rs +++ b/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() { - handle.spawn(async move { - 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 - ); - } + // Use lock-free sync return - no runtime needed, guaranteed token return + bucket.return_tokens_sync(self.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-" diff --git a/sgl-model-gateway/src/observability/inflight_tracker.rs b/sgl-model-gateway/src/observability/inflight_tracker.rs index 19a0639f5..25115ab78 100644 --- a/sgl-model-gateway/src/observability/inflight_tracker.rs +++ b/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 = self.requests.iter().map(|entry| *entry.value()).collect(); - let mut non_cumulative_counts = [0usize; AGE_BUCKET_LABELS.len()]; - for inst in instants { - let age_secs = now.duration_since(inst).as_secs(); + for entry in self.requests.iter() { + let age_secs = now.duration_since(*entry.value()).as_secs(); let bucket_idx = AGE_BUCKET_BOUNDS .iter() .position(|&bound| age_secs <= bound) diff --git a/sgl-model-gateway/src/observability/metrics.rs b/sgl-model-gateway/src/observability/metrics.rs index 3cabf199c..2ec8fdb4e 100644 --- a/sgl-model-gateway/src/observability/metrics.rs +++ b/sgl-model-gateway/src/observability/metrics.rs @@ -31,6 +31,7 @@ static STRING_INTERNER: Lazy>> = 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 { // Fast path: check if already interned if let Some(entry) = STRING_INTERNER.get(s) { diff --git a/sgl-model-gateway/src/observability/otel_trace.rs b/sgl-model-gateway/src/observability/otel_trace.rs index 12e6b2cf3..76542c6c0 100644 --- a/sgl-model-gateway/src/observability/otel_trace.rs +++ b/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 = OnceLock::new(); static PROVIDER: OnceLock = 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"); } }