From f98f9615d610663b8059bb701dc67d04279c752a Mon Sep 17 00:00:00 2001
From: Brian Martin <brayniac@gmail.com>
Date: Mon, 15 Apr 2024 10:19:21 -0700
Subject: [PATCH] ratelimit: add tracking for number of dropped tokens

Adds a tracking counter for the number of tokens dropped due to
bucket overflow.
---
 ratelimit/Cargo.toml |  2 +-
 ratelimit/src/lib.rs | 20 ++++++++++++++++++--
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/ratelimit/Cargo.toml b/ratelimit/Cargo.toml
index 8b153f2..470668f 100644
--- a/ratelimit/Cargo.toml
+++ b/ratelimit/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "ratelimit"
-version = "0.9.0"
+version = "0.9.1"
 authors = ["Brian Martin <brian@pelikan.io>"]
 edition = "2021"
 license = "MIT OR Apache-2.0"
diff --git a/ratelimit/src/lib.rs b/ratelimit/src/lib.rs
index ff97a83..fff1fcb 100644
--- a/ratelimit/src/lib.rs
+++ b/ratelimit/src/lib.rs
@@ -79,6 +79,7 @@ struct Parameters {
 
 pub struct Ratelimiter {
     available: AtomicU64,
+    dropped: AtomicU64,
     parameters: RwLock<Parameters>,
     refill_at: AtomicInstant,
 }
@@ -177,10 +178,13 @@ impl Ratelimiter {
         }
     }
 
+    /// Returns the number of tokens currently available.
     pub fn available(&self) -> u64 {
         self.available.load(Ordering::Relaxed)
     }
 
+    /// Sets the number of tokens available to some amount. Returns an error if
+    /// the amount exceeds the bucket capacity.
     pub fn set_available(&self, amount: u64) -> Result<(), Error> {
         let parameters = self.parameters.read();
         if amount > parameters.capacity {
@@ -191,6 +195,12 @@ impl Ratelimiter {
         }
     }
 
+    /// Returns the number of tokens that have been dropped due to bucket
+    /// overflowing.
+    pub fn dropped(&self) -> u64 {
+        self.dropped.load(Ordering::Relaxed)
+    }
+
     /// Internal function to refill the token bucket. Called as part of
     /// `try_wait()`
     fn refill(&self, time: Instant) -> Result<(), core::time::Duration> {
@@ -236,8 +246,12 @@ impl Ratelimiter {
         let available = self.available.load(Ordering::Acquire);
 
         if available + amount >= parameters.capacity {
-            self.available
-                .fetch_add(parameters.capacity - available, Ordering::Release);
+            // we will fill the bucket up to the capacity
+            let to_add = parameters.capacity - available;
+            self.available.fetch_add(to_add, Ordering::Release);
+
+            // and increment the number of tokens dropped
+            self.dropped.fetch_add(amount - to_add, Ordering::Relaxed);
         } else {
             self.available.fetch_add(amount, Ordering::Release);
         }
@@ -395,6 +409,7 @@ impl Builder {
 
         Ok(Ratelimiter {
             available,
+            dropped: AtomicU64::new(0),
             parameters: parameters.into(),
             refill_at,
         })
@@ -459,6 +474,7 @@ mod tests {
         std::thread::sleep(Duration::from_millis(10));
         assert!(rl.try_wait().is_ok());
         assert!(rl.try_wait().is_err());
+        assert!(rl.dropped() >= 8);
     }
 
     // quick test that capacity acts as expected