drm_core/exchange/
rate_limit.rs

1use std::time::{Duration, Instant};
2use tokio::time::sleep;
3
4pub struct RateLimiter {
5    last_request: Instant,
6    min_interval: Duration,
7}
8
9impl RateLimiter {
10    pub fn new(requests_per_second: u32) -> Self {
11        let min_interval = if requests_per_second > 0 {
12            Duration::from_secs_f64(1.0 / requests_per_second as f64)
13        } else {
14            Duration::ZERO
15        };
16
17        Self {
18            last_request: Instant::now() - min_interval,
19            min_interval,
20        }
21    }
22
23    pub async fn wait(&mut self) {
24        let elapsed = self.last_request.elapsed();
25        if elapsed < self.min_interval {
26            let wait_time = self.min_interval - elapsed;
27            sleep(wait_time).await;
28        }
29        self.last_request = Instant::now();
30    }
31}
32
33pub async fn retry_with_backoff<T, E, F, Fut>(
34    max_attempts: u32,
35    initial_delay: Duration,
36    mut f: F,
37) -> Result<T, E>
38where
39    F: FnMut() -> Fut,
40    Fut: std::future::Future<Output = Result<T, E>>,
41{
42    let mut delay = initial_delay;
43
44    for attempt in 0..max_attempts {
45        match f().await {
46            Ok(result) => return Ok(result),
47            Err(_) if attempt + 1 < max_attempts => {
48                sleep(delay).await;
49                delay *= 2;
50                continue;
51            }
52            Err(e) => return Err(e),
53        }
54    }
55
56    unreachable!()
57}
58
59#[cfg(test)]
60mod tests {
61    use super::*;
62
63    #[tokio::test]
64    async fn test_rate_limiter_respects_interval() {
65        let mut limiter = RateLimiter::new(10);
66        let start = Instant::now();
67
68        limiter.wait().await;
69        limiter.wait().await;
70
71        let elapsed = start.elapsed();
72        assert!(elapsed >= Duration::from_millis(90));
73    }
74}
drm_core/exchange/rate_limit.rs

drm_core/exchange/
rate_limit.rs
