rate limit via env var

Author: cameron1024

catalyst-toolbox/src/http/mod.rs

@@ -2,15 +2,29 @@ use std::marker::PhantomData;
 
 use ::reqwest::{blocking::Response, StatusCode};
 use color_eyre::eyre::Result;
+use log::warn;
 use serde::Deserialize;
 
 use self::{rate_limit::RateLimitClient, reqwest::ReqwestClient};
 
 mod rate_limit;
 mod reqwest;
 
+const RATE_LIMIT_ENV_VAR: &str = "CATALYST_RATE_LIMIT_MS";
+
 pub fn default_http_client(api_key: &str) -> impl HttpClient {
-    RateLimitClient::new(ReqwestClient::new(api_key), 10_000)
+    let rate_limit = match std::env::var(RATE_LIMIT_ENV_VAR).map(|s| s.parse::<u64>()) {
+        Ok(Ok(rate_limit)) => rate_limit,
+        Ok(Err(_)) => {
+            warn!(
+                "{} could not be parsed as a u64, defaulting to no rate-limiting",
+                RATE_LIMIT_ENV_VAR
+            );
+            0
+        }
+        _ => 0,
+    };
+    RateLimitClient::new(ReqwestClient::new(api_key), rate_limit)
 }
 
 #[cfg(test)]
@@ -21,7 +35,7 @@ fn test_default_client_send_sync() {
 }
 
 /// Types which can make HTTP requests
-pub trait HttpClient:  Send + Sync + 'static {
+pub trait HttpClient: Send + Sync + 'static {
     fn get<T>(&self, path: &str) -> Result<HttpResponse<T>>
     where
         T: for<'a> Deserialize<'a>;

catalyst-toolbox/src/http/rate_limit.rs

@@ -14,13 +14,17 @@ use super::{HttpClient, HttpResponse};
 
 pub struct RateLimitClient<T: HttpClient> {
     inner: T,
-    limiter: RateLimiter<NotKeyed, InMemoryState, DefaultClock>,
+    limiter: Option<RateLimiter<NotKeyed, InMemoryState, DefaultClock>>,
 }
 
 impl<T: HttpClient> RateLimitClient<T> {
     pub fn new(inner: T, request_interval_ms: u64) -> Self {
-        let quota = Quota::with_period(Duration::from_millis(request_interval_ms)).unwrap();
-        let limiter = RateLimiter::direct(quota);
+        let limiter = if request_interval_ms == 0 {
+            None
+        } else {
+            let quota = Quota::with_period(Duration::from_millis(request_interval_ms)).unwrap();
+            Some(RateLimiter::direct(quota))
+        };
         Self { inner, limiter }
     }
 }
@@ -30,10 +34,12 @@ impl<T: HttpClient> HttpClient for RateLimitClient<T> {
     where
         S: for<'a> Deserialize<'a>,
     {
-        while let Err(e) = self.limiter.check() {
-            let time = e.wait_time_from(Instant::now());
-            debug!("waiting for {time:?}");
-            std::thread::sleep(time);
+        if let Some(limiter) = self.limiter {
+            while let Err(e) = limiter.check() {
+                let time = e.wait_time_from(Instant::now());
+                debug!("waiting for rate limit: {time:?}");
+                std::thread::sleep(time);
+            }
         }
         self.inner.get(path)
     }