Creates a `TokenBucket` wrapped in an `Option`. TokenBucket created is of `size` total capacity and takes `complete_refill_time_ms` milliseconds to go from zero tokens to total capacity. The `one_time_burst` is initial extra credit on top of total capacity, that does not replenish and which can be used for an initial burst of data. If the `size` or the `complete refill time` are zero, then `None
(size: u64, one_time_burst: u64, complete_refill_time_ms: u64)
| 127 | /// |
| 128 | /// If the `size` or the `complete refill time` are zero, then `None` is returned. |
| 129 | pub fn new(size: u64, one_time_burst: u64, complete_refill_time_ms: u64) -> Option<Self> { |
| 130 | // If either token bucket capacity or refill time is 0, disable limiting. |
| 131 | if size == 0 || complete_refill_time_ms == 0 { |
| 132 | return None; |
| 133 | } |
| 134 | // Formula for computing current refill amount: |
| 135 | // refill_token_count = (delta_time * size) / (complete_refill_time_ms * 1_000_000) |
| 136 | // In order to avoid overflows, simplify the fractions by computing greatest common divisor. |
| 137 | |
| 138 | let complete_refill_time_ns = complete_refill_time_ms * NANOSEC_IN_ONE_MILLISEC; |
| 139 | // Get the greatest common factor between `size` and `complete_refill_time_ns`. |
| 140 | let common_factor = gcd(size, complete_refill_time_ns); |
| 141 | // The division will be exact since `common_factor` is a factor of `size`. |
| 142 | let processed_capacity: u64 = size / common_factor; |
| 143 | // The division will be exact since `common_factor` is a factor of `complete_refill_time_ns`. |
| 144 | let processed_refill_time: u64 = complete_refill_time_ns / common_factor; |
| 145 | |
| 146 | Some(TokenBucket { |
| 147 | size, |
| 148 | one_time_burst, |
| 149 | refill_time: complete_refill_time_ms, |
| 150 | // Start off full. |
| 151 | budget: size, |
| 152 | // Last updated is now. |
| 153 | last_update: Instant::now(), |
| 154 | processed_capacity, |
| 155 | processed_refill_time, |
| 156 | }) |
| 157 | } |
| 158 | |
| 159 | /// Attempts to consume `tokens` from the bucket and returns whether the action succeeded. |
| 160 | // TODO (Issue #259): handle cases where a single request is larger than the full capacity |