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
pub mod gcra;
pub mod leaky_bucket;

pub use self::gcra::*;
pub use self::leaky_bucket::*;

use evmap::ShallowCopy;
use failure::Fail;
use std::fmt;
use std::num::NonZeroU32;
use std::time::{Duration, Instant};
use {InconsistentCapacity, NegativeMultiDecision};

/// The default rate limiting algorithm in this crate: The ["leaky
/// bucket"](leaky_bucket/struct.LeakyBucket.html).
///
/// The leaky bucket algorithm is fairly easy to understand and has
/// decent performance in most cases. If better threaded performance
/// is needed, this crate also offers the
/// [`GCRA`](gcra/struct.GCRA.html) algorithm.
pub type DefaultAlgorithm = LeakyBucket;

/// Provides additional information about non-conforming cells, most
/// importantly the earliest time until the next cell could be
/// considered conforming.
///
/// Since this does not account for effects like thundering herds,
/// users should always add random jitter to the times given.
pub trait NonConformance {
    /// Returns the earliest time at which a decision could be
    /// conforming (excluding conforming decisions made by the Decider
    /// that are made in the meantime).
    fn earliest_possible(&self) -> Instant;

    /// Returns the minimum amount of time from the time that the
    /// decision was made (relative to the `at` argument in a
    /// `Decider`'s `check_at` method) that must pass before a
    /// decision can be conforming. Since Durations can not be
    /// negative, a zero duration is returned if `from` is already
    /// after that duration.
    fn wait_time_from(&self, from: Instant) -> Duration;

    /// Returns the minimum amount of time (down to 0) that needs to
    /// pass from the current instant for the Decider to consider a
    /// cell conforming again.
    fn wait_time(&self) -> Duration {
        self.wait_time_from(Instant::now())
    }
}

/// The trait that implementations of metered rate-limiter algorithms
/// have to implement.
///
/// Implementing structures are expected to represent the "parameters"
/// (e.g., the allowed requests/s), and keep the information necessary
/// to make a decision, e.g. concrete usage statistics for an
/// in-memory rate limiter, in the associated structure
/// [`BucketState`](#associatedtype.BucketState).
pub trait Algorithm: Send + Sync + Sized + fmt::Debug {
    /// The state of a single rate limiting bucket.
    ///
    /// Every new rate limiting state is initialized as `Default`. The
    /// states must be safe to share across threads (this crate uses a
    /// `parking_lot` Mutex to allow that).
    type BucketState: RateLimitState<Self> + Default + Send + Sync + Eq + ShallowCopy + fmt::Debug;

    /// The type returned when a rate limiting decision for a single
    /// cell is negative. Each rate limiting algorithm can decide to
    /// return the type that suits it best, but most algorithms'
    /// decisions also implement
    /// [`NonConformance`](trait.NonConformance.html), to ease
    /// handling of how long to wait.
    type NegativeDecision: PartialEq + Fail;

    /// Constructs a rate limiter with the given parameters:
    /// `capacity` is the number of cells to allow, weighing
    /// `cell_weight`, every `per_time_unit`.
    fn construct(
        capacity: NonZeroU32,
        cell_weight: NonZeroU32,
        per_time_unit: Duration,
    ) -> Result<Self, InconsistentCapacity>;

    /// Tests if `n` cells can be accommodated in the rate limiter at
    /// the instant `at` and updates the rate-limiter state to account
    /// for the weight of the cells and updates the ratelimiter state.
    ///
    /// The update is all or nothing: Unless all n cells can be
    /// accommodated, the state of the rate limiter will not be
    /// updated.
    fn test_n_and_update(
        &self,
        state: &Self::BucketState,
        n: u32,
        at: Instant,
    ) -> Result<(), NegativeMultiDecision<Self::NegativeDecision>>;

    /// Tests if a single cell can be accommodated in the rate limiter
    /// at the instant `at` and updates the rate-limiter state to
    /// account for the weight of the cell.
    ///
    /// This method is provided by default, using the `n` test&update
    /// method.
    fn test_and_update(
        &self,
        state: &Self::BucketState,
        at: Instant,
    ) -> Result<(), Self::NegativeDecision> {
        match self.test_n_and_update(state, 1, at) {
            Ok(()) => Ok(()),
            Err(NegativeMultiDecision::BatchNonConforming(1, nc)) => Err(nc),
            Err(other) => unreachable!(
                "BUG: measuring a batch of size 1 reported insufficient capacity: {:?}",
                other
            ),
        }
    }
}

/// Trait that all rate limit states have to implement around
/// housekeeping in keyed rate limiters.
pub trait RateLimitState<P> {
    /// Returns the last time instant that the state had any relevance
    /// (i.e. the rate limiter would behave exactly as if it was a new
    /// rate limiter after this time).
    ///
    /// If the state has not been touched for a given amount of time,
    /// the keyed rate limiter will expire it.
    ///
    /// # Thread safety
    /// This uses a bucket state snapshot to determine eligibility;
    /// race conditions can occur.
    fn last_touched(&self, params: &P) -> Instant;
}