[model-gateway] allow refill rate to be zero (#14030)
This commit is contained in:
@@ -261,11 +261,12 @@ impl ConfigValidator {
|
|||||||
}
|
}
|
||||||
|
|
||||||
if let Some(tokens_per_second) = config.rate_limit_tokens_per_second {
|
if let Some(tokens_per_second) = config.rate_limit_tokens_per_second {
|
||||||
if tokens_per_second <= 0 {
|
// Allow 0 for pure concurrency limiting (semaphore behavior)
|
||||||
|
if tokens_per_second < 0 {
|
||||||
return Err(ConfigError::InvalidValue {
|
return Err(ConfigError::InvalidValue {
|
||||||
field: "rate_limit_tokens_per_second".to_string(),
|
field: "rate_limit_tokens_per_second".to_string(),
|
||||||
value: tokens_per_second.to_string(),
|
value: tokens_per_second.to_string(),
|
||||||
reason: "Must be > 0 when specified".to_string(),
|
reason: "Must be >= 0 when specified".to_string(),
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -33,10 +33,10 @@ impl TokenBucket {
|
|||||||
/// * `refill_rate` - Tokens added per second
|
/// * `refill_rate` - Tokens added per second
|
||||||
pub fn new(capacity: usize, refill_rate: usize) -> Self {
|
pub fn new(capacity: usize, refill_rate: usize) -> Self {
|
||||||
let capacity = capacity as f64;
|
let capacity = capacity as f64;
|
||||||
|
// Allow refill_rate=0 for pure concurrency limiting (semaphore behavior)
|
||||||
|
// When refill_rate=0, tokens are only returned via return_tokens()
|
||||||
let refill_rate = refill_rate as f64;
|
let refill_rate = refill_rate as f64;
|
||||||
|
|
||||||
let refill_rate = if refill_rate > 0.0 { refill_rate } else { 1.0 };
|
|
||||||
|
|
||||||
Self {
|
Self {
|
||||||
inner: Arc::new(Mutex::new(TokenBucketInner {
|
inner: Arc::new(Mutex::new(TokenBucketInner {
|
||||||
tokens: capacity,
|
tokens: capacity,
|
||||||
@@ -83,6 +83,25 @@ impl TokenBucket {
|
|||||||
return 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.
|
||||||
|
if self.refill_rate == 0.0 {
|
||||||
|
debug!(
|
||||||
|
"Token bucket: waiting indefinitely for {} tokens (refill_rate=0)",
|
||||||
|
tokens
|
||||||
|
);
|
||||||
|
|
||||||
|
loop {
|
||||||
|
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 wait_time = {
|
||||||
let inner = self.inner.lock().await;
|
let inner = self.inner.lock().await;
|
||||||
let tokens_needed = tokens - inner.tokens;
|
let tokens_needed = tokens - inner.tokens;
|
||||||
@@ -176,4 +195,49 @@ mod tests {
|
|||||||
let available = bucket.available_tokens().await;
|
let available = bucket.available_tokens().await;
|
||||||
assert!((4.0..=6.0).contains(&available));
|
assert!((4.0..=6.0).contains(&available));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_token_bucket_zero_refill_rate() {
|
||||||
|
// With refill_rate=0, tokens should only come back via return_tokens()
|
||||||
|
let bucket = TokenBucket::new(2, 0);
|
||||||
|
|
||||||
|
// Acquire both tokens
|
||||||
|
assert!(bucket.try_acquire(1.0).await.is_ok());
|
||||||
|
assert!(bucket.try_acquire(1.0).await.is_ok());
|
||||||
|
|
||||||
|
// No more tokens available
|
||||||
|
assert!(bucket.try_acquire(1.0).await.is_err());
|
||||||
|
|
||||||
|
// Wait - should NOT refill automatically
|
||||||
|
tokio::time::sleep(Duration::from_millis(500)).await;
|
||||||
|
assert!(bucket.try_acquire(1.0).await.is_err());
|
||||||
|
|
||||||
|
// Return a token - now we should be able to acquire
|
||||||
|
bucket.return_tokens(1.0).await;
|
||||||
|
assert!(bucket.try_acquire(1.0).await.is_ok());
|
||||||
|
|
||||||
|
// No more tokens again
|
||||||
|
assert!(bucket.try_acquire(1.0).await.is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_token_bucket_zero_refill_with_notify() {
|
||||||
|
// Test that acquire wakes up when tokens are returned
|
||||||
|
let bucket = Arc::new(TokenBucket::new(1, 0));
|
||||||
|
|
||||||
|
// Acquire the only token
|
||||||
|
assert!(bucket.try_acquire(1.0).await.is_ok());
|
||||||
|
|
||||||
|
let bucket_clone = bucket.clone();
|
||||||
|
|
||||||
|
// Spawn a task that will return the token after a delay
|
||||||
|
tokio::spawn(async move {
|
||||||
|
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||||
|
bucket_clone.return_tokens(1.0).await;
|
||||||
|
});
|
||||||
|
|
||||||
|
// This should wait and then succeed when token is returned
|
||||||
|
let result = bucket.acquire_timeout(1.0, Duration::from_secs(1)).await;
|
||||||
|
assert!(result.is_ok());
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user