governor/state/
direct.rs

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