diff --git a/crates/apollo_network_benchmark/src/node_args.rs b/crates/apollo_network_benchmark/src/node_args.rs index 92d4e8790ed..6b1e2b2c704 100644 --- a/crates/apollo_network_benchmark/src/node_args.rs +++ b/crates/apollo_network_benchmark/src/node_args.rs @@ -3,6 +3,16 @@ use std::fmt::Display; use clap::{Parser, ValueEnum}; use serde::{Deserialize, Serialize}; +#[derive(Debug, Clone, ValueEnum, PartialEq, Eq, Serialize, Deserialize)] +pub enum Mode { + /// All nodes broadcast messages + #[value(name = "all")] + AllBroadcast, + /// Only the node specified by --broadcaster broadcasts messages + #[value(name = "one")] + OneBroadcast, +} + #[derive(Debug, Clone, ValueEnum, PartialEq, Eq, Serialize, Deserialize)] pub enum NetworkProtocol { /// Use gossipsub for broadcasting (default) @@ -10,6 +20,12 @@ pub enum NetworkProtocol { Gossipsub, } +impl Display for Mode { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.to_possible_value().unwrap().get_name()) + } +} + impl Display for NetworkProtocol { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{}", self.to_possible_value().unwrap().get_name()) @@ -49,10 +65,18 @@ pub struct UserArgs { #[arg(long, env, default_value = "100000")] pub buffer_size: usize, + /// The mode to use for the stress test. + #[arg(long, env, default_value = "all")] + pub mode: Mode, + /// The network protocol to use for communication (default: gossipsub) #[arg(long, env, default_value = "gossipsub")] pub network_protocol: NetworkProtocol, + /// Which node ID should do the broadcasting - for OneBroadcast mode + #[arg(long, env, required_if_eq("mode", "one"))] + pub broadcaster: Option, + /// Size of StressTestMessage #[arg(long, env, default_value = "1024")] pub message_size_bytes: usize,