1use std::{collections::VecDeque, fmt, time::Duration};
7
8use parking_lot::Mutex;
9use tokio::time::{Instant, sleep};
10
11use crate::Error;
12
13const DEFAULT_HISTORY_REQUESTS: usize = 50;
14const DEFAULT_HISTORY_WINDOW: Duration = Duration::from_secs(30);
15const DEFAULT_GENERAL_REQUESTS: usize = 200;
16const DEFAULT_GENERAL_WINDOW: Duration = Duration::from_mins(1);
17const MAX_INITIAL_CAPACITY: usize = 1_024;
18
19#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
21#[non_exhaustive]
22pub enum RateLimitKind {
23 History,
25 General,
27}
28
29impl fmt::Display for RateLimitKind {
30 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
31 match self {
32 Self::History => formatter.write_str("history"),
33 Self::General => formatter.write_str("general"),
34 }
35 }
36}
37
38#[derive(Clone, Copy, Debug, Eq, PartialEq)]
40#[non_exhaustive]
41pub struct RateLimit {
42 max_requests: usize,
43 window: Duration,
44}
45
46impl RateLimit {
47 pub fn new(max_requests: usize, window: Duration) -> Result<Self, Error> {
54 if max_requests == 0
55 || window.is_zero()
56 || std::time::Instant::now().checked_add(window).is_none()
57 {
58 return Err(Error::Configuration(
59 "rate limits require a positive request count and representable window".to_owned(),
60 ));
61 }
62 Ok(Self {
63 max_requests,
64 window,
65 })
66 }
67
68 const fn trusted(max_requests: usize, window: Duration) -> Self {
69 Self {
70 max_requests,
71 window,
72 }
73 }
74
75 #[must_use]
77 pub const fn max_requests(self) -> usize {
78 self.max_requests
79 }
80
81 #[must_use]
83 pub const fn window(self) -> Duration {
84 self.window
85 }
86}
87
88#[derive(Clone, Copy, Debug, Eq, PartialEq)]
93#[non_exhaustive]
94pub struct RateLimitConfig {
95 history: RateLimit,
96 general: RateLimit,
97}
98
99impl RateLimitConfig {
100 #[must_use]
102 pub const fn new(history: RateLimit, general: RateLimit) -> Self {
103 Self { history, general }
104 }
105
106 #[must_use]
108 pub const fn history(self) -> RateLimit {
109 self.history
110 }
111
112 #[must_use]
114 pub const fn general(self) -> RateLimit {
115 self.general
116 }
117
118 pub(crate) const fn limit(self, kind: RateLimitKind) -> RateLimit {
119 match kind {
120 RateLimitKind::History => self.history,
121 RateLimitKind::General => self.general,
122 }
123 }
124}
125
126impl Default for RateLimitConfig {
127 fn default() -> Self {
128 Self {
129 history: RateLimit::trusted(DEFAULT_HISTORY_REQUESTS, DEFAULT_HISTORY_WINDOW),
130 general: RateLimit::trusted(DEFAULT_GENERAL_REQUESTS, DEFAULT_GENERAL_WINDOW),
131 }
132 }
133}
134
135pub(crate) struct RateLimits {
136 config: Option<RateLimitConfig>,
137 history: Option<WindowLimiter>,
138 general: Option<WindowLimiter>,
139}
140
141impl RateLimits {
142 pub(crate) fn new(config: Option<RateLimitConfig>) -> Self {
143 let history = config.map(|limits| WindowLimiter::new(limits.history));
144 let general = config.map(|limits| WindowLimiter::new(limits.general));
145 Self {
146 config,
147 history,
148 general,
149 }
150 }
151
152 pub(crate) const fn config(&self) -> Option<RateLimitConfig> {
153 self.config
154 }
155
156 pub(crate) fn limit(&self, kind: RateLimitKind) -> RateLimit {
157 self.config.unwrap_or_default().limit(kind)
158 }
159
160 pub(crate) async fn wait(&self, kind: RateLimitKind) {
161 let Some(limiter) = self.limiter(kind) else {
162 return;
163 };
164 loop {
165 match limiter.try_acquire(Instant::now()) {
166 Ok(()) => return,
167 Err(retry_after) => sleep(retry_after).await,
168 }
169 }
170 }
171
172 pub(crate) fn try_acquire(&self, kind: RateLimitKind) -> Result<(), Duration> {
173 self.limiter(kind)
174 .map_or(Ok(()), |limiter| limiter.try_acquire(Instant::now()))
175 }
176
177 pub(crate) fn cool_down(&self, kind: RateLimitKind, retry_after: Duration) {
178 if let Some(limiter) = self.limiter(kind) {
179 limiter.cool_down(Instant::now(), retry_after);
180 }
181 }
182
183 fn limiter(&self, kind: RateLimitKind) -> Option<&WindowLimiter> {
184 match kind {
185 RateLimitKind::History => self.history.as_ref(),
186 RateLimitKind::General => self.general.as_ref(),
187 }
188 }
189}
190
191struct WindowLimiter {
192 limit: RateLimit,
193 state: Mutex<WindowState>,
194}
195
196impl WindowLimiter {
197 fn new(limit: RateLimit) -> Self {
198 Self {
199 limit,
200 state: Mutex::new(WindowState::with_capacity(
201 limit.max_requests.min(MAX_INITIAL_CAPACITY),
202 )),
203 }
204 }
205
206 fn try_acquire(&self, now: Instant) -> Result<(), Duration> {
207 let mut state = self.state.lock();
208 state.expire(now, self.limit.window);
209 let retry_after = state.retry_after(now, self.limit);
210 if retry_after.is_zero() {
211 state.admitted.push_back(now);
212 Ok(())
213 } else {
214 Err(retry_after)
215 }
216 }
217
218 fn cool_down(&self, now: Instant, retry_after: Duration) {
219 let Some(deadline) = now.checked_add(retry_after) else {
220 return;
221 };
222 let mut state = self.state.lock();
223 if state
224 .cooldown_until
225 .is_none_or(|current| deadline > current)
226 {
227 state.cooldown_until = Some(deadline);
228 }
229 }
230}
231
232struct WindowState {
233 admitted: VecDeque<Instant>,
234 cooldown_until: Option<Instant>,
235}
236
237impl WindowState {
238 fn with_capacity(capacity: usize) -> Self {
239 Self {
240 admitted: VecDeque::with_capacity(capacity),
241 cooldown_until: None,
242 }
243 }
244
245 fn expire(&mut self, now: Instant, window: Duration) {
246 while self
247 .admitted
248 .front()
249 .is_some_and(|admitted| now.saturating_duration_since(*admitted) >= window)
250 {
251 self.admitted.pop_front();
252 }
253 if self.cooldown_until.is_some_and(|deadline| deadline <= now) {
254 self.cooldown_until = None;
255 }
256 }
257
258 fn retry_after(&self, now: Instant, limit: RateLimit) -> Duration {
259 let cooldown = self.cooldown_until.map_or(Duration::ZERO, |deadline| {
260 deadline.saturating_duration_since(now)
261 });
262 let capacity = if self.admitted.len() < limit.max_requests {
263 Duration::ZERO
264 } else {
265 self.admitted.front().map_or(limit.window, |oldest| {
266 oldest
267 .checked_add(limit.window)
268 .map_or(limit.window, |deadline| {
269 deadline.saturating_duration_since(now)
270 })
271 })
272 };
273 cooldown.max(capacity)
274 }
275}
276
277#[cfg(test)]
278mod tests {
279 use std::sync::Arc;
280
281 use tokio::time::{Instant, advance};
282
283 use super::*;
284
285 fn limits(requests: usize, window: Duration) -> Arc<RateLimits> {
286 let limit = RateLimit::new(requests, window)
287 .unwrap_or_else(|error| panic!("test rate limit must be valid: {error}"));
288 Arc::new(RateLimits::new(Some(RateLimitConfig::new(limit, limit))))
289 }
290
291 #[tokio::test(start_paused = true)]
292 async fn wait_resumes_when_the_rolling_window_releases_capacity() {
293 let window = Duration::from_secs(10);
294 let limits = limits(1, window);
295 assert_eq!(limits.try_acquire(RateLimitKind::General), Ok(()));
296
297 let started = Instant::now();
298 let waiting_limits = Arc::clone(&limits);
299 let waiter = tokio::spawn(async move {
300 waiting_limits.wait(RateLimitKind::General).await;
301 Instant::now()
302 });
303 tokio::task::yield_now().await;
304 assert!(!waiter.is_finished());
305
306 advance(window).await;
307 let admitted = waiter
308 .await
309 .unwrap_or_else(|error| panic!("waiting task must finish: {error}"));
310 assert_eq!(admitted.saturating_duration_since(started), window);
311 }
312
313 #[tokio::test(start_paused = true)]
314 async fn provider_cooldown_blocks_admission_until_its_deadline() {
315 let cooldown = Duration::from_secs(15);
316 let limits = limits(2, Duration::from_secs(30));
317 limits.cool_down(RateLimitKind::General, cooldown);
318
319 assert_eq!(limits.try_acquire(RateLimitKind::General), Err(cooldown));
320 advance(cooldown).await;
321 assert_eq!(limits.try_acquire(RateLimitKind::General), Ok(()));
322 }
323
324 #[tokio::test(start_paused = true)]
325 async fn rolling_window_expires_staggered_attempts_individually() {
326 let window = Duration::from_secs(10);
327 let limits = limits(2, window);
328 assert_eq!(limits.try_acquire(RateLimitKind::General), Ok(()));
329 advance(Duration::from_secs(5)).await;
330 assert_eq!(limits.try_acquire(RateLimitKind::General), Ok(()));
331
332 advance(Duration::from_secs(5)).await;
333 assert_eq!(limits.try_acquire(RateLimitKind::General), Ok(()));
334 assert_eq!(
335 limits.try_acquire(RateLimitKind::General),
336 Err(Duration::from_secs(5))
337 );
338 }
339
340 #[tokio::test(start_paused = true)]
341 async fn cancelled_waiter_leaves_no_ghost_reservation() {
342 let window = Duration::from_secs(10);
343 let limits = limits(1, window);
344 assert_eq!(limits.try_acquire(RateLimitKind::General), Ok(()));
345 let waiting_limits = Arc::clone(&limits);
346 let waiter = tokio::spawn(async move {
347 waiting_limits.wait(RateLimitKind::General).await;
348 });
349 tokio::task::yield_now().await;
350 waiter.abort();
351
352 advance(window).await;
353 assert_eq!(limits.try_acquire(RateLimitKind::General), Ok(()));
354 }
355
356 #[tokio::test(start_paused = true)]
357 async fn default_history_budget_reaches_its_limit_before_throttling() {
358 let limits = Arc::new(RateLimits::new(Some(RateLimitConfig::default())));
359 for request_number in 1..=50 {
360 assert_eq!(
361 limits.try_acquire(RateLimitKind::History),
362 Ok(()),
363 "history request {request_number} must be admitted"
364 );
365 }
366 assert_eq!(
367 limits.try_acquire(RateLimitKind::History),
368 Err(Duration::from_secs(30))
369 );
370 }
371
372 #[tokio::test(start_paused = true)]
373 async fn default_general_budget_reaches_its_limit_before_throttling() {
374 let limits = Arc::new(RateLimits::new(Some(RateLimitConfig::default())));
375 for request_number in 1..=200 {
376 assert_eq!(
377 limits.try_acquire(RateLimitKind::General),
378 Ok(()),
379 "general request {request_number} must be admitted"
380 );
381 }
382 assert_eq!(
383 limits.try_acquire(RateLimitKind::General),
384 Err(Duration::from_mins(1))
385 );
386 }
387
388 #[tokio::test(start_paused = true)]
389 async fn later_provider_cooldown_cannot_be_shortened() {
390 let limits = limits(200, Duration::from_mins(1));
391 limits.cool_down(RateLimitKind::General, Duration::from_secs(30));
392 limits.cool_down(RateLimitKind::General, Duration::from_secs(5));
393 assert_eq!(
394 limits.try_acquire(RateLimitKind::General),
395 Err(Duration::from_secs(30))
396 );
397 }
398
399 #[tokio::test(start_paused = true)]
400 async fn admission_uses_the_later_of_window_and_provider_cooldown() {
401 let limits = limits(1, Duration::from_secs(20));
402 assert_eq!(limits.try_acquire(RateLimitKind::General), Ok(()));
403 limits.cool_down(RateLimitKind::General, Duration::from_secs(5));
404 assert_eq!(
405 limits.try_acquire(RateLimitKind::General),
406 Err(Duration::from_secs(20))
407 );
408 }
409
410 #[test]
411 fn history_and_general_budgets_are_independent() {
412 let limits = limits(1, Duration::from_secs(30));
413 assert_eq!(limits.try_acquire(RateLimitKind::General), Ok(()));
414 assert_eq!(limits.try_acquire(RateLimitKind::History), Ok(()));
415 assert!(limits.try_acquire(RateLimitKind::General).is_err());
416 assert!(limits.try_acquire(RateLimitKind::History).is_err());
417 }
418
419 #[test]
420 fn public_defaults_match_the_provider_contract() {
421 let config = RateLimitConfig::default();
422 assert_eq!(config.history().max_requests(), 50);
423 assert_eq!(config.history().window(), Duration::from_secs(30));
424 assert_eq!(config.general().max_requests(), 200);
425 assert_eq!(config.general().window(), Duration::from_mins(1));
426 }
427
428 #[test]
429 fn public_rate_limit_rejects_zero_values() {
430 assert!(RateLimit::new(0, Duration::from_secs(1)).is_err());
431 assert!(RateLimit::new(1, Duration::ZERO).is_err());
432 }
433}