Extract block announce validation from ChainSync (#1170)

substrate/client/network/sync/src/futures_stream.rs

@@ -0,0 +1,134 @@
+// This file is part of Substrate.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+//! A wrapper for [`FuturesUnordered`] that wakes the task up once a new future is pushed
+//! for it to be polled automatically. It's [`Stream`] never terminates.
+
+use futures::{stream::FuturesUnordered, Future, Stream, StreamExt};
+use std::{
+	pin::Pin,
+	task::{Context, Poll, Waker},
+};
+
+/// Wrapper around [`FuturesUnordered`] that wakes a task up automatically.
+pub struct FuturesStream<F> {
+	futures: FuturesUnordered<F>,
+	waker: Option<Waker>,
+}
+
+/// Surprizingly, `#[derive(Default)]` doesn't work on [`FuturesStream`].
+impl<F> Default for FuturesStream<F> {
+	fn default() -> FuturesStream<F> {
+		FuturesStream { futures: Default::default(), waker: None }
+	}
+}
+
+impl<F> FuturesStream<F> {
+	/// Push a future for processing.
+	pub fn push(&mut self, future: F) {
+		self.futures.push(future);
+
+		if let Some(waker) = self.waker.take() {
+			waker.wake();
+		}
+	}
+
+	/// The number of futures in the stream.
+	pub fn len(&self) -> usize {
+		self.futures.len()
+	}
+}
+
+impl<F: Future> Stream for FuturesStream<F> {
+	type Item = <F as Future>::Output;
+
+	fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+		let Poll::Ready(Some(result)) = self.futures.poll_next_unpin(cx) else {
+			self.waker = Some(cx.waker().clone());
+
+			return Poll::Pending
+		};
+
+		Poll::Ready(Some(result))
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use futures::future::{BoxFuture, FutureExt};
+
+	/// [`Stream`] implementation for [`FuturesStream`] relies on the undocumented
+	/// feature that [`FuturesUnordered`] can be polled and repeatedly yield
+	/// `Poll::Ready(None)` before any futures are added into it.
+	#[tokio::test]
+	async fn empty_futures_unordered_can_be_polled() {
+		let mut unordered = FuturesUnordered::<BoxFuture<()>>::default();
+
+		futures::future::poll_fn(|cx| {
+			assert_eq!(unordered.poll_next_unpin(cx), Poll::Ready(None));
+			assert_eq!(unordered.poll_next_unpin(cx), Poll::Ready(None));
+
+			Poll::Ready(())
+		})
+		.await;
+	}
+
+	/// [`Stream`] implementation for [`FuturesStream`] relies on the undocumented
+	/// feature that [`FuturesUnordered`] can be polled and repeatedly yield
+	/// `Poll::Ready(None)` after all the futures in it have resolved.
+	#[tokio::test]
+	async fn deplenished_futures_unordered_can_be_polled() {
+		let mut unordered = FuturesUnordered::<BoxFuture<()>>::default();
+
+		unordered.push(futures::future::ready(()).boxed());
+		assert_eq!(unordered.next().await, Some(()));
+
+		futures::future::poll_fn(|cx| {
+			assert_eq!(unordered.poll_next_unpin(cx), Poll::Ready(None));
+			assert_eq!(unordered.poll_next_unpin(cx), Poll::Ready(None));
+
+			Poll::Ready(())
+		})
+		.await;
+	}
+
+	#[tokio::test]
+	async fn empty_futures_stream_yields_pending() {
+		let mut stream = FuturesStream::<BoxFuture<()>>::default();
+
+		futures::future::poll_fn(|cx| {
+			assert_eq!(stream.poll_next_unpin(cx), Poll::Pending);
+			Poll::Ready(())
+		})
+		.await;
+	}
+
+	#[tokio::test]
+	async fn futures_stream_resolves_futures_and_yields_pending() {
+		let mut stream = FuturesStream::default();
+		stream.push(futures::future::ready(17));
+
+		futures::future::poll_fn(|cx| {
+			assert_eq!(stream.poll_next_unpin(cx), Poll::Ready(Some(17)));
+			assert_eq!(stream.poll_next_unpin(cx), Poll::Pending);
+			Poll::Ready(())
+		})
+		.await;
+	}
+}