refactor+feat: allow subsystems to send only declared messages, generate graphviz (#5314)

Closes #3774
Closes #3826

polkadot/node/collation-generation/src/lib.rs

@@ -22,9 +22,9 @@ use futures::{channel::mpsc, future::FutureExt, join, select, sink::SinkExt, str
 use parity_scale_codec::Encode;
 use polkadot_node_primitives::{AvailableData, CollationGenerationConfig, PoV};
 use polkadot_node_subsystem::{
-	messages::{AllMessages, CollationGenerationMessage, CollatorProtocolMessage},
+	messages::{CollationGenerationMessage, CollatorProtocolMessage},
 	overseer, ActiveLeavesUpdate, FromOverseer, OverseerSignal, SpawnedSubsystem, SubsystemContext,
-	SubsystemError, SubsystemResult, SubsystemSender,
+	SubsystemError, SubsystemResult,
 };
 use polkadot_node_subsystem_util::{
 	request_availability_cores, request_persisted_validation_data, request_validation_code,
@@ -54,6 +54,7 @@ pub struct CollationGenerationSubsystem {
 	metrics: Metrics,
 }
 
+#[overseer::contextbounds(CollationGeneration, prefix = self::overseer)]
 impl CollationGenerationSubsystem {
 	/// Create a new instance of the `CollationGenerationSubsystem`.
 	pub fn new(metrics: Metrics) -> Self {
@@ -71,11 +72,7 @@ impl CollationGenerationSubsystem {
 	///
 	/// If `err_tx` is not `None`, errors are forwarded onto that channel as they occur.
 	/// Otherwise, most are logged and then discarded.
-	async fn run<Context>(mut self, mut ctx: Context)
-	where
-		Context: SubsystemContext<Message = CollationGenerationMessage>,
-		Context: overseer::SubsystemContext<Message = CollationGenerationMessage>,
-	{
+	async fn run<Context>(mut self, mut ctx: Context) {
 		// when we activate new leaves, we spawn a bunch of sub-tasks, each of which is
 		// expected to generate precisely one message. We don't want to block the main loop
 		// at any point waiting for them all, so instead, we create a channel on which they can
@@ -108,12 +105,8 @@ impl CollationGenerationSubsystem {
 		&mut self,
 		incoming: SubsystemResult<FromOverseer<<Context as SubsystemContext>::Message>>,
 		ctx: &mut Context,
-		sender: &mpsc::Sender<AllMessages>,
-	) -> bool
-	where
-		Context: SubsystemContext<Message = CollationGenerationMessage>,
-		Context: overseer::SubsystemContext<Message = CollationGenerationMessage>,
-	{
+		sender: &mpsc::Sender<overseer::CollationGenerationOutgoingMessages>,
+	) -> bool {
 		match incoming {
 			Ok(FromOverseer::Signal(OverseerSignal::ActiveLeaves(ActiveLeavesUpdate {
 				activated,
@@ -162,11 +155,8 @@ impl CollationGenerationSubsystem {
 	}
 }
 
-impl<Context> overseer::Subsystem<Context, SubsystemError> for CollationGenerationSubsystem
-where
-	Context: SubsystemContext<Message = CollationGenerationMessage>,
-	Context: overseer::SubsystemContext<Message = CollationGenerationMessage>,
-{
+#[overseer::subsystem(CollationGeneration, error=SubsystemError, prefix=self::overseer)]
+impl<Context> CollationGenerationSubsystem {
 	fn start(self, ctx: Context) -> SpawnedSubsystem {
 		let future = async move {
 			self.run(ctx).await;
@@ -178,12 +168,13 @@ where
 	}
 }
 
-async fn handle_new_activations<Context: SubsystemContext>(
+#[overseer::contextbounds(CollationGeneration, prefix = self::overseer)]
+async fn handle_new_activations<Context>(
 	config: Arc<CollationGenerationConfig>,
 	activated: impl IntoIterator<Item = Hash>,
 	ctx: &mut Context,
 	metrics: Metrics,
-	sender: &mpsc::Sender<AllMessages>,
+	sender: &mpsc::Sender<overseer::CollationGenerationOutgoingMessages>,
 ) -> crate::error::Result<()> {
 	// follow the procedure from the guide:
 	// https://w3f.github.io/parachain-implementers-guide/node/collators/collation-generation.html
@@ -393,9 +384,10 @@ async fn handle_new_activations<Context: SubsystemContext>(
 					metrics.on_collation_generated();
 
 					if let Err(err) = task_sender
-						.send(AllMessages::CollatorProtocol(
-							CollatorProtocolMessage::DistributeCollation(ccr, pov, result_sender),
-						))
+						.send(
+							CollatorProtocolMessage::DistributeCollation(ccr, pov, result_sender)
+								.into(),
+						)
 						.await
 					{
 						gum::warn!(
@@ -417,7 +409,7 @@ async fn obtain_current_validation_code_hash(
 	relay_parent: Hash,
 	para_id: ParaId,
 	assumption: OccupiedCoreAssumption,
-	sender: &mut impl SubsystemSender,
+	sender: &mut impl overseer::CollationGenerationSenderTrait,
 ) -> Result<Option<ValidationCodeHash>, crate::error::Error> {
 	use polkadot_node_subsystem::RuntimeApiError;
 

polkadot/node/collation-generation/src/tests.rs

@@ -296,7 +296,7 @@ mod handle_new_activations {
 			*subsystem_sent_messages.lock().await = rx.collect().await;
 		});
 
-		let sent_messages = Arc::try_unwrap(sent_messages)
+		let mut sent_messages = Arc::try_unwrap(sent_messages)
 			.expect("subsystem should have shut down by now")
 			.into_inner();
 
@@ -328,7 +328,7 @@ mod handle_new_activations {
 		};
 
 		assert_eq!(sent_messages.len(), 1);
-		match &sent_messages[0] {
+		match AllMessages::from(sent_messages.pop().unwrap()) {
 			AllMessages::CollatorProtocol(CollatorProtocolMessage::DistributeCollation(
 				CandidateReceipt { descriptor, .. },
 				_pov,
@@ -356,7 +356,7 @@ mod handle_new_activations {
 					expect_descriptor.erasure_root = descriptor.erasure_root.clone();
 					expect_descriptor
 				};
-				assert_eq!(descriptor, &expect_descriptor);
+				assert_eq!(descriptor, expect_descriptor);
 			},
 			_ => panic!("received wrong message type"),
 		}
@@ -470,11 +470,13 @@ mod handle_new_activations {
 
 		assert_eq!(sent_messages.len(), 1);
 		match &sent_messages[0] {
-			AllMessages::CollatorProtocol(CollatorProtocolMessage::DistributeCollation(
-				CandidateReceipt { descriptor, .. },
-				_pov,
-				..,
-			)) => {
+			overseer::CollationGenerationOutgoingMessages::CollatorProtocolMessage(
+				CollatorProtocolMessage::DistributeCollation(
+					CandidateReceipt { descriptor, .. },
+					_pov,
+					..,
+				),
+			) => {
 				assert_eq!(expect_validation_code_hash, descriptor.validation_code_hash);
 			},
 			_ => panic!("received wrong message type"),