From 30e0e248d6766d4e133336d272951e3a966b5f67 Mon Sep 17 00:00:00 2001
From: Paul Masurel <paul@quickwit.io>
Date: Wed, 8 May 2024 13:12:11 +0900
Subject: [PATCH] Re-added acquire_with_duration in rate_limiter.

---
 quickwit/quickwit-common/src/rate_limiter.rs | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/quickwit/quickwit-common/src/rate_limiter.rs b/quickwit/quickwit-common/src/rate_limiter.rs
index 0b80a7e17db..0e748d65913 100644
--- a/quickwit/quickwit-common/src/rate_limiter.rs
+++ b/quickwit/quickwit-common/src/rate_limiter.rs
@@ -101,6 +101,23 @@ impl RateLimiter {
         }
     }
 
+    /// Acquires some permits from the rate limiter.
+    /// If the permits are not available, returns the duration to wait before trying again.
+    ///
+    /// This method is currently only used in simian.
+    pub fn acquire_with_duration(&mut self, num_permits: u64) -> Result<(), Duration> {
+        if self.acquire_inner(num_permits) {
+            return Ok(());
+        }
+        self.refill(Instant::now());
+        if self.acquire_inner(num_permits) {
+            return Ok(());
+        }
+        let missing = num_permits - self.available_permits;
+        let wait = Duration::from_micros(missing * self.refill_period_micros / self.refill_amount);
+        Err(wait)
+    }
+
     /// Acquires some permits expressed in bytes from the rate limiter. Returns whether the permits
     /// were acquired.
     pub fn acquire_bytes(&mut self, bytes: ByteSize) -> bool {