salvo_rate_limiter/
quota.rs

1use std::borrow::Borrow;
2use std::convert::Infallible;
3use std::error::Error as StdError;
4use std::hash::Hash;
5
6use serde::{Deserialize, Serialize};
7use time::Duration;
8
9/// Used to get quota and you can config users' quota config in database.
10pub trait QuotaGetter<Key>: Send + Sync + 'static {
11    /// Quota type.
12    type Quota: Clone + Send + Sync + 'static;
13    /// Error type.
14    type Error: StdError;
15
16    /// Get quota.
17    fn get<Q>(&self, key: &Q) -> impl Future<Output = Result<Self::Quota, Self::Error>> + Send
18    where
19        Key: Borrow<Q>,
20        Q: Hash + Eq + Sync;
21}
22
23/// A basic quota.
24#[non_exhaustive]
25#[derive(Serialize, Deserialize, Eq, PartialEq, Clone, Debug)]
26pub struct BasicQuota {
27    /// The limit of requests.
28    pub limit: usize,
29    /// The period of requests.
30    pub period: Duration,
31}
32impl BasicQuota {
33    /// Create new `BasicQuota`.
34    pub const fn new(limit: usize, period: Duration) -> Self {
35        Self { limit, period }
36    }
37
38    /// Sets the limit of the quota per second.
39    pub const fn per_second(limit: usize) -> Self {
40        Self::new(limit, Duration::seconds(1))
41    }
42    /// Sets the limit of the quota seconds.
43    pub const fn set_seconds(limit: usize, seconds: i64) -> Self {
44        Self::new(limit, Duration::seconds(seconds))
45    }
46
47    /// Sets the limit of the quota per minute.
48    pub const fn per_minute(limit: usize) -> Self {
49        Self::new(limit, Duration::seconds(60))
50    }
51    /// Sets the limit of the quota minutes.
52    pub const fn set_minutes(limit: usize, minutes: i64) -> Self {
53        Self::new(limit, Duration::seconds(60 * minutes))
54    }
55
56    /// Sets the limit of the quota per hour.
57    pub const fn per_hour(limit: usize) -> Self {
58        Self::new(limit, Duration::seconds(3600))
59    }
60    /// Sets the limit of the quota hours.
61    pub const fn set_hours(limit: usize, hours: i64) -> Self {
62        Self::new(limit, Duration::seconds(3600 * hours))
63    }
64}
65
66/// A common used quota has cells field.
67#[non_exhaustive]
68#[derive(Serialize, Deserialize, Eq, PartialEq, Clone, Debug)]
69pub struct CelledQuota {
70    /// The limit of requests.
71    pub limit: usize,
72    /// The period of requests.
73    pub period: Duration,
74    /// The cells of this period splitted to.
75    pub cells: usize,
76}
77impl CelledQuota {
78    /// Create new `CelledQuota`.
79    pub const fn new(limit: usize, cells: usize, period: Duration) -> Self {
80        Self {
81            limit,
82            cells,
83            period,
84        }
85    }
86
87    /// Sets the limit of the quota per second.
88    pub const fn per_second(limit: usize, cells: usize) -> Self {
89        Self::new(limit, cells, Duration::seconds(1))
90    }
91    /// Sets the limit of the quota seconds.
92    pub const fn set_seconds(limit: usize, cells: usize, seconds: i64) -> Self {
93        Self::new(limit, cells, Duration::seconds(seconds))
94    }
95
96    /// Sets the limit of the quota per minute.
97    pub const fn per_minute(limit: usize, cells: usize) -> Self {
98        Self::new(limit, cells, Duration::seconds(60))
99    }
100    /// Sets the limit of the quota minutes.
101    pub const fn set_minutes(limit: usize, cells: usize, minutes: i64) -> Self {
102        Self::new(limit, cells, Duration::seconds(60 * minutes))
103    }
104
105    /// Sets the limit of the quota per hour.
106    pub const fn per_hour(limit: usize, cells: usize) -> Self {
107        Self::new(limit, cells, Duration::seconds(3600))
108    }
109    /// Sets the limit of the quota hours.
110    pub const fn set_hours(limit: usize, cells: usize, hours: i64) -> Self {
111        Self::new(limit, cells, Duration::seconds(3600 * hours))
112    }
113}
114
115impl<Key, T> QuotaGetter<Key> for T
116where
117    Key: Hash + Eq + Send + Sync + 'static,
118    T: Clone + Send + Sync + 'static,
119{
120    type Quota = T;
121    type Error = Infallible;
122
123    async fn get<Q>(&self, _key: &Q) -> Result<Self::Quota, Self::Error>
124    where
125        Key: Borrow<Q>,
126        Q: Hash + Eq + Sync,
127    {
128        Ok(self.clone())
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135
136    #[test]
137    fn test_basic_quota() {
138        let quota = BasicQuota::per_second(10);
139        assert_eq!(quota.limit, 10);
140        assert_eq!(quota.period, Duration::seconds(1));
141
142        let quota = BasicQuota::set_seconds(15, 2);
143        assert_eq!(quota.limit, 15);
144        assert_eq!(quota.period, Duration::seconds(2));
145
146        let quota = BasicQuota::per_minute(10);
147        assert_eq!(quota.limit, 10);
148        assert_eq!(quota.period, Duration::seconds(60));
149
150        let quota = BasicQuota::set_minutes(15, 2);
151        assert_eq!(quota.limit, 15);
152        assert_eq!(quota.period, Duration::seconds(120));
153
154        let quota = BasicQuota::per_hour(10);
155        assert_eq!(quota.limit, 10);
156        assert_eq!(quota.period, Duration::seconds(3600));
157
158        let quota = BasicQuota::set_hours(15, 2);
159        assert_eq!(quota.limit, 15);
160        assert_eq!(quota.period, Duration::seconds(7200));
161    }
162
163    #[test]
164    fn test_celled_quota() {
165        let quota = CelledQuota::per_second(10, 3);
166        assert_eq!(quota.limit, 10);
167        assert_eq!(quota.cells, 3);
168        assert_eq!(quota.period, Duration::seconds(1));
169
170        let quota = CelledQuota::set_seconds(15, 7, 2);
171        assert_eq!(quota.limit, 15);
172        assert_eq!(quota.cells, 7);
173        assert_eq!(quota.period, Duration::seconds(2));
174
175        let quota = CelledQuota::per_minute(10, 9);
176        assert_eq!(quota.limit, 10);
177        assert_eq!(quota.cells, 9);
178        assert_eq!(quota.period, Duration::seconds(60));
179
180        let quota = CelledQuota::set_minutes(15, 7, 2);
181        assert_eq!(quota.limit, 15);
182        assert_eq!(quota.cells, 7);
183        assert_eq!(quota.period, Duration::seconds(120));
184
185        let quota = CelledQuota::per_hour(10, 3);
186        assert_eq!(quota.limit, 10);
187        assert_eq!(quota.cells, 3);
188        assert_eq!(quota.period, Duration::seconds(3600));
189
190        let quota = CelledQuota::set_hours(15, 6, 2);
191        assert_eq!(quota.limit, 15);
192        assert_eq!(quota.cells, 6);
193        assert_eq!(quota.period, Duration::seconds(7200));
194    }
195}
salvo_rate_limiter/quota.rs

salvo_rate_limiter/
quota.rs
