test(limiter): stop racing the token bucket refill clock (#2752)

* test(limiter): stop racing the token bucket refill clock

test_initial_state drains a 1000 byte bucket that refills at 1000
bytes/s and then asserts that one more token is rejected. try_consume
refills from Instant::now() first, so the assertion only holds while
the thread runs within one millisecond of the drain; one millisecond
of scheduling delay is enough for the token to appear.

Freeze the lazy refill by moving last_refill into the future so no
token can appear between the two calls, and apply the same freeze to
test_consumption, which had the same race with a 100 ms window.

* test(limiter): pass the refill instant explicitly instead of faking the clock

Moving last_refill into the future only widened the window: once that
offset passes the bucket refills again, so the tests still depended on
the wall clock.

Split the locked consumption into try_consume_locked(state, tokens, now)
and have production take the lock first, then read the clock. Both
tests take one instant and reuse it for every call, so no refill can
happen between the assertions.
This commit is contained in:
KKRainbow
2026-10-25 12:57:13 +08:00
committed by GitHub
parent ffcfba4afd
commit 5565251562
+23 -6
View File
@@ -139,7 +139,16 @@ impl TokenBucket {
let tokens = tokens.min(self.config.capacity);
let mut state = self.state.lock();
self.refill(&mut state, Instant::now());
let now = Instant::now();
self.try_consume_locked(&mut state, tokens, now)
}
/// Consume from an already locked state at an explicit instant.
///
/// Production takes the lock first and then reads the clock, so tests can
/// pass the same instant twice to observe the bucket without any refill.
fn try_consume_locked(&self, state: &mut BucketState, tokens: u64, now: Instant) -> bool {
self.refill(state, now);
if state.available_tokens < tokens {
return false;
}
@@ -313,10 +322,14 @@ mod tests {
#[tokio::test]
async fn test_initial_state() {
let bucket = TokenBucket::new(1000, 1000);
// Refill is time based; reuse one instant so no token can appear
// between the two calls regardless of scheduling.
let now = Instant::now();
let mut state = bucket.state.lock();
// Should have full capacity initially
assert!(bucket.try_consume(1000));
assert!(!bucket.try_consume(1)); // Should be empty now
assert!(bucket.try_consume_locked(&mut state, 1000, now));
assert!(!bucket.try_consume_locked(&mut state, 1, now)); // Should be empty now
}
/// Test token consumption behavior
@@ -324,14 +337,18 @@ mod tests {
async fn test_consumption() {
let bucket = TokenBucket::new(1500, 1000);
// Same instant for every call: no refill can happen between them.
let now = Instant::now();
let mut state = bucket.state.lock();
// First packet should succeed
assert!(bucket.try_consume(1000));
assert!(bucket.try_consume_locked(&mut state, 1000, now));
// Second packet should fail (only 500 left)
assert!(!bucket.try_consume(600));
assert!(!bucket.try_consume_locked(&mut state, 600, now));
// Should be able to take remaining tokens
assert!(bucket.try_consume(500));
assert!(bucket.try_consume_locked(&mut state, 500, now));
}
#[tokio::test]