From 3a72e35ca33fe2e4191ab35a73ad2c2c08470258 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Fri, 21 Aug 2026 13:39:33 -0700 Subject: [PATCH] fix(middleware): enforce the admission queue bound with FIFO waiters --queue-size never bound the number of queued requests: the queue processor drained the size-capped channel and spawned a detached waiter task per request, so every overflow request waited anyway (production read smg_admission_queue_depth=1288 with --queue-size 512), and notify_waiters woke the whole herd to race for each returned token. Replace the channel-plus-processor design with a slot semaphore held for the duration of the wait: at most queue_size requests park, the next arrival sheds 429 admission_queue_full immediately, and a waiter that outlives queue_timeout_secs sheds 503 admission_queue_timeout. Waiters park directly on the token bucket, which now hands returned tokens to the oldest waiter when refill_rate=0: FIFO order, no lost wakeups, no barging past the queue, and cancelled or timed-out waiters return any uncollected grant. Shed bodies, Retry-After, the permit-rides-response-body lifecycle, and the admission metrics are unchanged; smg_admission_queue_depth can no longer exceed queue_size. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- .../benches/wasm_middleware_latency.rs | 2 +- model_gateway/src/middleware/concurrency.rs | 368 ++++++++---------- model_gateway/src/middleware/mod.rs | 4 +- model_gateway/src/middleware/token_bucket.rs | 269 +++++++++++-- model_gateway/src/server.rs | 51 ++- model_gateway/tests/common/test_app.rs | 4 +- model_gateway/tests/wasm_test.rs | 2 +- 7 files changed, 421 insertions(+), 279 deletions(-) diff --git a/model_gateway/benches/wasm_middleware_latency.rs b/model_gateway/benches/wasm_middleware_latency.rs index feb11f9e21..154eecf549 100644 --- a/model_gateway/benches/wasm_middleware_latency.rs +++ b/model_gateway/benches/wasm_middleware_latency.rs @@ -78,7 +78,7 @@ fn bench_wasm_middleware_buffering(c: &mut Criterion) { router: Arc::new(MockRouter), probe_state: ProbeState::new(context.inflight_tracker.clone()), context: Arc::new(context), - concurrency_queue_tx: None, + admission_queue: None, router_manager: None, mesh_handler: None, mesh_adapters: None, diff --git a/model_gateway/src/middleware/concurrency.rs b/model_gateway/src/middleware/concurrency.rs index d71f0daa82..c220974ca3 100644 --- a/model_gateway/src/middleware/concurrency.rs +++ b/model_gateway/src/middleware/concurrency.rs @@ -1,17 +1,17 @@ //! Per-request concurrency limiting via a token bucket, with optional //! queuing for backpressure. //! -//! `ConcurrencyLimiter` wires a bounded `mpsc` channel that -//! `concurrency_limit_middleware` uses to enqueue requests when the -//! bucket is empty; `QueueProcessor` drains that channel and hands tokens -//! back to waiters. `TokenGuardBody` wraps the response body so the token -//! is only released after the entire stream has been delivered. +//! `AdmissionQueue` caps concurrent waiters at `queue_size`; a request +//! that finds every slot taken sheds immediately. Waiters park directly +//! on the token bucket, which serves them in FIFO order. `TokenGuardBody` +//! wraps the response body so the token is only released after the entire +//! stream has been delivered. use std::{ pin::Pin, sync::Arc, task::{Context, Poll}, - time::{Duration, Instant}, + time::Duration, }; use axum::{ @@ -19,15 +19,12 @@ use axum::{ extract::{Request, State}, http::{header::RETRY_AFTER, HeaderValue, StatusCode}, middleware::Next, - response::{IntoResponse, Response}, + response::Response, }; use bytes::Bytes; use http_body::Frame; -use tokio::{ - sync::{mpsc, oneshot}, - time::error::Elapsed, -}; -use tracing::{debug, error, warn}; +use tokio::{sync::Semaphore, time::error::Elapsed}; +use tracing::{debug, warn}; use super::{token_bucket::TokenBucket, SHED_RETRY_AFTER_SECS}; use crate::{ @@ -149,108 +146,21 @@ async fn run_with_permit(next: Next, request: Request, permit: TokenPermit Response::from_parts(parts, Body::new(body)) } -/// Request queue entry -pub struct QueuedRequest { - /// Time when the request was queued - queued_at: Instant, - /// Channel to send the permit back when acquired - permit_tx: oneshot::Sender>, -} - -/// Queue processor that handles queued requests -pub struct QueueProcessor { - token_bucket: Arc, - queue_rx: mpsc::Receiver, +/// Bounded admission queue: at most `queue_size` requests wait for a token; +/// the next arrival sheds immediately. FIFO order comes from the token +/// bucket's waiter queue. +pub struct AdmissionQueue { + slots: Semaphore, queue_timeout: Duration, } -impl QueueProcessor { - pub fn new( - token_bucket: Arc, - queue_rx: mpsc::Receiver, - queue_timeout: Duration, - ) -> Self { +impl AdmissionQueue { + pub fn new(queue_size: usize, queue_timeout: Duration) -> Self { Self { - token_bucket, - queue_rx, + slots: Semaphore::new(queue_size), queue_timeout, } } - - pub async fn run(mut self) { - debug!("Starting concurrency queue processor"); - - // Process requests in a single task to reduce overhead - while let Some(queued) = self.queue_rx.recv().await { - // Check timeout immediately - let elapsed = queued.queued_at.elapsed(); - if elapsed >= self.queue_timeout { - warn!("Request already timed out in queue"); - let _ = queued.permit_tx.send(Err(StatusCode::SERVICE_UNAVAILABLE)); - continue; - } - - let remaining_timeout = self.queue_timeout - elapsed; - - // Try to acquire token for this request - if let Ok(permit) = TokenPermit::try_acquire(self.token_bucket.clone(), 1.0) { - // Got token immediately - debug!("Queue: acquired token immediately for queued request"); - let _ = queued.permit_tx.send(Ok(permit)); - } else { - // Need to wait for token - let token_bucket = self.token_bucket.clone(); - - // Spawn task only when we actually need to wait - #[expect( - clippy::disallowed_methods, - reason = "fire-and-forget permit acquisition: task is bounded by remaining_timeout and communicates via oneshot; dropping the JoinHandle detaches the task but it self-terminates" - )] - tokio::spawn(async move { - if let Ok(permit) = - TokenPermit::acquire_timeout(token_bucket, 1.0, remaining_timeout).await - { - debug!("Queue: acquired token after waiting"); - let _ = queued.permit_tx.send(Ok(permit)); - } else { - warn!("Queue: request timed out waiting for token"); - let _ = queued.permit_tx.send(Err(StatusCode::SERVICE_UNAVAILABLE)); - } - }); - } - } - - warn!("Concurrency queue processor shutting down"); - } -} - -/// State for the concurrency limiter -pub struct ConcurrencyLimiter { - pub queue_tx: Option>, -} - -impl ConcurrencyLimiter { - /// Create new concurrency limiter with optional queue - pub fn new( - token_bucket: Option>, - queue_size: usize, - queue_timeout: Duration, - ) -> (Self, Option) { - match (token_bucket, queue_size) { - (None, _) => (Self { queue_tx: None }, None), - (Some(bucket), size) if size > 0 => { - let (queue_tx, queue_rx) = mpsc::channel(size); - let processor = QueueProcessor::new(bucket, queue_rx, queue_timeout); - ( - Self { - queue_tx: Some(queue_tx), - }, - Some(processor), - ) - } - (Some(_), _) => (Self { queue_tx: None }, None), - } - } } /// Middleware function for concurrency limiting with optional queuing @@ -277,71 +187,52 @@ pub async fn concurrency_limit_middleware( if let Ok(permit) = TokenPermit::try_acquire(token_bucket.clone(), 1.0) { debug!("Acquired token immediately"); Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_ALLOWED); - run_with_permit(next, request, permit).await - } else { - // No tokens available, try to queue if enabled - if let Some(queue_tx) = &app_state.concurrency_queue_tx { - debug!("No tokens available, attempting to queue request"); - - // Create a channel for the token response - let (permit_tx, permit_rx) = oneshot::channel(); - - let queued = QueuedRequest { - queued_at: Instant::now(), - permit_tx, - }; - - // Try to send to queue - match queue_tx.try_send(queued) { - Ok(()) => { - // Wait for token from queue processor - let permit_result = { - let _queued = QueueDepthGuard::enter(); - permit_rx.await - }; - match permit_result { - Ok(Ok(permit)) => { - debug!("Acquired token from queue"); - Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_ALLOWED); - run_with_permit(next, request, permit).await - } - Ok(Err(status)) => { - warn!("Queue returned error status: {}", status); - Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_REJECTED); - Metrics::record_admission_rejected( - metrics_labels::ADMISSION_REJECTED_TIMEOUT, - ); - shed_response( - status, - "admission_queue_timeout", - "timed out waiting for an admission slot", - ) - } - Err(_) => { - error!("Queue response channel closed"); - Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_REJECTED); - StatusCode::INTERNAL_SERVER_ERROR.into_response() - } - } - } - Err(_) => { - warn!("Request queue is full, returning 429"); - Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_REJECTED); - Metrics::record_admission_rejected(metrics_labels::ADMISSION_REJECTED_FULL); - shed_response( - StatusCode::TOO_MANY_REQUESTS, - "admission_queue_full", - "admission queue is at capacity", - ) - } - } - } else { - warn!("No tokens available and queuing is disabled, returning 429"); + return run_with_permit(next, request, permit).await; + } + + let Some(queue) = &app_state.admission_queue else { + warn!("No tokens available and queuing is disabled, returning 429"); + Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_REJECTED); + return shed_response( + StatusCode::TOO_MANY_REQUESTS, + "admission_queue_full", + "concurrency limit reached and queuing is disabled", + ); + }; + + let Ok(slot) = queue.slots.try_acquire() else { + warn!("Request queue is full, returning 429"); + Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_REJECTED); + Metrics::record_admission_rejected(metrics_labels::ADMISSION_REJECTED_FULL); + return shed_response( + StatusCode::TOO_MANY_REQUESTS, + "admission_queue_full", + "admission queue is at capacity", + ); + }; + + debug!("No tokens available, waiting in admission queue"); + let permit_result = { + let _queued = QueueDepthGuard::enter(); + TokenPermit::acquire_timeout(token_bucket, 1.0, queue.queue_timeout).await + }; + match permit_result { + Ok(permit) => { + // Free the queue slot before running the request: admitted + // requests are bounded by the bucket, not the queue. + drop(slot); + debug!("Acquired token from queue"); + Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_ALLOWED); + run_with_permit(next, request, permit).await + } + Err(_) => { + warn!("Request timed out in admission queue"); Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_REJECTED); + Metrics::record_admission_rejected(metrics_labels::ADMISSION_REJECTED_TIMEOUT); shed_response( - StatusCode::TOO_MANY_REQUESTS, - "admission_queue_full", - "concurrency limit reached and queuing is disabled", + StatusCode::SERVICE_UNAVAILABLE, + "admission_queue_timeout", + "timed out waiting for an admission slot", ) } } @@ -361,7 +252,7 @@ mod tests { use smg_data_connector::{ MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage, }; - use tokio::sync::Notify; + use tokio::sync::{mpsc, oneshot, Notify}; use tokio_stream::wrappers::ReceiverStream; use tower::ServiceExt; @@ -373,7 +264,7 @@ mod tests { fn test_app_state( bucket: Arc, - queue_tx: Option>, + admission_queue: Option>, ) -> Arc { let router_config = RouterConfig::default(); let context = Arc::new( @@ -403,7 +294,7 @@ mod tests { )), probe_state: ProbeState::new(context.inflight_tracker.clone()), context, - concurrency_queue_tx: queue_tx, + admission_queue, router_manager: None, mesh_handler: None, mesh_adapters: None, @@ -471,6 +362,18 @@ mod tests { ); } + async fn wait_for_slots(queue: &AdmissionQueue, expected: usize) { + while queue.slots.available_permits() != expected { + tokio::task::yield_now().await; + } + } + + async fn wait_for_bucket_waiters(bucket: &TokenBucket, expected: usize) { + while bucket.waiter_count() != expected { + tokio::task::yield_now().await; + } + } + /// Sets `polled` on the first poll, so a test can detect any body /// collection that happens while the request is parked at admission. struct PollRecordingBody { @@ -629,26 +532,87 @@ mod tests { assert_eq!(bucket.available_tokens(), 1.0); } + /// With the cap saturated, exactly `queue_size` requests park; the next + /// arrival sheds 429 immediately, and the parked ones admit FIFO once + /// capacity frees. #[tokio::test] - async fn cancelled_queued_request_returns_acquired_token() { + #[expect( + clippy::disallowed_methods, + reason = "test parks requests in background tasks while asserting the bound" + )] + async fn queue_admits_exactly_queue_size_waiters_then_sheds() { let bucket = Arc::new(TokenBucket::new(1, 0)); - let (queue_tx, queue_rx) = mpsc::channel(1); - let (permit_tx, permit_rx) = oneshot::channel(); - drop(permit_rx); - - assert!(queue_tx - .send(QueuedRequest { - queued_at: Instant::now(), - permit_tx, - }) + let queue = Arc::new(AdmissionQueue::new(2, Duration::from_secs(5))); + let app = echo_app(test_app_state(bucket.clone(), Some(queue.clone()))); + + let held = TokenPermit::try_acquire(bucket.clone(), 1.0).unwrap(); + + let first = tokio::spawn(app.clone().oneshot(echo_request(Body::from("first")))); + wait_for_slots(&queue, 1).await; + wait_for_bucket_waiters(&bucket, 1).await; + let second = tokio::spawn(app.clone().oneshot(echo_request(Body::from("second")))); + wait_for_slots(&queue, 0).await; + wait_for_bucket_waiters(&bucket, 2).await; + + let shed = app + .clone() + .oneshot(echo_request(Body::from("overflow"))) + .await + .unwrap(); + assert_eq!(shed.status(), StatusCode::TOO_MANY_REQUESTS); + assert_eq!( + shed.headers().get(RETRY_AFTER).unwrap(), + &HeaderValue::from(SHED_RETRY_AFTER_SECS) + ); + assert!(!first.is_finished()); + assert!(!second.is_finished()); + + drop(held); + let first_response = first.await.unwrap().unwrap(); + assert_eq!(first_response.status(), StatusCode::OK); + assert!( + !second.is_finished(), + "second waiter must stay parked until the first releases" + ); + let echoed = axum::body::to_bytes(first_response.into_body(), usize::MAX) .await - .is_ok()); - drop(queue_tx); + .unwrap(); + assert_eq!(&echoed[..], b"first"); - QueueProcessor::new(bucket.clone(), queue_rx, Duration::from_secs(1)) - .run() - .await; - assert_eq!(bucket.available_tokens(), 1.0); + let second_response = second.await.unwrap().unwrap(); + assert_eq!(second_response.status(), StatusCode::OK); + let echoed = axum::body::to_bytes(second_response.into_body(), usize::MAX) + .await + .unwrap(); + assert_eq!(&echoed[..], b"second"); + + wait_for_slots(&queue, 2).await; + } + + /// A parked request that is cancelled (client disconnect) frees its + /// queue slot for the next arrival. + #[tokio::test] + #[expect( + clippy::disallowed_methods, + reason = "test parks a request in a background task and aborts it" + )] + async fn cancelled_parked_request_frees_queue_slot() { + let bucket = Arc::new(TokenBucket::new(1, 0)); + let queue = Arc::new(AdmissionQueue::new(1, Duration::from_secs(5))); + let app = echo_app(test_app_state(bucket.clone(), Some(queue.clone()))); + + let held = TokenPermit::try_acquire(bucket, 1.0).unwrap(); + + let parked = tokio::spawn(app.clone().oneshot(echo_request(Body::from("parked")))); + wait_for_slots(&queue, 0).await; + + parked.abort(); + let _ = parked.await; + wait_for_slots(&queue, 1).await; + + drop(held); + let admitted = app.oneshot(echo_request(Body::from("next"))).await.unwrap(); + assert_eq!(admitted.status(), StatusCode::OK); } /// A request parked in the admission queue must keep its body unpolled; @@ -660,9 +624,7 @@ mod tests { )] async fn queued_request_body_stays_unpolled_until_admitted() { let bucket = Arc::new(TokenBucket::new(1, 0)); - let (limiter, processor) = - ConcurrencyLimiter::new(Some(bucket.clone()), 4, Duration::from_secs(5)); - tokio::spawn(processor.unwrap().run()); + let queue = Arc::new(AdmissionQueue::new(4, Duration::from_secs(5))); let hold_gate = Arc::new(Notify::new()); let entered = Arc::new(Notify::new()); @@ -683,7 +645,7 @@ mod tests { .route("/hold", post(hold)) .route("/echo", post(|body: Bytes| async move { body })) .layer(axum::middleware::from_fn_with_state( - test_app_state(bucket, limiter.queue_tx), + test_app_state(bucket, Some(queue)), concurrency_limit_middleware, )); @@ -748,16 +710,12 @@ mod tests { .unwrap(); rt.block_on(async { let bucket = Arc::new(TokenBucket::new(1, 0)); - let (limiter, processor) = - ConcurrencyLimiter::new(Some(bucket.clone()), 1, Duration::from_secs(5)); - // Keep the queue channel open but undrained: the first - // request parks and the second finds the queue full. - let _processor = processor.unwrap(); - let app = echo_app(test_app_state(bucket.clone(), limiter.queue_tx)); + let queue = Arc::new(AdmissionQueue::new(1, Duration::from_secs(5))); + let app = echo_app(test_app_state(bucket.clone(), Some(queue.clone()))); let held = TokenPermit::try_acquire(bucket, 1.0).unwrap(); let parked = tokio::spawn(app.clone().oneshot(echo_request(Body::from("parked")))); - tokio::time::sleep(Duration::from_millis(50)).await; + wait_for_slots(&queue, 0).await; assert!(!parked.is_finished()); let rejected = app @@ -794,10 +752,6 @@ mod tests { /// A request that outlives the queue timeout sheds as 503 + Retry-After /// with reason="timeout" and releases its queue-depth slot. #[test] - #[expect( - clippy::disallowed_methods, - reason = "queue processor runs as a background task" - )] fn admission_metrics_record_timeout_rejections() { let recorder = PrometheusBuilder::new().build_recorder(); let handle = recorder.handle(); @@ -809,10 +763,8 @@ mod tests { .unwrap(); rt.block_on(async { let bucket = Arc::new(TokenBucket::new(1, 0)); - let (limiter, processor) = - ConcurrencyLimiter::new(Some(bucket.clone()), 4, Duration::from_millis(100)); - tokio::spawn(processor.unwrap().run()); - let app = echo_app(test_app_state(bucket.clone(), limiter.queue_tx)); + let queue = Arc::new(AdmissionQueue::new(4, Duration::from_millis(100))); + let app = echo_app(test_app_state(bucket.clone(), Some(queue))); let held = TokenPermit::try_acquire(bucket, 1.0).unwrap(); let response = app.oneshot(echo_request(Body::from("late"))).await.unwrap(); diff --git a/model_gateway/src/middleware/mod.rs b/model_gateway/src/middleware/mod.rs index fa20f0972d..6df4113438 100644 --- a/model_gateway/src/middleware/mod.rs +++ b/model_gateway/src/middleware/mod.rs @@ -16,9 +16,7 @@ pub mod token_bucket; pub mod wasm; pub use auth::{auth_middleware, deny_all_middleware, AuthConfig}; -pub use concurrency::{ - concurrency_limit_middleware, ConcurrencyLimiter, QueueProcessor, QueuedRequest, TokenGuardBody, -}; +pub use concurrency::{concurrency_limit_middleware, AdmissionQueue, TokenGuardBody}; pub use logging::{ create_logging_layer, ProbeResponse, RequestLogger, RequestSpan, ResponseLogger, StreamFailureLogger, diff --git a/model_gateway/src/middleware/token_bucket.rs b/model_gateway/src/middleware/token_bucket.rs index 289a545065..efc4d3387b 100644 --- a/model_gateway/src/middleware/token_bucket.rs +++ b/model_gateway/src/middleware/token_bucket.rs @@ -1,4 +1,5 @@ use std::{ + collections::{HashSet, VecDeque}, sync::Arc, time::{Duration, Instant}, }; @@ -12,7 +13,10 @@ use tracing::{debug, trace}; /// This implementation provides: /// - Smooth rate limiting with configurable refill rate /// - Burst capacity handling -/// - Fair queuing for waiting requests via Notify +/// - FIFO waiter handoff when `refill_rate=0` (pure concurrency limiting): +/// returned tokens are granted directly to the oldest waiter, so waiters +/// are served in arrival order, wakeups cannot be lost, and new arrivals +/// cannot barge past the queue /// - Sync token return for Drop handlers (via `return_tokens_sync`) /// /// Uses `parking_lot::Mutex` for sync-compatible locking (no async required). @@ -27,6 +31,41 @@ pub struct TokenBucket { struct TokenBucketInner { tokens: f64, last_refill: Instant, + next_waiter_id: u64, + /// FIFO waiters (refill_rate=0 acquire path only). + waiters: VecDeque, + /// Granted-but-uncollected waiter ids; their tokens are already + /// deducted from the pool. + granted: HashSet, +} + +struct Waiter { + id: u64, + tokens: f64, + notify: Arc, +} + +/// Removes a cancelled waiter; returns an uncollected grant to the pool. +struct FifoWaiterGuard<'a> { + bucket: &'a TokenBucket, + id: u64, + tokens: f64, + armed: bool, +} + +impl Drop for FifoWaiterGuard<'_> { + fn drop(&mut self) { + if !self.armed { + return; + } + let mut inner = self.bucket.inner.lock(); + if inner.granted.remove(&self.id) { + inner.tokens = (inner.tokens + self.tokens).min(self.bucket.capacity); + TokenBucket::grant_waiters_locked(&mut inner); + } else if let Some(pos) = inner.waiters.iter().position(|w| w.id == self.id) { + inner.waiters.remove(pos); + } + } } impl TokenBucket { @@ -45,6 +84,9 @@ impl TokenBucket { inner: Arc::new(Mutex::new(TokenBucketInner { tokens: capacity, last_refill: Instant::now(), + next_waiter_id: 0, + waiters: VecDeque::new(), + granted: HashSet::new(), })), notify: Arc::new(Notify::new()), capacity, @@ -52,9 +94,33 @@ impl TokenBucket { } } + fn refill_locked(&self, inner: &mut TokenBucketInner) { + let now = Instant::now(); + let elapsed = now.duration_since(inner.last_refill).as_secs_f64(); + inner.tokens = (inner.tokens + elapsed * self.refill_rate).min(self.capacity); + inner.last_refill = now; + } + + /// Hand tokens to FIFO waiters, oldest first. + fn grant_waiters_locked(inner: &mut TokenBucketInner) { + loop { + match inner.waiters.front() { + Some(front) if inner.tokens >= front.tokens => {} + _ => break, + } + let Some(waiter) = inner.waiters.pop_front() else { + break; + }; + inner.tokens -= waiter.tokens; + inner.granted.insert(waiter.id); + waiter.notify.notify_one(); + } + } + /// Try to acquire tokens immediately. /// - /// Returns `Ok(())` if tokens were acquired, `Err(())` if insufficient tokens. + /// Returns `Ok(())` if tokens were acquired, `Err(())` if insufficient tokens + /// or FIFO waiters are queued ahead. #[expect( clippy::result_unit_err, reason = "Try-acquire pattern: callers only need success/failure, a custom error type adds no information" @@ -71,12 +137,8 @@ impl TokenBucket { ); let mut inner = self.inner.lock(); - let now = Instant::now(); - let elapsed = now.duration_since(inner.last_refill).as_secs_f64(); - let refill_amount = elapsed * self.refill_rate; - - inner.tokens = (inner.tokens + refill_amount).min(self.capacity); - inner.last_refill = now; + self.refill_locked(&mut inner); + Self::grant_waiters_locked(&mut inner); trace!( "Token bucket: {} tokens available, requesting {}", @@ -84,7 +146,7 @@ impl TokenBucket { tokens ); - if inner.tokens >= tokens { + if inner.waiters.is_empty() && inner.tokens >= tokens { inner.tokens -= tokens; debug!( "Token bucket: acquired {} tokens, {} remaining", @@ -98,28 +160,19 @@ impl TokenBucket { /// 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. + /// When `refill_rate=0`, waits in FIFO order (indefinitely) for tokens to + /// be returned via `return_tokens()`. Use `acquire_timeout()` to set an + /// appropriate timeout; a timed-out or cancelled waiter leaves the queue + /// and returns any uncollected grant. pub async fn acquire(&self, tokens: f64) -> Result<(), Elapsed> { - if self.try_acquire(tokens).is_ok() { - return Ok(()); - } - // When refill_rate=0 (pure concurrency limiting), tokens only come back - // via return_tokens(), so we wait on notify signal only. + // via return_tokens(), which hands them to the oldest waiter directly. if self.refill_rate == 0.0 { - debug!( - "Token bucket: waiting indefinitely for {} tokens (refill_rate=0)", - tokens - ); - - loop { - self.notify.notified().await; + return self.acquire_fifo(tokens).await; + } - if self.try_acquire(tokens).is_ok() { - return Ok(()); - } - } + if self.try_acquire(tokens).is_ok() { + return Ok(()); } let wait_time = { @@ -150,6 +203,49 @@ impl TokenBucket { Ok(()) } + async fn acquire_fifo(&self, tokens: f64) -> Result<(), Elapsed> { + let (id, notify) = { + let mut inner = self.inner.lock(); + if inner.waiters.is_empty() && inner.tokens >= tokens { + inner.tokens -= tokens; + return Ok(()); + } + let id = inner.next_waiter_id; + inner.next_waiter_id += 1; + let notify = Arc::new(Notify::new()); + inner.waiters.push_back(Waiter { + id, + tokens, + notify: Arc::clone(¬ify), + }); + (id, notify) + }; + + debug!( + "Token bucket: waiting in FIFO queue for {} tokens (refill_rate=0)", + tokens + ); + + let mut guard = FifoWaiterGuard { + bucket: self, + id, + tokens, + armed: true, + }; + loop { + { + let mut inner = self.inner.lock(); + if inner.granted.remove(&id) { + guard.armed = false; + return Ok(()); + } + } + // notify_one on grant stores a permit, so a grant between the + // check above and this await still wakes us. + notify.notified().await; + } + } + /// Acquire tokens with custom timeout. pub async fn acquire_timeout(&self, tokens: f64, timeout: Duration) -> Result<(), Elapsed> { tokio::time::timeout(timeout, self.acquire(tokens)).await? @@ -167,6 +263,7 @@ impl TokenBucket { { let mut inner = self.inner.lock(); inner.tokens = (inner.tokens + tokens).min(self.capacity); + Self::grant_waiters_locked(&mut inner); debug!( "Token bucket: returned {} tokens, {} available", tokens, inner.tokens @@ -183,16 +280,14 @@ impl TokenBucket { /// Get current available tokens (for monitoring). pub fn available_tokens(&self) -> f64 { let mut inner = self.inner.lock(); - - let now = Instant::now(); - let elapsed = now.duration_since(inner.last_refill).as_secs_f64(); - let refill_amount = elapsed * self.refill_rate; - - inner.tokens = (inner.tokens + refill_amount).min(self.capacity); - inner.last_refill = now; - + self.refill_locked(&mut inner); inner.tokens } + + #[cfg(test)] + pub(crate) fn waiter_count(&self) -> usize { + self.inner.lock().waiters.len() + } } #[cfg(test)] @@ -286,4 +381,108 @@ mod tests { bucket.return_tokens_sync(1.0); assert!(bucket.try_acquire(1.0).is_ok()); } + + async fn registered_waiter( + bucket: &Arc, + expected_waiters: usize, + ) -> tokio::task::JoinHandle> { + let waiter_bucket = bucket.clone(); + #[expect( + clippy::disallowed_methods, + reason = "Test helper: waiter tasks are joined or aborted before the test ends" + )] + let handle = tokio::spawn(async move { + waiter_bucket + .acquire_timeout(1.0, Duration::from_secs(5)) + .await + }); + while bucket.waiter_count() < expected_waiters { + tokio::task::yield_now().await; + } + handle + } + + #[tokio::test] + async fn test_fifo_waiters_served_in_arrival_order() { + let bucket = Arc::new(TokenBucket::new(1, 0)); + assert!(bucket.try_acquire(1.0).is_ok()); + + let first = registered_waiter(&bucket, 1).await; + let second = registered_waiter(&bucket, 2).await; + let third = registered_waiter(&bucket, 3).await; + + bucket.return_tokens(1.0); + first.await.expect("first waiter").expect("first grant"); + assert!(!second.is_finished()); + assert!(!third.is_finished()); + + bucket.return_tokens(1.0); + second.await.expect("second waiter").expect("second grant"); + assert!(!third.is_finished()); + + bucket.return_tokens(1.0); + third.await.expect("third waiter").expect("third grant"); + assert_eq!(bucket.available_tokens(), 0.0); + } + + #[tokio::test] + async fn test_try_acquire_cannot_barge_past_waiters() { + let bucket = Arc::new(TokenBucket::new(1, 0)); + assert!(bucket.try_acquire(1.0).is_ok()); + + let waiter = registered_waiter(&bucket, 1).await; + + bucket.return_tokens(1.0); + // The returned token is already granted to the waiter. + assert!(bucket.try_acquire(1.0).is_err()); + waiter.await.expect("waiter").expect("grant"); + } + + #[tokio::test] + async fn test_cancelled_waiter_leaves_queue_and_returns_grant() { + let bucket = Arc::new(TokenBucket::new(1, 0)); + assert!(bucket.try_acquire(1.0).is_ok()); + + let cancelled = registered_waiter(&bucket, 1).await; + let survivor = registered_waiter(&bucket, 2).await; + + cancelled.abort(); + let _ = cancelled.await; + assert_eq!(bucket.waiter_count(), 1); + + // The survivor is now first in line. + bucket.return_tokens(1.0); + survivor.await.expect("survivor").expect("grant"); + } + + #[tokio::test] + async fn test_timed_out_waiter_does_not_leak_tokens() { + let bucket = Arc::new(TokenBucket::new(1, 0)); + assert!(bucket.try_acquire(1.0).is_ok()); + + let result = bucket.acquire_timeout(1.0, Duration::from_millis(20)).await; + assert!(result.is_err()); + assert_eq!(bucket.waiter_count(), 0); + + bucket.return_tokens(1.0); + assert_eq!(bucket.available_tokens(), 1.0); + } + + #[tokio::test] + async fn test_no_lost_wakeup_under_tight_release_acquire_loop() { + let bucket = Arc::new(TokenBucket::new(1, 0)); + + for _ in 0..200 { + assert!(bucket.try_acquire(1.0).is_ok()); + let waiter = registered_waiter(&bucket, 1).await; + bucket.return_tokens(1.0); + waiter + .await + .expect("waiter task") + .expect("waiter must be granted, not time out"); + // Release the grant the waiter still holds. + bucket.return_tokens(1.0); + } + assert_eq!(bucket.available_tokens(), 1.0); + } } diff --git a/model_gateway/src/server.rs b/model_gateway/src/server.rs index a4fce0dad9..eec2417e7f 100644 --- a/model_gateway/src/server.rs +++ b/model_gateway/src/server.rs @@ -40,7 +40,7 @@ use rustls::crypto::ring; use serde::Deserialize; use serde_json::Value; use smg_mesh::{MeshServerBuilder, MeshServerConfig, MeshServerHandler}; -use tokio::{signal, spawn, sync::mpsc}; +use tokio::{signal, spawn}; use tracing::{debug, error, info, warn, Level}; use wfaas::LoggingSubscriber; @@ -48,7 +48,7 @@ use crate::{ app_context::AppContext, config::RouterConfig, mesh::MeshAdapters, - middleware::{self, AuthConfig, QueuedRequest}, + middleware::{self, AdmissionQueue, AuthConfig}, observability::{ logging::{self, LoggingConfig}, metrics::{self, PrometheusConfig}, @@ -77,7 +77,7 @@ use crate::{ pub struct AppState { pub router: Arc, pub context: Arc, - pub concurrency_queue_tx: Option>, + pub admission_queue: Option>, pub router_manager: Option>, pub mesh_handler: Option>, pub mesh_adapters: Option>, @@ -1295,34 +1295,27 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box { - #[expect( - clippy::disallowed_methods, - reason = "request queue processor runs for the lifetime of the server" - )] - spawn(proc.run()); + let admission_queue = + if app_context.rate_limiter.is_some() && config.router_config.queue_size > 0 { debug!( - "Started request queue (size: {}, timeout: {}s)", + "Admission queue enabled (size: {}, timeout: {}s)", config.router_config.queue_size, config.router_config.queue_timeout_secs ); - } - None => { - debug!( - "Rate limiting enabled (max_concurrent_requests = {}, queue disabled)", - config.router_config.max_concurrent_requests - ); - } + Some(Arc::new(AdmissionQueue::new( + config.router_config.queue_size, + Duration::from_secs(config.router_config.queue_timeout_secs), + ))) + } else { + None + }; + + if app_context.rate_limiter.is_none() { + info!("Rate limiting is disabled (max_concurrent_requests = -1)"); + } else if admission_queue.is_none() { + debug!( + "Rate limiting enabled (max_concurrent_requests = {}, queue disabled)", + config.router_config.max_concurrent_requests + ); } // Get mesh cluster state and port before moving mesh_handler into app_state @@ -1358,7 +1351,7 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box (axum::Router, Arc, TempDir) router, probe_state: ProbeState::new(app_context.inflight_tracker.clone()), context: app_context.clone(), - concurrency_queue_tx: None, + admission_queue: None, router_manager: None, mesh_handler: None, mesh_adapters: None,