apollo_network: using propeller in network stress test

Author: sirandreww-starkware

crates/apollo_network/src/bin/broadcast_network_stress_test_node/args.rs

@@ -33,6 +33,9 @@ pub enum NetworkProtocol {
     /// Use Reversed SQMR where receivers initiate requests to broadcasters
     #[value(name = "reversed-sqmr")]
     ReveresedSqmr,
+    /// Use Propeller for leader-based erasure-coded broadcasting
+    #[value(name = "propeller")]
+    Propeller,
 }
 
 impl Display for Mode {
@@ -47,6 +50,41 @@ impl Display for NetworkProtocol {
     }
 }
 
+impl NetworkProtocol {
+    /// Validates protocol-specific requirements for the given message size.
+    /// Returns an error message if validation fails, or Ok(()) if validation passes.
+    pub fn validate_message_size(&self, message_size_bytes: Option<usize>) -> Result<(), String> {
+        match self {
+            NetworkProtocol::Propeller => {
+                if let Some(message_size) = message_size_bytes {
+                    let default_config = libp2p::propeller::Config::default();
+                    let shreds =
+                        default_config.fec_coding_shreds() + default_config.fec_data_shreds();
+                    if message_size % shreds != 0 {
+                        return Err(format!(
+                            "Propeller protocol requires message size to be a multiple of 64 \
+                             bytes, got {}",
+                            message_size
+                        ));
+                    }
+                    let shred_size = message_size / shreds;
+                    if !(shred_size).is_multiple_of(2) {
+                        return Err(format!(
+                            "Propeller protocol requires shred size to be a multiple of 2, got {}",
+                            shred_size
+                        ));
+                    }
+                }
+                Ok(())
+            }
+            NetworkProtocol::Gossipsub | NetworkProtocol::Sqmr | NetworkProtocol::ReveresedSqmr => {
+                // These protocols don't have specific message size requirements
+                Ok(())
+            }
+        }
+    }
+}
+
 #[derive(Parser, Debug, Clone)]
 #[command(version, about, long_about = None)]
 pub struct Args {

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

@@ -49,6 +49,13 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         *METADATA_SIZE
     );
 
+    // Protocol-specific validation
+    if let Err(validation_error) =
+        args.network_protocol.validate_message_size(args.message_size_bytes)
+    {
+        panic!("{}", validation_error);
+    }
+
     // Set up metrics
     let builder = PrometheusBuilder::new().with_http_listener(SocketAddr::V4(SocketAddrV4::new(
         Ipv4Addr::UNSPECIFIED,

crates/apollo_network/src/bin/broadcast_network_stress_test_node/message_handling.rs

@@ -2,6 +2,10 @@ use apollo_network::network_manager::{
     BroadcastTopicClient,
     BroadcastTopicClientTrait,
     BroadcastTopicServer,
+    PropellerClient,
+    PropellerClientTrait,
+    PropellerMessageServer,
+    ReceivedPropellerMessage,
     SqmrClientSender,
     SqmrServerReceiver,
 };
@@ -16,6 +20,30 @@ pub enum MessageSender {
     Gossipsub(BroadcastTopicClient<TopicType>),
     Sqmr(SqmrClientSender<TopicType, TopicType>),
     ReveresedSqmr(ReveresedSqmrSender),
+    Propeller(PropellerSender),
+}
+
+/// Wrapper for Propeller client that handles message ID generation
+pub struct PropellerSender {
+    client: PropellerClient<TopicType>,
+    next_message_id: u64,
+}
+
+impl PropellerSender {
+    pub fn new(client: PropellerClient<TopicType>) -> Self {
+        Self { client, next_message_id: 1 }
+    }
+
+    async fn send_message(&mut self, message: TopicType) {
+        let message_id = self.next_message_id;
+        self.next_message_id += 1;
+
+        if let Err(e) = self.client.send_message(message, message_id).await {
+            error!("Failed to send Propeller message {}: {:?}", message_id, e);
+        } else {
+            trace!("Sent Propeller message with ID: {}", message_id);
+        }
+    }
 }
 
 /// Wrapper for ReveresedSqmr that maintains the last active query
@@ -88,6 +116,9 @@ impl MessageSender {
                 // Then broadcast the message to all active queries
                 sender.broadcast_to_queries(message).await;
             }
+            MessageSender::Propeller(sender) => {
+                sender.send_message(message).await;
+            }
         }
     }
 }
@@ -96,6 +127,7 @@ pub enum MessageReceiver {
     Gossipsub(BroadcastTopicServer<TopicType>),
     Sqmr(SqmrServerReceiver<TopicType, TopicType>),
     ReveresedSqmr(SqmrClientSender<TopicType, TopicType>),
+    Propeller(PropellerMessageServer<TopicType>),
 }
 
 impl MessageReceiver {
@@ -146,6 +178,26 @@ impl MessageReceiver {
                     }
                 }
             },
+            MessageReceiver::Propeller(receiver) => {
+                receiver
+                    .for_each(
+                        |(message_id, result): ReceivedPropellerMessage<TopicType>| async move {
+                            match result {
+                                Ok(message) => {
+                                    trace!("Received Propeller message with ID: {}", message_id);
+                                    f(message);
+                                }
+                                Err(e) => {
+                                    error!(
+                                        "Failed to deserialize Propeller message {}: {:?}",
+                                        message_id, e
+                                    );
+                                }
+                            }
+                        },
+                    )
+                    .await
+            }
         }
     }
 }

crates/apollo_network/src/bin/broadcast_network_stress_test_node/network_channels.rs

@@ -3,19 +3,73 @@ use apollo_network::network_manager::{
     BroadcastTopicClient,
     BroadcastTopicServer,
     NetworkManager,
+    PropellerChannels,
+    PropellerClient,
+    PropellerMessageServer,
     SqmrClientSender,
     SqmrServerReceiver,
 };
 use apollo_network::NetworkConfig;
 
 use crate::args::NetworkProtocol;
-use crate::message_handling::{MessageReceiver, MessageSender};
+use crate::message_handling::{MessageReceiver, MessageSender, PropellerSender};
 use crate::metrics::{create_network_metrics, TOPIC};
 
 pub type TopicType = Vec<u8>;
 
 pub const SQMR_PROTOCOL_NAME: &str = "/stress-test/1.0.0";
 
+/// Creates peer configurations for Propeller protocol from bootstrap addresses
+/// Each peer gets equal weight (1000) for simplicity
+fn create_propeller_peers_from_bootstrap(bootstrap_addrs: &[String]) -> Vec<(libp2p::PeerId, u64)> {
+    use std::str::FromStr;
+
+    use libp2p::Multiaddr;
+
+    bootstrap_addrs
+        .iter()
+        .filter_map(|addr_str| {
+            // Parse the multiaddr and extract the peer ID
+            match Multiaddr::from_str(addr_str.trim()) {
+                Ok(multiaddr) => {
+                    // Extract peer ID from the multiaddr (last component should be /p2p/<peer_id>)
+                    for protocol in multiaddr.iter() {
+                        if let libp2p::multiaddr::Protocol::P2p(peer_id) = protocol {
+                            return Some((peer_id, 1000)); // Equal weight for all peers
+                        }
+                    }
+                    tracing::warn!("No peer ID found in bootstrap address: {}", addr_str);
+                    None
+                }
+                Err(e) => {
+                    tracing::error!("Failed to parse bootstrap address '{}': {}", addr_str, e);
+                    None
+                }
+            }
+        })
+        .collect()
+}
+
+/// Fallback function to create peer configurations when no bootstrap addresses are provided
+/// This creates deterministic peer IDs from node indices for testing
+fn create_propeller_peers_fallback(num_nodes: u64) -> Vec<(libp2p::PeerId, u64)> {
+    use libp2p::identity::Keypair;
+
+    (0..num_nodes)
+        .map(|node_id| {
+            // Create deterministic peer ID from node ID (same as stress test node)
+            let mut private_key = [0u8; 32];
+            private_key[0..8].copy_from_slice(&node_id.to_le_bytes());
+
+            let keypair = Keypair::ed25519_from_bytes(private_key)
+                .expect("Failed to create keypair from private key");
+            let peer_id = keypair.public().to_peer_id();
+
+            (peer_id, 1000) // Equal weight for all peers
+        })
+        .collect()
+}
+
 /// Network communication channels for different protocols
 pub enum NetworkChannels {
     Gossipsub {
@@ -30,6 +84,10 @@ pub enum NetworkChannels {
         sqmr_client: Option<SqmrClientSender<TopicType, TopicType>>,
         sqmr_server: Option<SqmrServerReceiver<TopicType, TopicType>>,
     },
+    Propeller {
+        propeller_client: Option<PropellerClient<TopicType>>,
+        propeller_messages_receiver: Option<PropellerMessageServer<TopicType>>,
+    },
 }
 
 impl NetworkChannels {
@@ -49,6 +107,11 @@ impl NetworkChannels {
                     sqmr_server.take().expect("sqmr_server should be available"),
                 ))
             }
+            NetworkChannels::Propeller { propeller_client, propeller_messages_receiver: _ } => {
+                MessageSender::Propeller(PropellerSender::new(
+                    propeller_client.take().expect("propeller_client should be available"),
+                ))
+            }
         }
     }
 
@@ -70,6 +133,13 @@ impl NetworkChannels {
                     sqmr_client.take().expect("sqmr_client should be available"),
                 )
             }
+            NetworkChannels::Propeller { propeller_messages_receiver, propeller_client: _ } => {
+                MessageReceiver::Propeller(
+                    propeller_messages_receiver
+                        .take()
+                        .expect("propeller_messages_receiver should be available"),
+                )
+            }
         }
     }
 }
@@ -80,6 +150,9 @@ pub fn create_network_manager_with_channels(
     network_config: &NetworkConfig,
     buffer_size: usize,
     protocol: &NetworkProtocol,
+    num_nodes: u64,
+    _current_node_id: u64,
+    bootstrap_addrs: &[String],
 ) -> (NetworkManager, NetworkChannels) {
     let network_metrics = create_network_metrics();
     let mut network_manager =
@@ -90,10 +163,8 @@ pub fn create_network_manager_with_channels(
             let network_channels = network_manager
                 .register_broadcast_topic::<TopicType>(TOPIC.clone(), buffer_size)
                 .expect("Failed to register broadcast topic");
-            let BroadcastTopicChannels {
-                broadcasted_messages_receiver,
-                broadcast_topic_client,
-            } = network_channels;
+            let BroadcastTopicChannels { broadcasted_messages_receiver, broadcast_topic_client } =
+                network_channels;
 
             NetworkChannels::Gossipsub {
                 broadcast_topic_client: Some(broadcast_topic_client),
@@ -112,10 +183,7 @@ pub fn create_network_manager_with_channels(
                     buffer_size,
                 );
 
-            NetworkChannels::Sqmr {
-                sqmr_client: Some(sqmr_client),
-                sqmr_server: Some(sqmr_server),
-            }
+            NetworkChannels::Sqmr { sqmr_client: Some(sqmr_client), sqmr_server: Some(sqmr_server) }
         }
         NetworkProtocol::ReveresedSqmr => {
             let sqmr_client = network_manager
@@ -134,6 +202,43 @@ pub fn create_network_manager_with_channels(
                 sqmr_server: Some(sqmr_server),
             }
         }
+        NetworkProtocol::Propeller => {
+            // Create peer configurations from bootstrap addresses, or fallback to generated peers
+            let peers = if !bootstrap_addrs.is_empty() {
+                tracing::info!(
+                    "Creating Propeller peers from {} bootstrap addresses",
+                    bootstrap_addrs.len()
+                );
+                let peers = create_propeller_peers_from_bootstrap(bootstrap_addrs);
+                tracing::info!(
+                    "Successfully created {} Propeller peers from bootstrap addresses",
+                    peers.len()
+                );
+                for (peer_id, weight) in &peers {
+                    tracing::debug!("Propeller peer: {} (weight: {})", peer_id, weight);
+                }
+                peers
+            } else {
+                tracing::info!(
+                    "No bootstrap addresses provided, using fallback peer generation for {} nodes",
+                    num_nodes
+                );
+                create_propeller_peers_fallback(num_nodes)
+            };
+
+            // Register propeller channels
+            let propeller_channels = network_manager
+                .register_propeller_channels::<TopicType>(buffer_size, peers)
+                .expect("Failed to register propeller channels");
+
+            let PropellerChannels { propeller_messages_receiver, propeller_client } =
+                propeller_channels;
+
+            NetworkChannels::Propeller {
+                propeller_client: Some(propeller_client),
+                propeller_messages_receiver: Some(propeller_messages_receiver),
+            }
+        }
     };
 
     (network_manager, channels)

crates/apollo_network/src/bin/broadcast_network_stress_test_node/run/args.py

@@ -41,7 +41,7 @@ def add_shared_args_to_parser(parser: argparse.ArgumentParser):
     parser.add_argument(
         "--network-protocol",
         help="The network protocol to use for communication.",
-        choices=["gossipsub", "sqmr", "reversed-sqmr"],
+        choices=["gossipsub", "sqmr", "reversed-sqmr", "propeller"],
         default="gossipsub",
     )
     parser.add_argument(