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
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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168

#![cfg(all(feature = "std", feature = "dashmap"))]

use governor::{
    clock::{Clock, FakeRelativeClock},
    Quota, RateLimiter,
};
use governor::{middleware::NoOpMiddleware, state::keyed::DashMapStateStore};
use nonzero_ext::nonzero;
use std::hash::Hash;
use std::time::Duration;

const KEYS: &[u32] = &[1u32, 2u32];

#[test]
fn accepts_first_cell() {
    let clock = FakeRelativeClock::default();
    let lb = RateLimiter::dashmap_with_clock(Quota::per_second(nonzero!(5u32)), clock.clone());
    for key in KEYS {
        assert_eq!(Ok(()), lb.check_key(&key), "key {key:?}");
    }
}

#[test]
fn rejects_too_many() {
    let clock = FakeRelativeClock::default();
    let lb = RateLimiter::dashmap_with_clock(Quota::per_second(nonzero!(2u32)), clock.clone());
    let ms = Duration::from_millis(1);

    for key in KEYS {
        // use up our burst capacity (2 in the first second):
        assert_eq!(Ok(()), lb.check_key(key), "Now: {:?}", clock.now());
        clock.advance(ms);
        assert_eq!(Ok(()), lb.check_key(key), "Now: {:?}", clock.now());

        clock.advance(ms);
        assert_ne!(Ok(()), lb.check_key(key), "Now: {:?}", clock.now());

        // should be ok again in 1s:
        clock.advance(ms * 1000);
        assert_eq!(Ok(()), lb.check_key(key), "Now: {:?}", clock.now());
        clock.advance(ms);
        assert_eq!(Ok(()), lb.check_key(key));

        clock.advance(ms);
        assert_ne!(Ok(()), lb.check_key(key), "{lb:?}");
    }
}

fn retained_keys<T: Clone + Hash + Eq + Copy + Ord>(
    keys: &[T],
    limiter: RateLimiter<
        T,
        DashMapStateStore<T>,
        FakeRelativeClock,
        NoOpMiddleware<<FakeRelativeClock as Clock>::Instant>,
    >,
) -> Vec<T> {
    let state = limiter.into_state_store();
    let mut keys: Vec<T> = keys
        .iter()
        .copied()
        .filter(|&k| state.contains_key(&k))
        .collect();
    keys.sort();
    keys
}

#[test]
fn expiration() {
    let clock = FakeRelativeClock::default();
    let ms = Duration::from_millis(1);

    let make_bucket = || {
        let lim = RateLimiter::dashmap_with_clock(Quota::per_second(nonzero!(1u32)), clock.clone());
        lim.check_key(&"foo").unwrap();
        clock.advance(ms * 200);
        lim.check_key(&"bar").unwrap();
        clock.advance(ms * 600);
        lim.check_key(&"baz").unwrap();
        lim
    };
    let keys = &["bar", "baz", "foo"];

    // clean up all keys that are indistinguishable from unoccupied keys:
    let lim_shrunk = make_bucket();
    lim_shrunk.retain_recent();
    assert_eq!(retained_keys(keys, lim_shrunk), keys);

    let lim_later = make_bucket();
    clock.advance(ms * 1200);
    lim_later.retain_recent();
    assert_eq!(retained_keys(keys, lim_later), vec!["bar", "baz"]);

    let lim_later = make_bucket();
    clock.advance(ms * (1200 + 200));
    lim_later.retain_recent();
    assert_eq!(retained_keys(keys, lim_later), vec!["baz"]);

    let lim_later = make_bucket();
    clock.advance(ms * (1200 + 200 + 600));
    lim_later.retain_recent();
    assert_eq!(retained_keys(keys, lim_later), Vec::<&str>::new());
}

#[test]
fn actual_threadsafety() {
    use crossbeam;

    let clock = FakeRelativeClock::default();
    let lim = RateLimiter::dashmap_with_clock(Quota::per_second(nonzero!(20u32)), clock.clone());
    let ms = Duration::from_millis(1);

    for key in KEYS {
        crossbeam::scope(|scope| {
            for _i in 0..20 {
                scope.spawn(|_| {
                    assert_eq!(Ok(()), lim.check_key(key));
                });
            }
        })
        .unwrap();

        clock.advance(ms * 2);
        assert_ne!(Ok(()), lim.check_key(key));
        clock.advance(ms * 998);
        assert_eq!(Ok(()), lim.check_key(key));
    }
}

#[test]
fn dashmap_length() {
    let lim = RateLimiter::dashmap(Quota::per_second(nonzero!(1u32)));
    assert_eq!(lim.len(), 0);
    assert!(lim.is_empty());

    lim.check_key(&"foo").unwrap();
    assert_eq!(lim.len(), 1);
    assert!(!lim.is_empty(),);

    lim.check_key(&"bar").unwrap();
    assert_eq!(lim.len(), 2);
    assert!(!lim.is_empty());

    lim.check_key(&"baz").unwrap();
    assert_eq!(lim.len(), 3);
    assert!(!lim.is_empty());
}

#[test]
fn dashmap_shrink_to_fit() {
    let clock = FakeRelativeClock::default();
    // a steady rate of 3ms between elements:
    let lim = RateLimiter::dashmap_with_clock(Quota::per_second(nonzero!(20u32)), clock.clone());
    let ms = Duration::from_millis(1);

    assert_eq!(
        lim.check_key_n(&"long-lived".to_string(), nonzero!(10_u32)),
        Ok(Ok(()))
    );
    assert_eq!(lim.check_key(&"short-lived".to_string()), Ok(()));

    // Move the clock forward far enough that the short-lived key gets dropped:
    clock.advance(ms * 300);
    lim.retain_recent();
    lim.shrink_to_fit();

    assert_eq!(lim.len(), 1);
}