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

//! async-rate-limiter implements a [token bucket
//! algorithm](https://en.wikipedia.org/wiki/Token_bucket) that can be used to
//! limit API access frequency.
//!
//! ## Example
//!
//! Update your `Cargo.toml`:
//!
//! ```toml
//! [dependencies]
//! # Change features to ["rt-async-std"] if you are using async-std runtime.
//! async-rate-limiter = { version = "1.39.2", features = ["rt-tokio"] }
//! ```
//!
//! Thanks to Rust’s async functionality, this crate is very simple to use.
//! Just put your function call after [`RateLimiter::acquire()`].`await`, then
//! the function will be called with the specified rate limit.
//!
//! Here is a simple example:
//!
//! ```rust
//! use async_rate_limiter::RateLimiter;
//! use std::time::Duration;
//!
//! #[tokio::main]
//! async fn main() {
//!     let mut rl = RateLimiter::new(3, 5);
//!     
//!     let ok = rl.acquire(None).await;
//!     assert!(ok);
//!     println!("Do something that you want to limit the rate ...");
//!
//!     // acquire with a timeout
//!     let ok = rl.acquire(Some(Duration::from_secs(10))).await;
//!     if ok {
//!         println!("Do something that you want to limit the rate ...");
//!     }
//! }
//!
//! ```
//!
//! async-rate-limiter can support different async runtimes, tokio & async-std
//! are supported currently. You can use features to switch async runtimes.

mod token_bucket;
pub use token_bucket::TokenBucketRateLimiter as RateLimiter;

mod rt;

#[cfg(test)]
mod tests {
    use std::time::{Duration, Instant};

    use super::*;

    #[tokio::test]
    #[cfg(any(feature = "rt-tokio", feature = "rt-async-std"))]
    async fn test_rate_limit() {
        let mut rl = RateLimiter::new(3, 5);

        let start = Instant::now();
        let res = rl.acquire(None).await;
        assert!(res);
        let res = rl.acquire(None).await;
        assert!(res);
        let res = rl.acquire(None).await;
        assert!(res);
        assert!(start.elapsed() < Duration::from_millis(100));

        let res = rl.acquire(Some(Duration::from_secs(1))).await;
        assert!(res);
        let res = rl.acquire(None).await;
        assert!(res);
        let res = rl.acquire(None).await;
        assert!(res);
        assert!(start.elapsed() >= Duration::from_secs(1));
        assert!(start.elapsed() < Duration::from_secs(1) + Duration::from_millis(100));

        let res = rl.acquire(Some(Duration::from_millis(10))).await;
        assert!(!res);
        assert!(start.elapsed() < Duration::from_secs(1) + Duration::from_millis(100));
    }
}