diff --git a/polkadot/node/network/approval-distribution/src/lib.rs b/polkadot/node/network/approval-distribution/src/lib.rs index 3462aaef1f696..134586253ce16 100644 --- a/polkadot/node/network/approval-distribution/src/lib.rs +++ b/polkadot/node/network/approval-distribution/src/lib.rs @@ -2508,7 +2508,9 @@ 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; + } }, } @@ -2516,6 +2518,8 @@ impl ApprovalDistribution { } /// 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, A: overseer::SubsystemSender, @@ -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< diff --git a/polkadot/node/network/approval-distribution/src/tests.rs b/polkadot/node/network/approval-distribution/src/tests.rs index 3ea722c51a927..1ca571721ea90 100644 --- a/polkadot/node/network/approval-distribution/src/tests.rs +++ b/polkadot/node/network/approval-distribution/src/tests.rs @@ -59,9 +59,13 @@ fn test_harness>( 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>( .timeout(TIMEOUT) .await .expect("Conclude send timeout"); + let _ = + rx.timeout(Duration::from_secs(2)).await.expect("Subsystem did not conclude"); }, subsystem, ));