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

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 => {

eden/mononoke/edenapi_service/src/middleware/rate_limiter.rs

@@ -69,15 +69,22 @@ impl Middleware for ThrottleMiddleware {
             None
         })?;
 
-        let category = rate_limiter.category();
-        let limit = rate_limiter.find_rate_limit(Metric::EdenApiQps)?;
+        let identities = state.try_borrow::<MetadataState>()?.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,
             RateLimitStatus::Tracked => false,
             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;

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<crate::RateLimit> {
-        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<MononokeIdentitySet>,
+        main_id: Option<&str>,
+    ) -> Option<crate::RateLimit> {
+        // 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
     }
 }
 

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<RateLimit>;
+    fn find_rate_limit(
+        &self,
+        metric: Metric,
+        identities: Option<MononokeIdentitySet>,
+        main_id: Option<&str>,
+    ) -> Option<RateLimit>;
 }
 
 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<LoadShedLimit>,
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
 pub struct RateLimit {
     pub body: RateLimitBody,
     #[allow(dead_code)]
-    target: Option<Target>,
+    pub target: Option<Target>,
     #[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<i32> for SlicePct {
@@ -253,7 +258,7 @@ impl TryFrom<i32> 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)
+        );
+    }
 }

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<RateLimit> {
+    fn find_rate_limit(
+        &self,
+        metric: Metric,
+        identities: Option<MononokeIdentitySet>,
+        main_id: Option<&str>,
+    ) -> Option<crate::RateLimit> {
         None
     }
 }