diff --git a/foundations/Cargo.toml b/foundations/Cargo.toml index 03a8a1e..5d1653a 100644 --- a/foundations/Cargo.toml +++ b/foundations/Cargo.toml @@ -105,7 +105,6 @@ tokio-runtime-metrics = [ # Enables logging functionality. logging = [ "ratelimit", - "dep:governor", "dep:parking_lot", "dep:slog-async", "dep:slog-json", @@ -119,8 +118,8 @@ logging = [ # Enables distributed tracing functionality. tracing = [ + "ratelimit", "dep:foundations-macros", - "dep:governor", "dep:parking_lot", "dep:rand", "dep:cf-rustracing-jaeger", @@ -173,7 +172,7 @@ cli = ["settings", "dep:clap"] testing = ["dep:foundations-macros"] # Enables the ratelimit! utility macro. -ratelimit = ["dep:governor"] +ratelimit = ["dep:governor", "dep:crossbeam-utils"] # Enables panicking when too much nesting is reached on the logger panic_on_too_much_logger_nesting = [] diff --git a/foundations/src/lib.rs b/foundations/src/lib.rs index 8827819..2b27491 100644 --- a/foundations/src/lib.rs +++ b/foundations/src/lib.rs @@ -34,6 +34,7 @@ //! - **metrics**: Enables metrics functionality. //! - **logging**: Enables logging functionality. //! - **tracing**: Enables distributed tracing functionality. +//! - **ratelimit**: Enables helpers to simplify rate-limiting your code. //! - **testing**: Enables testing-related functionality. //! - **security**: Enables security features. Available only on Linux (x86_64, aarch64) with the `libclang-dev` package installed (for bindgen). //! - **jemalloc**: Enables [jemalloc] memory allocator which is known to perform much better than @@ -99,12 +100,14 @@ pub mod telemetry; ))] pub mod security; +#[doc(hidden)] +#[cfg(feature = "ratelimit")] +pub mod ratelimit; + #[doc(hidden)] pub mod reexports_for_macros { #[cfg(feature = "tracing")] pub use cf_rustracing; - #[cfg(feature = "ratelimit")] - pub use governor; #[cfg(feature = "security")] pub use once_cell; #[cfg(feature = "metrics")] @@ -157,7 +160,7 @@ pub type Result = std::result::Result; #[cfg(feature = "ratelimit")] #[doc(inline)] -pub use utils::ratelimit; +pub use self::ratelimit::ratelimit; /// Basic service information. #[derive(Clone, Debug, Default)] diff --git a/foundations/src/ratelimit.rs b/foundations/src/ratelimit.rs new file mode 100644 index 0000000..72fb210 --- /dev/null +++ b/foundations/src/ratelimit.rs @@ -0,0 +1,169 @@ +use crossbeam_utils::CachePadded; +use governor::clock::{Clock, QuantaClock}; +use std::sync::OnceLock; + +// Reexport for macro +pub use governor::Quota; + +/// Wrapper for sharing a `&'static QuantaClock` between many RateLimiter instances. +/// +/// [`QuantaClock`] is 40 bytes but all state is the same across all instances, so keeping +/// a single static instance around saves some memory and helps with cache locality. +#[derive(Debug, Clone, Copy)] +pub struct StaticQuantaClock(&'static QuantaClock); + +impl Default for StaticQuantaClock { + fn default() -> Self { + static CLOCK: CachePadded> = CachePadded::new(OnceLock::new()); + Self(CLOCK.get_or_init(Default::default)) + } +} + +impl Clock for StaticQuantaClock { + type Instant = ::Instant; + + #[inline] + fn now(&self) -> Self::Instant { + Clock::now(self.0) + } +} + +/// Type alias for a [`governor::RateLimiter`] using our [`StaticQuantaClock`]. +pub type DirectRateLimiter = governor::RateLimiter< + governor::state::NotKeyed, + governor::state::InMemoryState, + StaticQuantaClock, + governor::middleware::NoOpMiddleware<::Instant>, +>; + +/// Applies a rate limit to the evaluation of an expression. +/// +/// The macro takes two arguments, separated by a `;`. The first is the quota to use +/// for the ratelimit. This can either be a const expression evaluating to a +/// [`governor::Quota`], or a rate specifier like `200/s`, `10/m`, or `5/h`. The latter +/// three are equivalent to [`Quota`]'s `per_second`/`per_minute`/`per_hour` constructors. +/// +/// The second argument is the expression to evaluate if the rate limit has not been +/// reached yet. The expression's result will be discarded. +/// +/// # Examples +/// ```rust +/// # fn expensive_computation() -> u32 { 42 } +/// # +/// use foundations::telemetry::log; +/// use governor::Quota; +/// use std::num::NonZeroU32; +/// +/// foundations::ratelimit!(10/s; println!("frequently failing operation failed!") ); +/// +/// // You can return data from the expression with an Option: +/// let mut output = None; +/// foundations::ratelimit!(1/h; output.insert(expensive_computation()) ); +/// assert_eq!(output, Some(42)); +/// +/// // A quota expression allows customizing the burst size. By default, +/// // it is equivalent to the rate per time unit (i.e., 10/m yields a burst size of 10). +/// // Note: you could also reference a `const` declared somewhere else here. +/// foundations::ratelimit!( +/// Quota::per_hour(NonZeroU32::new(100).unwrap()).allow_burst(NonZeroU32::new(1).unwrap()); +/// println!("this will be printed only once before the rate limit kicks in") +/// ); +/// +/// // Here the rate limit kicks in after the initial burst of 60 iterations: +/// let mut counter = 0; +/// for _ in 0..1000 { +/// foundations::ratelimit!(60/h; counter += 1); +/// } +/// assert_eq!(counter, 60); +/// ``` +/// +/// [`Quota`]: governor::Quota +#[macro_export] +#[doc(hidden)] +macro_rules! __ratelimit { + ($limit:literal / s ; $expr:expr) => { + $crate::__ratelimit!( + $crate::ratelimit::Quota::per_second(::std::num::NonZeroU32::new($limit).unwrap()); + $expr + ) + }; + + ($limit:literal / m ; $expr:expr) => { + $crate::__ratelimit!( + $crate::ratelimit::Quota::per_minute(::std::num::NonZeroU32::new($limit).unwrap()); + $expr + ) + }; + + ($limit:literal / h ; $expr:expr) => { + $crate::__ratelimit!( + $crate::ratelimit::Quota::per_hour(::std::num::NonZeroU32::new($limit).unwrap()); + $expr + ) + }; + + ($quota:expr ; $expr:expr) => {{ + const QUOTA: $crate::ratelimit::Quota = $quota; + static LIMITER: ::std::sync::LazyLock<$crate::ratelimit::DirectRateLimiter> = ::std::sync::LazyLock::new( + || $crate::ratelimit::DirectRateLimiter::direct_with_clock(QUOTA, ::std::default::Default::default()) + ); + if LIMITER.check().is_ok() { + $expr; + } + }}; +} + +pub use __ratelimit as ratelimit; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_ratelimit() { + use governor::Quota; + use std::num::NonZeroU32; + + const CUSTOM_QUOTA: Quota = + Quota::per_hour(NonZeroU32::new(60).unwrap()).allow_burst(NonZeroU32::new(20).unwrap()); + + // Burst size is only 20 for this quota, despite the refill rate being 60/h + let mut res_custom = 0; + for _ in 0..200 { + ratelimit!(CUSTOM_QUOTA; res_custom += 1); + } + + assert_eq!(res_custom, 20); + + // Cells may refill as the loop executes already, so a value >20 is possible + let mut res_sec = 0; + for _ in 0..100 { + ratelimit!(20/s; res_sec += 1); + } + + assert!(res_sec >= 20); + assert!(res_sec < 100); + + // This should execute exactly 3 times; we don't expect any cells to refill + let mut res_minute = 1; + for _ in 0..20 { + ratelimit!(3/m; res_minute *= 2); + } + + assert_eq!(res_minute, 1 << 3); + + let mut res_hour_a = 0; + let mut res_hour_b = 0; + + for _ in 0..1000 { + ratelimit!(100/h; { + res_hour_a += 1; + res_hour_b += 2; + }); + } + + assert!(res_hour_a >= 100); + assert!(res_hour_a < 1000); + assert_eq!(res_hour_b, 2 * res_hour_a); + } +} diff --git a/foundations/src/sentry/hook.rs b/foundations/src/sentry/hook.rs index 38da384..323e979 100644 --- a/foundations/src/sentry/hook.rs +++ b/foundations/src/sentry/hook.rs @@ -1,6 +1,7 @@ //! Sentry hook implementation for tracking sentry events and rate-limiting them. use super::SentrySettings; +use crate::ratelimit::StaticQuantaClock; use governor::{Quota, RateLimiter}; use std::borrow::Cow; use std::num::NonZeroU32; @@ -34,9 +35,9 @@ pub fn install_hook_with_settings( options: &mut sentry_core::ClientOptions, settings: &SentrySettings, ) { - let rate_limiter = settings - .max_events_per_second - .map(|rl| RateLimiter::::dashmap(Quota::per_second(rl))); + let rate_limiter = settings.max_events_per_second.map(|rl| { + RateLimiter::dashmap_with_clock(Quota::per_second(rl), StaticQuantaClock::default()) + }); let previous = options.before_send.take(); diff --git a/foundations/src/telemetry/log/rate_limit.rs b/foundations/src/telemetry/log/rate_limit.rs index e381f58..782c063 100644 --- a/foundations/src/telemetry/log/rate_limit.rs +++ b/foundations/src/telemetry/log/rate_limit.rs @@ -1,11 +1,11 @@ +use crate::ratelimit::{DirectRateLimiter, Quota, StaticQuantaClock}; use crate::telemetry::settings::RateLimitingSettings; -use governor::{DefaultDirectRateLimiter, Quota, RateLimiter}; use slog::{Drain, OwnedKVList, Record}; use std::num::NonZeroU32; pub(crate) struct RateLimitingDrain { inner: D, - rate_limiter: Option, + rate_limiter: Option, } impl RateLimitingDrain { @@ -13,7 +13,10 @@ impl RateLimitingDrain { let rate_limiter = if settings.enabled && let Some(rate) = NonZeroU32::new(settings.max_events_per_second) { - Some(RateLimiter::direct(Quota::per_second(rate))) + Some(DirectRateLimiter::direct_with_clock( + Quota::per_second(rate), + StaticQuantaClock::default(), + )) } else { None }; diff --git a/foundations/src/telemetry/tracing/rate_limit.rs b/foundations/src/telemetry/tracing/rate_limit.rs index 8bf898e..4baaded 100644 --- a/foundations/src/telemetry/tracing/rate_limit.rs +++ b/foundations/src/telemetry/tracing/rate_limit.rs @@ -1,16 +1,14 @@ +use crate::ratelimit::{DirectRateLimiter, Quota, StaticQuantaClock}; use crate::telemetry::settings::ActiveSamplingSettings; use cf_rustracing::sampler::Sampler; use cf_rustracing::span::CandidateSpan; use cf_rustracing::{Result, sampler::ProbabilisticSampler}; -use governor::clock::DefaultClock; -use governor::middleware::NoOpMiddleware; -use governor::state::{InMemoryState, NotKeyed}; -use governor::{Quota, RateLimiter}; +use std::num::NonZeroU32; #[derive(Debug)] pub(crate) struct RateLimitingProbabilisticSampler { inner: ProbabilisticSampler, - rate_limiter: Option>, + rate_limiter: Option, } impl Default for RateLimitingProbabilisticSampler { @@ -27,13 +25,13 @@ impl RateLimitingProbabilisticSampler { /// If `sampling_rate` is not in the range `0.0...1.0`, /// it will return an error with the kind `ErrorKind::InvalidInput`. pub(crate) fn new(settings: &ActiveSamplingSettings) -> Result { - let rate_limiter = if settings.rate_limit.enabled { - settings - .rate_limit - .max_events_per_second - .try_into() - .ok() - .map(|r| RateLimiter::direct(Quota::per_second(r))) + let rate_limiter = if settings.rate_limit.enabled + && let Some(rate) = NonZeroU32::new(settings.rate_limit.max_events_per_second) + { + Some(DirectRateLimiter::direct_with_clock( + Quota::per_second(rate), + StaticQuantaClock::default(), + )) } else { None }; diff --git a/foundations/src/utils.rs b/foundations/src/utils.rs index f956056..ef17f55 100644 --- a/foundations/src/utils.rs +++ b/foundations/src/utils.rs @@ -15,143 +15,3 @@ macro_rules! feature_use { // NOTE: don't complain about unused macro for feature combinations that don't use it. #[allow(unused_imports)] pub(crate) use feature_use; - -/// Applies a rate limit to the evaluation of an expression. -/// -/// The macro takes two arguments, separated by a `;`. The first is the quota to use -/// for the ratelimit. This can either be a const expression evaluating to a -/// [`governor::Quota`], or a rate specifier like `200/s`, `10/m`, or `5/h`. The latter -/// three are equivalent to [`Quota`]'s `per_second`/`per_minute`/`per_hour` constructors. -/// -/// The second argument is the expression to evaluate if the rate limit has not been -/// reached yet. The expression's result will be discarded. -/// -/// # Examples -/// ```rust -/// # fn expensive_computation() -> u32 { 42 } -/// # -/// use foundations::telemetry::log; -/// use governor::Quota; -/// use std::num::NonZeroU32; -/// -/// foundations::ratelimit!(10/s; println!("frequently failing operation failed!") ); -/// -/// // You can return data from the expression with an Option: -/// let mut output = None; -/// foundations::ratelimit!(1/h; output.insert(expensive_computation()) ); -/// assert_eq!(output, Some(42)); -/// -/// // A quota expression allows customizing the burst size. By default, -/// // it is equivalent to the rate per time unit (i.e., 10/m yields a burst size of 10). -/// // Note: you could also reference a `const` declared somewhere else here. -/// foundations::ratelimit!( -/// Quota::per_hour(NonZeroU32::new(100).unwrap()).allow_burst(NonZeroU32::new(1).unwrap()); -/// println!("this will be printed only once before the rate limit kicks in") -/// ); -/// -/// // Here the rate limit kicks in after the initial burst of 60 iterations: -/// let mut counter = 0; -/// for _ in 0..1000 { -/// foundations::ratelimit!(60/h; counter += 1); -/// } -/// assert_eq!(counter, 60); -/// ``` -/// -/// [`Quota`]: governor::Quota -#[cfg(feature = "ratelimit")] -#[macro_export] -#[doc(hidden)] -macro_rules! __ratelimit { - ($limit:literal / s ; $expr:expr) => { - $crate::__ratelimit!( - $crate::reexports_for_macros::governor::Quota::per_second( - ::std::num::NonZeroU32::new($limit).unwrap() - ); - $expr - ) - }; - - ($limit:literal / m ; $expr:expr) => { - $crate::__ratelimit!( - $crate::reexports_for_macros::governor::Quota::per_minute( - ::std::num::NonZeroU32::new($limit).unwrap() - ); - $expr - ) - }; - - ($limit:literal / h ; $expr:expr) => { - $crate::__ratelimit!( - $crate::reexports_for_macros::governor::Quota::per_hour( - ::std::num::NonZeroU32::new($limit).unwrap() - ); - $expr - ) - }; - - ($quota:expr ; $expr:expr) => {{ - const QUOTA: $crate::reexports_for_macros::governor::Quota = $quota; - static LIMITER: ::std::sync::LazyLock<$crate::reexports_for_macros::governor::DefaultDirectRateLimiter> = - ::std::sync::LazyLock::new(|| $crate::reexports_for_macros::governor::RateLimiter::direct(QUOTA)); - if LIMITER.check().is_ok() { - $expr; - } - }}; -} - -#[cfg(feature = "ratelimit")] -pub use __ratelimit as ratelimit; - -#[cfg(test)] -mod tests { - use super::*; - - #[cfg(feature = "ratelimit")] - #[test] - fn test_ratelimit() { - use governor::Quota; - use std::num::NonZeroU32; - - const CUSTOM_QUOTA: Quota = - Quota::per_hour(NonZeroU32::new(60).unwrap()).allow_burst(NonZeroU32::new(20).unwrap()); - - // Burst size is only 20 for this quota, despite the refill rate being 60/h - let mut res_custom = 0; - for _ in 0..200 { - ratelimit!(CUSTOM_QUOTA; res_custom += 1); - } - - assert_eq!(res_custom, 20); - - // Cells may refill as the loop executes already, so a value >20 is possible - let mut res_sec = 0; - for _ in 0..100 { - ratelimit!(20/s; res_sec += 1); - } - - assert!(res_sec >= 20); - assert!(res_sec < 100); - - // This should execute exactly 3 times; we don't expect any cells to refill - let mut res_minute = 1; - for _ in 0..20 { - ratelimit!(3/m; res_minute *= 2); - } - - assert_eq!(res_minute, 1 << 3); - - let mut res_hour_a = 0; - let mut res_hour_b = 0; - - for _ in 0..1000 { - ratelimit!(100/h; { - res_hour_a += 1; - res_hour_b += 2; - }); - } - - assert!(res_hour_a >= 100); - assert!(res_hour_a < 1000); - assert_eq!(res_hour_b, 2 * res_hour_a); - } -}