From 81e72ba61bca72e3c8740109e9fa85b8b50f8f5e Mon Sep 17 00:00:00 2001 From: Andrej Mihajlov Date: Tue, 14 Oct 2025 20:01:47 +0200 Subject: [PATCH] Add shutdown event handle --- Cargo.lock | 1 - .../src/mixnet_listener.rs | 20 +++++++------- nym-ip-packet-client/src/connect.rs | 12 ++++----- sdk/rust/nym-sdk/src/mixnet/native_client.rs | 27 ++++++++++++++++--- 4 files changed, 37 insertions(+), 23 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 20a519f085..01206ee22f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5746,7 +5746,6 @@ dependencies = [ "criterion", "ff", "group", - "lazy_static", "nym-contracts-common", "nym-pemstore", "rand 0.8.5", diff --git a/nym-authenticator-client/src/mixnet_listener.rs b/nym-authenticator-client/src/mixnet_listener.rs index 9cb6828e83..e6564693a1 100644 --- a/nym-authenticator-client/src/mixnet_listener.rs +++ b/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, } @@ -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 diff --git a/nym-ip-packet-client/src/connect.rs b/nym-ip-packet-client/src/connect.rs index 5a37cf6d82..e8c49e66ab 100644 --- a/nym-ip-packet-client/src/connect.rs +++ b/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); diff --git a/sdk/rust/nym-sdk/src/mixnet/native_client.rs b/sdk/rust/nym-sdk/src/mixnet/native_client.rs index c9495680b2..3724274d7a 100644 --- a/sdk/rust/nym-sdk/src/mixnet/native_client.rs +++ b/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> { 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 + } +}