_ Git - dnsseed-rust/blob - src/timeout_stream.rs

   1 use tokio::prelude::*;
   2 use tokio::timer::Delay;
   3
   4 use std::time::{Duration, Instant};
   5
   6 pub struct TimeoutStream<S> where S : Stream {
   7         stream: S,
   8         next_deadline: Delay,
   9         extend_on_recv: bool,
  10         timeout: Duration,
  11 }
  12
  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);
  16                 Self {
  17                         stream,
  18                         next_deadline,
  19                         extend_on_recv: true,
  20                         timeout,
  21                 }
  22         }
  23
  24         pub fn new_timeout(stream: S, timeout: Instant) -> Self {
  25                 let next_deadline = Delay::new(timeout);
  26                 Self {
  27                         stream,
  28                         next_deadline,
  29                         extend_on_recv: false,
  30                         timeout: Duration::from_secs(0),
  31                 }
  32         }
  33 }
  34
  35 impl<S> Stream for TimeoutStream<S> where S : Stream {
  36         type Item = S::Item;
  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);
  46                                                 }
  47                                                 Ok(Async::Ready(v))
  48                                         },
  49                                         Ok(Async::NotReady) => Ok(Async::NotReady),
  50                                         Err(e) => Err(e),
  51                                 }
  52                         },
  53                         Err(_) => Ok(Async::Ready(None)), // TODO: If I want to upstream TimeoutStream this is gonna need some love
  54                 }
  55         }
  56 }