Return join handle separately

nym-authenticator-client/src/mixnet_listener.rs

@@ -25,7 +25,7 @@ pub type MixnetMessageInputReceiver = mpsc::Receiver<InputMessage>; // 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}");
-        }
-    }
 }

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 {

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,
             },
         )))

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<dyn BandwidthTicketProvider>,
 }