rusty-rl 0.1.0

simple ratelimit helper
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

use std::{collections::HashMap, time::Instant};

use crate::{Selfb, TokenBucket};

/// bucket
/// ```
///     let b = TokenBucket::new(600, 10.0);
/// ```
///
impl TokenBucket {
    /// takes capacity to store max request per second and rates to stores request per second
    pub fn new(capacity: u64, rates: f64) -> Self {
        // println!("{}", rates);
        // token per milli second
        let rates = rates / 1000.0;
        // println!("{}", rates);
        Self {
            capacity,
            rates,
            last_update: HashMap::new(),
        }
    }
    fn bucket(&mut self, key: String) {
        let now = Instant::now();
        let (tokens, b) = self.last_update.entry(key).or_insert((0.0, now));
        let current = now.duration_since(*b).as_millis();
        // token per current time
        let btoken = (self.rates * current as f64).ceil();
        // token = min(old+new,capacity)
        *tokens = (*tokens + btoken).min(self.capacity as f64);
        *b = now;
    }
    pub fn allow(&mut self, key: impl Into<String> + Copy, token: impl Into<f64>) -> Selfb {
        let key = key.into();
        let token = token.into();
        self.bucket(key.clone());
        let (tokens, _) = self.last_update.get_mut(&key).unwrap();
        if token <= *tokens {
            *tokens -= token;
            Selfb {
                is_allowed: true,
                remaning_count: *tokens as u64,
            }
        } else {
            Selfb {
                is_allowed: false,
                remaning_count: *tokens as u64,
            }
        }
    }
}