fix: handle rate limiting

Signed-off-by: Rahul Baradol <[email protected]>

Author: Rahul-Baradol

crates/ofrep/Cargo.toml

@@ -6,7 +6,7 @@ edition = "2024"
 [dev-dependencies]
 wiremock = "0.6.3"
 test-log = { version = "0.2", features = ["trace"] }
-tokio = { version = "1.45", features = ["full"] }
+serial_test = "3.2.0"
 
 [dependencies]
 async-trait = "0.1.88"
@@ -16,8 +16,10 @@ reqwest = { version = "0.12", default-features = false, features = [
   "stream",
   "rustls-tls",
 ] }
-serde = "1.0.219"
 serde_json = "1.0.140"
 tracing = "0.1.41"
 thiserror = "2.0"
 anyhow = "1.0.98"
+chrono = "0.4"
+once_cell = "1.18"
+tokio = { version = "1.45", features = ["full"] }

crates/ofrep/src/resolver.rs

@@ -1,16 +1,22 @@
 use async_trait::async_trait;
+use chrono::{DateTime, Duration, Utc};
+use once_cell::sync::Lazy;
 use open_feature::provider::{FeatureProvider, ProviderMetadata, ResolutionDetails};
 use open_feature::{
     EvaluationContext, EvaluationContextFieldValue, EvaluationError, EvaluationErrorCode,
     EvaluationResult, StructValue, Value,
 };
 use reqwest::Client;
 use reqwest::StatusCode;
+use reqwest::header::RETRY_AFTER;
 use std::any;
+use tokio::sync::Mutex;
 use tracing::{debug, error, instrument};
 
 use crate::OfrepOptions;
 
+static CURRENT_RETRY_AFTER: Lazy<Mutex<DateTime<Utc>>> = Lazy::new(|| Mutex::new(Utc::now()));
+
 #[derive(Debug)]
 pub struct Resolver {
     base_url: String,
@@ -31,13 +37,51 @@ impl Resolver {
         }
     }
 
+    async fn parse_retry_after(retry_after: &str) -> DateTime<Utc> {
+        let now = Utc::now();
+
+        if retry_after.trim().is_empty() {
+            return now;
+        }
+
+        if let Ok(seconds) = retry_after.trim().parse::<i64>() {
+            return now + Duration::seconds(seconds);
+        }
+
+        if let Ok(parsed_date) = retry_after.trim().parse::<DateTime<Utc>>() {
+            return parsed_date.with_timezone(&Utc);
+        }
+
+        debug!("Failed to parse Retry-After header : {}", retry_after);
+        now
+    }
+
+    async fn update_retry_after(new_retry_after: DateTime<Utc>) {
+        let mut retry_after = CURRENT_RETRY_AFTER.lock().await;
+        *retry_after = new_retry_after;
+    }
+
+    async fn is_rate_limit_exceeded() -> bool {
+        let retry_after = CURRENT_RETRY_AFTER.lock().await;
+        Utc::now() < *retry_after
+    }
+
     #[instrument(skip(self, evaluation_context), fields(flag_key = %flag_key))]
     async fn resolve_value<T: std::fmt::Debug>(
         &self,
         flag_key: &str,
         evaluation_context: &EvaluationContext,
         convertor: fn(serde_json::Value) -> Option<T>,
     ) -> EvaluationResult<ResolutionDetails<T>> {
+        if Resolver::is_rate_limit_exceeded().await {
+            return Err(EvaluationError {
+                code: EvaluationErrorCode::General("Rate limit exceeded".to_string()),
+                message: Some(
+                    "Rate limit exceeded. Please wait before making another request.".to_string(),
+                ),
+            });
+        }
+
         debug!("Resolving {} flag", std::any::type_name::<T>());
         let payload = serde_json::json!({
             "context": context_to_json(evaluation_context)
@@ -87,6 +131,28 @@ impl Resolver {
                     message: Some(format!("Flag: {flag_key} not found")),
                 });
             }
+            StatusCode::TOO_MANY_REQUESTS => {
+                let header_retry_after: Option<&str> = response
+                    .headers()
+                    .get(RETRY_AFTER)
+                    .and_then(|value| value.to_str().ok());
+
+                if let Some(header_retry_after) = header_retry_after {
+                    let new_retry_after: DateTime<Utc> =
+                        Resolver::parse_retry_after(header_retry_after).await;
+                    Resolver::update_retry_after(new_retry_after).await;
+                } else {
+                    debug!("Couldn't parse the retry-after header.");
+                    let mut retry_after = CURRENT_RETRY_AFTER.lock().await;
+                    *retry_after = Utc::now();
+                }
+
+                let retry_after = CURRENT_RETRY_AFTER.lock().await;
+                return Err(EvaluationError {
+                    code: EvaluationErrorCode::General("Rate limit exceeded".to_string()),
+                    message: Some(format!("Rate limit exceeded. Retry after {}", *retry_after)),
+                });
+            }
             _ => {
                 let result = response.json::<serde_json::Value>().await.map_err(|e| {
                     error!(error = %e, "Failed to parse {} response", any::type_name::<T>());
@@ -242,9 +308,15 @@ mod tests {
     use super::*;
     use serde_json::json;
     use test_log::test;
+    use tokio::time::{Duration, sleep};
     use wiremock::matchers::{method, path};
     use wiremock::{Mock, MockServer, ResponseTemplate};
 
+    async fn reset_states() {
+        let mut retry_after = CURRENT_RETRY_AFTER.lock().await;
+        *retry_after = Utc::now();
+    }
+
     async fn setup_mock_server() -> (MockServer, Resolver) {
         let mock_server = MockServer::start().await;
         let options = OfrepOptions {
@@ -256,7 +328,9 @@ mod tests {
     }
 
     #[test(tokio::test)]
+    #[serial_test::serial]
     async fn test_resolve_bool_value() {
+        reset_states().await;
         let (mock_server, resolver) = setup_mock_server().await;
 
         Mock::given(method("POST"))
@@ -281,7 +355,9 @@ mod tests {
     }
 
     #[test(tokio::test)]
+    #[serial_test::serial]
     async fn test_resolve_string_value() {
+        reset_states().await;
         let (mock_server, resolver) = setup_mock_server().await;
 
         Mock::given(method("POST"))
@@ -306,7 +382,9 @@ mod tests {
     }
 
     #[test(tokio::test)]
+    #[serial_test::serial]
     async fn test_resolve_float_value() {
+        reset_states().await;
         let (mock_server, resolver) = setup_mock_server().await;
 
         Mock::given(method("POST"))
@@ -331,7 +409,9 @@ mod tests {
     }
 
     #[test(tokio::test)]
+    #[serial_test::serial]
     async fn test_resolve_int_value() {
+        reset_states().await;
         let (mock_server, resolver) = setup_mock_server().await;
 
         Mock::given(method("POST"))
@@ -356,7 +436,9 @@ mod tests {
     }
 
     #[test(tokio::test)]
+    #[serial_test::serial]
     async fn test_resolve_struct_value() {
+        reset_states().await;
         let (mock_server, resolver) = setup_mock_server().await;
 
         Mock::given(method("POST"))
@@ -401,7 +483,9 @@ mod tests {
     }
 
     #[test(tokio::test)]
+    #[serial_test::serial]
     async fn test_error_400() {
+        reset_states().await;
         let (mock_server, resolver) = setup_mock_server().await;
 
         Mock::given(method("POST"))
@@ -446,7 +530,9 @@ mod tests {
     }
 
     #[test(tokio::test)]
+    #[serial_test::serial]
     async fn test_error_401() {
+        reset_states().await;
         let (mock_server, resolver) = setup_mock_server().await;
 
         Mock::given(method("POST"))
@@ -492,7 +578,9 @@ mod tests {
     }
 
     #[test(tokio::test)]
+    #[serial_test::serial]
     async fn test_error_403() {
+        reset_states().await;
         let (mock_server, resolver) = setup_mock_server().await;
 
         Mock::given(method("POST"))
@@ -538,7 +626,9 @@ mod tests {
     }
 
     #[test(tokio::test)]
+    #[serial_test::serial]
     async fn test_error_404() {
+        reset_states().await;
         let (mock_server, resolver) = setup_mock_server().await;
 
         Mock::given(method("POST"))
@@ -596,4 +686,67 @@ mod tests {
             "Flag: test-flag not found"
         );
     }
+
+    #[test(tokio::test)]
+    #[serial_test::serial]
+    async fn test_error_429() {
+        reset_states().await;
+        let (mock_server, resolver) = setup_mock_server().await;
+
+        Mock::given(method("POST"))
+            .and(path("/ofrep/v1/evaluate/flags/test-flag"))
+            .respond_with(
+                ResponseTemplate::new(429)
+                    .insert_header("Retry-After", "3")
+                    .set_body_json(json!({})),
+            )
+            .mount(&mock_server)
+            .await;
+
+        let context = EvaluationContext::default();
+
+        let result_bool = resolver.resolve_bool_value("test-flag", &context).await;
+        let result_bool_2 = resolver.resolve_bool_value("test-flag", &context).await;
+        
+        assert!(result_bool.is_err());
+        let result_bool_error = result_bool.unwrap_err();
+        assert_eq!(
+            result_bool_error.code,
+            EvaluationErrorCode::General("Rate limit exceeded".to_string())
+        );
+        assert!(
+            result_bool_error
+                .message
+                .unwrap()
+                .starts_with("Rate limit exceeded. Retry after")
+        );
+
+        assert!(result_bool_2.is_err());
+        let result_bool_error_2 = result_bool_2.unwrap_err();
+        assert_eq!(
+            result_bool_error_2.code,
+            EvaluationErrorCode::General("Rate limit exceeded".to_string())
+        );
+        assert_eq!(
+            result_bool_error_2.message.unwrap(),
+            "Rate limit exceeded. Please wait before making another request."
+        );
+
+        sleep(Duration::from_secs(3)).await;
+
+        let result_bool_3 = resolver.resolve_bool_value("test-flag", &context).await;
+        assert!(result_bool_3.is_err());
+
+        let result_bool_error_3 = result_bool_3.unwrap_err();
+        assert_eq!(
+            result_bool_error_3.code,
+            EvaluationErrorCode::General("Rate limit exceeded".to_string())
+        );
+        assert!(
+            result_bool_error_3
+                .message
+                .unwrap()
+                .starts_with("Rate limit exceeded. Retry after")
+        );
+    }
 }