rusty_rl/
fixed_bucket.rs

1use std::{
2    collections::HashMap,
3    iter::Inspect,
4    time::{Duration, Instant},
5};
6
7use crate::{FixedBucket, Selfb};
8
9/// fixed bucket
10///
11/// ```
12///     let b = FixedBucket::new(600, 10).bucket(90, 10.0)
13/// ```
14impl FixedBucket {
15    pub fn new(total_request: u64, minutes: u64) -> Self {
16        FixedBucket {
17            request: total_request,
18            duration: Duration::from_secs(minutes * 60),
19            capacity: 0,
20            rates: 0.0,
21            last_update: HashMap::new(),
22            requests: HashMap::new(),
23        }
24    }
25    pub fn bucket(self, capacity: u64, rates: f64) -> Self {
26        let rates = rates / 1000.0;
27        FixedBucket {
28            request: self.request,
29            duration: self.duration,
30            capacity: capacity,
31            rates: rates,
32            last_update: self.last_update,
33            requests: self.requests,
34        }
35    }
36    fn bucket_update(&mut self, key: String) {
37        let now = Instant::now();
38        let (tokens, b) = self.last_update.entry(key).or_insert((0.0, now));
39        println!("{}", tokens);
40        let current_time = now.duration_since(*b).as_millis();
41        let token = (self.rates * current_time as f64).ceil();
42        *tokens = (*tokens + token).min(self.capacity as f64);
43        *b = now;
44    }
45    fn fixed_bucket(&mut self, key: String) -> bool {
46        let now = Instant::now();
47        let (tokens, b) = self.requests.entry(key).or_insert((0, now));
48        let current = now.duration_since(*b);
49        // println!("{:?} {:?} {:?} {}", b, current, self.duration, tokens);
50        if current > self.duration {
51            *b = now;
52            *tokens = 0;
53        }
54        if *tokens <= self.request {
55            *tokens += 1;
56            true
57        } else {
58            false
59        }
60    }
61    pub fn allow(&mut self, ip: impl Into<String> + Copy, token: f64) -> Selfb {
62        let ip = ip.into();
63        let request = self.fixed_bucket(ip.clone());
64        if request {
65            self.bucket_update(ip.clone());
66            let (tokens, _) = self.last_update.get_mut(&ip).unwrap();
67            if *tokens >= token {
68                *tokens -= token;
69                Selfb {
70                    is_allowed: true,
71                    remaning_count: *tokens as u64,
72                }
73            } else {
74                Selfb {
75                    is_allowed: false,
76                    remaning_count: *tokens as u64,
77                }
78            }
79        } else {
80            Selfb {
81                is_allowed: false,
82                remaning_count: 0,
83            }
84        }
85    }
86}
rusty_rl/fixed_bucket.rs

rusty_rl/
fixed_bucket.rs
