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

use std::prelude::v1::*;

use nonzero_ext::nonzero;
use std::num::NonZeroU32;
use std::time::Duration;

/// A rate-limiting quota.
///
/// Quotas are expressed in a positive number of "cells" (the maximum number of positive decisions /
/// allowed items until the rate limiter needs to replenish) and the amount of time for the rate
/// limiter to replenish a single cell.
///
/// Neither the number of cells nor the replenishment unit of time may be zero.
///
/// # Burst sizes
/// There are multiple ways of expressing the same quota: a quota given as `Quota::per_second(1)`
/// allows, on average, the same number of cells through as a quota given as `Quota::per_minute(60)`.
/// However, the quota of `Quota::per_minute(60)` has a burst size of 60 cells, meaning it is possible
/// to accommodate 60 cells in one go, followed by a minute of waiting.
///
/// Burst size gets really important when you construct a rate limiter that should allow multiple
/// elements through at one time (using [`RateLimiter.check_n`](struct.RateLimiter.html#method.check_n)
/// and its related functions): Only
/// at most as many cells can be let through in one call as are given as the burst size.
///
/// In other words, the burst size is the maximum number of cells that the rate limiter will ever
/// allow through without replenishing them.
///
/// # Examples
///
/// Construct a quota that allows 50 cells per second (replenishing at a rate of one cell
/// per 20 milliseconds), with a burst size of 50 cells, allowing a full rate limiter to allow 50
/// cells through at a time:
/// ```rust
/// # use governor::Quota;
/// # use nonzero_ext::nonzero;
/// # use std::time::Duration;
/// let q = Quota::per_second(nonzero!(50u32));
/// assert_eq!(q, Quota::per_second(nonzero!(50u32)).allow_burst(nonzero!(50u32)));
/// assert_eq!(q.replenish_interval(), Duration::from_millis(20));
/// assert_eq!(q.burst_size().get(), 50);
/// // The Quota::new constructor is deprecated, but this constructs the equivalent quota:
/// #[allow(deprecated)]
/// assert_eq!(q, Quota::new(nonzero!(50u32), Duration::from_secs(1)).unwrap());
/// ```
///
/// Construct a quota that allows 2 cells per hour through (replenishing at a rate of one cell
/// per 30min), but allows bursting up to 90 cells at once:
/// ```rust
/// # use governor::Quota;
/// # use nonzero_ext::nonzero;
/// # use std::time::Duration;
/// let q = Quota::per_hour(nonzero!(2u32)).allow_burst(nonzero!(90u32));
/// assert_eq!(q.replenish_interval(), Duration::from_secs(30 * 60));
/// assert_eq!(q.burst_size().get(), 90);
/// // The entire maximum burst size will be restored if no cells are let through for 45 hours:
/// assert_eq!(q.burst_size_replenished_in(), Duration::from_secs(60 * 60 * (90 / 2)));
/// ```
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub struct Quota {
    pub(crate) max_burst: NonZeroU32,
    pub(crate) replenish_1_per: Duration,
}

/// Constructors for Quotas
impl Quota {
    /// Construct a quota for a number of cells per second. The given number of cells is also
    /// assumed to be the maximum burst size.
    pub const fn per_second(max_burst: NonZeroU32) -> Quota {
        let replenish_interval_ns = Duration::from_secs(1).as_nanos() / (max_burst.get() as u128);
        Quota {
            max_burst,
            replenish_1_per: Duration::from_nanos(replenish_interval_ns as u64),
        }
    }

    /// Construct a quota for a number of cells per 60-second period. The given number of cells is
    /// also assumed to be the maximum burst size.
    pub const fn per_minute(max_burst: NonZeroU32) -> Quota {
        let replenish_interval_ns = Duration::from_secs(60).as_nanos() / (max_burst.get() as u128);
        Quota {
            max_burst,
            replenish_1_per: Duration::from_nanos(replenish_interval_ns as u64),
        }
    }

    /// Construct a quota for a number of cells per 60-minute (3600-second) period. The given number
    /// of cells is also assumed to be the maximum burst size.
    pub const fn per_hour(max_burst: NonZeroU32) -> Quota {
        let replenish_interval_ns =
            Duration::from_secs(60 * 60).as_nanos() / (max_burst.get() as u128);
        Quota {
            max_burst,
            replenish_1_per: Duration::from_nanos(replenish_interval_ns as u64),
        }
    }

    /// Construct a quota that replenishes one cell in a given
    /// interval.
    ///
    /// This constructor is meant to replace [`::new`](#method.new),
    /// in cases where a longer refresh period than 1 cell/hour is
    /// necessary.
    ///
    /// If the time interval is zero, returns `None`.
    ///
    /// # Example
    /// ```rust
    /// # use nonzero_ext::nonzero;
    /// # use governor::Quota;
    /// # use std::time::Duration;
    /// // Replenish one cell per day, with a burst capacity of 10 cells:
    /// let _quota = Quota::with_period(Duration::from_secs(60 * 60 * 24))
    ///     .unwrap()
    ///     .allow_burst(nonzero!(10u32));
    /// ```
    pub fn with_period(replenish_1_per: Duration) -> Option<Quota> {
        if replenish_1_per.as_nanos() == 0 {
            None
        } else {
            Some(Quota {
                max_burst: nonzero!(1u32),
                replenish_1_per,
            })
        }
    }

    /// Adjusts the maximum burst size for a quota to construct a rate limiter with a capacity
    /// for at most the given number of cells.
    pub const fn allow_burst(self, max_burst: NonZeroU32) -> Quota {
        Quota { max_burst, ..self }
    }

    /// Construct a quota for a given burst size, replenishing the entire burst size in that
    /// given unit of time.
    ///
    /// Returns `None` if the duration is zero.
    ///
    /// This constructor allows greater control over the resulting
    /// quota, but doesn't make as much intuitive sense as other
    /// methods of constructing the same quotas. Unless your quotas
    /// are given as "max burst size, and time it takes to replenish
    /// that burst size", you are better served by the
    /// [`Quota::per_second`](#method.per_second) (and similar)
    /// constructors with the [`allow_burst`](#method.allow_burst)
    /// modifier.
    #[deprecated(
        since = "0.2.0",
        note = "This constructor is often confusing and non-intuitive. \
    Use the `per_(interval)` / `with_period` and `max_burst` constructors instead."
    )]
    pub fn new(max_burst: NonZeroU32, replenish_all_per: Duration) -> Option<Quota> {
        if replenish_all_per.as_nanos() == 0 {
            None
        } else {
            Some(Quota {
                max_burst,
                replenish_1_per: replenish_all_per / max_burst.get(),
            })
        }
    }
}

/// Retrieving information about a quota
impl Quota {
    /// The time it takes for a rate limiter with an exhausted burst budget to replenish
    /// a single element.
    pub const fn replenish_interval(&self) -> Duration {
        self.replenish_1_per
    }

    /// The maximum number of cells that can be allowed in one burst.
    pub const fn burst_size(&self) -> NonZeroU32 {
        self.max_burst
    }

    /// The time it takes to replenish the entire maximum burst size.
    pub const fn burst_size_replenished_in(&self) -> Duration {
        let fill_in_ns = self.replenish_1_per.as_nanos() * self.max_burst.get() as u128;
        Duration::from_nanos(fill_in_ns as u64)
    }
}