quota 0.1.1

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
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
87
88
89
90
91
92
93
94

use std::cell::{Cell, RefCell};
use std::collections::hash_map::Entry;
use std::collections::HashMap;
use crate::{Quota, QuotaResult};
use crate::QuotaPolicy;

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

impl QuotaPool {
    #[inline]
    pub fn new(policy: QuotaPolicy, initial_available_tokens: u64) -> QuotaPool {
        QuotaPool {
            quotas: Cell::new(Vec::new()),
            last_access_times_ns: Cell::new(Vec::new()),
            policy,
            key_map: Cell::new(HashMap::new()),
            initial_available_tokens,
            epoch: Cell::new(std::time::Instant::now())
        }
    }

    #[inline]
    pub fn start(&mut self) {
        let now = std::time::Instant::now();
        self.epoch = Cell::new(now);
        self.last_access_times_ns.get_mut().fill(now);
    }

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

    #[inline]
    pub fn quota_with_index(&self, key: &'static str) -> (&Quota, usize) {
        unsafe {
            let key_map = &mut *self.key_map.as_ptr();
            let quotas = &mut *self.quotas.as_ptr();
            match key_map.entry(key) {
                Entry::Occupied(occupied_entry) => {
                    let i = *occupied_entry.get();
                    (&quotas[i], i)
                }
                Entry::Vacant(vacant_entry) => {
                    let last_access_times_ns = &mut *self.last_access_times_ns.as_ptr();
                    let new_index = quotas.len();
                    vacant_entry.insert(new_index);
                    quotas.push(Quota::with_initial_tokens(self.initial_available_tokens));
                    last_access_times_ns.push(std::time::Instant::now());
                    (&quotas.get_unchecked(new_index), new_index)
                }
            }
        }
    }

    /// Ticks the Quota with the time difference since last access
    /// and then consumes the cost and returns the result
    #[inline]
    pub fn consume(&self, key: &'static str, cost: u64) -> QuotaResult {
        let (quota, index) = self.quota_with_index(key);
        let last_access_time_ns = unsafe { (&mut *self.last_access_times_ns.as_ptr())[index] };

        let delta = last_access_time_ns.elapsed();

        self.policy.tick(delta.as_nanos() as u64, &quota);
        quota.consume(cost)
    }
}

#[cfg(test)]
mod tests;