quota 0.1.0

High-performance Rate-limiter
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

use std::collections::hash_map::Entry;
use std::collections::HashMap;
use crate::Quota;
use crate::QuotaPolicy;

pub struct QuotaPool {
    quotas: Vec<Quota>,
    last_access_times_ns: Vec<u64>,
    policy: QuotaPolicy,
    key_map: HashMap<&'static str, usize>,
    initial_available_tokens: u64,
}

impl QuotaPool {
    #[inline]
    pub fn new(policy: QuotaPolicy) -> QuotaPool {
        QuotaPool {
            quotas: Vec::new(),
            last_access_times_ns: Vec::new(),
            policy,
            key_map: HashMap::new(),
            initial_available_tokens: 0,
        }
    }

    pub fn quota(&mut self, key: &'static str) -> &Quota {
        match self.key_map.entry(key) {
            Entry::Occupied(occupied_entry) => {
                &self.quotas[*occupied_entry.get()]
            }
            Entry::Vacant(vacant_entry) => unsafe {
                let new_index = self.quotas.len();
                vacant_entry.insert(new_index);
                self.quotas.push(Quota::with_initial_tokens(self.initial_available_tokens));
                self.last_access_times_ns.push(0);
                &self.quotas.get_unchecked(new_index)
            }
        }
    }
}

#[cfg(test)]
mod tests;