From 0cee59612d01911e054ddb7865caca7d4a615fbb Mon Sep 17 00:00:00 2001 From: Wez Furlong Date: Thu, 5 Sep 2024 13:47:58 -0700 Subject: [PATCH] throttle: add force_local flag I want to use a local-only throttle in an experiment, so this commit prepares the throttle crate for that. You can now use `local:200/hr` to request a local-only throttle, even when redis throttles have been enabled. While I'm in here, since this changes the debug repr and there are a number of snapshot tests that are affected by this change, I've added Debug and Display impls that match the canonical textual format. --- crates/kumo-api-types/src/shaping.rs | 48 +++------------- crates/throttle/src/lib.rs | 84 ++++++++++++++++++++++++++-- 2 files changed, 88 insertions(+), 44 deletions(-) diff --git a/crates/kumo-api-types/src/shaping.rs b/crates/kumo-api-types/src/shaping.rs index 26218b92..4f5dddad 100644 --- a/crates/kumo-api-types/src/shaping.rs +++ b/crates/kumo-api-types/src/shaping.rs @@ -1189,11 +1189,7 @@ additional_message_rate_throttles = {"second"="100/hr"} resolved.params.additional_message_rate_throttles, r#" { - "second": ThrottleSpec { - limit: 100, - period: 3600, - max_burst: None, - }, + "second": 100/h, } "# ); @@ -1233,11 +1229,7 @@ provider_max_message_rate = "120/s" resolved.params.additional_message_rate_throttles, r#" { - "shaping-provider-Office 365-invalid.source-rate": ThrottleSpec { - limit: 120, - period: 1, - max_burst: None, - }, + "shaping-provider-Office 365-invalid.source-rate": 120/s, } "# ); @@ -1359,19 +1351,11 @@ MergedEntry { smtp_auth_plain_password: None, allow_smtp_auth_plain_without_tls: false, max_message_rate: Some( - ThrottleSpec { - limit: 100, - period: 1, - max_burst: None, - }, + 100/s, ), additional_message_rate_throttles: {}, max_connection_rate: Some( - ThrottleSpec { - limit: 100, - period: 60, - max_burst: None, - }, + 100/m, ), max_deliveries_per_connection: 100, prohibited_hosts: CidrSet( @@ -1527,19 +1511,11 @@ MergedEntry { smtp_auth_plain_password: None, allow_smtp_auth_plain_without_tls: false, max_message_rate: Some( - ThrottleSpec { - limit: 100, - period: 1, - max_burst: None, - }, + 100/s, ), additional_message_rate_throttles: {}, max_connection_rate: Some( - ThrottleSpec { - limit: 100, - period: 60, - max_burst: None, - }, + 100/m, ), max_deliveries_per_connection: 100, prohibited_hosts: CidrSet( @@ -1776,19 +1752,11 @@ MergedEntry { smtp_auth_plain_password: None, allow_smtp_auth_plain_without_tls: false, max_message_rate: Some( - ThrottleSpec { - limit: 100, - period: 1, - max_burst: None, - }, + 100/s, ), additional_message_rate_throttles: {}, max_connection_rate: Some( - ThrottleSpec { - limit: 100, - period: 60, - max_burst: None, - }, + 100/m, ), max_deliveries_per_connection: 20, prohibited_hosts: CidrSet( diff --git a/crates/throttle/src/lib.rs b/crates/throttle/src/lib.rs index cfc3dfdf..64a5388f 100644 --- a/crates/throttle/src/lib.rs +++ b/crates/throttle/src/lib.rs @@ -38,13 +38,14 @@ pub enum Error { NonExistentLease, } -#[derive(Debug, Eq, PartialEq, Clone, Copy, Serialize, Deserialize, Hash)] +#[derive(Eq, PartialEq, Clone, Copy, Serialize, Deserialize, Hash)] #[serde(try_from = "String")] pub struct ThrottleSpec { pub limit: u64, /// Period, in seconds pub period: u64, pub max_burst: Option, + pub force_local: bool, } #[cfg(feature = "impl")] @@ -55,7 +56,33 @@ impl ThrottleSpec { let period = self.period; let max_burst = self.max_burst.unwrap_or(limit); let key = format!("{key}:{limit}:{max_burst}:{period}"); - throttle(&key, limit, Duration::from_secs(period), max_burst, Some(1)).await + throttle( + &key, + limit, + Duration::from_secs(period), + max_burst, + Some(1), + self.force_local, + ) + .await + } +} + +impl std::fmt::Debug for ThrottleSpec { + fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result { + match self.as_string() { + Ok(s) => write!(fmt, "{}", s), + Err(_) => Err(std::fmt::Error), + } + } +} + +impl std::fmt::Display for ThrottleSpec { + fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result { + match self.as_string() { + Ok(s) => write!(fmt, "{}", s), + Err(_) => Err(std::fmt::Error), + } } } @@ -72,7 +99,11 @@ impl ThrottleSpec { return Err(format!("cannot represent max_burst {burst} as string")); } - Ok(format!("{}/{period}", self.limit)) + Ok(format!( + "{}{}/{period}", + if self.force_local { "local:" } else { "" }, + self.limit + )) } } @@ -86,6 +117,10 @@ impl TryFrom for ThrottleSpec { impl TryFrom<&str> for ThrottleSpec { type Error = String; fn try_from(s: &str) -> Result { + let (force_local, s) = match s.strip_prefix("local:") { + Some(s) => (true, s), + None => (false, s), + }; let (limit, period) = s .split_once("/") .ok_or_else(|| format!("expected 'limit/period', got {s}"))?; @@ -121,6 +156,7 @@ impl TryFrom<&str> for ThrottleSpec { limit, period, max_burst: None, + force_local, }) } } @@ -243,6 +279,8 @@ async fn redis_throttle( /// to spread out across time. /// * `quantity` - how many tokens to add to the throttle. If omitted, /// 1 token is added. +/// * `force_local` - if true, always use the in-memory store on the local +/// machine even if the redis backend has been configured. #[cfg(feature = "impl")] pub async fn throttle( key: &str, @@ -250,8 +288,11 @@ pub async fn throttle( period: Duration, max_burst: u64, quantity: Option, + force_local: bool, ) -> Result { - if let Some(redis) = REDIS.get().cloned() { + if force_local { + local_throttle(key, limit, period, max_burst, quantity) + } else if let Some(redis) = REDIS.get().cloned() { redis_throttle(redis, key, limit, period, max_burst, quantity).await } else { local_throttle(key, limit, period, max_burst, quantity) @@ -340,14 +381,49 @@ mod test { limit: 100, period: 3600, max_burst: None, + force_local: false, } ); + assert_eq!( + ThrottleSpec::try_from("local:100/hr").unwrap(), + ThrottleSpec { + limit: 100, + period: 3600, + max_burst: None, + force_local: true, + } + ); + + assert_eq!( + ThrottleSpec { + limit: 100, + period: 3600, + max_burst: None, + force_local: false, + } + .as_string() + .unwrap(), + "100/h" + ); + assert_eq!( + ThrottleSpec { + limit: 100, + period: 3600, + max_burst: None, + force_local: true, + } + .as_string() + .unwrap(), + "local:100/h" + ); + assert_eq!( ThrottleSpec::try_from("1_0,0/hour").unwrap(), ThrottleSpec { limit: 100, period: 3600, max_burst: None, + force_local: false, } ); assert_eq!(