rama_http/layer/retry/
mod.rs

1//! Middleware for retrying "failed" requests.
2
3use crate::dep::http_body::Body as HttpBody;
4use crate::dep::http_body_util::BodyExt;
5use crate::Request;
6use rama_core::error::BoxError;
7use rama_core::{Context, Service};
8use rama_utils::macros::define_inner_service_accessors;
9
10mod layer;
11mod policy;
12
13mod body;
14#[doc(inline)]
15pub use body::RetryBody;
16
17pub mod managed;
18pub use managed::ManagedPolicy;
19
20#[cfg(test)]
21mod tests;
22
23pub use self::layer::RetryLayer;
24pub use self::policy::{Policy, PolicyResult};
25
26/// Configure retrying requests of "failed" responses.
27///
28/// A [`Policy`] classifies what is a "failed" response.
29pub struct Retry<P, S> {
30    policy: P,
31    inner: S,
32}
33
34impl<P, S> std::fmt::Debug for Retry<P, S>
35where
36    P: std::fmt::Debug,
37    S: std::fmt::Debug,
38{
39    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
40        f.debug_struct("Retry")
41            .field("policy", &self.policy)
42            .field("inner", &self.inner)
43            .finish()
44    }
45}
46
47impl<P, S> Clone for Retry<P, S>
48where
49    P: Clone,
50    S: Clone,
51{
52    fn clone(&self) -> Self {
53        Retry {
54            policy: self.policy.clone(),
55            inner: self.inner.clone(),
56        }
57    }
58}
59
60// ===== impl Retry =====
61
62impl<P, S> Retry<P, S> {
63    /// Retry the inner service depending on this [`Policy`].
64    pub const fn new(policy: P, service: S) -> Self {
65        Retry {
66            policy,
67            inner: service,
68        }
69    }
70
71    define_inner_service_accessors!();
72}
73
74#[derive(Debug)]
75/// Error type for [`Retry`]
76pub struct RetryError {
77    kind: RetryErrorKind,
78    inner: Option<BoxError>,
79}
80
81#[derive(Debug)]
82enum RetryErrorKind {
83    BodyConsume,
84    Service,
85}
86
87impl std::fmt::Display for RetryError {
88    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
89        match &self.inner {
90            Some(inner) => write!(f, "{}: {}", self.kind, inner),
91            None => write!(f, "{}", self.kind),
92        }
93    }
94}
95
96impl std::fmt::Display for RetryErrorKind {
97    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
98        match self {
99            RetryErrorKind::BodyConsume => write!(f, "failed to consume body"),
100            RetryErrorKind::Service => write!(f, "service error"),
101        }
102    }
103}
104
105impl std::error::Error for RetryError {
106    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
107        self.inner.as_ref().and_then(|e| e.source())
108    }
109}
110
111impl<P, S, State, Body> Service<State, Request<Body>> for Retry<P, S>
112where
113    P: Policy<State, S::Response, S::Error>,
114    S: Service<State, Request<RetryBody>, Error: Into<BoxError>>,
115    State: Clone + Send + Sync + 'static,
116    Body: HttpBody<Data: Send + 'static, Error: Into<BoxError>> + Send + 'static,
117{
118    type Response = S::Response;
119    type Error = RetryError;
120
121    async fn serve(
122        &self,
123        ctx: Context<State>,
124        request: Request<Body>,
125    ) -> Result<Self::Response, Self::Error> {
126        let mut ctx = ctx;
127
128        // consume body so we can clone the request if desired
129        let (parts, body) = request.into_parts();
130        let body = body.collect().await.map_err(|e| RetryError {
131            kind: RetryErrorKind::BodyConsume,
132            inner: Some(e.into()),
133        })?;
134        let body = RetryBody::new(body.to_bytes());
135        let mut request = Request::from_parts(parts, body);
136
137        let mut cloned = self.policy.clone_input(&ctx, &request);
138
139        loop {
140            let resp = self.inner.serve(ctx, request).await;
141            match cloned.take() {
142                Some((cloned_ctx, cloned_req)) => {
143                    let (cloned_ctx, cloned_req) =
144                        match self.policy.retry(cloned_ctx, cloned_req, resp).await {
145                            PolicyResult::Abort(result) => {
146                                return result.map_err(|e| RetryError {
147                                    kind: RetryErrorKind::Service,
148                                    inner: Some(e.into()),
149                                })
150                            }
151                            PolicyResult::Retry { ctx, req } => (ctx, req),
152                        };
153
154                    cloned = self.policy.clone_input(&cloned_ctx, &cloned_req);
155                    ctx = cloned_ctx;
156                    request = cloned_req;
157                }
158                // no clone was made, so no possibility to retry
159                None => {
160                    return resp.map_err(|e| RetryError {
161                        kind: RetryErrorKind::Service,
162                        inner: Some(e.into()),
163                    })
164                }
165            }
166        }
167    }
168}
169
170#[cfg(test)]
171mod test {
172    use super::*;
173    use crate::{
174        layer::retry::managed::DoNotRetry, BodyExtractExt, IntoResponse, Response, StatusCode,
175    };
176    use rama_core::{service::service_fn, Context, Layer};
177    use rama_utils::{backoff::ExponentialBackoff, rng::HasherRng};
178    use std::{
179        sync::{atomic::AtomicUsize, Arc},
180        time::Duration,
181    };
182
183    #[tokio::test]
184    async fn test_service_with_managed_retry() {
185        let backoff = ExponentialBackoff::new(
186            Duration::from_millis(1),
187            Duration::from_millis(5),
188            0.1,
189            HasherRng::default,
190        )
191        .unwrap();
192
193        #[derive(Debug, Clone)]
194        struct State {
195            retry_counter: Arc<AtomicUsize>,
196        }
197
198        async fn retry<E>(
199            ctx: Context<State>,
200            result: Result<Response, E>,
201        ) -> (Context<State>, Result<Response, E>, bool) {
202            if ctx.contains::<DoNotRetry>() {
203                panic!("unexpected retry: should be disabled");
204            }
205
206            match result {
207                Ok(ref res) => {
208                    if res.status().is_server_error() {
209                        ctx.state()
210                            .retry_counter
211                            .fetch_add(1, std::sync::atomic::Ordering::AcqRel);
212                        (ctx, result, true)
213                    } else {
214                        (ctx, result, false)
215                    }
216                }
217                Err(_) => {
218                    ctx.state()
219                        .retry_counter
220                        .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
221                    (ctx, result, true)
222                }
223            }
224        }
225
226        let retry_policy = ManagedPolicy::new(retry).with_backoff(backoff);
227
228        let service = RetryLayer::new(retry_policy).layer(service_fn(
229            |_ctx, req: Request<RetryBody>| async {
230                let txt = req.try_into_string().await.unwrap();
231                match txt.as_str() {
232                    "internal" => Ok(StatusCode::INTERNAL_SERVER_ERROR.into_response()),
233                    "error" => Err(rama_core::error::BoxError::from("custom error")),
234                    _ => Ok(txt.into_response()),
235                }
236            },
237        ));
238
239        fn request(s: &'static str) -> Request {
240            Request::builder().body(s.into()).unwrap()
241        }
242
243        fn ctx() -> Context<State> {
244            Context::with_state(State {
245                retry_counter: Arc::new(AtomicUsize::new(0)),
246            })
247        }
248
249        fn ctx_do_not_retry() -> Context<State> {
250            let mut ctx = ctx();
251            ctx.insert(DoNotRetry::default());
252            ctx
253        }
254
255        async fn assert_serve_ok<E: std::fmt::Debug>(
256            msg: &'static str,
257            input: &'static str,
258            output: &'static str,
259            ctx: Context<State>,
260            retried: bool,
261            service: &impl Service<State, Request, Response = Response, Error = E>,
262        ) {
263            let state = ctx.state_clone();
264
265            let fut = service.serve(ctx, request(input));
266            let res = fut.await.unwrap();
267
268            let body = res.try_into_string().await.unwrap();
269            assert_eq!(body, output, "{msg}");
270            if retried {
271                assert!(
272                    state
273                        .retry_counter
274                        .load(std::sync::atomic::Ordering::Acquire)
275                        > 0,
276                    "{msg}"
277                );
278            } else {
279                assert_eq!(
280                    state
281                        .retry_counter
282                        .load(std::sync::atomic::Ordering::Acquire),
283                    0,
284                    "{msg}"
285                );
286            }
287        }
288
289        async fn assert_serve_err<E: std::fmt::Debug>(
290            msg: &'static str,
291            input: &'static str,
292            ctx: Context<State>,
293            retried: bool,
294            service: &impl Service<State, Request, Response = Response, Error = E>,
295        ) {
296            let state = ctx.state_clone();
297
298            let fut = service.serve(ctx, request(input));
299            let res = fut.await;
300
301            assert!(res.is_err(), "{msg}");
302            if retried {
303                assert!(
304                    state
305                        .retry_counter
306                        .load(std::sync::atomic::Ordering::Acquire)
307                        > 0,
308                    "{msg}"
309                );
310            } else {
311                assert_eq!(
312                    state
313                        .retry_counter
314                        .load(std::sync::atomic::Ordering::Acquire),
315                    0,
316                    "{msg}"
317                )
318            }
319        }
320
321        assert_serve_ok(
322            "ok response should be aborted as response without retry",
323            "hello",
324            "hello",
325            ctx(),
326            false,
327            &service,
328        )
329        .await;
330        assert_serve_ok(
331            "internal will trigger 500 with a retry",
332            "internal",
333            "",
334            ctx(),
335            true,
336            &service,
337        )
338        .await;
339        assert_serve_err(
340            "error will trigger an actual non-http error with a retry",
341            "error",
342            ctx(),
343            true,
344            &service,
345        )
346        .await;
347
348        assert_serve_ok(
349            "normally internal will trigger a 500 with retry, but using DoNotRetry will disable retrying",
350            "internal",
351            "",
352            ctx_do_not_retry(),
353            false,
354            &service,
355        ).await;
356    }
357}
rama_http/layer/retry/mod.rs

rama_http/layer/retry/
mod.rs
