From 6b47c294f570b141a2349d5b6da537c0b64d165d Mon Sep 17 00:00:00 2001
From: Alex Auvolat <alex@adnab.me>
Date: Wed, 27 Oct 2021 10:36:04 +0200
Subject: Refactoring on repair commands

---
 src/util/token_bucket.rs | 40 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)
 create mode 100644 src/util/token_bucket.rs

(limited to 'src/util/token_bucket.rs')

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;
+	}
+}
-- 
cgit v1.2.3