timvw_rust_runtime/
retry.rs

1//! Retries for HTTP APIs: [`RetryPolicy`], a tower retry policy, and [`RetryLayer`], which adds to it a look at the
2//! body of a `429` to tell a rate limit (retried) from an exhausted quota (not retried).
3//!
4//! `Client::default()` already retries with `RetryLayer::new(RetryPolicy::default())`. For another policy, start from
5//! the bare service:
6//!
7//! ```no_run
8//! # #[cfg(feature = "reqwest")]
9//! # fn demo() {
10//! use timvw_rust_runtime::retry::{self, RetryLayer, RetryPolicy};
11//! use timvw_rust_runtime::{Client, ReqwestService};
12//!
13//! let client = Client::from_service(ReqwestService::default())
14//!     .layer(RetryLayer::new(RetryPolicy::new(5)).permanent_429(retry::openai_insufficient_quota));
15//! # }
16//! ```
17
18use std::fmt;
19use std::sync::Arc;
20use std::task::{Context, Poll};
21use std::time::Duration;
22
23use bytes::{Bytes, BytesMut};
24use http::header::RETRY_AFTER;
25use http::{HeaderMap, StatusCode};
26use http_body_util::{BodyDataStream, BodyExt};
27use tower::retry::Policy;
28use tower::{Layer, Service};
29
30use crate::maybe::BoxFuture;
31use crate::timer::{since_epoch, Instant, Sleep};
32use crate::{Body, BoxError, MaybeSend, MaybeSync};
33
34/// The header OpenAI (and others) send next to `Retry-After`, in milliseconds.
35const RETRY_AFTER_MS: &str = "retry-after-ms";
36
37/// How much of a `429` body [`RetryLayer`] reads by default to classify it: 64 KiB.
38pub const DEFAULT_BODY_LIMIT: usize = 64 * 1024;
39
40/// Retries `429 Too Many Requests`, `408 Request Timeout` and `5xx` responses, and failures to connect.
41///
42/// - It waits as long as the server asks, with `retry-after-ms` (milliseconds) or `Retry-After` (seconds or an HTTP
43///   date), up to [`RetryPolicy::max_retry_after`] (60 s by default). The wait is read once into a deadline, counted
44///   from when the response arrived, and the policy sleeps until it: time spent in between (reading a `429` body for
45///   [`RetryLayer::permanent_429`]) is not added. A server that asks for longer than `max_retry_after` gets its
46///   response returned at once: waiting less than it asked would likely fail again, and waiting an hour is rarely
47///   what a caller wants. A date in the past means retry now.
48/// - With [`RetryPolicy::server_directed_429`], a `429` is retried only when the server says when.
49/// - Without those headers it backs off exponentially from [`RetryPolicy::base_delay`] (200 ms), doubling on each retry
50///   up to [`RetryPolicy::max_delay`] (8 s), with jitter: each delay is picked at random between half and all of it, so
51///   that clients that failed together do not retry together. [`RetryPolicy::jitter`] turns that off.
52/// - A request that timed out (a [`TimeoutLayer`](crate::timeout::TimeoutLayer) inside the retries, or a reqwest
53///   timeout) is retried only with [`RetryPolicy::retry_timeouts`]: it may have reached the server.
54/// - It never retries a response marked [`Permanent`] (see [`RetryLayer::permanent_429`]), nor a request whose body is
55///   a stream (a streamed file upload), which cannot be sent again.
56#[derive(Clone, Debug)]
57pub struct RetryPolicy {
58    max_retries: usize,
59    attempts: usize,
60    base_delay: Duration,
61    max_delay: Duration,
62    max_retry_after: Duration,
63    jitter: bool,
64    server_directed_429: bool,
65    retry_timeouts: bool,
66}
67
68impl RetryPolicy {
69    /// A policy that retries at most `max_retries` times (after the first attempt).
70    pub fn new(max_retries: usize) -> Self {
71        Self {
72            max_retries,
73            attempts: 0,
74            base_delay: Duration::from_millis(200),
75            max_delay: Duration::from_secs(8),
76            max_retry_after: Duration::from_secs(60),
77            jitter: true,
78            server_directed_429: false,
79            retry_timeouts: false,
80        }
81    }
82
83    /// The first backoff delay, doubled on each retry (default 200 ms).
84    pub fn base_delay(mut self, delay: Duration) -> Self {
85        self.base_delay = delay;
86        self
87    }
88
89    /// The longest backoff delay (default 8 s). It does not limit how long a server can ask to wait with
90    /// `Retry-After`: that is [`RetryPolicy::max_retry_after`].
91    pub fn max_delay(mut self, delay: Duration) -> Self {
92        self.max_delay = delay;
93        self
94    }
95
96    /// The longest wait a server can ask for with `Retry-After` or `retry-after-ms` (default 60 s). A response that
97    /// asks for longer is not retried but returned.
98    pub fn max_retry_after(mut self, delay: Duration) -> Self {
99        self.max_retry_after = delay;
100        self
101    }
102
103    /// Whether each backoff delay is picked at random between half and all of it (default `true`). The waits a
104    /// server asks for are never jittered.
105    pub fn jitter(mut self, jitter: bool) -> Self {
106        self.jitter = jitter;
107        self
108    }
109
110    /// Whether a `429` is retried only when the server says when, with `Retry-After` or `retry-after-ms` (default
111    /// `false`: every `429` is retried, with backoff when the server does not say when).
112    ///
113    /// Some APIs answer a `429` without advice when waiting will not help (a quota, a request too large for a rate
114    /// limit), and with it when it will. `408`, `5xx` and connect failures are retried either way.
115    pub fn server_directed_429(mut self, server_directed: bool) -> Self {
116        self.server_directed_429 = server_directed;
117        self
118    }
119
120    /// Whether a request that timed out is retried (default `false`), like a failure to connect: with a
121    /// [`TimeoutLayer`](crate::timeout::TimeoutLayer) inside the retries, an attempt that takes too long is given up
122    /// and tried again. A request that timed out may have reached the server, so turn this on for requests that are
123    /// safe to repeat (or that the API deduplicates, with an idempotency key).
124    pub fn retry_timeouts(mut self, retry: bool) -> Self {
125        self.retry_timeouts = retry;
126        self
127    }
128
129    /// The number of retries allowed.
130    pub fn max_retries(&self) -> usize {
131        self.max_retries
132    }
133
134    /// When to retry: the server's deadline, or `now` plus the backoff; `None` when the server asks to wait longer than
135    /// `max_retry_after`.
136    fn deadline(&self, retry_after: Option<Instant>, now: Instant) -> Option<Instant> {
137        if let Some(deadline) = retry_after {
138            return (deadline.saturating_duration_since(now) <= self.max_retry_after)
139                .then_some(deadline);
140        }
141        Some(now + self.backoff())
142    }
143
144    /// The backoff before the next retry, jittered unless turned off.
145    fn backoff(&self) -> Duration {
146        let backoff = self
147            .base_delay
148            .saturating_mul(1u32 << self.attempts.min(16))
149            .min(self.max_delay);
150        if self.jitter {
151            jittered(backoff, random_unit())
152        } else {
153            backoff
154        }
155    }
156}
157
158/// Three retries.
159impl Default for RetryPolicy {
160    fn default() -> Self {
161        Self::new(3)
162    }
163}
164
165/// `backoff` scaled to between half and all of it, by `unit` in `[0, 1)`.
166fn jittered(backoff: Duration, unit: f64) -> Duration {
167    backoff.mul_f64(0.5 + 0.5 * unit)
168}
169
170/// A random number in `[0, 1)`, from the standard library's randomly keyed hasher: good enough to spread retries.
171fn random_unit() -> f64 {
172    use std::hash::{BuildHasher, Hasher};
173    let hash = std::collections::hash_map::RandomState::new()
174        .build_hasher()
175        .finish();
176    (hash >> 11) as f64 / (1u64 << 53) as f64
177}
178
179/// When a response arrived, set by [`RetryLayer`] on a `429` before it reads the body, so that the time spent reading
180/// is not added to the server's `Retry-After`.
181#[derive(Clone, Copy, Debug)]
182struct ReceivedAt(Instant);
183
184/// The clocks a `Retry-After` is read against: `Retry-After` in seconds counts from when the response arrived, an HTTP
185/// date against the wall clock (the time since the Unix epoch), read together with the monotonic one.
186#[derive(Clone, Copy, Debug)]
187struct Clocks {
188    received_at: Instant,
189    now: Instant,
190    wall: Duration,
191}
192
193impl Clocks {
194    fn now<B>(response: &http::Response<B>) -> Self {
195        // The wall clock first, so that a pause between the two reads cannot shorten an HTTP date's deadline.
196        let wall = since_epoch();
197        let now = Instant::now();
198        let received_at = response
199            .extensions()
200            .get::<ReceivedAt>()
201            .map_or(now, |received| received.0);
202        Self {
203            received_at,
204            now,
205            wall,
206        }
207    }
208}
209
210/// Whether a response or error is worth retrying, and the deadline the server set, if any.
211fn retryable<B>(
212    result: &Result<http::Response<B>, BoxError>,
213    server_directed_429: bool,
214    retry_timeouts: bool,
215) -> Option<Option<Instant>> {
216    match result {
217        Ok(response) => {
218            let status = response.status();
219            let transient = status == StatusCode::TOO_MANY_REQUESTS
220                || status == StatusCode::REQUEST_TIMEOUT
221                || status.is_server_error();
222            let permanent = response.extensions().get::<Permanent>().is_some();
223            if !transient || permanent {
224                return None;
225            }
226            let deadline = retry_after(response.headers(), Clocks::now(response));
227            if server_directed_429 && status == StatusCode::TOO_MANY_REQUESTS && deadline.is_none()
228            {
229                #[cfg(feature = "tracing")]
230                tracing::debug!(
231                    target: "timvw_rust_runtime::trace",
232                    status = status.as_u16(),
233                    "not retrying: a 429 without Retry-After (server_directed_429)"
234                );
235                return None;
236            }
237            Some(deadline)
238        }
239        Err(error) => (is_connect(error.as_ref())
240            || (retry_timeouts && crate::error::is_timeout(error.as_ref())))
241        .then_some(None),
242    }
243}
244
245/// The deadline a server sets: `retry-after-ms` in milliseconds, else `Retry-After` in seconds (both may have a
246/// fraction) or as an HTTP date. A date in the past is now.
247fn retry_after(headers: &HeaderMap, clocks: Clocks) -> Option<Instant> {
248    let value = |name| Some(headers.get(name)?.to_str().ok()?.trim());
249    let delay = |value: &str, scale: f64| {
250        let delay = Duration::try_from_secs_f64(value.parse::<f64>().ok()? * scale).ok()?;
251        clocks.received_at.checked_add(delay)
252    };
253    if let Some(deadline) = value(RETRY_AFTER_MS).and_then(|value| delay(value, 0.001)) {
254        return Some(deadline);
255    }
256    let value = value(RETRY_AFTER.as_str())?;
257    delay(value, 1.0).or_else(|| {
258        let date = httpdate::parse_http_date(value).ok()?;
259        let date = date
260            .duration_since(std::time::UNIX_EPOCH)
261            .unwrap_or_default();
262        let delay = date.saturating_sub(clocks.wall);
263        clocks.now.checked_add(delay)
264    })
265}
266
267/// Whether the request failed to connect, so it was not sent and can be sent again. Never on WebAssembly: `fetch`
268/// fails the same way whether or not the request reached the server (and for CORS), so a failure there is not retried.
269#[allow(unused_variables)]
270fn is_connect(error: &(dyn std::error::Error + 'static)) -> bool {
271    #[cfg(all(feature = "reqwest", not(target_family = "wasm")))]
272    {
273        let mut source = Some(error);
274        while let Some(error) = source {
275            if error
276                .downcast_ref::<reqwest::Error>()
277                .is_some_and(reqwest::Error::is_connect)
278            {
279                return true;
280            }
281            source = error.source();
282        }
283    }
284    false
285}
286
287impl<B> Policy<http::Request<Body>, http::Response<B>, BoxError> for RetryPolicy {
288    type Future = BoxFuture<'static, ()>;
289
290    fn retry(
291        &mut self,
292        _request: &mut http::Request<Body>,
293        result: &mut Result<http::Response<B>, BoxError>,
294    ) -> Option<Self::Future> {
295        if self.attempts >= self.max_retries {
296            return None;
297        }
298        let retry_after = retryable(result, self.server_directed_429, self.retry_timeouts)?;
299        let now = Instant::now();
300        let deadline = self.deadline(retry_after, now);
301        #[cfg(feature = "tracing")]
302        if deadline.is_none() {
303            tracing::debug!(
304                target: "timvw_rust_runtime::trace",
305                retry_after_ms = retry_after
306                    .map_or(0, |deadline| deadline.saturating_duration_since(now).as_millis() as u64),
307                max_retry_after_ms = self.max_retry_after.as_millis() as u64,
308                status = result.as_ref().ok().map(|response| response.status().as_u16()),
309                "not retrying: the server asks to wait longer than max_retry_after"
310            );
311        }
312        let deadline = deadline?;
313        self.attempts += 1;
314        #[cfg(feature = "tracing")]
315        tracing::debug!(
316            target: "timvw_rust_runtime::trace",
317            attempt = self.attempts,
318            max_retries = self.max_retries,
319            delay_ms = deadline.saturating_duration_since(now).as_millis() as u64,
320            status = result.as_ref().ok().map(|response| response.status().as_u16()),
321            "retrying"
322        );
323        Some(Box::pin(Sleep::until(deadline)))
324    }
325
326    fn clone_request(&mut self, request: &http::Request<Body>) -> Option<http::Request<Body>> {
327        let body = request.body().try_clone()?;
328        let mut clone = http::Request::new(body);
329        *clone.method_mut() = request.method().clone();
330        *clone.uri_mut() = request.uri().clone();
331        *clone.version_mut() = request.version();
332        *clone.headers_mut() = request.headers().clone();
333        if let Some(applied) = request
334            .extensions()
335            .get::<crate::request_options::Applied>()
336        {
337            clone.extensions_mut().insert(applied.clone());
338        }
339        Some(clone)
340    }
341}
342
343/// A response extension that tells [`RetryPolicy`] not to retry a response it would otherwise retry.
344///
345/// [`RetryLayer::permanent_429`] sets it on a `429` whose body says the error is permanent. A layer of your own, inside
346/// the retries, can set it too: `response.extensions_mut().insert(Permanent)`.
347#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
348pub struct Permanent;
349
350/// Whether the body of a `429` says the error is permanent, such as [`openai_insufficient_quota`].
351type Classifier = Arc<dyn Fn(&[u8]) -> bool + Send + Sync>;
352
353/// OpenAI's "quota exhausted": a JSON body whose `error.code` or `error.type` is `insufficient_quota`.
354///
355/// OpenAI answers `429` both when requests come too fast, which passes, and when the account's quota or credit is
356/// used up, which does not pass by waiting. Use it with [`RetryLayer::permanent_429`].
357pub fn openai_insufficient_quota(body: &[u8]) -> bool {
358    const INSUFFICIENT_QUOTA: &str = "insufficient_quota";
359    let Ok(value) = serde_json::from_slice::<serde_json::Value>(body) else {
360        return false;
361    };
362    ["code", "type"].iter().any(|field| {
363        value
364            .pointer(&format!("/error/{field}"))
365            .and_then(serde_json::Value::as_str)
366            == Some(INSUFFICIENT_QUOTA)
367    })
368}
369
370/// Retries with a [`RetryPolicy`], and optionally reads the body of a `429` to tell a permanent error from a rate
371/// limit first.
372///
373/// `RetryLayer::new(policy)` retries exactly like `tower::retry::RetryLayer::new(policy)`. With
374/// [`RetryLayer::permanent_429`], the body of each `429` is read (up to [`RetryLayer::body_limit`]) and given to a
375/// function; when it returns `true`, the response is marked [`Permanent`] and returned at once. The body is put back,
376/// so the caller (and a generated client's error) still sees all of it.
377///
378/// ```no_run
379/// # #[cfg(feature = "reqwest")]
380/// # fn demo() {
381/// use timvw_rust_runtime::retry::{openai_insufficient_quota, RetryLayer, RetryPolicy};
382/// use timvw_rust_runtime::{Client, ReqwestService};
383///
384/// let client = Client::from_service(ReqwestService::default())
385///     .layer(RetryLayer::new(RetryPolicy::default()).permanent_429(openai_insufficient_quota));
386/// # }
387/// ```
388#[derive(Clone)]
389pub struct RetryLayer {
390    policy: RetryPolicy,
391    classify: Option<Classifier>,
392    body_limit: usize,
393}
394
395impl RetryLayer {
396    /// Retries with `policy`, without reading any body.
397    pub fn new(policy: RetryPolicy) -> Self {
398        Self {
399            policy,
400            classify: None,
401            body_limit: DEFAULT_BODY_LIMIT,
402        }
403    }
404
405    /// Reads the body of each `429` and does not retry it when `is_permanent` returns `true` for the body's bytes.
406    ///
407    /// A body longer than [`RetryLayer::body_limit`], or one that fails to read, is not given to `is_permanent`: the
408    /// response is retried as a rate limit, and the caller still reads the whole body (or its error).
409    pub fn permanent_429<F>(mut self, is_permanent: F) -> Self
410    where
411        F: Fn(&[u8]) -> bool + Send + Sync + 'static,
412    {
413        self.classify = Some(Arc::new(is_permanent));
414        self
415    }
416
417    /// How many bytes of a `429` body are read to classify it (default [`DEFAULT_BODY_LIMIT`], 64 KiB).
418    pub fn body_limit(mut self, limit: usize) -> Self {
419        self.body_limit = limit;
420        self
421    }
422
423    /// The retry policy.
424    pub fn policy(&self) -> &RetryPolicy {
425        &self.policy
426    }
427}
428
429/// Retries with [`RetryPolicy::default`].
430impl Default for RetryLayer {
431    fn default() -> Self {
432        Self::new(RetryPolicy::default())
433    }
434}
435
436impl fmt::Debug for RetryLayer {
437    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
438        f.debug_struct("RetryLayer")
439            .field("policy", &self.policy)
440            .field("permanent_429", &self.classify.is_some())
441            .field("body_limit", &self.body_limit)
442            .finish()
443    }
444}
445
446impl<S> Layer<S> for RetryLayer {
447    type Service = tower::retry::Retry<RetryPolicy, Classify429<S>>;
448
449    fn layer(&self, inner: S) -> Self::Service {
450        tower::retry::Retry::new(
451            self.policy.clone(),
452            Classify429 {
453                inner,
454                classify: self.classify.clone(),
455                body_limit: self.body_limit,
456            },
457        )
458    }
459}
460
461/// The service inside the retries of a [`RetryLayer`]: marks a `429` [`Permanent`] when its body says so.
462#[derive(Clone)]
463pub struct Classify429<S> {
464    inner: S,
465    classify: Option<Classifier>,
466    body_limit: usize,
467}
468
469impl<S: fmt::Debug> fmt::Debug for Classify429<S> {
470    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
471        f.debug_struct("Classify429")
472            .field("inner", &self.inner)
473            .field("permanent_429", &self.classify.is_some())
474            .field("body_limit", &self.body_limit)
475            .finish()
476    }
477}
478
479impl<S, R, B> Service<R> for Classify429<S>
480where
481    S: Service<R, Response = http::Response<B>>,
482    S::Future: MaybeSend + 'static,
483    B: http_body::Body + MaybeSend + MaybeSync + 'static,
484    B::Data: Into<Bytes>,
485    B::Error: Into<BoxError>,
486{
487    type Response = http::Response<Body>;
488    type Error = S::Error;
489    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
490
491    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
492        self.inner.poll_ready(cx)
493    }
494
495    fn call(&mut self, request: R) -> Self::Future {
496        let response = self.inner.call(request);
497        let classify = self.classify.clone();
498        let body_limit = self.body_limit;
499        Box::pin(async move {
500            let mut response = response.await?.map(Body::wrap);
501            Ok(match classify {
502                Some(classify) if response.status() == StatusCode::TOO_MANY_REQUESTS => {
503                    response.extensions_mut().insert(ReceivedAt(Instant::now()));
504                    classify_429(response, &*classify, body_limit).await
505                }
506                _ => response,
507            })
508        })
509    }
510}
511
512/// Reads a `429` body, up to `limit` bytes, marks the response [`Permanent`] if `classify` says so, and puts the body
513/// back.
514async fn classify_429(
515    response: http::Response<Body>,
516    classify: &(dyn Fn(&[u8]) -> bool + Send + Sync),
517    limit: usize,
518) -> http::Response<Body> {
519    let (mut parts, mut body) = response.into_parts();
520    let mut read = BytesMut::new();
521    let body = loop {
522        match body.frame().await {
523            None => {
524                let read = read.freeze();
525                if classify(&read) {
526                    parts.extensions.insert(Permanent);
527                    #[cfg(feature = "tracing")]
528                    tracing::debug!(
529                        target: "timvw_rust_runtime::trace",
530                        status = parts.status.as_u16(),
531                        "not retrying: the body says the error is permanent"
532                    );
533                }
534                break Body::from(read);
535            }
536            // Trailers on a 429 are dropped.
537            Some(Ok(frame)) => {
538                let Ok(data) = frame.into_data() else {
539                    continue;
540                };
541                read.extend_from_slice(&data);
542                if read.len() > limit {
543                    let read = futures_util::stream::once(async move { Ok(read.freeze()) });
544                    break Body::wrap_stream(futures_util::StreamExt::chain(
545                        read,
546                        BodyDataStream::new(body),
547                    ));
548                }
549            }
550            Some(Err(error)) => {
551                break Body::wrap_stream(futures_util::stream::iter([
552                    Ok(read.freeze()),
553                    Err::<Bytes, BoxError>(error),
554                ]));
555            }
556        }
557    };
558    http::Response::from_parts(parts, body)
559}
560
561#[cfg(test)]
562mod tests {
563    use super::*;
564
565    fn response(status: u16, headers: &[(&str, &str)]) -> Result<http::Response<()>, BoxError> {
566        let mut builder = http::Response::builder().status(status);
567        for (name, value) in headers {
568            builder = builder.header(*name, *value);
569        }
570        Ok(builder.body(()).unwrap())
571    }
572
573    /// Whether `result` is retried, and after how long, measured from now.
574    fn wait(
575        result: &Result<http::Response<()>, BoxError>,
576        server_directed_429: bool,
577    ) -> Option<Option<Duration>> {
578        let now = Instant::now();
579        retryable(result, server_directed_429, false)
580            .map(|deadline| deadline.map(|deadline| deadline.saturating_duration_since(now)))
581    }
582
583    #[tokio::test(start_paused = true)]
584    async fn retries_transient_statuses_only() {
585        assert_eq!(wait(&response(429, &[]), false), Some(None));
586        assert_eq!(
587            wait(&response(503, &[("retry-after", "2")]), false),
588            Some(Some(Duration::from_secs(2)))
589        );
590        assert_eq!(wait(&response(408, &[]), false), Some(None));
591        assert_eq!(wait(&response(400, &[]), false), None);
592        assert_eq!(wait(&response(200, &[]), false), None);
593        assert_eq!(wait(&Err("boom".into()), false), None);
594    }
595
596    #[tokio::test(start_paused = true)]
597    async fn server_directed_429_needs_advice() {
598        assert_eq!(wait(&response(429, &[]), true), None);
599        assert_eq!(wait(&response(429, &[("retry-after", "soon")]), true), None);
600        assert_eq!(
601            wait(&response(429, &[("retry-after", "0")]), true),
602            Some(Some(Duration::ZERO))
603        );
604        assert_eq!(
605            wait(&response(429, &[("retry-after-ms", "300")]), true),
606            Some(Some(Duration::from_millis(300)))
607        );
608        let past = httpdate::fmt_http_date(std::time::SystemTime::now() - Duration::from_secs(60));
609        assert_eq!(
610            wait(&response(429, &[("retry-after", &past)]), true),
611            Some(Some(Duration::ZERO))
612        );
613        // Only 429s: other transient statuses back off without advice.
614        assert_eq!(wait(&response(503, &[]), true), Some(None));
615        assert_eq!(wait(&response(408, &[]), true), Some(None));
616    }
617
618    #[tokio::test(start_paused = true)]
619    async fn permanent_responses_are_not_retried() {
620        let mut result = response(429, &[("retry-after", "1")]);
621        result.as_mut().unwrap().extensions_mut().insert(Permanent);
622        assert_eq!(wait(&result, false), None);
623    }
624
625    fn clocks() -> Clocks {
626        let now = Instant::now();
627        Clocks {
628            received_at: now,
629            now,
630            wall: httpdate::parse_http_date("Wed, 21 Oct 2015 07:28:00 GMT")
631                .unwrap()
632                .duration_since(std::time::UNIX_EPOCH)
633                .unwrap(),
634        }
635    }
636
637    fn parse(headers: &[(&str, &str)], clocks: Clocks) -> Option<Duration> {
638        let mut map = HeaderMap::new();
639        for (name, value) in headers {
640            map.append(
641                http::HeaderName::from_bytes(name.as_bytes()).unwrap(),
642                value.parse().unwrap(),
643            );
644        }
645        retry_after(&map, clocks).map(|deadline| deadline.saturating_duration_since(clocks.now))
646    }
647
648    #[tokio::test(start_paused = true)]
649    async fn parses_retry_after_and_retry_after_ms() {
650        let wait = |headers: &[(&str, &str)]| parse(headers, clocks());
651        assert_eq!(
652            wait(&[("retry-after", " 20 ")]),
653            Some(Duration::from_secs(20))
654        );
655        assert_eq!(
656            wait(&[("retry-after", "1.5")]),
657            Some(Duration::from_millis(1500))
658        );
659        assert_eq!(
660            wait(&[("retry-after-ms", "250"), ("retry-after", "1")]),
661            Some(Duration::from_millis(250))
662        );
663        assert_eq!(
664            wait(&[("retry-after-ms", "12.5")]),
665            Some(Duration::from_micros(12_500))
666        );
667        assert_eq!(wait(&[("retry-after", "-1")]), None);
668        assert_eq!(wait(&[("retry-after", "")]), None);
669        assert_eq!(wait(&[("retry-after", "later")]), None);
670        assert_eq!(
671            wait(&[("retry-after-ms", "nan"), ("retry-after", "3")]),
672            Some(Duration::from_secs(3))
673        );
674        assert_eq!(wait(&[]), None);
675    }
676
677    #[tokio::test(start_paused = true)]
678    async fn parses_http_dates_against_the_wall_clock() {
679        let wait = |value: &str| parse(&[("retry-after", value)], clocks());
680        // The wall clock of `clocks()` is 07:28:00.
681        assert_eq!(
682            wait("Wed, 21 Oct 2015 07:28:30 GMT"),
683            Some(Duration::from_secs(30))
684        );
685        assert_eq!(wait("Wed, 21 Oct 2015 07:28:00 GMT"), Some(Duration::ZERO));
686        assert_eq!(wait("Wed, 21 Oct 2015 07:00:00 GMT"), Some(Duration::ZERO));
687        // The obsolete formats HTTP still accepts.
688        assert_eq!(
689            wait("Wednesday, 21-Oct-15 07:29:00 GMT"),
690            Some(Duration::from_secs(60))
691        );
692        assert_eq!(
693            wait("Wed Oct 21 07:28:05 2015"),
694            Some(Duration::from_secs(5))
695        );
696    }
697
698    #[tokio::test(start_paused = true)]
699    async fn seconds_count_from_receipt_and_dates_from_now() {
700        let mut clocks = clocks();
701        clocks.now = clocks.received_at + Duration::from_secs(2); // two seconds reading the body
702        assert_eq!(
703            parse(&[("retry-after", "5")], clocks),
704            Some(Duration::from_secs(3))
705        );
706        assert_eq!(
707            parse(&[("retry-after-ms", "1000")], clocks),
708            Some(Duration::ZERO)
709        );
710        assert_eq!(
711            parse(&[("retry-after", "Wed, 21 Oct 2015 07:28:05 GMT")], clocks),
712            Some(Duration::from_secs(5))
713        );
714    }
715
716    #[tokio::test(start_paused = true)]
717    async fn backs_off_and_caps() {
718        let mut policy = RetryPolicy::new(10)
719            .base_delay(Duration::from_millis(100))
720            .max_delay(Duration::from_secs(1))
721            .jitter(false);
722        let now = Instant::now();
723        assert_eq!(
724            policy.deadline(None, now),
725            Some(now + Duration::from_millis(100))
726        );
727        policy.attempts = 2;
728        assert_eq!(
729            policy.deadline(None, now),
730            Some(now + Duration::from_millis(400))
731        );
732        policy.attempts = 8;
733        assert_eq!(
734            policy.deadline(None, now),
735            Some(now + Duration::from_secs(1))
736        );
737    }
738
739    #[tokio::test(start_paused = true)]
740    async fn honours_retry_after_beyond_max_delay_up_to_max_retry_after() {
741        let policy = RetryPolicy::new(3).max_delay(Duration::from_secs(1));
742        let now = Instant::now();
743        let after = |secs| Some(now + Duration::from_secs(secs));
744        assert_eq!(policy.deadline(after(30), now), after(30));
745        assert_eq!(policy.deadline(after(60), now), after(60));
746        assert_eq!(policy.deadline(after(61), now), None);
747        // A deadline already passed is kept: retry now.
748        assert_eq!(
749            policy.deadline(Some(now), now + Duration::from_secs(1)),
750            Some(now)
751        );
752        let policy = policy.max_retry_after(Duration::from_secs(5));
753        assert_eq!(policy.deadline(after(6), now), None);
754    }
755
756    #[tokio::test(start_paused = true)]
757    async fn jitter_stays_between_half_and_all_of_the_backoff() {
758        let backoff = Duration::from_millis(800);
759        assert_eq!(jittered(backoff, 0.0), Duration::from_millis(400));
760        assert_eq!(jittered(backoff, 0.5), Duration::from_millis(600));
761        assert!(jittered(backoff, 0.999_999) <= backoff);
762
763        let mut policy = RetryPolicy::new(10)
764            .base_delay(Duration::from_millis(100))
765            .max_delay(Duration::from_secs(1));
766        policy.attempts = 3; // 800 ms before jitter
767        let delays: Vec<_> = (0..200).map(|_| policy.backoff()).collect();
768        for delay in &delays {
769            assert!(
770                (Duration::from_millis(400)..=backoff).contains(delay),
771                "{delay:?}"
772            );
773        }
774        assert!(delays.iter().any(|delay| *delay != delays[0]), "random");
775        // The waits a server asks for are not jittered.
776        let now = Instant::now();
777        let deadline = Some(now + Duration::from_secs(2));
778        assert_eq!(policy.deadline(deadline, now), deadline);
779    }
780
781    #[test]
782    fn random_unit_is_in_range() {
783        for _ in 0..1000 {
784            assert!((0.0..1.0).contains(&random_unit()));
785        }
786    }
787
788    #[test]
789    fn recognises_openai_insufficient_quota() {
790        assert!(openai_insufficient_quota(
791            br#"{"error":{"message":"You exceeded your current quota","type":"insufficient_quota","param":null,"code":"insufficient_quota"}}"#
792        ));
793        assert!(openai_insufficient_quota(
794            br#"{"error":{"type":"insufficient_quota"}}"#
795        ));
796        assert!(openai_insufficient_quota(
797            br#"{"error":{"type":"requests","code":"insufficient_quota"}}"#
798        ));
799        assert!(!openai_insufficient_quota(
800            br#"{"error":{"message":"Rate limit reached","type":"requests","code":"rate_limit_exceeded"}}"#
801        ));
802        assert!(!openai_insufficient_quota(b"insufficient_quota"));
803        assert!(!openai_insufficient_quota(b""));
804    }
805
806    #[test]
807    fn only_bodies_in_memory_are_cloned() {
808        let mut policy = RetryPolicy::default();
809        let mut clone_request = |request: &http::Request<Body>| {
810            Policy::<_, http::Response<()>, BoxError>::clone_request(&mut policy, request)
811        };
812        let request = http::Request::post("https://example.com")
813            .header("x", "1")
814            .body(Body::from("a"))
815            .unwrap();
816        let clone = clone_request(&request).unwrap();
817        assert_eq!(clone.uri(), request.uri());
818        assert_eq!(clone.headers()["x"], "1");
819        let chunks =
820            futures_util::stream::iter(vec![Ok::<_, std::io::Error>(bytes::Bytes::from("a"))]);
821        let streamed = http::Request::post("https://example.com")
822            .body(Body::wrap_stream(chunks))
823            .unwrap();
824        assert!(clone_request(&streamed).is_none());
825    }
826
827    fn body_429(chunks: Vec<Result<&'static str, &'static str>>) -> http::Response<Body> {
828        let stream = futures_util::stream::iter(
829            chunks
830                .into_iter()
831                .map(|chunk| chunk.map(Bytes::from).map_err(BoxError::from)),
832        );
833        http::Response::builder()
834            .status(429)
835            .body(Body::wrap_stream(stream))
836            .unwrap()
837    }
838
839    #[tokio::test]
840    async fn classifies_the_whole_body_and_puts_it_back() {
841        let response = body_429(vec![
842            Ok(r#"{"error":{"type":"#),
843            Ok(r#""insufficient_quota"}}"#),
844        ]);
845        let response = classify_429(response, &openai_insufficient_quota, 1024).await;
846        assert!(response.extensions().get::<Permanent>().is_some());
847        let body = response.into_body().collect_bytes().await.unwrap();
848        assert_eq!(body, r#"{"error":{"type":"insufficient_quota"}}"#);
849    }
850
851    #[tokio::test]
852    async fn a_body_over_the_limit_is_not_classified_but_kept_whole() {
853        let response = body_429(vec![Ok("0123"), Ok("4567"), Ok("89")]);
854        let response = classify_429(response, &|_: &[u8]| true, 5).await;
855        assert!(response.extensions().get::<Permanent>().is_none());
856        let body = response.into_body().collect_bytes().await.unwrap();
857        assert_eq!(body, "0123456789");
858    }
859
860    #[tokio::test]
861    async fn a_body_that_fails_is_not_classified_and_keeps_its_error() {
862        let response = body_429(vec![Ok("01"), Err("reset")]);
863        let response = classify_429(response, &|_: &[u8]| true, 1024).await;
864        assert!(response.extensions().get::<Permanent>().is_none());
865        let error = response.into_body().collect_bytes().await.unwrap_err();
866        assert_eq!(error.to_string(), "reset");
867    }
868}
timvw_rust_runtime/retry.rs

timvw_rust_runtime/
retry.rs
