muffy 0.4.3

The static website validator
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

use core::{
    sync::atomic::{AtomicU64, Ordering},
    time::Duration,
};
use tokio::time::{Instant, sleep};

// TODO Optimize the ordering.

/// A token bucket rate limiter.
pub struct RateLimiter {
    token_count: AtomicU64,
    window_count: AtomicU64,
    time: Instant,
    supply: u64,
    window: Duration,
}

impl RateLimiter {
    /// Creates a rate limiter.
    pub fn new(supply: u64, window: Duration) -> Self {
        Self {
            token_count: AtomicU64::new(supply),
            window_count: Default::default(),
            time: Instant::now(),
            supply,
            window,
        }
    }

    /// Runs a future with rate limiting.
    pub async fn run<T>(&self, future: impl Future<Output = T>) -> T {
        while {
            self.add_supply();

            self.token_count
                .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |count| {
                    if count == 0 { None } else { Some(count - 1) }
                })
                .is_err()
        } {
            sleep(
                self.window
                    - Duration::from_nanos(
                        (self.time.elapsed().as_nanos() % self.window.as_nanos()) as u64,
                    ),
            )
            .await;
        }

        future.await
    }

    fn add_supply(&self) {
        let old = self.window_count.load(Ordering::Relaxed);
        let new = (self.time.elapsed().as_nanos() / self.window.as_nanos()) as _;

        if new > old
            && self
                .window_count
                .compare_exchange(old, new, Ordering::SeqCst, Ordering::SeqCst)
                .is_ok()
        {
            self.token_count.store(self.supply, Ordering::SeqCst);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use alloc::sync::Arc;
    use futures::future::join_all;
    use tokio::{spawn, time::advance};

    #[tokio::test]
    async fn run_once() {
        let limiter = RateLimiter::new(1, Duration::from_secs(1));

        assert_eq!(limiter.run(async { 42 }).await, 42);
    }

    #[tokio::test(start_paused = true)]
    async fn run_many_times() {
        const REQUEST_COUNT: u64 = 10_000;
        const SUPPLY: u64 = 1000;
        const WINDOW: Duration = Duration::from_millis(100);

        let time = Instant::now();
        let limiter = Arc::new(RateLimiter::new(SUPPLY, WINDOW));
        let mut futures = vec![];

        for _ in 0..REQUEST_COUNT {
            let limiter = limiter.clone();

            futures.push(spawn(async move { limiter.run(async { 42 }).await }));
        }

        join_all(futures).await;

        let duration = WINDOW * (REQUEST_COUNT / SUPPLY - 1) as _;

        assert!(time.elapsed() >= duration);
        assert!(time.elapsed() < duration.mul_f64(1.5));
    }

    #[tokio::test(start_paused = true)]
    async fn run_after_many_windows() {
        const WINDOW: Duration = Duration::from_micros(1);

        let limiter = RateLimiter::new(1, WINDOW);

        limiter.run(async {}).await;

        advance(WINDOW * u32::MAX).await;

        limiter.run(async {}).await;
        limiter.run(async {}).await;
    }
}