diff --git a/src/protocol/notification/mod.rs b/src/protocol/notification/mod.rs index cfc40120..57742fbd 100644 --- a/src/protocol/notification/mod.rs +++ b/src/protocol/notification/mod.rs @@ -1606,7 +1606,9 @@ impl NotificationProtocol { } /// Handle next notification event. - async fn next_event(&mut self) { + /// + /// Returns `true` when the user command stream was dropped. + async fn next_event(&mut self) -> bool { // biased select is used because the substream events must be prioritized above other events // that is because a closed substream is detected by either `substreams` or `negotiation` // and if that event is not handled with priority but, e.g., inbound substream is @@ -1783,9 +1785,16 @@ impl NotificationProtocol { ); } } + + // User commands. command = self.command_rx.recv() => match command { None => { - tracing::debug!(target: LOG_TARGET, "user protocol has exited, exiting"); + tracing::debug!( + target: LOG_TARGET, + protocol = %self.protocol, + "user protocol has exited, exiting" + ); + return true; } Some(command) => match command { NotificationCommand::OpenSubstream { peers } => { @@ -1811,14 +1820,14 @@ impl NotificationProtocol { } }, } + + false } /// Start [`NotificationProtocol`] event loop. pub(crate) async fn run(mut self) { tracing::debug!(target: LOG_TARGET, "starting notification event loop"); - loop { - self.next_event().await; - } + while !self.next_event().await {} } } diff --git a/src/protocol/notification/tests/notification.rs b/src/protocol/notification/tests/notification.rs index 763cb5ea..dc34f7ec 100644 --- a/src/protocol/notification/tests/notification.rs +++ b/src/protocol/notification/tests/notification.rs @@ -1117,3 +1117,22 @@ async fn second_inbound_substream_opened_while_outbound_substream_was_opening() state => panic!("invalid state for peer: {state:?}"), } } + +#[tokio::test] +async fn drop_handle_exits_protocol() { + let _ = tracing_subscriber::fmt() + .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) + .try_init(); + + let (mut protocol, handle, _sender, _tx) = make_notification_protocol(); + + // Simulate a handle drop. + drop(handle); + + // Call `next_event` and ensure it returns true. + let result = protocol.next_event().await; + assert!( + result, + "Expected `next_event` to return true when `command_rx` is dropped" + ); +}