apollo_network: broadcast_network_stress_test_node updated message

Author: sirandreww-starkware

crates/apollo_network/src/bin/broadcast_network_stress_test_node/converters.rs

@@ -1,52 +1,88 @@
-use std::mem::size_of;
-use std::str::FromStr;
-use std::time::{Duration, SystemTime};
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
-use libp2p::PeerId;
+use lazy_static::lazy_static;
 
-pub const METADATA_SIZE: usize = size_of::<u32>() + size_of::<u64>() + size_of::<u32>() + 38;
+lazy_static! {
+    // Calculate actual metadata size based on serialized empty message
+    pub static ref METADATA_SIZE: usize = {
+        let empty_message = StressTestMessage::new(0, 0, vec![]);
+        let serialized: Vec<u8> = empty_message.into();
+        serialized.len()
+    };
+}
+
+#[derive(Debug, Clone, Copy)]
+pub struct StressTestMessageMetadata {
+    pub sender_id: u64,
+    pub message_index: u64,
+    pub time: SystemTime,
+}
 
 #[derive(Debug, Clone)]
 pub struct StressTestMessage {
-    pub id: u32,
+    pub metadata: StressTestMessageMetadata,
     pub payload: Vec<u8>,
-    pub time: SystemTime,
-    pub peer_id: String,
 }
 
 impl StressTestMessage {
-    pub fn new(id: u32, payload: Vec<u8>, peer_id: String) -> Self {
-        StressTestMessage { id, payload, time: SystemTime::now(), peer_id }
+    pub fn new(sender_id: u64, message_index: u64, payload: Vec<u8>) -> Self {
+        StressTestMessage {
+            metadata: StressTestMessageMetadata {
+                sender_id,
+                message_index,
+                time: SystemTime::now(),
+            },
+            payload,
+        }
+    }
+
+    #[cfg(test)]
+    pub fn slow_len(self) -> usize {
+        let seq = Vec::<u8>::from(self);
+        seq.len()
+    }
+
+    pub fn len(&self) -> usize {
+        *METADATA_SIZE + self.payload.len()
     }
 }
 
 impl From<StressTestMessage> for Vec<u8> {
     fn from(value: StressTestMessage) -> Self {
-        let StressTestMessage { id, mut payload, time, peer_id } = value;
-        let id = id.to_be_bytes().to_vec();
-        let time = time.duration_since(SystemTime::UNIX_EPOCH).unwrap();
-        let seconds = time.as_secs().to_be_bytes().to_vec();
-        let nanos = time.subsec_nanos().to_be_bytes().to_vec();
-        let peer_id = PeerId::from_str(&peer_id).unwrap().to_bytes();
-        payload.extend(id);
-        payload.extend(seconds);
-        payload.extend(nanos);
-        payload.extend(peer_id);
-        payload
+        let payload_len: u64 = value.payload.len().try_into().unwrap();
+        let duration = value.metadata.time.duration_since(UNIX_EPOCH).unwrap();
+        [
+            &value.metadata.sender_id.to_be_bytes()[..],
+            &value.metadata.message_index.to_be_bytes()[..],
+            &duration.as_secs().to_be_bytes()[..],
+            &duration.subsec_nanos().to_be_bytes()[..],
+            &payload_len.to_be_bytes()[..],
+            &value.payload[..],
+        ]
+        .concat()
     }
 }
 
 impl From<Vec<u8>> for StressTestMessage {
-    // This auto implements TryFrom<Vec<u8>> for StressTestMessage
-    fn from(mut value: Vec<u8>) -> Self {
-        let vec_size = value.len();
-        let payload_size = vec_size - METADATA_SIZE;
-        let id_and_time = value.split_off(payload_size);
-        let id = u32::from_be_bytes(id_and_time[0..4].try_into().unwrap());
-        let seconds = u64::from_be_bytes(id_and_time[4..12].try_into().unwrap());
-        let nanos = u32::from_be_bytes(id_and_time[12..16].try_into().unwrap());
-        let time = SystemTime::UNIX_EPOCH + Duration::new(seconds, nanos);
-        let peer_id = PeerId::from_bytes(&id_and_time[16..]).unwrap().to_string();
-        StressTestMessage { id, payload: value, time, peer_id }
+    fn from(bytes: Vec<u8>) -> Self {
+        let mut i = 0;
+        let mut get = |n: usize| {
+            let r = &bytes[i..i + n];
+            i += n;
+            r
+        };
+
+        let sender_id = u64::from_be_bytes(get(8).try_into().unwrap());
+        let message_index = u64::from_be_bytes(get(8).try_into().unwrap());
+        let secs = u64::from_be_bytes(get(8).try_into().unwrap());
+        let nanos = u32::from_be_bytes(get(4).try_into().unwrap());
+        let time = UNIX_EPOCH + Duration::new(secs, nanos);
+        let payload_len = u64::from_be_bytes(get(8).try_into().unwrap()).try_into().unwrap();
+        let payload = get(payload_len).to_vec();
+
+        StressTestMessage {
+            metadata: StressTestMessageMetadata { sender_id, message_index, time },
+            payload,
+        }
     }
 }

crates/apollo_network/src/bin/broadcast_network_stress_test_node/converters_test.rs

@@ -0,0 +1,52 @@
+use rstest::rstest;
+
+use crate::converters::{StressTestMessage, METADATA_SIZE};
+
+#[rstest]
+#[case::one_byte_len(10)]
+#[case::two_byte_len(300)]
+#[case::three_byte_len(20_000)]
+fn test_message_size(#[case] vec_len: usize) {
+    let payload = vec![0xAA; vec_len];
+    let message = StressTestMessage::new(1, 7070, payload.clone());
+    let expected_size = *METADATA_SIZE + vec_len;
+    assert_eq!(message.len(), expected_size);
+    assert_eq!(message.slow_len(), expected_size);
+}
+
+#[test]
+fn test_serialization_and_deserilization() {
+    let original_message =
+        StressTestMessage::new(u64::MAX - 1, u64::MAX - 2, vec![0xa1, 0xb2, 0xc3, 0xd4, 0xe5]);
+
+    // Serialize to bytes
+    let serialized_bytes: Vec<u8> = original_message.clone().into();
+
+    // Deserialize back to message
+    let deserialized_message: StressTestMessage = serialized_bytes.into();
+
+    // Verify all fields match
+    assert_eq!(deserialized_message.metadata.sender_id, original_message.metadata.sender_id);
+    assert_eq!(
+        deserialized_message.metadata.message_index,
+        original_message.metadata.message_index
+    );
+    assert_eq!(deserialized_message.payload, original_message.payload);
+    assert_eq!(deserialized_message.metadata.time, original_message.metadata.time);
+}
+
+#[test]
+fn test_empty_payload() {
+    let original_message = StressTestMessage::new(1, 42, vec![]);
+
+    let serialized_bytes: Vec<u8> = original_message.clone().into();
+    let deserialized_message: StressTestMessage = serialized_bytes.into();
+
+    assert_eq!(deserialized_message.metadata.sender_id, original_message.metadata.sender_id);
+    assert_eq!(
+        deserialized_message.metadata.message_index,
+        original_message.metadata.message_index
+    );
+    assert_eq!(deserialized_message.payload, original_message.payload);
+    assert_eq!(deserialized_message.metadata.time, original_message.metadata.time);
+}

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

@@ -26,6 +26,9 @@ use metrics_exporter_prometheus::PrometheusBuilder;
 use tokio::time::Duration;
 use tracing::{info, trace, Level};
 
+#[cfg(test)]
+mod converters_test;
+
 mod converters;
 mod utils;
 
@@ -77,18 +80,18 @@ struct Args {
 async fn send_stress_test_messages(
     mut broadcast_topic_client: BroadcastTopicClient<StressTestMessage>,
     args: &Args,
-    peer_id: String,
+    _peer_id: String,
 ) {
     let mut message = StressTestMessage::new(
         args.id.try_into().unwrap(),
-        vec![0; args.message_size_bytes - METADATA_SIZE],
-        peer_id.clone(),
+        0, // message_index, will be updated in loop
+        vec![0; args.message_size_bytes - *METADATA_SIZE],
     );
     let duration = Duration::from_millis(args.heartbeat_millis);
 
     for i in 0.. {
-        message.time = SystemTime::now();
-        // message.id = i;
+        message.metadata.time = SystemTime::now();
+        message.metadata.message_index = i;
         broadcast_topic_client.broadcast_message(message.clone()).await.unwrap();
         trace!("Sent message {i}: {:?}", message);
         counter!("sent_messages").increment(1);
@@ -103,7 +106,7 @@ fn receive_stress_test_message(
     let end_time = SystemTime::now();
 
     let received_message = message_result.unwrap();
-    let start_time = received_message.time;
+    let start_time = received_message.metadata.time;
     let duration = match end_time.duration_since(start_time) {
         Ok(duration) => duration,
         Err(_) => panic!("Got a negative duration, the clocks are not synced!"),
@@ -114,16 +117,19 @@ fn receive_stress_test_message(
 
     // TODO(AndrewL): Concentrate all string metrics to constants in a different file
     counter!("message_received").increment(1);
-    counter!(format!("message_received_from_{}", received_message.id)).increment(1);
+    counter!(format!("message_received_from_{}", received_message.metadata.sender_id)).increment(1);
 
     // TODO(AndrewL): This should be a historgram
     gauge!("message_received_delay_seconds").set(delay_seconds);
-    gauge!(format!("message_received_delay_seconds_from_{}", received_message.id))
+    gauge!(format!("message_received_delay_seconds_from_{}", received_message.metadata.sender_id))
         .set(delay_seconds);
 
     counter!("message_received_delay_micros_sum").increment(delay_micros);
-    counter!(format!("message_received_delay_micros_sum_from_{}", received_message.id))
-        .increment(delay_micros);
+    counter!(format!(
+        "message_received_delay_micros_sum_from_{}",
+        received_message.metadata.sender_id
+    ))
+    .increment(delay_micros);
     // TODO(AndrewL): Figure out what to log here
 }
 
@@ -168,8 +174,9 @@ async fn main() {
     println!("Starting network stress test with args:\n{args:?}");
 
     assert!(
-        args.message_size_bytes >= METADATA_SIZE,
-        "Message size must be at least {METADATA_SIZE} bytes"
+        args.message_size_bytes >= *METADATA_SIZE,
+        "Message size must be at least {} bytes",
+        *METADATA_SIZE
     );
 
     let builder = PrometheusBuilder::new().with_http_listener(SocketAddr::V4(SocketAddrV4::new(