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

Closes #3774
Closes #3826

polkadot/node/network/availability-recovery/src/lib.rs

@@ -51,9 +51,8 @@ use polkadot_node_subsystem::{
 	errors::RecoveryError,
 	jaeger,
 	messages::{AvailabilityRecoveryMessage, AvailabilityStoreMessage, NetworkBridgeMessage},
-	overseer::{self, Subsystem},
-	ActiveLeavesUpdate, FromOverseer, OverseerSignal, SpawnedSubsystem, SubsystemContext,
-	SubsystemError, SubsystemResult, SubsystemSender,
+	overseer, ActiveLeavesUpdate, FromOverseer, OverseerSignal, SpawnedSubsystem, SubsystemError,
+	SubsystemResult,
 };
 use polkadot_node_subsystem_util::request_session_info;
 use polkadot_primitives::v2::{
@@ -156,8 +155,8 @@ enum Source {
 
 /// A stateful reconstruction of availability data in reference to
 /// a candidate hash.
-struct RecoveryTask<S> {
-	sender: S,
+struct RecoveryTask<Sender> {
+	sender: Sender,
 
 	/// The parameters of the recovery process.
 	params: RecoveryParams,
@@ -177,7 +176,7 @@ impl RequestFromBackers {
 	async fn run(
 		&mut self,
 		params: &RecoveryParams,
-		sender: &mut impl SubsystemSender,
+		sender: &mut impl overseer::AvailabilityRecoverySenderTrait,
 	) -> Result<AvailableData, RecoveryError> {
 		gum::trace!(
 			target: LOG_TARGET,
@@ -199,13 +198,10 @@ impl RequestFromBackers {
 			);
 
 			sender
-				.send_message(
-					NetworkBridgeMessage::SendRequests(
-						vec![Requests::AvailableDataFetchingV1(req)],
-						IfDisconnected::ImmediateError,
-					)
-					.into(),
-				)
+				.send_message(NetworkBridgeMessage::SendRequests(
+					vec![Requests::AvailableDataFetchingV1(req)],
+					IfDisconnected::ImmediateError,
+				))
 				.await;
 
 			match response.await {
@@ -298,11 +294,13 @@ impl RequestChunksFromValidators {
 		)
 	}
 
-	async fn launch_parallel_requests(
+	async fn launch_parallel_requests<Sender>(
 		&mut self,
 		params: &RecoveryParams,
-		sender: &mut impl SubsystemSender,
-	) {
+		sender: &mut Sender,
+	) where
+		Sender: overseer::AvailabilityRecoverySenderTrait,
+	{
 		let num_requests = self.get_desired_request_count(params.threshold);
 		let candidate_hash = &params.candidate_hash;
 		let already_requesting_count = self.requesting_chunks.len();
@@ -358,9 +356,10 @@ impl RequestChunksFromValidators {
 		}
 
 		sender
-			.send_message(
-				NetworkBridgeMessage::SendRequests(requests, IfDisconnected::ImmediateError).into(),
-			)
+			.send_message(NetworkBridgeMessage::SendRequests(
+				requests,
+				IfDisconnected::ImmediateError,
+			))
 			.await;
 	}
 
@@ -483,20 +482,21 @@ impl RequestChunksFromValidators {
 		}
 	}
 
-	async fn run(
+	async fn run<Sender>(
 		&mut self,
 		params: &RecoveryParams,
-		sender: &mut impl SubsystemSender,
-	) -> Result<AvailableData, RecoveryError> {
+		sender: &mut Sender,
+	) -> Result<AvailableData, RecoveryError>
+	where
+		Sender: overseer::AvailabilityRecoverySenderTrait,
+	{
 		let metrics = &params.metrics;
 
 		// First query the store for any chunks we've got.
 		{
 			let (tx, rx) = oneshot::channel();
 			sender
-				.send_message(
-					AvailabilityStoreMessage::QueryAllChunks(params.candidate_hash, tx).into(),
-				)
+				.send_message(AvailabilityStoreMessage::QueryAllChunks(params.candidate_hash, tx))
 				.await;
 
 			match rx.await {
@@ -646,16 +646,19 @@ fn reconstructed_data_matches_root(
 	branches.root() == *expected_root
 }
 
-impl<S: SubsystemSender> RecoveryTask<S> {
+impl<Sender> RecoveryTask<Sender>
+where
+	Sender: overseer::AvailabilityRecoverySenderTrait,
+{
 	async fn run(mut self) -> Result<AvailableData, RecoveryError> {
 		// First just see if we have the data available locally.
 		{
 			let (tx, rx) = oneshot::channel();
 			self.sender
-				.send_message(
-					AvailabilityStoreMessage::QueryAvailableData(self.params.candidate_hash, tx)
-						.into(),
-				)
+				.send_message(AvailabilityStoreMessage::QueryAvailableData(
+					self.params.candidate_hash,
+					tx,
+				))
 				.await;
 
 			match rx.await {
@@ -799,11 +802,8 @@ impl Default for State {
 	}
 }
 
-impl<Context> Subsystem<Context, SubsystemError> for AvailabilityRecoverySubsystem
-where
-	Context: SubsystemContext<Message = AvailabilityRecoveryMessage>,
-	Context: overseer::SubsystemContext<Message = AvailabilityRecoveryMessage>,
-{
+#[overseer::subsystem(AvailabilityRecovery, error=SubsystemError, prefix=self::overseer)]
+impl<Context> AvailabilityRecoverySubsystem {
 	fn start(self, ctx: Context) -> SpawnedSubsystem {
 		let future = self
 			.run(ctx)
@@ -832,6 +832,7 @@ async fn handle_signal(state: &mut State, signal: OverseerSignal) -> SubsystemRe
 }
 
 /// Machinery around launching recovery tasks into the background.
+#[overseer::contextbounds(AvailabilityRecovery, prefix = self::overseer)]
 async fn launch_recovery_task<Context>(
 	state: &mut State,
 	ctx: &mut Context,
@@ -840,11 +841,7 @@ async fn launch_recovery_task<Context>(
 	backing_group: Option<GroupIndex>,
 	response_sender: oneshot::Sender<Result<AvailableData, RecoveryError>>,
 	metrics: &Metrics,
-) -> error::Result<()>
-where
-	Context: SubsystemContext<Message = AvailabilityRecoveryMessage>,
-	Context: overseer::SubsystemContext<Message = AvailabilityRecoveryMessage>,
-{
+) -> error::Result<()> {
 	let candidate_hash = receipt.hash();
 
 	let params = RecoveryParams {
@@ -885,6 +882,7 @@ where
 }
 
 /// Handles an availability recovery request.
+#[overseer::contextbounds(AvailabilityRecovery, prefix = self::overseer)]
 async fn handle_recover<Context>(
 	state: &mut State,
 	ctx: &mut Context,
@@ -893,11 +891,7 @@ async fn handle_recover<Context>(
 	backing_group: Option<GroupIndex>,
 	response_sender: oneshot::Sender<Result<AvailableData, RecoveryError>>,
 	metrics: &Metrics,
-) -> error::Result<()>
-where
-	Context: SubsystemContext<Message = AvailabilityRecoveryMessage>,
-	Context: overseer::SubsystemContext<Message = AvailabilityRecoveryMessage>,
-{
+) -> error::Result<()> {
 	let candidate_hash = receipt.hash();
 
 	let span = jaeger::Span::new(candidate_hash, "availbility-recovery")
@@ -953,14 +947,11 @@ where
 }
 
 /// Queries a chunk from av-store.
+#[overseer::contextbounds(AvailabilityRecovery, prefix = self::overseer)]
 async fn query_full_data<Context>(
 	ctx: &mut Context,
 	candidate_hash: CandidateHash,
-) -> error::Result<Option<AvailableData>>
-where
-	Context: SubsystemContext<Message = AvailabilityRecoveryMessage>,
-	Context: overseer::SubsystemContext<Message = AvailabilityRecoveryMessage>,
-{
+) -> error::Result<Option<AvailableData>> {
 	let (tx, rx) = oneshot::channel();
 	ctx.send_message(AvailabilityStoreMessage::QueryAvailableData(candidate_hash, tx))
 		.await;
@@ -968,6 +959,7 @@ where
 	Ok(rx.await.map_err(error::Error::CanceledQueryFullData)?)
 }
 
+#[overseer::contextbounds(AvailabilityRecovery, prefix = self::overseer)]
 impl AvailabilityRecoverySubsystem {
 	/// Create a new instance of `AvailabilityRecoverySubsystem` which starts with a fast path to
 	/// request data from backers.
@@ -986,11 +978,7 @@ impl AvailabilityRecoverySubsystem {
 		Self { fast_path: false, req_receiver, metrics }
 	}
 
-	async fn run<Context>(self, mut ctx: Context) -> SubsystemResult<()>
-	where
-		Context: SubsystemContext<Message = AvailabilityRecoveryMessage>,
-		Context: overseer::SubsystemContext<Message = AvailabilityRecoveryMessage>,
-	{
+	async fn run<Context>(self, mut ctx: Context) -> SubsystemResult<()> {
 		let mut state = State::default();
 		let Self { fast_path, mut req_receiver, metrics } = self;