ReceivedBuffer is now responsible for starting both controllers

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()
             );
         });
 

nym-client/src/client/received_buffer.rs

@@ -7,56 +7,83 @@ use std::sync::Arc;
 pub type BufferResponse = oneshot::Sender<Vec<Vec<u8>>>;
 
 pub(crate) struct ReceivedMessagesBuffer {
-    messages: Vec<Vec<u8>>,
+    inner: Arc<FMutex<Inner>>,
 }
 
 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<FMutex<Self>> {
-        Arc::new(FMutex::new(self))
-    }
+    pub(crate) async fn start_controllers(
+        self,
+        poller_rx: mpsc::UnboundedReceiver<Vec<Vec<u8>>>, // to receive new messages
+        query_receiver: mpsc::UnboundedReceiver<BufferResponse>, // 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<FMutex<Self>>, msgs: Vec<Vec<u8>>) {
-        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<FMutex<Self>>,
+        buf: Arc<FMutex<Inner>>,
         mut poller_rx: mpsc::UnboundedReceiver<Vec<Vec<u8>>>,
     ) {
         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<FMutex<Self>>) -> Vec<Vec<u8>> {
-        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<FMutex<Self>>,
+        buf: Arc<FMutex<Inner>>,
         mut query_receiver: mpsc::UnboundedReceiver<BufferResponse>,
     ) {
         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<Vec<u8>>,
+}
+
+impl Inner {
+    fn new() -> Self {
+        Inner {
+            messages: Vec::new(),
+        }
+    }
+
+    async fn add_new_messages(buf: &FMutex<Self>, msgs: Vec<Vec<u8>>) {
+        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<Self>) -> Vec<Vec<u8>> {
+        trace!("Emptying the buffer and returning all messages");
+        let mut unlocked = buf.lock().await;
+        std::mem::replace(&mut unlocked.messages, Vec::new())
+    }
+}