feat(network): broadcast stress test

Cargo.toml

@@ -106,6 +106,7 @@ clap = "4.5.4"
 colored = "2.1.0"
 const_format = "0.2.30"
 criterion = "0.5.1"
+csv = "1.3.1"
 deadqueue = "0.2.4"
 defaultmap = "0.5.0"
 derive_more = "0.99.17"

crates/papyrus_network/Cargo.toml

@@ -12,6 +12,8 @@ testing = []
 async-stream.workspace = true
 async-trait.workspace = true
 bytes.workspace = true
+clap.workspace = true
+csv.workspace = true
 derive_more.workspace = true
 futures.workspace = true
 lazy_static.workspace = true
@@ -54,5 +56,8 @@ tokio = { workspace = true, features = ["full", "sync", "test-util"] }
 tokio-stream.workspace = true
 void.workspace = true
 
+[build-dependencies]
+
+
 [lints]
 workspace = true

crates/papyrus_network/src/bin/network_stress_test/bootstrap_test_config.json

@@ -0,0 +1,102 @@
+{
+  "buffer_size": {
+    "description": "The buffer size for the network receiver.",
+    "privacy": "Public",
+    "value": 1000
+  },
+  "message_size": {
+    "description": "The size of the payload for the test messages.",
+    "privacy": "Public",
+    "value": 1000
+  },
+  "network_config.advertised_multiaddr": {
+    "description": "The external address other peers see this node. If this is set, the node will not try to find out which addresses it has and will write this address as external instead",
+    "privacy": "Public",
+    "value": ""
+  },
+  "network_config.advertised_multiaddr.#is_none": {
+    "description": "Flag for an optional field.",
+    "privacy": "TemporaryValue",
+    "value": true
+  },
+  "network_config.bootstrap_peer_multiaddr": {
+    "description": "The multiaddress of the peer node. It should include the peer's id. For more info: https://docs.libp2p.io/concepts/fundamentals/peers/",
+    "privacy": "Public",
+    "value": ""
+  },
+  "network_config.bootstrap_peer_multiaddr.#is_none": {
+    "description": "Flag for an optional field.",
+    "privacy": "TemporaryValue",
+    "value": true
+  },
+  "network_config.chain_id": {
+    "description": "The chain to follow. For more details see https://docs.starknet.io/documentation/architecture_and_concepts/Blocks/transactions/#chain-id.",
+    "privacy": "Public",
+    "value": "SN_MAIN"
+  },
+  "network_config.discovery_config.bootstrap_dial_retry_config.base_delay_millis": {
+    "description": "The base delay in milliseconds for the exponential backoff strategy.",
+    "privacy": "Public",
+    "value": 2
+  },
+  "network_config.discovery_config.bootstrap_dial_retry_config.factor": {
+    "description": "The factor for the exponential backoff strategy.",
+    "privacy": "Public",
+    "value": 5
+  },
+  "network_config.discovery_config.bootstrap_dial_retry_config.max_delay_seconds": {
+    "description": "The maximum delay in seconds for the exponential backoff strategy.",
+    "privacy": "Public",
+    "value": 5
+  },
+  "network_config.discovery_config.heartbeat_interval": {
+    "description": "The interval between each discovery (Kademlia) query in milliseconds.",
+    "privacy": "Public",
+    "value": 100
+  },
+  "network_config.idle_connection_timeout": {
+    "description": "Amount of time in seconds that a connection with no active sessions will stay alive.",
+    "privacy": "Public",
+    "value": 120
+  },
+  "network_config.peer_manager_config.malicious_timeout_seconds": {
+    "description": "The duration in seconds a peer is blacklisted after being marked as malicious.",
+    "privacy": "Public",
+    "value": 31536000
+  },
+  "network_config.peer_manager_config.unstable_timeout_millis": {
+    "description": "The duration in milliseconds a peer blacklisted after being reported as unstable.",
+    "privacy": "Public",
+    "value": 1000
+  },
+  "network_config.quic_port": {
+    "description": "The port that the node listens on for incoming quic connections.",
+    "privacy": "Public",
+    "value": 10001
+  },
+  "network_config.secret_key": {
+    "description": "The secret key used for building the peer id. If it's an empty string a random one will be used.",
+    "privacy": "Private",
+    "value": "0x0000000000000000000000000000000000000000000000000000000000000000"
+  },
+  "network_config.session_timeout": {
+    "description": "Maximal time in seconds that each session can take before failing on timeout.",
+    "privacy": "Public",
+    "value": 120
+  },
+  "network_config.tcp_port": {
+    "description": "The port that the node listens on for incoming tcp connections.",
+    "privacy": "Public",
+    "value": 10000
+  },
+  "num_messages": {
+    "description": "The amount of messages to send and receive.",
+    "privacy": "Public",
+    "value": 10000
+  },
+  "output_path": {
+    "description": "The path of the output file.",
+    "privacy": "Public",
+    "value": "crates/papyrus_network/src/bin/network_stress_test/bootstrap_output.csv"
+  }
+}

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

@@ -1,21 +1,37 @@
+use std::mem::size_of;
+use std::str::FromStr;
 use std::time::{Duration, SystemTime};
 
-struct StressTestMessage {
-    id: u32,
-    payload: Vec<u8>,
-    time: SystemTime,
+use libp2p::PeerId;
+
+pub const METADATA_SIZE: usize = size_of::<u32>() + size_of::<u64>() + size_of::<u32>() + 38;
+
+#[derive(Debug, Clone)]
+pub struct StressTestMessage {
+    pub id: u32,
+    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 }
+    }
 }
 
 impl From<StressTestMessage> for Vec<u8> {
     fn from(value: StressTestMessage) -> Self {
-        let StressTestMessage { id, mut payload, time } = value;
+        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
     }
 }
@@ -24,12 +40,13 @@ 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());
         let nanos = u32::from_be_bytes(id_and_time[12..16].try_into().unwrap());
         let time = SystemTime::UNIX_EPOCH + Duration::new(seconds, nanos);
-        StressTestMessage { id, payload: value, time }
+        let peer_id = PeerId::from_bytes(&id_and_time[16..]).unwrap().to_string();
+        StressTestMessage { id, payload: value, time, peer_id }
     }
 }

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

@@ -1,3 +1,89 @@
+use std::time::SystemTime;
+use std::vec;
+
+use clap::Command;
+use converters::{StressTestMessage, METADATA_SIZE};
+use futures::StreamExt;
+use libp2p::gossipsub::Topic;
+use papyrus_config::loading::load_and_process_config;
+use papyrus_network::network_manager::{BroadcastTopicClientTrait, NetworkManager};
+use tokio::time::timeout;
+use utils::{Record, TestConfig, BOOTSTRAP_CONFIG_FILE_PATH};
+
 mod converters;
+mod utils;
 
-fn main() {}
+#[tokio::main]
+async fn main() {
+    let args = std::env::args().collect::<Vec<String>>();
+    let default_path = BOOTSTRAP_CONFIG_FILE_PATH.to_string();
+    let config_path = args.get(1).unwrap_or(&default_path);
+    let file = std::fs::File::open(config_path).unwrap();
+    let TestConfig { network_config, buffer_size, message_size, num_messages, output_path } =
+        load_and_process_config(file, Command::new("Stress Test"), vec![]).unwrap();
+    let mut network_manager = NetworkManager::new(network_config, None);
+    let peer_id = network_manager.get_local_peer_id();
+    let mut network_channels = network_manager
+        .register_broadcast_topic::<StressTestMessage>(
+            Topic::new("stress_test_topic".to_string()),
+            buffer_size,
+        )
+        .unwrap();
+    let mut output_vector = Vec::<Record>::new();
+    tokio::select! {
+        _ = network_manager.run() => {}
+        _ = async {
+            let mut i = 0;
+            tokio::time::sleep(std::time::Duration::from_secs(20)).await;
+            loop {
+                let message = StressTestMessage::new(i, vec![0; message_size - METADATA_SIZE], peer_id.clone());
+                network_channels.broadcast_topic_client.broadcast_message(message).await.unwrap();
+                i += 1;
+                if i == num_messages {
+                    println!("Finished sending messages");
+                    futures::future::pending::<()>().await;
+                }
+            }
+        } => {}
+        _ = async {
+            let mut i = 0;
+            loop {
+                let maybe_response = timeout(
+                    std::time::Duration::from_secs(120),
+                    network_channels.broadcasted_messages_receiver.next(),
+                ).await;
+                match maybe_response {
+                    Err(_) => {
+                        println!("Timeout on message {}", i);
+                        break;
+                    }
+                    Ok(None) => break,
+                    Ok(Some((received_message, _report_callback))) => {
+                        let received_message = received_message.unwrap();
+                        output_vector.push(Record {
+                            peer_id: received_message.peer_id,
+                            id: received_message.id,
+                            start_time: received_message.time,
+                            end_time: SystemTime::now(),
+                            duration: SystemTime::now()
+                                .duration_since(received_message.time)
+                                .unwrap()
+                                .as_micros(),
+                        });
+                        i += 1;
+                        if i == num_messages * 4 {
+                            tokio::time::sleep(std::time::Duration::from_secs(20)).await;
+                            break;
+                        }
+                    }
+                }
+            }
+        } => {
+            println!("Finished receiving messages");
+            let mut wtr = csv::Writer::from_path(output_path).unwrap();
+            for record in output_vector {
+                wtr.serialize(record).unwrap();
+            }
+        }
+    }
+}