governor 0.10.4

A rate-limiting implementation in 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

use governor::{
    clock::{self, FakeRelativeClock},
    middleware::{RateLimitingMiddleware, StateInformationMiddleware, StateSnapshot},
    Quota, RateLimiter,
};
use nonzero_ext::nonzero;

#[derive(Debug)]
struct MyMW;

impl RateLimitingMiddleware<<FakeRelativeClock as clock::Clock>::Instant> for MyMW {
    type PositiveOutcome = u16;

    fn allow<K>(_key: &K, _state: impl Into<StateSnapshot>) -> Self::PositiveOutcome {
        666
    }

    type NegativeOutcome = ();

    fn disallow<K>(
        _key: &K,
        _limiter: impl Into<StateSnapshot>,
        _start_time: <FakeRelativeClock as clock::Clock>::Instant,
    ) -> Self::NegativeOutcome {
    }
}

#[test]
fn changes_allowed_type() {
    let clock = FakeRelativeClock::default();
    let lim = RateLimiter::direct_with_clock(Quota::per_hour(nonzero!(1_u32)), clock.clone())
        .with_middleware::<MyMW>();
    assert_eq!(Ok(666), lim.check());
    assert_eq!(Err(()), lim.check());
}

#[test]
fn state_information() {
    let clock = FakeRelativeClock::default();
    let lim = RateLimiter::direct_with_clock(Quota::per_second(nonzero!(4u32)), clock.clone())
        .with_middleware::<StateInformationMiddleware>();
    assert_eq!(
        Ok(3),
        lim.check()
            .map(|outcome| outcome.remaining_burst_capacity())
    );
    assert_eq!(
        Ok(2),
        lim.check()
            .map(|outcome| outcome.remaining_burst_capacity())
    );
    assert_eq!(
        Ok(1),
        lim.check()
            .map(|outcome| outcome.remaining_burst_capacity())
    );
    assert_eq!(
        Ok(0),
        lim.check()
            .map(|outcome| outcome.remaining_burst_capacity())
    );
    assert!(lim.check().is_err());
}

#[test]
#[cfg(feature = "std")]
fn mymw_derives() {
    assert_eq!(format!("{MyMW:?}"), "MyMW");
}

#[test]
fn state_snapshot_tracks_quota_accurately() {
    use governor::middleware::StateInformationMiddleware;
    use governor::{Quota, RateLimiter};
    use std::num::NonZeroU32;
    use std::time::Duration;

    let burst_size = NonZeroU32::new(2).unwrap();
    let period = Duration::from_millis(90);
    let quota = Quota::with_period(period).unwrap().allow_burst(burst_size);

    let clock = FakeRelativeClock::default();

    // First test
    let lim = RateLimiter::direct_with_clock(quota, clock.clone())
        .with_middleware::<StateInformationMiddleware>();

    assert_eq!(lim.check().unwrap().remaining_burst_capacity(), 1);
    assert_eq!(lim.check().unwrap().remaining_burst_capacity(), 0);
    assert_eq!(lim.check().map_err(|_| ()), Err(()), "should rate limit");

    clock.advance(Duration::from_secs(120));
    assert_eq!(lim.check().map(|s| s.remaining_burst_capacity()), Ok(1));
    assert_eq!(lim.check().map(|s| s.remaining_burst_capacity()), Ok(0));
    assert_eq!(lim.check().map_err(|_| ()), Err(()), "should rate limit");
}

#[test]
#[cfg(feature = "std")]
fn state_snapshot_tracks_quota_accurately_with_real_clock() {
    use governor::middleware::StateInformationMiddleware;
    use governor::{Quota, RateLimiter};
    use std::num::NonZeroU32;
    use std::time::Duration;

    let burst_size = NonZeroU32::new(2).unwrap();
    let period = Duration::from_millis(90);
    let quota = Quota::with_period(period).unwrap().allow_burst(burst_size);
    let lim = RateLimiter::direct(quota).with_middleware::<StateInformationMiddleware>();

    assert_eq!(lim.check().unwrap().remaining_burst_capacity(), 1);
    assert_eq!(lim.check().unwrap().remaining_burst_capacity(), 0);
    assert_eq!(lim.check().map_err(|_| ()), Err(()), "should rate limit");
}