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 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208
//! An in-memory rate limiter that can make decisions for a single
//! situation.
// #![feature(trace_macros)]
// trace_macros!(true);
use std::marker::PhantomData;
use std::num::NonZeroU32;
use std::time::{Duration, Instant};
use {
algorithms::{Algorithm, DefaultAlgorithm},
InconsistentCapacity, NegativeMultiDecision,
};
/// 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 server's rate
/// limit.
#[derive(Debug, Clone)]
pub struct DirectRateLimiter<A: Algorithm = DefaultAlgorithm> {
state: A::BucketState,
algorithm: A,
}
impl<A> DirectRateLimiter<A>
where
A: Algorithm,
{
/// Construct a new rate limiter that allows `capacity` cells per
/// time unit through.
/// # Examples
/// You can construct a GCRA rate limiter like so:
/// ```
/// # use std::num::NonZeroU32;
/// # use std::time::Duration;
/// use ratelimit_meter::{DirectRateLimiter, GCRA};
/// # #[macro_use] extern crate nonzero_ext;
/// # extern crate ratelimit_meter;
/// # fn main () {
/// let _gcra = DirectRateLimiter::<GCRA>::new(nonzero!(100u32), Duration::from_secs(5));
/// # }
/// ```
///
/// and similarly, for a leaky bucket:
/// ```
/// # use std::time::Duration;
/// use ratelimit_meter::{DirectRateLimiter, LeakyBucket};
/// # #[macro_use] extern crate nonzero_ext;
/// # extern crate ratelimit_meter;
/// # fn main () {
/// let _lb = DirectRateLimiter::<LeakyBucket>::new(nonzero!(100u32), Duration::from_secs(5));
/// # }
/// ```
pub fn new(capacity: NonZeroU32, per_time_unit: Duration) -> Self {
DirectRateLimiter {
state: <A as Algorithm>::BucketState::default(),
algorithm: <A as Algorithm>::construct(capacity, nonzero!(1u32), per_time_unit)
.unwrap(),
}
}
/// Construct a new rate limiter that allows `capacity` cells per
/// second.
/// # Examples
/// Constructing a GCRA rate limiter that lets through 100 cells per second:
/// ```
/// # use std::time::Duration;
/// use ratelimit_meter::{DirectRateLimiter, GCRA};
/// # #[macro_use] extern crate nonzero_ext;
/// # extern crate ratelimit_meter;
/// # fn main () {
/// let _gcra = DirectRateLimiter::<GCRA>::per_second(nonzero!(100u32));
/// # }
/// ```
///
/// and a leaky bucket:
/// ```
/// # use std::time::Duration;
/// use ratelimit_meter::{DirectRateLimiter, LeakyBucket};
/// # #[macro_use] extern crate nonzero_ext;
/// # extern crate ratelimit_meter;
/// # fn main () {
/// let _gcra = DirectRateLimiter::<LeakyBucket>::per_second(nonzero!(100u32));
/// # }
/// ```
pub fn per_second(capacity: NonZeroU32) -> Self {
Self::new(capacity, Duration::from_secs(1))
}
/// Return a builder that can be used to construct a rate limiter using
/// the parameters passed to the Builder.
pub fn build_with_capacity(capacity: NonZeroU32) -> Builder<A> {
Builder {
capacity,
cell_weight: nonzero!(1u32),
time_unit: Duration::from_secs(1),
end_result: PhantomData,
}
}
/// Tests if a single cell can be accommodated at
/// `Instant::now()`. If it can be, `check` updates the rate
/// limiter state to account for the conforming cell and returns
/// `Ok(())`.
///
/// If the cell is non-conforming (i.e., it can't be accomodated
/// at this time stamp), `check_at` returns `Err` with information
/// about the earliest time at which a cell could be considered
/// conforming.
pub fn check(&mut self) -> Result<(), <A as Algorithm>::NegativeDecision> {
self.algorithm.test_and_update(&self.state, Instant::now())
}
/// Tests if `n` cells can be accommodated at the current time
/// stamp. If (and only if) all cells in the batch can be
/// accomodated, the `MultiDecider` updates the internal state to
/// account for all cells and returns `Ok(())`.
///
/// If the entire batch of cells would not be conforming but the
/// rate limiter has the capacity to accomodate the cells at any
/// point in time, `check_n_at` returns error
/// [`NegativeMultiDecision::BatchNonConforming`](../../enum.NegativeMultiDecision.html#variant.BatchNonConforming),
/// holding the number of cells the rate limiter's negative
/// outcome result.
///
/// If `n` exceeds the bucket capacity, `check_n_at` returns
/// [`NegativeMultiDecision::InsufficientCapacity`](../../enum.NegativeMultiDecision.html#variant.InsufficientCapacity),
/// indicating that a batch of this many cells can never succeed.
pub fn check_n(
&mut self,
n: u32,
) -> Result<(), NegativeMultiDecision<<A as Algorithm>::NegativeDecision>> {
self.algorithm
.test_n_and_update(&self.state, n, Instant::now())
}
/// Tests whether a single cell can be accommodated at the given
/// time stamp. See [`check`](#method.check).
pub fn check_at(&mut self, at: Instant) -> Result<(), <A as Algorithm>::NegativeDecision> {
self.algorithm.test_and_update(&self.state, at)
}
/// Tests if `n` cells can be accommodated at the given time
/// (`Instant::now()`), using [`check_n`](#method.check_n)
pub fn check_n_at(
&mut self,
n: u32,
at: Instant,
) -> Result<(), NegativeMultiDecision<<A as Algorithm>::NegativeDecision>> {
self.algorithm.test_n_and_update(&self.state, n, at)
}
}
/// An object that allows incrementally constructing rate Limiter
/// objects.
pub struct Builder<T>
where
T: Algorithm + Sized,
{
capacity: NonZeroU32,
cell_weight: NonZeroU32,
time_unit: Duration,
end_result: PhantomData<T>,
}
impl<A> Builder<A>
where
A: Algorithm + Sized,
{
/// Sets the "weight" of each cell being checked against the
/// bucket. Each cell fills the bucket by this much.
pub fn cell_weight(
&mut self,
weight: NonZeroU32,
) -> Result<&mut Builder<A>, InconsistentCapacity> {
if self.cell_weight > self.capacity {
return Err(InconsistentCapacity {
capacity: self.capacity,
cell_weight: self.cell_weight,
});
}
self.cell_weight = weight;
Ok(self)
}
/// Sets the "unit of time" within which the bucket drains.
///
/// The assumption is that in a period of `time_unit` (if no cells
/// are being checked), the bucket is fully drained.
pub fn per(&mut self, time_unit: Duration) -> &mut Builder<A> {
self.time_unit = time_unit;
self
}
/// Builds a rate limiter of the specified type.
pub fn build(&self) -> Result<DirectRateLimiter<A>, InconsistentCapacity> {
Ok(DirectRateLimiter {
state: <A as Algorithm>::BucketState::default(),
algorithm: <A as Algorithm>::construct(
self.capacity,
self.cell_weight,
self.time_unit,
)?,
})
}
}