r8limit 0.2.2

Dead simple rate limiter for Rust
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
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145

//! An extremely simple window rate limiter.
//!
//! # Usage
//! In your `Cargo.toml`:
//! ```toml
//! [dependencies]
//! r8limit = "0.1"
//! ```
//!
//! In your code:
//! ```
//! use std::time::Duration;
//!
//! // Allow 3 attempts every 5 seconds
//! let mut limiter = r8limit::RateLimiter::new(3, Duration::from_secs(5));
//!
//! println!("{}", limiter.attempt()); // true
//! println!("{}", limiter.attempt()); // true
//! println!("{}", limiter.attempt()); // true
//! println!("{}", limiter.attempt()); // false
//! ```

use std::time::{Duration, Instant};
use crate::RefillPolicy::Full;

pub struct RateLimiter {
    max_executions_per_window: f64,
    executions_left_in_window: f64,
    window_start_time: Instant,
    window_duration: Duration,
    refill_policy: RefillPolicy
}

impl RateLimiter {
    pub fn new(max_executions_per_internal: u64, interval: Duration) -> RateLimiter {
        RateLimiter {
            max_executions_per_window: max_executions_per_internal as f64,
            executions_left_in_window: max_executions_per_internal as f64 ,
            window_start_time: Instant::now(),
            window_duration: interval,
            refill_policy: Full
        }
    }

    pub fn with_refill_policy(mut self, refill_policy: RefillPolicy) -> RateLimiter {
        self.refill_policy = refill_policy;
        self
    }

    pub fn attempt(&mut self) -> bool {
        match &self.refill_policy {
            RefillPolicy::Full => {
                if self.window_start_time.elapsed() > self.window_duration {
                    self.window_start_time = Instant::now();
                    self.executions_left_in_window = self.max_executions_per_window;
                }
            },
            RefillPolicy::Gradual => {
                if self.executions_left_in_window != self.max_executions_per_window {
                    if self.window_start_time.elapsed() > self.window_duration {
                        // Been a while since the rate limiter was triggered, so we'll reset everything
                        self.window_start_time = Instant::now();
                        self.executions_left_in_window = self.max_executions_per_window;
                    } else {
                        // Need to move window and update the executions_left_in_window
                        let percentage_done = self.window_start_time.elapsed().as_secs_f64()/self.window_duration.as_secs_f64();
                        let executions_due = percentage_done * self.max_executions_per_window;
                        self.window_start_time = Instant::now();
                        self.executions_left_in_window += executions_due;
                    }
                }
            }
        }
        if self.executions_left_in_window < 1.0 {
            return false;
        }
        self.executions_left_in_window -= 1.0;
        return true;
    }
}

pub enum RefillPolicy {
    Full,
    Gradual
}

#[cfg(test)]
mod tests {
    use std::time::Duration;
    use std::thread::sleep;
    use crate::{RateLimiter, RefillPolicy};

    #[test]
    fn rate_limiter_full() {
        let mut limiter = RateLimiter::new(3, Duration::from_secs(1));
        assert_eq!(limiter.attempt(), true); // executions remaining for window: 2
        assert_eq!(limiter.attempt(), true); // executions remaining for window: 1
        assert_eq!(limiter.attempt(), true); // executions remaining for window: 0
        assert_eq!(limiter.attempt(), false); // executions remaining for window: 0
        // Remember that the interval is set to 1s
        sleep(Duration::from_millis(500)); // executions remaining for window: 0
        // As you can see, even though half of the interval has passed, there are still 0 executions available.
        assert_eq!(limiter.attempt(), false);
        // That is what the default refill policy, RefillPolicy::Full, does.
        // We'll sleep for the remainder of the window
        sleep(Duration::from_millis(500)); // executions remaining for window: 3
        assert_eq!(limiter.attempt(), true); // executions remaining for window: 2
        assert_eq!(limiter.attempt(), true); // executions remaining for window: 1
        assert_eq!(limiter.attempt(), true); // executions remaining for window: 0
        assert_eq!(limiter.attempt(), false); // executions remaining for window: 0
        assert_eq!(limiter.attempt(), false); // executions remaining for window: 0
    }

    #[test]
    fn rate_limiter_gradual() {
        let mut limiter = RateLimiter::new(4, Duration::from_millis(500)).with_refill_policy(RefillPolicy::Gradual);
        assert_eq!(limiter.attempt(), true);
        assert_eq!(limiter.attempt(), true);
        assert_eq!(limiter.attempt(), true);
        assert_eq!(limiter.attempt(), true);
        assert_eq!(limiter.attempt(), false);
        sleep(Duration::from_millis(250));
        assert_eq!(limiter.attempt(), true);
        assert_eq!(limiter.attempt(), true);
        assert_eq!(limiter.attempt(), false);
        sleep(Duration::from_millis(500));
        assert_eq!(limiter.attempt(), true);
        assert_eq!(limiter.attempt(), true);
        assert_eq!(limiter.attempt(), true);
        assert_eq!(limiter.attempt(), true);
        assert_eq!(limiter.attempt(), false);
        sleep(Duration::from_millis(125));
        assert_eq!(limiter.attempt(), true);
        assert_eq!(limiter.attempt(), false);
        sleep(Duration::from_millis(125));
        assert_eq!(limiter.attempt(), true);
        assert_eq!(limiter.attempt(), false);
        sleep(Duration::from_millis(750));
        assert_eq!(limiter.attempt(), true);
        assert_eq!(limiter.attempt(), true);
        assert_eq!(limiter.attempt(), true);
        assert_eq!(limiter.attempt(), true);
        assert_eq!(limiter.attempt(), false);
    }
}