Merge version peek and reset

Implementing peeking

Add more peek functions

Fix failing tests

Implemented peek

Fix n key peek function

Added peek tests for n advances

Implemented reset

Added reset test for dashmap

stable formatting

fix failing no)_std tests

Fix formatting

Clippy lints

Add test for peek

Add tests

Add tests

Add dashmap tests

assert coverage test

Address coverage

formatting

Author: jmfrank63

governor/src/gcra.rs

@@ -129,6 +129,46 @@ impl Gcra {
         })
     }
 
+    pub(crate) fn peek_test<
+        K,
+        P: clock::Reference,
+        S: StateStore<Key = K>,
+        MW: RateLimitingMiddleware<P>,
+    >(
+        &self,
+        start: P,
+        key: &K,
+        state: &S,
+        t0: P,
+    ) -> Result<MW::PositiveOutcome, MW::NegativeOutcome> {
+        let t0 = t0.duration_since(start);
+        let tau = self.tau;
+        let t = self.t;
+        match state.measure_and_peek(key, |tat| {
+            let tat = tat.unwrap_or_else(|| self.starting_state(t0));
+            let earliest_time = tat.saturating_sub(tau);
+            if t0 < earliest_time {
+                Err(MW::disallow(
+                    key,
+                    StateSnapshot::new(self.t, self.tau, earliest_time, earliest_time),
+                    start,
+                ))
+            } else {
+                let next = cmp::max(tat, t0) + t;
+                Ok((
+                    MW::allow(key, StateSnapshot::new(self.t, self.tau, t0, next)),
+                    next,
+                ))
+            }
+        }) {
+            Some(outcome) => outcome,
+            None => Ok(MW::allow(
+                key,
+                StateSnapshot::new(self.t, self.tau, t0, tau),
+            )),
+        }
+    }
+
     /// Tests whether all `n` cells could be accommodated and updates the rate limiter state, if so.
     pub(crate) fn test_n_all_and_update<
         K,
@@ -176,12 +216,68 @@ impl Gcra {
             }
         })
     }
+
+    pub(crate) fn test_n_all_peek<
+        K,
+        P: clock::Reference,
+        S: StateStore<Key = K>,
+        MW: RateLimitingMiddleware<P>,
+    >(
+        &self,
+        start: P,
+        key: &K,
+        n: NonZeroU32,
+        state: &S,
+        t0: P,
+    ) -> Result<MW::PositiveOutcome, NegativeMultiDecision<MW::NegativeOutcome>> {
+        let t0 = t0.duration_since(start);
+        let tau = self.tau;
+        let t = self.t;
+        let additional_weight = t * (n.get() - 1) as u64;
+
+        // check that we can allow enough cells through. Note that `additional_weight` is the
+        // value of the cells *in addition* to the first cell - so add that first cell back.
+        if additional_weight + t > tau {
+            return Err(NegativeMultiDecision::InsufficientCapacity(
+                (tau.as_u64() / t.as_u64()) as u32,
+            ));
+        }
+        match state.measure_and_peek(key, |tat| {
+            let tat = tat.unwrap_or_else(|| self.starting_state(t0));
+            let earliest_time = (tat + additional_weight).saturating_sub(tau);
+            if t0 < earliest_time {
+                Err(NegativeMultiDecision::BatchNonConforming(
+                    n.get(),
+                    MW::disallow(
+                        key,
+                        StateSnapshot::new(self.t, self.tau, earliest_time, earliest_time),
+                        start,
+                    ),
+                ))
+            } else {
+                let next = cmp::max(tat, t0) + t + additional_weight;
+                Ok((
+                    MW::allow(key, StateSnapshot::new(self.t, self.tau, t0, next)),
+                    next,
+                ))
+            }
+        }) {
+            Some(outcome) => outcome,
+            None => Ok(MW::allow(
+                key,
+                StateSnapshot::new(self.t, self.tau, t0, tau),
+            )),
+        }
+    }
 }
 
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::Quota;
+    use crate::RateLimiter;
+    use crate::{clock::FakeRelativeClock, Quota};
+    use no_std_compat::prelude::v1::*;
+    use nonzero_ext::nonzero;
     use std::num::NonZeroU32;
 
     use proptest::prelude::*;
@@ -190,34 +286,40 @@ mod test {
     #[cfg(feature = "std")]
     #[test]
     fn gcra_derives() {
-        use all_asserts::assert_gt;
         use nonzero_ext::nonzero;
 
         let g = Gcra::new(Quota::per_second(nonzero!(1u32)));
         let g2 = Gcra::new(Quota::per_second(nonzero!(2u32)));
         assert_eq!(g, g);
         assert_ne!(g, g2);
-        assert_gt!(format!("{:?}", g).len(), 0);
+        assert!(!format!("{:?}", g).is_empty());
     }
 
     /// Exercise derives and convenience impls on NotUntil to make coverage happy
     #[cfg(feature = "std")]
     #[test]
     fn notuntil_impls() {
-        use crate::RateLimiter;
-        use all_asserts::assert_gt;
-        use clock::FakeRelativeClock;
-        use nonzero_ext::nonzero;
-
         let clock = FakeRelativeClock::default();
         let quota = Quota::per_second(nonzero!(1u32));
         let lb = RateLimiter::direct_with_clock(quota, &clock);
+        for _ in 0..2 {
+            assert!(lb.peek().is_ok());
+        }
         assert!(lb.check().is_ok());
+        assert!(lb
+            .peek()
+            .map_err(|nu| {
+                assert_eq!(nu.earliest_possible(), Nanos::from(1_000_000_000));
+                assert!(!format!("{:?}", nu).is_empty());
+                assert_eq!(format!("{}", nu), "rate-limited until Nanos(1s)");
+                assert_eq!(nu.quota(), quota);
+            })
+            .is_err());
         assert!(lb
             .check()
             .map_err(|nu| {
-                assert_eq!(nu, nu);
-                assert_gt!(format!("{:?}", nu).len(), 0);
+                assert_eq!(nu.earliest_possible(), Nanos::from(1_000_000_000));
+                assert!(!format!("{:?}", nu).is_empty());
                 assert_eq!(format!("{}", nu), "rate-limited until Nanos(1s)");
                 assert_eq!(nu.quota(), quota);
             })
@@ -255,4 +357,139 @@ mod test {
             assert_eq!(quota, back);
         })
     }
+
+    #[test]
+    fn peek_key_test_and_update_works() {
+        let clock = FakeRelativeClock::default();
+        let quota = Quota::per_second(nonzero!(1u32));
+        let lk = RateLimiter::hashmap_with_clock(quota, &clock);
+        let key = 1u32;
+        let key2 = 2u32;
+        for _ in 0..2 {
+            assert!(lk.peek_key(&key).is_ok());
+        }
+        for _ in 0..2 {
+            assert!(lk.peek_key(&key2).is_ok());
+        }
+        assert!(lk.check_key(&key).is_ok());
+        for _ in 0..2 {
+            assert!(lk.peek_key(&key2).is_ok());
+        }
+        assert!(lk
+            .check_key(&key)
+            .map_err(|nu| {
+                assert_eq!(nu.earliest_possible(), Nanos::from(1_000_000_000));
+                assert!(!format!("{:?}", nu).is_empty());
+                assert_eq!(format!("{}", nu), "rate-limited until Nanos(1s)");
+                assert_eq!(nu.quota(), quota);
+            })
+            .is_err());
+        assert!(lk
+            .check_key(&key)
+            .map_err(|nu| {
+                assert_eq!(nu.earliest_possible(), Nanos::from(1_000_000_000));
+                assert!(!format!("{:?}", nu).is_empty());
+                assert_eq!(format!("{}", nu), "rate-limited until Nanos(1s)");
+                assert_eq!(nu.quota(), quota);
+            })
+            .is_err());
+        assert!(lk.check_key(&key2).is_ok());
+        assert!(lk
+            .check_key(&key2)
+            .map_err(|nu| {
+                assert_eq!(nu.earliest_possible(), Nanos::from(1_000_000_000));
+                assert!(!format!("{:?}", nu).is_empty());
+                assert_eq!(format!("{}", nu), "rate-limited until Nanos(1s)");
+                assert_eq!(nu.quota(), quota);
+            })
+            .is_err());
+        assert!(lk
+            .check_key(&key2)
+            .map_err(|nu| {
+                assert_eq!(nu.earliest_possible(), Nanos::from(1_000_000_000));
+                assert!(!format!("{:?}", nu).is_empty());
+                assert_eq!(format!("{}", nu), "rate-limited until Nanos(1s)");
+                assert_eq!(nu.quota(), quota);
+            })
+            .is_err());
+        clock.advance(Duration::from_millis(500));
+        assert!(lk
+            .peek_key(&key)
+            .map_err(|nu| {
+                assert_eq!(nu.earliest_possible(), Nanos::from(1_000_000_000));
+                assert!(!format!("{:?}", nu).is_empty());
+                assert_eq!(format!("{}", nu), "rate-limited until Nanos(1s)");
+                assert_eq!(nu.quota(), quota);
+            })
+            .is_err());
+        assert!(lk
+            .peek_key(&key)
+            .map_err(|nu| {
+                assert_eq!(nu.earliest_possible(), Nanos::from(1_000_000_000));
+                assert!(!format!("{:?}", nu).is_empty());
+                assert_eq!(format!("{}", nu), "rate-limited until Nanos(1s)");
+                assert_eq!(nu.quota(), quota);
+            })
+            .is_err());
+        clock.advance(Duration::from_millis(500));
+        for _ in 0..2 {
+            assert!(lk.peek_key(&key).is_ok());
+        }
+        for _ in 0..2 {
+            assert!(lk.peek_key(&key2).is_ok());
+        }
+        assert!(lk.check_key(&key).is_ok());
+        assert!(lk.check_key(&key2).is_ok());
+        assert!(lk.check_key(&key).is_err());
+        assert!(lk.check_key(&key2).is_err());
+        assert_eq!(lk.reset_key(&key), ());
+        assert_eq!(lk.reset_key(&key2), ());
+        assert_eq!(lk.reset_key(&3), ());
+        assert!(lk.check_key(&key).is_ok());
+        assert!(lk.check_key(&key2).is_ok());
+    }
+
+    #[test]
+    fn peek_n_key_test_and_update_works() {
+        let clock = FakeRelativeClock::default();
+        let quota = Quota::per_second(nonzero!(2u32));
+        let lk = RateLimiter::hashmap_with_clock(quota, &clock);
+        let key = 1u32;
+        let key2 = 2u32;
+        for _ in 0..2 {
+            assert!(lk.peek_key_n(&key, nonzero!(2u32)).is_ok());
+        }
+        for _ in 0..2 {
+            assert!(lk.peek_key_n(&key2, nonzero!(2u32)).is_ok());
+        }
+        for _ in 0..2 {
+            assert!(lk
+                .peek_key_n(&key, nonzero!(3u32))
+                .map_err(|nu| {
+                    assert_eq!(nu, NegativeMultiDecision::InsufficientCapacity(2));
+                })
+                .is_err());
+        }
+        for _ in 0..2 {
+            assert!(lk
+                .peek_key_n(&key2, nonzero!(3u32))
+                .map_err(|nu| {
+                    assert_eq!(nu, NegativeMultiDecision::InsufficientCapacity(2));
+                })
+                .is_err());
+        }
+        assert!(lk.check_key_n(&key, nonzero!(2u32)).is_ok());
+        assert!(lk.check_key_n(&key, nonzero!(1u32)).is_err());
+        assert!(lk.check_key_n(&key2, nonzero!(2u32)).is_ok());
+        assert!(lk.check_key_n(&key2, nonzero!(1u32)).is_err());
+        for _ in 0..2 {
+            assert!(lk.peek_key_n(&key, nonzero!(1u32)).is_err());
+        }
+        for _ in 0..2 {
+            assert!(lk.peek_key_n(&key2, nonzero!(1u32)).is_err());
+        }
+        assert_eq!(lk.reset_key(&key), ());
+        assert!(lk.check_key_n(&key, nonzero!(2u32)).is_ok());
+        // TODO: impl Reference for FakeRelativeClock and test returned error
+    }
 }

governor/src/state.rs

@@ -46,6 +46,13 @@ pub trait StateStore {
     fn measure_and_replace<T, F, E>(&self, key: &Self::Key, f: F) -> Result<T, E>
     where
         F: Fn(Option<Nanos>) -> Result<(T, Nanos), E>;
+    /// `measure_and_peek` does the same as `measure_and_replace` except
+    /// it will not change the state. This is useful if you need to know the next
+    /// decision in advance
+    fn measure_and_peek<T, F, E>(&self, key: &Self::Key, f: F) -> Option<Result<T, E>>
+    where
+        F: Fn(Option<Nanos>) -> Result<(T, Nanos), E>;
+    fn reset(&self, key: &Self::Key);
 }
 
 /// A rate limiter.
@@ -134,12 +141,11 @@ where
 mod test {
     use super::*;
     use crate::Quota;
-    use all_asserts::assert_gt;
     use nonzero_ext::nonzero;
 
     #[test]
     fn ratelimiter_impl_coverage() {
         let lim = RateLimiter::direct(Quota::per_second(nonzero!(3u32)));
-        assert_gt!(format!("{:?}", lim).len(), 0);
+        assert!(!format!("{:?}", lim).is_empty());
     }
 }

governor/src/state/direct.rs

@@ -80,6 +80,15 @@ where
         )
     }
 
+    pub fn peek(&self) -> Result<MW::PositiveOutcome, MW::NegativeOutcome> {
+        self.gcra.peek_test::<NotKeyed, C::Instant, S, MW>(
+            self.start,
+            &NotKeyed::NonKey,
+            &self.state,
+            self.clock.now(),
+        )
+    }
+
     /// Allow *only all* `n` cells through the rate limiter.
     ///
     /// This method can succeed in only one way and fail in two ways:
@@ -107,6 +116,19 @@ where
                 self.clock.now(),
             )
     }
+
+    pub fn peek_n(
+        &self,
+        n: NonZeroU32,
+    ) -> Result<MW::PositiveOutcome, NegativeMultiDecision<MW::NegativeOutcome>> {
+        self.gcra.test_n_all_peek::<NotKeyed, C::Instant, S, MW>(
+            self.start,
+            &NotKeyed::NonKey,
+            n,
+            &self.state,
+            self.clock.now(),
+        )
+    }
 }
 
 #[cfg(feature = "std")]

governor/src/state/direct/future.rs

@@ -118,14 +118,13 @@ where
 
 #[cfg(test)]
 mod test {
-    use all_asserts::assert_gt;
 
     use super::*;
 
     #[test]
     fn insufficient_capacity_impl_coverage() {
         let i = InsufficientCapacity(1);
         assert_eq!(i.0, i.clone().0);
-        assert_gt!(format!("{}", i).len(), 0);
+        assert!(!format!("{}", i).is_empty());
     }
 }