Fix a bug where tokens aren't replenished if there is no activity at all

Author: airhorns

spec/aimd-bucket.spec.ts

@@ -513,45 +513,48 @@ describe("AIMDBucket", () => {
       bucket = new AIMDBucket({ initialRate: 10 });
     });
 
-    it("should process pending requests when tokens are completed with sufficient time gaps", async () => {
-      // This test verifies the fix for the lockup bug where pending requests get stuck
+    it("should process pending requests when no events trigger refill (lockup bug fix)", async () => {
       bucket = new AIMDBucket({
-        initialRate: 4, // Start with low rate like user's case
-        tokenReturnTimeoutMs: 10000, // Long timeout so tokens don't auto-expire
+        initialRate: 3, // Low rate
+        tokenReturnTimeoutMs: 30000, // Long timeout so no auto-timeouts interfere
       });
 
-      // Acquire tokens that will be resolved immediately (up to initial capacity)
-      const immediateTokens = await Promise.all([bucket.acquire(), bucket.acquire(), bucket.acquire(), bucket.acquire()]);
-      expect(immediateTokens).toHaveLength(4);
+      // Get initial 3 tokens immediately (up to initial capacity)
+      const token1 = await bucket.acquire();
+      const token2 = await bucket.acquire();
+      const token3 = await bucket.acquire();
 
-      // These should become pending since we've exhausted the bucket
-      const pendingTokenPromises = [bucket.acquire(), bucket.acquire(), bucket.acquire(), bucket.acquire(), bucket.acquire()];
+      // Advance time to get 1 more token through refill
+      await vi.advanceTimersByTimeAsync(334); // ~1/3 second = 1 token at rate 3
+      const token4 = await bucket.acquire();
 
-      // Advance time slightly to see initial state
-      await vi.advanceTimersByTimeAsync(100);
       let stats = bucket.getStatistics();
       expect(stats.tokensIssued).toBe(4);
-      expect(stats.pendingCount).toBe(5);
+      expect(stats.currentRate).toBe(3);
 
-      // Complete the immediate tokens successfully to trigger rate increase
-      // Add some time between completions to allow bucket to refill
-      for (let i = 0; i < immediateTokens.length; i++) {
-        immediateTokens[i].success();
-        await vi.advanceTimersByTimeAsync(500); // Allow some refill time
-      }
+      token1.success();
+      token2.success();
+      token3.success();
+      token4.success();
 
-      // At this point, some pending requests should have been processed during token completions
+      const pendingPromises = [bucket.acquire(), bucket.acquire(), bucket.acquire(), bucket.acquire(), bucket.acquire()];
+
+      await vi.advanceTimersByTimeAsync(50);
       stats = bucket.getStatistics();
+      expect(stats.currentRate).toBe(3);
+      expect(stats.tokensIssued).toBe(4);
       expect(stats.successCount).toBe(4);
+      expect(stats.pendingCount).toBe(5);
 
-      // All pending requests should have been processed due to the fix
+      // Now advance time significantly, but make NO new acquire() calls and don't complete any more tokens. This a lockup scenario. Time passes, bucket should refill, but pending requests never get processed because no events trigger _refill()
+      await vi.advanceTimersByTimeAsync(5000); // 5 seconds = 15 tokens should be available
+
+      stats = bucket.getStatistics();
       expect(stats.pendingCount).toBe(0);
-      expect(stats.tokensIssued).toBe(9); // 4 immediate + 5 pending
+      expect(stats.tokensIssued).toBe(9);
 
-      // Verify that all pending promises were resolved
-      const resolvedTokens = await Promise.all(pendingTokenPromises);
+      const resolvedTokens = await Promise.all(pendingPromises);
       expect(resolvedTokens).toHaveLength(5);
-      resolvedTokens.forEach((token) => expect(token).toBeInstanceOf(AIMDBucketToken));
     });
 
     it("should process pending requests when tokens timeout automatically", async () => {

src/index.ts

@@ -153,6 +153,7 @@ export class AIMDBucket {
   private tokensIssued = 0; // Only track total issued for reporting
   private pending: { resolve: (token: AIMDBucketToken) => void; reject: (error: Error) => void; timestamp: number }[] = [];
   private isShutdown = false;
+  private pendingTimer?: NodeJS.Timeout;
 
   private config: Required<AIMDBucketConfig>;
 
@@ -215,6 +216,9 @@ export class AIMDBucket {
           },
           timestamp: Date.now(),
         });
+
+        // Set a one-shot timer to process pending requests if no events occur
+        this._schedulePendingCheck();
       });
     }
   }
@@ -260,6 +264,12 @@ export class AIMDBucket {
   async shutdown(): Promise<void> {
     this.isShutdown = true;
 
+    // Clear pending timer
+    if (this.pendingTimer) {
+      clearTimeout(this.pendingTimer);
+      this.pendingTimer = undefined;
+    }
+
     // Reject all pending acquisitions
     const pendingRequests = [...this.pending];
     this.pending = [];
@@ -318,6 +328,36 @@ export class AIMDBucket {
       this.tokensIssued++;
       request.resolve(new AIMDBucketToken(this, this.config.tokenReturnTimeoutMs));
     }
+
+    // Clear pending timer if no more pending requests
+    if (this.pending.length === 0 && this.pendingTimer) {
+      clearTimeout(this.pendingTimer);
+      this.pendingTimer = undefined;
+    }
+  }
+
+  /**
+   * Schedule a one-shot timer to process pending requests if no events occur
+   */
+  private _schedulePendingCheck(): void {
+    // Don't schedule if already scheduled or shutting down
+    if (this.pendingTimer || this.isShutdown) {
+      return;
+    }
+
+    // Calculate next expected refill time based on current rate
+    // At minimum, check every 100ms, but ideally check when next token should be available
+    const nextTokenTime = Math.max(100, 1000 / this.rate);
+
+    this.pendingTimer = setTimeout(() => {
+      this.pendingTimer = undefined;
+      this._refill();
+
+      // If there are still pending requests after refill, schedule another check
+      if (this.pending.length > 0 && !this.isShutdown) {
+        this._schedulePendingCheck();
+      }
+    }, nextTokenTime).unref(); // Use unref() so this doesn't keep the process alive
   }
 
   private _adjustRate(): void {