X-Git-Url: http://git.bitcoin.ninja/index.cgi?a=blobdiff_plain;f=src%2Ftimeout_stream.rs;h=0d0caea3d1829fb1c5adc4c8b1c364c01698535c;hb=1ad3d257518100cc938b4dabfff8636414131ce0;hp=57664f8f5de7fe0fb89b11d4b44180747be15fbd;hpb=8808d87877909d7d4f97d7ff320d8d165ece5800;p=dnsseed-rust

diff --git a/src/timeout_stream.rs b/src/timeout_stream.rs
index 57664f8..0d0caea 100644
--- 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),