[router] refactor router and worker management 3/n (#10727)

sgl-router/src/core/token_bucket.rs

@@ -32,16 +32,11 @@ impl TokenBucket {
         let capacity = capacity as f64;
         let refill_rate = refill_rate as f64;
 
-        // Ensure refill_rate is not zero to prevent division by zero
-        let refill_rate = if refill_rate > 0.0 {
-            refill_rate
-        } else {
-            1.0 // Default to 1 token per second if zero
-        };
+        let refill_rate = if refill_rate > 0.0 { refill_rate } else { 1.0 };
 
         Self {
             inner: Arc::new(Mutex::new(TokenBucketInner {
-                tokens: capacity, // Start full
+                tokens: capacity,
                 last_refill: Instant::now(),
             })),
             notify: Arc::new(Notify::new()),
@@ -54,7 +49,6 @@ impl TokenBucket {
     pub async fn try_acquire(&self, tokens: f64) -> Result<(), ()> {
         let mut inner = self.inner.lock().await;
 
-        // Refill tokens based on elapsed time
         let now = Instant::now();
         let elapsed = now.duration_since(inner.last_refill).as_secs_f64();
         let refill_amount = elapsed * self.refill_rate;
@@ -82,12 +76,10 @@ impl TokenBucket {
 
     /// Acquire tokens, waiting if necessary
     pub async fn acquire(&self, tokens: f64) -> Result<(), tokio::time::error::Elapsed> {
-        // First try to acquire immediately
         if self.try_acquire(tokens).await.is_ok() {
             return Ok(());
         }
 
-        // Calculate wait time
         let wait_time = {
             let inner = self.inner.lock().await;
             let tokens_needed = tokens - inner.tokens;
@@ -100,15 +92,12 @@ impl TokenBucket {
             wait_time, tokens
         );
 
-        // Wait for tokens to be available
         tokio::time::timeout(wait_time, async {
             loop {
-                // Check if we can acquire now
                 if self.try_acquire(tokens).await.is_ok() {
                     return;
                 }
 
-                // Wait for notification or small interval
                 tokio::select! {
                     _ = self.notify.notified() => {},
                     _ = tokio::time::sleep(Duration::from_millis(10)) => {},
@@ -144,7 +133,6 @@ impl TokenBucket {
     pub async fn available_tokens(&self) -> f64 {
         let mut inner = self.inner.lock().await;
 
-        // Refill before checking
         let now = Instant::now();
         let elapsed = now.duration_since(inner.last_refill).as_secs_f64();
         let refill_amount = elapsed * self.refill_rate;
@@ -162,33 +150,26 @@ mod tests {
 
     #[tokio::test]
     async fn test_token_bucket_basic() {
-        let bucket = TokenBucket::new(10, 5); // 10 capacity, 5 per second
+        let bucket = TokenBucket::new(10, 5);
 
-        // Should succeed - bucket starts full
         assert!(bucket.try_acquire(5.0).await.is_ok());
         assert!(bucket.try_acquire(5.0).await.is_ok());
 
-        // Should fail - no tokens left
         assert!(bucket.try_acquire(1.0).await.is_err());
 
-        // Wait for refill
         tokio::time::sleep(Duration::from_millis(300)).await;
 
-        // Should have ~1.5 tokens now
         assert!(bucket.try_acquire(1.0).await.is_ok());
     }
 
     #[tokio::test]
     async fn test_token_bucket_refill() {
-        let bucket = TokenBucket::new(10, 10); // 10 capacity, 10 per second
+        let bucket = TokenBucket::new(10, 10);
 
-        // Use all tokens
         assert!(bucket.try_acquire(10.0).await.is_ok());
 
-        // Wait for partial refill
         tokio::time::sleep(Duration::from_millis(500)).await;
 
-        // Should have ~5 tokens
         let available = bucket.available_tokens().await;
         assert!((4.0..=6.0).contains(&available));
     }