apollo_network_benchmark: added the message that will be used by the benchmark

Author: sirandreww-starkware

Cargo.lock

@@ -10,11 +10,12 @@ testing = []
 
 [dependencies]
 clap = { workspace = true, features = ["derive"] }
+lazy_static.workspace = true
 tokio = { workspace = true, features = ["full", "sync"] }
 
 
 [dev-dependencies]
-
+rstest.workspace = true
 
 [lints]
 workspace = true

crates/apollo_network_benchmark/Cargo.toml

@@ -1,7 +1,12 @@
 //! Runs a node that stress tests the p2p communication of the network.
 
 use clap::Parser;
+
+#[cfg(test)]
+mod message_test;
+
 mod args;
+mod message;
 
 use args::Args;
 

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

@@ -0,0 +1,88 @@
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
+
+use lazy_static::lazy_static;
+
+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 metadata: StressTestMessageMetadata,
+    pub payload: Vec<u8>,
+}
+
+impl StressTestMessage {
+    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 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 {
+    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_benchmark/src/bin/broadcast_network_stress_test_node/message.rs

@@ -0,0 +1,52 @@
+use rstest::rstest;
+
+use crate::message::{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);
+}