leaky-bucket 1.1.2

A token-based rate limiter based on the leaky bucket algorithm.
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

use std::pin::pin;
use std::sync::Arc;
use tokio::time::{self, Duration, Instant};

use leaky_bucket::RateLimiter;
use tokio::task::JoinSet;

#[tokio::test(start_paused = true)]
async fn test_drop() -> anyhow::Result<()> {
    let limiter = Arc::new(
        RateLimiter::builder()
            .initial(0)
            .refill(10)
            .interval(Duration::from_millis(50))
            .max(100)
            .build(),
    );

    let limiter = limiter.clone();

    let mut task = pin!(Some(limiter.acquire(10000)));

    let mut join_set = JoinSet::new();

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

        join_set.spawn(async move {
            limiter.acquire(10).await;
        });
    }

    tokio::select! {
        _ = time::sleep(Duration::from_millis(1000)) => {
            // Drop the task
            task.set(None);
        }
        _ = Option::as_pin_mut(task.as_mut()).unwrap() => {
        }
        // Should never complete, because we have a giant task waiting.
        _ = join_set.join_next() => {
        }
    }

    let mut released = 0;

    let start = Instant::now();

    while join_set.join_next().await.is_some() {
        released += 1;
    }

    assert!(Instant::now().duration_since(start).as_millis() <= 10);
    assert_eq!(released, 10);
    Ok(())
}