rama_http/layer/timeout/
service.rs

1use super::{DeadlineBody, TimeoutBody};
2use crate::{Request, Response, StatusCode};
3use rama_core::extensions::{Extension, ExtensionsRef};
4use rama_core::{Layer, Service, telemetry::tracing};
5use rama_http_types::body::OptionalBody;
6use rama_utils::macros::define_inner_service_accessors;
7use std::time::Duration;
8
9/// Cause attached as a response extension when [`Timeout`] synthesises a
10/// timeout response.
11///
12/// The timeout layer maps a request that did not complete in time to a
13/// configured status code (default `408 Request Timeout`). The original
14/// timing information is preserved here so callers can log or react to the
15/// cause without parsing the response status.
16#[derive(Debug, Clone, Copy, Extension)]
17#[extension(tags(http))]
18pub struct TimeoutCause {
19    /// Configured timeout that elapsed before the inner service produced
20    /// a response.
21    pub elapsed: Duration,
22}
23
24/// Layer that applies the [`Timeout`] middleware which apply a timeout to requests.
25///
26/// See the [module docs](super) for an example.
27#[derive(Debug, Clone)]
28pub struct TimeoutLayer {
29    timeout: Duration,
30    status_code: StatusCode,
31}
32
33impl TimeoutLayer {
34    /// Creates a new [`TimeoutLayer`].
35    ///
36    /// By default, it will return a `408 Request Timeout` response if the request does not complete within the specified timeout.
37    /// To customize the response status code, use the `with_status_code` method.
38    #[must_use]
39    #[inline(always)]
40    pub const fn new(timeout: Duration) -> Self {
41        Self::with_status_code(StatusCode::REQUEST_TIMEOUT, timeout)
42    }
43
44    /// Creates a new [`TimeoutLayer`] with the specified status code for the timeout response.
45    #[must_use]
46    pub const fn with_status_code(status_code: StatusCode, timeout: Duration) -> Self {
47        Self {
48            timeout,
49            status_code,
50        }
51    }
52}
53
54impl<S> Layer<S> for TimeoutLayer {
55    type Service = Timeout<S>;
56
57    fn layer(&self, inner: S) -> Self::Service {
58        Timeout::with_status_code(inner, self.status_code, self.timeout)
59    }
60}
61
62/// Middleware which apply a timeout to requests.
63///
64/// See the [module docs](super) for an example.
65#[derive(Debug, Clone)]
66pub struct Timeout<S> {
67    inner: S,
68    timeout: Duration,
69    status_code: StatusCode,
70}
71
72impl<S> Timeout<S> {
73    #[inline(always)]
74    /// Creates a new [`Timeout`].
75    ///
76    /// By default, it will return a `408 Request Timeout` response if the request does not complete within the specified timeout.
77    /// To customize the response status code, use the `with_status_code` method.
78    pub const fn new(inner: S, timeout: Duration) -> Self {
79        Self::with_status_code(inner, StatusCode::REQUEST_TIMEOUT, timeout)
80    }
81
82    /// Creates a new [`Timeout`] with the specified status code for the timeout response.
83    pub const fn with_status_code(inner: S, status_code: StatusCode, timeout: Duration) -> Self {
84        Self {
85            inner,
86            timeout,
87            status_code,
88        }
89    }
90
91    define_inner_service_accessors!();
92}
93
94impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for Timeout<S>
95where
96    S: Service<Request<ReqBody>, Output = Response<ResBody>>,
97    ReqBody: Send + 'static,
98    ResBody: Send + 'static,
99{
100    type Output = Response<OptionalBody<ResBody>>;
101    type Error = S::Error;
102
103    async fn serve(&self, req: Request<ReqBody>) -> Result<Self::Output, Self::Error> {
104        // Capture method + uri before moving the request into the inner
105        // service so we can attach them to the timeout log if it fires.
106        let method = req.method().clone();
107        let uri = req.uri().clone();
108        tokio::select! {
109            res = self.inner.serve(req) => Ok(res?.map(OptionalBody::some)),
110            _ = tokio::time::sleep(self.timeout) => {
111                let elapsed_ms = u64::try_from(self.timeout.as_millis()).unwrap_or(u64::MAX);
112                tracing::warn!(
113                    target: "rama_http::timeout",
114                    elapsed_ms,
115                    status_code = self.status_code.as_u16(),
116                    http.method = %method,
117                    url.path = %uri.path_or_root().as_ref(),
118                    "request did not complete within configured timeout; synthesising response",
119                );
120                let mut res = Response::new(OptionalBody::none());
121                *res.status_mut() = self.status_code;
122                res.extensions().insert(TimeoutCause {
123                    elapsed: self.timeout,
124                });
125                Ok(res)
126            }
127        }
128    }
129}
130
131/// Applies a [`TimeoutBody`] to the request body.
132#[derive(Clone, Debug)]
133pub struct RequestBodyTimeoutLayer {
134    timeout: Duration,
135}
136
137impl RequestBodyTimeoutLayer {
138    /// Creates a new [`RequestBodyTimeoutLayer`].
139    #[must_use]
140    pub fn new(timeout: Duration) -> Self {
141        Self { timeout }
142    }
143}
144
145impl<S> Layer<S> for RequestBodyTimeoutLayer {
146    type Service = RequestBodyTimeout<S>;
147
148    fn layer(&self, inner: S) -> Self::Service {
149        RequestBodyTimeout::new(inner, self.timeout)
150    }
151}
152
153/// Applies a [`TimeoutBody`] to the request body.
154#[derive(Clone, Debug)]
155pub struct RequestBodyTimeout<S> {
156    inner: S,
157    timeout: Duration,
158}
159
160impl<S> RequestBodyTimeout<S> {
161    /// Creates a new [`RequestBodyTimeout`].
162    pub fn new(service: S, timeout: Duration) -> Self {
163        Self {
164            inner: service,
165            timeout,
166        }
167    }
168
169    /// Returns a new [`Layer`] that wraps services with a [`RequestBodyTimeoutLayer`] middleware.
170    #[must_use]
171    pub fn layer(timeout: Duration) -> RequestBodyTimeoutLayer {
172        RequestBodyTimeoutLayer::new(timeout)
173    }
174
175    define_inner_service_accessors!();
176}
177
178impl<S, ReqBody> Service<Request<ReqBody>> for RequestBodyTimeout<S>
179where
180    S: Service<Request<TimeoutBody<ReqBody>>>,
181    ReqBody: Send + 'static,
182{
183    type Output = S::Output;
184    type Error = S::Error;
185
186    async fn serve(&self, req: Request<ReqBody>) -> Result<Self::Output, Self::Error> {
187        let req = req.map(|body| TimeoutBody::new(self.timeout, body));
188        self.inner.serve(req).await
189    }
190}
191
192/// Applies a [`TimeoutBody`] to the response body.
193#[derive(Clone)]
194pub struct ResponseBodyTimeoutLayer {
195    timeout: Duration,
196}
197
198impl ResponseBodyTimeoutLayer {
199    /// Creates a new [`ResponseBodyTimeoutLayer`].
200    #[must_use]
201    pub fn new(timeout: Duration) -> Self {
202        Self { timeout }
203    }
204}
205
206impl<S> Layer<S> for ResponseBodyTimeoutLayer {
207    type Service = ResponseBodyTimeout<S>;
208
209    fn layer(&self, inner: S) -> Self::Service {
210        ResponseBodyTimeout::new(inner, self.timeout)
211    }
212}
213
214/// Applies a [`TimeoutBody`] to the response body.
215#[derive(Clone)]
216pub struct ResponseBodyTimeout<S> {
217    inner: S,
218    timeout: Duration,
219}
220
221impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for ResponseBodyTimeout<S>
222where
223    S: Service<Request<ReqBody>, Output = Response<ResBody>>,
224    ReqBody: Send + 'static,
225    ResBody: Send + 'static,
226{
227    type Output = Response<TimeoutBody<ResBody>>;
228    type Error = S::Error;
229
230    async fn serve(&self, req: Request<ReqBody>) -> Result<Self::Output, Self::Error> {
231        let res = self.inner.serve(req).await?;
232        let res = res.map(|body| TimeoutBody::new(self.timeout, body));
233        Ok(res)
234    }
235}
236
237impl<S> ResponseBodyTimeout<S> {
238    /// Creates a new [`ResponseBodyTimeout`].
239    pub fn new(service: S, timeout: Duration) -> Self {
240        Self {
241            inner: service,
242            timeout,
243        }
244    }
245
246    /// Returns a new [`Layer`] that wraps services with a [`ResponseBodyTimeoutLayer`] middleware.
247    #[must_use]
248    pub fn layer(timeout: Duration) -> ResponseBodyTimeoutLayer {
249        ResponseBodyTimeoutLayer::new(timeout)
250    }
251
252    define_inner_service_accessors!();
253}
254
255/// Applies a [`DeadlineBody`] to the request body.
256///
257/// Unlike [`RequestBodyTimeoutLayer`], which resets on each frame, this enforces a hard
258/// deadline on the entire body transfer.
259#[derive(Clone, Debug)]
260pub struct RequestBodyDeadlineLayer {
261    timeout: Duration,
262}
263
264impl RequestBodyDeadlineLayer {
265    /// Creates a new [`RequestBodyDeadlineLayer`].
266    #[must_use]
267    pub fn new(timeout: Duration) -> Self {
268        Self { timeout }
269    }
270}
271
272impl<S> Layer<S> for RequestBodyDeadlineLayer {
273    type Service = RequestBodyDeadline<S>;
274
275    fn layer(&self, inner: S) -> Self::Service {
276        RequestBodyDeadline::new(inner, self.timeout)
277    }
278}
279
280/// Applies a [`DeadlineBody`] to the request body.
281#[derive(Clone, Debug)]
282pub struct RequestBodyDeadline<S> {
283    inner: S,
284    timeout: Duration,
285}
286
287impl<S> RequestBodyDeadline<S> {
288    /// Creates a new [`RequestBodyDeadline`].
289    pub fn new(service: S, timeout: Duration) -> Self {
290        Self {
291            inner: service,
292            timeout,
293        }
294    }
295
296    /// Returns a new [`Layer`] that wraps services with a [`RequestBodyDeadlineLayer`] middleware.
297    #[must_use]
298    pub fn layer(timeout: Duration) -> RequestBodyDeadlineLayer {
299        RequestBodyDeadlineLayer::new(timeout)
300    }
301
302    define_inner_service_accessors!();
303}
304
305impl<S, ReqBody> Service<Request<ReqBody>> for RequestBodyDeadline<S>
306where
307    S: Service<Request<DeadlineBody<ReqBody>>>,
308    ReqBody: Send + 'static,
309{
310    type Output = S::Output;
311    type Error = S::Error;
312
313    async fn serve(&self, req: Request<ReqBody>) -> Result<Self::Output, Self::Error> {
314        let req = req.map(|body| DeadlineBody::new(self.timeout, body));
315        self.inner.serve(req).await
316    }
317}
318
319/// Applies a [`DeadlineBody`] to the response body.
320///
321/// Unlike [`ResponseBodyTimeoutLayer`], which resets on each frame, this enforces a hard
322/// deadline on the entire body transfer.
323#[derive(Clone, Debug)]
324pub struct ResponseBodyDeadlineLayer {
325    timeout: Duration,
326}
327
328impl ResponseBodyDeadlineLayer {
329    /// Creates a new [`ResponseBodyDeadlineLayer`].
330    #[must_use]
331    pub fn new(timeout: Duration) -> Self {
332        Self { timeout }
333    }
334}
335
336impl<S> Layer<S> for ResponseBodyDeadlineLayer {
337    type Service = ResponseBodyDeadline<S>;
338
339    fn layer(&self, inner: S) -> Self::Service {
340        ResponseBodyDeadline::new(inner, self.timeout)
341    }
342}
343
344/// Applies a [`DeadlineBody`] to the response body.
345#[derive(Clone, Debug)]
346pub struct ResponseBodyDeadline<S> {
347    inner: S,
348    timeout: Duration,
349}
350
351impl<S> ResponseBodyDeadline<S> {
352    /// Creates a new [`ResponseBodyDeadline`].
353    pub fn new(service: S, timeout: Duration) -> Self {
354        Self {
355            inner: service,
356            timeout,
357        }
358    }
359
360    /// Returns a new [`Layer`] that wraps services with a [`ResponseBodyDeadlineLayer`] middleware.
361    #[must_use]
362    pub fn layer(timeout: Duration) -> ResponseBodyDeadlineLayer {
363        ResponseBodyDeadlineLayer::new(timeout)
364    }
365
366    define_inner_service_accessors!();
367}
368
369impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for ResponseBodyDeadline<S>
370where
371    S: Service<Request<ReqBody>, Output = Response<ResBody>>,
372    ReqBody: Send + 'static,
373    ResBody: Send + 'static,
374{
375    type Output = Response<DeadlineBody<ResBody>>;
376    type Error = S::Error;
377
378    async fn serve(&self, req: Request<ReqBody>) -> Result<Self::Output, Self::Error> {
379        let res = self.inner.serve(req).await?;
380        let res = res.map(|body| DeadlineBody::new(self.timeout, body));
381        Ok(res)
382    }
383}
384
385#[cfg(test)]
386mod tests {
387    use std::convert::Infallible;
388
389    use super::*;
390    use crate::{Body, body::util::BodyExt};
391    use rama_core::service::service_fn;
392
393    #[tokio::test]
394    async fn request_completes_within_timeout() {
395        let service =
396            TimeoutLayer::with_status_code(StatusCode::GATEWAY_TIMEOUT, Duration::from_secs(1))
397                .into_layer(service_fn(fast_handler));
398
399        let request = Request::get("/").body(Body::empty()).unwrap();
400        let res = service.serve(request).await.unwrap();
401
402        assert_eq!(res.status(), StatusCode::OK);
403    }
404
405    #[tokio::test]
406    async fn timeout_middleware_with_custom_status_code() {
407        let service = Timeout::with_status_code(
408            service_fn(slow_handler),
409            StatusCode::REQUEST_TIMEOUT,
410            Duration::from_millis(10),
411        );
412
413        let request = Request::get("/").body(Body::empty()).unwrap();
414        let res = service.serve(request).await.unwrap();
415
416        assert_eq!(res.status(), StatusCode::REQUEST_TIMEOUT);
417    }
418
419    #[tokio::test]
420    async fn timeout_response_has_empty_body() {
421        let service =
422            TimeoutLayer::with_status_code(StatusCode::GATEWAY_TIMEOUT, Duration::from_millis(10))
423                .into_layer(service_fn(slow_handler));
424
425        let request = Request::get("/").body(Body::empty()).unwrap();
426        let res = service.serve(request).await.unwrap();
427
428        assert_eq!(res.status(), StatusCode::GATEWAY_TIMEOUT);
429
430        // Verify the body is empty (default)
431        let body = res.into_body();
432        let bytes = body.collect().await.unwrap().to_bytes();
433        assert!(bytes.is_empty());
434    }
435
436    #[tokio::test]
437    async fn deprecated_new_method_compatibility() {
438        let layer = TimeoutLayer::new(Duration::from_millis(10));
439        let service = layer.into_layer(service_fn(slow_handler));
440
441        let request = Request::get("/").body(Body::empty()).unwrap();
442        let res = service.serve(request).await.unwrap();
443
444        // Should use default 408 status code
445        assert_eq!(res.status(), StatusCode::REQUEST_TIMEOUT);
446    }
447
448    #[tokio::test]
449    async fn timeout_response_carries_timeout_cause_extension() {
450        let timeout = Duration::from_millis(10);
451        let service = Timeout::new(service_fn(slow_handler), timeout);
452        let request = Request::get("/").body(Body::empty()).unwrap();
453        let res = service.serve(request).await.unwrap();
454
455        assert_eq!(res.status(), StatusCode::REQUEST_TIMEOUT);
456        let cause = res
457            .extensions()
458            .get_ref::<TimeoutCause>()
459            .expect("TimeoutCause extension");
460        assert_eq!(cause.elapsed, timeout);
461    }
462
463    #[tokio::test]
464    async fn fast_response_has_no_timeout_cause_extension() {
465        let service = Timeout::new(service_fn(fast_handler), Duration::from_secs(1));
466        let request = Request::get("/").body(Body::empty()).unwrap();
467        let res = service.serve(request).await.unwrap();
468
469        assert_eq!(res.status(), StatusCode::OK);
470        assert!(res.extensions().get_ref::<TimeoutCause>().is_none());
471    }
472
473    async fn slow_handler(_req: Request<Body>) -> Result<Response<Body>, Infallible> {
474        tokio::time::sleep(Duration::from_secs(10)).await;
475        Ok(Response::builder()
476            .status(StatusCode::OK)
477            .body(Body::empty())
478            .unwrap())
479    }
480
481    async fn fast_handler(_req: Request<Body>) -> Result<Response<Body>, Infallible> {
482        Ok(Response::builder()
483            .status(StatusCode::OK)
484            .body(Body::empty())
485            .unwrap())
486    }
487}
rama_http/layer/timeout/service.rs

rama_http/layer/timeout/
service.rs
