aboutsummaryrefslogtreecommitdiff
path: root/src/util/token_bucket.rs
diff options
context:
space:
mode:
authorAlex Auvolat <alex@adnab.me>2021-10-27 10:36:04 +0200
committerAlex Auvolat <alex@adnab.me>2021-10-27 11:14:55 +0200
commit6b47c294f570b141a2349d5b6da537c0b64d165d (patch)
tree9716a373adc6b51a249af015e2a461904210d4ba /src/util/token_bucket.rs
parent28c015d9ffcb3255295572465654fdca680b1964 (diff)
downloadgarage-6b47c294f570b141a2349d5b6da537c0b64d165d.tar.gz
garage-6b47c294f570b141a2349d5b6da537c0b64d165d.zip
Refactoring on repair commandscli-verify-integrity
Diffstat (limited to 'src/util/token_bucket.rs')
-rw-r--r--src/util/token_bucket.rs40
1 files changed, 40 insertions, 0 deletions
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;
+ }
+}