Skip to main content

hypersync_client/
rate_limit.rs

1/// Rate limit information extracted from response headers.
2///
3/// Envoy's rate limiter returns these headers in the IETF draft format:
4/// - `x-ratelimit-limit`: e.g. `"50, 50;w=60"` (total quota for the window)
5/// - `x-ratelimit-remaining`: e.g. `"40"` (remaining budget in window)
6/// - `x-ratelimit-reset`: e.g. `"41"` (seconds until window resets)
7/// - `x-ratelimit-cost`: e.g. `"10"` (budget consumed per request)
8#[derive(Debug, Clone, Default)]
9pub struct RateLimitInfo {
10    /// Total request quota for the current window.
11    ///
12    /// Parsed from `x-ratelimit-limit`. For IETF draft format like `"50, 50;w=60"`,
13    /// the first integer before the comma is used.
14    pub limit: Option<u64>,
15    /// Remaining budget in the current window.
16    ///
17    /// Parsed from `x-ratelimit-remaining`. Note this is budget units, not request count.
18    /// Divide by [`cost`](Self::cost) to get the number of requests remaining.
19    pub remaining: Option<u64>,
20    /// Seconds until the rate limit window resets.
21    ///
22    /// Parsed from `x-ratelimit-reset`.
23    pub reset_secs: Option<u64>,
24    /// Budget consumed per request.
25    ///
26    /// Parsed from `x-ratelimit-cost`. For example, if `limit` is 50 and `cost` is 10,
27    /// you can make 5 requests per window.
28    pub cost: Option<u64>,
29}
30
31impl std::fmt::Display for RateLimitInfo {
32    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33        let mut parts = Vec::new();
34        if let (Some(remaining), Some(limit)) = (self.remaining, self.limit) {
35            let cost = self.cost.unwrap_or(1);
36            parts.push(format!(
37                "remaining={}/{} reqs ({}/{} budget, cost={})",
38                remaining / cost,
39                limit / cost,
40                remaining,
41                limit,
42                cost
43            ));
44        } else {
45            if let Some(remaining) = self.remaining {
46                parts.push(format!("remaining={remaining}"));
47            }
48            if let Some(limit) = self.limit {
49                parts.push(format!("limit={limit}"));
50            }
51        }
52        if let Some(reset) = self.reset_secs {
53            parts.push(format!("resets_in={reset}s"));
54        }
55        write!(f, "{}", parts.join(", "))
56    }
57}
58
59impl RateLimitInfo {
60    /// Extracts rate limit information from HTTP response headers.
61    ///
62    /// All parsing is best-effort: missing or unparseable headers become `None`.
63    pub(crate) fn from_response(res: &reqwest::Response) -> Self {
64        Self {
65            limit: Self::parse_limit_header(res),
66            remaining: Self::parse_u64_header(res, "x-ratelimit-remaining"),
67            reset_secs: Self::parse_u64_header(res, "x-ratelimit-reset"),
68            cost: Self::parse_u64_header(res, "x-ratelimit-cost"),
69        }
70    }
71
72    /// Returns `true` if the rate limit quota has been exhausted.
73    pub fn is_rate_limited(&self) -> bool {
74        self.remaining == Some(0)
75    }
76
77    /// Returns the suggested number of seconds to wait before making another request.
78    pub fn suggested_wait_secs(&self) -> Option<u64> {
79        self.reset_secs
80    }
81
82    /// Parses `x-ratelimit-limit` which uses IETF draft format: `"60, 60;w=60"`.
83    /// Extracts the first integer before the comma.
84    fn parse_limit_header(res: &reqwest::Response) -> Option<u64> {
85        let value = res.headers().get("x-ratelimit-limit")?.to_str().ok()?;
86        // Take first value before comma: "60, 60;w=60" -> "60"
87        let first = value.split(',').next()?.trim();
88        first.parse().ok()
89    }
90
91    /// Parses a simple u64 header value.
92    fn parse_u64_header(res: &reqwest::Response, name: &str) -> Option<u64> {
93        res.headers().get(name)?.to_str().ok()?.trim().parse().ok()
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    #[test]
102    fn test_is_rate_limited() {
103        let info = RateLimitInfo {
104            remaining: Some(0),
105            ..Default::default()
106        };
107        assert!(info.is_rate_limited());
108
109        let info = RateLimitInfo {
110            remaining: Some(5),
111            ..Default::default()
112        };
113        assert!(!info.is_rate_limited());
114
115        let info = RateLimitInfo::default();
116        assert!(!info.is_rate_limited());
117    }
118
119    #[test]
120    fn test_from_response_header_case_insensitive() {
121        // Build an http::Response with mixed-case headers, then convert to reqwest::Response.
122        // This confirms that HeaderMap normalizes names so our lowercase lookups match.
123        let http_resp = http::Response::builder()
124            .header("X-RateLimit-Remaining", "42")
125            .header("X-RATELIMIT-RESET", "30")
126            .header("X-Ratelimit-Limit", "100, 100;w=60")
127            .header("X-Ratelimit-Cost", "10")
128            .body("")
129            .unwrap();
130        let resp: reqwest::Response = http_resp.into();
131
132        let info = RateLimitInfo::from_response(&resp);
133        assert_eq!(info.limit, Some(100));
134        assert_eq!(info.remaining, Some(42));
135        assert_eq!(info.reset_secs, Some(30));
136        assert_eq!(info.cost, Some(10));
137    }
138
139    #[test]
140    fn test_suggested_wait_secs() {
141        // Uses reset_secs
142        let info = RateLimitInfo {
143            reset_secs: Some(30),
144            ..Default::default()
145        };
146        assert_eq!(info.suggested_wait_secs(), Some(30));
147
148        // None when no info
149        let info = RateLimitInfo::default();
150        assert_eq!(info.suggested_wait_secs(), None);
151    }
152
153    #[test]
154    fn test_display_full() {
155        let info = RateLimitInfo {
156            limit: Some(50),
157            remaining: Some(0),
158            reset_secs: Some(59),
159            cost: Some(10),
160        };
161        assert_eq!(
162            info.to_string(),
163            "remaining=0/5 reqs (0/50 budget, cost=10), resets_in=59s"
164        );
165    }
166
167    #[test]
168    fn test_display_partial() {
169        let info = RateLimitInfo {
170            remaining: Some(3),
171            reset_secs: Some(30),
172            ..Default::default()
173        };
174        assert_eq!(info.to_string(), "remaining=3, resets_in=30s");
175    }
176
177    #[test]
178    fn test_display_empty() {
179        let info = RateLimitInfo::default();
180        assert_eq!(info.to_string(), "");
181    }
182}