Merge pull request #213 from AcalaNetwork/add_missing_methods

add keyed future missing methods

governor/src/state/keyed/future.rs

@@ -1,11 +1,11 @@
 use std::prelude::v1::*;
 
 use crate::{
-    clock, middleware::RateLimitingMiddleware, state::keyed::KeyedStateStore, Jitter, NotUntil,
-    RateLimiter,
+    clock, errors::InsufficientCapacity, middleware::RateLimitingMiddleware,
+    state::keyed::KeyedStateStore, Jitter, NotUntil, RateLimiter,
 };
 use futures_timer::Delay;
-use std::hash::Hash;
+use std::{hash::Hash, num::NonZeroU32};
 
 #[cfg(feature = "std")]
 /// # Keyed rate limiters - `async`/`await`
@@ -57,4 +57,47 @@ where
             }
         }
     }
+
+    /// Asynchronously resolves as soon as the rate limiter allows it.
+    ///
+    /// This is similar to `until_key_ready` except it waits for an abitrary number
+    /// of `n` cells to be available.
+    ///
+    /// Returns `InsufficientCapacity` if the `n` provided exceeds the maximum
+    /// capacity of the rate limiter.
+    pub async fn until_key_n_ready(
+        &self,
+        key: &K,
+        n: NonZeroU32,
+    ) -> Result<MW::PositiveOutcome, InsufficientCapacity> {
+        self.until_key_n_ready_with_jitter(key, n, Jitter::NONE)
+            .await
+    }
+
+    /// Asynchronously resolves as soon as the rate limiter allows it, with a
+    /// randomized wait period.
+    ///
+    /// This is similar to `until_key_ready_with_jitter` except it waits for an
+    /// abitrary number of `n` cells to be available.
+    ///
+    /// Returns `InsufficientCapacity` if the `n` provided exceeds the maximum
+    /// capacity of the rate limiter.
+    pub async fn until_key_n_ready_with_jitter(
+        &self,
+        key: &K,
+        n: NonZeroU32,
+        jitter: Jitter,
+    ) -> Result<MW::PositiveOutcome, InsufficientCapacity> {
+        loop {
+            match self.check_key_n(key, n)? {
+                Ok(x) => {
+                    return Ok(x);
+                }
+                Err(negative) => {
+                    let delay = Delay::new(jitter + negative.wait_time_from(self.clock.now()));
+                    delay.await;
+                }
+            }
+        }
+    }
 }

governor/tests/future.rs

@@ -55,6 +55,18 @@ fn pauses_keyed() {
     assert_ge!(i.elapsed(), Duration::from_millis(100));
 }
 
+#[test]
+fn pauses_keyed_n() {
+    let lim = RateLimiter::keyed(Quota::per_second(nonzero!(10u32)));
+
+    for _ in 0..6 {
+        lim.check_key(&1u32).unwrap();
+    }
+    let i = Instant::now();
+    block_on(lim.until_key_n_ready(&1u32, nonzero!(5u32))).unwrap();
+    assert_ge!(i.elapsed(), Duration::from_millis(100));
+}
+
 #[test]
 fn proceeds() {
     let lim = RateLimiter::direct(Quota::per_second(nonzero!(2u32)));
@@ -79,6 +91,14 @@ fn proceeds_keyed() {
     assert_le!(i.elapsed(), MAX_TEST_RUN_DURATION);
 }
 
+#[test]
+fn proceeds_keyed_n() {
+    let lim = RateLimiter::keyed(Quota::per_second(nonzero!(3u32)));
+    let i = Instant::now();
+    block_on(lim.until_key_n_ready(&1u32, nonzero!(2u32))).unwrap();
+    assert_le!(i.elapsed(), MAX_TEST_RUN_DURATION);
+}
+
 #[test]
 fn multiple() {
     let lim = Arc::new(RateLimiter::direct(Quota::per_second(nonzero!(10u32))));
@@ -125,4 +145,8 @@ fn errors_on_exceeded_capacity() {
     let lim = RateLimiter::direct(Quota::per_second(nonzero!(10u32)));
 
     block_on(lim.until_n_ready(nonzero!(11u32))).unwrap_err();
+
+    let lim = RateLimiter::keyed(Quota::per_second(nonzero!(10u32)));
+
+    block_on(lim.until_key_n_ready(&1u32, nonzero!(11u32))).unwrap_err();
 }