From 5565251562987f10a9c75f41ca750e7de4f79d62 Mon Sep 17 00:00:00 2001
From: KKRainbow <5665404+KKRainbow@users.noreply.github.com>
Date: Sun, 25 Oct 2026 12:57:13 +0800
Subject: [PATCH] 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.
---
 easytier-core/src/foundation/token_bucket.rs | 29 ++++++++++++++++----
 1 file changed, 23 insertions(+), 6 deletions(-)

diff --git a/easytier-core/src/foundation/token_bucket.rs b/easytier-core/src/foundation/token_bucket.rs
index a22c441d..66838594 100644
--- a/easytier-core/src/foundation/token_bucket.rs
+++ b/easytier-core/src/foundation/token_bucket.rs
@@ -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]