tork_core/middleware/
timeout.rs

1//! Request-timeout middleware.
2
3use std::time::Duration;
4
5use crate::error::{Error, Result};
6use crate::middleware::{DuplicatePolicy, Middleware, Next, Request};
7use crate::response::Response;
8use crate::router::BoxFuture;
9
10/// Fails a request with `504 Gateway Timeout` if it runs longer than a deadline.
11pub struct Timeout {
12    duration: Duration,
13}
14
15impl Timeout {
16    /// Creates a timeout from a [`Duration`].
17    pub fn new(duration: Duration) -> Self {
18        Self { duration }
19    }
20
21    /// Creates a timeout of `secs` seconds.
22    pub fn seconds(secs: u64) -> Self {
23        Self::new(Duration::from_secs(secs))
24    }
25
26    /// Creates a timeout of `millis` milliseconds.
27    pub fn millis(millis: u64) -> Self {
28        Self::new(Duration::from_millis(millis))
29    }
30}
31
32impl Middleware for Timeout {
33    fn handle(&self, request: Request, next: Next) -> BoxFuture<'static, Result<Response>> {
34        let duration = self.duration;
35        Box::pin(async move {
36            match tokio::time::timeout(duration, next.run(request)).await {
37                Ok(result) => result,
38                Err(_elapsed) => Err(Error::gateway_timeout("request timed out")),
39            }
40        })
41    }
42
43    fn name(&self) -> &'static str {
44        "Timeout"
45    }
46
47    fn duplicate_policy(&self) -> DuplicatePolicy {
48        DuplicatePolicy::Reject
49    }
50}
51
52#[cfg(test)]
53mod tests {
54    use super::*;
55
56    #[test]
57    fn constructors_preserve_requested_duration() {
58        assert_eq!(
59            Timeout::new(Duration::from_secs(3)).duration,
60            Duration::from_secs(3)
61        );
62        assert_eq!(Timeout::seconds(2).duration, Duration::from_secs(2));
63        assert_eq!(Timeout::millis(25).duration, Duration::from_millis(25));
64    }
65}
tork_core/middleware/timeout.rs

tork_core/middleware/
timeout.rs
