viz_middleware/
timeout.rs1use 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#[derive(Debug)]
10pub struct Timeout {
11 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 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}