From 102cd1033c414d3fd93f01bbee2b9bbbe021cb5b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C4=99drzej=20Stuczy=C5=84ski?= Date: Mon, 13 Jan 2025 09:13:13 +0000 Subject: [PATCH] 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 --- Cargo.lock | 3 + common/task/Cargo.toml | 3 + common/task/src/cancellation.rs | 366 ++++++++++++++++++ common/task/src/lib.rs | 9 +- common/task/src/manager.rs | 20 + common/verloc/src/measurements/listener.rs | 81 ++-- common/verloc/src/measurements/measurer.rs | 23 +- common/verloc/src/measurements/sender.rs | 15 +- nym-node/src/error.rs | 6 + nym-node/src/node/http/mod.rs | 47 +-- nym-node/src/node/http/router/mod.rs | 6 +- nym-node/src/node/metrics/aggregator.rs | 7 +- nym-node/src/node/metrics/console_logger.rs | 11 +- nym-node/src/node/mixnet/handler.rs | 13 +- nym-node/src/node/mixnet/listener.rs | 16 +- nym-node/src/node/mixnet/packet_forwarding.rs | 7 +- nym-node/src/node/mixnet/shared/mod.rs | 16 +- nym-node/src/node/mod.rs | 57 +-- 18 files changed, 515 insertions(+), 191 deletions(-) create mode 100644 common/task/src/cancellation.rs diff --git a/Cargo.lock b/Cargo.lock index 21988cff4c..41fc2357d0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6659,10 +6659,13 @@ dependencies = [ name = "nym-task" version = "0.1.0" dependencies = [ + "cfg-if", "futures", "log", "thiserror", "tokio", + "tokio-util", + "tracing", "wasm-bindgen", "wasm-bindgen-futures", "wasmtimer", diff --git a/common/task/Cargo.toml b/common/task/Cargo.toml index 6f0f6b7579..6047373c53 100644 --- a/common/task/Cargo.toml +++ b/common/task/Cargo.toml @@ -8,10 +8,13 @@ license.workspace = true repository.workspace = true [dependencies] +cfg-if = { workspace = true } futures = { workspace = true } log = { workspace = true } thiserror = { workspace = true } tokio = { workspace = true, features = ["macros", "sync"] } +tokio-util = { workspace = true, features = ["rt"] } +tracing = { workspace = true } [target."cfg(not(target_arch = \"wasm32\"))".dependencies.tokio] workspace = true diff --git a/common/task/src/cancellation.rs b/common/task/src/cancellation.rs new file mode 100644 index 0000000000..bb7f52a0bf --- /dev/null +++ b/common/task/src/cancellation.rs @@ -0,0 +1,366 @@ +// Copyright 2025 - Nym Technologies SA +// SPDX-License-Identifier: Apache-2.0 + +use crate::{TaskClient, TaskManager}; +use futures::stream::FuturesUnordered; +use futures::StreamExt; +use std::future::Future; +use std::ops::Deref; +use std::pin::Pin; +use std::time::Duration; +use tokio::task::JoinSet; +use tokio::time::sleep; +use tokio_util::sync::{CancellationToken, DropGuard}; +use tokio_util::task::TaskTracker; +use tracing::{debug, info, trace}; + +#[cfg(unix)] +use tokio::signal::unix::{signal, SignalKind}; + +pub const DEFAULT_MAX_SHUTDOWN_DURATION: Duration = Duration::from_secs(5); + +pub fn token_name(name: &Option) -> String { + name.clone().unwrap_or_else(|| "unknown".to_string()) +} + +// a wrapper around tokio's CancellationToken that adds optional `name` information to more easily +// track down sources of shutdown +#[derive(Debug, Default)] +pub struct ShutdownToken { + name: Option, + inner: CancellationToken, +} + +impl Clone for ShutdownToken { + fn clone(&self) -> Self { + // make sure to not accidentally overflow the stack if we keep cloning the handle + let name = if let Some(name) = &self.name { + if name != Self::OVERFLOW_NAME && name.len() < Self::MAX_NAME_LENGTH { + Some(format!("{name}-child")) + } else { + Some(Self::OVERFLOW_NAME.to_string()) + } + } else { + None + }; + + ShutdownToken { + name, + inner: self.inner.clone(), + } + } +} + +impl Deref for ShutdownToken { + type Target = CancellationToken; + + fn deref(&self) -> &Self::Target { + &self.inner + } +} + +impl ShutdownToken { + const MAX_NAME_LENGTH: usize = 128; + const OVERFLOW_NAME: &'static str = "reached maximum ShutdownToken children name depth"; + + pub fn new(name: impl Into) -> Self { + ShutdownToken { + name: Some(name.into()), + inner: CancellationToken::new(), + } + } + + // Creates a ShutdownToken which will get cancelled whenever the current token gets cancelled. + // Unlike a cloned/forked ShutdownToken, cancelling a child token does not cancel the parent token. + #[must_use] + pub fn child_token>(&self, child_suffix: S) -> Self { + let suffix = child_suffix.into(); + let child_name = if let Some(base) = &self.name { + format!("{base}-{suffix}") + } else { + format!("unknown-{suffix}") + }; + + ShutdownToken { + name: Some(child_name), + inner: self.inner.child_token(), + } + } + + // Creates a clone of the ShutdownToken which will get cancelled whenever the current token gets cancelled, and vice versa. + #[must_use] + pub fn clone_with_suffix>(&self, child_suffix: S) -> Self { + let mut child = self.clone(); + let suffix = child_suffix.into(); + let child_name = if let Some(base) = &self.name { + format!("{base}-{suffix}") + } else { + format!("unknown-{suffix}") + }; + + child.name = Some(child_name); + child + } + + // exposed method with the old name for easier migration + // it will eventually be removed so please try to use `.clone_with_suffix` instead + #[must_use] + pub fn fork>(&self, child_suffix: S) -> Self { + self.clone_with_suffix(child_suffix) + } + + // exposed method with the old name for easier migration + // it will eventually be removed so please try to use `.clone().named(name)` instead + #[must_use] + pub fn fork_named>(&self, name: S) -> Self { + self.clone().named(name) + } + + #[must_use] + pub fn named>(mut self, name: S) -> Self { + self.name = Some(name.into()); + self + } + + #[must_use] + pub fn add_suffix>(self, suffix: S) -> Self { + let suffix = suffix.into(); + let name = if let Some(base) = &self.name { + format!("{base}-{suffix}") + } else { + format!("unknown-{suffix}") + }; + self.named(name) + } + + // Returned guard will cancel this token (and all its children) on drop unless disarmed. + pub fn drop_guard(self) -> ShutdownDropGuard { + ShutdownDropGuard { + name: self.name, + inner: self.inner.drop_guard(), + } + } + + pub fn name(&self) -> String { + token_name(&self.name) + } + + pub async fn run_until_cancelled(&self, fut: F) -> Option + where + F: Future, + { + let res = self.inner.run_until_cancelled(fut).await; + trace!("'{}' got cancelled", self.name()); + res + } +} + +pub struct ShutdownDropGuard { + name: Option, + inner: DropGuard, +} + +impl Deref for ShutdownDropGuard { + type Target = DropGuard; + + fn deref(&self) -> &Self::Target { + &self.inner + } +} + +impl ShutdownDropGuard { + pub fn disarm(self) -> ShutdownToken { + ShutdownToken { + name: self.name, + inner: self.inner.disarm(), + } + } + + pub fn name(&self) -> String { + token_name(&self.name) + } +} + +pub struct ShutdownManager { + pub root_token: ShutdownToken, + + legacy_task_manager: Option, + + shutdown_signals: JoinSet<()>, + + // the reason I'm not using a `JoinSet` is because it forces us to use futures with the same `::Output` type + tracker: TaskTracker, + + max_shutdown_duration: Duration, +} + +impl Deref for ShutdownManager { + type Target = TaskTracker; + + fn deref(&self) -> &Self::Target { + &self.tracker + } +} + +impl ShutdownManager { + pub fn new(root_token_name: impl Into) -> Self { + let manager = ShutdownManager { + root_token: ShutdownToken::new(root_token_name), + legacy_task_manager: None, + shutdown_signals: Default::default(), + tracker: Default::default(), + max_shutdown_duration: Default::default(), + }; + + // we need to add an explicit watcher for the cancellation token being cancelled + // so that we could cancel all legacy tasks + let cancel_watcher = manager.root_token.clone(); + manager.with_shutdown(async move { cancel_watcher.cancelled().await }) + } + + pub fn with_legacy_task_manager(mut self) -> Self { + let mut legacy_manager = + TaskManager::default().named(format!("{}-legacy", self.root_token.name())); + let mut legacy_error_rx = legacy_manager.task_return_error_rx(); + let mut legacy_drop_rx = legacy_manager.task_drop_rx(); + + self.legacy_task_manager = Some(legacy_manager); + + // add a task that listens for legacy task clients being dropped to trigger cancellation + self.with_shutdown(async move { + tokio::select! { + _ = legacy_error_rx.recv() => (), + _ = legacy_drop_rx.recv() => (), + } + + info!("received legacy shutdown signal"); + }) + } + + #[cfg(not(target_arch = "wasm32"))] + pub fn with_default_shutdown_signals(self) -> std::io::Result { + cfg_if::cfg_if! { + if #[cfg(unix)] { + self.with_interrupt_signal() + .with_terminate_signal()? + .with_quit_signal() + } else { + Ok(self.with_interrupt_signal()) + } + } + } + + #[must_use] + pub fn with_shutdown(mut self, shutdown: F) -> Self + where + F: Future, + F: Send + 'static, + { + let shutdown_token = self.root_token.clone(); + self.shutdown_signals.spawn(async move { + shutdown.await; + + info!("sending cancellation after receiving shutdown signal"); + shutdown_token.cancel(); + }); + self + } + + #[cfg(unix)] + pub fn with_shutdown_signal(self, signal_kind: SignalKind) -> std::io::Result { + let mut sig = signal(signal_kind)?; + Ok(self.with_shutdown(async move { + sig.recv().await; + })) + } + + #[cfg(not(target_arch = "wasm32"))] + pub fn with_interrupt_signal(self) -> Self { + self.with_shutdown(async move { + let _ = tokio::signal::ctrl_c().await; + }) + } + + #[cfg(unix)] + pub fn with_terminate_signal(self) -> std::io::Result { + self.with_shutdown_signal(SignalKind::terminate()) + } + + #[cfg(unix)] + pub fn with_quit_signal(self) -> std::io::Result { + self.with_shutdown_signal(SignalKind::quit()) + } + + #[must_use] + pub fn with_shutdown_duration(mut self, duration: Duration) -> Self { + self.max_shutdown_duration = duration; + self + } + + pub fn child_token>(&self, child_suffix: S) -> ShutdownToken { + self.root_token.child_token(child_suffix) + } + + pub fn clone_token>(&self, child_suffix: S) -> ShutdownToken { + self.root_token.clone_with_suffix(child_suffix) + } + + #[must_use] + pub fn subscribe_legacy>(&self, child_suffix: S) -> TaskClient { + // alternatively we could have set self.legacy_task_manager = Some(TaskManager::default()); + // on demand if it wasn't unavailable, but then we'd have to use mutable reference + #[allow(clippy::expect_used)] + self.legacy_task_manager + .as_ref() + .expect("did not enable legacy shutdown support") + .subscribe_named(child_suffix) + } + + async fn finish_shutdown(mut self) { + let mut wait_futures = FuturesUnordered::>>>::new(); + + // force shutdown via ctrl-c + wait_futures.push(Box::pin(async move { + #[cfg(not(target_arch = "wasm32"))] + let interrupt_future = tokio::signal::ctrl_c(); + + #[cfg(target_arch = "wasm32")] + let interrupt_future = futures::future::pending::<()>(); + + let _ = interrupt_future.await; + info!("received interrupt - forcing shutdown"); + })); + + // timeout + wait_futures.push(Box::pin(async move { + sleep(self.max_shutdown_duration).await; + info!("timeout reached, forcing shutdown"); + })); + + // graceful + wait_futures.push(Box::pin(async move { + self.tracker.wait().await; + debug!("migrated tasks successfully shutdown"); + if let Some(legacy) = self.legacy_task_manager.as_mut() { + legacy.wait_for_graceful_shutdown().await; + debug!("legacy tasks successfully shutdown"); + } + + info!("all registered tasks successfully shutdown") + })); + + wait_futures.next().await; + } + + pub async fn wait_for_shutdown_signal(mut self) { + self.shutdown_signals.join_next().await; + + if let Some(legacy_manager) = self.legacy_task_manager.as_mut() { + info!("attempting to shutdown legacy tasks"); + let _ = legacy_manager.signal_shutdown(); + } + + info!("waiting for tasks to finish... (press ctrl-c to force)"); + self.finish_shutdown().await; + } +} diff --git a/common/task/src/lib.rs b/common/task/src/lib.rs index 0ed54166d4..edbf0afef4 100644 --- a/common/task/src/lib.rs +++ b/common/task/src/lib.rs @@ -1,6 +1,7 @@ // Copyright 2022 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 +pub mod cancellation; pub mod connections; pub mod event; pub mod manager; @@ -8,9 +9,11 @@ pub mod manager; pub mod signal; pub mod spawn; +pub use cancellation::{ShutdownDropGuard, ShutdownManager, ShutdownToken}; pub use event::{StatusReceiver, StatusSender, TaskStatus, TaskStatusEvent}; pub use manager::{TaskClient, TaskHandle, TaskManager}; -#[cfg(not(target_arch = "wasm32"))] -pub use signal::wait_for_signal_and_error; - pub use spawn::{spawn, spawn_with_report_error}; +pub use tokio_util::task::TaskTracker; + +#[cfg(not(target_arch = "wasm32"))] +pub use signal::{wait_for_signal, wait_for_signal_and_error}; diff --git a/common/task/src/manager.rs b/common/task/src/manager.rs index 7d8ba3a780..5df2012c11 100644 --- a/common/task/src/manager.rs +++ b/common/task/src/manager.rs @@ -185,6 +185,19 @@ impl TaskManager { } } + // used for compatibility with the ShutdownManager + pub(crate) fn task_return_error_rx(&mut self) -> ErrorReceiver { + self.task_return_error_rx + .take() + .expect("unable to get error channel: attempt to wait twice?") + } + + pub(crate) fn task_drop_rx(&mut self) -> ErrorReceiver { + self.task_drop_rx + .take() + .expect("unable to get task drop channel: attempt to wait twice?") + } + pub async fn wait_for_error(&mut self) -> Option { let mut error_rx = self .task_return_error_rx @@ -208,6 +221,13 @@ impl TaskManager { } } + pub(crate) async fn wait_for_graceful_shutdown(&mut self) { + if let Some(notify_rx) = self.notify_rx.take() { + drop(notify_rx); + } + self.notify_tx.closed().await + } + pub async fn wait_for_shutdown(&mut self) { log::debug!("Waiting for shutdown"); if let Some(notify_rx) = self.notify_rx.take() { diff --git a/common/verloc/src/measurements/listener.rs b/common/verloc/src/measurements/listener.rs index 77c1739068..09cf483179 100644 --- a/common/verloc/src/measurements/listener.rs +++ b/common/verloc/src/measurements/listener.rs @@ -6,7 +6,7 @@ use crate::measurements::packet::{EchoPacket, ReplyPacket}; use bytes::{BufMut, BytesMut}; use futures::StreamExt; use nym_crypto::asymmetric::identity; -use nym_task::TaskClient; +use nym_task::ShutdownToken; use std::net::SocketAddr; use std::sync::Arc; use std::{io, process}; @@ -19,19 +19,19 @@ use tracing::{debug, error, info, trace, warn}; pub struct PacketListener { address: SocketAddr, connection_handler: Arc, - shutdown: TaskClient, + shutdown_token: ShutdownToken, } impl PacketListener { pub fn new( address: SocketAddr, identity: Arc, - shutdown: TaskClient, + shutdown_token: ShutdownToken, ) -> Self { PacketListener { address, connection_handler: Arc::new(ConnectionHandler { identity }), - shutdown, + shutdown_token, } } } @@ -51,26 +51,22 @@ impl PacketListener { info!("Started listening for echo packets on {}", self.address); - let mut shutdown_listener = self.shutdown.clone(); - - while !shutdown_listener.is_shutdown() { + while !self.shutdown_token.is_cancelled() { // cloning the arc as each accepted socket is handled in separate task let connection_handler = Arc::clone(&self.connection_handler); - let mut handler_shutdown_listener = self.shutdown.clone(); - handler_shutdown_listener.disarm(); tokio::select! { socket = listener.accept() => { match socket { Ok((socket, remote_addr)) => { - debug!("New verloc connection from {}", remote_addr); - - tokio::spawn(connection_handler.handle_connection(socket, remote_addr, handler_shutdown_listener)); + debug!("New verloc connection from {remote_addr}"); + let cancel = self.shutdown_token.child_token(format!("handler_{remote_addr}")); + tokio::spawn(async move { cancel.run_until_cancelled(connection_handler.handle_connection(socket, remote_addr)).await }); } Err(err) => warn!("Failed to accept incoming connection - {err}"), } }, - _ = shutdown_listener.recv() => { + _ = self.shutdown_token.cancelled() => { trace!("PacketListener: Received shutdown"); } } @@ -88,50 +84,29 @@ impl ConnectionHandler { packet.construct_reply(self.identity.private_key()) } - pub(crate) async fn handle_connection( - self: Arc, - conn: TcpStream, - remote: SocketAddr, - mut shutdown_listener: TaskClient, - ) { - debug!("Starting connection handler for {:?}", remote); + pub(crate) async fn handle_connection(self: Arc, conn: TcpStream, remote: SocketAddr) { + debug!("Starting connection handler for {remote}"); let mut framed_conn = Framed::new(conn, EchoPacketCodec); - while !shutdown_listener.is_shutdown() { - tokio::select! { - biased; - _ = shutdown_listener.recv() => { - trace!("ConnectionHandler: Shutdown received"); + while let Some(echo_packet) = framed_conn.next().await { + let reply_packet = match echo_packet { + Ok(echo_packet) => self.handle_echo_packet(echo_packet), + Err(err) => { + debug!( + "The socket connection got corrupted with error: {err}. Closing the socket" + ); + return; } - maybe_echo_packet = framed_conn.next() => { - // handle echo packet - let reply_packet = match maybe_echo_packet { - Some(Ok(echo_packet)) => self.handle_echo_packet(echo_packet), - Some(Err(err)) => { - debug!( - "The socket connection got corrupted with error: {err}. Closing the socket", - ); - return; - } - None => { - debug!("The socket connection got terminated by the remote!"); - return; - } - }; + }; - // write back the reply (note the lack of framing) - if let Err(err) = framed_conn - .get_mut() - .write_all(reply_packet.to_bytes().as_ref()) - .await - { - debug!( - "Failed to write reply packet back to the sender - {}. Closing the socket on our end", - err - ); - return; - } - }, + // write back the reply (note the lack of framing) + if let Err(err) = framed_conn + .get_mut() + .write_all(reply_packet.to_bytes().as_ref()) + .await + { + debug!("Failed to write reply packet back to the sender: {err}. Closing the socket on our end"); + return; } } } diff --git a/common/verloc/src/measurements/measurer.rs b/common/verloc/src/measurements/measurer.rs index e56f35d086..573a2f0258 100644 --- a/common/verloc/src/measurements/measurer.rs +++ b/common/verloc/src/measurements/measurer.rs @@ -8,7 +8,7 @@ use crate::models::VerlocNodeResult; use futures::stream::FuturesUnordered; use futures::StreamExt; use nym_crypto::asymmetric::identity; -use nym_task::TaskClient; +use nym_task::ShutdownToken; use nym_validator_client::models::NymNodeDescription; use nym_validator_client::NymApiClient; use rand::prelude::SliceRandom; @@ -23,7 +23,7 @@ pub struct VerlocMeasurer { config: Config, packet_sender: Arc, packet_listener: Arc, - shutdown_listener: TaskClient, + shutdown_token: ShutdownToken, state: SharedVerlocStats, } @@ -31,7 +31,7 @@ impl VerlocMeasurer { pub fn new( config: Config, identity: Arc, - shutdown_listener: TaskClient, + shutdown_token: ShutdownToken, ) -> Self { VerlocMeasurer { packet_sender: Arc::new(PacketSender::new( @@ -40,14 +40,14 @@ impl VerlocMeasurer { config.packet_timeout, config.connection_timeout, config.delay_between_packets, - shutdown_listener.clone().named("VerlocPacketSender"), + shutdown_token.clone_with_suffix("packet_sender"), )), packet_listener: Arc::new(PacketListener::new( config.listening_address, Arc::clone(&identity), - shutdown_listener.clone().named("VerlocPacketListener"), + shutdown_token.clone_with_suffix("packet_listener"), )), - shutdown_listener, + shutdown_token, config, state: SharedVerlocStats::default(), } @@ -69,9 +69,6 @@ impl VerlocMeasurer { return MeasurementOutcome::Done; } - let mut shutdown_listener = self.shutdown_listener.clone().named("VerlocMeasurement"); - shutdown_listener.disarm(); - for chunk in nodes_to_test.chunks(self.config.tested_nodes_batch_size) { let mut chunk_results = Vec::with_capacity(chunk.len()); @@ -95,7 +92,7 @@ impl VerlocMeasurer { .collect::>(); // exhaust the results - while !shutdown_listener.is_shutdown() { + while !self.shutdown_token.is_cancelled() { tokio::select! { measurement_result = measurement_chunk.next() => { let Some(result) = measurement_result else { @@ -120,7 +117,7 @@ impl VerlocMeasurer { }; chunk_results.push(VerlocNodeResult::new(identity, measurement_result)); }, - _ = shutdown_listener.recv() => { + _ = self.shutdown_token.cancelled() => { trace!("Shutdown received while measuring"); return MeasurementOutcome::Shutdown; } @@ -155,7 +152,7 @@ impl VerlocMeasurer { pub async fn run(&mut self) { self.start_listening(); - while !self.shutdown_listener.is_shutdown() { + while !self.shutdown_token.is_cancelled() { info!("Starting verloc measurements"); // TODO: should we also measure gateways? @@ -209,7 +206,7 @@ impl VerlocMeasurer { tokio::select! { _ = sleep(self.config.testing_interval) => {}, - _ = self.shutdown_listener.recv() => { + _ = self.shutdown_token.cancelled() => { trace!("Shutdown received while sleeping"); } } diff --git a/common/verloc/src/measurements/sender.rs b/common/verloc/src/measurements/sender.rs index 292cd8d8c0..3c8b6e494d 100644 --- a/common/verloc/src/measurements/sender.rs +++ b/common/verloc/src/measurements/sender.rs @@ -5,7 +5,7 @@ use crate::error::VerlocError; use crate::measurements::packet::{EchoPacket, ReplyPacket}; use crate::models::VerlocMeasurement; use nym_crypto::asymmetric::ed25519; -use nym_task::TaskClient; +use nym_task::ShutdownToken; use rand::{thread_rng, Rng}; use std::net::SocketAddr; use std::sync::Arc; @@ -45,7 +45,7 @@ pub struct PacketSender { packet_timeout: Duration, connection_timeout: Duration, delay_between_packets: Duration, - shutdown_listener: TaskClient, + shutdown_token: ShutdownToken, } impl PacketSender { @@ -55,7 +55,7 @@ impl PacketSender { packet_timeout: Duration, connection_timeout: Duration, delay_between_packets: Duration, - shutdown_listener: TaskClient, + shutdown_token: ShutdownToken, ) -> Self { PacketSender { identity, @@ -63,7 +63,7 @@ impl PacketSender { packet_timeout, connection_timeout, delay_between_packets, - shutdown_listener, + shutdown_token, } } @@ -83,9 +83,6 @@ impl PacketSender { self: Arc, tested_node: TestedNode, ) -> Result { - let mut shutdown_listener = self.shutdown_listener.fork(tested_node.address.to_string()); - shutdown_listener.disarm(); - let mut conn = match tokio::time::timeout( self.connection_timeout, TcpStream::connect(tested_node.address), @@ -148,7 +145,7 @@ impl PacketSender { Ok(Ok(_)) => {} } }, - _ = shutdown_listener.recv() => { + _ = self.shutdown_token.cancelled() => { trace!("PacketSender: Received shutdown while sending"); return Err(VerlocError::ShutdownReceived); }, @@ -190,7 +187,7 @@ impl PacketSender { } } }, - _ = shutdown_listener.recv() => { + _ = self.shutdown_token.cancelled() => { trace!("PacketSender: Received shutdown while waiting for reply"); return Err(VerlocError::ShutdownReceived); } diff --git a/nym-node/src/error.rs b/nym-node/src/error.rs index 4794b12961..3c50451e07 100644 --- a/nym-node/src/error.rs +++ b/nym-node/src/error.rs @@ -50,6 +50,12 @@ pub enum NymNodeError { #[error("this binary version no longer supports migration from legacy mixnodes and gateways")] UnsupportedMigration, + #[error("failed to initialise shutdown signals: {source}")] + ShutdownSignalFailure { + #[source] + source: io::Error, + }, + #[error("could not find an existing config file at '{}' and fresh node initialisation has been disabled", config_path.display())] ForbiddenInitialisation { config_path: PathBuf }, diff --git a/nym-node/src/node/http/mod.rs b/nym-node/src/node/http/mod.rs index f05ec0355a..3b7f8d9b45 100644 --- a/nym-node/src/node/http/mod.rs +++ b/nym-node/src/node/http/mod.rs @@ -6,9 +6,7 @@ use axum::extract::ConnectInfo; use axum::middleware::AddExtension; use axum::serve::Serve; use axum::Router; -use nym_task::TaskClient; use std::net::SocketAddr; -use tracing::{debug, error}; pub use router::{api, HttpServerConfig, NymNodeRouter}; @@ -19,47 +17,4 @@ pub mod state; type InnerService = IntoMakeServiceWithConnectInfo; type ConnectInfoExt = AddExtension>; -pub type ServeService = Serve; - -pub struct NymNodeHttpServer { - task_client: Option, - inner: ServeService, -} - -impl NymNodeHttpServer { - pub(crate) fn new(inner: ServeService) -> Self { - NymNodeHttpServer { - task_client: None, - inner, - } - } - - #[must_use] - pub fn with_task_client(mut self, task_client: TaskClient) -> Self { - self.task_client = Some(task_client); - self - } - - async fn run_server_forever(server: ServeService) { - if let Err(err) = server.await { - error!("the HTTP server has terminated with the error: {err}"); - } else { - error!("the HTTP server has terminated with producing any errors"); - } - } - - pub async fn run(self) { - if let Some(mut task_client) = self.task_client { - tokio::select! { - _ = task_client.recv_with_delay() => { - debug!("NymNodeHTTPServer: Received shutdown"); - } - _ = Self::run_server_forever(self.inner) => { } - } - } else { - Self::run_server_forever(self.inner).await - } - - debug!("NymNodeHTTPServer: Exiting"); - } -} +pub type NymNodeHttpServer = Serve; diff --git a/nym-node/src/node/http/router/mod.rs b/nym-node/src/node/http/router/mod.rs index 3d40e9a67e..9750b1c0c1 100644 --- a/nym-node/src/node/http/router/mod.rs +++ b/nym-node/src/node/http/router/mod.rs @@ -175,12 +175,10 @@ impl NymNodeRouter { source, })?; - let axum_server = axum::serve( + Ok(axum::serve( listener, self.inner .into_make_service_with_connect_info::(), - ); - - Ok(NymNodeHttpServer::new(axum_server)) + )) } } diff --git a/nym-node/src/node/metrics/aggregator.rs b/nym-node/src/node/metrics/aggregator.rs index 8f76b14edb..3b3cc17eb3 100644 --- a/nym-node/src/node/metrics/aggregator.rs +++ b/nym-node/src/node/metrics/aggregator.rs @@ -6,6 +6,7 @@ use futures::StreamExt; use nym_node_metrics::events::{ events_channels, MetricEventsReceiver, MetricEventsSender, MetricsEvent, }; +use nym_task::ShutdownToken; use std::any; use std::any::TypeId; use std::collections::HashMap; @@ -24,11 +25,11 @@ pub(crate) struct MetricsAggregator { // registered_handlers: HashMap>, event_sender: MetricEventsSender, event_receiver: MetricEventsReceiver, - shutdown: nym_task::TaskClient, + shutdown: ShutdownToken, } impl MetricsAggregator { - pub fn new(handlers_update_interval: Duration, shutdown: nym_task::TaskClient) -> Self { + pub fn new(handlers_update_interval: Duration, shutdown: ShutdownToken) -> Self { let (event_sender, event_receiver) = events_channels(); MetricsAggregator { @@ -113,7 +114,7 @@ impl MetricsAggregator { loop { tokio::select! { biased; - _ = self.shutdown.recv() => { + _ = self.shutdown.cancelled() => { debug!("MetricsAggregator: Received shutdown"); break; } diff --git a/nym-node/src/node/metrics/console_logger.rs b/nym-node/src/node/metrics/console_logger.rs index 37ad46b7f9..c7db72960b 100644 --- a/nym-node/src/node/metrics/console_logger.rs +++ b/nym-node/src/node/metrics/console_logger.rs @@ -4,7 +4,7 @@ use human_repr::HumanCount; use human_repr::HumanThroughput; use nym_node_metrics::NymNodeMetrics; -use nym_task::TaskClient; +use nym_task::ShutdownToken; use std::time::Duration; use time::OffsetDateTime; use tokio::task::JoinHandle; @@ -49,14 +49,14 @@ pub(crate) struct ConsoleLogger { logging_delay: Duration, at_last_update: AtLastUpdate, metrics: NymNodeMetrics, - shutdown: TaskClient, + shutdown: ShutdownToken, } impl ConsoleLogger { pub(crate) fn new( logging_delay: Duration, metrics: NymNodeMetrics, - shutdown: TaskClient, + shutdown: ShutdownToken, ) -> Self { ConsoleLogger { logging_delay, @@ -126,11 +126,12 @@ impl ConsoleLogger { async fn run(&mut self) { trace!("Starting ConsoleLogger"); let mut interval = interval_at(Instant::now() + self.logging_delay, self.logging_delay); - while !self.shutdown.is_shutdown() { + loop { tokio::select! { biased; - _ = self.shutdown.recv() => { + _ = self.shutdown.cancelled() => { trace!("ConsoleLogger: Received shutdown"); + break } _ = interval.tick() => self.log_running_stats().await, }; diff --git a/nym-node/src/node/mixnet/handler.rs b/nym-node/src/node/mixnet/handler.rs index e645636c96..3c634d7c9c 100644 --- a/nym-node/src/node/mixnet/handler.rs +++ b/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"); } } diff --git a/nym-node/src/node/mixnet/listener.rs b/nym-node/src/node/mixnet/listener.rs index 34ae2918fe..738f591c97 100644 --- a/nym-node/src/node/mixnet/listener.rs +++ b/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() => { diff --git a/nym-node/src/node/mixnet/packet_forwarding.rs b/nym-node/src/node/mixnet/packet_forwarding.rs index 5e70841c86..1db0f72652 100644 --- a/nym-node/src/node/mixnet/packet_forwarding.rs +++ b/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 { packet_sender: MixForwardingSender, packet_receiver: MixForwardingReceiver, - shutdown: nym_task::TaskClient, + shutdown: ShutdownToken, } impl PacketForwarder { - 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 PacketForwarder { loop { tokio::select! { biased; - _ = self.shutdown.recv() => { + _ = self.shutdown.cancelled() => { debug!("PacketForwarder: Received shutdown"); break; } diff --git a/nym-node/src/node/mixnet/shared/mod.rs b/nym-node/src/node/mixnet/shared/mod.rs index 9387704a53..8d9f68f528 100644 --- a/nym-node/src/node/mixnet/shared/mod.rs +++ b/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(); } } diff --git a/nym-node/src/node/mod.rs b/nym-node/src/node/mod.rs index 23fdd2adab..8b8eb75212 100644 --- a/nym-node/src/node/mod.rs +++ b/nym-node/src/node/mod.rs @@ -43,7 +43,7 @@ use nym_node_metrics::NymNodeMetrics; use nym_node_requests::api::v1::node::models::{AnnouncePorts, NodeDescription}; use nym_sphinx_acknowledgements::AckKey; use nym_sphinx_addressing::Recipient; -use nym_task::{TaskClient, TaskManager}; +use nym_task::{ShutdownManager, ShutdownToken, TaskClient}; use nym_validator_client::client::NymApiClientExt; use nym_validator_client::models::NodeRefreshBody; use nym_validator_client::{NymApiClient, UserAgent}; @@ -351,6 +351,7 @@ impl From for nym_wireguard::WireguardData { pub(crate) struct NymNode { config: Config, accepted_operator_terms_and_conditions: bool, + shutdown_manager: ShutdownManager, description: NodeDescription, @@ -447,6 +448,10 @@ impl NymNode { wireguard: Some(wireguard_data), config, accepted_operator_terms_and_conditions: false, + shutdown_manager: ShutdownManager::new("NymNode") + .with_legacy_task_manager() + .with_default_shutdown_signals() + .map_err(|source| NymNodeError::ShutdownSignalFailure { source })?, }) } @@ -823,7 +828,7 @@ impl NymNode { } } - pub(crate) fn start_verloc_measurements(&self, shutdown: TaskClient) { + pub(crate) fn start_verloc_measurements(&self) { info!( "Starting the [verloc] round-trip-time measurer on {} ...", self.config.verloc.bind_address @@ -845,8 +850,11 @@ impl NymNode { .retry_timeout(self.config.verloc.debug.retry_timeout) .build(); - let mut verloc_measurer = - VerlocMeasurer::new(config, self.ed25519_identity_keys.clone(), shutdown); + let mut verloc_measurer = VerlocMeasurer::new( + config, + self.ed25519_identity_keys.clone(), + self.shutdown_manager.clone_token("verloc"), + ); verloc_measurer.set_shared_state(self.verloc_stats.clone()); tokio::spawn(async move { verloc_measurer.run().await }); } @@ -855,14 +863,14 @@ impl NymNode { &self, active_clients_store: ActiveClientsStore, active_egress_mixnet_connections: ActiveConnections, - shutdown: TaskClient, + shutdown: ShutdownToken, ) -> MetricEventsSender { info!("setting up node metrics..."); // aggregator (to listen for any metrics events) let mut metrics_aggregator = MetricsAggregator::new( self.config.metrics.debug.aggregator_update_rate, - shutdown.fork("aggregator"), + shutdown.clone_with_suffix("aggregator"), ); // >>>> START: register all relevant handlers for custom events @@ -924,12 +932,9 @@ impl NymNode { ConsoleLogger::new( self.config.metrics.debug.console_logging_update_interval, self.metrics.clone(), - shutdown.named("metrics-console-logger"), + shutdown.clone_with_suffix("metrics-console-logger"), ) .start(); - } else { - let mut shutdown = shutdown; - shutdown.disarm() } let events_sender = metrics_aggregator.sender(); @@ -943,7 +948,7 @@ impl NymNode { pub(crate) fn start_mixnet_listener( &self, active_clients_store: &ActiveClientsStore, - shutdown: TaskClient, + shutdown: ShutdownToken, ) -> (MixForwardingSender, ActiveConnections) { let processing_config = ProcessingConfig::new(&self.config); @@ -972,7 +977,7 @@ impl NymNode { let mut packet_forwarder = PacketForwarder::new( mixnet_client, self.metrics.clone(), - shutdown.fork("mix-packet-forwarder"), + shutdown.clone_with_suffix("mix-packet-forwarder"), ); let mix_packet_sender = packet_forwarder.sender(); tokio::spawn(async move { packet_forwarder.run().await }); @@ -1005,45 +1010,47 @@ impl NymNode { ); debug!("config: {:#?}", self.config); - let mut task_manager = TaskManager::default().named("NymNode"); - let http_server = self - .build_http_server() - .await? - .with_task_client(task_manager.subscribe_named("http-server")); + let http_server = self.build_http_server().await?; let bind_address = self.config.http.bind_address; - tokio::spawn(async move { + let server_shutdown = self.shutdown_manager.clone_token("http-server"); + + self.shutdown_manager.spawn(async move { { - info!("started NymNodeHTTPServer on {bind_address}"); - http_server.run().await + info!("starting NymNodeHTTPServer on {bind_address}"); + http_server + .with_graceful_shutdown(async move { server_shutdown.cancelled().await }) + .await } }); self.try_refresh_remote_nym_api_cache().await; - self.start_verloc_measurements(task_manager.subscribe_named("verloc-measurements")); + self.start_verloc_measurements(); let active_clients_store = ActiveClientsStore::new(); let (mix_packet_sender, active_egress_mixnet_connections) = self.start_mixnet_listener( &active_clients_store, - task_manager.subscribe_named("mixnet-traffic"), + self.shutdown_manager.clone_token("mixnet-traffic"), ); let metrics_sender = self.setup_metrics_backend( active_clients_store.clone(), active_egress_mixnet_connections, - task_manager.subscribe_named("metrics"), + self.shutdown_manager.clone_token("metrics"), ); self.start_gateway_tasks( metrics_sender, active_clients_store, mix_packet_sender, - task_manager.subscribe_named("gateway-tasks"), + self.shutdown_manager.subscribe_legacy("gateway-tasks"), ) .await?; - let _ = task_manager.catch_interrupt().await; + self.shutdown_manager.close(); + self.shutdown_manager.wait_for_shutdown_signal().await; + Ok(()) } }