Add a back-pressure-friendly alternative to NetworkService::write_notifications 🎉 (#6692)

* Add NetworkService::send_notifications

* Doc

* Doc

* API adjustment

* Address concerns

* Make it compile

* Start implementation

* Progress in the implementation

* Change implementation strategy again

* More work before weekend

* Finish changes

* Minor doc fix

* Revert some minor changes

* Apply suggestions from code review

* GroupError -> NotifsHandlerError

* Apply suggestions from code review

Co-authored-by: Roman Borschel <romanb@users.noreply.github.com>

* state_transition_waker -> close_waker

* Apply suggestions from code review

Co-authored-by: Roman Borschel <romanb@users.noreply.github.com>

* Finish renames in service.rs

* More renames

* More review suggestsions applied

* More review addressing

* Final change

* 512 -> 2048

Co-authored-by: Roman Borschel <romanb@users.noreply.github.com>

substrate/client/network/src/protocol/generic_proto/handler/group.rs

@@ -63,11 +63,21 @@ use libp2p::swarm::{
 	SubstreamProtocol,
 	NegotiatedSubstream,
 };
+use futures::{
+	channel::mpsc,
+	lock::{Mutex as FuturesMutex, MutexGuard as FuturesMutexGuard},
+	prelude::*
+};
 use log::{debug, error};
-use parking_lot::RwLock;
-use prometheus_endpoint::HistogramVec;
+use parking_lot::{Mutex, RwLock};
 use std::{borrow::Cow, error, io, str, sync::Arc, task::{Context, Poll}};
 
+/// Number of pending notifications in asynchronous contexts.
+/// See [`NotificationsSink::reserve_notification`] for context.
+const ASYNC_NOTIFICATIONS_BUFFER_SIZE: usize = 8;
+/// Number of pending notifications in synchronous contexts.
+const SYNC_NOTIFICATIONS_BUFFER_SIZE: usize = 2048;
+
 /// Implements the `IntoProtocolsHandler` trait of libp2p.
 ///
 /// Every time a connection with a remote starts, an instance of this struct is created and
@@ -107,6 +117,18 @@ pub struct NotifsHandler {
 	/// we push the corresponding index here and process them when the handler
 	/// gets enabled/disabled.
 	pending_in: Vec<usize>,
+
+	/// If `Some`, contains the two `Receiver`s connected to the [`NotificationsSink`] that has
+	/// been sent out. The notifications to send out can be pulled from this receivers.
+	/// We use two different channels in order to have two different channel sizes, but from the
+	/// receiving point of view, the two channels are the same.
+	/// The receivers are fused in case the user drops the [`NotificationsSink`] entirely.
+	notifications_sink_rx: Option<
+		stream::Select<
+			stream::Fuse<mpsc::Receiver<NotificationsSinkMessage>>,
+			stream::Fuse<mpsc::Receiver<NotificationsSinkMessage>>
+		>
+	>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -140,6 +162,7 @@ impl IntoProtocolsHandler for NotifsHandlerProto {
 			legacy: self.legacy.into_handler(remote_peer_id, connected_point),
 			enabled: EnabledState::Initial,
 			pending_in: Vec::new(),
+			notifications_sink_rx: None,
 		}
 	}
 }
@@ -152,32 +175,6 @@ pub enum NotifsHandlerIn {
 
 	/// The node should stop using custom protocols.
 	Disable,
-
-	/// Sends a message through the custom protocol substream.
-	///
-	/// > **Note**: This must **not** be a `ConsensusMessage`, `Transactions`, or
-	/// > 			`BlockAnnounce` message.
-	SendLegacy {
-		/// The message to send.
-		message: Vec<u8>,
-	},
-
-	/// Sends a notifications message.
-	SendNotification {
-		/// Name of the protocol for the message.
-		///
-		/// Must match one of the registered protocols. For backwards-compatibility reasons, if
-		/// the remote doesn't support this protocol, we use the legacy substream.
-		protocol_name: Cow<'static, [u8]>,
-
-		/// Message to send on the legacy substream if the protocol isn't available.
-		///
-		/// This corresponds to what you would have sent with `SendLegacy`.
-		encoded_fallback_message: Vec<u8>,
-
-		/// The message to send.
-		message: Vec<u8>,
-	},
 }
 
 /// Event that can be emitted by a `NotifsHandler`.
@@ -190,6 +187,8 @@ pub enum NotifsHandlerOut {
 		/// Handshake that was sent to us.
 		/// This is normally a "Status" message, but this out of the concern of this code.
 		received_handshake: Vec<u8>,
+		/// How notifications can be sent to this node.
+		notifications_sink: NotificationsSink,
 	},
 
 	/// The connection is closed for custom protocols.
@@ -227,19 +226,160 @@ pub enum NotifsHandlerOut {
 	},
 }
 
+/// Sink connected directly to the node background task. Allows sending notifications to the peer.
+///
+/// Can be cloned in order to obtain multiple references to the same peer.
+#[derive(Debug, Clone)]
+pub struct NotificationsSink {
+	inner: Arc<NotificationsSinkInner>,
+}
+
+#[derive(Debug)]
+struct NotificationsSinkInner {
+	/// Sender to use in asynchronous contexts. Uses an asynchronous mutex.
+	async_channel: FuturesMutex<mpsc::Sender<NotificationsSinkMessage>>,
+	/// Sender to use in synchronous contexts. Uses a synchronous mutex.
+	/// This channel has a large capacity and is meant to be used in contexts where
+	/// back-pressure cannot be properly exerted.
+	/// It will be removed in a future version.
+	sync_channel: Mutex<mpsc::Sender<NotificationsSinkMessage>>,
+}
+
+/// Message emitted through the [`NotificationsSink`] and processed by the background task
+/// dedicated to the peer.
+#[derive(Debug)]
+enum NotificationsSinkMessage {
+	/// Message emitted by [`NotificationsSink::send_legacy`].
+	Legacy {
+		message: Vec<u8>,
+	},
+
+	/// Message emitted by [`NotificationsSink::reserve_notification`] and
+	/// [`NotificationsSink::write_notification_now`].
+	Notification {
+		protocol_name: Vec<u8>,
+		encoded_fallback_message: Vec<u8>,
+		message: Vec<u8>,
+	},
+
+	/// Must close the connection.
+	ForceClose,
+}
+
+impl NotificationsSink {
+	/// Sends a message to the peer using the legacy substream.
+	///
+	/// If too many messages are already buffered, the message is silently discarded and the
+	/// connection to the peer will be closed shortly after.
+	///
+	/// This method will be removed in a future version.
+	pub fn send_legacy<'a>(&'a self, message: impl Into<Vec<u8>>) {
+		let mut lock = self.inner.sync_channel.lock();
+		let result = lock.try_send(NotificationsSinkMessage::Legacy {
+			message: message.into()
+		});
+
+		if result.is_err() {
+			// Cloning the `mpsc::Sender` guarantees the allocation of an extra spot in the
+			// buffer, and therefore that `try_send` will succeed.
+			let _result2 = lock.clone().try_send(NotificationsSinkMessage::ForceClose);
+			debug_assert!(_result2.map(|()| true).unwrap_or_else(|err| err.is_disconnected()));
+		}
+	}
+
+	/// Sends a notification to the peer.
+	///
+	/// If too many messages are already buffered, the notification is silently discarded and the
+	/// connection to the peer will be closed shortly after.
+	///
+	/// The protocol name is expected to be checked ahead of calling this method. It is a logic
+	/// error to send a notification using an unknown protocol.
+	///
+	/// This method will be removed in a future version.
+	pub fn send_sync_notification<'a>(
+		&'a self,
+		protocol_name: &[u8],
+		encoded_fallback_message: impl Into<Vec<u8>>,
+		message: impl Into<Vec<u8>>
+	) {
+		let mut lock = self.inner.sync_channel.lock();
+		let result = lock.try_send(NotificationsSinkMessage::Notification {
+			protocol_name: protocol_name.to_owned(),
+			encoded_fallback_message: encoded_fallback_message.into(),
+			message: message.into()
+		});
+
+		if result.is_err() {
+			// Cloning the `mpsc::Sender` guarantees the allocation of an extra spot in the
+			// buffer, and therefore that `try_send` will succeed.
+			let _result2 = lock.clone().try_send(NotificationsSinkMessage::ForceClose);
+			debug_assert!(_result2.map(|()| true).unwrap_or_else(|err| err.is_disconnected()));
+		}
+	}
+
+	/// Wait until the remote is ready to accept a notification.
+	///
+	/// Returns an error in the case where the connection is closed.
+	///
+	/// The protocol name is expected to be checked ahead of calling this method. It is a logic
+	/// error to send a notification using an unknown protocol.
+	pub async fn reserve_notification<'a>(&'a self, protocol_name: &[u8]) -> Result<Ready<'a>, ()> {
+		let mut lock = self.inner.async_channel.lock().await;
+
+		let poll_ready = future::poll_fn(|cx| lock.poll_ready(cx)).await;
+		if poll_ready.is_ok() {
+			Ok(Ready { protocol_name: protocol_name.to_owned(), lock })
+		} else {
+			Err(())
+		}
+	}
+}
+
+/// Notification slot is reserved and the notification can actually be sent.
+#[must_use]
+#[derive(Debug)]
+pub struct Ready<'a> {
+	/// Guarded channel. The channel inside is guaranteed to not be full.
+	lock: FuturesMutexGuard<'a, mpsc::Sender<NotificationsSinkMessage>>,
+	/// Name of the protocol. Should match one of the protocols passed at initialization.
+	protocol_name: Vec<u8>,
+}
+
+impl<'a> Ready<'a> {
+	/// Consumes this slots reservation and actually queues the notification.
+	///
+	/// Returns an error if the substream has been closed.
+	pub fn send(
+		mut self,
+		encoded_fallback_message: impl Into<Vec<u8>>,
+		notification: impl Into<Vec<u8>>
+	) -> Result<(), ()> {
+		self.lock.start_send(NotificationsSinkMessage::Notification {
+			protocol_name: self.protocol_name,
+			encoded_fallback_message: encoded_fallback_message.into(),
+			message: notification.into(),
+		}).map_err(|_| ())
+	}
+}
+
+/// Error specific to the collection of protocols.
+#[derive(Debug, derive_more::Display, derive_more::Error)]
+pub enum NotifsHandlerError {
+	/// Channel of synchronous notifications is full.
+	SyncNotificationsClogged,
+	/// Error in legacy protocol.
+	Legacy(<LegacyProtoHandler as ProtocolsHandler>::Error),
+}
+
 impl NotifsHandlerProto {
 	/// Builds a new handler.
 	///
 	/// `list` is a list of notification protocols names, and the message to send as part of the
 	/// handshake. At the moment, the message is always the same whether we open a substream
 	/// ourselves or respond to handshake from the remote.
-	///
-	/// The `queue_size_report` is an optional Prometheus metric that can report the size of the
-	/// messages queue. If passed, it must have one label for the protocol name.
 	pub fn new(
 		legacy: RegisteredProtocol,
 		list: impl Into<Vec<(Cow<'static, [u8]>, Arc<RwLock<Vec<u8>>>)>>,
-		queue_size_report: Option<HistogramVec>
 	) -> Self {
 		let list = list.into();
 
@@ -247,16 +387,7 @@ impl NotifsHandlerProto {
 			.clone()
 			.into_iter()
 			.map(|(proto_name, initial_message)| {
-				let queue_size_report = queue_size_report.as_ref().and_then(|qs| {
-					if let Ok(utf8) = str::from_utf8(&proto_name) {
-						Some(qs.with_label_values(&[utf8]))
-					} else {
-						log::warn!("Ignoring Prometheus metric because {:?} isn't UTF-8", proto_name);
-						None
-					}
-				});
-
-				(NotifsOutHandlerProto::new(proto_name, queue_size_report), initial_message)
+				(NotifsOutHandlerProto::new(proto_name), initial_message)
 			}).collect();
 
 		let in_handlers = list.clone()
@@ -275,13 +406,7 @@ impl NotifsHandlerProto {
 impl ProtocolsHandler for NotifsHandler {
 	type InEvent = NotifsHandlerIn;
 	type OutEvent = NotifsHandlerOut;
-	type Error = EitherError<
-		EitherError<
-			<NotifsInHandler as ProtocolsHandler>::Error,
-			<NotifsOutHandler as ProtocolsHandler>::Error,
-		>,
-		<LegacyProtoHandler as ProtocolsHandler>::Error,
-	>;
+	type Error = NotifsHandlerError;
 	type InboundProtocol = SelectUpgrade<UpgradeCollec<NotificationsIn>, RegisteredProtocol>;
 	type OutboundProtocol = EitherUpgrade<NotificationsOut, RegisteredProtocol>;
 	// Index within the `out_handlers`; None for legacy
@@ -363,24 +488,6 @@ impl ProtocolsHandler for NotifsHandler {
 					self.in_handlers[num].0.inject_event(NotifsInHandlerIn::Refuse);
 				}
 			},
-			NotifsHandlerIn::SendLegacy { message } =>
-				self.legacy.inject_event(LegacyProtoHandlerIn::SendCustomMessage { message }),
-			NotifsHandlerIn::SendNotification { message, encoded_fallback_message, protocol_name } => {
-				for (handler, _) in &mut self.out_handlers {
-					if handler.protocol_name() != &protocol_name[..] {
-						continue;
-					}
-
-					if handler.is_open() {
-						handler.inject_event(NotifsOutHandlerIn::Send(message));
-						return;
-					}
-				}
-
-				self.legacy.inject_event(LegacyProtoHandlerIn::SendCustomMessage {
-					message: encoded_fallback_message,
-				});
-			},
 		}
 	}
 
@@ -461,6 +568,60 @@ impl ProtocolsHandler for NotifsHandler {
 	) -> Poll<
 		ProtocolsHandlerEvent<Self::OutboundProtocol, Self::OutboundOpenInfo, Self::OutEvent, Self::Error>
 	> {
+		if let Some(notifications_sink_rx) = &mut self.notifications_sink_rx {
+			'poll_notifs_sink: loop {
+				// Before we poll the notifications sink receiver, check that all the notification
+				// channels are ready to send a message.
+				// TODO: it is planned that in the future we switch to one `NotificationsSink` per
+				// protocol, in which case each sink should wait only for its corresponding handler
+				// to be ready, and not all handlers
+				// see https://github.com/paritytech/substrate/issues/5670
+				for (out_handler, _) in &mut self.out_handlers {
+					match out_handler.poll_ready(cx) {
+						Poll::Ready(_) => {},
+						Poll::Pending => break 'poll_notifs_sink,
+					}
+				}
+
+				let message = match notifications_sink_rx.poll_next_unpin(cx) {
+					Poll::Ready(Some(msg)) => msg,
+					Poll::Ready(None) | Poll::Pending => break,
+				};
+
+				match message {
+					NotificationsSinkMessage::Legacy { message } => {
+						self.legacy.inject_event(LegacyProtoHandlerIn::SendCustomMessage {
+							message
+						});
+					}
+					NotificationsSinkMessage::Notification {
+						protocol_name,
+						encoded_fallback_message,
+						message
+					} => {
+						for (handler, _) in &mut self.out_handlers {
+							if handler.protocol_name() != &protocol_name[..] {
+								continue;
+							}
+
+							if handler.is_open() {
+								handler.send_or_discard(message);
+							}
+
+							continue 'poll_notifs_sink;
+						}
+
+						self.legacy.inject_event(LegacyProtoHandlerIn::SendCustomMessage {
+							message: encoded_fallback_message,
+						});
+					}
+					NotificationsSinkMessage::ForceClose => {
+						return Poll::Ready(ProtocolsHandlerEvent::Close(NotifsHandlerError::SyncNotificationsClogged));
+					}
+				}
+			}
+		}
+
 		if let Poll::Ready(ev) = self.legacy.poll(cx) {
 			return match ev {
 				ProtocolsHandlerEvent::OutboundSubstreamRequest { protocol, info: () } =>
@@ -468,14 +629,37 @@ impl ProtocolsHandler for NotifsHandler {
 						protocol: protocol.map_upgrade(EitherUpgrade::B),
 						info: None,
 					}),
-				ProtocolsHandlerEvent::Custom(LegacyProtoHandlerOut::CustomProtocolOpen { endpoint, received_handshake, .. }) =>
+				ProtocolsHandlerEvent::Custom(LegacyProtoHandlerOut::CustomProtocolOpen {
+					endpoint,
+					received_handshake,
+					..
+				}) => {
+					let (async_tx, async_rx) = mpsc::channel(ASYNC_NOTIFICATIONS_BUFFER_SIZE);
+					let (sync_tx, sync_rx) = mpsc::channel(SYNC_NOTIFICATIONS_BUFFER_SIZE);
+					let notifications_sink = NotificationsSink {
+						inner: Arc::new(NotificationsSinkInner {
+							async_channel: FuturesMutex::new(async_tx),
+							sync_channel: Mutex::new(sync_tx),
+						}),
+					};
+
+					debug_assert!(self.notifications_sink_rx.is_none());
+					self.notifications_sink_rx = Some(stream::select(async_rx.fuse(), sync_rx.fuse()));
+
 					Poll::Ready(ProtocolsHandlerEvent::Custom(
-						NotifsHandlerOut::Open { endpoint, received_handshake }
-					)),
-				ProtocolsHandlerEvent::Custom(LegacyProtoHandlerOut::CustomProtocolClosed { endpoint, reason }) =>
+						NotifsHandlerOut::Open { endpoint, received_handshake, notifications_sink }
+					))
+				},
+				ProtocolsHandlerEvent::Custom(LegacyProtoHandlerOut::CustomProtocolClosed { endpoint, reason }) => {
+					// We consciously drop the receivers despite notifications being potentially
+					// still buffered up.
+					debug_assert!(self.notifications_sink_rx.is_some());
+					self.notifications_sink_rx = None;
+
 					Poll::Ready(ProtocolsHandlerEvent::Custom(
 						NotifsHandlerOut::Closed { endpoint, reason }
-					)),
+					))
+				},
 				ProtocolsHandlerEvent::Custom(LegacyProtoHandlerOut::CustomMessage { message }) =>
 					Poll::Ready(ProtocolsHandlerEvent::Custom(
 						NotifsHandlerOut::CustomMessage { message }
@@ -485,7 +669,7 @@ impl ProtocolsHandler for NotifsHandler {
 						NotifsHandlerOut::ProtocolError { is_severe, error }
 					)),
 				ProtocolsHandlerEvent::Close(err) =>
-					Poll::Ready(ProtocolsHandlerEvent::Close(EitherError::B(err))),
+					Poll::Ready(ProtocolsHandlerEvent::Close(NotifsHandlerError::Legacy(err))),
 			}
 		}
 

substrate/client/network/src/protocol/generic_proto/handler/notif_out.rs

@@ -34,8 +34,10 @@ use libp2p::swarm::{
 	NegotiatedSubstream,
 };
 use log::{debug, warn, error};
-use prometheus_endpoint::Histogram;
-use std::{borrow::Cow, collections::VecDeque, fmt, mem, pin::Pin, task::{Context, Poll}, time::Duration};
+use std::{
+	borrow::Cow, collections::VecDeque, fmt, mem, pin::Pin, task::{Context, Poll, Waker},
+	time::Duration
+};
 use wasm_timer::Instant;
 
 /// Maximum duration to open a substream and receive the handshake message. After that, we
@@ -56,17 +58,14 @@ const INITIAL_KEEPALIVE_TIME: Duration = Duration::from_secs(5);
 pub struct NotifsOutHandlerProto {
 	/// Name of the protocol to negotiate.
 	protocol_name: Cow<'static, [u8]>,
-	/// Optional Prometheus histogram to report message queue size variations.
-	queue_size_report: Option<Histogram>,
 }
 
 impl NotifsOutHandlerProto {
 	/// Builds a new [`NotifsOutHandlerProto`]. Will use the given protocol name for the
 	/// notifications substream.
-	pub fn new(protocol_name: impl Into<Cow<'static, [u8]>>, queue_size_report: Option<Histogram>) -> Self {
+	pub fn new(protocol_name: impl Into<Cow<'static, [u8]>>) -> Self {
 		NotifsOutHandlerProto {
 			protocol_name: protocol_name.into(),
-			queue_size_report,
 		}
 	}
 }
@@ -78,14 +77,12 @@ impl IntoProtocolsHandler for NotifsOutHandlerProto {
 		DeniedUpgrade
 	}
 
-	fn into_handler(self, peer_id: &PeerId, _: &ConnectedPoint) -> Self::Handler {
+	fn into_handler(self, _: &PeerId, _: &ConnectedPoint) -> Self::Handler {
 		NotifsOutHandler {
 			protocol_name: self.protocol_name,
 			when_connection_open: Instant::now(),
-			queue_size_report: self.queue_size_report,
 			state: State::Disabled,
 			events_queue: VecDeque::new(),
-			peer_id: peer_id.clone(),
 		}
 	}
 }
@@ -108,17 +105,11 @@ pub struct NotifsOutHandler {
 	/// When the connection with the remote has been successfully established.
 	when_connection_open: Instant,
 
-	/// Optional prometheus histogram to report message queue sizes variations.
-	queue_size_report: Option<Histogram>,
-
 	/// Queue of events to send to the outside.
 	///
 	/// This queue must only ever be modified to insert elements at the back, or remove the first
 	/// element.
 	events_queue: VecDeque<ProtocolsHandlerEvent<NotificationsOut, (), NotifsOutHandlerOut, void::Void>>,
-
-	/// Who we are connected to.
-	peer_id: PeerId,
 }
 
 /// Our relationship with the node we're connected to.
@@ -153,6 +144,11 @@ enum State {
 	Open {
 		/// Substream that is currently open.
 		substream: NotificationsOutSubstream<NegotiatedSubstream>,
+		/// Waker for the last task that got `Poll::Pending` from `poll_ready`, to notify
+		/// when the open substream closes due to being disabled or encountering an
+		/// error, i.e. to notify the task as soon as the substream becomes unavailable,
+		/// without waiting for an underlying I/O task wakeup.
+		close_waker: Option<Waker>,
 		/// The initial message that we sent. Necessary if we need to re-open a substream.
 		initial_message: Vec<u8>,
 	},
@@ -173,11 +169,6 @@ pub enum NotifsOutHandlerIn {
 
 	/// Disables the notifications substream for this node. This is the default state.
 	Disable,
-
-	/// Sends a message on the notifications substream. Ignored if the substream isn't open.
-	///
-	/// It is only valid to send this if the notifications substream has been enabled.
-	Send(Vec<u8>),
 }
 
 /// Event that can be emitted by a `NotifsOutHandler`.
@@ -216,6 +207,41 @@ impl NotifsOutHandler {
 	pub fn protocol_name(&self) -> &[u8] {
 		&self.protocol_name
 	}
+
+	/// Polls whether the outbound substream is ready to send a notification.
+	///
+	/// - Returns `Poll::Pending` if the substream is open but not ready to send a notification.
+	/// - Returns `Poll::Ready(true)` if the substream is ready to send a notification.
+	/// - Returns `Poll::Ready(false)` if the substream is closed.
+	///
+	pub fn poll_ready(&mut self, cx: &mut Context) -> Poll<bool> {
+		if let State::Open { substream, close_waker, .. } = &mut self.state {
+			match substream.poll_ready_unpin(cx) {
+				Poll::Ready(Ok(())) => Poll::Ready(true),
+				Poll::Ready(Err(_)) => Poll::Ready(false),
+				Poll::Pending => {
+					*close_waker = Some(cx.waker().clone());
+					Poll::Pending
+				}
+			}
+		} else {
+			Poll::Ready(false)
+		}
+	}
+
+	/// Sends out a notification.
+	///
+	/// If the substream is closed, or not ready to send out a notification yet, then the
+	/// notification is silently discarded.
+	///
+	/// You are encouraged to call [`NotifsOutHandler::poll_ready`] beforehand to determine
+	/// whether this will succeed. If `Poll::Ready(true)` is returned, then this method will send
+	/// out a notification.
+	pub fn send_or_discard(&mut self, notification: Vec<u8>) {
+		if let State::Open { substream, .. } = &mut self.state {
+			let _ = substream.start_send_unpin(notification);
+		}
+	}
 }
 
 impl ProtocolsHandler for NotifsOutHandler {
@@ -247,7 +273,7 @@ impl ProtocolsHandler for NotifsOutHandler {
 			State::Opening { initial_message } => {
 				let ev = NotifsOutHandlerOut::Open { handshake: handshake_msg };
 				self.events_queue.push_back(ProtocolsHandlerEvent::Custom(ev));
-				self.state = State::Open { substream, initial_message };
+				self.state = State::Open { substream, initial_message, close_waker: None };
 			},
 			// If the handler was disabled while we were negotiating the protocol, immediately
 			// close it.
@@ -310,31 +336,15 @@ impl ProtocolsHandler for NotifsOutHandler {
 					}
 					State::Opening { .. } => self.state = State::DisabledOpening,
 					State::Refused => self.state = State::Disabled,
-					State::Open { substream, .. } => self.state = State::DisabledOpen(substream),
+					State::Open { substream, close_waker, .. } => {
+						if let Some(close_waker) = close_waker {
+							close_waker.wake();
+						}
+						self.state = State::DisabledOpen(substream)
+					},
 					State::Poisoned => error!("☎️ Notifications handler in a poisoned state"),
 				}
 			}
-
-			NotifsOutHandlerIn::Send(msg) =>
-				if let State::Open { substream, .. } = &mut self.state {
-					if substream.push_message(msg).is_err() {
-						warn!(
-							target: "sub-libp2p",
-							"📞 Notifications queue with peer {} is full, dropped message (protocol: {:?})",
-							self.peer_id,
-							self.protocol_name,
-						);
-					}
-					if let Some(metric) = &self.queue_size_report {
-						metric.observe(substream.queue_len() as f64);
-					}
-				} else {
-					// This is an API misuse.
-					warn!(
-						target: "sub-libp2p",
-						"📞 Tried to send a notification on a disabled handler"
-					);
-				},
 		}
 	}
 
@@ -375,10 +385,14 @@ impl ProtocolsHandler for NotifsOutHandler {
 		}
 
 		match &mut self.state {
-			State::Open { substream, initial_message } =>
+			State::Open { substream, initial_message, close_waker } =>
 				match Sink::poll_flush(Pin::new(substream), cx) {
 					Poll::Pending | Poll::Ready(Ok(())) => {},
 					Poll::Ready(Err(_)) => {
+						if let Some(close_waker) = close_waker.take() {
+							close_waker.wake();
+						}
+
 						// We try to re-open a substream.
 						let initial_message = mem::replace(initial_message, Vec::new());
 						self.state = State::Opening { initial_message: initial_message.clone() };