apollo_network_benchmark: added message receiving task

Author: sirandreww-starkware

crates/apollo_network_benchmark/src/bin/broadcast_network_stress_test_node/handlers.rs

@@ -0,0 +1,20 @@
+use std::time::SystemTime;
+
+use libp2p::PeerId;
+use tracing::trace;
+
+use crate::message::StressTestMessage;
+
+pub fn receive_stress_test_message(received_message: Vec<u8>, sender_peer_id: Option<PeerId>) {
+    let end_time = SystemTime::now();
+
+    let received_message: StressTestMessage = received_message.into();
+    let start_time = received_message.metadata.time;
+    let delay_seconds = end_time
+        .duration_since(start_time)
+        .expect("End time should be after start time (Probably clock misalignment)")
+        .as_secs_f64();
+
+    // TODO(AndrewL): Replace this with metric updates
+    trace!("Received stress test message from {sender_peer_id:?} in {delay_seconds} seconds");
+}

crates/apollo_network_benchmark/src/bin/broadcast_network_stress_test_node/main.rs

@@ -10,6 +10,7 @@ use tracing::Level;
 #[cfg(test)]
 mod message_test;
 
+mod handlers;
 mod message;
 mod protocol;
 mod stress_test_node;

crates/apollo_network_benchmark/src/bin/broadcast_network_stress_test_node/stress_test_node.rs

@@ -10,6 +10,7 @@ use libp2p::Multiaddr;
 use tokio::task::JoinHandle;
 use tracing::{info, warn};
 
+use crate::handlers::receive_stress_test_message;
 use crate::protocol::{register_protocol_channels, MessageReceiver, MessageSender};
 
 /// The main stress test node that manages network communication and monitoring
@@ -19,8 +20,6 @@ pub struct BroadcastNetworkStressTestNode {
     // TODO(AndrewL): Remove this once they are used
     #[allow(dead_code)]
     message_sender: Option<MessageSender>,
-    // TODO(AndrewL): Remove this once they are used
-    #[allow(dead_code)]
     message_receiver: Option<MessageReceiver>,
 }
 
@@ -83,10 +82,24 @@ impl BroadcastNetworkStressTestNode {
         .boxed()
     }
 
+    /// Starts the message receiving task
+    pub async fn make_message_receiver_task(&mut self) -> BoxFuture<'static, ()> {
+        let message_receiver =
+            self.message_receiver.take().expect("message_receiver should be available");
+
+        async move {
+            info!("Starting message receiver");
+            message_receiver.for_each(receive_stress_test_message).await;
+            info!("Message receiver task ended");
+        }
+        .boxed()
+    }
+
     /// Gets all the tasks that need to be run
     async fn get_tasks(&mut self) -> Vec<BoxFuture<'static, ()>> {
         let mut tasks = Vec::new();
         tasks.push(self.start_network_manager().await);
+        tasks.push(self.make_message_receiver_task().await);
 
         tasks
     }