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

use crate::gcra::{NotUntil, Tat, GCRA};
use crate::lib::*;
use crate::{clock, NegativeMultiDecision, Quota};

/// An in-memory rate limiter that makes direct (un-keyed)
/// rate-limiting decisions. Direct rate limiters can be used to
/// e.g. regulate the transmission of packets on a single connection,
/// or to ensure that an API client stays within a service's rate
/// limit.
#[derive(Debug)]
pub struct DirectRateLimiter<C: clock::Clock = clock::DefaultClock> {
    state: Tat,
    gcra: GCRA<C::Instant>,
    clock: C,
}

/// The default constructor in `std` mode.
#[cfg(feature = "std")]
impl DirectRateLimiter<clock::DefaultClock> {
    /// Construct a new direct rate limiter for a quota with the default clock.
    pub fn new(quota: Quota) -> Self {
        let clock = clock::DefaultClock::default();
        DirectRateLimiter::new_with_clock(quota, &clock)
    }
}

/// Manually checking cells against a rate limit.
impl<C: clock::Clock> DirectRateLimiter<C> {
    /// Allow a single cell through the rate limiter.
    ///
    /// If the rate limit is reached, `check` returns information about the earliest
    /// time that a cell might be allowed through again.  
    pub fn check(&self) -> Result<(), NotUntil<C::Instant>> {
        self.gcra.test_and_update(&self.state, self.clock.now())
    }

    /// Allow *only all* `n` cells through the rate limiter.
    ///
    /// This method can succeed in only one way and fail in two ways:
    /// * Success: If all `n` cells can be accommodated, it returns `Ok(())`.
    /// * Failure (but ok): Not all cells can make it through at the current time.
    ///   The result is `Err(NegativeMultiDecision::BatchNonConforming(NotUntil))`, which can
    ///   be interrogated about when the batch might next conform.
    /// * Failure (the batch can never go through): The rate limit is too low for the given number
    ///   of cells.
    ///
    /// # Performance
    /// This method diverges a little from the GCRA algorithm, using
    /// multiplication to determine the next theoretical arrival time, and so
    /// is not as fast as checking a single cell.  
    pub fn check_all(
        &self,
        n: NonZeroU32,
    ) -> Result<(), NegativeMultiDecision<NotUntil<C::Instant>>> {
        self.gcra
            .test_n_all_and_update(n, &self.state, self.clock.now())
    }
}

impl<C: clock::Clock> DirectRateLimiter<C> {
    /// Construct a new direct rate limiter with a custom clock.
    pub fn new_with_clock(quota: Quota, clock: &C) -> DirectRateLimiter<C> {
        let gcra: GCRA<C::Instant> = GCRA::new(clock.now(), quota);
        let clock = clock.clone();
        let state = gcra.new_state(clock.now());
        DirectRateLimiter { state, clock, gcra }
    }

    /// Returns a reference to the rate limiter's clock.
    pub fn get_clock(&self) -> &C {
        &self.clock
    }
}

#[cfg(feature = "std")]
mod future;
#[cfg(feature = "std")]
pub use future::*;

#[cfg(feature = "std")]
mod sinks;
#[cfg(feature = "std")]
pub use sinks::*;

#[cfg(feature = "std")]
mod streams;
#[cfg(feature = "std")]
pub use streams::*;