stop hardcoding time_window

Summary: Previously we had hardcoded the time window to look at when fetching a counter, but instead we hould use the window provided in the config Reviewed By: andreacampi Differential Revision: D66699998 fbshipit-source-id: a1cc10bf83af7cb09d5cd05e599531f80e86d5fc
facebook · Dec 4, 2024 · 062e8cc · 062e8cc
1 parent 23b89bf
commit 062e8cc
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/eden/mononoke/edenapi_service/src/middleware/rate_limiter.rs b/eden/mononoke/edenapi_service/src/middleware/rate_limiter.rs
@@ -80,12 +80,14 @@ impl Middleware for ThrottleMiddleware {
         };
         let counter = build_counter(&ctx, category, EDENAPI_QPS_LIMIT, &main_client_id);
         let max_value = limit.body.raw_config.limit;
+        let time_window = limit.body.window.as_secs() as u32;
 
         match counter_check_and_bump(
             &ctx,
             counter,
             EDENAPI_QPS_LIMIT,
             max_value,
+            time_window,
             enforced,
             hashmap! {"main_client_id" => main_client_id.as_str() },
         )

diff --git a/eden/mononoke/edenapi_service/src/utils/rate_limit.rs b/eden/mononoke/edenapi_service/src/utils/rate_limit.rs
@@ -51,6 +51,7 @@ pub async fn counter_check_and_bump<'a>(
     counter: BoxGlobalTimeWindowCounter,
     rate_limit_name: &'a str,
     max_value: f64,
+    time_window: u32,
     enforced: bool,
     scuba_extras: HashMap<&'a str, &'a str>,
 ) -> Result<(), Error> {
@@ -59,7 +60,7 @@ pub async fn counter_check_and_bump<'a>(
         scuba.add(key, val);
     }
 
-    match timeout(RATELIM_FETCH_TIMEOUT, counter.get(1)).await {
+    match timeout(RATELIM_FETCH_TIMEOUT, counter.get(time_window)).await {
         Ok(Ok(count)) => {
             let new_value = count + 1.0;
             if new_value <= max_value {
@@ -161,6 +162,7 @@ mod test {
             counter,
             rate_limit_name,
             max_value,
+            1,
             true,
             scuba_extras.clone(),
         )
@@ -175,6 +177,7 @@ mod test {
             counter,
             rate_limit_name,
             max_value,
+            1,
             true,
             scuba_extras.clone(),
         )
@@ -189,6 +192,7 @@ mod test {
             counter,
             rate_limit_name,
             max_value,
+            1,
             true,
             scuba_extras.clone(),
         )
@@ -204,6 +208,7 @@ mod test {
             counter,
             rate_limit_name,
             max_value,
+            1,
             false,
             scuba_extras.clone(),
         )