tower_async_http/timeout/
service.rs

1use http::{Request, Response, StatusCode};
2use std::time::Duration;
3use tower_async_layer::Layer;
4use tower_async_service::Service;
5
6/// Layer that applies the [`Timeout`] middleware which apply a timeout to requests.
7///
8/// See the [module docs](super) for an example.
9#[derive(Debug, Clone, Copy)]
10pub struct TimeoutLayer {
11    timeout: Duration,
12}
13
14impl TimeoutLayer {
15    /// Creates a new [`TimeoutLayer`].
16    pub fn new(timeout: Duration) -> Self {
17        TimeoutLayer { timeout }
18    }
19}
20
21impl<S> Layer<S> for TimeoutLayer {
22    type Service = Timeout<S>;
23
24    fn layer(&self, inner: S) -> Self::Service {
25        Timeout::new(inner, self.timeout)
26    }
27}
28
29/// Middleware which apply a timeout to requests.
30///
31/// If the request does not complete within the specified timeout it will be aborted and a `408
32/// Request Timeout` response will be sent.
33///
34/// See the [module docs](super) for an example.
35#[derive(Debug, Clone, Copy)]
36pub struct Timeout<S> {
37    inner: S,
38    timeout: Duration,
39}
40
41impl<S> Timeout<S> {
42    /// Creates a new [`Timeout`].
43    pub fn new(inner: S, timeout: Duration) -> Self {
44        Self { inner, timeout }
45    }
46
47    define_inner_service_accessors!();
48
49    /// Returns a new [`Layer`] that wraps services with a `Timeout` middleware.
50    ///
51    /// [`Layer`]: tower_async_layer::Layer
52    pub fn layer(timeout: Duration) -> TimeoutLayer {
53        TimeoutLayer::new(timeout)
54    }
55}
56
57impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for Timeout<S>
58where
59    S: Service<Request<ReqBody>, Response = Response<ResBody>>,
60    ResBody: Default,
61{
62    type Response = S::Response;
63    type Error = S::Error;
64
65    async fn call(&self, req: Request<ReqBody>) -> Result<Self::Response, Self::Error> {
66        tokio::select! {
67            res = self.inner.call(req) => res,
68            _ = tokio::time::sleep(self.timeout) => {
69                let mut res = Response::new(ResBody::default());
70                *res.status_mut() = StatusCode::REQUEST_TIMEOUT;
71                Ok(res)
72            }
73        }
74    }
75}