diff --git a/common/client-core/src/client/base_client/helpers.rs b/common/client-core/src/client/base_client/helpers.rs index 68c0cb61bc..4259ead046 100644 --- a/common/client-core/src/client/base_client/helpers.rs +++ b/common/client-core/src/client/base_client/helpers.rs @@ -1,9 +1,7 @@ // Copyright 2023 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 -use crate::error::ClientCoreError; use crate::{client::replies::reply_storage, config::DebugConfig}; -use nym_task::{ShutdownManager, ShutdownToken, ShutdownTracker}; pub fn setup_empty_reply_surb_backend(debug_config: &DebugConfig) -> reply_storage::Empty { reply_storage::Empty { @@ -15,49 +13,3 @@ pub fn setup_empty_reply_surb_backend(debug_config: &DebugConfig) -> reply_stora .maximum_reply_surb_storage_threshold, } } - -// old 'TaskHandle' -pub(crate) enum ShutdownHelper { - Internal(ShutdownManager), - External(ShutdownTracker), -} - -fn new_shutdown_manager() -> Result { - cfg_if::cfg_if! { - if #[cfg(not(target_arch = "wasm32"))] { - Ok(ShutdownManager::build_new_default()?) - } else { - Ok(ShutdownManager::new_without_signals()) - } - } -} - -impl ShutdownHelper { - pub(crate) fn new(shutdown_tracker: Option) -> Result { - match shutdown_tracker { - None => Ok(ShutdownHelper::Internal(new_shutdown_manager()?)), - Some(shutdown_tracker) => Ok(ShutdownHelper::External(shutdown_tracker)), - } - } - - pub(crate) fn into_internal(self) -> Option { - match self { - ShutdownHelper::Internal(manager) => Some(manager), - ShutdownHelper::External(_) => None, - } - } - - pub(crate) fn shutdown_token(&self) -> ShutdownToken { - match self { - ShutdownHelper::External(shutdown) => shutdown.clone_shutdown_token(), - ShutdownHelper::Internal(shutdown) => shutdown.clone_shutdown_token(), - } - } - - pub(crate) fn tracker(&self) -> &ShutdownTracker { - match self { - ShutdownHelper::External(shutdown) => shutdown, - ShutdownHelper::Internal(shutdown) => shutdown.shutdown_tracker(), - } - } -} diff --git a/common/client-core/src/client/base_client/mod.rs b/common/client-core/src/client/base_client/mod.rs index 2b187ca05c..2a3e30bf64 100644 --- a/common/client-core/src/client/base_client/mod.rs +++ b/common/client-core/src/client/base_client/mod.rs @@ -4,7 +4,6 @@ use super::mix_traffic::ClientRequestSender; use super::received_buffer::ReceivedBufferMessage; use super::statistics_control::StatisticsControl; -use crate::client::base_client::helpers::ShutdownHelper; use crate::client::base_client::storage::helpers::store_client_keys; use crate::client::base_client::storage::MixnetClientStorage; use crate::client::cover_traffic_stream::LoopCoverTrafficStream; @@ -53,7 +52,7 @@ use nym_sphinx::receiver::{ReconstructedMessage, SphinxMessageReceiver}; use nym_statistics_common::clients::ClientStatsSender; use nym_statistics_common::generate_client_stats_id; use nym_task::connections::{ConnectionCommandReceiver, ConnectionCommandSender, LaneQueueLengths}; -use nym_task::{ShutdownManager, ShutdownTracker}; +use nym_task::ShutdownTracker; use nym_topology::provider_trait::TopologyProvider; use nym_topology::HardcodedTopologyProvider; use nym_validator_client::nym_api::NymApiClientExt; @@ -881,8 +880,12 @@ where let shared_topology_accessor = TopologyAccessor::new(self.config.debug.topology.ignore_egress_epoch_role); - // Shutdown notifier for signalling tasks to stop - let shutdown = ShutdownHelper::new(self.shutdown)?; + // Create a shutdown tracker for this client - either as a child of provided tracker + // or get one from the registry + let shutdown_tracker = match self.shutdown { + Some(parent_tracker) => parent_tracker.child_tracker(), + None => nym_task::get_sdk_shutdown_tracker()?, + }; // channels responsible for dealing with reply-related fun let (reply_controller_sender, reply_controller_receiver) = @@ -914,7 +917,7 @@ where self.user_agent.clone(), generate_client_stats_id(*self_address.identity()), input_sender.clone(), - shutdown.tracker(), + &shutdown_tracker.child_tracker(), ); // needs to be started as the first thing to block if required waiting for the gateway @@ -924,14 +927,14 @@ where shared_topology_accessor.clone(), self_address.gateway(), self.wait_for_gateway, - shutdown.tracker(), + &shutdown_tracker.child_tracker(), ) .await?; let gateway_packet_router = PacketRouter::new( ack_sender, mixnet_messages_sender, - shutdown.shutdown_token(), + shutdown_tracker.clone_shutdown_token(), ); let gateway_transceiver = Self::setup_gateway_transceiver( @@ -944,7 +947,7 @@ where stats_reporter.clone(), #[cfg(unix)] self.connection_fd_callback, - shutdown.tracker(), + &shutdown_tracker.child_tracker(), ) .await?; let gateway_ws_fd = gateway_transceiver.ws_fd(); @@ -952,7 +955,7 @@ where let reply_storage = Self::setup_persistent_reply_storage( reply_storage_backend, key_rotation_config, - shutdown.tracker(), + &shutdown_tracker.child_tracker(), ) .await?; @@ -963,7 +966,7 @@ where reply_storage.key_storage(), reply_controller_sender.clone(), stats_reporter.clone(), - shutdown.tracker(), + &shutdown_tracker.child_tracker(), ); // The message_sender is the transmitter for any component generating sphinx packets @@ -971,8 +974,10 @@ where // traffic stream. // The MixTrafficController then sends the actual traffic - let (message_sender, client_request_sender) = - Self::start_mix_traffic_controller(gateway_transceiver, shutdown.tracker()); + let (message_sender, client_request_sender) = Self::start_mix_traffic_controller( + gateway_transceiver, + &shutdown_tracker.child_tracker(), + ); // Channels that the websocket listener can use to signal downstream to the real traffic // controller that connections are closed. @@ -1001,7 +1006,7 @@ where shared_lane_queue_lengths.clone(), client_connection_rx, stats_reporter.clone(), - shutdown.tracker(), + &shutdown_tracker.child_tracker(), ); if !self @@ -1017,7 +1022,7 @@ where shared_topology_accessor.clone(), message_sender, stats_reporter.clone(), - shutdown.tracker(), + &shutdown_tracker.child_tracker(), ); } @@ -1045,7 +1050,7 @@ where gateway_connection: GatewayConnection { gateway_ws_fd }, }, stats_reporter, - shutdown_handle: shutdown.into_internal(), + shutdown_handle: Some(shutdown_tracker), // The primary tracker for this client client_request_sender, forget_me: self.config.debug.forget_me, remember_me: self.config.debug.remember_me, @@ -1061,7 +1066,7 @@ pub struct BaseClient { pub client_state: ClientState, pub stats_reporter: ClientStatsSender, pub client_request_sender: ClientRequestSender, - pub shutdown_handle: Option, + pub shutdown_handle: Option, pub forget_me: ForgetMe, pub remember_me: RememberMe, } diff --git a/common/client-core/src/error.rs b/common/client-core/src/error.rs index c76b8a3c88..dc6e4fc1b2 100644 --- a/common/client-core/src/error.rs +++ b/common/client-core/src/error.rs @@ -4,6 +4,7 @@ use crate::client::mix_traffic::transceiver::ErasedGatewayError; use nym_crypto::asymmetric::ed25519::Ed25519RecoveryError; use nym_gateway_client::error::GatewayClientError; +use nym_task::RegistryAccessError; use nym_topology::node::RoutingNodeError; use nym_topology::{NodeId, NymTopologyError}; use nym_validator_client::nym_api::error::NymAPIError; @@ -242,6 +243,9 @@ pub enum ClientCoreError { #[error("failed to select valid gateway due to incomputable latency")] GatewaySelectionFailure { source: WeightedError }, + + #[error("Could not access task registry, {0}")] + RegistryAccess(#[from] RegistryAccessError), } impl From for ClientCoreError { diff --git a/common/task/src/cancellation/token.rs b/common/task/src/cancellation/token.rs index 18d6648219..e9b3778bf7 100644 --- a/common/task/src/cancellation/token.rs +++ b/common/task/src/cancellation/token.rs @@ -17,6 +17,12 @@ pub struct ShutdownToken { inner: CancellationToken, } +impl From for ShutdownToken { + fn from(inner: CancellationToken) -> Self { + ShutdownToken { inner } + } +} + impl ShutdownToken { /// A drop in no-op replacement for `send_status_msg` for easier migration from [TaskClient](crate::TaskClient). #[deprecated] @@ -45,6 +51,13 @@ impl ShutdownToken { &self.inner } + /// Get an owned [CancellationToken](tokio_util::sync::CancellationToken) for public API use. + /// This is useful when you need to expose cancellation to SDK users without + /// exposing the internal ShutdownToken type. + pub fn to_cancellation_token(&self) -> CancellationToken { + self.inner.clone() + } + /// Creates a `ShutdownToken` which will get cancelled whenever the /// current token gets cancelled. Unlike a cloned `ShutdownToken`, /// cancelling a child token does not cancel the parent token. diff --git a/common/task/src/cancellation/tracker.rs b/common/task/src/cancellation/tracker.rs index c48b9c71ff..12e2538dee 100644 --- a/common/task/src/cancellation/tracker.rs +++ b/common/task/src/cancellation/tracker.rs @@ -314,4 +314,40 @@ impl ShutdownTracker { pub fn clone_shutdown_token(&self) -> ShutdownToken { self.root_cancellation_token.clone() } + + /// Create a child ShutdownTracker that inherits cancellation from this tracker + /// but has its own TaskTracker for managing sub-tasks. + /// + /// This enables hierarchical task management where: + /// - Parent cancellation flows to all children + /// - Each level tracks its own tasks independently + /// - Components can wait for their specific sub-tasks to complete + pub fn child_tracker(&self) -> ShutdownTracker { + // Child token inherits cancellation from parent + let child_token = self.root_cancellation_token.child_token(); + + // New TaskTracker for this level's tasks + let child_task_tracker = TaskTracker::new(); + + ShutdownTracker { + root_cancellation_token: child_token, + tracker: child_task_tracker, + } + } + + /// Convenience method to perform a complete shutdown sequence. + /// This method: + /// 1. Signals cancellation to all tasks + /// 2. Closes the tracker to prevent new tasks + /// 3. Waits for all existing tasks to complete + pub async fn shutdown(self) { + // Signal cancellation to all tasks + self.root_cancellation_token.cancel(); + + // Close the tracker to prevent new tasks from being spawned + self.tracker.close(); + + // Wait for all existing tasks to complete + self.tracker.wait().await; + } } diff --git a/common/task/src/lib.rs b/common/task/src/lib.rs index 7be94f6c46..3eaf27cbde 100644 --- a/common/task/src/lib.rs +++ b/common/task/src/lib.rs @@ -5,6 +5,7 @@ pub mod cancellation; pub mod connections; pub mod event; pub mod manager; +pub(crate) mod runtime_registry; #[cfg(not(target_arch = "wasm32"))] pub mod signal; pub mod spawn; @@ -18,3 +19,11 @@ pub use tokio_util::task::TaskTracker; #[cfg(not(target_arch = "wasm32"))] pub use signal::{wait_for_signal, wait_for_signal_and_error}; + +pub use crate::runtime_registry::RegistryAccessError; + +/// Get or create a ShutdownTracker for SDK use. +/// This provides automatic task management without requiring manual setup. +pub fn get_sdk_shutdown_tracker() -> Result { + Ok(runtime_registry::RuntimeRegistry::get_or_create_sdk()?.shutdown_tracker_owned()) +} diff --git a/common/task/src/runtime_registry.rs b/common/task/src/runtime_registry.rs new file mode 100644 index 0000000000..c74fcfde1b --- /dev/null +++ b/common/task/src/runtime_registry.rs @@ -0,0 +1,96 @@ +// Copyright 2025 - Nym Technologies SA +// SPDX-License-Identifier: Apache-2.0 + +use thiserror::Error; + +use crate::ShutdownManager; +use std::sync::RwLock; +use std::sync::{Arc, LazyLock}; + +/// Global registry that manages ShutdownManagers transparently. +/// This allows SDK components to get automatic task management without +/// exposing the complexity to end users. +pub(crate) struct RuntimeRegistry { + // For SDK clients: auto-created manager without signal handling + sdk_manager: RwLock>>, +} + +#[derive(Debug, Error)] +pub enum RegistryAccessError { + #[error("the runtime registry is poisoned")] + Poisoned, +} + +impl RuntimeRegistry { + /// Get or create a ShutdownManager for SDK use. + /// This manager doesn't listen to OS signals, making it suitable for library use. + pub(crate) fn get_or_create_sdk() -> Result, RegistryAccessError> { + let guard = REGISTRY + .sdk_manager + .read() + .map_err(|_| RegistryAccessError::Poisoned)?; + if let Some(manager) = guard.as_ref() { + return Ok(manager.clone()); + } + drop(guard); + + let mut guard = REGISTRY + .sdk_manager + .write() + .map_err(|_| RegistryAccessError::Poisoned)?; + Ok(guard + .get_or_insert_with(|| Arc::new(ShutdownManager::new_without_signals())) + .clone()) + } + + /// Check if an SDK manager has been created. + /// Useful for testing and debugging. + #[allow(dead_code)] + pub(crate) fn has_sdk_manager() -> Result { + Ok(REGISTRY + .sdk_manager + .read() + .map_err(|_| RegistryAccessError::Poisoned)? + .is_some()) + } + + /// Clear the SDK manager. + /// This is primarily for testing to ensure isolation between tests. + #[cfg(test)] + pub(crate) async fn clear() -> Result<(), RegistryAccessError> { + *REGISTRY + .sdk_manager + .write() + .map_err(|_| RegistryAccessError::Poisoned)? = None; + Ok(()) + } +} + +/// Global instance of the runtime registry. +/// Uses LazyLock for on-demand initialization. +static REGISTRY: LazyLock = LazyLock::new(|| RuntimeRegistry { + sdk_manager: RwLock::new(None), +}); + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_get_or_create_sdk() { + // Clear any existing manager + let _ = RuntimeRegistry::clear().await; + + assert!(!RuntimeRegistry::has_sdk_manager().unwrap()); + + let manager1 = RuntimeRegistry::get_or_create_sdk().unwrap(); + assert!(RuntimeRegistry::has_sdk_manager().unwrap()); + + let manager2 = RuntimeRegistry::get_or_create_sdk().unwrap(); + // Should return the same instance + assert!(Arc::ptr_eq(&manager1, &manager2)); + + let _ = RuntimeRegistry::clear().await; + assert!(!RuntimeRegistry::has_sdk_manager().unwrap()); + } +} diff --git a/sdk/rust/nym-sdk/src/error.rs b/sdk/rust/nym-sdk/src/error.rs index c44e40d88d..af7706ec92 100644 --- a/sdk/rust/nym-sdk/src/error.rs +++ b/sdk/rust/nym-sdk/src/error.rs @@ -96,6 +96,9 @@ pub enum Error { #[error(transparent)] Bincode(#[from] bincode::Error), + + #[error("Failed to get shutdown tracker from the task runtime registry: {0}")] + RegistryAccess(#[from] nym_task::RegistryAccessError), } impl Error { diff --git a/sdk/rust/nym-sdk/src/mixnet/client.rs b/sdk/rust/nym-sdk/src/mixnet/client.rs index 1949ef534a..bee2b086e3 100644 --- a/sdk/rust/nym-sdk/src/mixnet/client.rs +++ b/sdk/rust/nym-sdk/src/mixnet/client.rs @@ -37,6 +37,7 @@ use std::path::Path; use std::path::PathBuf; #[cfg(unix)] use std::sync::Arc; +use tokio_util::sync::CancellationToken; use url::Url; use zeroize::Zeroizing; @@ -118,6 +119,11 @@ where ::StorageError: Send + Sync, ::StorageError: Send + Sync, { + pub fn with_shutdown_token(self, token: CancellationToken) -> Self { + let shutdown_tracker = ShutdownTracker::new_from_external_shutdown_token(token.into()); + self.custom_shutdown(shutdown_tracker) + } + /// Creates a client builder with the provided client storage implementation. #[must_use] pub fn new_with_storage(storage: S) -> MixnetClientBuilder { @@ -686,9 +692,15 @@ where base_builder = base_builder.with_topology_provider(topology_provider); } - if let Some(custom_shutdown) = self.custom_shutdown { - base_builder = base_builder.with_shutdown(custom_shutdown) - } + // Use custom shutdown if provided, otherwise get from registry + let shutdown_tracker = match self.custom_shutdown { + Some(custom) => custom, + None => { + // Auto-create from registry for SDK use + nym_task::get_sdk_shutdown_tracker()? + } + }; + base_builder = base_builder.with_shutdown(shutdown_tracker); if let Some(gateway_transceiver) = self.custom_gateway_transceiver { base_builder = base_builder.with_gateway_transceiver(gateway_transceiver); @@ -740,7 +752,7 @@ where let (mut started_client, nym_address) = self.connect_to_mixnet_common().await?; // TODO: more graceful handling here, surely both variants should work... I think? - let Some(task_manager) = started_client.shutdown_handle else { + let Some(tracker) = started_client.shutdown_handle else { return Err(Error::new_unsupported( "connecting with socks5 is currently unsupported with custom shutdown", )); @@ -757,14 +769,14 @@ where client_output, client_state.clone(), nym_address, - task_manager.shutdown_tracker_owned(), + tracker.child_tracker(), packet_type, ); Ok(Socks5MixnetClient { nym_address, client_state, - task_handle: task_manager, + task_handle: tracker, socks5_config, }) } diff --git a/sdk/rust/nym-sdk/src/mixnet/native_client.rs b/sdk/rust/nym-sdk/src/mixnet/native_client.rs index d9c92030f4..4cd7f65890 100644 --- a/sdk/rust/nym-sdk/src/mixnet/native_client.rs +++ b/sdk/rust/nym-sdk/src/mixnet/native_client.rs @@ -17,10 +17,8 @@ use nym_gateway_requests::ClientRequest; use nym_sphinx::addressing::clients::Recipient; use nym_sphinx::{params::PacketType, receiver::ReconstructedMessage}; use nym_statistics_common::clients::{ClientStatsEvents, ClientStatsSender}; -use nym_task::{ - connections::{ConnectionCommandSender, LaneQueueLengths}, - ShutdownManager, -}; +use nym_task::connections::{ConnectionCommandSender, LaneQueueLengths}; +use nym_task::ShutdownTracker; use nym_topology::{NymRouteProvider, NymTopology}; use std::pin::Pin; use std::sync::Arc; @@ -54,7 +52,7 @@ pub struct MixnetClient { pub(crate) stats_events_reporter: ClientStatsSender, /// The task manager that controls all the spawned tasks that the clients uses to do it's job. - pub(crate) shutdown_handle: Option, + pub(crate) shutdown_handle: Option, pub(crate) packet_type: Option, // internal state used for the `Stream` implementation @@ -74,7 +72,7 @@ impl MixnetClient { client_state: ClientState, reconstructed_receiver: ReconstructedMessagesReceiver, stats_events_reporter: ClientStatsSender, - task_handle: Option, + task_handle: Option, packet_type: Option, client_request_sender: ClientRequestSender, forget_me: ForgetMe, @@ -240,8 +238,8 @@ impl MixnetClient { tokio::time::sleep(tokio::time::Duration::from_secs(2)).await; } - if let Some(mut task_manager) = self.shutdown_handle { - task_manager.perform_shutdown().await; + if let Some(tracker) = self.shutdown_handle { + tracker.shutdown().await; } } diff --git a/sdk/rust/nym-sdk/src/mixnet/socks5_client.rs b/sdk/rust/nym-sdk/src/mixnet/socks5_client.rs index cb4e826bd5..0a9ee37c90 100644 --- a/sdk/rust/nym-sdk/src/mixnet/socks5_client.rs +++ b/sdk/rust/nym-sdk/src/mixnet/socks5_client.rs @@ -1,7 +1,8 @@ use nym_client_core::client::base_client::ClientState; use nym_socks5_client_core::config::Socks5; use nym_sphinx::addressing::clients::Recipient; -use nym_task::{connections::LaneQueueLengths, ShutdownManager}; +use nym_task::connections::LaneQueueLengths; +use nym_task::ShutdownTracker; use nym_topology::NymTopology; @@ -18,7 +19,7 @@ pub struct Socks5MixnetClient { pub(crate) client_state: ClientState, /// The task manager that controls all the spawned tasks that the clients uses to do it's job. - pub(crate) task_handle: ShutdownManager, + pub(crate) task_handle: ShutdownTracker, /// SOCKS5 configuration parameters. pub(crate) socks5_config: Socks5, @@ -81,7 +82,7 @@ impl Socks5MixnetClient { /// Disconnect from the mixnet. Currently it is not supported to reconnect a disconnected /// client. - pub async fn disconnect(mut self) { - self.task_handle.run_until_shutdown().await; + pub async fn disconnect(self) { + self.task_handle.shutdown().await; } } diff --git a/wasm/client/src/client.rs b/wasm/client/src/client.rs index 4e70186f73..0a9b0cc053 100644 --- a/wasm/client/src/client.rs +++ b/wasm/client/src/client.rs @@ -28,7 +28,7 @@ use wasm_client_core::helpers::{ add_gateway, generate_new_client_keys, parse_recipient, parse_sender_tag, }; use wasm_client_core::nym_task::connections::TransmissionLane; -use wasm_client_core::nym_task::ShutdownManager; +use wasm_client_core::nym_task::ShutdownTracker; use wasm_client_core::storage::core_client_traits::FullWasmClientStorage; use wasm_client_core::storage::wasm_client_traits::WasmClientStorage; use wasm_client_core::storage::ClientStorage; @@ -59,7 +59,7 @@ pub struct NymClient { // even though we don't use graceful shutdowns, other components rely on existence of this struct // and if it's dropped, everything will start going offline - _task_manager: ShutdownManager, + _task_manager: ShutdownTracker, packet_type: PacketType, } diff --git a/wasm/mix-fetch/src/client.rs b/wasm/mix-fetch/src/client.rs index e462d5835e..b68a6d0db1 100644 --- a/wasm/mix-fetch/src/client.rs +++ b/wasm/mix-fetch/src/client.rs @@ -20,7 +20,7 @@ use wasm_client_core::client::base_client::{BaseClientBuilder, ClientInput, Clie use wasm_client_core::client::inbound_messages::InputMessage; use wasm_client_core::helpers::{add_gateway, generate_new_client_keys}; use wasm_client_core::nym_task::connections::TransmissionLane; -use wasm_client_core::nym_task::ShutdownManager; +use wasm_client_core::nym_task::ShutdownTracker; use wasm_client_core::storage::core_client_traits::FullWasmClientStorage; use wasm_client_core::storage::wasm_client_traits::WasmClientStorage; use wasm_client_core::storage::ClientStorage; @@ -41,7 +41,7 @@ pub struct MixFetchClient { requests: ActiveRequests, // this has to be guarded by a mutex to be able to disconnect with an immutable reference - _shutdown_manager: Mutex, + _shutdown_manager: Mutex, } #[wasm_bindgen] @@ -233,11 +233,11 @@ impl MixFetchClient { self.invalidated.store(true, Ordering::Relaxed); console_log!("sending shutdown signal"); - let mut shutdown_guard = self._shutdown_manager.lock().await; - shutdown_guard.send_cancellation(); - + let shutdown_guard = self._shutdown_manager.lock().await; + shutdown_guard.clone_shutdown_token().cancel(); + shutdown_guard.close_tracker(); console_log!("waiting for shutdown to complete"); - shutdown_guard.run_until_shutdown().await; + shutdown_guard.wait_for_tracker().await; self.requests.invalidate_all().await;