viz_middleware/
timeout.rs

1use std::{future::Future, pin::Pin, time::Duration};
2
3use tokio::time::timeout;
4
5use viz_core::{http, Context, Middleware, Response, Result};
6use viz_utils::tracing;
7
8/// Timeout Middleware
9#[derive(Debug)]
10pub struct Timeout {
11    /// 0.256s
12    delay: Duration,
13}
14
15impl Default for Timeout {
16    fn default() -> Self {
17        Self::new(Duration::from_millis(256))
18    }
19}
20
21impl Timeout {
22    /// Creates Timeout Middleware
23    pub fn new(delay: Duration) -> Self {
24        Self { delay }
25    }
26
27    async fn run(&self, cx: &mut Context) -> Result<Response> {
28        let method = cx.method().to_owned();
29        let path = cx.path().to_owned();
30
31        match timeout(self.delay, cx.next()).await {
32            Ok(r) => r,
33            Err(e) => {
34                tracing::trace!(" {:>7} {} {}", method, path, e);
35                Ok(http::StatusCode::REQUEST_TIMEOUT.into())
36            }
37        }
38    }
39}
40
41impl<'a> Middleware<'a, Context> for Timeout {
42    type Output = Result<Response>;
43
44    #[must_use]
45    fn call(
46        &'a self,
47        cx: &'a mut Context,
48    ) -> Pin<Box<dyn Future<Output = Self::Output> + Send + 'a>> {
49        Box::pin(self.run(cx))
50    }
51}