diff --git a/clients/native/src/websocket/handler.rs b/clients/native/src/websocket/handler.rs index 3d199cc152..e26df51699 100644 --- a/clients/native/src/websocket/handler.rs +++ b/clients/native/src/websocket/handler.rs @@ -125,10 +125,17 @@ impl Handler { }; // get the number of pending replies waiting for reply surbs - let reply_queue_length = self + let reply_queue_length = match self .reply_controller_sender .get_lane_queue_length(connection_id) - .await; + .await + { + Ok(length) => length, + Err(err) => { + error!("Failed to get reply queue length for connection {connection_id}: {err}"); + 0 + } + }; let queue_length = base_length + reply_queue_length; diff --git a/common/client-core/src/client/base_client/mod.rs b/common/client-core/src/client/base_client/mod.rs index 61c40d9a16..bcaf48dcf7 100644 --- a/common/client-core/src/client/base_client/mod.rs +++ b/common/client-core/src/client/base_client/mod.rs @@ -384,8 +384,9 @@ where reply_key_storage, reply_controller_sender, metrics_reporter, + shutdown, ); - controller.start_with_shutdown(shutdown) + controller.start() } #[allow(clippy::too_many_arguments)] diff --git a/common/client-core/src/client/real_messages_control/acknowledgement_control/acknowledgement_listener.rs b/common/client-core/src/client/real_messages_control/acknowledgement_control/acknowledgement_listener.rs index abc8c72932..4b5d483ee5 100644 --- a/common/client-core/src/client/real_messages_control/acknowledgement_control/acknowledgement_listener.rs +++ b/common/client-core/src/client/real_messages_control/acknowledgement_control/acknowledgement_listener.rs @@ -11,6 +11,7 @@ use nym_sphinx::{ acknowledgements::{identifier::recover_identifier, AckKey}, chunking::fragment::{FragmentIdentifier, COVER_FRAG_ID}, }; +use nym_task::TaskClient; use std::sync::Arc; /// Module responsible for listening for any data resembling acknowledgements from the network @@ -20,6 +21,7 @@ pub(super) struct AcknowledgementListener { ack_receiver: AcknowledgementReceiver, action_sender: AckActionSender, stats_tx: ClientStatsSender, + task_client: TaskClient, } impl AcknowledgementListener { @@ -28,12 +30,14 @@ impl AcknowledgementListener { ack_receiver: AcknowledgementReceiver, action_sender: AckActionSender, stats_tx: ClientStatsSender, + task_client: TaskClient, ) -> Self { AcknowledgementListener { ack_key, ack_receiver, action_sender, stats_tx, + task_client, } } @@ -68,7 +72,9 @@ impl AcknowledgementListener { .action_sender .unbounded_send(Action::new_remove(frag_id)) { - error!("Failed to send remove action to action controller: {err}"); + if !self.task_client.is_shutdown_poll() { + error!("Failed to send remove action to action controller: {err}"); + } } } @@ -79,10 +85,10 @@ impl AcknowledgementListener { } } - pub(super) async fn run_with_shutdown(&mut self, mut shutdown: nym_task::TaskClient) { + pub(super) async fn run(&mut self) { debug!("Started AcknowledgementListener with graceful shutdown support"); - while !shutdown.is_shutdown() { + while !self.task_client.is_shutdown() { tokio::select! { acks = self.ack_receiver.next() => match acks { Some(acks) => self.handle_ack_receiver_item(acks).await, @@ -91,12 +97,12 @@ impl AcknowledgementListener { break; } }, - _ = shutdown.recv() => { + _ = self.task_client.recv() => { log::trace!("AcknowledgementListener: Received shutdown"); } } } - shutdown.recv_timeout().await; + self.task_client.recv_timeout().await; log::debug!("AcknowledgementListener: Exiting"); } } diff --git a/common/client-core/src/client/real_messages_control/acknowledgement_control/action_controller.rs b/common/client-core/src/client/real_messages_control/acknowledgement_control/action_controller.rs index 2ad7e5e753..2aca3d58c7 100644 --- a/common/client-core/src/client/real_messages_control/acknowledgement_control/action_controller.rs +++ b/common/client-core/src/client/real_messages_control/acknowledgement_control/action_controller.rs @@ -9,6 +9,7 @@ use log::*; use nym_nonexhaustive_delayqueue::{Expired, NonExhaustiveDelayQueue, QueueKey}; use nym_sphinx::chunking::fragment::FragmentIdentifier; use nym_sphinx::Delay as SphinxDelay; +use nym_task::TaskClient; use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; @@ -101,6 +102,8 @@ pub(super) struct ActionController { /// Channel for notifying `RetransmissionRequestListener` about expired acknowledgements. retransmission_sender: RetransmissionRequestSender, + + task_client: TaskClient, } impl ActionController { @@ -108,6 +111,7 @@ impl ActionController { config: Config, retransmission_sender: RetransmissionRequestSender, incoming_actions: AckActionReceiver, + task_client: TaskClient, ) -> Self { ActionController { config, @@ -115,6 +119,7 @@ impl ActionController { pending_acks_timers: NonExhaustiveDelayQueue::new(), incoming_actions, retransmission_sender, + task_client, } } @@ -216,11 +221,7 @@ impl ActionController { } // note: when the entry expires it's automatically removed from pending_acks_timers - fn handle_expired_ack_timer( - &mut self, - expired_ack: Expired, - task_client: &mut nym_task::TaskClient, - ) { + fn handle_expired_ack_timer(&mut self, expired_ack: Expired) { // I'm honestly not sure how to handle it, because getting it means other things in our // system are already misbehaving. If we ever see this panic, then I guess we should worry // about it. Perhaps just reschedule it at later point? @@ -238,15 +239,13 @@ impl ActionController { // downgrading an arc and then upgrading vs cloning is difference of 30ns vs 15ns // so it's literally a NO difference while it might prevent us from unnecessarily // resending data (in maybe 1 in 1 million cases, but it's something) - if self + if let Err(err) = self .retransmission_sender .unbounded_send(Arc::downgrade(pending_ack_data)) - .is_err() { - assert!( - task_client.is_shutdown_poll(), - "Failed to send pending ack for retransmission" - ); + if !self.task_client.is_shutdown_poll() { + log::error!("Failed to send pending ack for retransmission: {err}"); + } } } else { // this shouldn't cause any issues but shouldn't have happened to begin with! @@ -265,10 +264,10 @@ impl ActionController { } } - pub(super) async fn run_with_shutdown(&mut self, mut shutdown: nym_task::TaskClient) { + pub(super) async fn run(&mut self) { debug!("Started ActionController with graceful shutdown support"); - while !shutdown.is_shutdown() { + while !self.task_client.is_shutdown() { tokio::select! { action = self.incoming_actions.next() => match action { Some(action) => self.process_action(action), @@ -280,19 +279,19 @@ impl ActionController { } }, expired_ack = self.pending_acks_timers.next() => match expired_ack { - Some(expired_ack) => self.handle_expired_ack_timer(expired_ack, &mut shutdown), + Some(expired_ack) => self.handle_expired_ack_timer(expired_ack), None => { log::trace!("ActionController: Stopping since ack channel closed"); break; } }, - _ = shutdown.recv() => { + _ = self.task_client.recv() => { log::trace!("ActionController: Received shutdown"); break; } } } - shutdown.recv_timeout().await; + self.task_client.recv_timeout().await; log::debug!("ActionController: Exiting"); } } diff --git a/common/client-core/src/client/real_messages_control/acknowledgement_control/input_message_listener.rs b/common/client-core/src/client/real_messages_control/acknowledgement_control/input_message_listener.rs index 02bdab245d..566fe87402 100644 --- a/common/client-core/src/client/real_messages_control/acknowledgement_control/input_message_listener.rs +++ b/common/client-core/src/client/real_messages_control/acknowledgement_control/input_message_listener.rs @@ -11,6 +11,7 @@ use nym_sphinx::anonymous_replies::requests::AnonymousSenderTag; use nym_sphinx::forwarding::packet::MixPacket; use nym_sphinx::params::PacketType; use nym_task::connections::TransmissionLane; +use nym_task::TaskClient; use rand::{CryptoRng, Rng}; /// Module responsible for dealing with the received messages: splitting them, creating acknowledgements, @@ -23,6 +24,7 @@ where input_receiver: InputMessageReceiver, message_handler: MessageHandler, reply_controller_sender: ReplyControllerSender, + task_client: TaskClient, } impl InputMessageListener @@ -36,11 +38,13 @@ where input_receiver: InputMessageReceiver, message_handler: MessageHandler, reply_controller_sender: ReplyControllerSender, + task_client: TaskClient, ) -> Self { InputMessageListener { input_receiver, message_handler, reply_controller_sender, + task_client, } } @@ -63,8 +67,14 @@ where lane: TransmissionLane, ) { // offload reply handling to the dedicated task - self.reply_controller_sender + if let Err(err) = self + .reply_controller_sender .send_reply(recipient_tag, data, lane) + { + if !self.task_client.is_shutdown_poll() { + error!("failed to send a reply - {err}"); + } + } } async fn handle_plain_message( @@ -164,10 +174,10 @@ where }; } - pub(super) async fn run_with_shutdown(&mut self, mut shutdown: nym_task::TaskClient) { + pub(super) async fn run(&mut self) { debug!("Started InputMessageListener with graceful shutdown support"); - while !shutdown.is_shutdown() { + while !self.task_client.is_shutdown() { tokio::select! { input_msg = self.input_receiver.recv() => match input_msg { Some(input_msg) => { @@ -178,12 +188,12 @@ where break; } }, - _ = shutdown.recv() => { + _ = self.task_client.recv() => { log::trace!("InputMessageListener: Received shutdown"); } } } - shutdown.recv_timeout().await; + self.task_client.recv_timeout().await; log::debug!("InputMessageListener: Exiting"); } } diff --git a/common/client-core/src/client/real_messages_control/acknowledgement_control/mod.rs b/common/client-core/src/client/real_messages_control/acknowledgement_control/mod.rs index ab2f4e86cb..a6ab7d74a6 100644 --- a/common/client-core/src/client/real_messages_control/acknowledgement_control/mod.rs +++ b/common/client-core/src/client/real_messages_control/acknowledgement_control/mod.rs @@ -67,7 +67,7 @@ pub(crate) enum PacketDestination { /// Structure representing a data `Fragment` that is on-route to the specified `Recipient` #[derive(Debug)] -pub(crate) struct PendingAcknowledgement { +pub struct PendingAcknowledgement { message_chunk: Fragment, delay: SphinxDelay, destination: PacketDestination, @@ -204,7 +204,6 @@ where retransmission_request_listener: RetransmissionRequestListener, sent_notification_listener: SentNotificationListener, action_controller: ActionController, - task_client: TaskClient, } impl AcknowledgementController @@ -228,6 +227,7 @@ where action_config, retransmission_tx, connectors.ack_action_receiver, + task_client.fork("action_controller"), ); // will listen for any acks coming from the network @@ -236,6 +236,7 @@ where connectors.ack_receiver, connectors.ack_action_sender.clone(), stats_tx, + task_client.fork("acknowledgement_listener"), ); // will listen for any new messages from the client @@ -243,6 +244,7 @@ where connectors.input_receiver, message_handler.clone(), reply_controller_sender.clone(), + task_client.fork("input_message_listener"), ); // will listen for any ack timeouts and trigger retransmission @@ -252,12 +254,16 @@ where message_handler, retransmission_rx, reply_controller_sender, + task_client.fork("retransmission_request_listener"), ); // will listen for events indicating the packet was sent through the network so that // the retransmission timer should be started. - let sent_notification_listener = - SentNotificationListener::new(connectors.sent_notifier, connectors.ack_action_sender); + let sent_notification_listener = SentNotificationListener::new( + connectors.sent_notifier, + connectors.ack_action_sender, + task_client.fork("sent_notification_listener"), + ); AcknowledgementController { acknowledgement_listener, @@ -265,7 +271,6 @@ where retransmission_request_listener, sent_notification_listener, action_controller, - task_client, } } @@ -276,42 +281,28 @@ where let mut sent_notification_listener = self.sent_notification_listener; let mut action_controller = self.action_controller; - let shutdown_handle = self.task_client.fork("acknowledgement_listener"); spawn_future(async move { - acknowledgement_listener - .run_with_shutdown(shutdown_handle) - .await; + acknowledgement_listener.run().await; debug!("The acknowledgement listener has finished execution!"); }); - let shutdown_handle = self.task_client.fork("input_message_listener"); spawn_future(async move { - input_message_listener - .run_with_shutdown(shutdown_handle) - .await; + input_message_listener.run().await; debug!("The input listener has finished execution!"); }); - let shutdown_handle = self.task_client.fork("retransmission_request_listener"); spawn_future(async move { - retransmission_request_listener - .run_with_shutdown(shutdown_handle, packet_type) - .await; + retransmission_request_listener.run(packet_type).await; debug!("The retransmission request listener has finished execution!"); }); - let shutdown_handle = self.task_client.fork("sent_notification_listener"); spawn_future(async move { - sent_notification_listener - .run_with_shutdown(shutdown_handle) - .await; + sent_notification_listener.run().await; debug!("The sent notification listener has finished execution!"); }); spawn_future(async move { - action_controller - .run_with_shutdown(self.task_client.with_suffix("action_controller")) - .await; + action_controller.run().await; debug!("The controller has finished execution!"); }); } diff --git a/common/client-core/src/client/real_messages_control/acknowledgement_control/retransmission_request_listener.rs b/common/client-core/src/client/real_messages_control/acknowledgement_control/retransmission_request_listener.rs index cd729e2cf7..1c6feabf52 100644 --- a/common/client-core/src/client/real_messages_control/acknowledgement_control/retransmission_request_listener.rs +++ b/common/client-core/src/client/real_messages_control/acknowledgement_control/retransmission_request_listener.rs @@ -14,7 +14,7 @@ use log::*; use nym_sphinx::chunking::fragment::Fragment; use nym_sphinx::preparer::PreparedFragment; use nym_sphinx::{addressing::clients::Recipient, params::PacketType}; -use nym_task::connections::TransmissionLane; +use nym_task::{connections::TransmissionLane, TaskClient}; use rand::{CryptoRng, Rng}; use std::sync::{Arc, Weak}; @@ -25,6 +25,7 @@ pub(super) struct RetransmissionRequestListener { message_handler: MessageHandler, request_receiver: RetransmissionRequestReceiver, reply_controller_sender: ReplyControllerSender, + task_client: TaskClient, } impl RetransmissionRequestListener @@ -37,6 +38,7 @@ where message_handler: MessageHandler, request_receiver: RetransmissionRequestReceiver, reply_controller_sender: ReplyControllerSender, + task_client: TaskClient, ) -> Self { RetransmissionRequestListener { maximum_retransmissions, @@ -44,6 +46,7 @@ where message_handler, request_receiver, reply_controller_sender, + task_client, } } @@ -96,11 +99,16 @@ where } => { // if this is retransmission for reply, offload it to the dedicated task // that deals with all the surbs - return self.reply_controller_sender.send_retransmission_data( + if let Err(err) = self.reply_controller_sender.send_retransmission_data( *recipient_tag, weak_timed_out_ack, *extra_surb_request, - ); + ) { + if !self.task_client.is_shutdown_poll() { + error!("Failed to send retransmission data to the reply controller: {err}"); + } + } + return; } PacketDestination::KnownRecipient(recipient) => { self.prepare_normal_retransmission_chunk( @@ -151,7 +159,9 @@ where .action_sender .unbounded_send(Action::new_update_pending_ack(frag_id, new_delay)) { - error!("Failed to send update pending ack action to the controller: {err}"); + if !self.task_client.is_shutdown_poll() { + error!("Failed to send update pending ack action to the controller: {err}"); + } } // send to `OutQueueControl` to eventually send to the mix network @@ -166,14 +176,10 @@ where .await } - pub(super) async fn run_with_shutdown( - &mut self, - mut shutdown: nym_task::TaskClient, - packet_type: PacketType, - ) { + pub(super) async fn run(&mut self, packet_type: PacketType) { debug!("Started RetransmissionRequestListener with graceful shutdown support"); - while !shutdown.is_shutdown() { + while !self.task_client.is_shutdown() { tokio::select! { timed_out_ack = self.request_receiver.next() => match timed_out_ack { Some(timed_out_ack) => self.on_retransmission_request(timed_out_ack, packet_type).await, @@ -182,12 +188,12 @@ where break; } }, - _ = shutdown.recv() => { + _ = self.task_client.recv() => { log::trace!("RetransmissionRequestListener: Received shutdown"); } } } - shutdown.recv_timeout().await; + self.task_client.recv_timeout().await; log::debug!("RetransmissionRequestListener: Exiting"); } } diff --git a/common/client-core/src/client/real_messages_control/acknowledgement_control/sent_notification_listener.rs b/common/client-core/src/client/real_messages_control/acknowledgement_control/sent_notification_listener.rs index f492a4426a..0b563c375d 100644 --- a/common/client-core/src/client/real_messages_control/acknowledgement_control/sent_notification_listener.rs +++ b/common/client-core/src/client/real_messages_control/acknowledgement_control/sent_notification_listener.rs @@ -6,6 +6,7 @@ use super::SentPacketNotificationReceiver; use futures::StreamExt; use log::*; use nym_sphinx::chunking::fragment::{FragmentIdentifier, COVER_FRAG_ID}; +use nym_task::TaskClient; /// Module responsible for starting up retransmission timers. /// It is required because when we send our packet to the `real traffic stream` controlled @@ -14,16 +15,19 @@ use nym_sphinx::chunking::fragment::{FragmentIdentifier, COVER_FRAG_ID}; pub(super) struct SentNotificationListener { sent_notifier: SentPacketNotificationReceiver, action_sender: AckActionSender, + task_client: TaskClient, } impl SentNotificationListener { pub(super) fn new( sent_notifier: SentPacketNotificationReceiver, action_sender: AckActionSender, + task_client: TaskClient, ) -> Self { SentNotificationListener { sent_notifier, action_sender, + task_client, } } @@ -36,14 +40,16 @@ impl SentNotificationListener { .action_sender .unbounded_send(Action::new_start_timer(frag_id)) { - error!("Failed to send start timer action to action controller: {err}"); + if !self.task_client.is_shutdown_poll() { + error!("Failed to send start timer action to action controller: {err}"); + } } } - pub(super) async fn run_with_shutdown(&mut self, mut shutdown: nym_task::TaskClient) { + pub(super) async fn run(&mut self) { debug!("Started SentNotificationListener with graceful shutdown support"); - while !shutdown.is_shutdown() { + while !self.task_client.is_shutdown() { tokio::select! { frag_id = self.sent_notifier.next() => match frag_id { Some(frag_id) => { @@ -54,13 +60,13 @@ impl SentNotificationListener { break; } }, - _ = shutdown.recv() => { + _ = self.task_client.recv() => { log::trace!("SentNotificationListener: Received shutdown"); break; } } } - assert!(shutdown.is_shutdown_poll()); + assert!(self.task_client.is_shutdown_poll()); log::debug!("SentNotificationListener: Exiting"); } } diff --git a/common/client-core/src/client/real_messages_control/real_traffic_stream.rs b/common/client-core/src/client/real_messages_control/real_traffic_stream.rs index e2b3045521..2789228dcf 100644 --- a/common/client-core/src/client/real_messages_control/real_traffic_stream.rs +++ b/common/client-core/src/client/real_messages_control/real_traffic_stream.rs @@ -224,10 +224,7 @@ where } } - async fn on_message( - &mut self, - next_message: StreamMessage, - ) { + async fn on_message(&mut self, next_message: StreamMessage) { trace!("created new message"); let (next_message, fragment_id, packet_size) = match next_message { diff --git a/common/client-core/src/client/received_buffer.rs b/common/client-core/src/client/received_buffer.rs index cd3e5585ae..a5a206dd64 100644 --- a/common/client-core/src/client/received_buffer.rs +++ b/common/client-core/src/client/received_buffer.rs @@ -20,6 +20,7 @@ use nym_sphinx::message::{NymMessage, PlainMessage}; use nym_sphinx::params::ReplySurbKeyDigestAlgorithm; use nym_sphinx::receiver::{MessageReceiver, MessageRecoveryError, ReconstructedMessage}; use nym_statistics_common::clients::{packet_statistics::PacketStatisticsEvent, ClientStatsSender}; +use nym_task::TaskClient; use std::collections::HashSet; use std::sync::Arc; @@ -152,6 +153,7 @@ struct ReceivedMessagesBuffer { inner: Arc>>, reply_key_storage: SentReplyKeys, reply_controller_sender: ReplyControllerSender, + task_client: TaskClient, } impl ReceivedMessagesBuffer { @@ -160,6 +162,7 @@ impl ReceivedMessagesBuffer { reply_key_storage: SentReplyKeys, reply_controller_sender: ReplyControllerSender, stats_tx: ClientStatsSender, + task_client: TaskClient, ) -> Self { ReceivedMessagesBuffer { inner: Arc::new(Mutex::new(ReceivedMessagesBufferInner { @@ -172,6 +175,7 @@ impl ReceivedMessagesBuffer { })), reply_key_storage, reply_controller_sender, + task_client, } } @@ -257,11 +261,15 @@ impl ReceivedMessagesBuffer { } }; - self.reply_controller_sender.send_additional_surbs( + if let Err(err) = self.reply_controller_sender.send_additional_surbs( msg.sender_tag, reply_surbs, from_surb_request, - ) + ) { + if !self.task_client.is_shutdown_poll() { + error!("{err}"); + } + } } reconstructed } @@ -276,8 +284,14 @@ impl ReceivedMessagesBuffer { ReplyMessageContent::Data { message } => reconstructed.push(message.into()), ReplyMessageContent::SurbRequest { recipient, amount } => { debug!("received request for {amount} additional reply SURBs from {recipient}"); - self.reply_controller_sender - .send_additional_surbs_request(*recipient, amount); + if let Err(err) = self + .reply_controller_sender + .send_additional_surbs_request(*recipient, amount) + { + if !self.task_client.is_shutdown_poll() { + error!("{err}"); + } + } } } } @@ -399,16 +413,19 @@ pub enum ReceivedBufferMessage { struct RequestReceiver { received_buffer: ReceivedMessagesBuffer, query_receiver: ReceivedBufferRequestReceiver, + task_client: TaskClient, } impl RequestReceiver { fn new( received_buffer: ReceivedMessagesBuffer, query_receiver: ReceivedBufferRequestReceiver, + task_client: TaskClient, ) -> Self { RequestReceiver { received_buffer, query_receiver, + task_client, } } @@ -423,12 +440,12 @@ impl RequestReceiver { } } - async fn run_with_shutdown(&mut self, mut shutdown: nym_task::TaskClient) { + async fn run(&mut self) { debug!("Started RequestReceiver with graceful shutdown support"); - while !shutdown.is_shutdown() { + while !self.task_client.is_shutdown() { tokio::select! { biased; - _ = shutdown.recv() => { + _ = self.task_client.recv() => { log::trace!("RequestReceiver: Received shutdown"); } request = self.query_receiver.next() => { @@ -441,7 +458,7 @@ impl RequestReceiver { }, } } - shutdown.recv().await; + self.task_client.recv().await; log::debug!("RequestReceiver: Exiting"); } } @@ -449,25 +466,25 @@ impl RequestReceiver { struct FragmentedMessageReceiver { received_buffer: ReceivedMessagesBuffer, mixnet_packet_receiver: MixnetMessageReceiver, + task_client: TaskClient, } impl FragmentedMessageReceiver { fn new( received_buffer: ReceivedMessagesBuffer, mixnet_packet_receiver: MixnetMessageReceiver, + task_client: TaskClient, ) -> Self { FragmentedMessageReceiver { received_buffer, mixnet_packet_receiver, + task_client, } } - async fn run_with_shutdown( - &mut self, - mut shutdown: nym_task::TaskClient, - ) -> Result<(), MessageRecoveryError> { + async fn run(&mut self) -> Result<(), MessageRecoveryError> { debug!("Started FragmentedMessageReceiver with graceful shutdown support"); - while !shutdown.is_shutdown() { + while !self.task_client.is_shutdown() { tokio::select! { new_messages = self.mixnet_packet_receiver.next() => { if let Some(new_messages) = new_messages { @@ -477,12 +494,12 @@ impl FragmentedMessageReceiver { break; } }, - _ = shutdown.recv_with_delay() => { + _ = self.task_client.recv_with_delay() => { log::trace!("FragmentedMessageReceiver: Received shutdown"); } } } - shutdown.recv_timeout().await; + self.task_client.recv_timeout().await; log::debug!("FragmentedMessageReceiver: Exiting"); Ok(()) } @@ -501,41 +518,42 @@ impl ReceivedMessagesBufferControll reply_key_storage: SentReplyKeys, reply_controller_sender: ReplyControllerSender, metrics_reporter: ClientStatsSender, + task_client: TaskClient, ) -> Self { let received_buffer = ReceivedMessagesBuffer::new( local_encryption_keypair, reply_key_storage, reply_controller_sender, metrics_reporter, + task_client.fork("received_messages_buffer"), ); ReceivedMessagesBufferController { fragmented_message_receiver: FragmentedMessageReceiver::new( received_buffer.clone(), mixnet_packet_receiver, + task_client.fork("fragmented_message_receiver"), + ), + request_receiver: RequestReceiver::new( + received_buffer, + query_receiver, + task_client.with_suffix("request_receiver"), ), - request_receiver: RequestReceiver::new(received_buffer, query_receiver), } } - pub fn start_with_shutdown(self, shutdown: nym_task::TaskClient) { + pub fn start(self) { let mut fragmented_message_receiver = self.fragmented_message_receiver; let mut request_receiver = self.request_receiver; - let shutdown_handle = shutdown.fork("fragmented_message_receiver"); spawn_future(async move { - match fragmented_message_receiver - .run_with_shutdown(shutdown_handle) - .await - { + match fragmented_message_receiver.run().await { Ok(_) => {} Err(e) => error!("{e}"), } }); spawn_future(async move { - request_receiver - .run_with_shutdown(shutdown.with_suffix("request_receiver")) - .await; + request_receiver.run().await; }); } } diff --git a/common/client-core/src/client/replies/reply_controller/requests.rs b/common/client-core/src/client/replies/reply_controller/requests.rs index ffb6486023..a1555329e5 100644 --- a/common/client-core/src/client/replies/reply_controller/requests.rs +++ b/common/client-core/src/client/replies/reply_controller/requests.rs @@ -15,6 +15,27 @@ pub(crate) fn new_control_channels() -> (ReplyControllerSender, ReplyControllerR (tx.into(), rx) } +#[derive(Debug, thiserror::Error)] +pub enum ReplyControllerSenderError { + #[error("failed to send retransmission data to reply controller")] + SendRetransmissionData(#[source] mpsc::TrySendError), + + #[error("failed to send reply to reply controller")] + SendReply(#[source] mpsc::TrySendError), + + #[error("failed to send additional surbs to reply controller")] + AdditionalSurbs(#[source] mpsc::TrySendError), + + #[error("failed to send additional surbs request to reply controller")] + AdditionalSurbsRequest(#[source] mpsc::TrySendError), + + #[error("failed to request lane queue length from reply controller")] + LaneQueueLength(#[source] mpsc::TrySendError), + + #[error("response channel was dropped before we could receive the response")] + ResponseChannelDropped(#[source] oneshot::Canceled), +} + #[derive(Debug, Clone)] pub struct ReplyControllerSender(mpsc::UnboundedSender); @@ -30,17 +51,14 @@ impl ReplyControllerSender { recipient: AnonymousSenderTag, timed_out_ack: Weak, extra_surb_request: bool, - ) { - if let Err(err) = self - .0 + ) -> Result<(), ReplyControllerSenderError> { + self.0 .unbounded_send(ReplyControllerMessage::RetransmitReply { recipient, timed_out_ack, extra_surb_request, }) - { - error!("Failed to send retransmission data to reply controller: {err}",); - } + .map_err(ReplyControllerSenderError::SendRetransmissionData) } pub(crate) fn send_reply( @@ -48,14 +66,14 @@ impl ReplyControllerSender { recipient: AnonymousSenderTag, message: Vec, lane: TransmissionLane, - ) { - if let Err(err) = self.0.unbounded_send(ReplyControllerMessage::SendReply { - recipient, - message, - lane, - }) { - error!("Failed to send reply to reply controller: {err}",); - } + ) -> Result<(), ReplyControllerSenderError> { + self.0 + .unbounded_send(ReplyControllerMessage::SendReply { + recipient, + message, + lane, + }) + .map_err(ReplyControllerSenderError::SendReply) } pub(crate) fn send_additional_surbs( @@ -63,32 +81,33 @@ impl ReplyControllerSender { sender_tag: AnonymousSenderTag, reply_surbs: Vec, from_surb_request: bool, - ) { - if let Err(err) = self - .0 + ) -> Result<(), ReplyControllerSenderError> { + self.0 .unbounded_send(ReplyControllerMessage::AdditionalSurbs { sender_tag, reply_surbs, from_surb_request, }) - { - error!("Failed to send additional surbs to reply controller: {err}",); - } + .map_err(ReplyControllerSenderError::AdditionalSurbs) } - pub(crate) fn send_additional_surbs_request(&self, recipient: Recipient, amount: u32) { - if let Err(err) = self - .0 + pub(crate) fn send_additional_surbs_request( + &self, + recipient: Recipient, + amount: u32, + ) -> Result<(), ReplyControllerSenderError> { + self.0 .unbounded_send(ReplyControllerMessage::AdditionalSurbsRequest { recipient: Box::new(recipient), amount, }) - { - error!("Failed to send additional surbs request to reply controller: {err}"); - } + .map_err(ReplyControllerSenderError::AdditionalSurbsRequest) } - pub async fn get_lane_queue_length(&self, connection_id: ConnectionId) -> usize { + pub async fn get_lane_queue_length( + &self, + connection_id: ConnectionId, + ) -> Result { let (response_tx, response_rx) = oneshot::channel(); if let Err(err) = self .0 @@ -97,17 +116,12 @@ impl ReplyControllerSender { response_channel: response_tx, }) { - error!("Failed to send lane queue length request to reply controller: {err}"); + return Err(ReplyControllerSenderError::LaneQueueLength(err)); } - match response_rx.await { - Ok(length) => length, - Err(_) => { - error!("The reply controller has dropped our response channel!"); - // TODO: should we panic here instead? this message implies something weird and unrecoverable has happened - 0 - } - } + response_rx + .await + .map_err(ReplyControllerSenderError::ResponseChannelDropped) } } @@ -122,7 +136,10 @@ impl ReplyQueueLengths { } } - pub async fn get_lane_queue_length(&self, connection_id: ConnectionId) -> usize { + pub async fn get_lane_queue_length( + &self, + connection_id: ConnectionId, + ) -> Result { self.reply_controller_sender .get_lane_queue_length(connection_id) .await @@ -132,7 +149,7 @@ impl ReplyQueueLengths { pub(crate) type ReplyControllerReceiver = mpsc::UnboundedReceiver; #[derive(Debug)] -pub(crate) enum ReplyControllerMessage { +pub enum ReplyControllerMessage { RetransmitReply { recipient: AnonymousSenderTag, timed_out_ack: Weak,