notification: Unregister dropped protocols (#391)

Part of the investigation from:
- paritytech/polkadot-sdk#8474

---------

Signed-off-by: Alexandru Vasile <[email protected]>

Author: lexnv

src/protocol/notification/mod.rs

@@ -1795,6 +1795,9 @@ impl NotificationProtocol {
                         protocol = %self.protocol,
                         "user protocol has exited, exiting"
                     );
+
+                    self.service.unregister_protocol();
+
                     return true;
                 }
                 Some(command) => match command {

src/protocol/transport_service.rs

@@ -544,6 +544,14 @@ impl TransportService {
     pub fn local_peer_id(&self) -> PeerId {
         self.local_peer_id
     }
+
+    /// Dynamically unregister a protocol.
+    ///
+    /// This must be called when a protocol is no longer needed (e.g. user dropped the protocol
+    /// handle).
+    pub fn unregister_protocol(&self) {
+        self.transport_handle.unregister_protocol(self.protocol.clone());
+    }
 }
 
 impl Stream for TransportService {

src/transport/manager/handle.rs

@@ -60,6 +60,11 @@ pub enum InnerTransportManagerCommand {
         /// Remote address.
         address: Multiaddr,
     },
+
+    UnregisterProtocol {
+        /// Protocol name.
+        protocol: ProtocolName,
+    },
 }
 
 /// Handle for communicating with [`crate::transport::manager::TransportManager`].
@@ -285,6 +290,29 @@ impl TransportManagerHandle {
                 TrySendError::Closed(_) => ImmediateDialError::TaskClosed,
             })
     }
+
+    /// Dynamically unregister a protocol.
+    ///
+    /// This must be called when a protocol is no longer needed (e.g. user dropped the protocol
+    /// handle).
+    pub fn unregister_protocol(&self, protocol: ProtocolName) {
+        tracing::info!(
+            target: LOG_TARGET,
+            ?protocol,
+            "Unregistering user protocol on handle drop"
+        );
+
+        if let Err(err) = self
+            .cmd_tx
+            .try_send(InnerTransportManagerCommand::UnregisterProtocol { protocol })
+        {
+            tracing::error!(
+                target: LOG_TARGET,
+                ?err,
+                "Failed to unregister protocol"
+            );
+        }
+    }
 }
 
 pub struct TransportHandle {

src/transport/manager/mod.rs

@@ -360,6 +360,26 @@ impl TransportManager {
         service
     }
 
+    /// Unregister a protocol in response of the user dropping the protocol handle.
+    fn unregister_protocol(&mut self, protocol: ProtocolName) {
+        let Some(context) = self.protocols.remove(&protocol) else {
+            tracing::error!(target: LOG_TARGET, ?protocol, "Cannot unregister protocol, not registered");
+            return;
+        };
+
+        for fallback in &context.fallback_names {
+            if !self.protocol_names.remove(fallback) {
+                tracing::error!(target: LOG_TARGET, ?fallback, ?protocol, "Cannot unregister fallback protocol, not registered");
+            }
+        }
+
+        tracing::info!(
+            target: LOG_TARGET,
+            ?protocol,
+            "Protocol fully unregistered"
+        );
+    }
+
     /// Acquire `TransportHandle`.
     pub fn transport_handle(&self, executor: Arc<dyn Executor>) -> TransportHandle {
         TransportHandle {
@@ -1042,6 +1062,9 @@ impl TransportManager {
                                 tracing::debug!(target: LOG_TARGET, ?error, "failed to dial peer")
                             }
                         }
+                        InnerTransportManagerCommand::UnregisterProtocol { protocol } => {
+                            self.unregister_protocol(protocol);
+                        }
                     }
                 },
 
@@ -1256,7 +1279,7 @@ impl TransportManager {
                                                     ?peer,
                                                     %protocol,
                                                     ?connection_id,
-                                                    "call to protocol would, block try sending in a blocking way",
+                                                    "call to protocol would block try sending in a blocking way",
                                                 );
 
                                                 context