2 use tokio::timer::Delay;
4 use std::time::{Duration, Instant};
6 pub struct TimeoutStream<S> where S : Stream {
13 impl<S> TimeoutStream<S> where S : Stream {
14 pub fn new_persistent(stream: S, timeout: Duration) -> Self {
15 let next_deadline = Delay::new(Instant::now() + timeout);
24 pub fn new_timeout(stream: S, timeout: Instant) -> Self {
25 let next_deadline = Delay::new(timeout);
29 extend_on_recv: false,
30 timeout: Duration::from_secs(0),
35 impl<S> Stream for TimeoutStream<S> where S : Stream {
37 type Error = S::Error;
38 fn poll(&mut self) -> Result<Async<Option<S::Item>>, S::Error> {
39 match self.next_deadline.poll() {
40 Ok(Async::Ready(_)) => Ok(Async::Ready(None)),
41 Ok(Async::NotReady) => {
42 match self.stream.poll() {
43 Ok(Async::Ready(v)) => {
44 if self.extend_on_recv {
45 self.next_deadline.reset(Instant::now() + self.timeout);
49 Ok(Async::NotReady) => Ok(Async::NotReady),
53 Err(_) => Ok(Async::Ready(None)), // TODO: If I want to upstream TimeoutStream this is gonna need some love