chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Simo Lin
2026-01-04 10:40:07 -08:00
committed by GitHub
parent cf6800f6cc
commit 4436dc0f6c
5 changed files with 87 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

86
sgl-model-gateway/src/core/token_bucket.rs
View File

@@ -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.
// 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());
}
}

19
sgl-model-gateway/src/middleware.rs
View File

@@ -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-"

6
sgl-model-gateway/src/observability/inflight_tracker.rs
View File

@@ -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 {
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)

1
sgl-model-gateway/src/observability/metrics.rs
View File

@@ -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) {

24
sgl-model-gateway/src/observability/otel_trace.rs
View File

@@ -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");
}
}