metered mpsc channels (#2235)

polkadot/node/subsystem-util/src/lib.rs

@@ -50,6 +50,7 @@ use streamunordered::{StreamUnordered, StreamYield};
 use thiserror::Error;
 
 pub mod validator_discovery;
+pub use metered_channel as metered;
 
 /// These reexports are required so that external crates can use the `delegated_subsystem` macro properly.
 pub mod reexports {
@@ -987,9 +988,64 @@ impl<F: Future> Future for Timeout<F> {
 	}
 }
 
+
+#[derive(Copy, Clone)]
+enum MetronomeState {
+	Snooze,
+	SetAlarm,
+}
+
+/// Create a stream of ticks with a defined cycle duration.
+pub struct Metronome {
+	delay: Delay,
+	period: Duration,
+	state: MetronomeState,
+}
+
+impl Metronome
+{
+	/// Create a new metronome source with a defined cycle duration.
+	pub fn new(cycle: Duration) -> Self {
+		let period = cycle.into();
+		Self {
+			period,
+			delay: Delay::new(period),
+			state: MetronomeState::Snooze,
+		}
+	}
+}
+
+impl futures::Stream for Metronome
+{
+	type Item = ();
+	fn poll_next(
+		mut self: Pin<&mut Self>,
+		cx: &mut Context<'_>
+	) -> Poll<Option<Self::Item>> {
+		loop {
+			match self.state {
+				MetronomeState::SetAlarm => {
+					let val = self.period.clone();
+					self.delay.reset(val);
+					self.state = MetronomeState::Snooze;
+				}
+				MetronomeState::Snooze => {
+					if !Pin::new(&mut self.delay).poll(cx).is_ready() {
+						break
+					}
+					self.state = MetronomeState::SetAlarm;
+					return Poll::Ready(Some(()));
+				}
+			}
+		}
+		Poll::Pending
+	}
+}
+
 #[cfg(test)]
 mod tests {
 	use super::*;
+	use executor::block_on;
 	use thiserror::Error;
 	use polkadot_node_subsystem::{
 		messages::{AllMessages, CandidateSelectionMessage}, ActiveLeavesUpdate, FromOverseer, OverseerSignal,
@@ -999,7 +1055,7 @@ mod tests {
 	use futures::{channel::mpsc, executor, StreamExt, future, Future, FutureExt, SinkExt};
 	use polkadot_primitives::v1::Hash;
 	use polkadot_node_subsystem_test_helpers::{self as test_helpers, make_subsystem_context};
-	use std::{pin::Pin, time::Duration, sync::Arc};
+	use std::{pin::Pin, sync::{Arc, atomic::{AtomicUsize, Ordering}}, time::Duration};
 
 	// basic usage: in a nutshell, when you want to define a subsystem, just focus on what its jobs do;
 	// you can leave the subsystem itself to the job manager.
@@ -1183,4 +1239,46 @@ mod tests {
 			FakeCandidateSelectionSubsystem::new(pool, false, ()).start(context);
 		assert_eq!(name, "FakeCandidateSelection");
 	}
+
+
+	#[test]
+	fn tick_tack_metronome() {
+		let n = Arc::new(AtomicUsize::default());
+
+		let (tick, mut block) = mpsc::unbounded();
+
+		let metronome = {
+			let n = n.clone();
+			let stream = Metronome::new(Duration::from_millis(137_u64));
+			stream.for_each(move |_res| {
+				let _ = n.fetch_add(1, Ordering::Relaxed);
+				let mut tick = tick.clone();
+				async move {
+					tick.send(()).await.expect("Test helper channel works. qed");
+				}
+			}).fuse()
+		};
+
+		let f2 = async move {
+			block.next().await;
+			assert_eq!(n.load(Ordering::Relaxed), 1_usize);
+			block.next().await;
+			assert_eq!(n.load(Ordering::Relaxed), 2_usize);
+			block.next().await;
+			assert_eq!(n.load(Ordering::Relaxed), 3_usize);
+			block.next().await;
+			assert_eq!(n.load(Ordering::Relaxed), 4_usize);
+		}.fuse();
+
+		futures::pin_mut!(f2);
+		futures::pin_mut!(metronome);
+
+		block_on(async move {
+			// futures::join!(metronome, f2)
+			futures::select!(
+				_ = metronome => unreachable!("Metronome never stops. qed"),
+				_ = f2 => (),
+			)
+		});
+	}
 }