From 5a5d3c0b659062959654e2b649d9b5d291efc0bf Mon Sep 17 00:00:00 2001
From: Daniel Cadenas <dcadenas@gmail.com>
Date: Fri, 27 Sep 2024 10:15:58 -0300
Subject: [PATCH] Don't overcharge tokens

---
 src/domain/followee_notification_factory.rs |  3 +--
 src/domain/notification_factory.rs          | 10 +++++-----
 src/rate_limiter.rs                         | 11 +++++++----
 3 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/src/domain/followee_notification_factory.rs b/src/domain/followee_notification_factory.rs
index b5217cb..9b69b7b 100644
--- a/src/domain/followee_notification_factory.rs
+++ b/src/domain/followee_notification_factory.rs
@@ -126,8 +126,7 @@ impl FolloweeNotificationFactory {
                 .collect();
 
             let tokens_needed = messages.len() as f64;
-
-            self.rate_limiter.overcharge(tokens_needed);
+            self.rate_limiter.consume(tokens_needed);
 
             if let Some(follow_change) = followers.first() {
                 info!(
diff --git a/src/domain/notification_factory.rs b/src/domain/notification_factory.rs
index bb5888e..f1d9a4e 100644
--- a/src/domain/notification_factory.rs
+++ b/src/domain/notification_factory.rs
@@ -334,13 +334,13 @@ mod tests {
         let messages = notification_factory.flush();
         assert_eq!(messages.len(), 0);
 
-        // And we finally hit the min period
+        // And we finally hit the min period so we get one token
         advance(Duration::from_secs(6u64)).await;
         let messages = notification_factory.flush();
 
-        // But it's retained, there are not tokens
-        assert_eq!(messages.len(), 0);
-        assert_eq!(notification_factory.follow_changes_len(), 1);
+        // The token is used by the last change
+        assert_eq!(messages.len(), 1);
+        assert_eq!(notification_factory.follow_changes_len(), 0);
 
         // And another change arrives
         insert_new_follower(&mut notification_factory, followee);
@@ -351,7 +351,7 @@ mod tests {
         let messages = notification_factory.flush();
         // The new one is flushed, the old one is retained
         assert_eq!(messages.len(), 1);
-        assert_eq!(notification_factory.follow_changes_len(), 2);
+        assert_eq!(notification_factory.follow_changes_len(), 1);
 
         advance(Duration::from_secs(min_seconds_between_messages as u64 + 1)).await;
         let messages = notification_factory.flush();
diff --git a/src/rate_limiter.rs b/src/rate_limiter.rs
index 6800ab1..50aaf07 100644
--- a/src/rate_limiter.rs
+++ b/src/rate_limiter.rs
@@ -42,20 +42,23 @@ impl RateLimiter {
         self.tokens >= tokens_needed
     }
 
-    /// Attempts to consume the specified number of tokens.
+    /// Consumes the specified number of tokens.
     pub fn consume(&mut self, tokens_needed: f64) -> bool {
         self.refill_tokens();
+
         if self.tokens >= tokens_needed {
             self.tokens -= tokens_needed;
             true
         } else {
+            self.tokens = 0.0;
             false
         }
     }
 
-    /// Consumes tokens regardless of availability (can result in negative token count).
+    /// Consumes the specified number of tokens, allows deficit.
     pub fn overcharge(&mut self, tokens_needed: f64) {
         self.refill_tokens();
+
         self.tokens -= tokens_needed;
 
         if self.tokens < -self.max_negative_tokens {
@@ -116,8 +119,8 @@ mod tests {
         // Attempt to consume more tokens than available
         let result = rate_limiter.consume(15.0);
         assert!(!result);
-        // Tokens should remain unchanged since consume failed
-        assert_eq!(rate_limiter.tokens, capacity);
+        // Should have consumed all it could
+        assert_eq!(rate_limiter.tokens, 0.0);
     }
 
     #[tokio::test]