aboutsummaryrefslogtreecommitdiff
path: root/src/util
diff options
context:
space:
mode:
Diffstat (limited to 'src/util')
-rw-r--r--src/util/lib.rs1
-rw-r--r--src/util/token_bucket.rs40
2 files changed, 41 insertions, 0 deletions
diff --git a/src/util/lib.rs b/src/util/lib.rs
index c080e3a3..e2e01785 100644
--- a/src/util/lib.rs
+++ b/src/util/lib.rs
@@ -9,3 +9,4 @@ pub mod data;
pub mod error;
pub mod persister;
pub mod time;
+pub mod token_bucket;
diff --git a/src/util/token_bucket.rs b/src/util/token_bucket.rs
new file mode 100644
index 00000000..cc0dfa1f
--- /dev/null
+++ b/src/util/token_bucket.rs
@@ -0,0 +1,40 @@
+use std::time::{Duration, Instant};
+
+use tokio::time::sleep;
+
+pub struct TokenBucket {
+ // Replenish rate: number of tokens per second
+ replenish_rate: u64,
+ // Current number of tokens
+ tokens: u64,
+ // Last replenish time
+ last_replenish: Instant,
+}
+
+impl TokenBucket {
+ pub fn new(replenish_rate: u64) -> Self {
+ Self {
+ replenish_rate,
+ tokens: 0,
+ last_replenish: Instant::now(),
+ }
+ }
+
+ pub async fn take(&mut self, tokens: u64) {
+ while self.tokens < tokens {
+ let needed = tokens - self.tokens;
+ let delay = (needed as f64) / (self.replenish_rate as f64);
+ sleep(Duration::from_secs_f64(delay)).await;
+ self.replenish();
+ }
+ self.tokens -= tokens;
+ }
+
+ pub fn replenish(&mut self) {
+ let now = Instant::now();
+ let new_tokens =
+ ((now - self.last_replenish).as_secs_f64() * (self.replenish_rate as f64)) as u64;
+ self.tokens += new_tokens;
+ self.last_replenish = now;
+ }
+}