diff --git a/nym-authenticator-client/src/mixnet_listener.rs b/nym-authenticator-client/src/mixnet_listener.rs index cf4815ba19..89c7c07d58 100644 --- a/nym-authenticator-client/src/mixnet_listener.rs +++ b/nym-authenticator-client/src/mixnet_listener.rs @@ -25,7 +25,7 @@ pub type MixnetMessageInputReceiver = mpsc::Receiver; // This coul pub fn spawn( mut mixnet_client: MixnetClient, shutdown_token: CancellationToken, -) -> AuthClientMixnetListenerHandle { +) -> (JoinHandle<()>, AuthClientMixnetListenerHandle) { // Broadcast channel for the messages that we re-broadcast to the AuthClients let (message_broadcast, _) = broadcast::channel(100); // Channel for message to send to the mixnet @@ -80,17 +80,19 @@ pub fn spawn( tracing::debug!("AuthClientMixnetListener is shutting down"); }); - AuthClientMixnetListenerHandle { - message_broadcast: cloned_message_broadcast, - message_sender: cloned_message_sender, + ( join_handle, - } + AuthClientMixnetListenerHandle { + message_broadcast: cloned_message_broadcast, + message_sender: cloned_message_sender, + }, + ) } +/// Handle that enables the consumer to subscribe or submit messages to the mixnet. pub struct AuthClientMixnetListenerHandle { message_broadcast: MixnetMessageBroadcastSender, message_sender: MixnetMessageInputSender, - join_handle: JoinHandle<()>, } impl AuthClientMixnetListenerHandle { @@ -101,13 +103,4 @@ impl AuthClientMixnetListenerHandle { pub fn subscribe(&self) -> MixnetMessageBroadcastReceiver { self.message_broadcast.subscribe() } - - /// Join on listener handle to wait until it stops - /// - /// Important: Use cancellation token to initiate the shutdown - pub async fn join(self) { - if let Err(e) = self.join_handle.await { - tracing::error!("Error waiting for auth clients mixnet listener to stop: {e}"); - } - } } diff --git a/nym-gateway-probe/src/lib.rs b/nym-gateway-probe/src/lib.rs index e66a1c9dd1..680ffd0898 100644 --- a/nym-gateway-probe/src/lib.rs +++ b/nym-gateway-probe/src/lib.rs @@ -401,7 +401,7 @@ impl Probe { { // Start the mixnet listener that the auth clients use to receive messages. let mixnet_listener_cancel_token = CancellationToken::new(); - let mixnet_listener_task = + let (mixnet_listener_join_handle, mixnet_listener_task) = mixnet_listener::spawn(mixnet_client, mixnet_listener_cancel_token.child_token()); let auth_client = AuthenticatorClient::new( @@ -443,7 +443,9 @@ impl Probe { .unwrap_or_default(); mixnet_listener_cancel_token.cancel(); - mixnet_listener_task.join().await; + if let Err(e) = mixnet_listener_join_handle.await { + tracing::error!("Failed to join on mixnet listener: {e}"); + } outcome } else { diff --git a/nym-registration-client/src/lib.rs b/nym-registration-client/src/lib.rs index 3fddc076ae..339b9add94 100644 --- a/nym-registration-client/src/lib.rs +++ b/nym-registration-client/src/lib.rs @@ -85,7 +85,7 @@ impl RegistrationClient { // Start the auth client mixnet listener, which will listen for incoming messages from the // mixnet and rebroadcast them to the auth clients. - let mixnet_listener = mixnet_listener::spawn( + let (mixnet_listener_join_handle, mixnet_listener) = mixnet_listener::spawn( self.mixnet_client, // todo: are we sure we want to clone the cancel token here? self.cancel_token.clone(), @@ -142,6 +142,7 @@ impl RegistrationClient { entry_gateway_data: entry, exit_gateway_data: exit, authenticator_listener_handle: mixnet_listener, + mixnet_listener_join_handle, bw_controller: self.bandwidth_controller, }, ))) diff --git a/nym-registration-client/src/types.rs b/nym-registration-client/src/types.rs index 14588c3b3b..2d2ebb7555 100644 --- a/nym-registration-client/src/types.rs +++ b/nym-registration-client/src/types.rs @@ -24,5 +24,6 @@ pub struct WireguardRegistrationResult { pub entry_gateway_data: GatewayData, pub exit_gateway_data: GatewayData, pub authenticator_listener_handle: AuthClientMixnetListenerHandle, + pub mixnet_listener_join_handle: tokio::task::JoinHandle<()>, pub bw_controller: Box, }