Total runtime timeout, fix saves

[dnsseed-rust] / src / timeout_stream.rs

diff --git a/src/timeout_stream.rs b/src/timeout_stream.rs
index 57664f8f5de7fe0fb89b11d4b44180747be15fbd..0d0caea3d1829fb1c5adc4c8b1c364c01698535c 100644 (file)
--- a/src/timeout_stream.rs
+++ b/src/timeout_stream.rs
@@ -6,18 +6,30 @@ use std::time::{Duration, Instant};
 pub struct TimeoutStream<S> where S : Stream {
        stream: S,
        next_deadline: Delay,
+       extend_on_recv: bool,
        timeout: Duration,
 }
 
 impl<S> TimeoutStream<S> where S : Stream {
-       pub fn new(stream: S, timeout: Duration) -> Self {
+       pub fn new_persistent(stream: S, timeout: Duration) -> Self {
                let next_deadline = Delay::new(Instant::now() + timeout);
                Self {
                        stream,
                        next_deadline,
+                       extend_on_recv: true,
                        timeout,
                }
        }
+
+       pub fn new_timeout(stream: S, timeout: Instant) -> Self {
+               let next_deadline = Delay::new(timeout);
+               Self {
+                       stream,
+                       next_deadline,
+                       extend_on_recv: false,
+                       timeout: Duration::from_secs(0),
+               }
+       }
 }
 
 impl<S> Stream for TimeoutStream<S> where S : Stream {
@@ -29,7 +41,9 @@ impl<S> Stream for TimeoutStream<S> where S : Stream {
                        Ok(Async::NotReady) => {
                                match self.stream.poll() {
                                        Ok(Async::Ready(v)) => {
-                                               self.next_deadline.reset(Instant::now() + self.timeout);
+                                               if self.extend_on_recv {
+                                                       self.next_deadline.reset(Instant::now() + self.timeout);
+                                               }
                                                Ok(Async::Ready(v))
                                        },
                                        Ok(Async::NotReady) => Ok(Async::NotReady),