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
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
use std::prelude::v1::*;

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

use crate::nanos::Nanos;

/// 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 accomodate 60 cells in one go, after which the equivalent of a minute of inactivity
/// is required for the burst allowance to be fully restored.
///
/// 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)
    }
}

impl Quota {
    /// A way to reconstruct a Quota from an in-use Gcra.
    ///
    /// This is useful mainly for [`crate::middleware::RateLimitingMiddleware`]
    /// where custom code may want to construct information based on
    /// the amount of burst balance remaining.
    pub(crate) fn from_gcra_parameters(t: Nanos, tau: Nanos) -> Quota {
        // Safety assurance: As we're calling this from this crate
        // only, and we do not allow creating a Gcra from 0
        // parameters, this is, in fact, safe.
        //
        // The casts may look a little sketch, but again, they're
        // constructed from parameters that came from the crate
        // exactly like that.
        let max_burst = unsafe { NonZeroU32::new_unchecked((tau.as_u64() / t.as_u64()) as u32) };
        let replenish_1_per = t.into();
        Quota {
            max_burst,
            replenish_1_per,
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use nonzero_ext::nonzero;

    #[test]
    fn time_multiples() {
        let hourly = Quota::per_hour(nonzero!(1u32));
        let minutely = Quota::per_minute(nonzero!(1u32));
        let secondly = Quota::per_second(nonzero!(1u32));

        assert_eq!(
            hourly.replenish_interval() / 60,
            minutely.replenish_interval()
        );
        assert_eq!(
            minutely.replenish_interval() / 60,
            secondly.replenish_interval()
        );
    }

    #[test]
    fn period_error_cases() {
        assert!(Quota::with_period(Duration::from_secs(0)).is_none());

        #[allow(deprecated)]
        {
            assert!(Quota::new(nonzero!(1u32), Duration::from_secs(0)).is_none());
        }
    }
}