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

Closes #3774
Closes #3826

polkadot/node/network/dispute-distribution/src/lib.rs

@@ -33,7 +33,7 @@ use polkadot_node_network_protocol::request_response::{incoming::IncomingRequest
 use polkadot_node_primitives::DISPUTE_WINDOW;
 use polkadot_node_subsystem::{
 	messages::DisputeDistributionMessage, overseer, FromOverseer, OverseerSignal, SpawnedSubsystem,
-	SubsystemContext, SubsystemError,
+	SubsystemError,
 };
 use polkadot_node_subsystem_util::{runtime, runtime::RuntimeInfo};
 
@@ -114,12 +114,11 @@ pub struct DisputeDistributionSubsystem<AD> {
 	metrics: Metrics,
 }
 
-impl<Context, AD> overseer::Subsystem<Context, SubsystemError> for DisputeDistributionSubsystem<AD>
+#[overseer::subsystem(DisputeDistribution, error = SubsystemError, prefix = self::overseer)]
+impl<Context, AD> DisputeDistributionSubsystem<AD>
 where
-	Context: SubsystemContext<Message = DisputeDistributionMessage>
-		+ overseer::SubsystemContext<Message = DisputeDistributionMessage>
-		+ Sync
-		+ Send,
+	<Context as overseer::DisputeDistributionContextTrait>::Sender:
+		overseer::DisputeDistributionSenderTrait + Sync + Send,
 	AD: AuthorityDiscovery + Clone,
 {
 	fn start(self, ctx: Context) -> SpawnedSubsystem {
@@ -132,6 +131,7 @@ where
 	}
 }
 
+#[overseer::contextbounds(DisputeDistribution, prefix = self::overseer)]
 impl<AD> DisputeDistributionSubsystem<AD>
 where
 	AD: AuthorityDiscovery + Clone,
@@ -160,13 +160,7 @@ where
 	}
 
 	/// Start processing work as passed on from the Overseer.
-	async fn run<Context>(mut self, mut ctx: Context) -> std::result::Result<(), FatalError>
-	where
-		Context: SubsystemContext<Message = DisputeDistributionMessage>
-			+ overseer::SubsystemContext<Message = DisputeDistributionMessage>
-			+ Sync
-			+ Send,
-	{
+	async fn run<Context>(mut self, mut ctx: Context) -> std::result::Result<(), FatalError> {
 		let receiver = DisputesReceiver::new(
 			ctx.sender().clone(),
 			self.req_receiver
@@ -205,7 +199,7 @@ where
 	}
 
 	/// Handle overseer signals.
-	async fn handle_signals<Context: SubsystemContext>(
+	async fn handle_signals<Context>(
 		&mut self,
 		ctx: &mut Context,
 		signal: OverseerSignal,
@@ -221,7 +215,7 @@ where
 	}
 
 	/// Handle `DisputeDistributionMessage`s.
-	async fn handle_subsystem_message<Context: SubsystemContext>(
+	async fn handle_subsystem_message<Context>(
 		&mut self,
 		ctx: &mut Context,
 		msg: DisputeDistributionMessage,
@@ -243,10 +237,10 @@ enum MuxedMessage {
 	Sender(Option<TaskFinish>),
 }
 
+#[overseer::contextbounds(DisputeDistribution, prefix = self::overseer)]
 impl MuxedMessage {
-	async fn receive(
-		ctx: &mut (impl SubsystemContext<Message = DisputeDistributionMessage>
-		          + overseer::SubsystemContext<Message = DisputeDistributionMessage>),
+	async fn receive<Context>(
+		ctx: &mut Context,
 		from_sender: &mut mpsc::Receiver<TaskFinish>,
 	) -> Self {
 		// We are only fusing here to make `select` happy, in reality we will quit if the stream

polkadot/node/network/dispute-distribution/src/receiver/mod.rs

@@ -40,8 +40,8 @@ use polkadot_node_network_protocol::{
 };
 use polkadot_node_primitives::DISPUTE_WINDOW;
 use polkadot_node_subsystem::{
-	messages::{AllMessages, DisputeCoordinatorMessage, ImportStatementsResult},
-	SubsystemSender,
+	messages::{DisputeCoordinatorMessage, ImportStatementsResult},
+	overseer,
 };
 use polkadot_node_subsystem_util::{runtime, runtime::RuntimeInfo};
 
@@ -132,9 +132,10 @@ impl MuxedMessage {
 	}
 }
 
-impl<Sender: SubsystemSender, AD> DisputesReceiver<Sender, AD>
+impl<Sender, AD> DisputesReceiver<Sender, AD>
 where
 	AD: AuthorityDiscovery,
+	Sender: overseer::DisputeDistributionSenderTrait,
 {
 	/// Create a new receiver which can be `run`.
 	pub fn new(
@@ -265,15 +266,13 @@ where
 		let (pending_confirmation, confirmation_rx) = oneshot::channel();
 		let candidate_hash = candidate_receipt.hash();
 		self.sender
-			.send_message(AllMessages::DisputeCoordinator(
-				DisputeCoordinatorMessage::ImportStatements {
-					candidate_hash,
-					candidate_receipt,
-					session: valid_vote.0.session_index(),
-					statements: vec![valid_vote, invalid_vote],
-					pending_confirmation: Some(pending_confirmation),
-				},
-			))
+			.send_message(DisputeCoordinatorMessage::ImportStatements {
+				candidate_hash,
+				candidate_receipt,
+				session: valid_vote.0.session_index(),
+				statements: vec![valid_vote, invalid_vote],
+				pending_confirmation: Some(pending_confirmation),
+			})
 			.await;
 
 		self.pending_imports.push(peer, confirmation_rx, pending_response);

polkadot/node/network/dispute-distribution/src/sender/mod.rs

@@ -20,10 +20,7 @@ use futures::channel::{mpsc, oneshot};
 
 use polkadot_node_network_protocol::request_response::v1::DisputeRequest;
 use polkadot_node_primitives::{CandidateVotes, DisputeMessage, SignedDisputeStatement};
-use polkadot_node_subsystem::{
-	messages::{AllMessages, DisputeCoordinatorMessage},
-	ActiveLeavesUpdate, SubsystemContext,
-};
+use polkadot_node_subsystem::{messages::DisputeCoordinatorMessage, overseer, ActiveLeavesUpdate};
 use polkadot_node_subsystem_util::runtime::RuntimeInfo;
 use polkadot_primitives::v2::{CandidateHash, DisputeStatement, Hash, SessionIndex};
 
@@ -66,6 +63,7 @@ pub struct DisputeSender {
 	metrics: Metrics,
 }
 
+#[overseer::contextbounds(DisputeDistribution, prefix = self::overseer)]
 impl DisputeSender {
 	/// Create a new `DisputeSender` which can be used to start dispute sendings.
 	pub fn new(tx: mpsc::Sender<TaskFinish>, metrics: Metrics) -> Self {
@@ -79,7 +77,7 @@ impl DisputeSender {
 	}
 
 	/// Create a `SendTask` for a particular new dispute.
-	pub async fn start_sender<Context: SubsystemContext>(
+	pub async fn start_sender<Context>(
 		&mut self,
 		ctx: &mut Context,
 		runtime: &mut RuntimeInfo,
@@ -114,7 +112,7 @@ impl DisputeSender {
 	/// - Get new authorities to send messages to.
 	/// - Get rid of obsolete tasks and disputes.
 	/// - Get dispute sending started in case we missed one for some reason (e.g. on node startup)
-	pub async fn update_leaves<Context: SubsystemContext>(
+	pub async fn update_leaves<Context>(
 		&mut self,
 		ctx: &mut Context,
 		runtime: &mut RuntimeInfo,
@@ -179,7 +177,7 @@ impl DisputeSender {
 	/// Call `start_sender` on all passed in disputes.
 	///
 	/// Recover necessary votes for building up `DisputeMessage` and start sending for all of them.
-	async fn start_send_for_dispute<Context: SubsystemContext>(
+	async fn start_send_for_dispute<Context>(
 		&mut self,
 		ctx: &mut Context,
 		runtime: &mut RuntimeInfo,
@@ -303,7 +301,7 @@ impl DisputeSender {
 	/// Make active sessions correspond to currently active heads.
 	///
 	/// Returns: true if sessions changed.
-	async fn refresh_sessions<Context: SubsystemContext>(
+	async fn refresh_sessions<Context>(
 		&mut self,
 		ctx: &mut Context,
 		runtime: &mut RuntimeInfo,
@@ -321,7 +319,8 @@ impl DisputeSender {
 /// Retrieve the currently active sessions.
 ///
 /// List is all indices of all active sessions together with the head that was used for the query.
-async fn get_active_session_indices<Context: SubsystemContext>(
+#[overseer::contextbounds(DisputeDistribution, prefix = self::overseer)]
+async fn get_active_session_indices<Context>(
 	ctx: &mut Context,
 	runtime: &mut RuntimeInfo,
 	active_heads: &Vec<Hash>,
@@ -336,27 +335,29 @@ async fn get_active_session_indices<Context: SubsystemContext>(
 }
 
 /// Retrieve Set of active disputes from the dispute coordinator.
-async fn get_active_disputes<Context: SubsystemContext>(
+#[overseer::contextbounds(DisputeDistribution, prefix = self::overseer)]
+async fn get_active_disputes<Context>(
 	ctx: &mut Context,
 ) -> JfyiErrorResult<Vec<(SessionIndex, CandidateHash)>> {
 	let (tx, rx) = oneshot::channel();
+
 	// Caller scope is in `update_leaves` and this is bounded by fork count.
-	ctx.send_unbounded_message(AllMessages::DisputeCoordinator(
-		DisputeCoordinatorMessage::ActiveDisputes(tx),
-	));
+	ctx.send_unbounded_message(DisputeCoordinatorMessage::ActiveDisputes(tx));
 	rx.await.map_err(|_| JfyiError::AskActiveDisputesCanceled)
 }
 
 /// Get all locally available dispute votes for a given dispute.
-async fn get_candidate_votes<Context: SubsystemContext>(
+#[overseer::contextbounds(DisputeDistribution, prefix = self::overseer)]
+async fn get_candidate_votes<Context>(
 	ctx: &mut Context,
 	session_index: SessionIndex,
 	candidate_hash: CandidateHash,
 ) -> JfyiErrorResult<Option<CandidateVotes>> {
 	let (tx, rx) = oneshot::channel();
 	// Caller scope is in `update_leaves` and this is bounded by fork count.
-	ctx.send_unbounded_message(AllMessages::DisputeCoordinator(
-		DisputeCoordinatorMessage::QueryCandidateVotes(vec![(session_index, candidate_hash)], tx),
+	ctx.send_unbounded_message(DisputeCoordinatorMessage::QueryCandidateVotes(
+		vec![(session_index, candidate_hash)],
+		tx,
 	));
 	rx.await
 		.map(|v| v.get(0).map(|inner| inner.to_owned().2))

polkadot/node/network/dispute-distribution/src/sender/send_task.rs

@@ -26,10 +26,7 @@ use polkadot_node_network_protocol::{
 	},
 	IfDisconnected,
 };
-use polkadot_node_subsystem::{
-	messages::{AllMessages, NetworkBridgeMessage},
-	SubsystemContext,
-};
+use polkadot_node_subsystem::{messages::NetworkBridgeMessage, overseer};
 use polkadot_node_subsystem_util::{metrics, runtime::RuntimeInfo};
 use polkadot_primitives::v2::{
 	AuthorityDiscoveryId, CandidateHash, Hash, SessionIndex, ValidatorIndex,
@@ -100,9 +97,10 @@ impl TaskResult {
 	}
 }
 
+#[overseer::contextbounds(DisputeDistribution, prefix = self::overseer)]
 impl SendTask {
 	/// Initiates sending a dispute message to peers.
-	pub async fn new<Context: SubsystemContext>(
+	pub async fn new<Context>(
 		ctx: &mut Context,
 		runtime: &mut RuntimeInfo,
 		active_sessions: &HashMap<SessionIndex, Hash>,
@@ -120,7 +118,7 @@ impl SendTask {
 	///
 	/// This function is called at construction and should also be called whenever a session change
 	/// happens and on a regular basis to ensure we are retrying failed attempts.
-	pub async fn refresh_sends<Context: SubsystemContext>(
+	pub async fn refresh_sends<Context>(
 		&mut self,
 		ctx: &mut Context,
 		runtime: &mut RuntimeInfo,
@@ -197,7 +195,8 @@ impl SendTask {
 	///
 	/// This is all parachain validators of the session the candidate occurred and all authorities
 	/// of all currently active sessions, determined by currently active heads.
-	async fn get_relevant_validators<Context: SubsystemContext>(
+
+	async fn get_relevant_validators<Context>(
 		&self,
 		ctx: &mut Context,
 		runtime: &mut RuntimeInfo,
@@ -241,7 +240,8 @@ impl SendTask {
 /// Start sending of the given message to all given authorities.
 ///
 /// And spawn tasks for handling the response.
-async fn send_requests<Context: SubsystemContext>(
+#[overseer::contextbounds(DisputeDistribution, prefix = self::overseer)]
+async fn send_requests<Context>(
 	ctx: &mut Context,
 	tx: mpsc::Sender<TaskFinish>,
 	receivers: Vec<AuthorityDiscoveryId>,
@@ -271,7 +271,7 @@ async fn send_requests<Context: SubsystemContext>(
 	}
 
 	let msg = NetworkBridgeMessage::SendRequests(reqs, IfDisconnected::ImmediateError);
-	ctx.send_message(AllMessages::NetworkBridge(msg)).await;
+	ctx.send_message(msg).await;
 	Ok(statuses)
 }