Add a back-pressure-friendly alternative to NetworkService::write_notifications 🎉

client/network/src/behaviour.rs

@@ -17,10 +17,11 @@
 use crate::{
 	config::{ProtocolId, Role}, block_requests, light_client_handler, finality_requests,
 	peer_info, discovery::{DiscoveryBehaviour, DiscoveryConfig, DiscoveryOut},
-	protocol::{message::{self, Roles}, CustomMessageOutcome, Protocol},
-	Event, ObservedRole, DhtEvent, ExHashT,
+	protocol::{message::{self, Roles}, CustomMessageOutcome, NotificationsSink, Protocol},
+	ObservedRole, DhtEvent, ExHashT,
 };
 
+use bytes::Bytes;
 use codec::Encode as _;
 use libp2p::NetworkBehaviour;
 use libp2p::core::{Multiaddr, PeerId, PublicKey};
@@ -98,11 +99,53 @@ pub enum BehaviourOut<B: BlockT> {
 		request_duration: Duration,
 	},
 
-	/// Any event represented by the [`Event`] enum.
+	/// Opened a substream with the given node with the given notifications protocol.
 	///
-	/// > **Note**: The [`Event`] enum contains the events that are available through the public
-	/// > API of the library.
-	Event(Event),
+	/// The protocol is always one of the notification protocols that have been registered.
+	NotificationStreamOpened {
+		/// Node we opened the substream with.
+		remote: PeerId,
+		/// The concerned protocol. Each protocol uses a different substream.
+		engine_id: ConsensusEngineId,
+		/// Object that permits sending notifications to the peer.
+		notifications_sink: NotificationsSink,
+		/// Role of the remote.
+		role: ObservedRole,
+	},
+
+	/// The [`NotificationsSink`] object used to send notifications with the given peer must be
+	/// replaced with a new one.
+	///
+	/// This event is typically emitted when a transport-level connection is closed and we fall
+	/// back to a secondary connection.
+	NotificationStreamReplaced {
+		/// Id of the peer we are connected to.
+		remote: PeerId,
+		/// The concerned protocol. Each protocol uses a different substream.
+		engine_id: ConsensusEngineId,
+		/// Replacement for the previous [`NotificationsSink`].
+		notifications_sink: NotificationsSink,
+	},
+
+	/// Closed a substream with the given node. Always matches a corresponding previous
+	/// `NotificationStreamOpened` message.
+	NotificationStreamClosed {
+		/// Node we closed the substream with.
+		remote: PeerId,
+		/// The concerned protocol. Each protocol uses a different substream.
+		engine_id: ConsensusEngineId,
+	},
+
+	/// Received one or more messages from the given node using the given protocol.
+	NotificationsReceived {
+		/// Node we received the message from.
+		remote: PeerId,
+		/// Concerned protocol and associated message.
+		messages: Vec<(ConsensusEngineId, Bytes)>,
+	},
+
+	/// Event generated by a DHT.
+	Dht(DhtEvent),
 }
 
 impl<B: BlockT, H: ExHashT> Behaviour<B, H> {
@@ -165,8 +208,6 @@ impl<B: BlockT, H: ExHashT> Behaviour<B, H> {
 
 	/// Registers a new notifications protocol.
 	///
-	/// After that, you can call `write_notifications`.
-	///
 	/// Please call `event_stream` before registering a protocol, otherwise you may miss events
 	/// about the protocol that you have registered.
 	///
@@ -182,14 +223,14 @@ impl<B: BlockT, H: ExHashT> Behaviour<B, H> {
 		let handshake_message = Roles::from(&self.role).encode();
 
 		let list = self.substrate.register_notifications_protocol(engine_id, protocol_name, handshake_message);
-		for (remote, roles) in list {
+		for (remote, roles, notifications_sink) in list {
 			let role = reported_roles_to_observed_role(&self.role, remote, roles);
-			let ev = Event::NotificationStreamOpened {
+			self.events.push_back(BehaviourOut::NotificationStreamOpened {
 				remote: remote.clone(),
 				engine_id,
 				role,
-			};
-			self.events.push_back(BehaviourOut::Event(ev));
+				notifications_sink: notifications_sink.clone(),
+			});
 		}
 	}
 
@@ -278,26 +319,34 @@ Behaviour<B, H> {
 			CustomMessageOutcome::FinalityProofRequest { target, block_hash, request } => {
 				self.finality_proof_requests.send_request(&target, block_hash, request);
 			},
-			CustomMessageOutcome::NotificationStreamOpened { remote, protocols, roles } => {
+			CustomMessageOutcome::NotificationStreamOpened { remote, protocols, roles, notifications_sink } => {
 				let role = reported_roles_to_observed_role(&self.role, &remote, roles);
 				for engine_id in protocols {
-					self.events.push_back(BehaviourOut::Event(Event::NotificationStreamOpened {
+					self.events.push_back(BehaviourOut::NotificationStreamOpened {
 						remote: remote.clone(),
 						engine_id,
 						role: role.clone(),
-					}));
+						notifications_sink: notifications_sink.clone(),
+					});
 				}
 			},
+			CustomMessageOutcome::NotificationStreamReplaced { remote, protocols, notifications_sink } =>
+				for engine_id in protocols {
+					self.events.push_back(BehaviourOut::NotificationStreamReplaced {
+						remote: remote.clone(),
+						engine_id,
+						notifications_sink: notifications_sink.clone(),
+					});
+				},
 			CustomMessageOutcome::NotificationStreamClosed { remote, protocols } =>
 				for engine_id in protocols {
-					self.events.push_back(BehaviourOut::Event(Event::NotificationStreamClosed {
+					self.events.push_back(BehaviourOut::NotificationStreamClosed {
 						remote: remote.clone(),
 						engine_id,
-					}));
+					});
 				},
 			CustomMessageOutcome::NotificationsReceived { remote, messages } => {
-				let ev = Event::NotificationsReceived { remote, messages };
-				self.events.push_back(BehaviourOut::Event(ev));
+				self.events.push_back(BehaviourOut::NotificationsReceived { remote, messages });
 			},
 			CustomMessageOutcome::PeerNewBest(peer_id, number) => {
 				self.light_client_handler.update_best_block(&peer_id, number);
@@ -393,16 +442,16 @@ impl<B: BlockT, H: ExHashT> NetworkBehaviourEventProcess<DiscoveryOut>
 				self.substrate.add_discovered_nodes(iter::once(peer_id));
 			}
 			DiscoveryOut::ValueFound(results) => {
-				self.events.push_back(BehaviourOut::Event(Event::Dht(DhtEvent::ValueFound(results))));
+				self.events.push_back(BehaviourOut::Dht(DhtEvent::ValueFound(results)));
 			}
 			DiscoveryOut::ValueNotFound(key) => {
-				self.events.push_back(BehaviourOut::Event(Event::Dht(DhtEvent::ValueNotFound(key))));
+				self.events.push_back(BehaviourOut::Dht(DhtEvent::ValueNotFound(key)));
 			}
 			DiscoveryOut::ValuePut(key) => {
-				self.events.push_back(BehaviourOut::Event(Event::Dht(DhtEvent::ValuePut(key))));
+				self.events.push_back(BehaviourOut::Dht(DhtEvent::ValuePut(key)));
 			}
 			DiscoveryOut::ValuePutFailed(key) => {
-				self.events.push_back(BehaviourOut::Event(Event::Dht(DhtEvent::ValuePutFailed(key))));
+				self.events.push_back(BehaviourOut::Dht(DhtEvent::ValuePutFailed(key)));
 			}
 			DiscoveryOut::RandomKademliaStarted(protocols) => {
 				for protocol in protocols {

client/network/src/protocol.rs

@@ -47,8 +47,8 @@ use sp_runtime::traits::{
 };
 use sp_arithmetic::traits::SaturatedConversion;
 use message::{BlockAnnounce, Message};
-use message::generic::{Message as GenericMessage, ConsensusMessage, Roles};
-use prometheus_endpoint::{Registry, Gauge, Counter, GaugeVec, HistogramVec, PrometheusError, Opts, register, U64};
+use message::generic::{Message as GenericMessage, Roles};
+use prometheus_endpoint::{Registry, Gauge, Counter, GaugeVec, PrometheusError, Opts, register, U64};
 use sync::{ChainSync, SyncState};
 use std::borrow::Cow;
 use std::collections::{BTreeMap, HashMap, HashSet, VecDeque, hash_map::Entry};
@@ -67,7 +67,7 @@ pub mod message;
 pub mod event;
 pub mod sync;
 
-pub use generic_proto::LegacyConnectionKillError;
+pub use generic_proto::{NotificationsSink, Ready, NotifsHandlerError, LegacyConnectionKillError};
 
 const REQUEST_TIMEOUT_SEC: u64 = 40;
 /// Interval at which we perform time based maintenance
@@ -388,7 +388,6 @@ impl<B: BlockT, H: ExHashT> Protocol<B, H> {
 		block_announce_validator: Box<dyn BlockAnnounceValidator<B> + Send>,
 		metrics_registry: Option<&Registry>,
 		boot_node_ids: Arc<HashSet<PeerId>>,
-		queue_size_report: Option<HistogramVec>,
 	) -> error::Result<(Protocol<B, H>, sc_peerset::PeersetHandle)> {
 		let info = chain.info();
 		let sync = ChainSync::new(
@@ -417,7 +416,6 @@ impl<B: BlockT, H: ExHashT> Protocol<B, H> {
 			versions,
 			build_status_message(&config, &chain),
 			peerset,
-			queue_size_report,
 		);
 
 		let mut legacy_equiv_by_name = HashMap::new();
@@ -948,7 +946,12 @@ impl<B: BlockT, H: ExHashT> Protocol<B, H> {
 	}
 
 	/// Called on receipt of a status message via the legacy protocol on the first connection between two peers.
-	pub fn on_peer_connected(&mut self, who: PeerId, status: message::Status<B>) -> CustomMessageOutcome<B> {
+	pub fn on_peer_connected(
+		&mut self,
+		who: PeerId,
+		status: message::Status<B>,
+		notifications_sink: NotificationsSink,
+	) -> CustomMessageOutcome<B> {
 		trace!(target: "sync", "New peer {} {:?}", who, status);
 		let _protocol_version = {
 			if self.context_data.peers.contains_key(&who) {
@@ -1060,32 +1063,7 @@ impl<B: BlockT, H: ExHashT> Protocol<B, H> {
 			remote: who,
 			protocols: self.protocol_name_by_engine.keys().cloned().collect(),
 			roles: info.roles,
-		}
-	}
-
-	/// Send a notification to the given peer we're connected to.
-	///
-	/// Doesn't do anything if we don't have a notifications substream for that protocol with that
-	/// peer.
-	pub fn write_notification(
-		&mut self,
-		target: PeerId,
-		engine_id: ConsensusEngineId,
-		message: impl Into<Vec<u8>>,
-	) {
-		if let Some(protocol_name) = self.protocol_name_by_engine.get(&engine_id) {
-			let message = message.into();
-			let fallback = GenericMessage::<(), (), (), ()>::Consensus(ConsensusMessage {
-				engine_id,
-				data: message.clone(),
-			}).encode();
-			self.behaviour.write_notification(&target, protocol_name.clone(), message, fallback);
-		} else {
-			error!(
-				target: "sub-libp2p",
-				"Sending a notification with a protocol that wasn't registered: {:?}",
-				engine_id
-			);
+			notifications_sink,
 		}
 	}
 
@@ -1099,7 +1077,7 @@ impl<B: BlockT, H: ExHashT> Protocol<B, H> {
 		engine_id: ConsensusEngineId,
 		protocol_name: impl Into<Cow<'static, [u8]>>,
 		handshake_message: Vec<u8>,
-	) -> impl ExactSizeIterator<Item = (&'a PeerId, Roles)> + 'a {
+	) -> impl Iterator<Item = (&'a PeerId, Roles, &'a NotificationsSink)> + 'a {
 		let protocol_name = protocol_name.into();
 		if self.protocol_name_by_engine.insert(engine_id, protocol_name.clone()).is_some() {
 			error!(target: "sub-libp2p", "Notifications protocol already registered: {:?}", protocol_name);
@@ -1108,8 +1086,15 @@ impl<B: BlockT, H: ExHashT> Protocol<B, H> {
 			self.legacy_equiv_by_name.insert(protocol_name, Fallback::Consensus(engine_id));
 		}
 
-		self.context_data.peers.iter()
-			.map(|(peer_id, peer)| (peer_id, peer.info.roles))
+		let behaviour = &self.behaviour;
+		self.context_data.peers.iter().filter_map(move |(peer_id, peer)| {
+			if let Some(notifications_sink) = behaviour.notifications_sink(peer_id) {
+				Some((peer_id, peer.info.roles, notifications_sink))
+			} else {
+				log::error!("State mismatch: no notifications sink for opened peer {:?}", peer_id);
+				None
+			}
+		})
 	}
 
 	/// Called when peer sends us new transactions
@@ -1863,7 +1848,18 @@ pub enum CustomMessageOutcome<B: BlockT> {
 	JustificationImport(Origin, B::Hash, NumberFor<B>, Justification),
 	FinalityProofImport(Origin, B::Hash, NumberFor<B>, Vec<u8>),
 	/// Notification protocols have been opened with a remote.
-	NotificationStreamOpened { remote: PeerId, protocols: Vec<ConsensusEngineId>, roles: Roles },
+	NotificationStreamOpened {
+		remote: PeerId,
+		protocols: Vec<ConsensusEngineId>,
+		roles: Roles,
+		notifications_sink: NotificationsSink
+	},
+	/// The [`NotificationsSink`] of some notification protocols need an update.
+	NotificationStreamReplaced {
+		remote: PeerId,
+		protocols: Vec<ConsensusEngineId>,
+		notifications_sink: NotificationsSink,
+	},
 	/// Notification protocols have been closed with a remote.
 	NotificationStreamClosed { remote: PeerId, protocols: Vec<ConsensusEngineId> },
 	/// Messages have been received on one or more notifications protocols.
@@ -2028,9 +2024,10 @@ impl<B: BlockT, H: ExHashT> NetworkBehaviour for Protocol<B, H> {
 		};
 
 		let outcome = match event {
-			GenericProtoOut::CustomProtocolOpen { peer_id, received_handshake, .. } => {
+			GenericProtoOut::CustomProtocolOpen { peer_id, received_handshake, notifications_sink, .. } => {
 				match <Message<B> as Decode>::decode(&mut &received_handshake[..]) {
-					Ok(GenericMessage::Status(handshake)) => self.on_peer_connected(peer_id, handshake),
+					Ok(GenericMessage::Status(handshake)) =>
+						self.on_peer_connected(peer_id, handshake, notifications_sink),
 					Ok(msg) => {
 						debug!(
 							target: "sync",
@@ -2054,6 +2051,13 @@ impl<B: BlockT, H: ExHashT> NetworkBehaviour for Protocol<B, H> {
 					}
 				}
 			}
+			GenericProtoOut::CustomProtocolReplaced { peer_id, notifications_sink, .. } => {
+				CustomMessageOutcome::NotificationStreamReplaced {
+					remote: peer_id,
+					protocols: self.protocol_name_by_engine.keys().cloned().collect(),
+					notifications_sink,
+				}
+			},
 			GenericProtoOut::CustomProtocolClosed { peer_id, .. } => {
 				self.on_peer_disconnected(peer_id)
 			},

client/network/src/protocol/generic_proto.rs

@@ -21,7 +21,7 @@
 //! network, then performs the Substrate protocol handling on top.
 
 pub use self::behaviour::{GenericProto, GenericProtoOut};
-pub use self::handler::LegacyConnectionKillError;
+pub use self::handler::{NotifsHandlerError, NotificationsSink, Ready, LegacyConnectionKillError};
 
 mod behaviour;
 mod handler;