Add ratelimit! utility macro

foundations/Cargo.toml

@@ -76,7 +76,7 @@ telemetry = [
 ]
 
 # Enables a subset of telemetry features suitable for usage in clients.
-client-telemetry = ["logging", "metrics", "tracing"]
+client-telemetry = ["logging", "metrics", "tracing", "ratelimit"]
 
 # Enables the telemetry server.
 telemetry-server = [
@@ -105,6 +105,7 @@ tokio-runtime-metrics = [
 
 # Enables logging functionality.
 logging = [
+    "ratelimit",
     "dep:governor",
     "dep:parking_lot",
     "dep:slog-async",
@@ -172,6 +173,9 @@ cli = ["settings", "dep:clap"]
 # Enables testing-related functionality.
 testing = ["dep:foundations-macros"]
 
+# Enables the ratelimit! utility macro.
+ratelimit = ["dep:governor"]
+
 # Enables panicking when too much nesting is reached on the logger
 panic_on_too_much_logger_nesting = []
 

foundations/src/lib.rs

@@ -103,6 +103,8 @@ pub mod security;
 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")]
@@ -153,6 +155,10 @@ pub type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
 /// Operational (post-initialization) result that has [`Error`] as an error variant.
 pub type Result<T> = std::result::Result<T, Error>;
 
+#[cfg(feature = "ratelimit")]
+#[doc(inline)]
+pub use utils::ratelimit;
+
 /// Basic service information.
 #[derive(Clone, Debug, Default)]
 pub struct ServiceInfo {

foundations/src/telemetry/log/mod.rs

@@ -178,6 +178,11 @@ macro_rules! __add_fields {
 /// // and fields by `;`.
 /// log::error!("Answer: {}", 42; "foo" => "bar", "baz" => 1337);
 ///
+/// // Log messages can be rate-limited with an extra keyword
+/// for _ in 0..2 {
+///     log::error!(ratelimit=1/h; "Answer: {}", 42; "foo" => "bar");
+/// }
+///
 /// assert_eq!(*ctx.log_records(), &[
 ///     TestLogRecord {
 ///         level: Level::Error,
@@ -196,6 +201,13 @@ macro_rules! __add_fields {
 ///             ("baz".into(), "1337".into()),
 ///             ("foo".into(), "bar".into())
 ///         ]
+///     },
+///     TestLogRecord {
+///         level: Level::Error,
+///         message: "Answer: 42".into(),
+///         fields: vec![
+///             ("foo".into(), "bar".into())
+///         ]
 ///     }
 /// ]);
 /// ```
@@ -204,6 +216,14 @@ macro_rules! __add_fields {
 #[macro_export]
 #[doc(hidden)]
 macro_rules! __error {
+    ( ratelimit=$limit:literal / $unit:tt ; $($args:tt)+ ) => {
+        $crate::ratelimit!($limit / $unit;
+            $crate::reexports_for_macros::slog::error!(
+                $crate::telemetry::log::internal::current_log().read(),
+                $($args)+
+            )
+        );
+    };
     ( $($args:tt)+ ) => {
         $crate::reexports_for_macros::slog::error!(
             $crate::telemetry::log::internal::current_log().read(),
@@ -240,6 +260,11 @@ macro_rules! __error {
 /// // and fields by `;`.
 /// log::warn!("Answer: {}", 42; "foo" => "bar", "baz" => 1337);
 ///
+/// // Log messages can be rate-limited with an extra keyword
+/// for _ in 0..2 {
+///     log::warn!(ratelimit=1/h; "Answer: {}", 42; "foo" => "bar");
+/// }
+///
 /// assert_eq!(*ctx.log_records(), &[
 ///     TestLogRecord {
 ///         level: Level::Warning,
@@ -258,6 +283,13 @@ macro_rules! __error {
 ///             ("baz".into(), "1337".into()),
 ///             ("foo".into(), "bar".into())
 ///         ]
+///     },
+///     TestLogRecord {
+///         level: Level::Warning,
+///         message: "Answer: 42".into(),
+///         fields: vec![
+///             ("foo".into(), "bar".into())
+///         ]
 ///     }
 /// ]);
 /// ```
@@ -266,6 +298,14 @@ macro_rules! __error {
 #[doc(hidden)]
 #[macro_export]
 macro_rules! __warn {
+    ( ratelimit=$limit:literal / $unit:tt ; $($args:tt)+ ) => {
+        $crate::ratelimit!($limit / $unit;
+            $crate::reexports_for_macros::slog::warn!(
+                $crate::telemetry::log::internal::current_log().read(),
+                $($args)+
+            )
+        );
+    };
     ( $($args:tt)+ ) => {
         $crate::reexports_for_macros::slog::warn!(
             $crate::telemetry::log::internal::current_log().read(),
@@ -302,6 +342,11 @@ macro_rules! __warn {
 /// // and fields by `;`.
 /// log::debug!("Answer: {}", 42; "foo" => "bar", "baz" => 1337);
 ///
+/// // Log messages can be rate-limited with an extra keyword
+/// for _ in 0..2 {
+///     log::debug!(ratelimit=1/h; "Answer: {}", 42; "foo" => "bar");
+/// }
+///
 /// assert_eq!(*ctx.log_records(), &[
 ///     TestLogRecord {
 ///         level: Level::Debug,
@@ -320,6 +365,13 @@ macro_rules! __warn {
 ///             ("baz".into(), "1337".into()),
 ///             ("foo".into(), "bar".into())
 ///         ]
+///     },
+///     TestLogRecord {
+///         level: Level::Debug,
+///         message: "Answer: 42".into(),
+///         fields: vec![
+///             ("foo".into(), "bar".into())
+///         ]
 ///     }
 /// ]);
 /// ```
@@ -328,6 +380,14 @@ macro_rules! __warn {
 #[macro_export]
 #[doc(hidden)]
 macro_rules! __debug {
+    ( ratelimit=$limit:literal / $unit:tt ; $($args:tt)+ ) => {
+        $crate::ratelimit!($limit / $unit;
+            $crate::reexports_for_macros::slog::debug!(
+                $crate::telemetry::log::internal::current_log().read(),
+                $($args)+
+            )
+        );
+    };
     ( $($args:tt)+ ) => {
         $crate::reexports_for_macros::slog::debug!(
             $crate::telemetry::log::internal::current_log().read(),
@@ -364,6 +424,11 @@ macro_rules! __debug {
 /// // and fields by `;`.
 /// log::info!("Answer: {}", 42; "foo" => "bar", "baz" => 1337);
 ///
+/// // Log messages can be rate-limited with an extra keyword
+/// for _ in 0..2 {
+///     log::info!(ratelimit=1/h; "Answer: {}", 42; "foo" => "bar");
+/// }
+///
 /// assert_eq!(*ctx.log_records(), &[
 ///     TestLogRecord {
 ///         level: Level::Info,
@@ -382,6 +447,13 @@ macro_rules! __debug {
 ///             ("baz".into(), "1337".into()),
 ///             ("foo".into(), "bar".into())
 ///         ]
+///     },
+///     TestLogRecord {
+///         level: Level::Info,
+///         message: "Answer: 42".into(),
+///         fields: vec![
+///             ("foo".into(), "bar".into())
+///         ]
 ///     }
 /// ]);
 /// ```
@@ -390,6 +462,14 @@ macro_rules! __debug {
 #[macro_export]
 #[doc(hidden)]
 macro_rules! __info {
+    ( ratelimit=$limit:literal / $unit:tt ; $($args:tt)+ ) => {
+        $crate::ratelimit!($limit / $unit;
+            $crate::reexports_for_macros::slog::info!(
+                $crate::telemetry::log::internal::current_log().read(),
+                $($args)+
+            )
+        );
+    };
     ( $($args:tt)+ ) => {
         $crate::reexports_for_macros::slog::info!(
             $crate::telemetry::log::internal::current_log().read(),
@@ -426,6 +506,11 @@ macro_rules! __info {
 /// // and fields by `;`.
 /// log::trace!("Answer: {}", 42; "foo" => "bar", "baz" => 1337);
 ///
+/// // Log messages can be rate-limited with an extra keyword
+/// for _ in 0..2 {
+///     log::trace!(ratelimit=1/h; "Answer: {}", 42; "foo" => "bar");
+/// }
+///
 /// assert_eq!(*ctx.log_records(), &[
 ///     TestLogRecord {
 ///         level: Level::Trace,
@@ -444,6 +529,13 @@ macro_rules! __info {
 ///             ("baz".into(), "1337".into()),
 ///             ("foo".into(), "bar".into())
 ///         ]
+///     },
+///     TestLogRecord {
+///         level: Level::Trace,
+///         message: "Answer: 42".into(),
+///         fields: vec![
+///             ("foo".into(), "bar".into())
+///         ]
 ///     }
 /// ]);
 /// ```
@@ -452,6 +544,14 @@ macro_rules! __info {
 #[macro_export]
 #[doc(hidden)]
 macro_rules! __trace {
+    ( ratelimit=$limit:literal / $unit:tt ; $($args:tt)+ ) => {
+        $crate::ratelimit!($limit / $unit;
+            $crate::reexports_for_macros::slog::trace!(
+                $crate::telemetry::log::internal::current_log().read(),
+                $($args)+
+            )
+        );
+    };
     ( $($args:tt)+ ) => {
         $crate::reexports_for_macros::slog::trace!(
             $crate::telemetry::log::internal::current_log().read(),

foundations/src/utils.rs

@@ -15,3 +15,143 @@ 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);
+    }
+}