Add shutdown event handle

Cargo.lock

@@ -5746,7 +5746,6 @@ dependencies = [
  "criterion",
  "ff",
  "group",
- "lazy_static",
  "nym-contracts-common",
  "nym-pemstore",
  "rand 0.8.5",

nym-authenticator-client/src/mixnet_listener.rs

@@ -51,12 +51,12 @@ impl AuthClientMixnetListener {
     }
 
     async fn run(mut self) -> Self {
-        let mixnet_cancel_token = self.mixnet_client.cancellation_token();
         self.shutdown_token.run_until_cancelled(async {
+            let shutdown_event = self.mixnet_client.shutdown_event();
             loop {
                 tokio::select! {
                     biased;
-                    _ = mixnet_cancel_token.cancelled() => {
+                    _ = shutdown_event.wait() => {
                         tracing::debug!("AuthClientMixnetListener: mixnet client was shutdown");
                         break;
                     }
@@ -100,9 +100,7 @@ impl AuthClientMixnetListener {
 
     // Disconnects the mixnet client and effectively drop itself, since it doesn't work without one, and reconnecting isn't supported
     pub async fn disconnect_mixnet_client(self) {
-        if !self.mixnet_client.cancellation_token().is_cancelled() {
-            self.mixnet_client.disconnect().await;
-        }
+        self.mixnet_client.disconnect().await;
     }
 
     pub fn start(self) -> AuthClientMixnetListenerHandle {
@@ -110,14 +108,14 @@ impl AuthClientMixnetListener {
         let message_sender = self.input_message_tx.clone();
         // Allows stopping only this, e.g. if we don't need it in the new bandwidth controller
         let cancellation_token = self.shutdown_token.clone();
-        let mixnet_cancellation_token = self.mixnet_client.cancellation_token();
+        // let mixnet_cancellation_token = self.mixnet_client.cancellation_token();
         let handle = tokio::spawn(self.run());
 
         AuthClientMixnetListenerHandle {
             message_broadcast,
             message_sender,
             cancellation_token,
-            mixnet_cancellation_token,
+            // mixnet_cancellation_token,
             handle,
         }
     }
@@ -127,7 +125,7 @@ pub struct AuthClientMixnetListenerHandle {
     message_broadcast: MixnetMessageBroadcastSender,
     message_sender: MixnetMessageInputSender,
     cancellation_token: CancellationToken,
-    mixnet_cancellation_token: CancellationToken,
+    // mixnet_cancellation_token: CancellationToken,
     handle: JoinHandle<AuthClientMixnetListener>,
 }
 
@@ -140,9 +138,9 @@ impl AuthClientMixnetListenerHandle {
         self.message_broadcast.subscribe()
     }
 
-    pub fn mixnet_cancel_token(&self) -> CancellationToken {
-        self.mixnet_cancellation_token.clone()
-    }
+    // pub fn mixnet_cancel_token(&self) -> CancellationToken {
+    //     self.mixnet_cancellation_token.clone()
+    // }
 
     pub async fn stop(self) {
         // If shutdown was externally called, that call is a no-op

nym-ip-packet-client/src/connect.rs

@@ -133,8 +133,7 @@ impl IprClientConnect {
 
         let timeout = sleep(IPR_CONNECT_TIMEOUT);
         tokio::pin!(timeout);
-
-        let mixnet_cancel_token = self.mixnet_client.cancellation_token();
+        let shutdown_event = self.mixnet_client.shutdown_event();
 
         loop {
             tokio::select! {
@@ -142,15 +141,14 @@ impl IprClientConnect {
                     error!("Cancelled while waiting for reply to connect request");
                     return Err(Error::Cancelled);
                 },
-
-                _ = mixnet_cancel_token.cancelled() => {
-                    error!("Mixnet client stopped while waiting for reply to connect request");
-                    return Err(Error::Cancelled);
-                },
                 _ = &mut timeout => {
                     error!("Timed out waiting for reply to connect request");
                     return Err(Error::TimeoutWaitingForConnectResponse);
                 },
+                _ = shutdown_event.wait() => {
+                    error!("Mixnet client stopped while waiting for reply to connect request");
+                    return Err(Error::Cancelled);
+                },
                 msgs = self.mixnet_client.wait_for_messages() => match msgs {
                     None => {
                         return Err(Error::NoMixnetMessagesReceived);

sdk/rust/nym-sdk/src/mixnet/native_client.rs

@@ -24,7 +24,6 @@ use std::pin::Pin;
 use std::sync::Arc;
 use std::task::{Context, Poll};
 use tokio::sync::RwLockReadGuard;
-use tokio_util::sync::CancellationToken;
 
 /// Client connected to the Nym mixnet.
 pub struct MixnetClient {
@@ -124,9 +123,9 @@ impl MixnetClient {
     }
 
     /// Get a child token of the root, to monitor unexpected shutdown, without causing one
-    pub fn cancellation_token(&self) -> CancellationToken {
-        self.shutdown_handle.child_shutdown_token().inner().clone()
-    }
+    // pub fn cancellation_token(&self) -> CancellationToken {
+    //     self.shutdown_handle.child_shutdown_token().inner().clone()
+    // }
 
     pub fn client_request_sender(&self) -> ClientRequestSender {
         self.client_request_sender.clone()
@@ -199,6 +198,13 @@ impl MixnetClient {
         self.client_state.topology_accessor.release_manual_control()
     }
 
+    /// Returns a shutdown event handle that can be used for waiting for the client to shutdown.
+    pub fn shutdown_event(&self) -> ShutdownEventHandle {
+        ShutdownEventHandle {
+            shutdown_handle: self.shutdown_handle.clone(),
+        }
+    }
+
     /// Wait for messages from the mixnet
     pub async fn wait_for_messages(&mut self) -> Option<Vec<ReconstructedMessage>> {
         self.reconstructed_receiver.next().await
@@ -340,3 +346,16 @@ impl MixnetMessageSender for MixnetClientSender {
             .map_err(|_| Error::MessageSendingFailure)
     }
 }
+
+/// Handle for waiting on the shutdown event of the mixnet client.
+pub struct ShutdownEventHandle {
+    shutdown_handle: ShutdownTracker,
+}
+
+impl ShutdownEventHandle {
+    /// Returns once mixnet client has been shut down.
+    /// If mixnet client is already shut down, returns immediately.
+    pub async fn wait(&self) {
+        self.shutdown_handle.wait_for_tracker().await
+    }
+}