Implement separate signal and message flows

orchestra/proc-macro/src/impl_subsystem_ctx_sender.rs

@@ -494,6 +494,27 @@ pub(crate) fn impl_subsystem_context_trait_for(
 				}
 			}
 
+			async fn recv_signal(&mut self) -> ::std::result::Result<#signal, #error_ty> {
+				self.signals.next().await.ok_or(#support_crate ::OrchestraError::Context(
+					"Signal channel is terminated and empty.".to_owned(),
+				).into())
+			}
+
+			async fn recv_msg(&mut self) -> ::std::result::Result<Self::Message, #error_ty> {
+				loop {
+					if let Some((needs_signals_received, msg)) = self.pending_incoming.take() {
+						self.signals_received.wait_until(|v| v >= needs_signals_received).await;
+						return Ok(msg);
+					}
+					let msg = self.messages.next().await.ok_or(
+						#support_crate ::OrchestraError::Context(
+							"Message channel is terminated and empty.".to_owned()
+						)
+					)?;
+					self.pending_incoming = Some((msg.signals_received, msg.message));
+				}
+			}
+
 			fn sender(&mut self) -> &mut Self::Sender {
 				&mut self.to_subsystems
 			}

orchestra/src/lib.rs

@@ -75,7 +75,7 @@ pub use futures::{
 	future::{BoxFuture, Fuse, Future},
 	poll, select,
 	stream::{self, select, select_with_strategy, FuturesUnordered, PollNext},
-	task::{Context, Poll},
+	task::{AtomicWaker, Context, Poll},
 	FutureExt, StreamExt,
 };
 #[doc(hidden)]
@@ -217,22 +217,54 @@ pub type SubsystemIncomingMessages<M> = self::stream::SelectWithStrategy<
 	(),
 >;
 
+#[derive(Debug, Default)]
+struct SignalsReceivedInner {
+	waker: AtomicWaker,
+	value: AtomicUsize,
+}
+
+/// Future to wait on for the watermark predicate
+pub struct SignalsReceivedWaiter<'a, F: Fn(usize) -> bool> {
+	owner: &'a SignalsReceivedInner,
+	predicate: F,
+}
+
+impl<F: Fn(usize) -> bool> Future for SignalsReceivedWaiter<'_, F> {
+	type Output = usize;
+
+	fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+		self.owner.waker.register(cx.waker());
+		let value = self.owner.value.load(atomic::Ordering::Acquire);
+		if (self.predicate)(value) {
+			Poll::Ready(value)
+		} else {
+			Poll::Pending
+		}
+	}
+}
+
 /// Watermark to track the received signals.
 #[derive(Debug, Default, Clone)]
-pub struct SignalsReceived(Arc<AtomicUsize>);
+pub struct SignalsReceived(Arc<SignalsReceivedInner>);
 
 impl SignalsReceived {
 	/// Load the current value of received signals.
 	pub fn load(&self) -> usize {
 		// It's imperative that we prevent reading a stale value from memory because of reordering.
 		// Memory barrier to ensure that no reads or writes in the current thread before this load are reordered.
 		// All writes in other threads using release semantics become visible to the current thread.
-		self.0.load(atomic::Ordering::Acquire)
+		self.0.value.load(atomic::Ordering::Acquire)
 	}
 
 	/// Increase the number of signals by one.
 	pub fn inc(&self) {
-		let _previous = self.0.fetch_add(1, atomic::Ordering::AcqRel);
+		let _previous = self.0.value.fetch_add(1, atomic::Ordering::AcqRel);
+		self.0.waker.wake();
+	}
+
+	/// Wait until a predicate for the watermark is true.
+	pub fn wait_until<F: Fn(usize) -> bool>(&self, predicate: F) -> SignalsReceivedWaiter<F> {
+		SignalsReceivedWaiter { owner: &self.0, predicate }
 	}
 }
 
@@ -416,9 +448,15 @@ pub trait SubsystemContext: Send + 'static {
 	/// using `pending!()` macro you will end up with a busy loop!
 	async fn try_recv(&mut self) -> Result<Option<FromOrchestra<Self::Message, Self::Signal>>, ()>;
 
-	/// Receive a message.
+	/// Receive a signal or a message.
 	async fn recv(&mut self) -> Result<FromOrchestra<Self::Message, Self::Signal>, Self::Error>;
 
+	/// Receive a signal.
+	async fn recv_signal(&mut self) -> Result<Self::Signal, Self::Error>;
+
+	/// Receive a message.
+	async fn recv_msg(&mut self) -> Result<Self::Message, Self::Error>;
+
 	/// Spawn a child task on the executor.
 	fn spawn(
 		&mut self,