async_dnssd/
timeout_stream.rs

1use futures_core::Stream;
2use std::{
3	future::Future,
4	pin::Pin,
5	task::{
6		Context,
7		Poll,
8	},
9	time::Duration,
10};
11
12/// `Stream` extension to simplify building
13/// [`TimeoutStream`](struct.TimeoutStream.html)
14pub trait StreamTimeoutExt: Stream + Sized {
15	/// Create new [`TimeoutStream`](struct.TimeoutStream.html)
16	fn timeout(self, duration: Duration) -> TimeoutStream<Self>;
17}
18
19impl<S: Stream> StreamTimeoutExt for S {
20	fn timeout(self, duration: Duration) -> TimeoutStream<Self> {
21		TimeoutStream::new(self, duration)
22	}
23}
24
25/// Add a timeout to a stream; each time an item is received the timer
26/// is reset
27///
28/// If the timeout triggers the stream ends (without an error).
29#[must_use = "streams do nothing unless polled"]
30pub struct TimeoutStream<S> {
31	stream: S,
32	duration: Duration,
33	timeout: tokio::time::Sleep,
34}
35
36impl<S: Stream> TimeoutStream<S> {
37	pin_utils::unsafe_pinned!(stream: S);
38
39	pin_utils::unsafe_pinned!(timeout: tokio::time::Sleep);
40
41	/// Create new `TimeoutStream`.
42	///
43	/// Also see [`StreamTimeoutExt::timeout`](trait.StreamTimeoutExt.html#method.timeout).
44	pub fn new(stream: S, duration: Duration) -> Self {
45		Self {
46			stream,
47			duration,
48			timeout: tokio::time::sleep(duration),
49		}
50	}
51}
52
53impl<S: Stream> TimeoutStream<S> {
54	fn reset_timer(self: Pin<&mut Self>) {
55		let next = tokio::time::Instant::now() + self.duration;
56		self.timeout().reset(next);
57	}
58}
59
60impl<S: futures_core::TryStream> Stream for TimeoutStream<S> {
61	type Item = Result<S::Ok, S::Error>;
62
63	fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
64		match self.as_mut().stream().try_poll_next(cx) {
65			Poll::Ready(None) => Poll::Ready(None), // end of stream
66			Poll::Ready(Some(Ok(item))) => {
67				// not end of stream: reset timeout
68				self.reset_timer();
69				Poll::Ready(Some(Ok(item)))
70			},
71			Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
72			Poll::Pending => {
73				// check timeout
74				match self.timeout().poll(cx) {
75					// timed out?
76					Poll::Ready(()) => {
77						// not an error
78						Poll::Ready(None)
79					},
80					// still time left
81					Poll::Pending => Poll::Pending,
82				}
83			},
84		}
85	}
86}
async_dnssd/timeout_stream.rs

async_dnssd/
timeout_stream.rs
