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