1#[cfg(feature = "serde")]
18use serde::{Deserialize, Serialize};
19
20#[derive(Debug, Clone, PartialEq, Eq)]
33#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
34#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
35#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
36#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
37#[cfg_attr(feature = "proptest", derive(proptest_derive::Arbitrary))]
38pub struct RateLimitInfo {
39 pub limit: u64,
41
42 pub remaining: u64,
44
45 pub reset: u64,
47
48 #[cfg_attr(
50 feature = "serde",
51 serde(default, skip_serializing_if = "Option::is_none")
52 )]
53 pub retry_after: Option<u64>,
54}
55
56impl RateLimitInfo {
57 #[must_use]
70 pub fn new(limit: u64, remaining: u64, reset: u64) -> Self {
71 Self {
72 limit,
73 remaining,
74 reset,
75 retry_after: None,
76 }
77 }
78
79 #[must_use]
90 pub fn retry_after(mut self, seconds: u64) -> Self {
91 self.retry_after = Some(seconds);
92 self
93 }
94
95 #[must_use]
109 pub fn is_exceeded(&self) -> bool {
110 self.remaining == 0
111 }
112}
113
114#[cfg(feature = "http")]
119mod http_impl {
120 use super::RateLimitInfo;
121 #[cfg(not(feature = "std"))]
122 use alloc::string::ToString;
123 use http::{HeaderMap, HeaderValue};
124
125 pub const HEADER_LIMIT: &str = "x-ratelimit-limit";
127 pub const HEADER_REMAINING: &str = "x-ratelimit-remaining";
128 pub const HEADER_RESET: &str = "x-ratelimit-reset";
129 pub const HEADER_RETRY_AFTER: &str = "retry-after";
130
131 impl RateLimitInfo {
132 pub fn inject_headers(&self, headers: &mut HeaderMap) {
138 headers.insert(
140 HEADER_LIMIT,
141 HeaderValue::from_str(&self.limit.to_string())
142 .expect("u64 decimal is always a valid header value"),
143 );
144 headers.insert(
145 HEADER_REMAINING,
146 HeaderValue::from_str(&self.remaining.to_string())
147 .expect("u64 decimal is always a valid header value"),
148 );
149 headers.insert(
150 HEADER_RESET,
151 HeaderValue::from_str(&self.reset.to_string())
152 .expect("u64 decimal is always a valid header value"),
153 );
154 if let Some(secs) = self.retry_after {
155 headers.insert(
156 HEADER_RETRY_AFTER,
157 HeaderValue::from_str(&secs.to_string())
158 .expect("u64 decimal is always a valid header value"),
159 );
160 }
161 }
162
163 #[must_use]
169 pub fn from_headers(headers: &HeaderMap) -> Option<Self> {
170 let parse = |name| -> Option<u64> {
171 headers
172 .get(name)
173 .and_then(|v| v.to_str().ok())
174 .and_then(|s| s.parse().ok())
175 };
176
177 let limit = parse(HEADER_LIMIT)?;
178 let remaining = parse(HEADER_REMAINING)?;
179 let reset = parse(HEADER_RESET)?;
180 let retry_after = parse(HEADER_RETRY_AFTER);
181
182 Some(Self {
183 limit,
184 remaining,
185 reset,
186 retry_after,
187 })
188 }
189 }
190}
191
192#[cfg(test)]
197mod tests {
198 use super::*;
199
200 #[test]
205 fn new_sets_fields() {
206 let info = RateLimitInfo::new(100, 42, 1_700_000_000);
207 assert_eq!(info.limit, 100);
208 assert_eq!(info.remaining, 42);
209 assert_eq!(info.reset, 1_700_000_000);
210 assert!(info.retry_after.is_none());
211 }
212
213 #[test]
214 fn retry_after_builder() {
215 let info = RateLimitInfo::new(100, 0, 1_700_000_000).retry_after(60);
216 assert_eq!(info.retry_after, Some(60));
217 }
218
219 #[test]
224 fn is_exceeded_when_remaining_zero() {
225 let info = RateLimitInfo::new(100, 0, 1_700_000_000);
226 assert!(info.is_exceeded());
227 }
228
229 #[test]
230 fn is_not_exceeded_when_remaining_nonzero() {
231 let info = RateLimitInfo::new(100, 1, 1_700_000_000);
232 assert!(!info.is_exceeded());
233 }
234
235 #[cfg(feature = "serde")]
240 #[test]
241 fn serde_round_trip_without_retry_after() {
242 let info = RateLimitInfo::new(100, 50, 1_700_000_000);
243 let json = serde_json::to_value(&info).unwrap();
244 assert_eq!(json["limit"], 100);
245 assert_eq!(json["remaining"], 50);
246 assert_eq!(json["reset"], 1_700_000_000_u64);
247 assert!(json.get("retry_after").is_none());
248 let back: RateLimitInfo = serde_json::from_value(json).unwrap();
249 assert_eq!(back, info);
250 }
251
252 #[cfg(feature = "serde")]
253 #[test]
254 fn serde_round_trip_with_retry_after() {
255 let info = RateLimitInfo::new(100, 0, 1_700_000_000).retry_after(30);
256 let json = serde_json::to_value(&info).unwrap();
257 assert_eq!(json["retry_after"], 30);
258 let back: RateLimitInfo = serde_json::from_value(json).unwrap();
259 assert_eq!(back, info);
260 }
261
262 #[cfg(feature = "serde")]
263 #[test]
264 fn serde_omits_retry_after_when_none() {
265 let info = RateLimitInfo::new(10, 5, 999);
266 let json = serde_json::to_value(&info).unwrap();
267 assert!(json.get("retry_after").is_none());
268 }
269
270 #[cfg(feature = "http")]
275 mod http_tests {
276 use super::*;
277 use http::HeaderMap;
278
279 #[test]
280 fn inject_and_extract_without_retry_after() {
281 let info = RateLimitInfo::new(200, 150, 1_700_000_000);
282 let mut headers = HeaderMap::new();
283 info.inject_headers(&mut headers);
284
285 let extracted = RateLimitInfo::from_headers(&headers).unwrap();
286 assert_eq!(extracted.limit, 200);
287 assert_eq!(extracted.remaining, 150);
288 assert_eq!(extracted.reset, 1_700_000_000);
289 assert!(extracted.retry_after.is_none());
290 }
291
292 #[test]
293 fn inject_and_extract_with_retry_after() {
294 let info = RateLimitInfo::new(100, 0, 1_700_000_000).retry_after(45);
295 let mut headers = HeaderMap::new();
296 info.inject_headers(&mut headers);
297
298 let extracted = RateLimitInfo::from_headers(&headers).unwrap();
299 assert_eq!(extracted.retry_after, Some(45));
300 }
301
302 #[test]
303 fn from_headers_returns_none_on_missing_required_header() {
304 let headers = HeaderMap::new();
305 assert!(RateLimitInfo::from_headers(&headers).is_none());
306 }
307
308 #[test]
309 fn from_headers_returns_none_on_invalid_value() {
310 let mut headers = HeaderMap::new();
311 headers.insert("x-ratelimit-limit", "not-a-number".parse().unwrap());
312 headers.insert("x-ratelimit-remaining", "5".parse().unwrap());
313 headers.insert("x-ratelimit-reset", "999".parse().unwrap());
314 assert!(RateLimitInfo::from_headers(&headers).is_none());
315 }
316 }
317}