rusty_rl/
bucket.rs

1use std::{collections::HashMap, time::Instant};
2
3use crate::{Selfb, TokenBucket};
4
5/// bucket
6/// ```
7///     let b = TokenBucket::new(600, 10.0);
8/// ```
9///
10impl TokenBucket {
11    /// takes capacity to store max request per second and rates to stores request per second
12    pub fn new(capacity: u64, rates: f64) -> Self {
13        // println!("{}", rates);
14        // token per milli second
15        let rates = rates / 1000.0;
16        // println!("{}", rates);
17        Self {
18            capacity,
19            rates,
20            last_update: HashMap::new(),
21        }
22    }
23    fn bucket(&mut self, key: String) {
24        let now = Instant::now();
25        let (tokens, b) = self.last_update.entry(key).or_insert((0.0, now));
26        let current = now.duration_since(*b).as_millis();
27        // token per current time
28        let btoken = (self.rates * current as f64).ceil();
29        // token = min(old+new,capacity)
30        *tokens = (*tokens + btoken).min(self.capacity as f64);
31        *b = now;
32    }
33    pub fn allow(&mut self, key: impl Into<String> + Copy, token: impl Into<f64>) -> Selfb {
34        let key = key.into();
35        let token = token.into();
36        self.bucket(key.clone());
37        let (tokens, _) = self.last_update.get_mut(&key).unwrap();
38        if token <= *tokens {
39            *tokens -= token;
40            Selfb {
41                is_allowed: true,
42                remaning_count: *tokens as u64,
43            }
44        } else {
45            Selfb {
46                is_allowed: false,
47                remaning_count: *tokens as u64,
48            }
49        }
50    }
51}
rusty_rl/bucket.rs

rusty_rl/
bucket.rs
