feature: CancellationToken-based shutdowns (#5325)

* initial stub for ShutdownToken

* attempting to start using new ShutdownManager in NymNode

* migrated verloc tasks

* added custom shutdown signal registration

* integrated legacy task support

* migrated additional tasks inside nym-node

* removed import thats unused in wasm

* apply review comments

* windows fixes

nym-node/src/node/mixnet/handler.rs

@@ -38,10 +38,7 @@ impl ConnectionHandler {
         tcp_stream: TcpStream,
         remote_address: SocketAddr,
     ) -> Self {
-        let mut task_client = shared.task_client.fork(remote_address.to_string());
-        // we don't want dropped connections to cause global shutdown
-        task_client.disarm();
-
+        let shutdown = shared.shutdown.child_token(remote_address.to_string());
         shared.metrics.network.new_active_ingress_mixnet_client();
 
         ConnectionHandler {
@@ -51,7 +48,7 @@ impl ConnectionHandler {
                 mixnet_forwarder: shared.mixnet_forwarder.clone(),
                 final_hop: shared.final_hop.clone(),
                 metrics: shared.metrics.clone(),
-                task_client,
+                shutdown,
             },
             remote_address,
             mixnet_connection: Framed::new(tcp_stream, NymCodec),
@@ -165,11 +162,12 @@ impl ConnectionHandler {
         )
     )]
     pub(crate) async fn handle_stream(&mut self) {
-        while !self.shared.task_client.is_shutdown() {
+        loop {
             tokio::select! {
                 biased;
-                _ = self.shared.task_client.recv() => {
+                _ = self.shared.shutdown.cancelled() => {
                     trace!("connection handler: received shutdown");
+                    break
                 }
                 maybe_framed_nym_packet = self.mixnet_connection.next() => {
                     match maybe_framed_nym_packet {
@@ -186,6 +184,7 @@ impl ConnectionHandler {
                 }
             }
         }
+
         debug!("exiting and closing connection");
     }
 }

nym-node/src/node/mixnet/listener.rs

@@ -2,14 +2,14 @@
 // SPDX-License-Identifier: GPL-3.0-only
 
 use crate::node::mixnet::SharedData;
-use nym_task::TaskClient;
+use nym_task::ShutdownToken;
 use std::net::SocketAddr;
 use tokio::task::JoinHandle;
 use tracing::{error, info, trace};
 
 pub(crate) struct Listener {
     bind_address: SocketAddr,
-    shutdown: TaskClient,
+    shutdown: ShutdownToken,
     shared_data: SharedData,
 }
 
@@ -17,7 +17,7 @@ impl Listener {
     pub(crate) fn new(bind_address: SocketAddr, shared_data: SharedData) -> Self {
         Listener {
             bind_address,
-            shutdown: shared_data.task_client.fork("socket-listener"),
+            shutdown: shared_data.shutdown.clone_with_suffix("socket-listener"),
             shared_data,
         }
     }
@@ -29,19 +29,15 @@ impl Listener {
             Ok(listener) => listener,
             Err(err) => {
                 error!("Failed to bind to {}: {err}. Are you sure nothing else is running on the specified port and your user has sufficient permission to bind to the requested address?", self.bind_address);
-
-                // that's a bit gnarly, but we need to make sure we trigger shutdown
-                let mut shutdown_bomb = self.shutdown.fork("shutdown-bomb");
-                shutdown_bomb.rearm();
-                drop(shutdown_bomb);
+                self.shutdown.cancel();
                 return;
             }
         };
 
-        while !self.shutdown.is_shutdown() {
+        loop {
             tokio::select! {
                 biased;
-                _ = self.shutdown.recv() => {
+                _ = self.shutdown.cancelled() => {
                     trace!("mixnet listener: received shutdown");
                 }
                 connection = tcp_listener.accept() => {

nym-node/src/node/mixnet/packet_forwarding.rs

@@ -9,6 +9,7 @@ use nym_mixnet_client::SendWithoutResponse;
 use nym_node_metrics::NymNodeMetrics;
 use nym_nonexhaustive_delayqueue::{Expired, NonExhaustiveDelayQueue};
 use nym_sphinx_forwarding::packet::MixPacket;
+use nym_task::ShutdownToken;
 use std::io;
 use tokio::time::Instant;
 use tracing::{debug, error, trace, warn};
@@ -21,11 +22,11 @@ pub struct PacketForwarder<C> {
 
     packet_sender: MixForwardingSender,
     packet_receiver: MixForwardingReceiver,
-    shutdown: nym_task::TaskClient,
+    shutdown: ShutdownToken,
 }
 
 impl<C> PacketForwarder<C> {
-    pub fn new(client: C, metrics: NymNodeMetrics, shutdown: nym_task::TaskClient) -> Self {
+    pub fn new(client: C, metrics: NymNodeMetrics, shutdown: ShutdownToken) -> Self {
         let (packet_sender, packet_receiver) = mix_forwarding_channels();
 
         PacketForwarder {
@@ -123,7 +124,7 @@ impl<C> PacketForwarder<C> {
         loop {
             tokio::select! {
                 biased;
-                _ = self.shutdown.recv() => {
+                _ = self.shutdown.cancelled() => {
                     debug!("PacketForwarder: Received shutdown");
                     break;
                 }

nym-node/src/node/mixnet/shared/mod.rs

@@ -11,7 +11,7 @@ use nym_node_metrics::NymNodeMetrics;
 use nym_sphinx_forwarding::packet::MixPacket;
 use nym_sphinx_framing::processing::{MixProcessingResult, PacketProcessingError};
 use nym_sphinx_types::DestinationAddressBytes;
-use nym_task::TaskClient;
+use nym_task::ShutdownToken;
 use std::io;
 use std::net::{IpAddr, SocketAddr};
 use std::sync::Arc;
@@ -43,7 +43,6 @@ impl ProcessingConfig {
 }
 
 // explicitly do NOT derive clone as we want to manually apply relevant suffixes to the task clients
-// as well as immediately disarm them
 pub(crate) struct SharedData {
     pub(super) processing_config: ProcessingConfig,
     // TODO: this type is not `Zeroize` : (
@@ -56,7 +55,7 @@ pub(crate) struct SharedData {
     pub(super) final_hop: SharedFinalHopData,
 
     pub(super) metrics: NymNodeMetrics,
-    pub(super) task_client: TaskClient,
+    pub(super) shutdown: ShutdownToken,
 }
 
 impl SharedData {
@@ -66,7 +65,7 @@ impl SharedData {
         mixnet_forwarder: MixForwardingSender,
         final_hop: SharedFinalHopData,
         metrics: NymNodeMetrics,
-        task_client: TaskClient,
+        shutdown: ShutdownToken,
     ) -> Self {
         SharedData {
             processing_config,
@@ -74,7 +73,7 @@ impl SharedData {
             mixnet_forwarder,
             final_hop,
             metrics,
-            task_client,
+            shutdown,
         }
     }
 
@@ -144,13 +143,10 @@ impl SharedData {
             .mixnet_forwarder
             .forward_packet(PacketToForward::new(packet, delay_until))
             .is_err()
-            && !self.task_client.is_shutdown()
+            && !self.shutdown.is_cancelled()
         {
             error!("failed to forward sphinx packet on the channel while the process is not going through the shutdown!");
-            // this is a critical error, we're in uncharted lands, we have to shut down
-            let mut shutdown_bomb = self.task_client.fork("shutdown bomb");
-            shutdown_bomb.rearm();
-            drop(shutdown_bomb)
+            self.shutdown.cancel();
         }
     }