speedbump/
lib.rs

1use crate::{store::LocalStore, strategy::LimitStrategy};
2
3pub mod store;
4pub mod strategy;
5
6impl Limiter<Unset, Unset> {
7    #[must_use]
8    pub fn builder() -> LimiterBuilder<Unset, Unset> {
9        LimiterBuilder::new()
10    }
11}
12
13pub struct Limiter<Store, Strategy> {
14    store: Store,
15    strategy: Strategy,
16}
17
18#[derive(thiserror::Error, Debug)]
19pub enum Error<Str, Sto> {
20    Strategy(Str),
21    Store(Sto),
22}
23
24impl<Store, Strategy> Limiter<Store, Strategy>
25where
26    Strategy: LimitStrategy,
27    Store: LocalStore<<Strategy as LimitStrategy>::State>,
28{
29    /// Checks the limit for a given key
30    ///
31    /// # Errors
32    /// Errors when the underlying strategy or store errors
33    pub async fn limit(&self, key: &str) -> Result<bool, Error<Strategy::Error, Store::Error>> {
34        let mut state = self
35            .store
36            .get(key)
37            .await
38            .map_err(Error::Store)?
39            .unwrap_or(self.strategy.initialize_state());
40
41        let limited = self
42            .strategy
43            .check_limit(&mut state)
44            .map_err(Error::Strategy)?;
45
46        self.store.set(key, state).await.map_err(Error::Store)?;
47
48        Ok(limited.is_allowed())
49    }
50}
51pub struct Unset;
52
53pub struct LimiterBuilder<Store, Strategy> {
54    store: Store,
55    strategy: Strategy,
56}
57
58impl Default for LimiterBuilder<Unset, Unset> {
59    fn default() -> Self {
60        Self::new()
61    }
62}
63
64impl LimiterBuilder<Unset, Unset> {
65    #[must_use]
66    pub fn new() -> Self {
67        Self {
68            store: Unset,
69            strategy: Unset,
70        }
71    }
72}
73
74impl<Store, Strategy> LimiterBuilder<Store, Strategy> {
75    pub fn store<S>(self, store: S) -> LimiterBuilder<S, Strategy> {
76        LimiterBuilder {
77            store,
78            strategy: self.strategy,
79        }
80    }
81
82    pub fn strategy<S>(self, strategy: S) -> LimiterBuilder<Store, S> {
83        LimiterBuilder {
84            store: self.store,
85            strategy,
86        }
87    }
88}
89
90impl<Store, Strategy> LimiterBuilder<Store, Strategy>
91where
92    Store: LocalStore<<Strategy as LimitStrategy>::State>,
93    Strategy: LimitStrategy,
94{
95    pub fn build(self) -> Limiter<Store, Strategy> {
96        Limiter {
97            store: self.store,
98            strategy: self.strategy,
99        }
100    }
101}