From 0788aec800d2aa99561f2662e3dbd6573221dc14 Mon Sep 17 00:00:00 2001 From: Clara Rull Date: Tue, 21 Jan 2025 10:32:45 -0800 Subject: [PATCH] Find most specific ratelimit Summary: *What* This diff changes the `find_ratelimitt_method` to return the most specific one that applies to a certain client id and set of identitites: 1. If there is a ratelimit for the main_id that will be the most specific one 2. If none match, the most specific one will be the one that matches the largest amount of identities. Eventually it could be one that matches 0 identities, because it's empty. *Why* We will want a way to have several ratelimits for a metric aply to different subsets, and likely a default that applies to all. This diff enables that default Reviewed By: andreacampi Differential Revision: D68418006 fbshipit-source-id: 718be8c9fa12a4eac05a546f9bcbf62cb5caadb4 --- .../edenapi_service/src/handlers/commit.rs | 23 ++-- .../src/middleware/rate_limiter.rs | 11 +- eden/mononoke/rate_limiting/src/facebook.rs | 60 +++++++++-- eden/mononoke/rate_limiting/src/lib.rs | 100 ++++++++++++++++-- eden/mononoke/rate_limiting/src/oss.rs | 7 +- 5 files changed, 175 insertions(+), 26 deletions(-) diff --git a/eden/mononoke/edenapi_service/src/handlers/commit.rs b/eden/mononoke/edenapi_service/src/handlers/commit.rs index 3dc1eb89ea53c..036ed3e3d5a36 100644 --- a/eden/mononoke/edenapi_service/src/handlers/commit.rs +++ b/eden/mononoke/edenapi_service/src/handlers/commit.rs @@ -191,7 +191,20 @@ async fn ratelimit_commit_creation(ctx: CoreContext) -> Result<(), Error> { } }; let category = rate_limiter.category(); - let limit = match rate_limiter.find_rate_limit(Metric::CommitsPerUser) { + + let client_request_info = match ctx.client_request_info() { + Some(client_request_info) => client_request_info, + None => { + debug!(ctx.logger(), "No client request info found"); + return Ok(()); + } + }; + + let limit = match rate_limiter.find_rate_limit( + Metric::CommitsPerUser, + None, + client_request_info.main_id.as_deref(), + ) { Some(limit) => limit, None => { debug!(ctx.logger(), "No commits_per_user rate limit found"); @@ -208,14 +221,6 @@ async fn ratelimit_commit_creation(ctx: CoreContext) -> Result<(), Error> { let max_value = limit.body.raw_config.limit; let time_window = limit.fci_metric.window.as_secs() as u32; - let client_request_info = match ctx.client_request_info() { - Some(client_request_info) => client_request_info, - None => { - debug!(ctx.logger(), "No client request info found"); - return Ok(()); - } - }; - let main_client_id = match &client_request_info.main_id { Some(main_client_id) => main_client_id, None => { diff --git a/eden/mononoke/edenapi_service/src/middleware/rate_limiter.rs b/eden/mononoke/edenapi_service/src/middleware/rate_limiter.rs index e66563636f429..78626aa204998 100644 --- a/eden/mononoke/edenapi_service/src/middleware/rate_limiter.rs +++ b/eden/mononoke/edenapi_service/src/middleware/rate_limiter.rs @@ -69,8 +69,13 @@ impl Middleware for ThrottleMiddleware { None })?; - let category = rate_limiter.category(); - let limit = rate_limiter.find_rate_limit(Metric::EdenApiQps)?; + let identities = state.try_borrow::()?.metadata().identities(); + + let limit = rate_limiter.find_rate_limit( + Metric::EdenApiQps, + Some(identities.clone()), + Some(&main_client_id), + )?; let enforced = match limit.body.raw_config.status { RateLimitStatus::Disabled => return None, @@ -78,6 +83,8 @@ impl Middleware for ThrottleMiddleware { RateLimitStatus::Enforced => true, _ => panic!("Invalid limit status: {:?}", limit.body.raw_config.status), }; + + let category = rate_limiter.category(); let counter = build_counter(&ctx, category, EDENAPI_QPS_LIMIT, &main_client_id); let max_value = limit.body.raw_config.limit; let time_window = limit.fci_metric.window.as_secs() as u32; diff --git a/eden/mononoke/rate_limiting/src/facebook.rs b/eden/mononoke/rate_limiting/src/facebook.rs index ddaebbe167176..bb1304c17aab1 100644 --- a/eden/mononoke/rate_limiting/src/facebook.rs +++ b/eden/mononoke/rate_limiting/src/facebook.rs @@ -132,12 +132,60 @@ impl RateLimiter for MononokeRateLimits { &self.category } - fn find_rate_limit(&self, metric: Metric) -> Option { - self.config - .rate_limits - .iter() - .find(|r| r.fci_metric.metric == metric) - .cloned() + // Find the most specific rate limit that applies to the given identities or main_id + // If no rate limit applies, return None + // If multiple rate limits apply, return the most specific one + // The most specific rate limit is the one that strictly matches the main_id + // If none match, return the most specific rate limit that matches the most identities + fn find_rate_limit( + &self, + metric: Metric, + identities: Option, + main_id: Option<&str>, + ) -> Option { + // First, try to find a rate limit that matches the main client ID + if let Some(main_id) = main_id { + if let Some(rate_limit) = self + .config + .rate_limits + .iter() + .filter(|r| r.fci_metric.metric == metric) + .find(|r| { + if let Some(crate::Target::MainClientId(ref id)) = r.target { + id == main_id + } else { + false + } + }) + .cloned() + { + return Some(rate_limit); + } + } + + // If no main client ID match is found, find the most specific one that applies to the identities + let mut max_identities = 0; + let mut most_specific_rate_limit = None; + + if let Some(identities) = identities { + self.config + .rate_limits + .iter() + .filter(|r| r.fci_metric.metric == metric) + .filter(|r| r.applies_to_client(&identities, None)) + .for_each(|r| match &r.target { + Some(crate::Target::Identities(is)) => { + let num_identities = is.len(); + if num_identities >= max_identities { + max_identities = num_identities; + most_specific_rate_limit = Some(r.clone()); + } + } + _ => {} + }); + } + + most_specific_rate_limit } } diff --git a/eden/mononoke/rate_limiting/src/lib.rs b/eden/mononoke/rate_limiting/src/lib.rs index 6cbb858ca5ba0..17be167d639e7 100644 --- a/eden/mononoke/rate_limiting/src/lib.rs +++ b/eden/mononoke/rate_limiting/src/lib.rs @@ -63,7 +63,12 @@ pub trait RateLimiter { fn category(&self) -> &str; - fn find_rate_limit(&self, metric: Metric) -> Option; + fn find_rate_limit( + &self, + metric: Metric, + identities: Option, + main_id: Option<&str>, + ) -> Option; } define_stats! { @@ -97,7 +102,7 @@ impl RateLimitEnvironment { } } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Default, PartialEq)] pub struct RateLimitBody { pub raw_config: rate_limiting_config::RateLimitBody, } @@ -108,11 +113,11 @@ pub struct MononokeRateLimitConfig { pub load_shed_limits: Vec, } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq)] pub struct RateLimit { pub body: RateLimitBody, #[allow(dead_code)] - target: Option, + pub target: Option, #[allow(dead_code)] pub fci_metric: FciMetric, } @@ -231,14 +236,14 @@ pub enum RateLimitReason { LoadShedMetric(String, i64, i64), } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq)] pub enum Target { StaticSlice(StaticSlice), MainClientId(String), Identities(MononokeIdentitySet), } -#[derive(Debug, Copy, Clone)] +#[derive(Debug, Copy, Clone, PartialEq)] struct SlicePct(u8); impl TryFrom for SlicePct { @@ -253,7 +258,7 @@ impl TryFrom for SlicePct { } } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq)] pub struct StaticSlice { slice_pct: SlicePct, // This is hashed with a client's hostname to allow us to change @@ -262,7 +267,7 @@ pub struct StaticSlice { target: StaticSliceTarget, } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq)] pub enum StaticSliceTarget { Identities(MononokeIdentitySet), MainClientId(String), @@ -345,6 +350,8 @@ fn in_throttled_slice( #[cfg(test)] mod test { + use std::sync::Arc; + use mononoke_macros::mononoke; use permission_checker::MononokeIdentity; @@ -463,4 +470,81 @@ mod test { // 100% of SERVICE_IDENTITY: bar assert!(hundred_pct_service_identity.matches_client(idents2.as_ref(), None)); } + + #[cfg(fbcode_build)] + #[mononoke::fbinit_test] + fn test_find_rate_limit(fb: FacebookInit) { + let main_client_id_rate_limit = RateLimit { + body: RateLimitBody::default(), + target: Some(Target::MainClientId("client_id".to_string())), + fci_metric: FciMetric { + metric: Metric::EgressBytes, + window: Duration::from_secs(60), + scope: Scope::Global, + }, + }; + + let identities_rate_limit = RateLimit { + body: RateLimitBody::default(), + target: Some(Target::Identities( + [MononokeIdentity::new("TIER", "foo")].into(), + )), + fci_metric: FciMetric { + metric: Metric::EgressBytes, + window: Duration::from_secs(60), + scope: Scope::Global, + }, + }; + + let empty_target_rate_limit = RateLimit { + body: RateLimitBody::default(), + target: Some(Target::Identities([].into())), + fci_metric: FciMetric { + metric: Metric::EgressBytes, + window: Duration::from_secs(60), + scope: Scope::Global, + }, + }; + + let rate_limiter = create_rate_limiter( + fb, + "test".to_string(), + Arc::new(MononokeRateLimitConfig { + rate_limits: vec![ + main_client_id_rate_limit.clone(), + identities_rate_limit.clone(), + empty_target_rate_limit.clone(), + ], + load_shed_limits: vec![], + }), + ); + + let mut idents = MononokeIdentitySet::new(); + idents.insert(MononokeIdentity::new("USER", "bar")); + + assert!( + rate_limiter.find_rate_limit( + Metric::EgressBytes, + Some(idents.clone()), + Some("non_matching_id") + ) == Some(empty_target_rate_limit) + ); + assert!( + rate_limiter.find_rate_limit( + Metric::EgressBytes, + Some(idents.clone()), + Some("client_id") + ) == Some(main_client_id_rate_limit.clone()) + ); + + idents.insert(MononokeIdentity::new("TIER", "foo")); + assert!( + rate_limiter.find_rate_limit(Metric::EgressBytes, Some(idents.clone()), None) + == Some(identities_rate_limit) + ); + assert!( + rate_limiter.find_rate_limit(Metric::EgressBytes, Some(idents), Some("client_id")) + == Some(main_client_id_rate_limit) + ); + } } diff --git a/eden/mononoke/rate_limiting/src/oss.rs b/eden/mononoke/rate_limiting/src/oss.rs index 2e94f98afab9d..067d704518539 100644 --- a/eden/mononoke/rate_limiting/src/oss.rs +++ b/eden/mononoke/rate_limiting/src/oss.rs @@ -65,7 +65,12 @@ impl RateLimiter for FakeLimiter { &self.category } - fn find_rate_limit(&self, metric: Metric) -> Option { + fn find_rate_limit( + &self, + metric: Metric, + identities: Option, + main_id: Option<&str>, + ) -> Option { None } }