approval-distribution: Fix handling of conclude

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

@@ -2508,14 +2508,18 @@ impl ApprovalDistribution {
 					};
 
 
-					self.handle_from_orchestra(message, &mut approval_voting_sender, &mut network_sender, state, rng).await;
+					if self.handle_from_orchestra(message, &mut approval_voting_sender, &mut network_sender, state, rng).await {
+						return;
+					}
 
 				},
 			}
 		}
 	}
 
 	/// Handles a from orchestra message received by approval distribution subystem.
+	///
+	/// Returns `true` if the subsystem should be stopped.
 	pub async fn handle_from_orchestra<
 		N: overseer::SubsystemSender<NetworkBridgeTxMessage>,
 		A: overseer::SubsystemSender<ApprovalVotingMessage>,
@@ -2526,7 +2530,7 @@ impl ApprovalDistribution {
 		network_sender: &mut N,
 		state: &mut State,
 		rng: &mut (impl CryptoRng + Rng),
-	) {
+	) -> bool {
 		match message {
 			FromOrchestra::Communication { msg } =>
 				Self::handle_incoming(
@@ -2555,8 +2559,9 @@ impl ApprovalDistribution {
 				gum::trace!(target: LOG_TARGET, number = %number, "finalized signal");
 				state.handle_block_finalized(network_sender, &self.metrics, number).await;
 			},
-			FromOrchestra::Signal(OverseerSignal::Conclude) => return,
+			FromOrchestra::Signal(OverseerSignal::Conclude) => return true,
 		}
+		false
 	}
 
 	async fn handle_incoming<

polkadot/node/network/approval-distribution/src/tests.rs

@@ -59,9 +59,13 @@ fn test_harness<T: Future<Output = VirtualOverseer>>(
 	let subsystem = ApprovalDistribution::new(Default::default());
 	{
 		let mut rng = rand_chacha::ChaCha12Rng::seed_from_u64(12345);
-
-		let subsystem =
-			subsystem.run_inner(context, &mut state, REPUTATION_CHANGE_TEST_INTERVAL, &mut rng);
+		let (tx, rx) = oneshot::channel();
+		let subsystem = async {
+			subsystem
+				.run_inner(context, &mut state, REPUTATION_CHANGE_TEST_INTERVAL, &mut rng)
+				.await;
+			tx.send(()).expect("Fail to notify subystem is done");
+		};
 
 		let test_fut = test_fn(virtual_overseer);
 
@@ -76,6 +80,8 @@ fn test_harness<T: Future<Output = VirtualOverseer>>(
 					.timeout(TIMEOUT)
 					.await
 					.expect("Conclude send timeout");
+				let _ =
+					rx.timeout(Duration::from_secs(2)).await.expect("Subsystem did not conclude");
 			},
 			subsystem,
 		));