diff --git a/src/protocol/notification/mod.rs b/src/protocol/notification/mod.rs index 40bb762d..363adcfa 100644 --- a/src/protocol/notification/mod.rs +++ b/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 { diff --git a/src/protocol/transport_service.rs b/src/protocol/transport_service.rs index 4e8bc93a..b729e931 100644 --- a/src/protocol/transport_service.rs +++ b/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 { diff --git a/src/transport/manager/handle.rs b/src/transport/manager/handle.rs index 70bbfd3b..8d99666e 100644 --- a/src/transport/manager/handle.rs +++ b/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 { diff --git a/src/transport/manager/mod.rs b/src/transport/manager/mod.rs index 5fc9e330..f44a07a6 100644 --- a/src/transport/manager/mod.rs +++ b/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) -> 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