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

Closes #3774
Closes #3826

polkadot/node/core/chain-selection/src/lib.rs

@@ -20,7 +20,8 @@ use polkadot_node_primitives::BlockWeight;
 use polkadot_node_subsystem::{
 	errors::ChainApiError,
 	messages::{ChainApiMessage, ChainSelectionMessage},
-	overseer, FromOverseer, OverseerSignal, SpawnedSubsystem, SubsystemContext, SubsystemError,
+	overseer::{self, SubsystemSender},
+	FromOverseer, OverseerSignal, SpawnedSubsystem, SubsystemError,
 };
 use polkadot_node_subsystem_util::database::Database;
 use polkadot_primitives::v2::{BlockNumber, ConsensusLog, Hash, Header};
@@ -328,11 +329,8 @@ impl ChainSelectionSubsystem {
 	}
 }
 
-impl<Context> overseer::Subsystem<Context, SubsystemError> for ChainSelectionSubsystem
-where
-	Context: SubsystemContext<Message = ChainSelectionMessage>,
-	Context: overseer::SubsystemContext<Message = ChainSelectionMessage>,
-{
+#[overseer::subsystem(ChainSelection, error = SubsystemError, prefix = self::overseer)]
+impl<Context> ChainSelectionSubsystem {
 	fn start(self, ctx: Context) -> SpawnedSubsystem {
 		let backend = db_backend::v1::DbBackend::new(
 			self.db,
@@ -348,14 +346,13 @@ where
 	}
 }
 
+#[overseer::contextbounds(ChainSelection, prefix = self::overseer)]
 async fn run<Context, B>(
 	mut ctx: Context,
 	mut backend: B,
 	stagnant_check_interval: StagnantCheckInterval,
 	clock: Box<dyn Clock + Send + Sync>,
 ) where
-	Context: SubsystemContext<Message = ChainSelectionMessage>,
-	Context: overseer::SubsystemContext<Message = ChainSelectionMessage>,
 	B: Backend,
 {
 	loop {
@@ -363,7 +360,7 @@ async fn run<Context, B>(
 		match res {
 			Err(e) => {
 				e.trace();
-				// All errors right now are considered fatal:
+				// All errors are considered fatal right now:
 				break
 			},
 			Ok(()) => {
@@ -379,6 +376,7 @@ async fn run<Context, B>(
 //
 // A return value of `Ok` indicates that an exit should be made, while non-fatal errors
 // lead to another call to this function.
+#[overseer::contextbounds(ChainSelection, prefix = self::overseer)]
 async fn run_until_error<Context, B>(
 	ctx: &mut Context,
 	backend: &mut B,
@@ -386,8 +384,6 @@ async fn run_until_error<Context, B>(
 	clock: &(dyn Clock + Sync),
 ) -> Result<(), Error>
 where
-	Context: SubsystemContext<Message = ChainSelectionMessage>,
-	Context: overseer::SubsystemContext<Message = ChainSelectionMessage>,
 	B: Backend,
 {
 	let mut stagnant_check_stream = stagnant_check_interval.timeout_stream();
@@ -402,7 +398,7 @@ where
 					FromOverseer::Signal(OverseerSignal::ActiveLeaves(update)) => {
 						for leaf in update.activated {
 							let write_ops = handle_active_leaf(
-								ctx,
+								ctx.sender(),
 								&*backend,
 								clock.timestamp_now() + STAGNANT_TIMEOUT,
 								leaf.hash,
@@ -419,7 +415,7 @@ where
 							handle_approved_block(backend, hash)?
 						}
 						ChainSelectionMessage::Leaves(tx) => {
-							let leaves = load_leaves(ctx, &*backend).await?;
+							let leaves = load_leaves(ctx.sender(), &*backend).await?;
 							let _ = tx.send(leaves);
 						}
 						ChainSelectionMessage::BestLeafContaining(required, tx) => {
@@ -446,11 +442,11 @@ where
 }
 
 async fn fetch_finalized(
-	ctx: &mut impl SubsystemContext,
+	sender: &mut impl SubsystemSender<ChainApiMessage>,
 ) -> Result<Option<(Hash, BlockNumber)>, Error> {
 	let (number_tx, number_rx) = oneshot::channel();
 
-	ctx.send_message(ChainApiMessage::FinalizedBlockNumber(number_tx)).await;
+	sender.send_message(ChainApiMessage::FinalizedBlockNumber(number_tx)).await;
 
 	let number = match number_rx.await? {
 		Ok(number) => number,
@@ -462,7 +458,7 @@ async fn fetch_finalized(
 
 	let (hash_tx, hash_rx) = oneshot::channel();
 
-	ctx.send_message(ChainApiMessage::FinalizedBlockHash(number, hash_tx)).await;
+	sender.send_message(ChainApiMessage::FinalizedBlockHash(number, hash_tx)).await;
 
 	match hash_rx.await? {
 		Err(err) => {
@@ -478,11 +474,11 @@ async fn fetch_finalized(
 }
 
 async fn fetch_header(
-	ctx: &mut impl SubsystemContext,
+	sender: &mut impl SubsystemSender<ChainApiMessage>,
 	hash: Hash,
 ) -> Result<Option<Header>, Error> {
 	let (tx, rx) = oneshot::channel();
-	ctx.send_message(ChainApiMessage::BlockHeader(hash, tx)).await;
+	sender.send_message(ChainApiMessage::BlockHeader(hash, tx)).await;
 
 	Ok(rx.await?.unwrap_or_else(|err| {
 		gum::warn!(target: LOG_TARGET, ?hash, ?err, "Missing hash for finalized block number");
@@ -491,11 +487,11 @@ async fn fetch_header(
 }
 
 async fn fetch_block_weight(
-	ctx: &mut impl SubsystemContext,
+	sender: &mut impl overseer::SubsystemSender<ChainApiMessage>,
 	hash: Hash,
 ) -> Result<Option<BlockWeight>, Error> {
 	let (tx, rx) = oneshot::channel();
-	ctx.send_message(ChainApiMessage::BlockWeight(hash, tx)).await;
+	sender.send_message(ChainApiMessage::BlockWeight(hash, tx)).await;
 
 	let res = rx.await?;
 
@@ -507,7 +503,7 @@ async fn fetch_block_weight(
 
 // Handle a new active leaf.
 async fn handle_active_leaf(
-	ctx: &mut impl SubsystemContext,
+	sender: &mut impl overseer::ChainSelectionSenderTrait,
 	backend: &impl Backend,
 	stagnant_at: Timestamp,
 	hash: Hash,
@@ -519,10 +515,10 @@ async fn handle_active_leaf(
 			// tree.
 			l.saturating_sub(1)
 		},
-		None => fetch_finalized(ctx).await?.map_or(1, |(_, n)| n),
+		None => fetch_finalized(sender).await?.map_or(1, |(_, n)| n),
 	};
 
-	let header = match fetch_header(ctx, hash).await? {
+	let header = match fetch_header(sender, hash).await? {
 		None => {
 			gum::warn!(target: LOG_TARGET, ?hash, "Missing header for new head");
 			return Ok(Vec::new())
@@ -531,7 +527,7 @@ async fn handle_active_leaf(
 	};
 
 	let new_blocks = polkadot_node_subsystem_util::determine_new_blocks(
-		ctx.sender(),
+		sender,
 		|h| backend.load_block_entry(h).map(|b| b.is_some()),
 		hash,
 		&header,
@@ -544,7 +540,7 @@ async fn handle_active_leaf(
 	// determine_new_blocks gives blocks in descending order.
 	// for this, we want ascending order.
 	for (hash, header) in new_blocks.into_iter().rev() {
-		let weight = match fetch_block_weight(ctx, hash).await? {
+		let weight = match fetch_block_weight(sender, hash).await? {
 			None => {
 				gum::warn!(
 					target: LOG_TARGET,
@@ -654,13 +650,13 @@ fn detect_stagnant(backend: &mut impl Backend, now: Timestamp) -> Result<(), Err
 // Load the leaves from the backend. If there are no leaves, then return
 // the finalized block.
 async fn load_leaves(
-	ctx: &mut impl SubsystemContext,
+	sender: &mut impl overseer::SubsystemSender<ChainApiMessage>,
 	backend: &impl Backend,
 ) -> Result<Vec<Hash>, Error> {
 	let leaves: Vec<_> = backend.load_leaves()?.into_hashes_descending().collect();
 
 	if leaves.is_empty() {
-		Ok(fetch_finalized(ctx).await?.map_or(Vec::new(), |(h, _)| vec![h]))
+		Ok(fetch_finalized(sender).await?.map_or(Vec::new(), |(h, _)| vec![h]))
 	} else {
 		Ok(leaves)
 	}