rl_core/

tracker.rs

use crate::Duration as _;

/// State for a single rate limit.
///
/// This tracks the state of a single rate limit. For example the rate limit for a particular IP.
///
/// This object is intentionally small and cheap so that it can be created for fine-grained tracking.
#[cfg_attr(test, derive(PartialEq))]
#[derive(Clone,Debug)]
pub enum Tracker<T: crate::Time = std::time::SystemTime> {
	#[doc(hidden)]
	Filling{
		/// The time when the tracker was last empty. It has been filling since then.
		empty_at: T,
	},
	#[doc(hidden)]
	Full{
		/// Extra credits over the regular burst capacity.
		overfull: u32,
	},
	#[doc(hidden)]
	Unlimited,
}

impl<T: crate::Time> Tracker<T> {
	/// Create a tracker that is full.
	///
	/// This creates a tracker that has it's full burst capacity available.
	pub fn full() -> Self {
		Tracker::Full { overfull: 0 }
	}

	/// Create a tracker that is overfull.
	///
	/// This tracker is not only full to its full burst capacity but also has extra capacity. This extra capacity will be consumed before dipping into the `burst` quota and regular rate limiting. This tracker will not fill itself up to the overfull state again, once the extra capacity is consumed it behaves as a regular tracker. (If you want it to fill further use a different `crate::Config` which has a larger burst capacity.)
	///
	/// Note: This allows for requests up to [`crate::Config::burst + capacity`](Config) to be fulfilled until the overfull capacity runs out.
	pub fn overfull(capacity: u32) -> Self {
		Tracker::Full { overfull: capacity }
	}

	/// Create a tracker that is unlimited.
	///
	/// This can be used to disable rate limiting for a specific tracker. It is semantically identical to being overfull with infinite capacity.
	pub fn unlimited() -> Self {
		Tracker::Unlimited
	}

	/// Create a tracker that starts at the provided time.
	///
	/// This is an alias for [`Tracker::empty_at()`](Tracker::empty_at()).
	#[deprecated(note="Use `empty_at()` which has a more descriptive name.")]
	pub fn new_at(now: T) -> Self {
		Self::empty_at(now)
	}

	/// Create a tracker that was empty at the provided time.
	///
	/// This means that the bucket will be empty at the indicated time and will fill starting from there. Note that this is **not** what you want for an limit that hasn't been used for a while and cleaned from the DB. For that case you want [`Tracker::default()`].
	///
	/// Note that if `now` is in the future it is possible to create a tracker that is empty and will remain empty until that point. This may be seen as a sort of "overdrawn" bucket that will take filling to return to zero.
	pub fn empty_at(now: T) -> Self {
		Tracker::Filling { empty_at: now }
	}

	/// Create a tracker that has the specified capacity at the specified time.
	pub fn with_capacity_at(
		config: &crate::Config<T>,
		capacity: u32,
		now: T
	) -> Self {
		if let Some(overfull) = capacity.checked_sub(config.burst) {
			Self::overfull(overfull)
		} else {
			Self::with_partial_capacity_at(config, capacity, now)
		}
	}

	/// Create a tracker that has the specified capacity up to the burst size.
	///
	/// This will never [overfill](Tracker::overfull()) the tracker, instead saturating at the burst capacity. If you want to allow overfilling use [`Tracker::with_capacity_at()`](Tracker::with_capacity_at()).
	pub fn with_limited_capacity_at(
		config: &crate::Config<T>,
		capacity: u32,
		now: T
	) -> Self {
		if capacity >= config.burst {
			Self::full()
		} else {
			Self::with_partial_capacity_at(config, capacity, now)
		}
	}

	fn with_partial_capacity_at(
		config: &crate::Config<T>,
		capacity: u32,
		now: T
	) -> Self {
		debug_assert!(capacity <= config.burst);

		Tracker::Filling {
			empty_at: now.sub(config.rate.clone().mul(capacity)),
		}
	}

	fn overfull_capacity(&self) -> u32 {
		match *self {
			Tracker::Full{overfull} => overfull,
			Tracker::Filling{..} => 0,
			Tracker::Unlimited => u32::MAX,
		}
	}

	/// Returns the tokens available at the specified time.
	///
	/// Note that history is not stored, this is just extrapolating from the current state. So if `config.rate` is 1/s advancing `now` by 10s will raise the capacity by 10 (up to `config.burst`), moving `now` back 10s will lower the capacity by 10 (down to 0). Moving `now` past the time where tokens were acquired will not restore those tokens to the capacity as returned by this function.
	pub fn capacity_at(
		&self,
		config: &crate::Config<T>,
		now: T,
	) -> u32 {
		debug_assert_ne!(config.rate, T::Duration::zero(), "rate must be greater than zero");

		match *self {
			Tracker::Filling{ref empty_at} => {
				match now.since(empty_at.clone()) {
					Some(elapsed) => {
						let periods = elapsed.div(config.rate.clone());
						periods.min(config.burst)
					}
					None => 0,
				}
			}
			Tracker::Full{overfull} => {
				config.burst.saturating_add(overfull)
			}
			Tracker::Unlimited => u32::MAX,
		}
	}

	/// Attempt to acquire `count` tokens from the rate limiter.
	///
	/// If the requested number of tokens are available the state is updated and `Ok(())` is returned. Otherwise an error describing the issue is returned.
	///
	/// Warning: To save memory a Tracker does not remember its crate::Config. This means that you can switch the `config` argument between calls on the same Tracker. This is **not recommended** as it is not immediately obvious what the result will be but is completely supported. The only guarantee provided is that the new rate limit will (eventually) take effect. The logical state of the Tracker when this occurs may be surprising. For example a large "burst" config may be immediately filled or a new low "rate" may result in a previous "burst" capacity disappearing. The exact behaviour is not guaranteed to be stable. If you want a specific behaviour it is best to extract the information you want with the old config and create a fresh tracker using the constructor functions to produce a tracker in the state that you expect.
	pub fn acquire_at(
		&mut self,
		config: &crate::Config<T>,
		count: u32,
		now: T,
	) -> Result<(), crate::Denied<T>> {
		self.acquire_range_at(config, count..=count, now)
			.map(|acquired| {
				debug_assert_eq!(acquired, count);
			})
	}

	/// Acquire tokens from the tracker.
	///
	/// Acquires at least `request.start()` tokens but no more than `request.end()`.
	///
	/// Returns the number of tokens acquired.
	///
	/// If the available capacity is less than `request.start()` `Err(crate::Denied::TooEarly)` is returned.
	///
	/// If `request.start()` is zero than this call always succeeds.
	pub fn acquire_range_at(
		&mut self,
		config: &crate::Config<T>,
		request: std::ops::RangeInclusive<u32>,
		now: T,
	) -> Result<u32, crate::Denied<T>> {
		debug_assert_ne!(config.rate, T::Duration::zero(), "rate must be greater than zero");

		if request.is_empty() {
			return Err(crate::Denied::EmptyRequest)
		}
		if request.start().saturating_sub(self.overfull_capacity()) > config.burst {
			return Err(crate::Denied::TooBig)
		}

		match *self {
			Tracker::Filling{ref mut empty_at} => {
				let mut fill_time = now.clone().since(empty_at.clone())
					.unwrap_or(T::Duration::zero());

				let max_fill_time = config.fill_time();
				if fill_time > max_fill_time {
					fill_time = max_fill_time.clone();
					*empty_at = now.sub(max_fill_time);
				}

				let min_duration = config.rate.clone().mul(*request.start());
				let max_duration = config.rate.clone().mul(*request.end());
				if fill_time >= max_duration {
					empty_at.add_assign(max_duration);
					Ok(*request.end())
				} else if fill_time < min_duration {
					let min_available_at = empty_at.clone().add(min_duration);
					Err(crate::Denied::TooEarly(crate::TooEarly{next: min_available_at}))
				} else {
					let capacity = fill_time.div(config.rate.clone());
					empty_at.add_assign(config.rate.clone().mul(capacity));
					Ok(capacity)
				}
			}
			Tracker::Full{ref mut overfull} => {
				if *overfull >= *request.end() {
					*overfull -= request.end();
					Ok(*request.end())
				} else {
					let total = overfull.saturating_add(config.burst);
					let acquired = total.min(*request.end());
					let remaining = total - acquired;
					*self = Tracker::Filling {
						empty_at: now.sub(config.rate.clone().mul(remaining)),
					};
					Ok(acquired)
				}
			}
			Tracker::Unlimited => Ok(*request.end())
		}
	}

	/// Acquire up to `count` tokens from the rate limiter.
	///
	/// Returns the number of tokens actually acquired. Fails if no tokens were acquired. If 0 is acceptable use `.unwrap_or(0)` to explicitly waive the error.
	///
	/// As a special case requesting `0` will always succeed with `Ok(0)`.
	///
	/// WARNING: This function will return `Ok(0)` on disabled configs. If you are using disabled configs please migrate to [`Tracker::acquire_up_to_at_2()`][Tacker::acquire_up_to_at_2()] for a proper error.
	///
	/// Equivalent to [`Tracker::acquire_range_at(config, 1..=count, now)`](Tracker::acquire_range_at()).
	#[deprecated(note="Use acquire_up_to_at_2.")]
	pub fn acquire_up_to_at(
		&mut self,
		config: &crate::Config<T>,
		count: u32,
		now: T,
	) -> Result<u32, crate::TooEarly<T>> {
		match self.acquire_up_to_at_2(config, count, now) {
			Ok(acquired) => Ok(acquired),
			Err(crate::Denied::TooEarly(err)) => {
				Err(err)
			}
			Err(crate::Denied::TooBig) if config.is_disabled() => {
				Ok(0)
			}
			Err(other) => {
				unreachable!("Unexpected error in acquire_up_to(): {}", other)
			}
		}
	}

	/// Acquire up to `count` tokens from the rate limiter.
	///
	/// Returns the number of tokens actually acquired. Fails if no tokens were acquired. If 0 is acceptable use `.unwrap_or(0)` to explicitly waive the error.
	///
	/// As a special case requesting `0` will always succeed with `Ok(0)`.
	///
	/// Equivalent to [`Tracker::acquire_range_at(config, 1..=count, now)`](Tracker::acquire_range_at()).
	pub fn acquire_up_to_at_2(
		&mut self,
		config: &crate::Config<T>,
		count: u32,
		now: T,
	) -> Result<u32, crate::Denied<T>> {
		if count == 0 {
			return Ok(0)
		}

		self.acquire_range_at(config, 1..=count, now)
	}

	/// Force acquire from the rate limiter.
	///
	/// This can not fail, the capacity can go "below zero" using this method. For example if the Config replenishes at 1 token per second and there are 2 tokens available, then acquiring 5 tokens will result in a capacity of 0 but requires 4 seconds for the next token to become available.
	///
	/// Note: While [`Tracker::capacity()`](Tracker::capacity()) will never report negative values the time reported in [`TooEarly.available_at()`](TooEarly.available_at()) will accurately reflect when the request will become available.
	pub fn force_acquire_at(
		&mut self,
		config: &crate::Config<T>,
		count: u32,
		now: T,
	) {
		match self {
			Tracker::Filling { empty_at } => {
				empty_at.add_assign(config.rate.clone().mul(count));
			}
			Tracker::Full { overfull } => {
				if *overfull >= count {
					*overfull -= count;
				} else {
					let remaining = count - *overfull;
					*self = Tracker::Filling {
						empty_at: now
							.add(config.rate.clone().mul(remaining))
							.sub(config.fill_time()),
					}
				}
			}
			Tracker::Unlimited => {}
		}
	}

	/// Add capacity.
	///
	/// This adds some number of tokens to the tracker.
	///
	/// Warning: This method will overfill the tracker. If this is not desired use [`add_limited_capacity_at()`][Tracker::add_limited_capacity_at()].
	///
	/// Note: When this method overfills the tracker any existing progress toward the next token is lost. For example if a 1s token interval was 1ns away from fully filling, adding 1 capacity will effectively add that 1ns of time, whereas if it was 999ms away from full it will add 999ms. This also means that `t.add_capacity_at(c, 1, now); t.acquire_at(c, 1, now)` will not always leave `t` in the same state as if the additional fills the tracker the incremental progress will be reset.
	pub fn add_capacity_at(
		&mut self,
		config: &crate::Config<T>,
		tokens: u32,
		now: T,
	) {
		match self {
			Tracker::Filling { empty_at } => {
				let min = now.clone().sub(config.fill_time());
				if *empty_at < min {
					*empty_at = min.clone()
				}
				let duration = config.rate.clone().mul(tokens);
				empty_at.sub_assign(duration);
				if let Some(overfull_duration) = min.clone().since(empty_at.clone()) {
					*self = Tracker::overfull(overfull_duration.div(config.rate.clone()));
				}
			}
			Tracker::Full{overfull} => {
				*overfull = overfull.saturating_add(tokens);
			}
			Tracker::Unlimited => {}
		}
	}

	/// Add capacity without overfilling.
	///
	/// This adds some number of tokens to the tracker without allowing it to overfill.
	///
	/// Note: If overfilling is desired use [`add_capacity()`][Tracker::add_capacity_at()].
	pub fn add_limited_capacity_at(
		&mut self,
		config: &crate::Config<T>,
		tokens: u32,
		_now: T,
	) {
		match self {
			Tracker::Filling { empty_at } => {
				let duration = config.rate.clone().mul(tokens);
				empty_at.sub_assign(duration);
			}
			Tracker::Full{..} => {}
			Tracker::Unlimited => {}
		}
	}

	/// Attempts to minimize the state.
	///
	/// If a rate limit hasn't been used in a long time the bucket will be full and the state can be simplified.
	///
	/// Returns `true` if the state is "simple" and equivalent to [`Tracker::default()`]. If this occurs you don't need to store the state and can use the default state if needed again.
	pub fn simplify_at(
		&mut self,
		config: &crate::Config<T>,
		now: T,
	) -> bool {
		if let Tracker::Filling{empty_at} = self {
			let min = now.sub(config.fill_time());
			if *empty_at <= min || config.is_disabled() {
				*self = Tracker::full();
			}
		}

		matches!(self, Tracker::Full { overfull: 0 })
	}
}