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
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

use std::{
    collections::HashMap,
    iter::Inspect,
    time::{Duration, Instant},
};

use crate::{FixedBucket, Selfb};

/// fixed bucket
///
/// ```
///     let b = FixedBucket::new(600, 10).bucket(90, 10.0)
/// ```
impl FixedBucket {
    pub fn new(total_request: u64, minutes: u64) -> Self {
        FixedBucket {
            request: total_request,
            duration: Duration::from_secs(minutes * 60),
            capacity: 0,
            rates: 0.0,
            last_update: HashMap::new(),
            requests: HashMap::new(),
        }
    }
    pub fn bucket(self, capacity: u64, rates: f64) -> Self {
        let rates = rates / 1000.0;
        FixedBucket {
            request: self.request,
            duration: self.duration,
            capacity: capacity,
            rates: rates,
            last_update: self.last_update,
            requests: self.requests,
        }
    }
    fn bucket_update(&mut self, key: String) {
        let now = Instant::now();
        let (tokens, b) = self.last_update.entry(key).or_insert((0.0, now));
        println!("{}", tokens);
        let current_time = now.duration_since(*b).as_millis();
        let token = (self.rates * current_time as f64).ceil();
        *tokens = (*tokens + token).min(self.capacity as f64);
        *b = now;
    }
    fn fixed_bucket(&mut self, key: String) -> bool {
        let now = Instant::now();
        let (tokens, b) = self.requests.entry(key).or_insert((0, now));
        let current = now.duration_since(*b);
        // println!("{:?} {:?} {:?} {}", b, current, self.duration, tokens);
        if current > self.duration {
            *b = now;
            *tokens = 0;
        }
        if *tokens <= self.request {
            *tokens += 1;
            true
        } else {
            false
        }
    }
    pub fn allow(&mut self, ip: impl Into<String> + Copy, token: f64) -> Selfb {
        let ip = ip.into();
        let request = self.fixed_bucket(ip.clone());
        if request {
            self.bucket_update(ip.clone());
            let (tokens, _) = self.last_update.get_mut(&ip).unwrap();
            if *tokens >= token {
                *tokens -= token;
                Selfb {
                    is_allowed: true,
                    remaning_count: *tokens as u64,
                }
            } else {
                Selfb {
                    is_allowed: false,
                    remaning_count: *tokens as u64,
                }
            }
        } else {
            Selfb {
                is_allowed: false,
                remaning_count: 0,
            }
        }
    }
}