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
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, } #[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) } } 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()) } /// 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 } } } #[cfg(feature = "std")] mod with_async; #[cfg(feature = "std")] pub use with_async::*;