Rework the event system of sc-network (#1370)

This commit introduces a new concept called `NotificationService` which allows Polkadot protocols to communicate with the underlying notification protocol implementation directly, without routing events through `NetworkWorker`. This implies that each protocol has its own service which it uses to communicate with remote peers and that each `NotificationService` is unique with respect to the underlying notification protocol, meaning `NotificationService` for the transaction protocol can only be used to send and receive transaction-related notifications. The `NotificationService` concept introduces two additional benefits: * allow protocols to start using custom handshakes * allow protocols to accept/reject inbound peers Previously the validation of inbound connections was solely the responsibility of `ProtocolController`. This caused issues with light peers and `SyncingEngine` as `ProtocolController` would accept more peers than `SyncingEngine` could accept which caused peers to have differing views of their own states. `SyncingEngine` would reject excess peers but these rejections were not properly communicated to those peers causing them to assume that they were accepted. With `NotificationService`, the local handshake is not sent to remote peer if peer is rejected which allows it to detect that it was rejected. This commit also deprecates the use of `NetworkEventStream` for all notification-related events and going forward only DHT events are provided through `NetworkEventStream`. If protocols wish to follow each other's events, they must introduce additional abtractions, as is done for GRANDPA and transactions protocols by following the syncing protocol through `SyncEventStream`. Fixes https://github.com/paritytech/polkadot-sdk/issues/512 Fixes https://github.com/paritytech/polkadot-sdk/issues/514 Fixes https://github.com/paritytech/polkadot-sdk/issues/515 Fixes https://github.com/paritytech/polkadot-sdk/issues/554 Fixes https://github.com/paritytech/polkadot-sdk/issues/556 --- These changes are transferred from https://github.com/paritytech/substrate/pull/14197 but there are no functional changes compared to that PR --------- Co-authored-by: Dmitry Markin <dmitry@markin.tech> Co-authored-by: Alexandru Vasile <60601340+lexnv@users.noreply.github.com>
2026-06-14 01:41:09 +00:00 · 2023-11-28 20:18:52 +02:00
parent ec3a61ed86
commit e71c484d5b
102 changed files with 5694 additions and 2603 deletions
@@ -83,6 +83,7 @@ pub(crate) enum WireMessage<M> {
 	ViewUpdate(View),
 }

+#[derive(Debug)]
 pub(crate) struct PeerData {
 	/// The Latest view sent by the peer.
 	view: View,
@@ -14,23 +14,24 @@
 // You should have received a copy of the GNU General Public License
 // along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.

-use std::{collections::HashSet, sync::Arc};
+use std::{
+	collections::{HashMap, HashSet},
+	sync::Arc,
+};

 use async_trait::async_trait;
-use futures::{prelude::*, stream::BoxStream};
+use parking_lot::Mutex;

 use parity_scale_codec::Encode;

 use sc_network::{
-	config::parse_addr, multiaddr::Multiaddr, types::ProtocolName, Event as NetworkEvent,
-	IfDisconnected, NetworkEventStream, NetworkNotification, NetworkPeers, NetworkRequest,
-	NetworkService, OutboundFailure, ReputationChange, RequestFailure,
+	config::parse_addr, multiaddr::Multiaddr, types::ProtocolName, IfDisconnected, MessageSink,
+	NetworkPeers, NetworkRequest, NetworkService, OutboundFailure, ReputationChange,
+	RequestFailure,
 };

 use polkadot_node_network_protocol::{
-	peer_set::{
-		CollationVersion, PeerSet, PeerSetProtocolNames, ProtocolVersion, ValidationVersion,
-	},
+	peer_set::{CollationVersion, PeerSet, ProtocolVersion, ValidationVersion},
 	request_response::{OutgoingRequest, Recipient, ReqProtocolNames, Requests},
 	v1 as protocol_v1, v2 as protocol_v2, vstaging as protocol_vstaging, PeerId,
 };
@@ -44,104 +45,94 @@ const LOG_TARGET: &'static str = "parachain::network-bridge-net";
 // Helper function to send a validation v1 message to a list of peers.
 // Messages are always sent via the main protocol, even legacy protocol messages.
 pub(crate) fn send_validation_message_v1(
-	net: &mut impl Network,
 	peers: Vec<PeerId>,
-	peerset_protocol_names: &PeerSetProtocolNames,
 	message: WireMessage<protocol_v1::ValidationProtocol>,
 	metrics: &Metrics,
+	notification_sinks: &Arc<Mutex<HashMap<(PeerSet, PeerId), Box<dyn MessageSink>>>>,
 ) {
 	gum::trace!(target: LOG_TARGET, ?peers, ?message, "Sending validation v1 message to peers",);

 	send_message(
-		net,
 		peers,
 		PeerSet::Validation,
 		ValidationVersion::V1.into(),
-		peerset_protocol_names,
 		message,
 		metrics,
+		notification_sinks,
 	);
 }

 // Helper function to send a validation vstaging message to a list of peers.
 // Messages are always sent via the main protocol, even legacy protocol messages.
 pub(crate) fn send_validation_message_vstaging(
-	net: &mut impl Network,
 	peers: Vec<PeerId>,
-	peerset_protocol_names: &PeerSetProtocolNames,
 	message: WireMessage<protocol_vstaging::ValidationProtocol>,
 	metrics: &Metrics,
+	notification_sinks: &Arc<Mutex<HashMap<(PeerSet, PeerId), Box<dyn MessageSink>>>>,
 ) {
 	gum::trace!(target: LOG_TARGET, ?peers, ?message, "Sending validation vstaging message to peers",);

 	send_message(
-		net,
 		peers,
 		PeerSet::Validation,
 		ValidationVersion::VStaging.into(),
-		peerset_protocol_names,
 		message,
 		metrics,
+		notification_sinks,
 	);
 }

 // Helper function to send a validation v2 message to a list of peers.
 // Messages are always sent via the main protocol, even legacy protocol messages.
 pub(crate) fn send_validation_message_v2(
-	net: &mut impl Network,
 	peers: Vec<PeerId>,
-	protocol_names: &PeerSetProtocolNames,
 	message: WireMessage<protocol_v2::ValidationProtocol>,
 	metrics: &Metrics,
+	notification_sinks: &Arc<Mutex<HashMap<(PeerSet, PeerId), Box<dyn MessageSink>>>>,
 ) {
 	send_message(
-		net,
 		peers,
 		PeerSet::Validation,
 		ValidationVersion::V2.into(),
-		protocol_names,
 		message,
 		metrics,
+		notification_sinks,
 	);
 }

 // Helper function to send a collation v1 message to a list of peers.
 // Messages are always sent via the main protocol, even legacy protocol messages.
 pub(crate) fn send_collation_message_v1(
-	net: &mut impl Network,
 	peers: Vec<PeerId>,
-	peerset_protocol_names: &PeerSetProtocolNames,
 	message: WireMessage<protocol_v1::CollationProtocol>,
 	metrics: &Metrics,
+	notification_sinks: &Arc<Mutex<HashMap<(PeerSet, PeerId), Box<dyn MessageSink>>>>,
 ) {
 	send_message(
-		net,
 		peers,
 		PeerSet::Collation,
 		CollationVersion::V1.into(),
-		peerset_protocol_names,
 		message,
 		metrics,
+		notification_sinks,
 	);
 }

 // Helper function to send a collation v2 message to a list of peers.
 // Messages are always sent via the main protocol, even legacy protocol messages.
 pub(crate) fn send_collation_message_v2(
-	net: &mut impl Network,
 	peers: Vec<PeerId>,
-	peerset_protocol_names: &PeerSetProtocolNames,
 	message: WireMessage<protocol_v2::CollationProtocol>,
 	metrics: &Metrics,
+	notification_sinks: &Arc<Mutex<HashMap<(PeerSet, PeerId), Box<dyn MessageSink>>>>,
 ) {
 	send_message(
-		net,
 		peers,
 		PeerSet::Collation,
 		CollationVersion::V2.into(),
-		peerset_protocol_names,
 		message,
 		metrics,
+		notification_sinks,
 	);
 }

@@ -151,19 +142,19 @@ pub(crate) fn send_collation_message_v2(
 /// messages that are compatible with the passed peer set, as that is currently not enforced by
 /// this function. These are messages of type `WireMessage` parameterized on the matching type.
 fn send_message<M>(
-	net: &mut impl Network,
 	mut peers: Vec<PeerId>,
 	peer_set: PeerSet,
 	version: ProtocolVersion,
-	protocol_names: &PeerSetProtocolNames,
 	message: M,
 	metrics: &super::Metrics,
+	network_notification_sinks: &Arc<Mutex<HashMap<(PeerSet, PeerId), Box<dyn MessageSink>>>>,
 ) where
 	M: Encode + Clone,
 {
 	if peers.is_empty() {
 		return
 	}
+
 	let message = {
 		let encoded = message.encode();
 		metrics.on_notification_sent(peer_set, version, encoded.len(), peers.len());
@@ -171,13 +162,13 @@ fn send_message<M>(
 		encoded
 	};

-	// optimization: generate the protocol name once.
-	let protocol_name = protocol_names.get_name(peer_set, version);
+	let notification_sinks = network_notification_sinks.lock();
+
 	gum::trace!(
 		target: LOG_TARGET,
 		?peers,
+		?peer_set,
 		?version,
-		?protocol_name,
 		?message,
 		"Sending message to peers",
 	);
@@ -185,29 +176,26 @@ fn send_message<M>(
 	// optimization: avoid cloning the message for the last peer in the
 	// list. The message payload can be quite large. If the underlying
 	// network used `Bytes` this would not be necessary.
+	//
+	// peer may have gotten disconnect by the time `send_message()` is called
+	// at which point the the sink is not available.
 	let last_peer = peers.pop();
-
-	// We always send messages on the "main" name even when a negotiated
-	// fallback is used. The libp2p implementation handles the fallback
-	// under the hood.
-	let protocol_name = protocol_names.get_main_name(peer_set);
 	peers.into_iter().for_each(|peer| {
-		net.write_notification(peer, protocol_name.clone(), message.clone());
+		if let Some(sink) = notification_sinks.get(&(peer_set, peer)) {
+			sink.send_sync_notification(message.clone());
+		}
 	});
+
 	if let Some(peer) = last_peer {
-		net.write_notification(peer, protocol_name, message);
+		if let Some(sink) = notification_sinks.get(&(peer_set, peer)) {
+			sink.send_sync_notification(message.clone());
+		}
 	}
 }

 /// An abstraction over networking for the purposes of this subsystem.
 #[async_trait]
 pub trait Network: Clone + Send + 'static {
-	/// Get a stream of all events occurring on the network. This may include events unrelated
-	/// to the Polkadot protocol - the user of this function should filter only for events related
-	/// to the [`VALIDATION_PROTOCOL_NAME`](VALIDATION_PROTOCOL_NAME)
-	/// or [`COLLATION_PROTOCOL_NAME`](COLLATION_PROTOCOL_NAME)
-	fn event_stream(&mut self) -> BoxStream<'static, NetworkEvent>;
-
 	/// Ask the network to keep a substream open with these nodes and not disconnect from them
 	/// until removed from the protocol's peer set.
 	/// Note that `out_peers` setting has no effect on this.
@@ -239,16 +227,12 @@ pub trait Network: Clone + Send + 'static {
 	/// Disconnect a given peer from the protocol specified without harming reputation.
 	fn disconnect_peer(&self, who: PeerId, protocol: ProtocolName);

-	/// Write a notification to a peer on the given protocol.
-	fn write_notification(&self, who: PeerId, protocol: ProtocolName, message: Vec<u8>);
+	/// Get peer role.
+	fn peer_role(&self, who: PeerId, handshake: Vec<u8>) -> Option<sc_network::ObservedRole>;
 }

 #[async_trait]
 impl Network for Arc<NetworkService<Block, Hash>> {
-	fn event_stream(&mut self) -> BoxStream<'static, NetworkEvent> {
-		NetworkService::event_stream(self, "polkadot-network-bridge").boxed()
-	}
-
 	async fn set_reserved_peers(
 		&mut self,
 		protocol: ProtocolName,
@@ -273,10 +257,6 @@ impl Network for Arc<NetworkService<Block, Hash>> {
 		NetworkService::disconnect_peer(&**self, who, protocol);
 	}

-	fn write_notification(&self, who: PeerId, protocol: ProtocolName, message: Vec<u8>) {
-		NetworkService::write_notification(&**self, who, protocol, message);
-	}
-
 	async fn start_request<AD: AuthorityDiscovery>(
 		&self,
 		authority_discovery: &mut AD,
@@ -348,6 +328,10 @@ impl Network for Arc<NetworkService<Block, Hash>> {
 			if_disconnected,
 		);
 	}
+
+	fn peer_role(&self, who: PeerId, handshake: Vec<u8>) -> Option<sc_network::ObservedRole> {
+		NetworkService::peer_role(self, who, handshake)
+	}
 }

 /// We assume one `peer_id` per `authority_id`.
@@ -15,7 +15,7 @@
 // along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.

 use super::*;
-use futures::{channel::oneshot, executor, stream::BoxStream};
+use futures::{channel::oneshot, executor};
 use overseer::jaeger;
 use polkadot_node_network_protocol::{self as net_protocol, OurView};
 use polkadot_node_subsystem::messages::NetworkBridgeEvent;
@@ -26,10 +26,13 @@ use parking_lot::Mutex;
 use std::{
 	collections::HashSet,
 	sync::atomic::{AtomicBool, Ordering},
-	task::Poll,
 };

-use sc_network::{Event as NetworkEvent, IfDisconnected, ProtocolName, ReputationChange};
+use sc_network::{
+	service::traits::{Direction, MessageSink, NotificationService},
+	IfDisconnected, Multiaddr, ObservedRole as SubstrateObservedRole, ProtocolName,
+	ReputationChange, Roles,
+};

 use polkadot_node_network_protocol::{
 	peer_set::PeerSetProtocolNames,
@@ -47,9 +50,8 @@ use polkadot_node_subsystem_test_helpers::{
 	mock::new_leaf, SingleItemSink, SingleItemStream, TestSubsystemContextHandle,
 };
 use polkadot_node_subsystem_util::metered;
-use polkadot_primitives::{AuthorityDiscoveryId, CandidateHash, Hash};
+use polkadot_primitives::{AuthorityDiscoveryId, Hash};

-use sc_network::Multiaddr;
 use sp_keyring::Sr25519Keyring;

 use crate::{network::Network, validator_discovery::AuthorityDiscovery};
@@ -64,10 +66,9 @@ pub enum NetworkAction {
 	WriteNotification(PeerId, PeerSet, Vec<u8>),
 }

-// The subsystem's view of the network - only supports a single call to `event_stream`.
+// The subsystem's view of the network.
 #[derive(Clone)]
 struct TestNetwork {
-	net_events: Arc<Mutex<Option<SingleItemStream<NetworkEvent>>>>,
 	action_tx: Arc<Mutex<metered::UnboundedMeteredSender<NetworkAction>>>,
 	protocol_names: Arc<PeerSetProtocolNames>,
 }
@@ -79,37 +80,42 @@ struct TestAuthorityDiscovery;
 // of `NetworkAction`s.
 struct TestNetworkHandle {
 	action_rx: metered::UnboundedMeteredReceiver<NetworkAction>,
-	net_tx: SingleItemSink<NetworkEvent>,
-	protocol_names: PeerSetProtocolNames,
+	validation_tx: SingleItemSink<NotificationEvent>,
+	collation_tx: SingleItemSink<NotificationEvent>,
 }

 fn new_test_network(
 	protocol_names: PeerSetProtocolNames,
-) -> (TestNetwork, TestNetworkHandle, TestAuthorityDiscovery) {
-	let (net_tx, net_rx) = polkadot_node_subsystem_test_helpers::single_item_sink();
+) -> (
+	TestNetwork,
+	TestNetworkHandle,
+	TestAuthorityDiscovery,
+	Box<dyn NotificationService>,
+	Box<dyn NotificationService>,
+) {
 	let (action_tx, action_rx) = metered::unbounded();
+	let (validation_tx, validation_rx) = polkadot_node_subsystem_test_helpers::single_item_sink();
+	let (collation_tx, collation_rx) = polkadot_node_subsystem_test_helpers::single_item_sink();
+	let action_tx = Arc::new(Mutex::new(action_tx));

 	(
 		TestNetwork {
-			net_events: Arc::new(Mutex::new(Some(net_rx))),
-			action_tx: Arc::new(Mutex::new(action_tx)),
+			action_tx: action_tx.clone(),
 			protocol_names: Arc::new(protocol_names.clone()),
 		},
-		TestNetworkHandle { action_rx, net_tx, protocol_names },
+		TestNetworkHandle { action_rx, validation_tx, collation_tx },
 		TestAuthorityDiscovery,
+		Box::new(TestNotificationService::new(
+			PeerSet::Validation,
+			action_tx.clone(),
+			validation_rx,
+		)),
+		Box::new(TestNotificationService::new(PeerSet::Collation, action_tx, collation_rx)),
 	)
 }

 #[async_trait]
 impl Network for TestNetwork {
-	fn event_stream(&mut self) -> BoxStream<'static, NetworkEvent> {
-		self.net_events
-			.lock()
-			.take()
-			.expect("Subsystem made more than one call to `event_stream`")
-			.boxed()
-	}
-
 	async fn set_reserved_peers(
 		&mut self,
 		_protocol: ProtocolName,
@@ -143,7 +149,8 @@ impl Network for TestNetwork {
 	}

 	fn disconnect_peer(&self, who: PeerId, protocol: ProtocolName) {
-		let (peer_set, _) = self.protocol_names.try_get_protocol(&protocol).unwrap();
+		let (peer_set, version) = self.protocol_names.try_get_protocol(&protocol).unwrap();
+		assert_eq!(version, peer_set.get_main_version());

 		self.action_tx
 			.lock()
@@ -151,13 +158,10 @@ impl Network for TestNetwork {
 			.unwrap();
 	}

-	fn write_notification(&self, who: PeerId, protocol: ProtocolName, message: Vec<u8>) {
-		let (peer_set, _) = self.protocol_names.try_get_protocol(&protocol).unwrap();
-
-		self.action_tx
-			.lock()
-			.unbounded_send(NetworkAction::WriteNotification(who, peer_set, message))
-			.unwrap();
+	fn peer_role(&self, _peer_id: PeerId, handshake: Vec<u8>) -> Option<SubstrateObservedRole> {
+		Roles::decode_all(&mut &handshake[..])
+			.ok()
+			.and_then(|role| Some(SubstrateObservedRole::from(role)))
 	}
 }

@@ -201,35 +205,85 @@ impl TestNetworkHandle {
 		peer_set: PeerSet,
 		role: ObservedRole,
 	) {
-		let protocol_version = ProtocolVersion::from(protocol_version);
-		self.send_network_event(NetworkEvent::NotificationStreamOpened {
-			remote: peer,
-			protocol: self.protocol_names.get_name(peer_set, protocol_version),
-			negotiated_fallback: None,
-			role: role.into(),
-			received_handshake: vec![],
-		})
-		.await;
+		fn observed_role_to_handshake(role: &ObservedRole) -> Vec<u8> {
+			match role {
+				&ObservedRole::Light => Roles::LIGHT.encode(),
+				&ObservedRole::Authority => Roles::AUTHORITY.encode(),
+				&ObservedRole::Full => Roles::FULL.encode(),
+			}
+		}
+
+		// because of how protocol negotiation works, if two peers support at least one common
+		// protocol, the protocol is negotiated over the main protocol (`ValidationVersion::V2`) but
+		// if either one of the peers used a fallback protocol for the negotiation (meaning they
+		// don't support the main protocol but some older version of it ), `negotiated_fallback` is
+		// set to that protocol.
+		let negotiated_fallback = match protocol_version {
+			ValidationVersion::V2 => None,
+			ValidationVersion::V1 => match peer_set {
+				PeerSet::Validation => Some(ProtocolName::from("/polkadot/validation/1")),
+				PeerSet::Collation => Some(ProtocolName::from("/polkadot/collation/1")),
+			},
+			ValidationVersion::VStaging => match peer_set {
+				PeerSet::Validation => Some(ProtocolName::from("/polkadot/validation/3")),
+				PeerSet::Collation => unreachable!(),
+			},
+		};
+
+		match peer_set {
+			PeerSet::Validation => {
+				self.validation_tx
+					.send(NotificationEvent::NotificationStreamOpened {
+						peer,
+						direction: Direction::Inbound,
+						handshake: observed_role_to_handshake(&role),
+						negotiated_fallback,
+					})
+					.await
+					.expect("subsystem concluded early");
+			},
+			PeerSet::Collation => {
+				self.collation_tx
+					.send(NotificationEvent::NotificationStreamOpened {
+						peer,
+						direction: Direction::Inbound,
+						handshake: observed_role_to_handshake(&role),
+						negotiated_fallback,
+					})
+					.await
+					.expect("subsystem concluded early");
+			},
+		}
 	}

 	async fn disconnect_peer(&mut self, peer: PeerId, peer_set: PeerSet) {
-		self.send_network_event(NetworkEvent::NotificationStreamClosed {
-			remote: peer,
-			protocol: self.protocol_names.get_main_name(peer_set),
-		})
-		.await;
+		match peer_set {
+			PeerSet::Validation => self
+				.validation_tx
+				.send(NotificationEvent::NotificationStreamClosed { peer })
+				.await
+				.expect("subsystem concluded early"),
+			PeerSet::Collation => self
+				.collation_tx
+				.send(NotificationEvent::NotificationStreamClosed { peer })
+				.await
+				.expect("subsystem concluded early"),
+		}
 	}

 	async fn peer_message(&mut self, peer: PeerId, peer_set: PeerSet, message: Vec<u8>) {
-		self.send_network_event(NetworkEvent::NotificationsReceived {
-			remote: peer,
-			messages: vec![(self.protocol_names.get_main_name(peer_set), message.into())],
-		})
-		.await;
-	}
-
-	async fn send_network_event(&mut self, event: NetworkEvent) {
-		self.net_tx.send(event).await.expect("subsystem concluded early");
+		match peer_set {
+			PeerSet::Validation => self
+				.validation_tx
+				.send(NotificationEvent::NotificationReceived { peer, notification: message })
+				.await
+				.expect("subsystem concluded early"),
+			PeerSet::Collation => self
+				.collation_tx
+				.send(NotificationEvent::NotificationReceived { peer, notification: message })
+				.await
+				.expect("subsystem concluded early"),
+		}
 	}
 }

@@ -240,6 +294,121 @@ fn assert_network_actions_contains(actions: &[NetworkAction], action: &NetworkAc
 	}
 }

+struct TestNotificationService {
+	peer_set: PeerSet,
+	action_tx: Arc<Mutex<metered::UnboundedMeteredSender<NetworkAction>>>,
+	rx: SingleItemStream<NotificationEvent>,
+}
+
+impl std::fmt::Debug for TestNotificationService {
+	fn fmt(&self, _: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		Ok(())
+	}
+}
+
+impl TestNotificationService {
+	pub fn new(
+		peer_set: PeerSet,
+		action_tx: Arc<Mutex<metered::UnboundedMeteredSender<NetworkAction>>>,
+		rx: SingleItemStream<NotificationEvent>,
+	) -> Self {
+		Self { peer_set, action_tx, rx }
+	}
+}
+
+struct TestMessageSink {
+	peer: PeerId,
+	peer_set: PeerSet,
+	action_tx: Arc<Mutex<metered::UnboundedMeteredSender<NetworkAction>>>,
+}
+
+impl TestMessageSink {
+	fn new(
+		peer: PeerId,
+		peer_set: PeerSet,
+		action_tx: Arc<Mutex<metered::UnboundedMeteredSender<NetworkAction>>>,
+	) -> TestMessageSink {
+		Self { peer, peer_set, action_tx }
+	}
+}
+
+#[async_trait::async_trait]
+impl MessageSink for TestMessageSink {
+	fn send_sync_notification(&self, notification: Vec<u8>) {
+		self.action_tx
+			.lock()
+			.unbounded_send(NetworkAction::WriteNotification(
+				self.peer,
+				self.peer_set,
+				notification,
+			))
+			.unwrap();
+	}
+
+	async fn send_async_notification(
+		&self,
+		_notification: Vec<u8>,
+	) -> Result<(), sc_network::error::Error> {
+		unimplemented!();
+	}
+}
+
+#[async_trait::async_trait]
+impl NotificationService for TestNotificationService {
+	/// Instruct `Notifications` to open a new substream for `peer`.
+	async fn open_substream(&mut self, _peer: PeerId) -> Result<(), ()> {
+		unimplemented!();
+	}
+
+	/// Instruct `Notifications` to close substream for `peer`.
+	async fn close_substream(&mut self, _peer: PeerId) -> Result<(), ()> {
+		unimplemented!();
+	}
+
+	/// Send synchronous `notification` to `peer`.
+	fn send_sync_notification(&self, _peer: &PeerId, _notification: Vec<u8>) {
+		unimplemented!();
+	}
+
+	/// Send asynchronous `notification` to `peer`, allowing sender to exercise backpressure.
+	async fn send_async_notification(
+		&self,
+		_peer: &PeerId,
+		_notification: Vec<u8>,
+	) -> Result<(), sc_network::error::Error> {
+		unimplemented!();
+	}
+
+	/// Set handshake for the notification protocol replacing the old handshake.
+	async fn set_handshake(&mut self, _handshake: Vec<u8>) -> Result<(), ()> {
+		unimplemented!();
+	}
+
+	fn try_set_handshake(&mut self, _handshake: Vec<u8>) -> Result<(), ()> {
+		unimplemented!();
+	}
+
+	/// Get next event from the `Notifications` event stream.
+	async fn next_event(&mut self) -> Option<NotificationEvent> {
+		self.rx.next().await
+	}
+
+	// Clone [`NotificationService`]
+	fn clone(&mut self) -> Result<Box<dyn NotificationService>, ()> {
+		unimplemented!();
+	}
+
+	/// Get protocol name.
+	fn protocol(&self) -> &ProtocolName {
+		unimplemented!();
+	}
+
+	/// Get notification sink of the peer.
+	fn message_sink(&self, peer: &PeerId) -> Option<Box<dyn MessageSink>> {
+		Some(Box::new(TestMessageSink::new(*peer, self.peer_set, self.action_tx.clone())))
+	}
+}
+
 #[derive(Clone)]
 struct TestSyncOracle {
 	is_major_syncing: Arc<AtomicBool>,
@@ -335,10 +504,11 @@ fn test_harness<T: Future<Output = VirtualOverseer>>(
 	let peerset_protocol_names = PeerSetProtocolNames::new(genesis_hash, fork_id);

 	let pool = sp_core::testing::TaskExecutor::new();
-	let (mut network, network_handle, discovery) = new_test_network(peerset_protocol_names.clone());
+	let (network, network_handle, discovery, validation_service, collation_service) =
+		new_test_network(peerset_protocol_names.clone());
 	let (context, virtual_overseer) =
 		polkadot_node_subsystem_test_helpers::make_subsystem_context(pool);
-	let network_stream = network.event_stream();
+	let notification_sinks = Arc::new(Mutex::new(HashMap::new()));
 	let shared = Shared::default();

 	let bridge = NetworkBridgeRx {
@@ -348,9 +518,12 @@ fn test_harness<T: Future<Output = VirtualOverseer>>(
 		sync_oracle,
 		shared: shared.clone(),
 		peerset_protocol_names,
+		validation_service,
+		collation_service,
+		notification_sinks,
 	};

-	let network_bridge = run_network_in(bridge, context, network_stream)
+	let network_bridge = run_network_in(bridge, context)
 		.map_err(|_| panic!("subsystem execution failed"))
 		.map(|_| ());

@@ -942,8 +1115,6 @@ fn relays_collation_protocol_messages() {
 			.await;
 		}

-		// peer A gets reported for sending a collation message.
-
 		let collator_protocol_message = protocol_v1::CollatorProtocolMessage::Declare(
 			Sr25519Keyring::Alice.public().into(),
 			Default::default(),
@@ -953,19 +1124,23 @@ fn relays_collation_protocol_messages() {
 		let message_v1 =
 			protocol_v1::CollationProtocol::CollatorProtocol(collator_protocol_message.clone());

-		network_handle
-			.peer_message(
-				peer_a,
-				PeerSet::Collation,
-				WireMessage::ProtocolMessage(message_v1.clone()).encode(),
-			)
-			.await;
+		// peer A gets reported for sending a collation message.
+		// NOTE: this is not possible since peer A cannot send
+		// a collation message if it has not opened a collation protocol

-		let actions = network_handle.next_network_actions(3).await;
-		assert_network_actions_contains(
-			&actions,
-			&NetworkAction::ReputationChange(peer_a, UNCONNECTED_PEERSET_COST.into()),
-		);
+		// network_handle
+		// 	.peer_message(
+		// 		peer_a,
+		// 		PeerSet::Collation,
+		// 		WireMessage::ProtocolMessage(message_v1.clone()).encode(),
+		// 	)
+		// 	.await;
+
+		// let actions = network_handle.next_network_actions(3).await;
+		// assert_network_actions_contains(
+		// 	&actions,
+		// 	&NetworkAction::ReputationChange(peer_a, UNCONNECTED_PEERSET_COST.into()),
+		// );

 		// peer B has the message relayed.

@@ -1212,7 +1387,7 @@ fn our_view_updates_decreasing_order_and_limited_to_max() {
 fn network_protocol_versioning_view_update() {
 	let (oracle, handle) = make_sync_oracle(false);
 	test_harness(Box::new(oracle), |test_harness| async move {
-		let TestHarness { mut network_handle, mut virtual_overseer, .. } = test_harness;
+		let TestHarness { mut network_handle, mut virtual_overseer, shared } = test_harness;

 		let peer_ids: Vec<_> = (0..4).map(|_| PeerId::random()).collect();
 		let peers = [
@@ -1231,12 +1406,22 @@ fn network_protocol_versioning_view_update() {

 		handle.await_mode_switch().await;

+		let mut total_validation_peers = 0;
+		let mut total_collation_peers = 0;
+
 		for &(peer_id, peer_set, version) in &peers {
 			network_handle
 				.connect_peer(peer_id, version, peer_set, ObservedRole::Full)
 				.await;
+
+			match peer_set {
+				PeerSet::Validation => total_validation_peers += 1,
+				PeerSet::Collation => total_collation_peers += 1,
+			}
 		}

+		await_peer_connections(&shared, total_validation_peers, total_collation_peers).await;
+
 		let view = view![head];
 		let actions = network_handle.next_network_actions(4).await;

@@ -1264,15 +1449,19 @@ fn network_protocol_versioning_view_update() {

 #[test]
 fn network_protocol_versioning_subsystem_msg() {
+	use polkadot_primitives::CandidateHash;
+	use std::task::Poll;
+
 	let (oracle, _handle) = make_sync_oracle(false);
 	test_harness(Box::new(oracle), |test_harness| async move {
-		let TestHarness { mut network_handle, mut virtual_overseer, .. } = test_harness;
+		let TestHarness { mut network_handle, mut virtual_overseer, shared } = test_harness;

 		let peer = PeerId::random();

 		network_handle
 			.connect_peer(peer, ValidationVersion::V2, PeerSet::Validation, ObservedRole::Full)
 			.await;
+		await_peer_connections(&shared, 1, 0).await;

 		// bridge will inform about all connected peers.
 		{
@@ -32,7 +32,7 @@ use polkadot_node_subsystem::{
 /// To be passed to [`FullNetworkConfiguration::add_notification_protocol`]().
 pub use polkadot_node_network_protocol::peer_set::{peer_sets_info, IsAuthority};
 use polkadot_node_network_protocol::request_response::Requests;
-use sc_network::ReputationChange;
+use sc_network::{MessageSink, ReputationChange};

 use crate::validator_discovery;

@@ -60,6 +60,7 @@ pub struct NetworkBridgeTx<N, AD> {
 	metrics: Metrics,
 	req_protocol_names: ReqProtocolNames,
 	peerset_protocol_names: PeerSetProtocolNames,
+	notification_sinks: Arc<Mutex<HashMap<(PeerSet, PeerId), Box<dyn MessageSink>>>>,
 }

 impl<N, AD> NetworkBridgeTx<N, AD> {
@@ -74,6 +75,7 @@ impl<N, AD> NetworkBridgeTx<N, AD> {
 		metrics: Metrics,
 		req_protocol_names: ReqProtocolNames,
 		peerset_protocol_names: PeerSetProtocolNames,
+		notification_sinks: Arc<Mutex<HashMap<(PeerSet, PeerId), Box<dyn MessageSink>>>>,
 	) -> Self {
 		Self {
 			network_service,
@@ -81,6 +83,7 @@ impl<N, AD> NetworkBridgeTx<N, AD> {
 			metrics,
 			req_protocol_names,
 			peerset_protocol_names,
+			notification_sinks,
 		}
 	}
 }
@@ -107,6 +110,7 @@ async fn handle_subsystem_messages<Context, N, AD>(
 	metrics: Metrics,
 	req_protocol_names: ReqProtocolNames,
 	peerset_protocol_names: PeerSetProtocolNames,
+	notification_sinks: Arc<Mutex<HashMap<(PeerSet, PeerId), Box<dyn MessageSink>>>>,
 ) -> Result<(), Error>
 where
 	N: Network,
@@ -130,6 +134,7 @@ where
 						&metrics,
 						&req_protocol_names,
 						&peerset_protocol_names,
+						&notification_sinks,
 					)
 					.await;
 			},
@@ -140,13 +145,14 @@ where
 #[overseer::contextbounds(NetworkBridgeTx, prefix = self::overseer)]
 async fn handle_incoming_subsystem_communication<Context, N, AD>(
 	_ctx: &mut Context,
-	mut network_service: N,
+	network_service: N,
 	validator_discovery: &mut validator_discovery::Service<N, AD>,
 	mut authority_discovery_service: AD,
 	msg: NetworkBridgeTxMessage,
 	metrics: &Metrics,
 	req_protocol_names: &ReqProtocolNames,
 	peerset_protocol_names: &PeerSetProtocolNames,
+	notification_sinks: &Arc<Mutex<HashMap<(PeerSet, PeerId), Box<dyn MessageSink>>>>,
 ) -> (N, AD)
 where
 	N: Network,
@@ -194,25 +200,22 @@ where

 			match msg {
 				Versioned::V1(msg) => send_validation_message_v1(
-					&mut network_service,
 					peers,
-					peerset_protocol_names,
 					WireMessage::ProtocolMessage(msg),
 					&metrics,
+					notification_sinks,
 				),
 				Versioned::VStaging(msg) => send_validation_message_vstaging(
-					&mut network_service,
 					peers,
-					peerset_protocol_names,
 					WireMessage::ProtocolMessage(msg),
 					&metrics,
+					notification_sinks,
 				),
 				Versioned::V2(msg) => send_validation_message_v2(
-					&mut network_service,
 					peers,
-					peerset_protocol_names,
 					WireMessage::ProtocolMessage(msg),
 					&metrics,
+					notification_sinks,
 				),
 			}
 		},
@@ -227,25 +230,22 @@ where
 			for (peers, msg) in msgs {
 				match msg {
 					Versioned::V1(msg) => send_validation_message_v1(
-						&mut network_service,
 						peers,
-						peerset_protocol_names,
 						WireMessage::ProtocolMessage(msg),
 						&metrics,
+						notification_sinks,
 					),
 					Versioned::VStaging(msg) => send_validation_message_vstaging(
-						&mut network_service,
 						peers,
-						peerset_protocol_names,
 						WireMessage::ProtocolMessage(msg),
 						&metrics,
+						notification_sinks,
 					),
 					Versioned::V2(msg) => send_validation_message_v2(
-						&mut network_service,
 						peers,
-						peerset_protocol_names,
 						WireMessage::ProtocolMessage(msg),
 						&metrics,
+						notification_sinks,
 					),
 				}
 			}
@@ -259,18 +259,16 @@ where

 			match msg {
 				Versioned::V1(msg) => send_collation_message_v1(
-					&mut network_service,
 					peers,
-					peerset_protocol_names,
 					WireMessage::ProtocolMessage(msg),
 					&metrics,
+					notification_sinks,
 				),
 				Versioned::V2(msg) | Versioned::VStaging(msg) => send_collation_message_v2(
-					&mut network_service,
 					peers,
-					peerset_protocol_names,
 					WireMessage::ProtocolMessage(msg),
 					&metrics,
+					notification_sinks,
 				),
 			}
 		},
@@ -284,18 +282,16 @@ where
 			for (peers, msg) in msgs {
 				match msg {
 					Versioned::V1(msg) => send_collation_message_v1(
-						&mut network_service,
 						peers,
-						peerset_protocol_names,
 						WireMessage::ProtocolMessage(msg),
 						&metrics,
+						notification_sinks,
 					),
 					Versioned::V2(msg) | Versioned::VStaging(msg) => send_collation_message_v2(
-						&mut network_service,
 						peers,
-						peerset_protocol_names,
 						WireMessage::ProtocolMessage(msg),
 						&metrics,
+						notification_sinks,
 					),
 				}
 			}
@@ -389,6 +385,7 @@ where
 		metrics,
 		req_protocol_names,
 		peerset_protocol_names,
+		notification_sinks,
 	} = bridge;

 	handle_subsystem_messages(
@@ -398,6 +395,7 @@ where
 		metrics,
 		req_protocol_names,
 		peerset_protocol_names,
+		notification_sinks,
 	)
 	.await?;

@@ -15,15 +15,18 @@
 // along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.

 use super::*;
-use futures::{executor, stream::BoxStream};
+use futures::executor;
 use polkadot_node_subsystem_util::TimeoutExt;

 use async_trait::async_trait;
 use parking_lot::Mutex;
 use std::collections::HashSet;

-use sc_network::{Event as NetworkEvent, IfDisconnected, ProtocolName, ReputationChange};
+use sc_network::{
+	IfDisconnected, ObservedRole as SubstrateObservedRole, ProtocolName, ReputationChange, Roles,
+};

+use parity_scale_codec::DecodeAll;
 use polkadot_node_network_protocol::{
 	peer_set::{PeerSetProtocolNames, ValidationVersion},
 	request_response::{outgoing::Requests, ReqProtocolNames},
@@ -51,10 +54,9 @@ pub enum NetworkAction {
 	WriteNotification(PeerId, PeerSet, Vec<u8>),
 }

-// The subsystem's view of the network - only supports a single call to `event_stream`.
+// The subsystem's view of the network.
 #[derive(Clone)]
 struct TestNetwork {
-	net_events: Arc<Mutex<Option<metered::MeteredReceiver<NetworkEvent>>>>,
 	action_tx: Arc<Mutex<metered::UnboundedMeteredSender<NetworkAction>>>,
 	peerset_protocol_names: Arc<PeerSetProtocolNames>,
 }
@@ -66,37 +68,78 @@ struct TestAuthorityDiscovery;
 // of `NetworkAction`s.
 struct TestNetworkHandle {
 	action_rx: metered::UnboundedMeteredReceiver<NetworkAction>,
-	net_tx: metered::MeteredSender<NetworkEvent>,
-	peerset_protocol_names: PeerSetProtocolNames,
+	_peerset_protocol_names: PeerSetProtocolNames,
+	notification_sinks: Arc<Mutex<HashMap<(PeerSet, PeerId), Box<dyn MessageSink>>>>,
+	action_tx: Arc<Mutex<metered::UnboundedMeteredSender<NetworkAction>>>,
+}
+
+struct TestMessageSink {
+	peer: PeerId,
+	peer_set: PeerSet,
+	action_tx: Arc<Mutex<metered::UnboundedMeteredSender<NetworkAction>>>,
+}
+
+impl TestMessageSink {
+	fn new(
+		peer: PeerId,
+		peer_set: PeerSet,
+		action_tx: Arc<Mutex<metered::UnboundedMeteredSender<NetworkAction>>>,
+	) -> TestMessageSink {
+		Self { peer, peer_set, action_tx }
+	}
+}
+
+#[async_trait::async_trait]
+impl MessageSink for TestMessageSink {
+	fn send_sync_notification(&self, notification: Vec<u8>) {
+		self.action_tx
+			.lock()
+			.unbounded_send(NetworkAction::WriteNotification(
+				self.peer,
+				self.peer_set,
+				notification,
+			))
+			.unwrap();
+	}
+
+	async fn send_async_notification(
+		&self,
+		_notification: Vec<u8>,
+	) -> Result<(), sc_network::error::Error> {
+		unimplemented!();
+	}
 }

 fn new_test_network(
 	peerset_protocol_names: PeerSetProtocolNames,
-) -> (TestNetwork, TestNetworkHandle, TestAuthorityDiscovery) {
-	let (net_tx, net_rx) = metered::channel(10);
+) -> (
+	TestNetwork,
+	TestNetworkHandle,
+	TestAuthorityDiscovery,
+	Arc<Mutex<HashMap<(PeerSet, PeerId), Box<dyn MessageSink>>>>,
+) {
 	let (action_tx, action_rx) = metered::unbounded();
+	let notification_sinks = Arc::new(Mutex::new(HashMap::new()));
+	let action_tx = Arc::new(Mutex::new(action_tx));

 	(
 		TestNetwork {
-			net_events: Arc::new(Mutex::new(Some(net_rx))),
-			action_tx: Arc::new(Mutex::new(action_tx)),
+			action_tx: action_tx.clone(),
 			peerset_protocol_names: Arc::new(peerset_protocol_names.clone()),
 		},
-		TestNetworkHandle { action_rx, net_tx, peerset_protocol_names },
+		TestNetworkHandle {
+			action_rx,
+			_peerset_protocol_names: peerset_protocol_names,
+			action_tx,
+			notification_sinks: notification_sinks.clone(),
+		},
 		TestAuthorityDiscovery,
+		notification_sinks,
 	)
 }

 #[async_trait]
 impl Network for TestNetwork {
-	fn event_stream(&mut self) -> BoxStream<'static, NetworkEvent> {
-		self.net_events
-			.lock()
-			.take()
-			.expect("Subsystem made more than one call to `event_stream`")
-			.boxed()
-	}
-
 	async fn set_reserved_peers(
 		&mut self,
 		_protocol: ProtocolName,
@@ -130,7 +173,8 @@ impl Network for TestNetwork {
 	}

 	fn disconnect_peer(&self, who: PeerId, protocol: ProtocolName) {
-		let (peer_set, _) = self.peerset_protocol_names.try_get_protocol(&protocol).unwrap();
+		let (peer_set, version) = self.peerset_protocol_names.try_get_protocol(&protocol).unwrap();
+		assert_eq!(version, peer_set.get_main_version());

 		self.action_tx
 			.lock()
@@ -138,13 +182,10 @@ impl Network for TestNetwork {
 			.unwrap();
 	}

-	fn write_notification(&self, who: PeerId, protocol: ProtocolName, message: Vec<u8>) {
-		let (peer_set, _) = self.peerset_protocol_names.try_get_protocol(&protocol).unwrap();
-
-		self.action_tx
-			.lock()
-			.unbounded_send(NetworkAction::WriteNotification(who, peer_set, message))
-			.unwrap();
+	fn peer_role(&self, _peer_id: PeerId, handshake: Vec<u8>) -> Option<SubstrateObservedRole> {
+		Roles::decode_all(&mut &handshake[..])
+			.ok()
+			.and_then(|role| Some(SubstrateObservedRole::from(role)))
 	}
 }

@@ -174,23 +215,14 @@ impl TestNetworkHandle {
 	async fn connect_peer(
 		&mut self,
 		peer: PeerId,
-		protocol_version: ValidationVersion,
+		_protocol_version: ValidationVersion,
 		peer_set: PeerSet,
-		role: ObservedRole,
+		_role: ObservedRole,
 	) {
-		let protocol_version = ProtocolVersion::from(protocol_version);
-		self.send_network_event(NetworkEvent::NotificationStreamOpened {
-			remote: peer,
-			protocol: self.peerset_protocol_names.get_name(peer_set, protocol_version),
-			negotiated_fallback: None,
-			role: role.into(),
-			received_handshake: vec![],
-		})
-		.await;
-	}
-
-	async fn send_network_event(&mut self, event: NetworkEvent) {
-		self.net_tx.send(event).await.expect("subsystem concluded early");
+		self.notification_sinks.lock().insert(
+			(peer_set, peer),
+			Box::new(TestMessageSink::new(peer, peer_set, self.action_tx.clone())),
+		);
 	}
 }

@@ -208,7 +240,8 @@ fn test_harness<T: Future<Output = VirtualOverseer>>(test: impl FnOnce(TestHarne
 	let peerset_protocol_names = PeerSetProtocolNames::new(genesis_hash, fork_id);

 	let pool = sp_core::testing::TaskExecutor::new();
-	let (network, network_handle, discovery) = new_test_network(peerset_protocol_names.clone());
+	let (network, network_handle, discovery, network_notification_sinks) =
+		new_test_network(peerset_protocol_names.clone());

 	let (context, virtual_overseer) =
 		polkadot_node_subsystem_test_helpers::make_subsystem_context(pool);
@@ -219,6 +252,7 @@ fn test_harness<T: Future<Output = VirtualOverseer>>(test: impl FnOnce(TestHarne
 		Metrics(None),
 		req_protocol_names,
 		peerset_protocol_names,
+		network_notification_sinks,
 	);

 	let network_bridge_out_fut = run_network_out(bridge_out, context)
@@ -364,9 +398,9 @@ fn network_protocol_versioning_send() {
 				approval_distribution_message.clone(),
 			);

-			// Note that bridge doesn't ensure neither peer's protocol version
-			// or peer set match the message.
-			let receivers = vec![peer_ids[0], peer_ids[3]];
+			// only `peer_ids[0]` opened the validation protocol v2
+			// so only they will be sent a notification
+			let receivers = vec![peer_ids[0]];
 			virtual_overseer
 				.send(FromOrchestra::Communication {
 					msg: NetworkBridgeTxMessage::SendValidationMessage(
@@ -406,7 +440,9 @@ fn network_protocol_versioning_send() {
 			let msg =
 				protocol_v2::CollationProtocol::CollatorProtocol(collator_protocol_message.clone());

-			let receivers = vec![peer_ids[1], peer_ids[2]];
+			// only `peer_ids[0]` opened the collation protocol v2
+			// so only they will be sent a notification
+			let receivers = vec![peer_ids[1]];

 			virtual_overseer
 				.send(FromOrchestra::Communication {
@@ -169,13 +169,12 @@ mod tests {
 	use crate::network::Network;

 	use async_trait::async_trait;
-	use futures::stream::BoxStream;
 	use polkadot_node_network_protocol::{
 		request_response::{outgoing::Requests, ReqProtocolNames},
 		PeerId,
 	};
 	use polkadot_primitives::Hash;
-	use sc_network::{Event as NetworkEvent, IfDisconnected, ProtocolName, ReputationChange};
+	use sc_network::{IfDisconnected, ProtocolName, ReputationChange};
 	use sp_keyring::Sr25519Keyring;
 	use std::collections::{HashMap, HashSet};

@@ -224,10 +223,6 @@ mod tests {

 	#[async_trait]
 	impl Network for TestNetwork {
-		fn event_stream(&mut self) -> BoxStream<'static, NetworkEvent> {
-			panic!()
-		}
-
 		async fn set_reserved_peers(
 			&mut self,
 			_protocol: ProtocolName,
@@ -263,7 +258,11 @@ mod tests {
 			panic!()
 		}

-		fn write_notification(&self, _: PeerId, _: ProtocolName, _: Vec<u8>) {
+		fn peer_role(
+			&self,
+			_peer_id: PeerId,
+			_handshake: Vec<u8>,
+		) -> Option<sc_network::ObservedRole> {
 			panic!()
 		}
 	}