diff --git a/nym-client/src/client/mod.rs b/nym-client/src/client/mod.rs index 376e3a102c..0be475c82e 100644 --- a/nym-client/src/client/mod.rs +++ b/nym-client/src/client/mod.rs @@ -1,5 +1,6 @@ use crate::built_info; use crate::client::mix_traffic::{MixMessage, MixTrafficController}; +use crate::client::received_buffer::ReceivedMessagesBuffer; use crate::sockets::tcp; use crate::sockets::ws; use crate::utils; @@ -15,7 +16,6 @@ use std::net::SocketAddr; use std::time::Duration; use tokio::runtime::Runtime; use topology::NymTopology; -use crate::client::received_buffer::ReceivedMessagesBuffer; mod mix_traffic; pub mod received_buffer; @@ -223,18 +223,10 @@ impl NymClient { let (received_messages_buffer_output_tx, received_messages_buffer_output_rx) = mpsc::unbounded(); - let received_messages_buffer = ReceivedMessagesBuffer::new().add_arc_futures_mutex(); - - let received_messages_buffer_input_controller_future = - rt.spawn(ReceivedMessagesBuffer::run_poller_input_controller( - received_messages_buffer.clone(), - poller_input_rx, - )); - let received_messages_buffer_output_controller_future = - rt.spawn(ReceivedMessagesBuffer::run_query_output_controller( - received_messages_buffer, - received_messages_buffer_output_rx, - )); + let received_messages_buffer_controllers_future = rt.spawn( + ReceivedMessagesBuffer::new() + .start_controllers(poller_input_rx, received_messages_buffer_output_rx), + ); let mix_traffic_future = rt.spawn(MixTrafficController::run(mix_rx)); let loop_cover_traffic_future = rt.spawn(NymClient::start_loop_cover_traffic_stream( @@ -279,8 +271,7 @@ impl NymClient { rt.block_on(async { let future_results = join!( - received_messages_buffer_input_controller_future, - received_messages_buffer_output_controller_future, + received_messages_buffer_controllers_future, mix_traffic_future, loop_cover_traffic_future, out_queue_control_future, @@ -293,7 +284,6 @@ impl NymClient { && future_results.2.is_ok() && future_results.3.is_ok() && future_results.4.is_ok() - && future_results.5.is_ok() ); }); diff --git a/nym-client/src/client/received_buffer.rs b/nym-client/src/client/received_buffer.rs index 3dcb0627f7..a6df580c9f 100644 --- a/nym-client/src/client/received_buffer.rs +++ b/nym-client/src/client/received_buffer.rs @@ -7,56 +7,83 @@ use std::sync::Arc; pub type BufferResponse = oneshot::Sender>>; pub(crate) struct ReceivedMessagesBuffer { - messages: Vec>, + inner: Arc>, } impl ReceivedMessagesBuffer { pub(crate) fn new() -> Self { ReceivedMessagesBuffer { - messages: Vec::new(), + inner: Arc::new(FMutex::new(Inner::new())), } } - pub(crate) fn add_arc_futures_mutex(self) -> Arc> { - Arc::new(FMutex::new(self)) - } + pub(crate) async fn start_controllers( + self, + poller_rx: mpsc::UnboundedReceiver>>, // to receive new messages + query_receiver: mpsc::UnboundedReceiver, // to receive requests to acquire all stored messages + ) { + let input_controller_future = tokio::spawn(Self::run_poller_input_controller( + self.inner.clone(), + poller_rx, + )); + let output_controller_future = tokio::spawn(Self::run_query_output_controller( + self.inner, + query_receiver, + )); - pub(crate) async fn add_new_messages(buf: Arc>, msgs: Vec>) { - trace!("Adding new messages to the buffer! {:?}", msgs); - let mut unlocked = buf.lock().await; - unlocked.messages.extend(msgs); + futures::future::select(input_controller_future, output_controller_future).await; + panic!("One of the received buffer controllers failed!") } pub(crate) async fn run_poller_input_controller( - buf: Arc>, + buf: Arc>, mut poller_rx: mpsc::UnboundedReceiver>>, ) { info!("Started Received Messages Buffer Input Controller"); + while let Some(new_messages) = poller_rx.next().await { - ReceivedMessagesBuffer::add_new_messages(buf.clone(), new_messages).await; + Inner::add_new_messages(&*buf, new_messages).await; } } - pub(crate) async fn acquire_and_empty(buf: Arc>) -> Vec> { - trace!("Emptying the buffer and returning all messages"); - let mut unlocked = buf.lock().await; - std::mem::replace(&mut unlocked.messages, Vec::new()) - } - pub(crate) async fn run_query_output_controller( - buf: Arc>, + buf: Arc>, mut query_receiver: mpsc::UnboundedReceiver, ) { info!("Started Received Messages Buffer Output Controller"); while let Some(request) = query_receiver.next().await { - let messages = ReceivedMessagesBuffer::acquire_and_empty(buf.clone()).await; + let messages = Inner::acquire_and_empty(&*buf).await; if let Err(failed_messages) = request.send(messages) { error!( "Failed to send the messages to the requester. Adding them back to the buffer" ); - ReceivedMessagesBuffer::add_new_messages(buf.clone(), failed_messages).await; + Inner::add_new_messages(&*buf, failed_messages).await; } } } } + +pub(crate) struct Inner { + messages: Vec>, +} + +impl Inner { + fn new() -> Self { + Inner { + messages: Vec::new(), + } + } + + async fn add_new_messages(buf: &FMutex, msgs: Vec>) { + trace!("Adding new messages to the buffer! {:?}", msgs); + let mut unlocked = buf.lock().await; + unlocked.messages.extend(msgs); + } + + async fn acquire_and_empty(buf: &FMutex) -> Vec> { + trace!("Emptying the buffer and returning all messages"); + let mut unlocked = buf.lock().await; + std::mem::replace(&mut unlocked.messages, Vec::new()) + } +}