Skip to main content

brainwires_providers/
http_client.rs

1//! Rate-limited HTTP client wrapper.
2//!
3//! Wraps `reqwest::Client` with an optional [`RateLimiter`] so that all
4//! outgoing requests are throttled according to the configured limit.
5
6use reqwest::{Client, RequestBuilder};
7use std::sync::Arc;
8
9use super::rate_limiter::RateLimiter;
10
11/// An HTTP client that optionally enforces rate limiting on every request.
12#[derive(Clone, Debug)]
13pub struct RateLimitedClient {
14    client: Client,
15    limiter: Option<Arc<RateLimiter>>,
16}
17
18impl RateLimitedClient {
19    /// Create a new client with no rate limiting.
20    pub fn new() -> Self {
21        Self {
22            client: Client::new(),
23            limiter: None,
24        }
25    }
26
27    /// Create a new client with the given requests-per-minute limit.
28    pub fn with_rate_limit(requests_per_minute: u32) -> Self {
29        Self {
30            client: Client::new(),
31            limiter: Some(Arc::new(RateLimiter::new(requests_per_minute))),
32        }
33    }
34
35    /// Create from an existing `reqwest::Client`, adding a rate limiter.
36    pub fn from_client(client: Client, requests_per_minute: Option<u32>) -> Self {
37        Self {
38            client,
39            limiter: requests_per_minute.map(|rpm| Arc::new(RateLimiter::new(rpm))),
40        }
41    }
42
43    /// Get a reference to the inner `reqwest::Client`.
44    pub fn inner(&self) -> &Client {
45        &self.client
46    }
47
48    /// Build a GET request, waiting for rate-limit clearance first.
49    pub async fn get(&self, url: &str) -> RequestBuilder {
50        self.wait_for_token().await;
51        self.client.get(url)
52    }
53
54    /// Build a POST request, waiting for rate-limit clearance first.
55    pub async fn post(&self, url: &str) -> RequestBuilder {
56        self.wait_for_token().await;
57        self.client.post(url)
58    }
59
60    /// Wait for a rate-limit token (no-op if no limiter is configured).
61    async fn wait_for_token(&self) {
62        if let Some(ref limiter) = self.limiter {
63            limiter.acquire().await;
64        }
65    }
66
67    /// Get the current number of available tokens (for diagnostics).
68    /// Returns `None` if no rate limiter is configured.
69    pub fn available_tokens(&self) -> Option<u32> {
70        self.limiter.as_ref().map(|l| l.available_tokens())
71    }
72}
73
74impl Default for RateLimitedClient {
75    fn default() -> Self {
76        Self::new()
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83
84    #[test]
85    fn test_default_client() {
86        let client = RateLimitedClient::new();
87        assert!(client.available_tokens().is_none());
88    }
89
90    #[test]
91    fn test_rate_limited_client() {
92        let client = RateLimitedClient::with_rate_limit(60);
93        assert_eq!(client.available_tokens(), Some(60));
94    }
95
96    #[tokio::test]
97    async fn test_post_consumes_token() {
98        let client = RateLimitedClient::with_rate_limit(10);
99        let _req = client.post("https://example.com").await;
100        assert_eq!(client.available_tokens(), Some(9));
101    }
102}