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>

substrate/client/network/src/protocol/message.rs

@@ -29,6 +29,7 @@ use sc_network_common::message::RequestId;
 use sp_runtime::traits::{Block as BlockT, Header as HeaderT};
 
 /// Type alias for using the message type using block type parameters.
+#[allow(unused)]
 pub type Message<B> = generic::Message<
 	<B as BlockT>::Header,
 	<B as BlockT>::Hash,

substrate/client/network/src/protocol/notifications.rs

@@ -22,9 +22,13 @@
 pub use self::{
 	behaviour::{Notifications, NotificationsOut, ProtocolConfig},
 	handler::{NotificationsSink, NotifsHandlerError, Ready},
+	service::{notification_service, ProtocolHandlePair},
 };
 
+pub(crate) use self::service::ProtocolHandle;
+
 mod behaviour;
 mod handler;
+mod service;
 mod tests;
 mod upgrade;

substrate/client/network/src/protocol/notifications/handler.rs

@@ -58,9 +58,12 @@
 //! [`NotifsHandlerIn::Open`] has gotten an answer.
 
 use crate::{
-	protocol::notifications::upgrade::{
-		NotificationsIn, NotificationsInSubstream, NotificationsOut, NotificationsOutSubstream,
-		UpgradeCollec,
+	protocol::notifications::{
+		service::metrics,
+		upgrade::{
+			NotificationsIn, NotificationsInSubstream, NotificationsOut, NotificationsOutSubstream,
+			UpgradeCollec,
+		},
 	},
 	types::ProtocolName,
 };
@@ -92,7 +95,7 @@ use std::{
 
 /// Number of pending notifications in asynchronous contexts.
 /// See [`NotificationsSink::reserve_notification`] for context.
-const ASYNC_NOTIFICATIONS_BUFFER_SIZE: usize = 8;
+pub(crate) const ASYNC_NOTIFICATIONS_BUFFER_SIZE: usize = 8;
 
 /// Number of pending notifications in synchronous contexts.
 const SYNC_NOTIFICATIONS_BUFFER_SIZE: usize = 2048;
@@ -126,11 +129,19 @@ pub struct NotifsHandler {
 	events_queue: VecDeque<
 		ConnectionHandlerEvent<NotificationsOut, usize, NotifsHandlerOut, NotifsHandlerError>,
 	>,
+
+	/// Metrics.
+	metrics: Option<Arc<metrics::Metrics>>,
 }
 
 impl NotifsHandler {
 	/// Creates new [`NotifsHandler`].
-	pub fn new(peer_id: PeerId, endpoint: ConnectedPoint, protocols: Vec<ProtocolConfig>) -> Self {
+	pub fn new(
+		peer_id: PeerId,
+		endpoint: ConnectedPoint,
+		protocols: Vec<ProtocolConfig>,
+		metrics: Option<metrics::Metrics>,
+	) -> Self {
 		Self {
 			protocols: protocols
 				.into_iter()
@@ -148,6 +159,7 @@ impl NotifsHandler {
 			endpoint,
 			when_connection_open: Instant::now(),
 			events_queue: VecDeque::with_capacity(16),
+			metrics: metrics.map_or(None, |metrics| Some(Arc::new(metrics))),
 		}
 	}
 }
@@ -303,6 +315,8 @@ pub enum NotifsHandlerOut {
 	OpenDesiredByRemote {
 		/// Index of the protocol in the list of protocols passed at initialization.
 		protocol_index: usize,
+		/// Received handshake.
+		handshake: Vec<u8>,
 	},
 
 	/// The remote would like the substreams to be closed. Send a [`NotifsHandlerIn::Close`] in
@@ -331,6 +345,36 @@ pub enum NotifsHandlerOut {
 #[derive(Debug, Clone)]
 pub struct NotificationsSink {
 	inner: Arc<NotificationsSinkInner>,
+	metrics: Option<Arc<metrics::Metrics>>,
+}
+
+impl NotificationsSink {
+	/// Create new [`NotificationsSink`].
+	/// NOTE: only used for testing but must be `pub` as other crates in `client/network` use this.
+	pub fn new(
+		peer_id: PeerId,
+	) -> (Self, mpsc::Receiver<NotificationsSinkMessage>, mpsc::Receiver<NotificationsSinkMessage>)
+	{
+		let (async_tx, async_rx) = mpsc::channel(ASYNC_NOTIFICATIONS_BUFFER_SIZE);
+		let (sync_tx, sync_rx) = mpsc::channel(SYNC_NOTIFICATIONS_BUFFER_SIZE);
+		(
+			NotificationsSink {
+				inner: Arc::new(NotificationsSinkInner {
+					peer_id,
+					async_channel: FuturesMutex::new(async_tx),
+					sync_channel: Mutex::new(Some(sync_tx)),
+				}),
+				metrics: None,
+			},
+			async_rx,
+			sync_rx,
+		)
+	}
+
+	/// Get reference to metrics.
+	pub fn metrics(&self) -> &Option<Arc<metrics::Metrics>> {
+		&self.metrics
+	}
 }
 
 #[derive(Debug)]
@@ -350,8 +394,8 @@ struct NotificationsSinkInner {
 
 /// Message emitted through the [`NotificationsSink`] and processed by the background task
 /// dedicated to the peer.
-#[derive(Debug)]
-enum NotificationsSinkMessage {
+#[derive(Debug, PartialEq, Eq)]
+pub enum NotificationsSinkMessage {
 	/// Message emitted by [`NotificationsSink::reserve_notification`] and
 	/// [`NotificationsSink::write_notification_now`].
 	Notification { message: Vec<u8> },
@@ -379,8 +423,8 @@ impl NotificationsSink {
 		let mut lock = self.inner.sync_channel.lock();
 
 		if let Some(tx) = lock.as_mut() {
-			let result =
-				tx.try_send(NotificationsSinkMessage::Notification { message: message.into() });
+			let message = message.into();
+			let result = tx.try_send(NotificationsSinkMessage::Notification { message });
 
 			if result.is_err() {
 				// Cloning the `mpsc::Sender` guarantees the allocation of an extra spot in the
@@ -476,7 +520,10 @@ impl ConnectionHandler for NotifsHandler {
 				match protocol_info.state {
 					State::Closed { pending_opening } => {
 						self.events_queue.push_back(ConnectionHandlerEvent::Custom(
-							NotifsHandlerOut::OpenDesiredByRemote { protocol_index },
+							NotifsHandlerOut::OpenDesiredByRemote {
+								protocol_index,
+								handshake: in_substream_open.handshake,
+							},
 						));
 
 						protocol_info.state = State::OpenDesiredByRemote {
@@ -531,6 +578,7 @@ impl ConnectionHandler for NotifsHandler {
 								async_channel: FuturesMutex::new(async_tx),
 								sync_channel: Mutex::new(Some(sync_tx)),
 							}),
+							metrics: self.metrics.clone(),
 						};
 
 						self.protocols[protocol_index].state = State::Open {
@@ -881,6 +929,7 @@ pub mod tests {
 					async_channel: FuturesMutex::new(async_tx),
 					sync_channel: Mutex::new(Some(sync_tx)),
 				}),
+				metrics: None,
 			};
 			let (in_substream, out_substream) = MockSubstream::new();
 
@@ -1040,6 +1089,7 @@ pub mod tests {
 			},
 			peer_id: PeerId::random(),
 			events_queue: VecDeque::new(),
+			metrics: None,
 		}
 	}
 
@@ -1545,6 +1595,7 @@ pub mod tests {
 				async_channel: FuturesMutex::new(async_tx),
 				sync_channel: Mutex::new(Some(sync_tx)),
 			}),
+			metrics: None,
 		};
 
 		handler.protocols[0].state = State::Open {
@@ -1597,7 +1648,7 @@ pub mod tests {
 			assert!(std::matches!(
 				handler.poll(cx),
 				Poll::Ready(ConnectionHandlerEvent::Custom(
-					NotifsHandlerOut::OpenDesiredByRemote { protocol_index: 0 },
+					NotifsHandlerOut::OpenDesiredByRemote { protocol_index: 0, .. },
 				))
 			));
 			assert!(std::matches!(

substrate/client/network/src/protocol/notifications/service/metrics.rs

@@ -0,0 +1,130 @@
+// This file is part of Substrate.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+use crate::types::ProtocolName;
+
+use prometheus_endpoint::{
+	self as prometheus, CounterVec, HistogramOpts, HistogramVec, Opts, PrometheusError, Registry,
+	U64,
+};
+
+use std::sync::Arc;
+
+/// Notification metrics.
+#[derive(Debug, Clone)]
+pub struct Metrics {
+	// Total number of opened substreams.
+	pub notifications_streams_opened_total: CounterVec<U64>,
+
+	/// Total number of closed substreams.
+	pub notifications_streams_closed_total: CounterVec<U64>,
+
+	/// In/outbound notification sizes.
+	pub notifications_sizes: HistogramVec,
+}
+
+impl Metrics {
+	fn register(registry: &Registry) -> Result<Self, PrometheusError> {
+		Ok(Self {
+			notifications_sizes: prometheus::register(
+				HistogramVec::new(
+					HistogramOpts {
+						common_opts: Opts::new(
+							"substrate_sub_libp2p_notifications_sizes",
+							"Sizes of the notifications send to and received from all nodes",
+						),
+						buckets: prometheus::exponential_buckets(64.0, 4.0, 8)
+							.expect("parameters are always valid values; qed"),
+					},
+					&["direction", "protocol"],
+				)?,
+				registry,
+			)?,
+			notifications_streams_closed_total: prometheus::register(
+				CounterVec::new(
+					Opts::new(
+						"substrate_sub_libp2p_notifications_streams_closed_total",
+						"Total number of notification substreams that have been closed",
+					),
+					&["protocol"],
+				)?,
+				registry,
+			)?,
+			notifications_streams_opened_total: prometheus::register(
+				CounterVec::new(
+					Opts::new(
+						"substrate_sub_libp2p_notifications_streams_opened_total",
+						"Total number of notification substreams that have been opened",
+					),
+					&["protocol"],
+				)?,
+				registry,
+			)?,
+		})
+	}
+}
+
+/// Register metrics.
+pub fn register(registry: &Registry) -> Result<Metrics, PrometheusError> {
+	Metrics::register(registry)
+}
+
+/// Register opened substream to Prometheus.
+pub fn register_substream_opened(metrics: &Option<Metrics>, protocol: &ProtocolName) {
+	if let Some(metrics) = metrics {
+		metrics.notifications_streams_opened_total.with_label_values(&[&protocol]).inc();
+	}
+}
+
+/// Register closed substream to Prometheus.
+pub fn register_substream_closed(metrics: &Option<Metrics>, protocol: &ProtocolName) {
+	if let Some(metrics) = metrics {
+		metrics
+			.notifications_streams_closed_total
+			.with_label_values(&[&protocol[..]])
+			.inc();
+	}
+}
+
+/// Register sent notification to Prometheus.
+pub fn register_notification_sent(
+	metrics: &Option<Arc<Metrics>>,
+	protocol: &ProtocolName,
+	size: usize,
+) {
+	if let Some(metrics) = metrics {
+		metrics
+			.notifications_sizes
+			.with_label_values(&["out", protocol])
+			.observe(size as f64);
+	}
+}
+
+/// Register received notification to Prometheus.
+pub fn register_notification_received(
+	metrics: &Option<Metrics>,
+	protocol: &ProtocolName,
+	size: usize,
+) {
+	if let Some(metrics) = metrics {
+		metrics
+			.notifications_sizes
+			.with_label_values(&["in", protocol])
+			.observe(size as f64);
+	}
+}

substrate/client/network/src/protocol/notifications/service/mod.rs

@@ -0,0 +1,634 @@
+// This file is part of Substrate.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+//! Notification service implementation.
+
+use crate::{
+	error,
+	protocol::notifications::handler::NotificationsSink,
+	service::traits::{
+		Direction, MessageSink, NotificationEvent, NotificationService, ValidationResult,
+	},
+	types::ProtocolName,
+};
+
+use futures::{
+	stream::{FuturesUnordered, Stream},
+	StreamExt,
+};
+use libp2p::PeerId;
+use parking_lot::Mutex;
+use tokio::sync::{mpsc, oneshot};
+use tokio_stream::wrappers::ReceiverStream;
+
+use sc_utils::mpsc::{tracing_unbounded, TracingUnboundedReceiver, TracingUnboundedSender};
+
+use std::{collections::HashMap, fmt::Debug, sync::Arc};
+
+pub(crate) mod metrics;
+
+#[cfg(test)]
+mod tests;
+
+/// Logging target for the file.
+const LOG_TARGET: &str = "sub-libp2p";
+
+/// Default command queue size.
+const COMMAND_QUEUE_SIZE: usize = 64;
+
+/// Type representing subscribers of a notification protocol.
+type Subscribers = Arc<Mutex<Vec<TracingUnboundedSender<InnerNotificationEvent>>>>;
+
+/// Type represending a distributable message sink.
+/// Detached message sink must carry the protocol name for registering metrics.
+///
+/// See documentation for [`PeerContext`] for more details.
+type NotificationSink = Arc<Mutex<(NotificationsSink, ProtocolName)>>;
+
+#[async_trait::async_trait]
+impl MessageSink for NotificationSink {
+	/// Send synchronous `notification` to the peer associated with this [`MessageSink`].
+	fn send_sync_notification(&self, notification: Vec<u8>) {
+		let sink = self.lock();
+
+		metrics::register_notification_sent(&sink.0.metrics(), &sink.1, notification.len());
+		sink.0.send_sync_notification(notification);
+	}
+
+	/// Send an asynchronous `notification` to the peer associated with this [`MessageSink`],
+	/// allowing sender to exercise backpressure.
+	///
+	/// Returns an error if the peer does not exist.
+	async fn send_async_notification(&self, notification: Vec<u8>) -> Result<(), error::Error> {
+		// notification sink must be cloned because the lock cannot be held across `.await`
+		// this makes the implementation less efficient but not prohibitively so as the same
+		// method is also used by `NetworkService` when sending notifications.
+		let notification_len = notification.len();
+		let sink = self.lock().clone();
+		let permit = sink
+			.0
+			.reserve_notification()
+			.await
+			.map_err(|_| error::Error::ConnectionClosed)?;
+
+		permit.send(notification).map_err(|_| error::Error::ChannelClosed).map(|res| {
+			metrics::register_notification_sent(&sink.0.metrics(), &sink.1, notification_len);
+			res
+		})
+	}
+}
+
+/// Inner notification event to deal with `NotificationsSinks` without exposing that
+/// implementation detail to [`NotificationService`] consumers.
+#[derive(Debug)]
+enum InnerNotificationEvent {
+	/// Validate inbound substream.
+	ValidateInboundSubstream {
+		/// Peer ID.
+		peer: PeerId,
+
+		/// Received handshake.
+		handshake: Vec<u8>,
+
+		/// `oneshot::Sender` for sending validation result back to `Notifications`
+		result_tx: oneshot::Sender<ValidationResult>,
+	},
+
+	/// Notification substream open to `peer`.
+	NotificationStreamOpened {
+		/// Peer ID.
+		peer: PeerId,
+
+		/// Direction of the substream.
+		direction: Direction,
+
+		/// Received handshake.
+		handshake: Vec<u8>,
+
+		/// Negotiated fallback.
+		negotiated_fallback: Option<ProtocolName>,
+
+		/// Notification sink.
+		sink: NotificationsSink,
+	},
+
+	/// Substream was closed.
+	NotificationStreamClosed {
+		/// Peer ID.
+		peer: PeerId,
+	},
+
+	/// Notification was received from the substream.
+	NotificationReceived {
+		/// Peer ID.
+		peer: PeerId,
+
+		/// Received notification.
+		notification: Vec<u8>,
+	},
+
+	/// Notification sink has been replaced.
+	NotificationSinkReplaced {
+		/// Peer ID.
+		peer: PeerId,
+
+		/// Notification sink.
+		sink: NotificationsSink,
+	},
+}
+
+/// Notification commands.
+///
+/// Sent by the installed protocols to `Notifications` to open/close/modify substreams.
+#[derive(Debug)]
+pub enum NotificationCommand {
+	/// Instruct `Notifications` to open a substream to peer.
+	#[allow(unused)]
+	OpenSubstream(PeerId),
+
+	/// Instruct `Notifications` to close the substream to peer.
+	#[allow(unused)]
+	CloseSubstream(PeerId),
+
+	/// Set handshake for the notifications protocol.
+	SetHandshake(Vec<u8>),
+}
+
+/// Context assigned to each peer.
+///
+/// Contains `NotificationsSink` used by [`NotificationService`] to send notifications
+/// and an additional, distributable `NotificationsSink` which the protocol may acquire
+/// if it wishes to send notifications through `NotificationsSink` directly.
+///
+/// The distributable `NoticationsSink` is wrapped in an `Arc<Mutex<>>` to allow
+/// `NotificationsService` to swap the underlying sink in case it's replaced.
+#[derive(Debug, Clone)]
+struct PeerContext {
+	/// Sink for sending notificaitons.
+	sink: NotificationsSink,
+
+	/// Distributable notification sink.
+	shared_sink: NotificationSink,
+}
+
+/// Handle that is passed on to the notifications protocol.
+#[derive(Debug)]
+pub struct NotificationHandle {
+	/// Protocol name.
+	protocol: ProtocolName,
+
+	/// TX channel for sending commands to `Notifications`.
+	tx: mpsc::Sender<NotificationCommand>,
+
+	/// RX channel for receiving events from `Notifications`.
+	rx: TracingUnboundedReceiver<InnerNotificationEvent>,
+
+	/// All subscribers of `NotificationEvent`s.
+	subscribers: Subscribers,
+
+	/// Connected peers.
+	peers: HashMap<PeerId, PeerContext>,
+}
+
+impl NotificationHandle {
+	/// Create new [`NotificationHandle`].
+	fn new(
+		protocol: ProtocolName,
+		tx: mpsc::Sender<NotificationCommand>,
+		rx: TracingUnboundedReceiver<InnerNotificationEvent>,
+		subscribers: Arc<Mutex<Vec<TracingUnboundedSender<InnerNotificationEvent>>>>,
+	) -> Self {
+		Self { protocol, tx, rx, subscribers, peers: HashMap::new() }
+	}
+}
+
+#[async_trait::async_trait]
+impl NotificationService for NotificationHandle {
+	/// Instruct `Notifications` to open a new substream for `peer`.
+	async fn open_substream(&mut self, _peer: PeerId) -> Result<(), ()> {
+		todo!("support for opening substreams not implemented yet");
+	}
+
+	/// Instruct `Notifications` to close substream for `peer`.
+	async fn close_substream(&mut self, _peer: PeerId) -> Result<(), ()> {
+		todo!("support for closing substreams not implemented yet, call `NetworkService::disconnect_peer()` instead");
+	}
+
+	/// Send synchronous `notification` to `peer`.
+	fn send_sync_notification(&self, peer: &PeerId, notification: Vec<u8>) {
+		if let Some(info) = self.peers.get(&peer) {
+			metrics::register_notification_sent(
+				&info.sink.metrics(),
+				&self.protocol,
+				notification.len(),
+			);
+
+			let _ = info.sink.send_sync_notification(notification);
+		}
+	}
+
+	/// Send asynchronous `notification` to `peer`, allowing sender to exercise backpressure.
+	async fn send_async_notification(
+		&self,
+		peer: &PeerId,
+		notification: Vec<u8>,
+	) -> Result<(), error::Error> {
+		let notification_len = notification.len();
+		let sink = &self.peers.get(&peer).ok_or_else(|| error::Error::PeerDoesntExist(*peer))?.sink;
+
+		sink.reserve_notification()
+			.await
+			.map_err(|_| error::Error::ConnectionClosed)?
+			.send(notification)
+			.map_err(|_| error::Error::ChannelClosed)
+			.map(|res| {
+				metrics::register_notification_sent(
+					&sink.metrics(),
+					&self.protocol,
+					notification_len,
+				);
+				res
+			})
+	}
+
+	/// Set handshake for the notification protocol replacing the old handshake.
+	async fn set_handshake(&mut self, handshake: Vec<u8>) -> Result<(), ()> {
+		log::trace!(target: LOG_TARGET, "{}: set handshake to {handshake:?}", self.protocol);
+
+		self.tx.send(NotificationCommand::SetHandshake(handshake)).await.map_err(|_| ())
+	}
+
+	/// Non-blocking variant of `set_handshake()` that attempts to update the handshake
+	/// and returns an error if the channel is blocked.
+	///
+	/// Technically the function can return an error if the channel to `Notifications` is closed
+	/// but that doesn't happen under normal operation.
+	fn try_set_handshake(&mut self, handshake: Vec<u8>) -> Result<(), ()> {
+		self.tx.try_send(NotificationCommand::SetHandshake(handshake)).map_err(|_| ())
+	}
+
+	/// Get next event from the `Notifications` event stream.
+	async fn next_event(&mut self) -> Option<NotificationEvent> {
+		loop {
+			match self.rx.next().await? {
+				InnerNotificationEvent::ValidateInboundSubstream { peer, handshake, result_tx } =>
+					return Some(NotificationEvent::ValidateInboundSubstream {
+						peer,
+						handshake,
+						result_tx,
+					}),
+				InnerNotificationEvent::NotificationStreamOpened {
+					peer,
+					handshake,
+					negotiated_fallback,
+					direction,
+					sink,
+				} => {
+					self.peers.insert(
+						peer,
+						PeerContext {
+							sink: sink.clone(),
+							shared_sink: Arc::new(Mutex::new((sink, self.protocol.clone()))),
+						},
+					);
+					return Some(NotificationEvent::NotificationStreamOpened {
+						peer,
+						handshake,
+						direction,
+						negotiated_fallback,
+					})
+				},
+				InnerNotificationEvent::NotificationStreamClosed { peer } => {
+					self.peers.remove(&peer);
+					return Some(NotificationEvent::NotificationStreamClosed { peer })
+				},
+				InnerNotificationEvent::NotificationReceived { peer, notification } =>
+					return Some(NotificationEvent::NotificationReceived { peer, notification }),
+				InnerNotificationEvent::NotificationSinkReplaced { peer, sink } => {
+					match self.peers.get_mut(&peer) {
+						None => log::error!(
+							"{}: notification sink replaced for {peer} but peer does not exist",
+							self.protocol
+						),
+						Some(context) => {
+							context.sink = sink.clone();
+							*context.shared_sink.lock() = (sink.clone(), self.protocol.clone());
+						},
+					}
+				},
+			}
+		}
+	}
+
+	// Clone [`NotificationService`]
+	fn clone(&mut self) -> Result<Box<dyn NotificationService>, ()> {
+		let mut subscribers = self.subscribers.lock();
+		let (event_tx, event_rx) = tracing_unbounded("mpsc-notification-to-protocol", 100_000);
+		subscribers.push(event_tx);
+
+		Ok(Box::new(NotificationHandle {
+			protocol: self.protocol.clone(),
+			tx: self.tx.clone(),
+			rx: event_rx,
+			peers: self.peers.clone(),
+			subscribers: self.subscribers.clone(),
+		}))
+	}
+
+	/// Get protocol name.
+	fn protocol(&self) -> &ProtocolName {
+		&self.protocol
+	}
+
+	/// Get message sink of the peer.
+	fn message_sink(&self, peer: &PeerId) -> Option<Box<dyn MessageSink>> {
+		match self.peers.get(peer) {
+			Some(context) => Some(Box::new(context.shared_sink.clone())),
+			None => None,
+		}
+	}
+}
+
+/// Channel pair which allows `Notifications` to interact with a protocol.
+#[derive(Debug)]
+pub struct ProtocolHandlePair {
+	/// Protocol name.
+	protocol: ProtocolName,
+
+	/// Subscribers of the notification protocol events.
+	subscribers: Subscribers,
+
+	// Receiver for notification commands received from the protocol implementation.
+	rx: mpsc::Receiver<NotificationCommand>,
+}
+
+impl ProtocolHandlePair {
+	/// Create new [`ProtocolHandlePair`].
+	fn new(
+		protocol: ProtocolName,
+		subscribers: Subscribers,
+		rx: mpsc::Receiver<NotificationCommand>,
+	) -> Self {
+		Self { protocol, subscribers, rx }
+	}
+
+	/// Consume `self` and split [`ProtocolHandlePair`] into a handle which allows it to send events
+	/// to the protocol and a stream of commands received from the protocol.
+	pub(crate) fn split(
+		self,
+	) -> (ProtocolHandle, Box<dyn Stream<Item = NotificationCommand> + Send + Unpin>) {
+		(
+			ProtocolHandle::new(self.protocol, self.subscribers),
+			Box::new(ReceiverStream::new(self.rx)),
+		)
+	}
+}
+
+/// Handle that is passed on to `Notifications` and allows it to directly communicate
+/// with the protocol.
+#[derive(Debug, Clone)]
+pub(crate) struct ProtocolHandle {
+	/// Protocol name.
+	protocol: ProtocolName,
+
+	/// Subscribers of the notification protocol.
+	subscribers: Subscribers,
+
+	/// Number of connected peers.
+	num_peers: usize,
+
+	/// Delegate validation to `Peerset`.
+	delegate_to_peerset: bool,
+
+	/// Prometheus metrics.
+	metrics: Option<metrics::Metrics>,
+}
+
+pub(crate) enum ValidationCallResult {
+	WaitForValidation(oneshot::Receiver<ValidationResult>),
+	Delegated,
+}
+
+impl ProtocolHandle {
+	/// Create new [`ProtocolHandle`].
+	fn new(protocol: ProtocolName, subscribers: Subscribers) -> Self {
+		Self { protocol, subscribers, num_peers: 0usize, metrics: None, delegate_to_peerset: false }
+	}
+
+	/// Set metrics.
+	pub fn set_metrics(&mut self, metrics: Option<metrics::Metrics>) {
+		self.metrics = metrics;
+	}
+
+	/// Delegate validation to `Peerset`.
+	///
+	/// Protocols that do not do any validation themselves and only rely on `Peerset` handling
+	/// validation can disable protocol-side validation entirely by delegating all validation to
+	/// `Peerset`.
+	pub fn delegate_to_peerset(&mut self, delegate: bool) {
+		self.delegate_to_peerset = delegate;
+	}
+
+	/// Report to the protocol that a substream has been opened and it must be validated by the
+	/// protocol.
+	///
+	/// Return `oneshot::Receiver` which allows `Notifications` to poll for the validation result
+	/// from protocol.
+	pub fn report_incoming_substream(
+		&self,
+		peer: PeerId,
+		handshake: Vec<u8>,
+	) -> Result<ValidationCallResult, ()> {
+		let subscribers = self.subscribers.lock();
+
+		log::trace!(
+			target: LOG_TARGET,
+			"{}: report incoming substream for {peer}, handshake {handshake:?}",
+			self.protocol
+		);
+
+		if self.delegate_to_peerset {
+			return Ok(ValidationCallResult::Delegated)
+		}
+
+		// if there is only one subscriber, `Notifications` can wait directly on the
+		// `oneshot::channel()`'s RX half without indirection
+		if subscribers.len() == 1 {
+			let (result_tx, rx) = oneshot::channel();
+			return subscribers[0]
+				.unbounded_send(InnerNotificationEvent::ValidateInboundSubstream {
+					peer,
+					handshake,
+					result_tx,
+				})
+				.map(|_| ValidationCallResult::WaitForValidation(rx))
+				.map_err(|_| ())
+		}
+
+		// if there are multiple subscribers, create a task which waits for all of the
+		// validations to finish and returns the combined result to `Notifications`
+		let mut results: FuturesUnordered<_> = subscribers
+			.iter()
+			.filter_map(|subscriber| {
+				let (result_tx, rx) = oneshot::channel();
+
+				subscriber
+					.unbounded_send(InnerNotificationEvent::ValidateInboundSubstream {
+						peer,
+						handshake: handshake.clone(),
+						result_tx,
+					})
+					.is_ok()
+					.then_some(rx)
+			})
+			.collect();
+
+		let (tx, rx) = oneshot::channel();
+		tokio::spawn(async move {
+			while let Some(event) = results.next().await {
+				match event {
+					Err(_) | Ok(ValidationResult::Reject) =>
+						return tx.send(ValidationResult::Reject),
+					Ok(ValidationResult::Accept) => {},
+				}
+			}
+
+			return tx.send(ValidationResult::Accept)
+		});
+
+		Ok(ValidationCallResult::WaitForValidation(rx))
+	}
+
+	/// Report to the protocol that a substream has been opened and that it can now use the handle
+	/// to send notifications to the remote peer.
+	pub fn report_substream_opened(
+		&mut self,
+		peer: PeerId,
+		direction: Direction,
+		handshake: Vec<u8>,
+		negotiated_fallback: Option<ProtocolName>,
+		sink: NotificationsSink,
+	) -> Result<(), ()> {
+		metrics::register_substream_opened(&self.metrics, &self.protocol);
+
+		let mut subscribers = self.subscribers.lock();
+		log::trace!(target: LOG_TARGET, "{}: substream opened for {peer:?}", self.protocol);
+
+		subscribers.retain(|subscriber| {
+			subscriber
+				.unbounded_send(InnerNotificationEvent::NotificationStreamOpened {
+					peer,
+					direction,
+					handshake: handshake.clone(),
+					negotiated_fallback: negotiated_fallback.clone(),
+					sink: sink.clone(),
+				})
+				.is_ok()
+		});
+		self.num_peers += 1;
+
+		Ok(())
+	}
+
+	/// Substream was closed.
+	pub fn report_substream_closed(&mut self, peer: PeerId) -> Result<(), ()> {
+		metrics::register_substream_closed(&self.metrics, &self.protocol);
+
+		let mut subscribers = self.subscribers.lock();
+		log::trace!(target: LOG_TARGET, "{}: substream closed for {peer:?}", self.protocol);
+
+		subscribers.retain(|subscriber| {
+			subscriber
+				.unbounded_send(InnerNotificationEvent::NotificationStreamClosed { peer })
+				.is_ok()
+		});
+		self.num_peers -= 1;
+
+		Ok(())
+	}
+
+	/// Notification was received from the substream.
+	pub fn report_notification_received(
+		&mut self,
+		peer: PeerId,
+		notification: Vec<u8>,
+	) -> Result<(), ()> {
+		metrics::register_notification_received(&self.metrics, &self.protocol, notification.len());
+
+		let mut subscribers = self.subscribers.lock();
+		log::trace!(target: LOG_TARGET, "{}: notification received from {peer:?}", self.protocol);
+
+		subscribers.retain(|subscriber| {
+			subscriber
+				.unbounded_send(InnerNotificationEvent::NotificationReceived {
+					peer,
+					notification: notification.clone(),
+				})
+				.is_ok()
+		});
+
+		Ok(())
+	}
+
+	/// Notification sink was replaced.
+	pub fn report_notification_sink_replaced(
+		&mut self,
+		peer: PeerId,
+		sink: NotificationsSink,
+	) -> Result<(), ()> {
+		let mut subscribers = self.subscribers.lock();
+
+		log::trace!(
+			target: LOG_TARGET,
+			"{}: notification sink replaced for {peer:?}",
+			self.protocol
+		);
+
+		subscribers.retain(|subscriber| {
+			subscriber
+				.unbounded_send(InnerNotificationEvent::NotificationSinkReplaced {
+					peer,
+					sink: sink.clone(),
+				})
+				.is_ok()
+		});
+
+		Ok(())
+	}
+
+	/// Get the number of connected peers.
+	pub fn num_peers(&self) -> usize {
+		self.num_peers
+	}
+}
+
+/// Create new (protocol, notification) handle pair.
+///
+/// Handle pair allows `Notifications` and the protocol to communicate with each other directly.
+pub fn notification_service(
+	protocol: ProtocolName,
+) -> (ProtocolHandlePair, Box<dyn NotificationService>) {
+	let (cmd_tx, cmd_rx) = mpsc::channel(COMMAND_QUEUE_SIZE);
+	let (event_tx, event_rx) = tracing_unbounded("mpsc-notification-to-protocol", 100_000);
+	let subscribers = Arc::new(Mutex::new(vec![event_tx]));
+
+	(
+		ProtocolHandlePair::new(protocol.clone(), subscribers.clone(), cmd_rx),
+		Box::new(NotificationHandle::new(protocol.clone(), cmd_tx, event_rx, subscribers)),
+	)
+}

substrate/client/network/src/protocol/notifications/service/tests.rs

@@ -0,0 +1,839 @@
+// This file is part of Substrate.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+use super::*;
+use crate::protocol::notifications::handler::{
+	NotificationsSinkMessage, ASYNC_NOTIFICATIONS_BUFFER_SIZE,
+};
+
+use std::future::Future;
+
+#[tokio::test]
+async fn validate_and_accept_substream() {
+	let (proto, mut notif) = notification_service("/proto/1".into());
+	let (handle, _stream) = proto.split();
+
+	let peer_id = PeerId::random();
+	let ValidationCallResult::WaitForValidation(result_rx) =
+		handle.report_incoming_substream(peer_id, vec![1, 3, 3, 7]).unwrap()
+	else {
+		panic!("peerset not enabled");
+	};
+
+	if let Some(NotificationEvent::ValidateInboundSubstream { peer, handshake, result_tx }) =
+		notif.next_event().await
+	{
+		assert_eq!(peer_id, peer);
+		assert_eq!(handshake, vec![1, 3, 3, 7]);
+		let _ = result_tx.send(ValidationResult::Accept).unwrap();
+	} else {
+		panic!("invalid event received");
+	}
+
+	assert_eq!(result_rx.await.unwrap(), ValidationResult::Accept);
+}
+
+#[tokio::test]
+async fn substream_opened() {
+	let (proto, mut notif) = notification_service("/proto/1".into());
+	let (sink, _, _) = NotificationsSink::new(PeerId::random());
+	let (mut handle, _stream) = proto.split();
+
+	let peer_id = PeerId::random();
+	handle
+		.report_substream_opened(peer_id, Direction::Inbound, vec![1, 3, 3, 7], None, sink)
+		.unwrap();
+
+	if let Some(NotificationEvent::NotificationStreamOpened {
+		peer,
+		negotiated_fallback,
+		handshake,
+		direction,
+	}) = notif.next_event().await
+	{
+		assert_eq!(peer_id, peer);
+		assert_eq!(negotiated_fallback, None);
+		assert_eq!(handshake, vec![1, 3, 3, 7]);
+		assert_eq!(direction, Direction::Inbound);
+	} else {
+		panic!("invalid event received");
+	}
+}
+
+#[tokio::test]
+async fn send_sync_notification() {
+	let (proto, mut notif) = notification_service("/proto/1".into());
+	let (sink, _, mut sync_rx) = NotificationsSink::new(PeerId::random());
+	let (mut handle, _stream) = proto.split();
+	let peer_id = PeerId::random();
+
+	// validate inbound substream
+	let ValidationCallResult::WaitForValidation(result_rx) =
+		handle.report_incoming_substream(peer_id, vec![1, 3, 3, 7]).unwrap()
+	else {
+		panic!("peerset not enabled");
+	};
+
+	if let Some(NotificationEvent::ValidateInboundSubstream { peer, handshake, result_tx }) =
+		notif.next_event().await
+	{
+		assert_eq!(peer_id, peer);
+		assert_eq!(handshake, vec![1, 3, 3, 7]);
+		let _ = result_tx.send(ValidationResult::Accept).unwrap();
+	} else {
+		panic!("invalid event received");
+	}
+	assert_eq!(result_rx.await.unwrap(), ValidationResult::Accept);
+
+	// report that a substream has been opened
+	handle
+		.report_substream_opened(peer_id, Direction::Inbound, vec![1, 3, 3, 7], None, sink)
+		.unwrap();
+
+	if let Some(NotificationEvent::NotificationStreamOpened {
+		peer,
+		negotiated_fallback,
+		handshake,
+		direction,
+	}) = notif.next_event().await
+	{
+		assert_eq!(peer_id, peer);
+		assert_eq!(negotiated_fallback, None);
+		assert_eq!(handshake, vec![1, 3, 3, 7]);
+		assert_eq!(direction, Direction::Inbound);
+	} else {
+		panic!("invalid event received");
+	}
+
+	notif.send_sync_notification(&peer_id, vec![1, 3, 3, 8]);
+	assert_eq!(
+		sync_rx.next().await,
+		Some(NotificationsSinkMessage::Notification { message: vec![1, 3, 3, 8] })
+	);
+}
+
+#[tokio::test]
+async fn send_async_notification() {
+	let (proto, mut notif) = notification_service("/proto/1".into());
+	let (sink, mut async_rx, _) = NotificationsSink::new(PeerId::random());
+	let (mut handle, _stream) = proto.split();
+	let peer_id = PeerId::random();
+
+	// validate inbound substream
+	let ValidationCallResult::WaitForValidation(result_rx) =
+		handle.report_incoming_substream(peer_id, vec![1, 3, 3, 7]).unwrap()
+	else {
+		panic!("peerset not enabled");
+	};
+
+	if let Some(NotificationEvent::ValidateInboundSubstream { peer, handshake, result_tx }) =
+		notif.next_event().await
+	{
+		assert_eq!(peer_id, peer);
+		assert_eq!(handshake, vec![1, 3, 3, 7]);
+		let _ = result_tx.send(ValidationResult::Accept).unwrap();
+	} else {
+		panic!("invalid event received");
+	}
+	assert_eq!(result_rx.await.unwrap(), ValidationResult::Accept);
+
+	// report that a substream has been opened
+	handle
+		.report_substream_opened(peer_id, Direction::Inbound, vec![1, 3, 3, 7], None, sink)
+		.unwrap();
+
+	if let Some(NotificationEvent::NotificationStreamOpened {
+		peer,
+		negotiated_fallback,
+		handshake,
+		direction,
+	}) = notif.next_event().await
+	{
+		assert_eq!(peer_id, peer);
+		assert_eq!(negotiated_fallback, None);
+		assert_eq!(handshake, vec![1, 3, 3, 7]);
+		assert_eq!(direction, Direction::Inbound);
+	} else {
+		panic!("invalid event received");
+	}
+
+	notif.send_async_notification(&peer_id, vec![1, 3, 3, 9]).await.unwrap();
+	assert_eq!(
+		async_rx.next().await,
+		Some(NotificationsSinkMessage::Notification { message: vec![1, 3, 3, 9] })
+	);
+}
+
+#[tokio::test]
+async fn send_sync_notification_to_non_existent_peer() {
+	let (proto, notif) = notification_service("/proto/1".into());
+	let (_sink, _, _sync_rx) = NotificationsSink::new(PeerId::random());
+	let (_handle, _stream) = proto.split();
+	let peer = PeerId::random();
+
+	// as per the original implementation, the call doesn't fail
+	notif.send_sync_notification(&peer, vec![1, 3, 3, 7])
+}
+
+#[tokio::test]
+async fn send_async_notification_to_non_existent_peer() {
+	let (proto, notif) = notification_service("/proto/1".into());
+	let (_sink, _, _sync_rx) = NotificationsSink::new(PeerId::random());
+	let (_handle, _stream) = proto.split();
+	let peer = PeerId::random();
+
+	if let Err(error::Error::PeerDoesntExist(peer_id)) =
+		notif.send_async_notification(&peer, vec![1, 3, 3, 7]).await
+	{
+		assert_eq!(peer, peer_id);
+	} else {
+		panic!("invalid error received from `send_async_notification()`");
+	}
+}
+
+#[tokio::test]
+async fn receive_notification() {
+	let (proto, mut notif) = notification_service("/proto/1".into());
+	let (sink, _, _sync_rx) = NotificationsSink::new(PeerId::random());
+	let (mut handle, _stream) = proto.split();
+	let peer_id = PeerId::random();
+
+	// validate inbound substream
+	let ValidationCallResult::WaitForValidation(result_rx) =
+		handle.report_incoming_substream(peer_id, vec![1, 3, 3, 7]).unwrap()
+	else {
+		panic!("peerset not enabled");
+	};
+
+	if let Some(NotificationEvent::ValidateInboundSubstream { peer, handshake, result_tx }) =
+		notif.next_event().await
+	{
+		assert_eq!(peer_id, peer);
+		assert_eq!(handshake, vec![1, 3, 3, 7]);
+		let _ = result_tx.send(ValidationResult::Accept).unwrap();
+	} else {
+		panic!("invalid event received");
+	}
+	assert_eq!(result_rx.await.unwrap(), ValidationResult::Accept);
+
+	// report that a substream has been opened
+	handle
+		.report_substream_opened(peer_id, Direction::Inbound, vec![1, 3, 3, 7], None, sink)
+		.unwrap();
+
+	if let Some(NotificationEvent::NotificationStreamOpened {
+		peer,
+		negotiated_fallback,
+		handshake,
+		direction,
+	}) = notif.next_event().await
+	{
+		assert_eq!(peer_id, peer);
+		assert_eq!(negotiated_fallback, None);
+		assert_eq!(handshake, vec![1, 3, 3, 7]);
+		assert_eq!(direction, Direction::Inbound);
+	} else {
+		panic!("invalid event received");
+	}
+
+	// notification is received
+	handle.report_notification_received(peer_id, vec![1, 3, 3, 8]).unwrap();
+
+	if let Some(NotificationEvent::NotificationReceived { peer, notification }) =
+		notif.next_event().await
+	{
+		assert_eq!(peer_id, peer);
+		assert_eq!(notification, vec![1, 3, 3, 8]);
+	} else {
+		panic!("invalid event received");
+	}
+}
+
+#[tokio::test]
+async fn backpressure_works() {
+	let (proto, mut notif) = notification_service("/proto/1".into());
+	let (sink, mut async_rx, _) = NotificationsSink::new(PeerId::random());
+	let (mut handle, _stream) = proto.split();
+	let peer_id = PeerId::random();
+
+	// validate inbound substream
+	let ValidationCallResult::WaitForValidation(result_rx) =
+		handle.report_incoming_substream(peer_id, vec![1, 3, 3, 7]).unwrap()
+	else {
+		panic!("peerset not enabled");
+	};
+
+	if let Some(NotificationEvent::ValidateInboundSubstream { peer, handshake, result_tx }) =
+		notif.next_event().await
+	{
+		assert_eq!(peer_id, peer);
+		assert_eq!(handshake, vec![1, 3, 3, 7]);
+		let _ = result_tx.send(ValidationResult::Accept).unwrap();
+	} else {
+		panic!("invalid event received");
+	}
+	assert_eq!(result_rx.await.unwrap(), ValidationResult::Accept);
+
+	// report that a substream has been opened
+	handle
+		.report_substream_opened(peer_id, Direction::Inbound, vec![1, 3, 3, 7], None, sink)
+		.unwrap();
+
+	if let Some(NotificationEvent::NotificationStreamOpened {
+		peer,
+		negotiated_fallback,
+		handshake,
+		direction,
+	}) = notif.next_event().await
+	{
+		assert_eq!(peer_id, peer);
+		assert_eq!(negotiated_fallback, None);
+		assert_eq!(handshake, vec![1, 3, 3, 7]);
+		assert_eq!(direction, Direction::Inbound);
+	} else {
+		panic!("invalid event received");
+	}
+
+	// fill the message buffer with messages
+	for i in 0..=ASYNC_NOTIFICATIONS_BUFFER_SIZE {
+		assert!(futures::poll!(notif.send_async_notification(&peer_id, vec![1, 3, 3, i as u8]))
+			.is_ready());
+	}
+
+	// try to send one more message and verify that the call blocks
+	assert!(futures::poll!(notif.send_async_notification(&peer_id, vec![1, 3, 3, 9])).is_pending());
+
+	// release one slot from the buffer for new message
+	assert_eq!(
+		async_rx.next().await,
+		Some(NotificationsSinkMessage::Notification { message: vec![1, 3, 3, 0] })
+	);
+
+	// verify that a message can be sent
+	assert!(futures::poll!(notif.send_async_notification(&peer_id, vec![1, 3, 3, 9])).is_ready());
+}
+
+#[tokio::test]
+async fn peer_disconnects_then_sync_notification_is_sent() {
+	let (proto, mut notif) = notification_service("/proto/1".into());
+	let (sink, _, sync_rx) = NotificationsSink::new(PeerId::random());
+	let (mut handle, _stream) = proto.split();
+	let peer_id = PeerId::random();
+
+	// validate inbound substream
+	let ValidationCallResult::WaitForValidation(result_rx) =
+		handle.report_incoming_substream(peer_id, vec![1, 3, 3, 7]).unwrap()
+	else {
+		panic!("peerset not enabled");
+	};
+
+	if let Some(NotificationEvent::ValidateInboundSubstream { peer, handshake, result_tx }) =
+		notif.next_event().await
+	{
+		assert_eq!(peer_id, peer);
+		assert_eq!(handshake, vec![1, 3, 3, 7]);
+		let _ = result_tx.send(ValidationResult::Accept).unwrap();
+	} else {
+		panic!("invalid event received");
+	}
+	assert_eq!(result_rx.await.unwrap(), ValidationResult::Accept);
+
+	// report that a substream has been opened
+	handle
+		.report_substream_opened(peer_id, Direction::Inbound, vec![1, 3, 3, 7], None, sink)
+		.unwrap();
+
+	if let Some(NotificationEvent::NotificationStreamOpened {
+		peer,
+		negotiated_fallback,
+		handshake,
+		direction,
+	}) = notif.next_event().await
+	{
+		assert_eq!(peer_id, peer);
+		assert_eq!(negotiated_fallback, None);
+		assert_eq!(handshake, vec![1, 3, 3, 7]);
+		assert_eq!(direction, Direction::Inbound);
+	} else {
+		panic!("invalid event received");
+	}
+
+	// report that a substream has been closed but don't poll `notif` to receive this
+	// information
+	handle.report_substream_closed(peer_id).unwrap();
+	drop(sync_rx);
+
+	// as per documentation, error is not reported but the notification is silently dropped
+	notif.send_sync_notification(&peer_id, vec![1, 3, 3, 7]);
+}
+
+#[tokio::test]
+async fn peer_disconnects_then_async_notification_is_sent() {
+	let (proto, mut notif) = notification_service("/proto/1".into());
+	let (sink, async_rx, _) = NotificationsSink::new(PeerId::random());
+	let (mut handle, _stream) = proto.split();
+	let peer_id = PeerId::random();
+
+	// validate inbound substream
+	let ValidationCallResult::WaitForValidation(result_rx) =
+		handle.report_incoming_substream(peer_id, vec![1, 3, 3, 7]).unwrap()
+	else {
+		panic!("peerset not enabled");
+	};
+
+	if let Some(NotificationEvent::ValidateInboundSubstream { peer, handshake, result_tx }) =
+		notif.next_event().await
+	{
+		assert_eq!(peer_id, peer);
+		assert_eq!(handshake, vec![1, 3, 3, 7]);
+		let _ = result_tx.send(ValidationResult::Accept).unwrap();
+	} else {
+		panic!("invalid event received");
+	}
+	assert_eq!(result_rx.await.unwrap(), ValidationResult::Accept);
+
+	// report that a substream has been opened
+	handle
+		.report_substream_opened(peer_id, Direction::Inbound, vec![1, 3, 3, 7], None, sink)
+		.unwrap();
+
+	if let Some(NotificationEvent::NotificationStreamOpened {
+		peer,
+		negotiated_fallback,
+		handshake,
+		direction,
+	}) = notif.next_event().await
+	{
+		assert_eq!(peer_id, peer);
+		assert_eq!(negotiated_fallback, None);
+		assert_eq!(handshake, vec![1, 3, 3, 7]);
+		assert_eq!(direction, Direction::Inbound);
+	} else {
+		panic!("invalid event received");
+	}
+
+	// report that a substream has been closed but don't poll `notif` to receive this
+	// information
+	handle.report_substream_closed(peer_id).unwrap();
+	drop(async_rx);
+
+	// as per documentation, error is not reported but the notification is silently dropped
+	if let Err(error::Error::ConnectionClosed) =
+		notif.send_async_notification(&peer_id, vec![1, 3, 3, 7]).await
+	{
+	} else {
+		panic!("invalid state after calling `send_async_notificatio()` on closed connection")
+	}
+}
+
+#[tokio::test]
+async fn cloned_service_opening_substream_works() {
+	let (proto, mut notif1) = notification_service("/proto/1".into());
+	let (_sink, _async_rx, _) = NotificationsSink::new(PeerId::random());
+	let (handle, _stream) = proto.split();
+	let mut notif2 = notif1.clone().unwrap();
+	let peer_id = PeerId::random();
+
+	// validate inbound substream
+	let ValidationCallResult::WaitForValidation(mut result_rx) =
+		handle.report_incoming_substream(peer_id, vec![1, 3, 3, 7]).unwrap()
+	else {
+		panic!("peerset not enabled");
+	};
+
+	// verify that `notif1` gets the event
+	if let Some(NotificationEvent::ValidateInboundSubstream { peer, handshake, result_tx }) =
+		notif1.next_event().await
+	{
+		assert_eq!(peer_id, peer);
+		assert_eq!(handshake, vec![1, 3, 3, 7]);
+		let _ = result_tx.send(ValidationResult::Accept).unwrap();
+	} else {
+		panic!("invalid event received");
+	}
+
+	// verify that because only one listener has thus far send their result, the result is
+	// pending
+	assert!(result_rx.try_recv().is_err());
+
+	// verify that `notif2` also gets the event
+	if let Some(NotificationEvent::ValidateInboundSubstream { peer, handshake, result_tx }) =
+		notif2.next_event().await
+	{
+		assert_eq!(peer_id, peer);
+		assert_eq!(handshake, vec![1, 3, 3, 7]);
+		result_tx.send(ValidationResult::Accept).unwrap();
+	} else {
+		panic!("invalid event received");
+	}
+
+	assert_eq!(result_rx.await.unwrap(), ValidationResult::Accept);
+}
+
+#[tokio::test]
+async fn cloned_service_one_service_rejects_substream() {
+	let (proto, mut notif1) = notification_service("/proto/1".into());
+	let (_sink, _async_rx, _) = NotificationsSink::new(PeerId::random());
+	let (handle, _stream) = proto.split();
+	let mut notif2 = notif1.clone().unwrap();
+	let mut notif3 = notif2.clone().unwrap();
+	let peer_id = PeerId::random();
+
+	// validate inbound substream
+	let ValidationCallResult::WaitForValidation(mut result_rx) =
+		handle.report_incoming_substream(peer_id, vec![1, 3, 3, 7]).unwrap()
+	else {
+		panic!("peerset not enabled");
+	};
+
+	for notif in vec![&mut notif1, &mut notif2] {
+		if let Some(NotificationEvent::ValidateInboundSubstream { peer, handshake, result_tx }) =
+			notif.next_event().await
+		{
+			assert_eq!(peer_id, peer);
+			assert_eq!(handshake, vec![1, 3, 3, 7]);
+			let _ = result_tx.send(ValidationResult::Accept).unwrap();
+		} else {
+			panic!("invalid event received");
+		}
+	}
+
+	// `notif3` has not yet sent their validation result
+	assert!(result_rx.try_recv().is_err());
+
+	if let Some(NotificationEvent::ValidateInboundSubstream { peer, handshake, result_tx }) =
+		notif3.next_event().await
+	{
+		assert_eq!(peer_id, peer);
+		assert_eq!(handshake, vec![1, 3, 3, 7]);
+		let _ = result_tx.send(ValidationResult::Reject).unwrap();
+	} else {
+		panic!("invalid event received");
+	}
+	assert_eq!(result_rx.await.unwrap(), ValidationResult::Reject);
+}
+
+#[tokio::test]
+async fn cloned_service_opening_substream_sending_and_receiving_notifications_work() {
+	let (proto, mut notif1) = notification_service("/proto/1".into());
+	let (sink, _, mut sync_rx) = NotificationsSink::new(PeerId::random());
+	let (mut handle, _stream) = proto.split();
+	let mut notif2 = notif1.clone().unwrap();
+	let mut notif3 = notif1.clone().unwrap();
+	let peer_id = PeerId::random();
+
+	// validate inbound substream
+	let ValidationCallResult::WaitForValidation(result_rx) =
+		handle.report_incoming_substream(peer_id, vec![1, 3, 3, 7]).unwrap()
+	else {
+		panic!("peerset not enabled");
+	};
+
+	for notif in vec![&mut notif1, &mut notif2, &mut notif3] {
+		// accept the inbound substream for all services
+		if let Some(NotificationEvent::ValidateInboundSubstream { peer, handshake, result_tx }) =
+			notif.next_event().await
+		{
+			assert_eq!(peer_id, peer);
+			assert_eq!(handshake, vec![1, 3, 3, 7]);
+			let _ = result_tx.send(ValidationResult::Accept).unwrap();
+		} else {
+			panic!("invalid event received");
+		}
+	}
+	assert_eq!(result_rx.await.unwrap(), ValidationResult::Accept);
+
+	// report that then notification stream has been opened
+	handle
+		.report_substream_opened(peer_id, Direction::Inbound, vec![1, 3, 3, 7], None, sink)
+		.unwrap();
+
+	for notif in vec![&mut notif1, &mut notif2, &mut notif3] {
+		if let Some(NotificationEvent::NotificationStreamOpened {
+			peer,
+			negotiated_fallback,
+			handshake,
+			direction,
+		}) = notif.next_event().await
+		{
+			assert_eq!(peer_id, peer);
+			assert_eq!(negotiated_fallback, None);
+			assert_eq!(handshake, vec![1, 3, 3, 7]);
+			assert_eq!(direction, Direction::Inbound);
+		} else {
+			panic!("invalid event received");
+		}
+	}
+	// receive a notification from peer and verify all services receive it
+	handle.report_notification_received(peer_id, vec![1, 3, 3, 8]).unwrap();
+
+	for notif in vec![&mut notif1, &mut notif2, &mut notif3] {
+		if let Some(NotificationEvent::NotificationReceived { peer, notification }) =
+			notif.next_event().await
+		{
+			assert_eq!(peer_id, peer);
+			assert_eq!(notification, vec![1, 3, 3, 8]);
+		} else {
+			panic!("invalid event received");
+		}
+	}
+
+	for (i, notif) in vec![&mut notif1, &mut notif2, &mut notif3].iter().enumerate() {
+		// send notification from each service and verify peer receives it
+		notif.send_sync_notification(&peer_id, vec![1, 3, 3, i as u8]);
+		assert_eq!(
+			sync_rx.next().await,
+			Some(NotificationsSinkMessage::Notification { message: vec![1, 3, 3, i as u8] })
+		);
+	}
+
+	// close the substream for peer and verify all services receive the event
+	handle.report_substream_closed(peer_id).unwrap();
+
+	for notif in vec![&mut notif1, &mut notif2, &mut notif3] {
+		if let Some(NotificationEvent::NotificationStreamClosed { peer }) = notif.next_event().await
+		{
+			assert_eq!(peer_id, peer);
+		} else {
+			panic!("invalid event received");
+		}
+	}
+}
+
+#[tokio::test]
+async fn sending_notifications_using_notifications_sink_works() {
+	let (proto, mut notif) = notification_service("/proto/1".into());
+	let (sink, mut async_rx, mut sync_rx) = NotificationsSink::new(PeerId::random());
+	let (mut handle, _stream) = proto.split();
+	let peer_id = PeerId::random();
+
+	// validate inbound substream
+	let ValidationCallResult::WaitForValidation(result_rx) =
+		handle.report_incoming_substream(peer_id, vec![1, 3, 3, 7]).unwrap()
+	else {
+		panic!("peerset not enabled");
+	};
+
+	if let Some(NotificationEvent::ValidateInboundSubstream { peer, handshake, result_tx }) =
+		notif.next_event().await
+	{
+		assert_eq!(peer_id, peer);
+		assert_eq!(handshake, vec![1, 3, 3, 7]);
+		let _ = result_tx.send(ValidationResult::Accept).unwrap();
+	} else {
+		panic!("invalid event received");
+	}
+	assert_eq!(result_rx.await.unwrap(), ValidationResult::Accept);
+
+	// report that a substream has been opened
+	handle
+		.report_substream_opened(peer_id, Direction::Inbound, vec![1, 3, 3, 7], None, sink)
+		.unwrap();
+
+	if let Some(NotificationEvent::NotificationStreamOpened {
+		peer,
+		negotiated_fallback,
+		handshake,
+		direction,
+	}) = notif.next_event().await
+	{
+		assert_eq!(peer_id, peer);
+		assert_eq!(negotiated_fallback, None);
+		assert_eq!(handshake, vec![1, 3, 3, 7]);
+		assert_eq!(direction, Direction::Inbound);
+	} else {
+		panic!("invalid event received");
+	}
+
+	// get a copy of the notification sink and send a synchronous notification using.
+	let sink = notif.message_sink(&peer_id).unwrap();
+	sink.send_sync_notification(vec![1, 3, 3, 6]);
+
+	// send an asynchronous notification using the acquired notifications sink.
+	let _ = sink.send_async_notification(vec![1, 3, 3, 7]).await.unwrap();
+
+	assert_eq!(
+		sync_rx.next().await,
+		Some(NotificationsSinkMessage::Notification { message: vec![1, 3, 3, 6] }),
+	);
+	assert_eq!(
+		async_rx.next().await,
+		Some(NotificationsSinkMessage::Notification { message: vec![1, 3, 3, 7] }),
+	);
+
+	// send notifications using the stored notification sink as well.
+	notif.send_sync_notification(&peer_id, vec![1, 3, 3, 8]);
+	notif.send_async_notification(&peer_id, vec![1, 3, 3, 9]).await.unwrap();
+
+	assert_eq!(
+		sync_rx.next().await,
+		Some(NotificationsSinkMessage::Notification { message: vec![1, 3, 3, 8] }),
+	);
+	assert_eq!(
+		async_rx.next().await,
+		Some(NotificationsSinkMessage::Notification { message: vec![1, 3, 3, 9] }),
+	);
+}
+
+#[test]
+fn try_to_get_notifications_sink_for_non_existent_peer() {
+	let (_proto, notif) = notification_service("/proto/1".into());
+	assert!(notif.message_sink(&PeerId::random()).is_none());
+}
+
+#[tokio::test]
+async fn notification_sink_replaced() {
+	let (proto, mut notif) = notification_service("/proto/1".into());
+	let (sink, mut async_rx, mut sync_rx) = NotificationsSink::new(PeerId::random());
+	let (mut handle, _stream) = proto.split();
+	let peer_id = PeerId::random();
+
+	// validate inbound substream
+	let ValidationCallResult::WaitForValidation(result_rx) =
+		handle.report_incoming_substream(peer_id, vec![1, 3, 3, 7]).unwrap()
+	else {
+		panic!("peerset not enabled");
+	};
+
+	if let Some(NotificationEvent::ValidateInboundSubstream { peer, handshake, result_tx }) =
+		notif.next_event().await
+	{
+		assert_eq!(peer_id, peer);
+		assert_eq!(handshake, vec![1, 3, 3, 7]);
+		let _ = result_tx.send(ValidationResult::Accept).unwrap();
+	} else {
+		panic!("invalid event received");
+	}
+	assert_eq!(result_rx.await.unwrap(), ValidationResult::Accept);
+
+	// report that a substream has been opened
+	handle
+		.report_substream_opened(peer_id, Direction::Inbound, vec![1, 3, 3, 7], None, sink)
+		.unwrap();
+
+	if let Some(NotificationEvent::NotificationStreamOpened {
+		peer,
+		negotiated_fallback,
+		handshake,
+		direction,
+	}) = notif.next_event().await
+	{
+		assert_eq!(peer_id, peer);
+		assert_eq!(negotiated_fallback, None);
+		assert_eq!(handshake, vec![1, 3, 3, 7]);
+		assert_eq!(direction, Direction::Inbound);
+	} else {
+		panic!("invalid event received");
+	}
+
+	// get a copy of the notification sink and send a synchronous notification using.
+	let sink = notif.message_sink(&peer_id).unwrap();
+	sink.send_sync_notification(vec![1, 3, 3, 6]);
+
+	// send an asynchronous notification using the acquired notifications sink.
+	let _ = sink.send_async_notification(vec![1, 3, 3, 7]).await.unwrap();
+
+	assert_eq!(
+		sync_rx.next().await,
+		Some(NotificationsSinkMessage::Notification { message: vec![1, 3, 3, 6] }),
+	);
+	assert_eq!(
+		async_rx.next().await,
+		Some(NotificationsSinkMessage::Notification { message: vec![1, 3, 3, 7] }),
+	);
+
+	// send notifications using the stored notification sink as well.
+	notif.send_sync_notification(&peer_id, vec![1, 3, 3, 8]);
+	notif.send_async_notification(&peer_id, vec![1, 3, 3, 9]).await.unwrap();
+
+	assert_eq!(
+		sync_rx.next().await,
+		Some(NotificationsSinkMessage::Notification { message: vec![1, 3, 3, 8] }),
+	);
+	assert_eq!(
+		async_rx.next().await,
+		Some(NotificationsSinkMessage::Notification { message: vec![1, 3, 3, 9] }),
+	);
+
+	// the initial connection was closed and `Notifications` switched to secondary connection
+	// and emitted `CustomProtocolReplaced` which informs the local `NotificationService` that
+	// the notification sink was replaced.
+	let (new_sink, mut new_async_rx, mut new_sync_rx) = NotificationsSink::new(PeerId::random());
+	handle.report_notification_sink_replaced(peer_id, new_sink).unwrap();
+
+	// drop the old sinks and poll `notif` once to register the sink replacement
+	drop(sync_rx);
+	drop(async_rx);
+
+	futures::future::poll_fn(|cx| {
+		let _ = std::pin::Pin::new(&mut notif.next_event()).poll(cx);
+		std::task::Poll::Ready(())
+	})
+	.await;
+
+	// verify that using the `NotificationService` API automatically results in using the correct
+	// sink
+	notif.send_sync_notification(&peer_id, vec![1, 3, 3, 8]);
+	notif.send_async_notification(&peer_id, vec![1, 3, 3, 9]).await.unwrap();
+
+	assert_eq!(
+		new_sync_rx.next().await,
+		Some(NotificationsSinkMessage::Notification { message: vec![1, 3, 3, 8] }),
+	);
+	assert_eq!(
+		new_async_rx.next().await,
+		Some(NotificationsSinkMessage::Notification { message: vec![1, 3, 3, 9] }),
+	);
+
+	// now send two notifications using the acquired message sink and verify that
+	// it's also updated
+	sink.send_sync_notification(vec![1, 3, 3, 6]);
+
+	// send an asynchronous notification using the acquired notifications sink.
+	let _ = sink.send_async_notification(vec![1, 3, 3, 7]).await.unwrap();
+
+	assert_eq!(
+		new_sync_rx.next().await,
+		Some(NotificationsSinkMessage::Notification { message: vec![1, 3, 3, 6] }),
+	);
+	assert_eq!(
+		new_async_rx.next().await,
+		Some(NotificationsSinkMessage::Notification { message: vec![1, 3, 3, 7] }),
+	);
+}
+
+#[tokio::test]
+async fn set_handshake() {
+	let (proto, mut notif) = notification_service("/proto/1".into());
+	let (_handle, mut stream) = proto.split();
+
+	assert!(notif.try_set_handshake(vec![1, 3, 3, 7]).is_ok());
+
+	match stream.next().await {
+		Some(NotificationCommand::SetHandshake(handshake)) => {
+			assert_eq!(handshake, vec![1, 3, 3, 7]);
+		},
+		_ => panic!("invalid event received"),
+	}
+
+	for _ in 0..COMMAND_QUEUE_SIZE {
+		assert!(notif.try_set_handshake(vec![1, 3, 3, 7]).is_ok());
+	}
+
+	assert!(notif.try_set_handshake(vec![1, 3, 3, 7]).is_err());
+}

substrate/client/network/src/protocol/notifications/tests.rs

@@ -22,6 +22,7 @@ use crate::{
 	peer_store::PeerStore,
 	protocol::notifications::{Notifications, NotificationsOut, ProtocolConfig},
 	protocol_controller::{ProtoSetConfig, ProtocolController, SetId},
+	service::traits::{NotificationEvent, ValidationResult},
 };
 
 use futures::{future::BoxFuture, prelude::*};
@@ -70,6 +71,8 @@ fn build_nodes() -> (Swarm<CustomProtoWithAddr>, Swarm<CustomProtoWithAddr>) {
 			.timeout(Duration::from_secs(20))
 			.boxed();
 
+		let (protocol_handle_pair, mut notif_service) =
+			crate::protocol::notifications::service::notification_service("/foo".into());
 		let peer_store = PeerStore::new(if index == 0 {
 			keypairs.iter().skip(1).map(|keypair| keypair.public().to_peer_id()).collect()
 		} else {
@@ -91,16 +94,22 @@ fn build_nodes() -> (Swarm<CustomProtoWithAddr>, Swarm<CustomProtoWithAddr>) {
 			Box::new(peer_store.handle()),
 		);
 
+		let (notif_handle, command_stream) = protocol_handle_pair.split();
 		let behaviour = CustomProtoWithAddr {
 			inner: Notifications::new(
 				vec![controller_handle],
 				from_controller,
-				iter::once(ProtocolConfig {
-					name: "/foo".into(),
-					fallback_names: Vec::new(),
-					handshake: Vec::new(),
-					max_notification_size: 1024 * 1024,
-				}),
+				&None,
+				iter::once((
+					ProtocolConfig {
+						name: "/foo".into(),
+						fallback_names: Vec::new(),
+						handshake: Vec::new(),
+						max_notification_size: 1024 * 1024,
+					},
+					notif_handle,
+					command_stream,
+				)),
 			),
 			peer_store_future: peer_store.run().boxed(),
 			protocol_controller_future: controller.run().boxed(),
@@ -118,6 +127,16 @@ fn build_nodes() -> (Swarm<CustomProtoWithAddr>, Swarm<CustomProtoWithAddr>) {
 		};
 
 		let runtime = tokio::runtime::Runtime::new().unwrap();
+		runtime.spawn(async move {
+			loop {
+				if let NotificationEvent::ValidateInboundSubstream { result_tx, .. } =
+					notif_service.next_event().await.unwrap()
+				{
+					result_tx.send(ValidationResult::Accept).unwrap();
+				}
+			}
+		});
+
 		let mut swarm = SwarmBuilder::with_executor(
 			transport,
 			behaviour,