reqwest_drive/

throttle_middleware.rs

use crate::{
    DriveCache,
    cache_middleware::{CacheBust, CacheBypass},
};
use async_trait::async_trait;
use http::Extensions;
use rand::RngExt;
use reqwest::{Request, Response};
use reqwest_middleware::{Error, Middleware, Next};
use std::sync::Arc;
use tokio::sync::Semaphore;
use tokio::time::{Duration, sleep};

/// Defines the throttling and backoff behavior for handling HTTP requests.
///
/// This policy determines the **rate limiting strategy** used for outgoing requests,
/// including fixed delays, adaptive backoff, and retry settings.
#[derive(Clone, Debug)]
pub struct ThrottlePolicy {
    /// The base delay (in milliseconds) applied before making a request.
    ///
    /// This ensures a **minimum delay** between consecutive requests.
    pub base_delay_ms: u64,

    /// The maximum random jitter (in milliseconds) added to the backoff delay.
    ///
    /// This prevents synchronization issues when multiple clients are making requests,
    /// reducing the likelihood of rate-limit collisions.
    pub adaptive_jitter_ms: u64,

    /// The maximum number of concurrent requests allowed at any given time.
    ///
    /// This controls **parallel request execution**, ensuring that no more than
    /// `max_concurrent` requests are in-flight simultaneously.
    pub max_concurrent: usize,

    /// The maximum number of retries allowed in case of failed requests.
    ///
    /// If a request fails (e.g., due to a **server error or rate limiting**),
    /// it will be retried up to `max_retries` times with exponential backoff.
    pub max_retries: usize,
}

impl Default for ThrottlePolicy {
    /// Provides a sensible default throttling policy.
    ///
    /// This default configuration is suitable for most API use cases and includes:
    /// - A **base delay** of 500ms between requests.
    /// - A **random jitter** of up to 250ms to avoid synchronization issues.
    /// - A **maximum of 5 concurrent requests** to prevent excessive load.
    /// - A **maximum of 3 retries** for failed requests.
    ///
    /// # Returns
    ///
    /// A `ThrottlePolicy` instance with preconfigured defaults.
    fn default() -> Self {
        Self {
            base_delay_ms: 500,      // 500ms base delay between requests
            adaptive_jitter_ms: 250, // Add up to 250ms random jitter
            max_concurrent: 5,       // Allow up to 5 concurrent requests
            max_retries: 3,          // Retry failed requests up to 3 times
        }
    }
}

/// Implements a throttling and exponential backoff middleware for HTTP requests.
///
/// This middleware **limits request concurrency** and applies **adaptive delays**
/// between retries, helping to prevent rate-limiting issues when interacting
/// with APIs that enforce request quotas.
///
/// It can run in two modes:
/// - **Cache-aware mode** via [`DriveThrottleBackoff::new`], where cached requests can bypass throttling.
/// - **Throttle-only mode** via [`DriveThrottleBackoff::without_cache`], where all requests are throttled.
///
/// Requests are throttled using a **semaphore-based** approach, ensuring that
/// the maximum number of concurrent requests does not exceed `max_concurrent`.
///
/// If a request fails, it enters a **retry loop** where each retry is delayed
/// according to an **exponential backoff strategy**.
pub struct DriveThrottleBackoff {
    /// Semaphore controlling the maximum number of concurrent requests.
    semaphore: Arc<Semaphore>,

    /// Defines the backoff and throttling behavior.
    policy: ThrottlePolicy,

    /// Optional cache layer for detecting previously cached responses.
    cache: Option<Arc<DriveCache>>,
}

impl DriveThrottleBackoff {
    /// Creates a new `DriveThrottleBackoff` middleware with the specified throttling policy.
    ///
    /// # Arguments
    ///
    /// * `policy` - The throttling configuration defining concurrency limits, delays, and retry behavior.
    /// * `cache` - The shared cache instance used for **detecting previously cached requests**.
    ///
    /// # Returns
    ///
    /// A new instance of `DriveThrottleBackoff`.
    pub fn new(policy: ThrottlePolicy, cache: Arc<DriveCache>) -> Self {
        Self {
            semaphore: Arc::new(Semaphore::new(policy.max_concurrent)),
            policy,
            cache: Some(cache),
        }
    }

    /// Creates a new `DriveThrottleBackoff` middleware without any cache integration.
    ///
    /// In this mode, every request is throttled based on the configured policy,
    /// and no cache checks are performed.
    pub fn without_cache(policy: ThrottlePolicy) -> Self {
        Self {
            semaphore: Arc::new(Semaphore::new(policy.max_concurrent)),
            policy,
            cache: None,
        }
    }

    pub fn available_permits(&self) -> usize {
        self.semaphore.available_permits()
    }
}

#[async_trait]
impl Middleware for DriveThrottleBackoff {
    /// Handles throttling and retry logic for HTTP requests.
    ///
    /// This method:
    /// 1. **Optionally checks the cache**: In cache-aware mode, cached requests bypass throttling.
    /// 2. **Enforces concurrency limits**: Ensures no more than `max_concurrent` requests are in flight.
    /// 3. **Applies an initial delay** before sending the request.
    /// 4. **Retries failed requests**: Uses **exponential backoff** with jitter for failed requests.
    ///
    /// # Arguments
    ///
    /// * `req` - The incoming HTTP request.
    /// * `extensions` - A mutable reference to request extensions, used for tracking metadata.
    /// * `next` - The next middleware in the request chain.
    ///
    /// # Returns
    ///
    /// A `Result<Response, Error>` containing either:
    /// - A successfully processed response.
    /// - An error if the request failed after exhausting all retries.
    ///
    /// # Behavior
    ///
    /// - In cache-aware mode, if the request is **already cached**, the middleware immediately forwards it.
    /// - In throttle-only mode, cache checks are skipped.
    /// - If **throttling is required**, it waits according to the configured delay.
    /// - If a request fails, **exponential backoff** is applied before retrying.
    async fn handle(
        &self,
        req: Request,
        extensions: &mut Extensions,
        next: Next<'_>,
    ) -> Result<Response, Error> {
        let url = req.url().to_string();
        let bypass_cache = extensions
            .get::<CacheBypass>()
            .map(|flag| flag.0)
            .unwrap_or(false);
        let bust_cache = extensions
            .get::<CacheBust>()
            .map(|flag| flag.0)
            .unwrap_or(false);

        let cache_key = format!("{} {}", req.method(), &url);

        if !bypass_cache
            && !bust_cache
            && let Some(cache) = &self.cache
        {
            if cache.is_cached(&req).await {
                tracing::debug!("Using cache for: {}", &cache_key);

                return next.run(req, extensions).await;
            } else {
                tracing::debug!("No cache found for: {}", &cache_key);
            }
        }

        // Use a custom throttle policy if provided, otherwise default to `self.policy`
        let custom_policy = extensions.get::<ThrottlePolicy>().cloned();
        let policy = custom_policy.unwrap_or_else(|| self.policy.clone()); // Use override if available

        // Log if the permit is not immediately available
        if self.semaphore.available_permits() == 0 {
            tracing::debug!("Waiting for permit... ({} in use)", policy.max_concurrent);
        }

        // Acquire the permit and log when granted
        let permit = self
            .semaphore
            .acquire()
            .await
            .map_err(|e| Error::Middleware(e.into()))?;

        tracing::debug!(
            "Permit granted: {} ({} permits left)",
            cache_key,
            self.semaphore.available_permits()
        );

        // Hold the permit until this function completes
        let _permit_guard = permit;

        sleep(Duration::from_millis(policy.base_delay_ms)).await;

        let mut attempt = 0;

        loop {
            let req_clone = req.try_clone().expect("Request cloning failed");
            let result = next.clone().run(req_clone, extensions).await;

            match result {
                Ok(resp) if resp.status().is_success() => return Ok(resp),
                result if attempt >= policy.max_retries => return result,
                _ => {
                    attempt += 1;

                    let backoff_duration = {
                        let mut rng = rand::rng();
                        Duration::from_millis(
                            policy.base_delay_ms * 2u64.pow(attempt as u32)
                                + rng.random_range(0..=policy.adaptive_jitter_ms),
                        )
                    };

                    tracing::debug!(
                        "Retry {}/{} for URL {} after {} ms",
                        attempt,
                        policy.max_retries,
                        url,
                        backoff_duration.as_millis()
                    );

                    sleep(backoff_duration).await;

                    if attempt >= policy.max_retries {
                        break;
                    }
                }
            }
        }

        next.run(req, extensions).await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use reqwest_middleware::ClientBuilder;

    #[test]
    fn throttle_policy_default_values_are_stable() {
        let policy = ThrottlePolicy::default();

        assert_eq!(policy.base_delay_ms, 500);
        assert_eq!(policy.adaptive_jitter_ms, 250);
        assert_eq!(policy.max_concurrent, 5);
        assert_eq!(policy.max_retries, 3);
    }

    #[tokio::test]
    async fn closed_semaphore_returns_middleware_error() {
        let throttle = Arc::new(DriveThrottleBackoff::without_cache(ThrottlePolicy {
            base_delay_ms: 1,
            adaptive_jitter_ms: 0,
            max_concurrent: 1,
            max_retries: 0,
        }));

        // Force semaphore acquire to fail so we cover the map_err path.
        throttle.semaphore.close();

        let client = ClientBuilder::new(reqwest::Client::new())
            .with_arc(throttle)
            .build();

        let error = client
            .get("https://example.test/closed-semaphore")
            .send()
            .await
            .expect_err("closed semaphore should return middleware error");

        assert!(
            matches!(error, reqwest_middleware::Error::Middleware(_)),
            "expected middleware error variant, got: {:?}",
            error
        );
    }
}