fix: Rate limiter TokenBucket::auto_replenish()

Frequently calling `auto_replenish` will reset `self.last_update` each
time and `tokens` may be a fractional value (0 since it is a `u64`), in
this case no tokens will replenished. To avoid this we only update
`self.last_update` when `tokens > 0`.

Signed-off-by: Jonathan Woollett-Light <jcawl@amazon.co.uk>

src/rate_limiter/src/lib.rs

@@ -161,14 +161,21 @@ impl TokenBucket {
     fn auto_replenish(&mut self) {
         // Compute time passed since last refill/update.
         let time_delta = self.last_update.elapsed().as_nanos() as u64;
-        self.last_update = Instant::now();
 
         // At each 'time_delta' nanoseconds the bucket should refill with:
         // refill_amount = (time_delta * size) / (complete_refill_time_ms * 1_000_000)
         // `processed_capacity` and `processed_refill_time` are the result of simplifying above
         // fraction formula with their greatest-common-factor.
         let tokens = (time_delta * self.processed_capacity) / self.processed_refill_time;
-        self.budget = std::cmp::min(self.budget + tokens, self.size);
+
+        // Frequently calling `auto_replenish` will reset `self.last_update` each
+        // time and `tokens` may be a fractional value (0 since it is a `u64`), in
+        // this case no tokens will be replenished. To avoid this we only update
+        // `self.last_update` when `tokens > 0`.
+        if tokens > 0 {
+            self.last_update = Instant::now();
+            self.budget = std::cmp::min(self.budget + tokens, self.size);
+        }
     }
 
     /// Attempts to consume `tokens` from the bucket and returns whether the action succeeded.
@@ -560,6 +567,42 @@ pub(crate) mod tests {
         }
     }
 
+    #[test]
+    fn test_token_bucket_auto_replenish() {
+        // These values will give 1 token every 100 milliseconds
+        const SIZE: u64 = 10;
+        const TIME: u64 = 1000;
+        let mut tb = TokenBucket::new(SIZE, 0, TIME).unwrap();
+        tb.reduce(SIZE);
+        assert_eq!(tb.budget(), 0);
+
+        // Auto-replenishing after 10 milliseconds should not yield any tokens
+        thread::sleep(Duration::from_millis(10));
+        tb.auto_replenish();
+        assert_eq!(tb.budget(), 0);
+
+        // Neither after 20.
+        thread::sleep(Duration::from_millis(10));
+        tb.auto_replenish();
+        assert_eq!(tb.budget(), 0);
+
+        // We should get 1 token after 100 millis
+        thread::sleep(Duration::from_millis(80));
+        tb.auto_replenish();
+        assert_eq!(tb.budget(), 1);
+
+        // So, 5 after 500 millis
+        thread::sleep(Duration::from_millis(400));
+        tb.auto_replenish();
+        assert_eq!(tb.budget(), 5);
+
+        // And be fully replenished after 1 second.
+        // Wait more here to make sure we do not overshoot
+        thread::sleep(Duration::from_millis(1000));
+        tb.auto_replenish();
+        assert_eq!(tb.budget(), 10);
+    }
+
     #[test]
     fn test_token_bucket_create() {
         let before = Instant::now();