chore(network): stress test utils

crates/papyrus_network/src/bin/network_stress_test/converters.rs

@@ -1,9 +1,19 @@
+use std::mem::size_of;
 use std::time::{Duration, SystemTime};
 
-struct StressTestMessage {
-    id: u32,
-    payload: Vec<u8>,
-    time: SystemTime,
+pub const METADATA_SIZE: usize = size_of::<SystemTime>() + size_of::<u32>();
+
+#[derive(Debug, Clone)]
+pub struct StressTestMessage {
+    pub id: u32,
+    pub payload: Vec<u8>,
+    pub time: SystemTime,
+}
+
+impl StressTestMessage {
+    pub fn new(id: u32, payload: Vec<u8>) -> Self {
+        StressTestMessage { id, payload, time: SystemTime::now() }
+    }
 }
 
 impl From<StressTestMessage> for Vec<u8> {
@@ -24,7 +34,7 @@ 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 - 12;
+        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());

crates/papyrus_network/src/bin/network_stress_test/utils.rs

@@ -0,0 +1,132 @@
+use std::collections::{BTreeMap, HashSet};
+use std::str::FromStr;
+use std::time::{SystemTime, UNIX_EPOCH};
+use std::vec;
+
+use libp2p::identity::Keypair;
+use libp2p::Multiaddr;
+use papyrus_config::dumping::{append_sub_config_name, ser_param, SerializeConfig};
+use papyrus_config::{ParamPath, ParamPrivacyInput, SerializedParam};
+use papyrus_network::NetworkConfig;
+use serde::{Deserialize, Serialize, Serializer};
+
+pub const BOOTSTRAP_CONFIG_FILE_PATH: &str =
+    "crates/papyrus_network/src/bin/network_stress_test/bootstrap_test_config.json";
+pub const BOOTSTRAP_OUTPUT_FILE_PATH: &str =
+    "crates/papyrus_network/src/bin/network_stress_test/bootstrap_output.csv";
+pub const DEFAULT_CONFIG_FILE_PATH: &str =
+    "crates/papyrus_network/src/bin/network_stress_test/test_config.json";
+pub const DEFAULT_OUTPUT_FILE_PATH: &str =
+    "crates/papyrus_network/src/bin/network_stress_test/output.csv";
+
+#[derive(Debug, Deserialize, Serialize)]
+pub struct TestConfig {
+    pub network_config: NetworkConfig,
+    pub buffer_size: usize,
+    pub message_size: usize,
+    pub num_messages: u32,
+    pub output_path: String,
+}
+
+impl SerializeConfig for TestConfig {
+    fn dump(&self) -> BTreeMap<ParamPath, SerializedParam> {
+        let mut config = BTreeMap::from_iter([
+            ser_param(
+                "buffer_size",
+                &self.buffer_size,
+                "The buffer size for the network receiver.",
+                ParamPrivacyInput::Public,
+            ),
+            ser_param(
+                "message_size",
+                &self.message_size,
+                "The size of the payload for the test messages.",
+                ParamPrivacyInput::Public,
+            ),
+            ser_param(
+                "num_messages",
+                &self.num_messages,
+                "The amount of messages to send and receive.",
+                ParamPrivacyInput::Public,
+            ),
+            ser_param(
+                "output_path",
+                &self.output_path,
+                "The path of the output file.",
+                ParamPrivacyInput::Public,
+            ),
+        ]);
+        config.extend(append_sub_config_name(self.network_config.dump(), "network_config"));
+        config
+    }
+}
+
+impl Default for TestConfig {
+    fn default() -> Self {
+        Self {
+            network_config: NetworkConfig::default(),
+            buffer_size: 1000,
+            message_size: 1000,
+            num_messages: 100000,
+            output_path: BOOTSTRAP_OUTPUT_FILE_PATH.to_string(),
+        }
+    }
+}
+
+impl TestConfig {
+    #[allow(dead_code)]
+    pub fn create_config_files() {
+        let secret_key = vec![0; 32];
+        let keypair = Keypair::ed25519_from_bytes(secret_key.clone()).unwrap();
+        let peer_id = keypair.public().to_peer_id();
+
+        let _ = TestConfig {
+            network_config: NetworkConfig {
+                tcp_port: 10000,
+                quic_port: 10001,
+                secret_key: Some(secret_key),
+                ..Default::default()
+            },
+            ..Default::default()
+        }
+        .dump_to_file(&vec![], &HashSet::new(), BOOTSTRAP_CONFIG_FILE_PATH);
+        let _ = TestConfig {
+            network_config: NetworkConfig {
+                tcp_port: 10002,
+                quic_port: 10003,
+                bootstrap_peer_multiaddr: Some(
+                    Multiaddr::from_str(&format!("/ip4/127.0.0.1/tcp/10000/p2p/{}", peer_id))
+                        .unwrap(),
+                ),
+                ..Default::default()
+            },
+            output_path: DEFAULT_OUTPUT_FILE_PATH.to_string(),
+            ..Default::default()
+        }
+        .dump_to_file(&vec![], &HashSet::new(), DEFAULT_CONFIG_FILE_PATH);
+    }
+}
+
+#[derive(Debug, Deserialize, Serialize)]
+pub struct Record {
+    pub id: u32,
+    #[serde(serialize_with = "serialize_system_time_as_u128_millis")]
+    pub start_time: SystemTime,
+    #[serde(serialize_with = "serialize_system_time_as_u128_millis")]
+    pub end_time: SystemTime,
+    pub duration: u128,
+}
+
+pub fn serialize_system_time_as_u128_millis<S>(
+    time: &SystemTime,
+    serializer: S,
+) -> Result<S::Ok, S::Error>
+where
+    S: Serializer,
+{
+    let duration_since_epoch =
+        time.duration_since(UNIX_EPOCH).map_err(serde::ser::Error::custom)?;
+    let millis =
+        duration_since_epoch.as_millis() + u128::from(duration_since_epoch.subsec_millis());
+    serializer.serialize_u128(millis)
+}