From eb07ec85805f83989647b369ffada6fc2b5fb91b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jon=20H=C3=A4ggblad?= Date: Wed, 14 Dec 2022 17:13:00 +0100 Subject: [PATCH] client: sort out shutdown procedure and harmonize with socks5-client (#2695) * common/task: rename ShutdownNotifier to TaskManager * nym-client: return boxed error * nym-client: enable graceful shutdown * nym-client: task wait on shutdown to instead exit on closed channel * Fix build * Fix unused * changelog: update --- CHANGELOG.md | 1 + .../client-core/src/client/base_client/mod.rs | 48 +++++++------- .../src/client/cover_traffic_stream.rs | 2 +- clients/client-core/src/client/mix_traffic.rs | 7 +- .../acknowledgement_listener.rs | 4 +- .../action_controller.rs | 4 +- .../input_message_listener.rs | 4 +- .../acknowledgement_control/mod.rs | 2 +- .../retransmission_request_listener.rs | 4 +- .../sent_notification_listener.rs | 4 +- .../src/client/real_messages_control/mod.rs | 2 +- .../real_traffic_stream.rs | 6 +- .../sending_delay_controller.rs | 2 + .../client-core/src/client/received_buffer.rs | 28 ++++---- .../src/client/replies/reply_controller.rs | 4 +- .../src/client/replies/reply_storage/mod.rs | 2 +- .../src/client/topology_control.rs | 2 +- clients/native/src/client/mod.rs | 41 +++++++----- clients/native/src/commands/mod.rs | 5 +- clients/native/src/commands/run.rs | 8 ++- clients/native/src/main.rs | 5 +- clients/native/src/websocket/handler.rs | 66 ++++++++++++------- clients/native/src/websocket/listener.rs | 25 +++++-- clients/socks5/src/client/mod.rs | 16 +++-- clients/socks5/src/socks/client.rs | 6 +- clients/socks5/src/socks/mixnet_responses.rs | 6 +- clients/socks5/src/socks/server.rs | 6 +- clients/webassembly/src/client/mod.rs | 6 +- .../client-libs/gateway-client/src/client.rs | 8 +-- .../gateway-client/src/packet_router.rs | 6 +- .../gateway-client/src/socket_state.rs | 4 +- common/mixnode-common/src/verloc/listener.rs | 8 +-- common/mixnode-common/src/verloc/mod.rs | 6 +- common/mixnode-common/src/verloc/sender.rs | 6 +- .../src/connection_controller.rs | 6 +- .../proxy-helpers/src/proxy_runner/inbound.rs | 4 +- .../proxy-helpers/src/proxy_runner/mod.rs | 6 +- .../src/proxy_runner/outbound.rs | 4 +- common/task/src/lib.rs | 5 +- common/task/src/{shutdown.rs => manager.rs} | 64 ++++++++++-------- common/task/src/signal.rs | 6 +- common/task/src/spawn.rs | 4 +- .../src/country_statistics/distribution.rs | 6 +- .../src/country_statistics/geolocate.rs | 6 +- explorer-api/src/main.rs | 6 +- explorer-api/src/tasks.rs | 6 +- .../node/listener/connection_handler/mod.rs | 4 +- mixnode/src/node/listener/mod.rs | 6 +- mixnode/src/node/mod.rs | 14 ++-- mixnode/src/node/node_statistics.rs | 20 +++--- mixnode/src/node/packet_delayforwarder.rs | 10 +-- nym-api/src/coconut/dkg/controller.rs | 4 +- nym-api/src/contract_cache/mod.rs | 4 +- nym-api/src/epoch_operations/mod.rs | 9 +-- nym-api/src/main.rs | 6 +- nym-api/src/network_monitor/mod.rs | 4 +- .../monitor/gateways_pinger.rs | 4 +- nym-api/src/network_monitor/monitor/mod.rs | 4 +- .../src/network_monitor/monitor/receiver.rs | 4 +- nym-api/src/network_monitor/monitor/sender.rs | 10 +-- nym-api/src/node_status_api/cache.rs | 4 +- nym-api/src/node_status_api/uptime_updater.rs | 4 +- .../network-requester/src/connection.rs | 4 +- .../network-requester/src/core.rs | 10 +-- 64 files changed, 325 insertions(+), 277 deletions(-) rename common/task/src/{shutdown.rs => manager.rs} (89%) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0b66b2305c..a40671598e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,7 @@ Post 1.0.0 release, the changelog format is based on [Keep a Changelog](https:// ### Changed - all-binaries: improved error logging ([#2686]) +- native client: bring shutdown logic up to the same level as socks5-client [#2686]: https://github.com/nymtech/nym/pull/2686 diff --git a/clients/client-core/src/client/base_client/mod.rs b/clients/client-core/src/client/base_client/mod.rs index 6db0280019..807170cc84 100644 --- a/clients/client-core/src/client/base_client/mod.rs +++ b/clients/client-core/src/client/base_client/mod.rs @@ -36,19 +36,19 @@ use nymsphinx::addressing::nodes::NodeIdentity; use std::sync::Arc; use std::time::Duration; use tap::TapFallible; -use task::{ShutdownListener, ShutdownNotifier}; +use task::{TaskClient, TaskManager}; use url::Url; #[cfg(all(not(target_arch = "wasm32"), feature = "fs-surb-storage"))] pub mod non_wasm_helpers; pub struct ClientInput { - pub shared_lane_queue_lengths: LaneQueueLengths, pub connection_command_sender: ConnectionCommandSender, pub input_sender: InputMessageSender, } pub struct ClientOutput { + pub shared_lane_queue_lengths: LaneQueueLengths, pub received_buffer_request_sender: ReceivedBufferRequestSender, } @@ -127,8 +127,8 @@ where debug_config, disabled_credentials, nym_api_endpoints, - bandwidth_controller, reply_storage_backend, + bandwidth_controller, key_manager, } } @@ -151,7 +151,7 @@ where self_address: Recipient, topology_accessor: TopologyAccessor, mix_tx: BatchMixMessageSender, - shutdown: ShutdownListener, + shutdown: TaskClient, ) { info!("Starting loop cover traffic stream..."); @@ -185,7 +185,7 @@ where reply_controller_receiver: ReplyControllerReceiver, lane_queue_lengths: LaneQueueLengths, client_connection_rx: ConnectionCommandReceiver, - shutdown: ShutdownListener, + shutdown: TaskClient, ) { info!("Starting real traffic stream..."); @@ -212,7 +212,7 @@ where mixnet_receiver: MixnetMessageReceiver, reply_key_storage: SentReplyKeys, reply_controller_sender: ReplyControllerSender, - shutdown: ShutdownListener, + shutdown: TaskClient, ) { info!("Starting received messages buffer controller..."); ReceivedMessagesBufferController::new( @@ -229,7 +229,7 @@ where &mut self, mixnet_message_sender: MixnetMessageSender, ack_sender: AcknowledgementSender, - shutdown: ShutdownListener, + shutdown: TaskClient, ) -> Result> { let gateway_id = self.gateway_config.gateway_id.clone(); if gateway_id.is_empty() { @@ -284,7 +284,7 @@ where nym_api_urls: Vec, refresh_rate: Duration, topology_accessor: TopologyAccessor, - shutdown: ShutdownListener, + shutdown: TaskClient, ) -> Result<(), ClientCoreError> { let topology_refresher_config = TopologyRefresherConfig::new( nym_api_urls, @@ -317,7 +317,7 @@ where // requests? fn start_mix_traffic_controller( gateway_client: GatewayClient, - shutdown: ShutdownListener, + shutdown: TaskClient, ) -> BatchMixMessageSender { info!("Starting mix traffic controller..."); let (mix_traffic_controller, mix_tx) = MixTrafficController::new(gateway_client); @@ -327,7 +327,7 @@ where async fn setup_persistent_reply_storage( backend: B, - shutdown: ShutdownListener, + shutdown: TaskClient, ) -> Result> { let persistent_storage = PersistentReplyStorage::new(backend); let mem_store = persistent_storage @@ -367,7 +367,7 @@ where let shared_topology_accessor = TopologyAccessor::new(); // Shutdown notifier for signalling tasks to stop - let shutdown = ShutdownNotifier::default(); + let task_manager = TaskManager::default(); // channels responsible for dealing with reply-related fun let (reply_controller_sender, reply_controller_receiver) = @@ -378,18 +378,20 @@ where // the components are started in very specific order. Unless you know what you are doing, // do not change that. let gateway_client = self - .start_gateway_client(mixnet_messages_sender, ack_sender, shutdown.subscribe()) + .start_gateway_client(mixnet_messages_sender, ack_sender, task_manager.subscribe()) .await?; - let reply_storage = - Self::setup_persistent_reply_storage(self.reply_storage_backend, shutdown.subscribe()) - .await?; + let reply_storage = Self::setup_persistent_reply_storage( + self.reply_storage_backend, + task_manager.subscribe(), + ) + .await?; Self::start_topology_refresher( self.nym_api_endpoints.clone(), self.debug_config.topology_refresh_rate, shared_topology_accessor.clone(), - shutdown.subscribe(), + task_manager.subscribe(), ) .await?; @@ -399,7 +401,7 @@ where mixnet_messages_receiver, reply_storage.key_storage(), reply_controller_sender.clone(), - shutdown.subscribe(), + task_manager.subscribe(), ); // The sphinx_message_sender is the transmitter for any component generating sphinx packets @@ -407,7 +409,7 @@ where // traffic stream. // The MixTrafficController then sends the actual traffic let sphinx_message_sender = - Self::start_mix_traffic_controller(gateway_client, shutdown.subscribe()); + Self::start_mix_traffic_controller(gateway_client, task_manager.subscribe()); // Channels that the websocket listener can use to signal downstream to the real traffic // controller that connections are closed. @@ -439,7 +441,7 @@ where reply_controller_receiver, shared_lane_queue_lengths.clone(), client_connection_rx, - shutdown.subscribe(), + task_manager.subscribe(), ); if !self.debug_config.disable_loop_cover_traffic_stream { @@ -449,7 +451,7 @@ where self_address, shared_topology_accessor, sphinx_message_sender, - shutdown.subscribe(), + task_manager.subscribe(), ); } @@ -459,17 +461,17 @@ where Ok(BaseClient { client_input: ClientInputStatus::AwaitingProducer { client_input: ClientInput { - shared_lane_queue_lengths, connection_command_sender: client_connection_tx, input_sender, }, }, client_output: ClientOutputStatus::AwaitingConsumer { client_output: ClientOutput { + shared_lane_queue_lengths, received_buffer_request_sender, }, }, - shutdown_notifier: shutdown, + task_manager, }) } } @@ -478,5 +480,5 @@ pub struct BaseClient { pub client_input: ClientInputStatus, pub client_output: ClientOutputStatus, - pub shutdown_notifier: ShutdownNotifier, + pub task_manager: TaskManager, } diff --git a/clients/client-core/src/client/cover_traffic_stream.rs b/clients/client-core/src/client/cover_traffic_stream.rs index e3fd0ca197..a8ecad86b2 100644 --- a/clients/client-core/src/client/cover_traffic_stream.rs +++ b/clients/client-core/src/client/cover_traffic_stream.rs @@ -213,7 +213,7 @@ impl LoopCoverTrafficStream { tokio::task::yield_now().await; } - pub fn start_with_shutdown(mut self, mut shutdown: task::ShutdownListener) { + pub fn start_with_shutdown(mut self, mut shutdown: task::TaskClient) { // we should set initial delay only when we actually start the stream let sampled = sample_poisson_duration(&mut self.rng, self.average_cover_message_sending_delay); diff --git a/clients/client-core/src/client/mix_traffic.rs b/clients/client-core/src/client/mix_traffic.rs index ccba8f86ca..94944cb481 100644 --- a/clients/client-core/src/client/mix_traffic.rs +++ b/clients/client-core/src/client/mix_traffic.rs @@ -67,11 +67,11 @@ impl MixTrafficController { } } - pub fn start_with_shutdown(mut self, mut shutdown: task::ShutdownListener) { + pub fn start_with_shutdown(mut self, mut shutdown: task::TaskClient) { spawn_future(async move { debug!("Started MixTrafficController with graceful shutdown support"); - while !shutdown.is_shutdown() { + loop { tokio::select! { mix_packets = self.mix_rx.recv() => match mix_packets { Some(mix_packets) => { @@ -82,8 +82,9 @@ impl MixTrafficController { break; } }, - _ = shutdown.recv() => { + _ = shutdown.recv_with_delay() => { log::trace!("MixTrafficController: Received shutdown"); + break; } } } diff --git a/clients/client-core/src/client/real_messages_control/acknowledgement_control/acknowledgement_listener.rs b/clients/client-core/src/client/real_messages_control/acknowledgement_control/acknowledgement_listener.rs index 95c431fd59..8febbfa85d 100644 --- a/clients/client-core/src/client/real_messages_control/acknowledgement_control/acknowledgement_listener.rs +++ b/clients/client-core/src/client/real_messages_control/acknowledgement_control/acknowledgement_listener.rs @@ -65,7 +65,7 @@ impl AcknowledgementListener { } } - pub(super) async fn run_with_shutdown(&mut self, mut shutdown: task::ShutdownListener) { + pub(super) async fn run_with_shutdown(&mut self, mut shutdown: task::TaskClient) { debug!("Started AcknowledgementListener with graceful shutdown support"); while !shutdown.is_shutdown() { @@ -77,7 +77,7 @@ impl AcknowledgementListener { break; } }, - _ = shutdown.recv() => { + _ = shutdown.recv_with_delay() => { log::trace!("AcknowledgementListener: Received shutdown"); } } diff --git a/clients/client-core/src/client/real_messages_control/acknowledgement_control/action_controller.rs b/clients/client-core/src/client/real_messages_control/acknowledgement_control/action_controller.rs index cb2e0cc7ca..c65592b2f6 100644 --- a/clients/client-core/src/client/real_messages_control/acknowledgement_control/action_controller.rs +++ b/clients/client-core/src/client/real_messages_control/acknowledgement_control/action_controller.rs @@ -249,7 +249,7 @@ impl ActionController { } } - pub(super) async fn run_with_shutdown(&mut self, mut shutdown: task::ShutdownListener) { + pub(super) async fn run_with_shutdown(&mut self, mut shutdown: task::TaskClient) { debug!("Started ActionController with graceful shutdown support"); while !shutdown.is_shutdown() { @@ -270,7 +270,7 @@ impl ActionController { break; } }, - _ = shutdown.recv() => { + _ = shutdown.recv_with_delay() => { log::trace!("ActionController: Received shutdown"); } } diff --git a/clients/client-core/src/client/real_messages_control/acknowledgement_control/input_message_listener.rs b/clients/client-core/src/client/real_messages_control/acknowledgement_control/input_message_listener.rs index cff6b6d719..f492de1aae 100644 --- a/clients/client-core/src/client/real_messages_control/acknowledgement_control/input_message_listener.rs +++ b/clients/client-core/src/client/real_messages_control/acknowledgement_control/input_message_listener.rs @@ -109,7 +109,7 @@ where }; } - pub(super) async fn run_with_shutdown(&mut self, mut shutdown: task::ShutdownListener) { + pub(super) async fn run_with_shutdown(&mut self, mut shutdown: task::TaskClient) { debug!("Started InputMessageListener with graceful shutdown support"); while !shutdown.is_shutdown() { @@ -123,7 +123,7 @@ where break; } }, - _ = shutdown.recv() => { + _ = shutdown.recv_with_delay() => { log::trace!("InputMessageListener: Received shutdown"); } } diff --git a/clients/client-core/src/client/real_messages_control/acknowledgement_control/mod.rs b/clients/client-core/src/client/real_messages_control/acknowledgement_control/mod.rs index 5587958f33..befc24b86a 100644 --- a/clients/client-core/src/client/real_messages_control/acknowledgement_control/mod.rs +++ b/clients/client-core/src/client/real_messages_control/acknowledgement_control/mod.rs @@ -249,7 +249,7 @@ where } } - pub(super) fn start_with_shutdown(self, shutdown: task::ShutdownListener) { + pub(super) fn start_with_shutdown(self, shutdown: task::TaskClient) { let mut acknowledgement_listener = self.acknowledgement_listener; let mut input_message_listener = self.input_message_listener; let mut retransmission_request_listener = self.retransmission_request_listener; diff --git a/clients/client-core/src/client/real_messages_control/acknowledgement_control/retransmission_request_listener.rs b/clients/client-core/src/client/real_messages_control/acknowledgement_control/retransmission_request_listener.rs index b5c6e6a2a8..50a930b94e 100644 --- a/clients/client-core/src/client/real_messages_control/acknowledgement_control/retransmission_request_listener.rs +++ b/clients/client-core/src/client/real_messages_control/acknowledgement_control/retransmission_request_listener.rs @@ -137,7 +137,7 @@ where .await } - pub(super) async fn run_with_shutdown(&mut self, mut shutdown: task::ShutdownListener) { + pub(super) async fn run_with_shutdown(&mut self, mut shutdown: task::TaskClient) { debug!("Started RetransmissionRequestListener with graceful shutdown support"); while !shutdown.is_shutdown() { @@ -149,7 +149,7 @@ where break; } }, - _ = shutdown.recv() => { + _ = shutdown.recv_with_delay() => { log::trace!("RetransmissionRequestListener: Received shutdown"); } } diff --git a/clients/client-core/src/client/real_messages_control/acknowledgement_control/sent_notification_listener.rs b/clients/client-core/src/client/real_messages_control/acknowledgement_control/sent_notification_listener.rs index c273182d59..a931ce4962 100644 --- a/clients/client-core/src/client/real_messages_control/acknowledgement_control/sent_notification_listener.rs +++ b/clients/client-core/src/client/real_messages_control/acknowledgement_control/sent_notification_listener.rs @@ -37,7 +37,7 @@ impl SentNotificationListener { .unwrap(); } - pub(super) async fn run_with_shutdown(&mut self, mut shutdown: task::ShutdownListener) { + pub(super) async fn run_with_shutdown(&mut self, mut shutdown: task::TaskClient) { debug!("Started SentNotificationListener with graceful shutdown support"); while !shutdown.is_shutdown() { @@ -51,7 +51,7 @@ impl SentNotificationListener { break; } }, - _ = shutdown.recv() => { + _ = shutdown.recv_with_delay() => { log::trace!("SentNotificationListener: Received shutdown"); } } diff --git a/clients/client-core/src/client/real_messages_control/mod.rs b/clients/client-core/src/client/real_messages_control/mod.rs index 8789cbdf6a..df04365c9b 100644 --- a/clients/client-core/src/client/real_messages_control/mod.rs +++ b/clients/client-core/src/client/real_messages_control/mod.rs @@ -263,7 +263,7 @@ impl RealMessagesController { } } - pub fn start_with_shutdown(self, shutdown: task::ShutdownListener) { + pub fn start_with_shutdown(self, shutdown: task::TaskClient) { let mut out_queue_control = self.out_queue_control; let ack_control = self.ack_control; let mut reply_control = self.reply_control; diff --git a/clients/client-core/src/client/real_messages_control/real_traffic_stream.rs b/clients/client-core/src/client/real_messages_control/real_traffic_stream.rs index ef1c03046d..a7d62910e5 100644 --- a/clients/client-core/src/client/real_messages_control/real_traffic_stream.rs +++ b/clients/client-core/src/client/real_messages_control/real_traffic_stream.rs @@ -469,7 +469,7 @@ where } #[cfg(not(target_arch = "wasm32"))] - fn log_status(&self, shutdown: &mut task::ShutdownListener) { + fn log_status(&self, shutdown: &mut task::TaskClient) { use crate::error::ClientCoreStatusMessage; let packets = self.transmission_buffer.total_size(); @@ -514,7 +514,7 @@ where } } - pub(super) async fn run_with_shutdown(&mut self, mut shutdown: task::ShutdownListener) { + pub(super) async fn run_with_shutdown(&mut self, mut shutdown: task::TaskClient) { debug!("Started OutQueueControl with graceful shutdown support"); #[cfg(not(target_arch = "wasm32"))] @@ -525,7 +525,7 @@ where while !shutdown.is_shutdown() { tokio::select! { biased; - _ = shutdown.recv() => { + _ = shutdown.recv_with_delay() => { log::trace!("OutQueueControl: Received shutdown"); } _ = status_timer.tick() => { diff --git a/clients/client-core/src/client/real_messages_control/real_traffic_stream/sending_delay_controller.rs b/clients/client-core/src/client/real_messages_control/real_traffic_stream/sending_delay_controller.rs index 4ae2632b85..b7fdff2299 100644 --- a/clients/client-core/src/client/real_messages_control/real_traffic_stream/sending_delay_controller.rs +++ b/clients/client-core/src/client/real_messages_control/real_traffic_stream/sending_delay_controller.rs @@ -77,10 +77,12 @@ impl SendingDelayController { self.current_multiplier } + #[cfg(not(target_arch = "wasm32"))] pub(crate) fn min_multiplier(&self) -> u32 { self.lower_bound } + #[cfg(not(target_arch = "wasm32"))] pub(crate) fn max_multiplier(&self) -> u32 { self.upper_bound } diff --git a/clients/client-core/src/client/received_buffer.rs b/clients/client-core/src/client/received_buffer.rs index 2e7876ae62..afafec44ca 100644 --- a/clients/client-core/src/client/received_buffer.rs +++ b/clients/client-core/src/client/received_buffer.rs @@ -399,21 +399,20 @@ impl RequestReceiver { } } - async fn run_with_shutdown(&mut self, mut shutdown: task::ShutdownListener) { + async fn run_with_shutdown(&mut self, mut shutdown: task::TaskClient) { debug!("Started RequestReceiver with graceful shutdown support"); while !shutdown.is_shutdown() { tokio::select! { biased; - _ = shutdown.recv() => { + _ = shutdown.recv_with_delay() => { log::trace!("RequestReceiver: Received shutdown"); } request = self.query_receiver.next() => { - match request { - Some(message) => self.handle_message(message).await, - None => { - log::trace!("RequestReceiver: Stopping since channel closed"); - break; - }, + if let Some(message) = request { + self.handle_message(message).await + } else { + log::trace!("RequestReceiver: Stopping since channel closed"); + break; } }, } @@ -439,20 +438,19 @@ impl FragmentedMessageReceiver { } } - async fn run_with_shutdown(&mut self, mut shutdown: task::ShutdownListener) { + async fn run_with_shutdown(&mut self, mut shutdown: task::TaskClient) { debug!("Started FragmentedMessageReceiver with graceful shutdown support"); while !shutdown.is_shutdown() { tokio::select! { - new_messages = self.mixnet_packet_receiver.next() => match new_messages { - Some(new_messages) => { + new_messages = self.mixnet_packet_receiver.next() => { + if let Some(new_messages) = new_messages { self.received_buffer.handle_new_received(new_messages).await; - } - None => { + } else { log::trace!("FragmentedMessageReceiver: Stopping since channel closed"); break; } }, - _ = shutdown.recv() => { + _ = shutdown.recv_with_delay() => { log::trace!("FragmentedMessageReceiver: Received shutdown"); } } @@ -490,7 +488,7 @@ impl ReceivedMessagesBufferController { } } - pub fn start_with_shutdown(self, shutdown: task::ShutdownListener) { + pub fn start_with_shutdown(self, shutdown: task::TaskClient) { let mut fragmented_message_receiver = self.fragmented_message_receiver; let mut request_receiver = self.request_receiver; diff --git a/clients/client-core/src/client/replies/reply_controller.rs b/clients/client-core/src/client/replies/reply_controller.rs index 99687527f0..8716203734 100644 --- a/clients/client-core/src/client/replies/reply_controller.rs +++ b/clients/client-core/src/client/replies/reply_controller.rs @@ -891,7 +891,7 @@ where return gloo_timers::future::IntervalStream::new(polling_rate.as_millis() as u32); } - pub(crate) async fn run_with_shutdown(&mut self, mut shutdown: task::ShutdownListener) { + pub(crate) async fn run_with_shutdown(&mut self, mut shutdown: task::TaskClient) { debug!("Started ReplyController with graceful shutdown support"); let polling_rate = Duration::from_secs(5); @@ -904,7 +904,7 @@ where while !shutdown.is_shutdown() { tokio::select! { biased; - _ = shutdown.recv() => { + _ = shutdown.recv_with_delay() => { log::trace!("ReplyController: Received shutdown"); }, req = self.request_receiver.next() => match req { diff --git a/clients/client-core/src/client/replies/reply_storage/mod.rs b/clients/client-core/src/client/replies/reply_storage/mod.rs index 83c0e0e34f..c74ad2bb6b 100644 --- a/clients/client-core/src/client/replies/reply_storage/mod.rs +++ b/clients/client-core/src/client/replies/reply_storage/mod.rs @@ -37,7 +37,7 @@ where pub async fn flush_on_shutdown( mut self, mem_state: CombinedReplyStorage, - mut shutdown: task::ShutdownListener, + mut shutdown: task::TaskClient, ) { use log::{debug, error, info, warn}; diff --git a/clients/client-core/src/client/topology_control.rs b/clients/client-core/src/client/topology_control.rs index 897d6a0e3b..c51fd350cd 100644 --- a/clients/client-core/src/client/topology_control.rs +++ b/clients/client-core/src/client/topology_control.rs @@ -304,7 +304,7 @@ impl TopologyRefresher { self.topology_accessor.ensure_is_routable().await } - pub fn start_with_shutdown(mut self, mut shutdown: task::ShutdownListener) { + pub fn start_with_shutdown(mut self, mut shutdown: task::TaskClient) { spawn_future(async move { debug!("Started TopologyRefresher with graceful shutdown support"); diff --git a/clients/native/src/client/mod.rs b/clients/native/src/client/mod.rs index 0d61b502ab..1a90c48b6f 100644 --- a/clients/native/src/client/mod.rs +++ b/clients/native/src/client/mod.rs @@ -1,6 +1,8 @@ // Copyright 2021-2022 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 +use std::error::Error; + use crate::client::config::Config; use crate::error::ClientError; use crate::websocket; @@ -18,7 +20,7 @@ use log::*; use nymsphinx::addressing::clients::Recipient; use nymsphinx::anonymous_replies::requests::AnonymousSenderTag; use nymsphinx::receiver::ReconstructedMessage; -use task::{wait_for_signal, ShutdownNotifier}; +use task::TaskManager; pub(crate) mod config; @@ -85,16 +87,19 @@ impl SocketClient { client_input: ClientInput, client_output: ClientOutput, self_address: Recipient, + shutdown: task::TaskClient, ) { info!("Starting websocket listener..."); let ClientInput { - shared_lane_queue_lengths, connection_command_sender, input_sender, } = client_input; - let received_buffer_request_sender = client_output.received_buffer_request_sender; + let ClientOutput { + shared_lane_queue_lengths, + received_buffer_request_sender, + } = client_output; let websocket_handler = websocket::Handler::new( input_sender, @@ -104,32 +109,26 @@ impl SocketClient { shared_lane_queue_lengths, ); - websocket::Listener::new(config.get_listening_port()).start(websocket_handler); + websocket::Listener::new(config.get_listening_port()).start(websocket_handler, shutdown); } /// blocking version of `start_socket` method. Will run forever (or until SIGINT is sent) - pub async fn run_socket_forever(self) -> Result<(), ClientError> { + pub async fn run_socket_forever(self) -> Result<(), Box> { let mut shutdown = self.start_socket().await?; - wait_for_signal().await; - println!( - "Received signal - the client will terminate now (threads are not yet nicely stopped, if you see stack traces that's alright)." - ); + let res = task::wait_for_signal_and_error(&mut shutdown).await; log::info!("Sending shutdown"); shutdown.signal_shutdown().ok(); - // Some of these components have shutdown signalling implemented as part of socks5 work, - // but since it's not fully implemented (yet) for all the components of the native client, - // we don't try to wait and instead just stop immediately. log::info!("Waiting for tasks to finish... (Press ctrl-c to force)"); shutdown.wait_for_shutdown().await; log::info!("Stopping nym-client"); - Ok(()) + res } - pub async fn start_socket(self) -> Result { + pub async fn start_socket(self) -> Result { if !self.config.get_socket_type().is_websocket() { return Err(ClientError::InvalidSocketMode); } @@ -150,12 +149,18 @@ impl SocketClient { let client_input = started_client.client_input.register_producer(); let client_output = started_client.client_output.register_consumer(); - Self::start_websocket_listener(&self.config, client_input, client_output, self_address); + Self::start_websocket_listener( + &self.config, + client_input, + client_output, + self_address, + started_client.task_manager.subscribe(), + ); info!("Client startup finished!"); info!("The address of this client is: {}", self_address); - Ok(started_client.shutdown_notifier) + Ok(started_client.task_manager) } pub async fn start_direct(self) -> Result { @@ -192,7 +197,7 @@ impl SocketClient { Ok(DirectClient { client_input, reconstructed_receiver, - _shutdown_notifier: started_client.shutdown_notifier, + _shutdown_notifier: started_client.task_manager, }) } } @@ -202,7 +207,7 @@ pub struct DirectClient { reconstructed_receiver: ReconstructedMessagesReceiver, // we need to keep reference to this guy otherwise things will start dropping - _shutdown_notifier: ShutdownNotifier, + _shutdown_notifier: TaskManager, } impl DirectClient { diff --git a/clients/native/src/commands/mod.rs b/clients/native/src/commands/mod.rs index 4ad9d414dc..cf40470556 100644 --- a/clients/native/src/commands/mod.rs +++ b/clients/native/src/commands/mod.rs @@ -1,8 +1,9 @@ // Copyright 2021 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 +use std::error::Error; + use crate::client::config::{Config, SocketType}; -use crate::error::ClientError; use clap::CommandFactory; use clap::{Parser, Subcommand}; use completions::{fig_generate, ArgShell}; @@ -86,7 +87,7 @@ pub(crate) struct OverrideConfig { enabled_credentials_mode: bool, } -pub(crate) async fn execute(args: &Cli) -> Result<(), ClientError> { +pub(crate) async fn execute(args: &Cli) -> Result<(), Box> { let bin_name = "nym-native-client"; match &args.command { diff --git a/clients/native/src/commands/run.rs b/clients/native/src/commands/run.rs index d5c7eccbdb..3fffa60e18 100644 --- a/clients/native/src/commands/run.rs +++ b/clients/native/src/commands/run.rs @@ -1,6 +1,8 @@ // Copyright 2021 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 +use std::error::Error; + use crate::{ client::{config::Config, SocketClient}, commands::{override_config, OverrideConfig}, @@ -89,14 +91,14 @@ fn version_check(cfg: &Config) -> bool { } } -pub(crate) async fn execute(args: &Run) -> Result<(), ClientError> { +pub(crate) async fn execute(args: &Run) -> Result<(), Box> { let id = &args.id; let mut config = match Config::load_from_file(Some(id)) { Ok(cfg) => cfg, Err(err) => { error!("Failed to load config for {}. Are you sure you have run `init` before? (Error was: {err})", id); - return Err(ClientError::FailedToLoadConfig(id.to_string())); + return Err(Box::new(ClientError::FailedToLoadConfig(id.to_string()))); } }; @@ -105,7 +107,7 @@ pub(crate) async fn execute(args: &Run) -> Result<(), ClientError> { if !version_check(&config) { error!("failed the local version check"); - return Err(ClientError::FailedLocalVersionCheck); + return Err(Box::new(ClientError::FailedLocalVersionCheck)); } SocketClient::new(config).run_socket_forever().await diff --git a/clients/native/src/main.rs b/clients/native/src/main.rs index 00529494e7..972eef1925 100644 --- a/clients/native/src/main.rs +++ b/clients/native/src/main.rs @@ -1,8 +1,9 @@ // Copyright 2021 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 +use std::error::Error; + use clap::{crate_version, Parser}; -use error::ClientError; use logging::setup_logging; use network_defaults::setup_env; @@ -12,7 +13,7 @@ pub mod error; pub mod websocket; #[tokio::main] -async fn main() -> Result<(), ClientError> { +async fn main() -> Result<(), Box> { setup_logging(); println!("{}", banner()); diff --git a/clients/native/src/websocket/handler.rs b/clients/native/src/websocket/handler.rs index 476eb0a567..f35058485e 100644 --- a/clients/native/src/websocket/handler.rs +++ b/clients/native/src/websocket/handler.rs @@ -43,29 +43,16 @@ pub(crate) struct Handler { socket: Option>, received_response_type: ReceivedResponseType, lane_queue_lengths: LaneQueueLengths, -} - -// clone is used to use handler on a new connection, which initially is `None` -impl Clone for Handler { - fn clone(&self) -> Self { - Handler { - msg_input: self.msg_input.clone(), - client_connection_tx: self.client_connection_tx.clone(), - buffer_requester: self.buffer_requester.clone(), - self_full_address: self.self_full_address, - socket: None, - received_response_type: Default::default(), - lane_queue_lengths: self.lane_queue_lengths.clone(), - } - } + is_active: bool, } impl Drop for Handler { fn drop(&mut self) { - if self - .buffer_requester - .unbounded_send(ReceivedBufferMessage::ReceiverDisconnect) - .is_err() + if self.is_active + && self + .buffer_requester + .unbounded_send(ReceivedBufferMessage::ReceiverDisconnect) + .is_err() { error!("we failed to disconnect the receiver from the buffer! presumably the shutdown procedure has been initiated!") } @@ -88,6 +75,22 @@ impl Handler { socket: None, received_response_type: Default::default(), lane_queue_lengths, + is_active: false, + } + } + + // Used to use handler on a new connection, which initially is `None` + // TODO: make sure we only ever have one active handler + pub fn create_active_handler(&self) -> Self { + Handler { + msg_input: self.msg_input.clone(), + client_connection_tx: self.client_connection_tx.clone(), + buffer_requester: self.buffer_requester.clone(), + self_full_address: self.self_full_address, + socket: None, + received_response_type: Default::default(), + lane_queue_lengths: self.lane_queue_lengths.clone(), + is_active: true, } } @@ -362,8 +365,12 @@ impl Handler { } } - async fn listen_for_requests(&mut self, mut msg_receiver: ReconstructedMessagesReceiver) { - loop { + async fn listen_for_requests( + &mut self, + mut msg_receiver: ReconstructedMessagesReceiver, + mut task_client: task::TaskClient, + ) { + while !task_client.is_shutdown() { tokio::select! { // we can either get a client request from the websocket socket_msg = self.next_websocket_request() => { @@ -402,12 +409,24 @@ impl Handler { break; } } + _ = task_client.recv() => { + log::trace!("Websocket handler: Received shutdown"); + } } } + log::debug!("Websocket handler: Exiting"); } // consume self to make sure `drop` is called after this is done - pub(crate) async fn handle_connection(mut self, socket: TcpStream) { + pub(crate) async fn handle_connection( + mut self, + socket: TcpStream, + mut task_client: task::TaskClient, + ) { + // We don't want a crash in the connection handler to trigger a shutdown of the whole + // process. + task_client.mark_as_success(); + let ws_stream = match accept_async(socket).await { Ok(ws_stream) => ws_stream, Err(err) => { @@ -426,7 +445,8 @@ impl Handler { )) .expect("the buffer request failed!"); - self.listen_for_requests(reconstructed_receiver).await; + self.listen_for_requests(reconstructed_receiver, task_client) + .await; } } diff --git a/clients/native/src/websocket/listener.rs b/clients/native/src/websocket/listener.rs index cf9c7cfa65..aef9150b43 100644 --- a/clients/native/src/websocket/listener.rs +++ b/clients/native/src/websocket/listener.rs @@ -32,7 +32,7 @@ impl Listener { } } - pub(crate) async fn run(&mut self, handler: Handler) { + pub(crate) async fn run(&mut self, handler: Handler, mut task_client: task::TaskClient) { let tcp_listener = match tokio::net::TcpListener::bind(self.address).await { Ok(listener) => listener, Err(err) => { @@ -45,10 +45,23 @@ impl Listener { loop { tokio::select! { + // When the handler finishes we check if shutdown is signalled _ = notify.notified() => { + if task_client.is_shutdown() { + log::trace!("Websocket listener: detected shutdown after connection closed"); + break; + } // our connection terminated - we are open to a new one now! self.state = State::AwaitingConnection; } + // ... but when there is no connected client at the time of shutdown being + // signalled, we handle it here. + _ = task_client.recv() => { + if !self.state.is_connected() { + log::trace!("Not connected: shutting down"); + break; + } + } new_conn = tcp_listener.accept() => { match new_conn { Ok((mut socket, remote_addr)) => { @@ -70,9 +83,10 @@ impl Listener { // it's done so that any new connections to this listener could be rejected rather than left // hanging because the executor doesn't come back here let notify_clone = Arc::clone(¬ify); - let fresh_handler = handler.clone(); + let fresh_handler = handler.create_active_handler(); + let task_client_handler = task_client.clone(); tokio::spawn(async move { - fresh_handler.handle_connection(socket).await; + fresh_handler.handle_connection(socket, task_client_handler).await; notify_clone.notify_one(); }); self.state = State::Connected; @@ -83,11 +97,12 @@ impl Listener { } } } + log::debug!("Websocket listener: Exiting"); } - pub(crate) fn start(mut self, handler: Handler) -> JoinHandle<()> { + pub(crate) fn start(mut self, handler: Handler, shutdown: task::TaskClient) -> JoinHandle<()> { info!("Running websocket on {:?}", self.address.to_string()); - tokio::spawn(async move { self.run(handler).await }) + tokio::spawn(async move { self.run(handler, shutdown).await }) } } diff --git a/clients/socks5/src/client/mod.rs b/clients/socks5/src/client/mod.rs index 45c9ffad7f..2f62354a01 100644 --- a/clients/socks5/src/client/mod.rs +++ b/clients/socks5/src/client/mod.rs @@ -19,7 +19,7 @@ use gateway_client::bandwidth::BandwidthController; use log::*; use nymsphinx::addressing::clients::Recipient; use std::error::Error; -use task::{wait_for_signal_and_error, ShutdownListener, ShutdownNotifier}; +use task::{wait_for_signal_and_error, TaskClient, TaskManager}; pub mod config; @@ -96,19 +96,21 @@ impl NymClient { client_input: ClientInput, client_output: ClientOutput, self_address: Recipient, - shutdown: ShutdownListener, + shutdown: TaskClient, ) { info!("Starting socks5 listener..."); let auth_methods = vec![AuthenticationMethods::NoAuth as u8]; let allowed_users: Vec = Vec::new(); let ClientInput { - shared_lane_queue_lengths, connection_command_sender, input_sender, } = client_input; - let received_buffer_request_sender = client_output.received_buffer_request_sender; + let ClientOutput { + shared_lane_queue_lengths, + received_buffer_request_sender, + } = client_output; let authenticator = Authenticator::new(auth_methods, allowed_users); let mut sphinx_socks = SphinxSocksServer::new( @@ -200,7 +202,7 @@ impl NymClient { res } - pub async fn start(self) -> Result { + pub async fn start(self) -> Result { let base_builder = BaseClientBuilder::new_from_base_config( self.config.get_base(), self.key_manager, @@ -222,12 +224,12 @@ impl NymClient { client_input, client_output, self_address, - started_client.shutdown_notifier.subscribe(), + started_client.task_manager.subscribe(), ); info!("Client startup finished!"); info!("The address of this client is: {}", self_address); - Ok(started_client.shutdown_notifier) + Ok(started_client.task_manager) } } diff --git a/clients/socks5/src/socks/client.rs b/clients/socks5/src/socks/client.rs index e5a8fc9d56..ff7caa1b6b 100644 --- a/clients/socks5/src/socks/client.rs +++ b/clients/socks5/src/socks/client.rs @@ -20,7 +20,7 @@ use socks5_requests::{ConnectionId, Message, RemoteAddress, Request}; use std::io; use std::net::SocketAddr; use std::pin::Pin; -use task::ShutdownListener; +use task::TaskClient; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf}; use tokio::{self, net::TcpStream}; @@ -164,7 +164,7 @@ pub(crate) struct SocksClient { self_address: Recipient, started_proxy: bool, lane_queue_lengths: LaneQueueLengths, - shutdown_listener: ShutdownListener, + shutdown_listener: TaskClient, } impl Drop for SocksClient { @@ -190,7 +190,7 @@ impl SocksClient { controller_sender: ControllerSender, self_address: &Recipient, lane_queue_lengths: LaneQueueLengths, - mut shutdown_listener: ShutdownListener, + mut shutdown_listener: TaskClient, ) -> Self { // If this task fails and exits, we don't want to send shutdown signal shutdown_listener.mark_as_success(); diff --git a/clients/socks5/src/socks/mixnet_responses.rs b/clients/socks5/src/socks/mixnet_responses.rs index d4fb8fd544..97835340df 100644 --- a/clients/socks5/src/socks/mixnet_responses.rs +++ b/clients/socks5/src/socks/mixnet_responses.rs @@ -9,13 +9,13 @@ use client_core::client::received_buffer::{ReceivedBufferMessage, ReceivedBuffer use nymsphinx::receiver::ReconstructedMessage; use proxy_helpers::connection_controller::{ControllerCommand, ControllerSender}; use socks5_requests::Message; -use task::ShutdownListener; +use task::TaskClient; pub(crate) struct MixnetResponseListener { buffer_requester: ReceivedBufferRequestSender, mix_response_receiver: ReconstructedMessagesReceiver, controller_sender: ControllerSender, - shutdown: ShutdownListener, + shutdown: TaskClient, } impl Drop for MixnetResponseListener { @@ -37,7 +37,7 @@ impl MixnetResponseListener { pub(crate) fn new( buffer_requester: ReceivedBufferRequestSender, controller_sender: ControllerSender, - shutdown: ShutdownListener, + shutdown: TaskClient, ) -> Self { let (mix_response_sender, mix_response_receiver) = mpsc::unbounded(); buffer_requester diff --git a/clients/socks5/src/socks/server.rs b/clients/socks5/src/socks/server.rs index a7e0a3c0e3..0b193dc9d4 100644 --- a/clients/socks5/src/socks/server.rs +++ b/clients/socks5/src/socks/server.rs @@ -13,7 +13,7 @@ use nymsphinx::addressing::clients::Recipient; use proxy_helpers::connection_controller::{BroadcastActiveConnections, Controller}; use std::net::SocketAddr; use tap::TapFallible; -use task::ShutdownListener; +use task::TaskClient; use tokio::net::TcpListener; /// A Socks5 server that listens for connections. @@ -24,7 +24,7 @@ pub struct SphinxSocksServer { self_address: Recipient, client_config: client::Config, lane_queue_lengths: LaneQueueLengths, - shutdown: ShutdownListener, + shutdown: TaskClient, } impl SphinxSocksServer { @@ -36,7 +36,7 @@ impl SphinxSocksServer { self_address: Recipient, lane_queue_lengths: LaneQueueLengths, client_config: client::Config, - shutdown: ShutdownListener, + shutdown: TaskClient, ) -> Self { // hardcode ip as we (presumably) ONLY want to listen locally. If we change it, we can // just modify the config diff --git a/clients/webassembly/src/client/mod.rs b/clients/webassembly/src/client/mod.rs index 9c9f20b2c6..bd9ceab576 100644 --- a/clients/webassembly/src/client/mod.rs +++ b/clients/webassembly/src/client/mod.rs @@ -14,7 +14,7 @@ use nymsphinx::addressing::clients::Recipient; use nymsphinx::anonymous_replies::requests::AnonymousSenderTag; use rand::rngs::OsRng; use std::sync::Arc; -use task::ShutdownNotifier; +use task::TaskManager; use wasm_bindgen::prelude::*; use wasm_bindgen_futures::future_to_promise; use wasm_utils::{console_error, console_log}; @@ -30,7 +30,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 - _shutdown: ShutdownNotifier, + _task_manager: TaskManager, } #[wasm_bindgen] @@ -121,7 +121,7 @@ impl NymClientBuilder { Ok(JsValue::from(NymClient { self_address, client_input: Arc::new(client_input), - _shutdown: started_client.shutdown_notifier, + _task_manager: started_client.task_manager, })) }) } diff --git a/common/client-libs/gateway-client/src/client.rs b/common/client-libs/gateway-client/src/client.rs index bb6fe0bbd1..23cae7e4a6 100644 --- a/common/client-libs/gateway-client/src/client.rs +++ b/common/client-libs/gateway-client/src/client.rs @@ -22,7 +22,7 @@ use rand::rngs::OsRng; use std::convert::TryFrom; use std::sync::Arc; use std::time::Duration; -use task::ShutdownListener; +use task::TaskClient; use tungstenite::protocol::Message; #[cfg(feature = "coconut")] @@ -67,7 +67,7 @@ pub struct GatewayClient { reconnection_backoff: Duration, /// Listen to shutdown messages. - shutdown: ShutdownListener, + shutdown: TaskClient, } impl GatewayClient { @@ -83,7 +83,7 @@ impl GatewayClient { ack_sender: AcknowledgementSender, response_timeout_duration: Duration, bandwidth_controller: Option>, - shutdown: ShutdownListener, + shutdown: TaskClient, ) -> Self { GatewayClient { authenticated: false, @@ -135,7 +135,7 @@ impl GatewayClient { // perfectly fine here, because it's not meant to be used let (ack_tx, _) = mpsc::unbounded(); let (mix_tx, _) = mpsc::unbounded(); - let shutdown = ShutdownListener::dummy(); + let shutdown = TaskClient::dummy(); let packet_router = PacketRouter::new(ack_tx, mix_tx, shutdown.clone()); GatewayClient { diff --git a/common/client-libs/gateway-client/src/packet_router.rs b/common/client-libs/gateway-client/src/packet_router.rs index c64fbf92f8..daeac0e548 100644 --- a/common/client-libs/gateway-client/src/packet_router.rs +++ b/common/client-libs/gateway-client/src/packet_router.rs @@ -9,7 +9,7 @@ use futures::channel::mpsc; use log::*; use nymsphinx::addressing::nodes::MAX_NODE_ADDRESS_UNPADDED_LEN; use nymsphinx::params::packet_sizes::PacketSize; -use task::ShutdownListener; +use task::TaskClient; pub type MixnetMessageSender = mpsc::UnboundedSender>>; pub type MixnetMessageReceiver = mpsc::UnboundedReceiver>>; @@ -21,14 +21,14 @@ pub type AcknowledgementReceiver = mpsc::UnboundedReceiver>>; pub struct PacketRouter { ack_sender: AcknowledgementSender, mixnet_message_sender: MixnetMessageSender, - shutdown: ShutdownListener, + shutdown: TaskClient, } impl PacketRouter { pub fn new( ack_sender: AcknowledgementSender, mixnet_message_sender: MixnetMessageSender, - shutdown: ShutdownListener, + shutdown: TaskClient, ) -> Self { PacketRouter { ack_sender, diff --git a/common/client-libs/gateway-client/src/socket_state.rs b/common/client-libs/gateway-client/src/socket_state.rs index 8c66b20906..2bedf1e8ed 100644 --- a/common/client-libs/gateway-client/src/socket_state.rs +++ b/common/client-libs/gateway-client/src/socket_state.rs @@ -10,7 +10,7 @@ use futures::{SinkExt, StreamExt}; use gateway_requests::registration::handshake::SharedKeys; use log::*; use std::sync::Arc; -use task::ShutdownListener; +use task::TaskClient; use tungstenite::Message; #[cfg(not(target_arch = "wasm32"))] @@ -84,7 +84,7 @@ impl PartiallyDelegated { conn: WsConn, packet_router: PacketRouter, shared_key: Arc, - mut shutdown: ShutdownListener, + mut shutdown: TaskClient, ) -> Self { // when called for, it NEEDS TO yield back the stream so that we could merge it and // read control request responses. diff --git a/common/mixnode-common/src/verloc/listener.rs b/common/mixnode-common/src/verloc/listener.rs index 537ec3aad5..812358e330 100644 --- a/common/mixnode-common/src/verloc/listener.rs +++ b/common/mixnode-common/src/verloc/listener.rs @@ -11,7 +11,7 @@ use std::fmt::{Display, Formatter}; use std::net::SocketAddr; use std::sync::Arc; use std::{fmt, io, process}; -use task::ShutdownListener; +use task::TaskClient; use tokio::io::AsyncWriteExt; use tokio::net::{TcpListener, TcpStream}; use tokio_util::codec::{Decoder, Encoder, Framed}; @@ -19,14 +19,14 @@ use tokio_util::codec::{Decoder, Encoder, Framed}; pub(crate) struct PacketListener { address: SocketAddr, connection_handler: Arc, - shutdown: ShutdownListener, + shutdown: TaskClient, } impl PacketListener { pub(crate) fn new( address: SocketAddr, identity: Arc, - shutdown: ShutdownListener, + shutdown: TaskClient, ) -> Self { PacketListener { address, @@ -91,7 +91,7 @@ impl ConnectionHandler { self: Arc, conn: TcpStream, remote: SocketAddr, - mut shutdown_listener: ShutdownListener, + mut shutdown_listener: TaskClient, ) { debug!("Starting connection handler for {:?}", remote); diff --git a/common/mixnode-common/src/verloc/mod.rs b/common/mixnode-common/src/verloc/mod.rs index fc7e19332a..21053dd7cc 100644 --- a/common/mixnode-common/src/verloc/mod.rs +++ b/common/mixnode-common/src/verloc/mod.rs @@ -12,7 +12,7 @@ use rand::thread_rng; use std::net::{SocketAddr, ToSocketAddrs}; use std::sync::Arc; use std::time::Duration; -use task::ShutdownListener; +use task::TaskClient; use tokio::task::JoinHandle; use tokio::time::sleep; use url::Url; @@ -168,7 +168,7 @@ pub struct VerlocMeasurer { config: Config, packet_sender: Arc, packet_listener: Arc, - shutdown_listener: ShutdownListener, + shutdown_listener: TaskClient, currently_used_api: usize, @@ -184,7 +184,7 @@ impl VerlocMeasurer { pub fn new( mut config: Config, identity: Arc, - shutdown_listener: ShutdownListener, + shutdown_listener: TaskClient, ) -> Self { config.nym_api_urls.shuffle(&mut thread_rng()); diff --git a/common/mixnode-common/src/verloc/sender.rs b/common/mixnode-common/src/verloc/sender.rs index bed34f4474..91f7ee9974 100644 --- a/common/mixnode-common/src/verloc/sender.rs +++ b/common/mixnode-common/src/verloc/sender.rs @@ -11,7 +11,7 @@ use std::net::SocketAddr; use std::sync::Arc; use std::time::Duration; use std::{fmt, io}; -use task::ShutdownListener; +use task::TaskClient; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpStream; use tokio::time::sleep; @@ -45,7 +45,7 @@ pub(crate) struct PacketSender { packet_timeout: Duration, connection_timeout: Duration, delay_between_packets: Duration, - shutdown_listener: ShutdownListener, + shutdown_listener: TaskClient, } impl PacketSender { @@ -55,7 +55,7 @@ impl PacketSender { packet_timeout: Duration, connection_timeout: Duration, delay_between_packets: Duration, - shutdown_listener: ShutdownListener, + shutdown_listener: TaskClient, ) -> Self { PacketSender { identity, diff --git a/common/socks5/proxy-helpers/src/connection_controller.rs b/common/socks5/proxy-helpers/src/connection_controller.rs index d7e2645190..66c8b94a7a 100644 --- a/common/socks5/proxy-helpers/src/connection_controller.rs +++ b/common/socks5/proxy-helpers/src/connection_controller.rs @@ -11,7 +11,7 @@ use std::{ collections::{HashMap, HashSet}, time::Duration, }; -use task::ShutdownListener; +use task::TaskClient; use tokio::time; /// A generic message produced after reading from a socket/connection. It includes data that was @@ -98,14 +98,14 @@ pub struct Controller { // un-order messages. Note we don't ever expect to have more than 1-2 messages per connection here pending_messages: HashMap, bool)>>, - shutdown: ShutdownListener, + shutdown: TaskClient, } impl Controller { pub fn new( client_connection_tx: ConnectionCommandSender, broadcast_connections: BroadcastActiveConnections, - shutdown: ShutdownListener, + shutdown: TaskClient, ) -> (Self, ControllerSender) { let (sender, receiver) = mpsc::unbounded(); ( diff --git a/common/socks5/proxy-helpers/src/proxy_runner/inbound.rs b/common/socks5/proxy-helpers/src/proxy_runner/inbound.rs index b6688657de..2065f7770b 100644 --- a/common/socks5/proxy-helpers/src/proxy_runner/inbound.rs +++ b/common/socks5/proxy-helpers/src/proxy_runner/inbound.rs @@ -15,7 +15,7 @@ use socks5_requests::ConnectionId; use std::fmt::Debug; use std::time::Duration; use std::{io, sync::Arc}; -use task::ShutdownListener; +use task::TaskClient; use tokio::select; use tokio::{net::tcp::OwnedReadHalf, sync::Notify, time::sleep}; @@ -170,7 +170,7 @@ pub(super) async fn run_inbound( adapter_fn: F, shutdown_notify: Arc, lane_queue_lengths: Option, - mut shutdown_listener: ShutdownListener, + mut shutdown_listener: TaskClient, ) -> OwnedReadHalf where F: Fn(ConnectionId, Vec, bool) -> S + Send + 'static, diff --git a/common/socks5/proxy-helpers/src/proxy_runner/mod.rs b/common/socks5/proxy-helpers/src/proxy_runner/mod.rs index defcbf4f5f..a87ada287e 100644 --- a/common/socks5/proxy-helpers/src/proxy_runner/mod.rs +++ b/common/socks5/proxy-helpers/src/proxy_runner/mod.rs @@ -6,7 +6,7 @@ use client_connections::LaneQueueLengths; use socks5_requests::ConnectionId; use std::fmt::Debug; use std::{sync::Arc, time::Duration}; -use task::ShutdownListener; +use task::TaskClient; use tokio::{net::TcpStream, sync::Notify}; mod inbound; @@ -50,7 +50,7 @@ pub struct ProxyRunner { lane_queue_lengths: Option, // Listens to shutdown commands from higher up - shutdown_listener: ShutdownListener, + shutdown_listener: TaskClient, } impl ProxyRunner @@ -66,7 +66,7 @@ where mix_sender: MixProxySender, connection_id: ConnectionId, lane_queue_lengths: Option, - shutdown_listener: ShutdownListener, + shutdown_listener: TaskClient, ) -> Self { ProxyRunner { mix_receiver: Some(mix_receiver), diff --git a/common/socks5/proxy-helpers/src/proxy_runner/outbound.rs b/common/socks5/proxy-helpers/src/proxy_runner/outbound.rs index f6c713e948..ed7f8d74c6 100644 --- a/common/socks5/proxy-helpers/src/proxy_runner/outbound.rs +++ b/common/socks5/proxy-helpers/src/proxy_runner/outbound.rs @@ -8,7 +8,7 @@ use futures::StreamExt; use log::*; use socks5_requests::ConnectionId; use std::{sync::Arc, time::Duration}; -use task::ShutdownListener; +use task::TaskClient; use tokio::io::AsyncWriteExt; use tokio::select; use tokio::{net::tcp::OwnedWriteHalf, sync::Notify, time::sleep, time::Instant}; @@ -51,7 +51,7 @@ pub(super) async fn run_outbound( mut mix_receiver: ConnectionReceiver, connection_id: ConnectionId, shutdown_notify: Arc, - mut shutdown_listener: ShutdownListener, + mut shutdown_listener: TaskClient, ) -> (OwnedWriteHalf, ConnectionReceiver) { let shutdown_future = shutdown_notify.notified().then(|_| sleep(SHUTDOWN_TIMEOUT)); tokio::pin!(shutdown_future); diff --git a/common/task/src/lib.rs b/common/task/src/lib.rs index 5d3d0ea2f5..7ac5ab8431 100644 --- a/common/task/src/lib.rs +++ b/common/task/src/lib.rs @@ -1,13 +1,12 @@ // Copyright 2022 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 -pub mod shutdown; +pub mod manager; #[cfg(not(target_arch = "wasm32"))] pub mod signal; pub mod spawn; -// WIP(JON): those both need to be public? -pub use shutdown::{ShutdownListener, ShutdownNotifier, StatusReceiver, StatusSender}; +pub use manager::{StatusReceiver, StatusSender, TaskClient, TaskManager}; #[cfg(not(target_arch = "wasm32"))] pub use signal::{wait_for_signal, wait_for_signal_and_error}; diff --git a/common/task/src/shutdown.rs b/common/task/src/manager.rs similarity index 89% rename from common/task/src/shutdown.rs rename to common/task/src/manager.rs index 4fe7a22bdd..3729d317a9 100644 --- a/common/task/src/shutdown.rs +++ b/common/task/src/manager.rs @@ -14,7 +14,6 @@ use tokio::{ const DEFAULT_SHUTDOWN_TIMER_SECS: u64 = 5; -// WIP(JON): can these be from futures channel too? pub(crate) type SentError = Box; type ErrorSender = mpsc::UnboundedSender; type ErrorReceiver = mpsc::UnboundedReceiver; @@ -29,9 +28,10 @@ enum TaskError { UnexpectedHalt, } -/// Used to notify other tasks to gracefully shutdown +/// Listens to status and error messages from tasks, as well as notifying them to gracefully +/// shutdown. Keeps track of if task stop unexpectedly, such as in a panic. #[derive(Debug)] -pub struct ShutdownNotifier { +pub struct TaskManager { // These channels have the dual purpose of signalling it's time to shutdown, but also to keep // track of which tasks we are still waiting for. notify_tx: watch::Sender<()>, @@ -55,7 +55,7 @@ pub struct ShutdownNotifier { task_status_rx: Option, } -impl Default for ShutdownNotifier { +impl Default for TaskManager { fn default() -> Self { let (notify_tx, notify_rx) = watch::channel(()); let (task_halt_tx, task_halt_rx) = mpsc::unbounded_channel(); @@ -77,11 +77,7 @@ impl Default for ShutdownNotifier { } } -impl ShutdownNotifier { - pub fn take_task_status_rx(&mut self) -> Option { - self.task_status_rx.take() - } - +impl TaskManager { pub fn new(shutdown_timer_secs: u64) -> Self { Self { shutdown_timer_secs, @@ -89,8 +85,8 @@ impl ShutdownNotifier { } } - pub fn subscribe(&self) -> ShutdownListener { - ShutdownListener::new( + pub fn subscribe(&self) -> TaskClient { + TaskClient::new( self.notify_rx .as_ref() .expect("Unable to subscribe to shutdown notifier that is already shutdown") @@ -173,9 +169,10 @@ impl ShutdownNotifier { } } -/// Listen for shutdown notifications +/// Listen for shutdown notifications, and can send error and status messages back to the +/// `TaskManager` #[derive(Clone, Debug)] -pub struct ShutdownListener { +pub struct TaskClient { // If a shutdown notification has been registered shutdown: bool, @@ -193,10 +190,10 @@ pub struct ShutdownListener { status_msg: StatusSender, // The current operating mode - mode: ShutdownListenerMode, + mode: ClientOperatingMode, } -impl ShutdownListener { +impl TaskClient { #[cfg(not(target_arch = "wasm32"))] const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); @@ -205,31 +202,31 @@ impl ShutdownListener { return_error: ErrorSender, drop_error: ErrorSender, status_msg: StatusSender, - ) -> ShutdownListener { - ShutdownListener { + ) -> TaskClient { + TaskClient { shutdown: false, notify, return_error, drop_error, status_msg, - mode: ShutdownListenerMode::Listening, + mode: ClientOperatingMode::Listening, } } // Create a dummy that will never report that we should shutdown. - pub fn dummy() -> ShutdownListener { + pub fn dummy() -> TaskClient { let (_notify_tx, notify_rx) = watch::channel(()); let (task_halt_tx, _task_halt_rx) = mpsc::unbounded_channel(); let (task_drop_tx, _task_drop_rx) = mpsc::unbounded_channel(); //let (task_status_tx, _task_status_rx) = futures::channel::mpsc::unbounded(); let (task_status_tx, _task_status_rx) = futures::channel::mpsc::channel(128); - ShutdownListener { + TaskClient { shutdown: false, notify: notify_rx, return_error: task_halt_tx, drop_error: task_drop_tx, status_msg: task_status_tx, - mode: ShutdownListenerMode::Dummy, + mode: ClientOperatingMode::Dummy, } } @@ -256,6 +253,15 @@ impl ShutdownListener { self.shutdown = true; } + pub async fn recv_with_delay(&mut self) { + self.recv() + .then(|msg| async move { + sleep(Duration::from_secs(1)).await; + msg + }) + .await + } + pub async fn recv_timeout(&mut self) { if self.mode.is_dummy() { return pending().await; @@ -315,7 +321,7 @@ impl ShutdownListener { } } -impl Drop for ShutdownListener { +impl Drop for TaskClient { fn drop(&mut self) { if !self.mode.should_signal_on_drop() { return; @@ -331,7 +337,7 @@ impl Drop for ShutdownListener { } #[derive(Clone, Debug, PartialEq, Eq)] -enum ShutdownListenerMode { +enum ClientOperatingMode { // Normal operations Listening, // Normal operations, but we don't report back if the we stop by getting dropped. @@ -340,20 +346,20 @@ enum ShutdownListenerMode { Dummy, } -impl ShutdownListenerMode { +impl ClientOperatingMode { fn is_dummy(&self) -> bool { - self == &ShutdownListenerMode::Dummy + self == &ClientOperatingMode::Dummy } fn should_signal_on_drop(&self) -> bool { match self { - ShutdownListenerMode::Listening => true, - ShutdownListenerMode::ListeningButDontReportHalt | ShutdownListenerMode::Dummy => false, + ClientOperatingMode::Listening => true, + ClientOperatingMode::ListeningButDontReportHalt | ClientOperatingMode::Dummy => false, } } fn set_should_not_signal_on_drop(&mut self) { - use ShutdownListenerMode::{Dummy, Listening, ListeningButDontReportHalt}; + use ClientOperatingMode::{Dummy, Listening, ListeningButDontReportHalt}; *self = match &self { ListeningButDontReportHalt | Listening => ListeningButDontReportHalt, Dummy => Dummy, @@ -367,7 +373,7 @@ mod tests { #[tokio::test] async fn signal_shutdown() { - let shutdown = ShutdownNotifier::default(); + let shutdown = TaskManager::default(); let mut listener = shutdown.subscribe(); let task = tokio::spawn(async move { diff --git a/common/task/src/signal.rs b/common/task/src/signal.rs index a04a0f5601..483c7c630b 100644 --- a/common/task/src/signal.rs +++ b/common/task/src/signal.rs @@ -1,4 +1,4 @@ -use crate::{shutdown::SentError, ShutdownNotifier}; +use crate::{manager::SentError, TaskManager}; #[cfg(unix)] pub async fn wait_for_signal() { @@ -29,7 +29,7 @@ pub async fn wait_for_signal() { } #[cfg(unix)] -pub async fn wait_for_signal_and_error(shutdown: &mut ShutdownNotifier) -> Result<(), SentError> { +pub async fn wait_for_signal_and_error(shutdown: &mut TaskManager) -> Result<(), SentError> { use tokio::signal::unix::{signal, SignalKind}; let mut sigterm = signal(SignalKind::terminate()).expect("Failed to setup SIGTERM channel"); @@ -56,7 +56,7 @@ pub async fn wait_for_signal_and_error(shutdown: &mut ShutdownNotifier) -> Resul } #[cfg(not(unix))] -pub async fn wait_for_signal_and_error(shutdown: &mut ShutdownNotifier) -> Result<(), SentError> { +pub async fn wait_for_signal_and_error(shutdown: &mut TaskManager) -> Result<(), SentError> { tokio::select! { _ = tokio::signal::ctrl_c() => { log::info!("Received SIGINT"); diff --git a/common/task/src/spawn.rs b/common/task/src/spawn.rs index ba66ea7f16..d30c9ac850 100644 --- a/common/task/src/spawn.rs +++ b/common/task/src/spawn.rs @@ -1,4 +1,4 @@ -use crate::ShutdownListener; +use crate::TaskClient; use std::future::Future; #[cfg(target_arch = "wasm32")] @@ -18,7 +18,7 @@ where tokio::spawn(future); } -pub fn spawn_with_report_error(future: F, mut shutdown: ShutdownListener) +pub fn spawn_with_report_error(future: F, mut shutdown: TaskClient) where F: Future> + Send + 'static, T: 'static, diff --git a/explorer-api/src/country_statistics/distribution.rs b/explorer-api/src/country_statistics/distribution.rs index f48e87eebb..93d12067b6 100644 --- a/explorer-api/src/country_statistics/distribution.rs +++ b/explorer-api/src/country_statistics/distribution.rs @@ -1,5 +1,5 @@ use log::info; -use task::ShutdownListener; +use task::TaskClient; use crate::country_statistics::country_nodes_distribution::CountryNodesDistribution; use crate::COUNTRY_DATA_REFRESH_INTERVAL; @@ -8,11 +8,11 @@ use crate::state::ExplorerApiStateContext; pub(crate) struct CountryStatisticsDistributionTask { state: ExplorerApiStateContext, - shutdown: ShutdownListener, + shutdown: TaskClient, } impl CountryStatisticsDistributionTask { - pub(crate) fn new(state: ExplorerApiStateContext, shutdown: ShutdownListener) -> Self { + pub(crate) fn new(state: ExplorerApiStateContext, shutdown: TaskClient) -> Self { CountryStatisticsDistributionTask { state, shutdown } } diff --git a/explorer-api/src/country_statistics/geolocate.rs b/explorer-api/src/country_statistics/geolocate.rs index c77a61a9a1..0ae79985c7 100644 --- a/explorer-api/src/country_statistics/geolocate.rs +++ b/explorer-api/src/country_statistics/geolocate.rs @@ -4,15 +4,15 @@ use crate::mix_nodes::location::Location; use crate::state::ExplorerApiStateContext; use log::{info, warn}; -use task::ShutdownListener; +use task::TaskClient; pub(crate) struct GeoLocateTask { state: ExplorerApiStateContext, - shutdown: ShutdownListener, + shutdown: TaskClient, } impl GeoLocateTask { - pub(crate) fn new(state: ExplorerApiStateContext, shutdown: ShutdownListener) -> Self { + pub(crate) fn new(state: ExplorerApiStateContext, shutdown: TaskClient) -> Self { GeoLocateTask { state, shutdown } } diff --git a/explorer-api/src/main.rs b/explorer-api/src/main.rs index 5154c87e78..e1a02c57ea 100644 --- a/explorer-api/src/main.rs +++ b/explorer-api/src/main.rs @@ -8,7 +8,7 @@ use dotenv::dotenv; use log::info; use logging::setup_logging; use network_defaults::setup_env; -use task::ShutdownNotifier; +use task::TaskManager; mod buy_terms; pub(crate) mod cache; @@ -57,7 +57,7 @@ impl ExplorerApi { let nym_api_url = self.state.inner.validator_client.api_endpoint(); info!("Using validator API - {}", nym_api_url); - let shutdown = ShutdownNotifier::default(); + let shutdown = TaskManager::default(); // spawn concurrent tasks crate::tasks::ExplorerApiTasks::new(self.state.clone(), shutdown.subscribe()).start(); @@ -78,7 +78,7 @@ impl ExplorerApi { self.wait_for_interrupt(shutdown).await } - async fn wait_for_interrupt(&self, mut shutdown: ShutdownNotifier) { + async fn wait_for_interrupt(&self, mut shutdown: TaskManager) { wait_for_signal().await; log::info!("Sending shutdown"); diff --git a/explorer-api/src/tasks.rs b/explorer-api/src/tasks.rs index bab29b798e..b7e72f72d2 100644 --- a/explorer-api/src/tasks.rs +++ b/explorer-api/src/tasks.rs @@ -4,7 +4,7 @@ use std::future::Future; use mixnet_contract_common::GatewayBond; -use task::ShutdownListener; +use task::TaskClient; use validator_client::models::MixNodeBondAnnotated; use validator_client::nymd::error::NymdError; use validator_client::nymd::{Paging, QueryNymdClient, ValidatorResponse}; @@ -15,11 +15,11 @@ use crate::state::ExplorerApiStateContext; pub(crate) struct ExplorerApiTasks { state: ExplorerApiStateContext, - shutdown: ShutdownListener, + shutdown: TaskClient, } impl ExplorerApiTasks { - pub(crate) fn new(state: ExplorerApiStateContext, shutdown: ShutdownListener) -> Self { + pub(crate) fn new(state: ExplorerApiStateContext, shutdown: TaskClient) -> Self { ExplorerApiTasks { state, shutdown } } diff --git a/mixnode/src/node/listener/connection_handler/mod.rs b/mixnode/src/node/listener/connection_handler/mod.rs index 08e8395917..52203c89ab 100644 --- a/mixnode/src/node/listener/connection_handler/mod.rs +++ b/mixnode/src/node/listener/connection_handler/mod.rs @@ -5,7 +5,7 @@ use crate::node::listener::connection_handler::packet_processing::{ MixProcessingResult, PacketProcessor, }; use crate::node::packet_delayforwarder::PacketDelayForwardSender; -use crate::node::ShutdownListener; +use crate::node::TaskClient; use futures::StreamExt; use log::{error, info}; use nymsphinx::forwarding::packet::MixPacket; @@ -74,7 +74,7 @@ impl ConnectionHandler { self, conn: TcpStream, remote: SocketAddr, - mut shutdown: ShutdownListener, + mut shutdown: TaskClient, ) { debug!("Starting connection handler for {:?}", remote); let mut framed_conn = Framed::new(conn, SphinxCodec); diff --git a/mixnode/src/node/listener/mod.rs b/mixnode/src/node/listener/mod.rs index 6b8ddbee1a..18ab26bc63 100644 --- a/mixnode/src/node/listener/mod.rs +++ b/mixnode/src/node/listener/mod.rs @@ -8,17 +8,17 @@ use std::process; use tokio::net::TcpListener; use tokio::task::JoinHandle; -use super::ShutdownListener; +use super::TaskClient; pub(crate) mod connection_handler; pub(crate) struct Listener { address: SocketAddr, - shutdown: ShutdownListener, + shutdown: TaskClient, } impl Listener { - pub(crate) fn new(address: SocketAddr, shutdown: ShutdownListener) -> Self { + pub(crate) fn new(address: SocketAddr, shutdown: TaskClient) -> Self { Listener { address, shutdown } } diff --git a/mixnode/src/node/mod.rs b/mixnode/src/node/mod.rs index e986c4fa67..9fdbff3ebb 100644 --- a/mixnode/src/node/mod.rs +++ b/mixnode/src/node/mod.rs @@ -26,7 +26,7 @@ use rand::thread_rng; use std::net::SocketAddr; use std::process; use std::sync::Arc; -use task::{wait_for_signal, ShutdownListener, ShutdownNotifier}; +use task::{wait_for_signal, TaskClient, TaskManager}; use version_checker::parse_version; mod http; @@ -153,7 +153,7 @@ impl MixNode { fn start_node_stats_controller( &self, - shutdown: ShutdownListener, + shutdown: TaskClient, ) -> (SharedNodeStats, node_statistics::UpdateSender) { info!("Starting node stats controller..."); let controller = node_statistics::Controller::new( @@ -171,7 +171,7 @@ impl MixNode { &self, node_stats_update_sender: node_statistics::UpdateSender, delay_forwarding_channel: PacketDelayForwardSender, - shutdown: ShutdownListener, + shutdown: TaskClient, ) { info!("Starting socket listener..."); @@ -191,7 +191,7 @@ impl MixNode { fn start_packet_delay_forwarder( &mut self, node_stats_update_sender: node_statistics::UpdateSender, - shutdown: ShutdownListener, + shutdown: TaskClient, ) -> PacketDelayForwardSender { info!("Starting packet delay-forwarder..."); @@ -215,7 +215,7 @@ impl MixNode { packet_sender } - fn start_verloc_measurements(&self, shutdown: ShutdownListener) -> AtomicVerlocResult { + fn start_verloc_measurements(&self, shutdown: TaskClient) -> AtomicVerlocResult { info!("Starting the round-trip-time measurer..."); // this is a sanity check to make sure we didn't mess up with the minimum version at some point @@ -288,7 +288,7 @@ impl MixNode { .map(|node| node.bond_information.mix_node.identity_key.clone()) } - async fn wait_for_interrupt(&self, mut shutdown: ShutdownNotifier) { + async fn wait_for_interrupt(&self, mut shutdown: TaskManager) { wait_for_signal().await; log::info!("Sending shutdown"); @@ -315,7 +315,7 @@ impl MixNode { } } - let shutdown = ShutdownNotifier::default(); + let shutdown = TaskManager::default(); let (node_stats_pointer, node_stats_update_sender) = self.start_node_stats_controller(shutdown.subscribe()); diff --git a/mixnode/src/node/node_statistics.rs b/mixnode/src/node/node_statistics.rs index 86f98401ab..2fa8c39b8e 100644 --- a/mixnode/src/node/node_statistics.rs +++ b/mixnode/src/node/node_statistics.rs @@ -9,7 +9,7 @@ use std::sync::Arc; use std::time::{Duration, SystemTime}; use tokio::sync::{RwLock, RwLockReadGuard}; -use super::ShutdownListener; +use super::TaskClient; // convenience aliases type PacketsMap = HashMap; @@ -211,14 +211,14 @@ impl CurrentPacketData { struct UpdateHandler { current_data: CurrentPacketData, update_receiver: PacketDataReceiver, - shutdown: ShutdownListener, + shutdown: TaskClient, } impl UpdateHandler { fn new( current_data: CurrentPacketData, update_receiver: PacketDataReceiver, - shutdown: ShutdownListener, + shutdown: TaskClient, ) -> Self { UpdateHandler { current_data, @@ -293,7 +293,7 @@ struct StatsUpdater { updating_delay: Duration, current_packet_data: CurrentPacketData, current_stats: SharedNodeStats, - shutdown: ShutdownListener, + shutdown: TaskClient, } impl StatsUpdater { @@ -301,7 +301,7 @@ impl StatsUpdater { updating_delay: Duration, current_packet_data: CurrentPacketData, current_stats: SharedNodeStats, - shutdown: ShutdownListener, + shutdown: TaskClient, ) -> Self { StatsUpdater { updating_delay, @@ -335,11 +335,11 @@ impl StatsUpdater { struct PacketStatsConsoleLogger { logging_delay: Duration, stats: SharedNodeStats, - shutdown: ShutdownListener, + shutdown: TaskClient, } impl PacketStatsConsoleLogger { - fn new(logging_delay: Duration, stats: SharedNodeStats, shutdown: ShutdownListener) -> Self { + fn new(logging_delay: Duration, stats: SharedNodeStats, shutdown: TaskClient) -> Self { PacketStatsConsoleLogger { logging_delay, stats, @@ -451,7 +451,7 @@ impl Controller { pub(crate) fn new( logging_delay: Duration, stats_updating_delay: Duration, - shutdown: ShutdownListener, + shutdown: TaskClient, ) -> Self { let (sender, receiver) = mpsc::unbounded(); let shared_packet_data = CurrentPacketData::new(); @@ -503,13 +503,13 @@ impl Controller { #[cfg(test)] mod tests { use super::*; - use task::ShutdownNotifier; + use task::TaskManager; #[tokio::test] async fn node_stats_reported_are_received() { let logging_delay = Duration::from_millis(20); let stats_updating_delay = Duration::from_millis(10); - let shutdown = ShutdownNotifier::default(); + let shutdown = TaskManager::default(); let node_stats_controller = Controller::new(logging_delay, stats_updating_delay, shutdown.subscribe()); diff --git a/mixnode/src/node/packet_delayforwarder.rs b/mixnode/src/node/packet_delayforwarder.rs index 81e0303afd..da098bcdf9 100644 --- a/mixnode/src/node/packet_delayforwarder.rs +++ b/mixnode/src/node/packet_delayforwarder.rs @@ -9,7 +9,7 @@ use nymsphinx::forwarding::packet::MixPacket; use std::io; use tokio::time::Instant; -use super::ShutdownListener; +use super::TaskClient; // Delay + MixPacket vs Instant + MixPacket @@ -28,7 +28,7 @@ where packet_sender: PacketDelayForwardSender, packet_receiver: PacketDelayForwardReceiver, node_stats_update_sender: UpdateSender, - shutdown: ShutdownListener, + shutdown: TaskClient, } impl DelayForwarder @@ -38,7 +38,7 @@ where pub(crate) fn new( client: C, node_stats_update_sender: UpdateSender, - shutdown: ShutdownListener, + shutdown: TaskClient, ) -> DelayForwarder { let (packet_sender, packet_receiver) = mpsc::unbounded(); @@ -134,7 +134,7 @@ mod tests { use std::sync::{Arc, Mutex}; use std::time::Duration; - use task::ShutdownNotifier; + use task::TaskManager; use nymsphinx::addressing::nodes::NymNodeRoutingAddress; use nymsphinx_params::packet_sizes::PacketSize; @@ -205,7 +205,7 @@ mod tests { let node_stats_update_sender = UpdateSender::new(stats_sender); let client = TestClient::default(); let client_packets_sent = client.packets_sent.clone(); - let shutdown = ShutdownNotifier::default(); + let shutdown = TaskManager::default(); let mut delay_forwarder = DelayForwarder::new(client, node_stats_update_sender, shutdown.subscribe()); let packet_sender = delay_forwarder.sender(); diff --git a/nym-api/src/coconut/dkg/controller.rs b/nym-api/src/coconut/dkg/controller.rs index 6a7268af47..c54074141d 100644 --- a/nym-api/src/coconut/dkg/controller.rs +++ b/nym-api/src/coconut/dkg/controller.rs @@ -19,7 +19,7 @@ use rand::rngs::OsRng; use rand::RngCore; use std::path::PathBuf; use std::time::Duration; -use task::ShutdownListener; +use task::TaskClient; use tokio::time::interval; use validator_client::nymd::SigningNymdClient; @@ -134,7 +134,7 @@ impl DkgController { } } - pub(crate) async fn run(mut self, mut shutdown: ShutdownListener) { + pub(crate) async fn run(mut self, mut shutdown: TaskClient) { let mut interval = interval(self.polling_rate); while !shutdown.is_shutdown() { tokio::select! { diff --git a/nym-api/src/contract_cache/mod.rs b/nym-api/src/contract_cache/mod.rs index 61a7150a02..0775ca5a85 100644 --- a/nym-api/src/contract_cache/mod.rs +++ b/nym-api/src/contract_cache/mod.rs @@ -21,7 +21,7 @@ use std::ops::Deref; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::time::Duration; -use task::ShutdownListener; +use task::TaskClient; use tokio::sync::{watch, RwLock}; use tokio::time; use validator_client::nymd::CosmWasmClient; @@ -198,7 +198,7 @@ impl ValidatorCacheRefresher { Ok(()) } - pub(crate) async fn run(&self, mut shutdown: ShutdownListener) + pub(crate) async fn run(&self, mut shutdown: TaskClient) where C: CosmWasmClient + Sync + Send, { diff --git a/nym-api/src/epoch_operations/mod.rs b/nym-api/src/epoch_operations/mod.rs index cf0233e453..7f109d3dea 100644 --- a/nym-api/src/epoch_operations/mod.rs +++ b/nym-api/src/epoch_operations/mod.rs @@ -35,7 +35,7 @@ mod helpers; use crate::epoch_operations::helpers::stake_to_f64; use crate::node_status_api::ONE_DAY; use error::RewardingError; -use task::ShutdownListener; +use task::TaskClient; #[derive(Debug, Clone, Copy)] pub(crate) struct MixnodeToReward { @@ -373,7 +373,7 @@ impl RewardedSetUpdater { Ok(()) } - async fn wait_until_epoch_end(&mut self, shutdown: &mut ShutdownListener) -> Option { + async fn wait_until_epoch_end(&mut self, shutdown: &mut TaskClient) -> Option { const POLL_INTERVAL: Duration = Duration::from_secs(120); loop { @@ -421,10 +421,7 @@ impl RewardedSetUpdater { } } - pub(crate) async fn run( - &mut self, - mut shutdown: ShutdownListener, - ) -> Result<(), RewardingError> { + pub(crate) async fn run(&mut self, mut shutdown: TaskClient) -> Result<(), RewardingError> { self.validator_cache.wait_for_initial_values().await; while !shutdown.is_shutdown() { diff --git a/nym-api/src/main.rs b/nym-api/src/main.rs index a47068f5fc..256bb527b9 100644 --- a/nym-api/src/main.rs +++ b/nym-api/src/main.rs @@ -31,7 +31,7 @@ use std::str::FromStr; use std::sync::Arc; use std::time::Duration; use std::{fs, process}; -use task::ShutdownNotifier; +use task::TaskManager; use tokio::sync::Notify; #[cfg(feature = "coconut")] use url::Url; @@ -206,7 +206,7 @@ fn parse_args() -> ArgMatches { base_app.get_matches() } -async fn wait_for_interrupt(mut shutdown: ShutdownNotifier) { +async fn wait_for_interrupt(mut shutdown: TaskManager) { wait_for_signal().await; log::info!("Sending shutdown"); @@ -527,7 +527,7 @@ async fn run_nym_api(matches: ArgMatches) -> Result<()> { let liftoff_notify = Arc::new(Notify::new()); // We need a bigger timeout - let shutdown = ShutdownNotifier::new(10); + let shutdown = TaskManager::new(10); #[cfg(feature = "coconut")] let coconut_keypair = coconut::keypair::KeyPair::new(); diff --git a/nym-api/src/network_monitor/mod.rs b/nym-api/src/network_monitor/mod.rs index 0048e2be02..ff4d6bd585 100644 --- a/nym-api/src/network_monitor/mod.rs +++ b/nym-api/src/network_monitor/mod.rs @@ -6,7 +6,7 @@ use crypto::asymmetric::{encryption, identity}; use futures::channel::mpsc; use gateway_client::bandwidth::BandwidthController; use std::sync::Arc; -use task::ShutdownNotifier; +use task::TaskManager; use validator_client::nymd::SigningNymdClient; use crate::config::Config; @@ -142,7 +142,7 @@ impl NetworkMonitorRunnables { // TODO: note, that is not exactly doing what we want, because when // `ReceivedProcessor` is constructed, it already spawns a future // this needs to be refactored! - pub(crate) fn spawn_tasks(self, shutdown: &ShutdownNotifier) { + pub(crate) fn spawn_tasks(self, shutdown: &TaskManager) { let mut packet_receiver = self.packet_receiver; let mut monitor = self.monitor; let shutdown_listener = shutdown.subscribe(); diff --git a/nym-api/src/network_monitor/monitor/gateways_pinger.rs b/nym-api/src/network_monitor/monitor/gateways_pinger.rs index 24cd581bf3..986c1db5fb 100644 --- a/nym-api/src/network_monitor/monitor/gateways_pinger.rs +++ b/nym-api/src/network_monitor/monitor/gateways_pinger.rs @@ -7,7 +7,7 @@ use crypto::asymmetric::identity; use crypto::asymmetric::identity::PUBLIC_KEY_LENGTH; use log::{debug, info, trace, warn}; use std::time::Duration; -use task::ShutdownListener; +use task::TaskClient; use tokio::time::{sleep, Instant}; // TODO: should it perhaps be moved to config along other timeout values? @@ -144,7 +144,7 @@ impl GatewayPinger { info!("Pinging all active gateways took {:?}", time_taken); } - pub(crate) async fn run(&self, mut shutdown: ShutdownListener) { + pub(crate) async fn run(&self, mut shutdown: TaskClient) { while !shutdown.is_shutdown() { tokio::select! { _ = sleep(self.pinging_interval) => { diff --git a/nym-api/src/network_monitor/monitor/mod.rs b/nym-api/src/network_monitor/monitor/mod.rs index c333c5e883..f0f17bad09 100644 --- a/nym-api/src/network_monitor/monitor/mod.rs +++ b/nym-api/src/network_monitor/monitor/mod.rs @@ -12,7 +12,7 @@ use crate::storage::NymApiStorage; use log::{debug, error, info}; use std::collections::{HashMap, HashSet}; use std::process; -use task::ShutdownListener; +use task::TaskClient; use tokio::time::{sleep, Duration, Instant}; pub(crate) mod gateway_clients_cache; @@ -301,7 +301,7 @@ impl Monitor { self.test_nonce += 1; } - pub(crate) async fn run(&mut self, mut shutdown: ShutdownListener) { + pub(crate) async fn run(&mut self, mut shutdown: TaskClient) { self.received_processor.start_receiving(); // wait for validator cache to be ready diff --git a/nym-api/src/network_monitor/monitor/receiver.rs b/nym-api/src/network_monitor/monitor/receiver.rs index 991da5fa05..c232458c6f 100644 --- a/nym-api/src/network_monitor/monitor/receiver.rs +++ b/nym-api/src/network_monitor/monitor/receiver.rs @@ -7,7 +7,7 @@ use crypto::asymmetric::identity; use futures::channel::mpsc; use futures::StreamExt; use gateway_client::{AcknowledgementReceiver, MixnetMessageReceiver}; -use task::ShutdownListener; +use task::TaskClient; pub(crate) type GatewayClientUpdateSender = mpsc::UnboundedSender; pub(crate) type GatewayClientUpdateReceiver = mpsc::UnboundedReceiver; @@ -56,7 +56,7 @@ impl PacketReceiver { .expect("packet processor seems to have crashed!"); } - pub(crate) async fn run(&mut self, mut shutdown: ShutdownListener) { + pub(crate) async fn run(&mut self, mut shutdown: TaskClient) { while !shutdown.is_shutdown() { tokio::select! { biased; diff --git a/nym-api/src/network_monitor/monitor/sender.rs b/nym-api/src/network_monitor/monitor/sender.rs index d2fa815a11..5c6f3df22d 100644 --- a/nym-api/src/network_monitor/monitor/sender.rs +++ b/nym-api/src/network_monitor/monitor/sender.rs @@ -24,7 +24,7 @@ use std::pin::Pin; use std::sync::Arc; use std::task::Poll; use std::time::Duration; -use task::ShutdownListener; +use task::TaskClient; use gateway_client::bandwidth::BandwidthController; @@ -166,11 +166,7 @@ impl PacketSender { } } - pub(crate) fn spawn_gateways_pinger( - &self, - pinging_interval: Duration, - shutdown: ShutdownListener, - ) { + pub(crate) fn spawn_gateways_pinger(&self, pinging_interval: Duration, shutdown: TaskClient) { let gateway_pinger = GatewayPinger::new( self.active_gateway_clients.clone(), self.fresh_gateway_client_data @@ -209,7 +205,7 @@ impl PacketSender { ack_sender, fresh_gateway_client_data.gateway_response_timeout, Some(fresh_gateway_client_data.bandwidth_controller.clone()), - task::ShutdownListener::dummy(), + task::TaskClient::dummy(), ); gateway_client diff --git a/nym-api/src/node_status_api/cache.rs b/nym-api/src/node_status_api/cache.rs index f485716e8e..5a76ffd85a 100644 --- a/nym-api/src/node_status_api/cache.rs +++ b/nym-api/src/node_status_api/cache.rs @@ -12,7 +12,7 @@ use nym_api_requests::models::{MixNodeBondAnnotated, MixnodeStatus}; use rocket::fairing::AdHoc; use std::collections::HashMap; use std::{sync::Arc, time::Duration}; -use task::ShutdownListener; +use task::TaskClient; use tokio::sync::RwLockReadGuard; use tokio::{ sync::{watch, RwLock}, @@ -174,7 +174,7 @@ impl NodeStatusCacheRefresher { } } - pub async fn run(&mut self, mut shutdown: ShutdownListener) { + pub async fn run(&mut self, mut shutdown: TaskClient) { let mut fallback_interval = time::interval(self.fallback_caching_interval); while !shutdown.is_shutdown() { tokio::select! { diff --git a/nym-api/src/node_status_api/uptime_updater.rs b/nym-api/src/node_status_api/uptime_updater.rs index cc22d75741..0c36c1a124 100644 --- a/nym-api/src/node_status_api/uptime_updater.rs +++ b/nym-api/src/node_status_api/uptime_updater.rs @@ -8,7 +8,7 @@ use crate::node_status_api::ONE_DAY; use crate::storage::NymApiStorage; use log::error; use std::time::Duration; -use task::ShutdownListener; +use task::TaskClient; use time::{OffsetDateTime, PrimitiveDateTime, Time}; use tokio::time::{interval, sleep}; @@ -68,7 +68,7 @@ impl HistoricalUptimeUpdater { Ok(()) } - pub(crate) async fn run(&self, mut shutdown: ShutdownListener) { + pub(crate) async fn run(&self, mut shutdown: TaskClient) { // update uptimes at 23:00 UTC each day so that we'd have data from the actual [almost] whole day // and so that we would avoid the edge case of starting validator API at 23:59 and having some // nodes update for different days diff --git a/service-providers/network-requester/src/connection.rs b/service-providers/network-requester/src/connection.rs index f4ffaca992..3457375737 100644 --- a/service-providers/network-requester/src/connection.rs +++ b/service-providers/network-requester/src/connection.rs @@ -7,7 +7,7 @@ use proxy_helpers::connection_controller::ConnectionReceiver; use proxy_helpers::proxy_runner::{MixProxySender, ProxyRunner}; use socks5_requests::{ConnectionId, Message as Socks5Message, RemoteAddress, Response}; use std::io; -use task::ShutdownListener; +use task::TaskClient; use tokio::net::TcpStream; /// A TCP connection between the Socks5 service provider, which makes @@ -42,7 +42,7 @@ impl Connection { mix_receiver: ConnectionReceiver, mix_sender: MixProxySender<(Socks5Message, ReturnAddress)>, lane_queue_lengths: LaneQueueLengths, - shutdown: ShutdownListener, + shutdown: TaskClient, ) { let stream = self.conn.take().unwrap(); let remote_source_address = "???".to_string(); // we don't know ip address of requester diff --git a/service-providers/network-requester/src/core.rs b/service-providers/network-requester/src/core.rs index 202ac3d1c1..42ded9fbd0 100644 --- a/service-providers/network-requester/src/core.rs +++ b/service-providers/network-requester/src/core.rs @@ -28,7 +28,7 @@ use socks5_requests::{ use statistics_common::collector::StatisticsSender; use std::path::PathBuf; use std::sync::atomic::{AtomicUsize, Ordering}; -use task::ShutdownListener; +use task::TaskClient; use tokio_tungstenite::tungstenite::protocol::Message; use websocket_requests::{requests::ClientRequest, responses::ServerResponse}; @@ -250,7 +250,7 @@ impl ServiceProvider { controller_sender: ControllerSender, mix_input_sender: MixProxySender<(Socks5Message, ReturnAddress)>, lane_queue_lengths: LaneQueueLengths, - shutdown: ShutdownListener, + shutdown: TaskClient, ) { let mut conn = match Connection::new(conn_id, remote_addr.clone(), return_address.clone()).await { @@ -313,7 +313,7 @@ impl ServiceProvider { lane_queue_lengths: LaneQueueLengths, sender_tag: Option, connect_req: Box, - shutdown: ShutdownListener, + shutdown: TaskClient, ) { let return_address = match ReturnAddress::new(connect_req.return_address, sender_tag) { Some(address) => address, @@ -379,7 +379,7 @@ impl ServiceProvider { mix_input_sender: &MixProxySender<(Socks5Message, ReturnAddress)>, lane_queue_lengths: LaneQueueLengths, stats_collector: Option, - shutdown: ShutdownListener, + shutdown: TaskClient, ) { let deserialized_msg = match Socks5Message::try_from_bytes(&message.message) { Ok(msg) => msg, @@ -445,7 +445,7 @@ impl ServiceProvider { tokio::sync::mpsc::channel::<(Socks5Message, ReturnAddress)>(1); // Used to notify tasks to shutdown. Not all tasks fully supports this (yet). - let shutdown = task::ShutdownNotifier::default(); + let shutdown = task::TaskManager::default(); // Channel for announcing client connection state by the controller. // The `mixnet_response_listener` will use this to either report closed connection to the