1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40

use tokio_service::Service;
use tokio_timer::{self as timer, Timer};
use std::time::Duration;

/// Abort requests that are taking too long
#[derive(Clone)]
pub struct Timeout<S> {
    upstream: S,
    timer: Timer,
    duration: Duration,
}

impl<S> Timeout<S> {
    /// Crate a new `Timeout` with the given `upstream` service.
    ///
    /// Requests will be limited to `duration` and aborted once the limit has
    /// been reached.
    pub fn new(upstream: S, timer: Timer, duration: Duration) -> Timeout<S> {
        Timeout {
            upstream: upstream,
            duration: duration,
            timer: timer,
        }
    }
}

impl<S, E> Service for Timeout<S>
    where S: Service<Error = E>,
          E: From<timer::TimeoutError<S::Future>>,
{
    type Request = S::Request;
    type Response = S::Response;
    type Error = S::Error;
    type Future = timer::Timeout<S::Future>;

    fn call(&self, request: Self::Request) -> Self::Future {
        let resp = self.upstream.call(request);
        self.timer.timeout(resp, self.duration)
    }
}