diff --git a/Cargo.lock b/Cargo.lock index 42b753c5908..1e96b564583 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2108,6 +2108,8 @@ name = "apollo_network_benchmark" version = "0.0.0" dependencies = [ "clap", + "lazy_static", + "rstest", "tokio", ] diff --git a/crates/apollo_network_benchmark/Cargo.toml b/crates/apollo_network_benchmark/Cargo.toml index 48e133b4073..0b2f86ea378 100644 --- a/crates/apollo_network_benchmark/Cargo.toml +++ b/crates/apollo_network_benchmark/Cargo.toml @@ -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 diff --git a/crates/apollo_network_benchmark/src/bin/broadcast_network_stress_test_node/main.rs b/crates/apollo_network_benchmark/src/bin/broadcast_network_stress_test_node/main.rs index 44bc06af1e8..ec9d32bdfac 100644 --- a/crates/apollo_network_benchmark/src/bin/broadcast_network_stress_test_node/main.rs +++ b/crates/apollo_network_benchmark/src/bin/broadcast_network_stress_test_node/main.rs @@ -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; diff --git a/crates/apollo_network_benchmark/src/bin/broadcast_network_stress_test_node/message.rs b/crates/apollo_network_benchmark/src/bin/broadcast_network_stress_test_node/message.rs new file mode 100644 index 00000000000..c520aec4d96 --- /dev/null +++ b/crates/apollo_network_benchmark/src/bin/broadcast_network_stress_test_node/message.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 = 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, +} + +impl StressTestMessage { + pub fn new(sender_id: u64, message_index: u64, payload: Vec) -> Self { + StressTestMessage { + metadata: StressTestMessageMetadata { + sender_id, + message_index, + time: SystemTime::now(), + }, + payload, + } + } + + #[cfg(test)] + pub fn slow_len(self) -> usize { + let seq = Vec::::from(self); + seq.len() + } + + pub fn len(&self) -> usize { + *METADATA_SIZE + self.payload.len() + } +} + +impl From for Vec { + 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> for StressTestMessage { + fn from(bytes: Vec) -> 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, + } + } +} diff --git a/crates/apollo_network_benchmark/src/bin/broadcast_network_stress_test_node/message_test.rs b/crates/apollo_network_benchmark/src/bin/broadcast_network_stress_test_node/message_test.rs new file mode 100644 index 00000000000..e28b427de83 --- /dev/null +++ b/crates/apollo_network_benchmark/src/bin/broadcast_network_stress_test_node/message_test.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 = 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 = 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); +}