Separate out Check And Cell Consumption

governor/src/gcra.rs

@@ -174,6 +174,125 @@ impl Gcra {
             }
         }))
     }
+
+    /// Tests whether all `n` cells could be accommodated.
+    /// No update would be made.
+    pub(crate) fn test_n_all_without_update<
+        K,
+        P: clock::Reference,
+        S: StateStore<Key = K>,
+        MW: RateLimitingMiddleware<P>,
+    >(
+        &self,
+        start: P,
+        key: &K,
+        n: NonZeroU32,
+        state: &S,
+        t0: P,
+    ) -> Result<Result<MW::PositiveOutcome, MW::NegativeOutcome>, InsufficientCapacity> {
+        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(InsufficientCapacity((tau.as_u64() / t.as_u64()) as u32));
+        }
+        Ok(state.measure(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(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),
+                ))
+            }
+        }))
+    }
+
+    /// Tests a single cell against the rate limiter state.
+    /// No update would be made.
+    pub(crate) fn test_without_update<
+        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;
+        state.measure(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),
+                ))
+            }
+        })
+    }
+
+    /// Update a single cell against the rate limiter state at the given key.
+    pub(crate) fn update<K, P: clock::Reference, S: StateStore<Key = K>>(
+        &self,
+        start: P,
+        key: &K,
+        state: &S,
+        t0: P,
+    ) {
+        let t0 = t0.duration_since(start);
+        let t = self.t;
+        let _ = state.measure_and_replace(key, |tat| {
+            let tat = tat.unwrap_or_else(|| self.starting_state(t0));
+            let next = cmp::max(tat, t0) + t;
+            // always ask state to update
+            Ok::<((), Nanos), ()>(((), next))
+        });
+    }
+
+    /// Update `n` cells for the rate limiter state.
+    pub(crate) fn update_n<K, P: clock::Reference, S: StateStore<Key = K>>(
+        &self,
+        start: P,
+        key: &K,
+        n: NonZeroU32,
+        state: &S,
+        t0: P,
+    ) {
+        let t0 = t0.duration_since(start);
+        let t = self.t;
+        let additional_weight = t * (n.get() - 1) as u64;
+
+        let _ = state.measure_and_replace(key, |tat| {
+            let tat = tat.unwrap_or_else(|| self.starting_state(t0));
+
+            let next = cmp::max(tat, t0) + t + additional_weight;
+            Ok::<((), Nanos), ()>(((), next))
+        });
+    }
 }
 
 #[cfg(test)]

governor/src/state.rs

@@ -46,6 +46,11 @@ 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>;
+
+    /// Same as [`measure_and_replace`](`StateStore::measure_and_replace`), but it would not replace the value at the key
+    fn measure<T, F, E>(&self, key: &Self::Key, f: F) -> Result<T, E>
+    where
+        F: Fn(Option<Nanos>) -> Result<T, E>;
 }
 
 /// A rate limiter.

governor/src/state/direct.rs

@@ -108,6 +108,57 @@ where
                 self.clock.now(),
             )
     }
+
+    /// same as `check`, but will not update internal state.
+    /// It would only query if the rate limit is reached.
+    pub fn check_only(&self) -> Result<MW::PositiveOutcome, MW::NegativeOutcome> {
+        self.gcra
+            .test_without_update::<NotKeyed, C::Instant, S, MW>(
+                self.start,
+                &NotKeyed::NonKey,
+                &self.state,
+                self.clock.now(),
+            )
+    }
+
+    /// same as `check_n`, but will not update internal state.
+    /// It would only query if all `n` cells can be accommodated.
+    pub fn check_n_only(
+        &self,
+        n: NonZeroU32,
+    ) -> Result<Result<MW::PositiveOutcome, MW::NegativeOutcome>, InsufficientCapacity> {
+        self.gcra
+            .test_n_all_without_update::<NotKeyed, C::Instant, S, MW>(
+                self.start,
+                &NotKeyed::NonKey,
+                n,
+                &self.state,
+                self.clock.now(),
+            )
+    }
+
+    /// Consume a cell through the rate limiter.
+    /// If no cell is available, it would "borrow" from future cells.
+    pub fn consume(&self) {
+        self.gcra.update::<NotKeyed, C::Instant, S>(
+            self.start,
+            &NotKeyed::NonKey,
+            &self.state,
+            self.clock.now(),
+        )
+    }
+
+    /// Consume n cells through the rate limiter.
+    /// If no cell is available, it would "borrow" from future cells.
+    pub fn consume_n(&self, n: NonZeroU32) {
+        self.gcra.update_n::<NotKeyed, C::Instant, S>(
+            self.start,
+            &NotKeyed::NonKey,
+            n,
+            &self.state,
+            self.clock.now(),
+        )
+    }
 }
 
 #[cfg(feature = "std")]

governor/src/state/in_memory.rs

@@ -44,6 +44,14 @@ impl InMemoryState {
         decision.map(|(result, _)| result)
     }
 
+    pub(crate) fn measure<T, F, E>(&self, mut f: F) -> Result<T, E>
+    where
+        F: FnMut(Option<Nanos>) -> Result<T, E>,
+    {
+        let prev = self.0.load(Ordering::Acquire);
+        f(NonZeroU64::new(prev).map(|n| n.get().into()))
+    }
+
     pub(crate) fn is_older_than(&self, nanos: Nanos) -> bool {
         self.0.load(Ordering::Relaxed) <= nanos.into()
     }
@@ -59,6 +67,13 @@ impl StateStore for InMemoryState {
     {
         self.measure_and_replace_one(f)
     }
+
+    fn measure<T, F, E>(&self, _key: &Self::Key, f: F) -> Result<T, E>
+    where
+        F: Fn(Option<Nanos>) -> Result<T, E>,
+    {
+        self.measure(f)
+    }
 }
 
 impl Debug for InMemoryState {

governor/src/state/keyed.rs

@@ -258,6 +258,13 @@ mod test {
             {
                 f(None).map(|(res, _)| res)
             }
+
+            fn measure<T, F, E>(&self, _key: &Self::Key, f: F) -> Result<T, E>
+            where
+                F: Fn(Option<Nanos>) -> Result<T, E>,
+            {
+                f(None)
+            }
         }
 
         impl<K: Hash + Eq + Clone> ShrinkableKeyedStateStore<K> for NaiveKeyedStateStore<K> {

governor/src/state/keyed/dashmap.rs

@@ -27,6 +27,19 @@ impl<K: Hash + Eq + Clone> StateStore for DashMapStateStore<K> {
         let entry = self.entry(key.clone()).or_default();
         (*entry).measure_and_replace_one(f)
     }
+
+    fn measure<T, F, E>(&self, key: &Self::Key, f: F) -> Result<T, E>
+    where
+        F: Fn(Option<Nanos>) -> Result<T, E>,
+    {
+        if let Some(v) = self.get(key) {
+            // fast path: measure existing entry
+            return v.measure(f);
+        }
+        // make an entry and measure that:
+        let entry = self.entry(key.clone()).or_default();
+        (*entry).measure(f)
+    }
 }
 
 /// # Keyed rate limiters - [`DashMap`]-backed

governor/src/state/keyed/hashmap.rs

@@ -35,6 +35,20 @@ impl<K: Hash + Eq + Clone> StateStore for HashMapStateStore<K> {
         let entry = (*map).entry(key.clone()).or_default();
         entry.measure_and_replace_one(f)
     }
+
+    fn measure<T, F, E>(&self, key: &Self::Key, f: F) -> Result<T, E>
+    where
+        F: Fn(Option<Nanos>) -> Result<T, E>,
+    {
+        let mut map = self.lock();
+        if let Some(v) = (*map).get(key) {
+            // fast path: a rate limiter is already present for the key.
+            return v.measure(f);
+        }
+        // not-so-fast path: make a new entry and measure it.
+        let entry = (*map).entry(key.clone()).or_default();
+        entry.measure(f)
+    }
 }
 
 impl<K: Hash + Eq + Clone> ShrinkableKeyedStateStore<K> for HashMapStateStore<K> {