http_rate/state/
direct.rs

1#![allow(unused)]
2
3use core::num::NonZeroU32;
4
5use crate::{
6    error::InsufficientCapacity,
7    gcra::NotUntil,
8    quota::Quota,
9    snapshot::RateSnapshot,
10    state::InMemoryState,
11    timer::{DefaultTimer, Timer},
12};
13
14/// The "this state store does not use keys" key type.
15///
16/// It's possible to use this to create a "direct" rate limiter. It explicitly does not implement
17/// [`Hash`][std::hash::Hash] so that it is possible to tell apart from "hashable" key types.
18#[derive(PartialEq, Debug, Eq)]
19pub enum NotKeyed {
20    /// The value given to state stores' methods.
21    NonKey,
22}
23
24/// A trait for state stores that only keep one rate limiting state.
25///
26/// This is blanket-implemented by all [`StateStore`]s with [`NotKeyed`] key associated types.
27pub(crate) trait DirectStateStore: StateStore<Key = NotKeyed> {}
28
29impl<T> DirectStateStore for T where T: StateStore<Key = NotKeyed> {}
30
31/// # Direct in-memory rate limiters - Constructors
32///
33/// Here we construct an in-memory rate limiter that makes direct (un-keyed)
34/// rate-limiting decisions. Direct rate limiters can be used to
35/// e.g. regulate the transmission of packets on a single connection,
36/// or to ensure that an API client stays within a service's rate
37/// limit.
38impl RateLimiter<NotKeyed, InMemoryState, DefaultTimer> {
39    /// Constructs a new in-memory direct rate limiter for a quota with the default real-time clock.
40    pub(crate) fn direct(quota: Quota) -> RateLimiter<NotKeyed, InMemoryState, DefaultTimer> {
41        let clock = DefaultTimer;
42        Self::direct_with_clock(quota, &clock)
43    }
44}
45
46impl<C> RateLimiter<NotKeyed, InMemoryState, C>
47where
48    C: Timer,
49{
50    /// Constructs a new direct rate limiter for a quota with a custom clock.
51    pub(crate) fn direct_with_clock(quota: Quota, clock: &C) -> Self {
52        let state: InMemoryState = Default::default();
53        RateLimiter::new(quota, state, clock)
54    }
55}
56
57/// # Direct rate limiters - Manually checking cells
58impl<S, C> RateLimiter<NotKeyed, S, C>
59where
60    S: DirectStateStore,
61    C: Timer,
62{
63    /// Allow a single cell through the rate limiter.
64    ///
65    /// If the rate limit is reached, `check` returns information about the earliest
66    /// time that a cell might be allowed through again.
67    pub(crate) fn check(&self) -> Result<RateSnapshot, NotUntil<C::Instant>> {
68        self.gcra.test_and_update::<NotKeyed, C::Instant, S>(
69            self.start,
70            &NotKeyed::NonKey,
71            &self.state,
72            self.clock.now(),
73        )
74    }
75
76    #[cfg(test)]
77    /// Allow *only all* `n` cells through the rate limiter.
78    ///
79    /// This method can succeed in only one way and fail in two ways:
80    /// * Success: If all `n` cells can be accommodated, it returns `Ok(())`.
81    /// * Failure (but ok): Not all cells can make it through at the current time.
82    ///   The result is `Err(NegativeMultiDecision::BatchNonConforming(NotUntil))`, which can
83    ///   be interrogated about when the batch might next conform.
84    /// * Failure (the batch can never go through): The rate limit quota's burst size is too low
85    ///   for the given number of cells to ever be allowed through.
86    ///
87    /// ### Performance
88    /// This method diverges a little from the GCRA algorithm, using
89    /// multiplication to determine the next theoretical arrival time, and so
90    /// is not as fast as checking a single cell.
91    pub(crate) fn check_n(
92        &self,
93        n: NonZeroU32,
94    ) -> Result<Result<RateSnapshot, NotUntil<C::Instant>>, InsufficientCapacity> {
95        self.gcra.test_n_all_and_update::<NotKeyed, C::Instant, S>(
96            self.start,
97            &NotKeyed::NonKey,
98            n,
99            &self.state,
100            self.clock.now(),
101        )
102    }
103}
104
105use crate::state::{RateLimiter, StateStore};
106
107#[cfg(test)]
108mod test {
109    use super::*;
110
111    #[test]
112    fn not_keyed_impls_coverage() {
113        assert_eq!(NotKeyed::NonKey, NotKeyed::NonKey);
114    }
115}