Merge pull request #950 from openmina/fix/pubsub-optimizations

fix(p2p): Pubsub optimizations

Author: tizoc

Cargo.lock

@@ -2,6 +2,7 @@ use std::{
     fs::File,
     io::{self, Read, Write},
     path::PathBuf,
+    sync::Arc,
 };
 
 use anyhow::{bail, format_err, Result};
@@ -244,7 +245,7 @@ macro_rules! converter {
 fn converters() -> Vec<Converter> {
     vec![
         converter!("gossip", GossipNetMessageV2 =>
-                   ("new-state", MinaBlockBlockStableV2),
+                   ("new-state", Arc<MinaBlockBlockStableV2>),
                    ("snark-pool-diff", NetworkPoolSnarkPoolDiffVersionedStableV2),
                    ("tx-pool-diff", NetworkPoolTransactionPoolDiffVersionedStableV2),
         ),

mina-p2p-messages/examples/mina-types-converter.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use binprot_derive::{BinProtRead, BinProtWrite};
 use derive_more::{From, TryInto};
 use serde::{Deserialize, Serialize};
@@ -9,7 +11,7 @@ use crate::{number::Int32, v2};
 )]
 #[serde(tag = "type", content = "message", rename_all = "snake_case")]
 pub enum GossipNetMessageV2 {
-    NewState(v2::MinaBlockBlockStableV2),
+    NewState(Arc<v2::MinaBlockBlockStableV2>),
     SnarkPoolDiff {
         message: v2::NetworkPoolSnarkPoolDiffVersionedStableV2,
         nonce: Int32,

mina-p2p-messages/src/gossip.rs

@@ -57,6 +57,17 @@ impl webrtc::P2pServiceWebrtc for NodeService {
         self.p2p.sec_key.decrypt(other_pk, encrypted)
     }
 
+    #[cfg(not(feature = "p2p-webrtc"))]
+    fn auth_encrypt_and_send(
+        &mut self,
+        peer_id: PeerId,
+        other_pub_key: &PublicKey,
+        auth: ConnectionAuth,
+    ) {
+        let _ = (peer_id, other_pub_key, auth);
+    }
+
+    #[cfg(feature = "p2p-webrtc")]
     fn auth_encrypt_and_send(
         &mut self,
         peer_id: PeerId,
@@ -65,7 +76,6 @@ impl webrtc::P2pServiceWebrtc for NodeService {
     ) {
         let encrypted = auth.encrypt(&self.p2p.sec_key, other_pub_key, &mut self.rng);
         if let Some(peer) = self.peers().get(&peer_id) {
-            #[cfg(feature = "p2p-webrtc")]
             let _ = peer
                 .cmd_sender
                 .send(webrtc::PeerCmd::ConnectionAuthorizationSend(encrypted));

node/common/src/service/p2p.rs

@@ -407,6 +407,7 @@ pub enum ActionKind {
     P2pNetworkPubsubNewStream,
     P2pNetworkPubsubOutgoingData,
     P2pNetworkPubsubOutgoingMessage,
+    P2pNetworkPubsubOutgoingMessageClear,
     P2pNetworkPubsubOutgoingMessageError,
     P2pNetworkPubsubPrune,
     P2pNetworkPubsubSign,
@@ -686,7 +687,7 @@ pub enum ActionKind {
 }
 
 impl ActionKind {
-    pub const COUNT: u16 = 577;
+    pub const COUNT: u16 = 578;
 }
 
 impl std::fmt::Display for ActionKind {
@@ -1898,6 +1899,7 @@ impl ActionKindGet for P2pNetworkPubsubAction {
             Self::SignError { .. } => ActionKind::P2pNetworkPubsubSignError,
             Self::BroadcastSigned { .. } => ActionKind::P2pNetworkPubsubBroadcastSigned,
             Self::OutgoingMessage { .. } => ActionKind::P2pNetworkPubsubOutgoingMessage,
+            Self::OutgoingMessageClear { .. } => ActionKind::P2pNetworkPubsubOutgoingMessageClear,
             Self::OutgoingMessageError { .. } => ActionKind::P2pNetworkPubsubOutgoingMessageError,
             Self::OutgoingData { .. } => ActionKind::P2pNetworkPubsubOutgoingData,
         }

node/src/action_kind.rs

@@ -9,6 +9,7 @@ use openmina_core::{
         global_sub_window, in_same_checkpoint_window, in_seed_update_range, relative_sub_window,
     },
 };
+use p2p::P2pNetworkPubsubAction;
 use redux::{callback, Dispatcher, Timestamp};
 
 use crate::{
@@ -357,7 +358,11 @@ impl BlockProducerEnabled {
                     bug_condition!("Invalid state for `BlockProducerAction::BlockInjected` expected: `BlockProducerCurrentState::Produced`, found: {:?}", state.current);
                 }
 
-                let dispatcher = state_context.into_dispatcher();
+                let (dispatcher, global_state) = state_context.into_dispatcher_and_state();
+
+                #[cfg(feature = "p2p-libp2p")]
+                broadcast_injected_block(global_state, dispatcher);
+
                 dispatcher.push(BlockProducerAction::WonSlotSearch);
             }
         }
@@ -748,6 +753,24 @@ impl BlockProducerEnabled {
     }
 }
 
+#[cfg(feature = "p2p-libp2p")]
+fn broadcast_injected_block(global_state: &State, dispatcher: &mut Dispatcher<Action, State>) {
+    use mina_p2p_messages::gossip::GossipNetMessageV2;
+
+    let Some(block) = global_state
+        .block_producer
+        .as_ref()
+        .and_then(|bp| bp.current.injected_block())
+        .map(|pb| pb.block.clone())
+    else {
+        // Should be impossible, we call this immediately after having injected the block.
+        return;
+    };
+
+    let message = Box::new(GossipNetMessageV2::NewState(block));
+    dispatcher.push(P2pNetworkPubsubAction::Broadcast { message });
+}
+
 fn can_apply_supercharged_coinbase(
     block_stake_winner: &v2::NonZeroCurvePoint,
     stake_proof_sparse_ledger: &v2::MinaBaseSparseLedgerBaseStableV2,

node/src/block_producer/block_producer_reducer.rs

@@ -372,6 +372,13 @@ impl BlockProducerCurrentState {
         }
     }
 
+    pub fn injected_block(&self) -> Option<&ArcBlockWithHash> {
+        match self {
+            Self::Injected { block, .. } => Some(block),
+            _ => None,
+        }
+    }
+
     pub fn produced_block_with_chain(&self) -> Option<(&ArcBlockWithHash, &[AppliedBlock])> {
         match self {
             Self::Produced { chain, block, .. } => Some((block, chain)),

node/src/block_producer/block_producer_state.rs

@@ -1,11 +1,20 @@
-use std::time::Duration;
+use std::{
+    str,
+    sync::{Arc, Mutex},
+    time::Duration,
+};
 
-use mina_p2p_messages::v2;
-use node::transition_frontier::genesis::{GenesisConfig, NonStakers};
+use mina_p2p_messages::{binprot::BinProtRead, gossip, v2};
+use node::{
+    p2p::{P2pNetworkAction, P2pNetworkPubsubAction, PeerId},
+    transition_frontier::genesis::{GenesisConfig, NonStakers},
+    Action, ActionWithMeta, P2pAction,
+};
 
 use crate::{
     node::Recorder,
     scenarios::{ClusterRunner, RunCfgAdvanceTime},
+    service::NodeTestingService,
     simulator::{Simulator, SimulatorConfig, SimulatorRunUntil},
 };
 
@@ -20,9 +29,58 @@ use crate::{
 pub struct MultiNodePubsubPropagateBlock;
 
 impl MultiNodePubsubPropagateBlock {
-    const WORKERS: usize = 20;
+    const WORKERS: usize = 10;
 
     pub async fn run(self, mut runner: ClusterRunner<'_>) {
+        let graph = Arc::new(Mutex::new("digraph {\n".to_owned()));
+        let factory = || {
+            let graph = graph.clone();
+            move |_id,
+                  state: &node::State,
+                  _service: &NodeTestingService,
+                  action: &ActionWithMeta| {
+                let this = state.p2p.my_id();
+
+                let cut = |peer_id: &PeerId| {
+                    let st = peer_id.to_string();
+                    let len = st.len();
+                    st[(len - 6)..len].to_owned()
+                };
+                let this = cut(&this);
+
+                match action.action() {
+                    Action::P2p(P2pAction::Network(P2pNetworkAction::Pubsub(
+                        P2pNetworkPubsubAction::OutgoingMessage { peer_id },
+                    ))) => {
+                        let pubsub_state =
+                            &state.p2p.ready().unwrap().network.scheduler.broadcast_state;
+                        let msg = &pubsub_state.clients.get(peer_id).unwrap().message;
+
+                        for publish_message in &msg.publish {
+                            let mut slice = &publish_message.data()[8..];
+                            if let Ok(gossip::GossipNetMessageV2::NewState(block)) =
+                                gossip::GossipNetMessageV2::binprot_read(&mut slice)
+                            {
+                                let height = block
+                                    .header
+                                    .protocol_state
+                                    .body
+                                    .consensus_state
+                                    .global_slot();
+                                let mut lock = graph.lock().unwrap();
+                                *lock = format!(
+                                    "{lock}  \"{this}\" -> \"{}\" [label=\"{height}\"];\n",
+                                    cut(peer_id)
+                                );
+                            }
+                        }
+                        false
+                    }
+                    _ => false,
+                }
+            }
+        };
+
         let initial_time = redux::Timestamp::global_now();
         let mut constants = v2::PROTOCOL_CONSTANTS.clone();
         constants.genesis_state_timestamp =
@@ -40,11 +98,15 @@ impl MultiNodePubsubPropagateBlock {
             snark_workers: 1,
             block_producers: 1,
             advance_time: RunCfgAdvanceTime::Rand(1..=200),
-            run_until: SimulatorRunUntil::BlockchainLength(3),
+            run_until: SimulatorRunUntil::BlockchainLength(4),
             run_until_timeout: Duration::from_secs(10 * 60),
             recorder: Recorder::StateWithInputActions,
         };
         let mut simulator = Simulator::new(initial_time, config);
-        simulator.setup_and_run(&mut runner).await;
+        simulator
+            .setup_and_run_with_listener(&mut runner, factory)
+            .await;
+
+        println!("{}}}\n", graph.lock().unwrap());
     }
 }

node/testing/src/scenarios/multi_node/pubsub_advanced.rs

@@ -7,13 +7,16 @@ use mina_p2p_messages::v2::{
 
 use std::{collections::BTreeSet, time::Duration};
 
-use node::{ActionKind, BlockProducerConfig, SnarkerConfig, SnarkerStrategy, State};
+use node::{
+    ActionKind, ActionWithMeta, BlockProducerConfig, SnarkerConfig, SnarkerStrategy, State,
+};
 
 use crate::{
     cluster::ClusterNodeId,
     node::{Node, RustNodeBlockProducerTestingConfig, RustNodeTestingConfig},
     scenario::ListenerNode,
     scenarios::{ClusterRunner, RunCfg},
+    service::NodeTestingService,
 };
 
 pub struct Simulator {
@@ -217,19 +220,46 @@ impl Simulator {
         self.wait_for_all_nodes_synced(runner).await;
     }
 
+    pub async fn setup_and_run_with_listener<'a, AL, ALF>(
+        &mut self,
+        runner: &mut ClusterRunner<'a>,
+        listener: ALF,
+    ) where
+        ALF: FnMut() -> AL,
+        AL: 'static
+            + Send
+            + FnMut(ClusterNodeId, &State, &NodeTestingService, &ActionWithMeta) -> bool,
+    {
+        self.setup(runner).await;
+        self.run_with_listener(runner, listener).await;
+    }
+
+    pub async fn setup_and_run<'a>(&mut self, runner: &mut ClusterRunner<'a>) {
+        self.setup(runner).await;
+        self.run_with_listener(runner, || |_, _, _, _| false).await;
+    }
+
     pub async fn setup<'a>(&mut self, runner: &mut ClusterRunner<'a>) {
         self.set_up_seed_nodes(runner).await;
         self.set_up_normal_nodes(runner).await;
         self.set_up_snark_worker_nodes(runner).await;
         self.set_up_block_producer_nodes(runner).await;
     }
 
-    pub async fn setup_and_run<'a>(&mut self, runner: &mut ClusterRunner<'a>) {
-        self.setup(runner).await;
-        self.run(runner).await;
+    pub async fn run<'a>(&mut self, runner: &mut ClusterRunner<'a>) {
+        self.run_with_listener(runner, || |_, _, _, _| false).await;
     }
 
-    pub async fn run<'a>(&mut self, runner: &mut ClusterRunner<'a>) {
+    pub async fn run_with_listener<'a, AL, ALF>(
+        &mut self,
+        runner: &mut ClusterRunner<'a>,
+        mut listener: ALF,
+    ) where
+        ALF: FnMut() -> AL,
+        AL: 'static
+            + Send
+            + FnMut(ClusterNodeId, &State, &NodeTestingService, &ActionWithMeta) -> bool,
+    {
         let run_until = self.config.run_until.clone();
         let advance_time = self.config.advance_time.clone();
         let start_t = *self.start_t.get_or_insert_with(redux::Instant::now);
@@ -238,13 +268,11 @@ impl Simulator {
 
         while start_t.elapsed() < self.config.run_until_timeout {
             tokio::task::yield_now().await;
-            let _ = runner
-                .run(
-                    RunCfg::default()
-                        .advance_time(advance_time.clone())
-                        .timeout(Duration::ZERO),
-                )
-                .await;
+            let cfg = RunCfg::default()
+                .advance_time(advance_time.clone())
+                .timeout(Duration::ZERO)
+                .action_handler(listener());
+            let _ = runner.run(cfg).await;
 
             let printed_elapsed_time = {
                 let state = runner.nodes_iter().next().unwrap().1.state();

node/testing/src/simulator/mod.rs

@@ -3,7 +3,7 @@ use redux::ActionWithMeta;
 
 use crate::{
     channels::{ChannelId, ChannelMsg, MsgId, P2pChannelsEffectfulAction},
-    P2pNetworkPubsubAction, P2pPeerAction, P2pState, PeerId,
+    P2pPeerAction, P2pState, PeerId,
 };
 
 use super::{
@@ -162,23 +162,6 @@ impl P2pChannelsBestTipState {
                     return Ok(());
                 }
 
-                #[cfg(feature = "p2p-libp2p")]
-                {
-                    use mina_p2p_messages::gossip::GossipNetMessageV2;
-                    let block = (*best_tip.block).clone();
-                    let message = Box::new(GossipNetMessageV2::NewState(block));
-                    // TODO(vlad): `P2pChannelsBestTipAction::ResponseSend`
-                    // action is dispatched for each peer. So `P2pNetworkPubsubAction::Broadcast`
-                    // will be called many times causing many duplicate
-                    // broadcasts. Either in pubsub state machine, we
-                    // need to filter out duplicate messages, or better,
-                    // have a simple action to send pubsub message to a
-                    // specific peer instead of sending to everyone.
-                    // That way we can avoid duplicate state, since we
-                    // already store last sent best tip here and we make
-                    // sure we don't send same block to same peer again.
-                    dispatcher.push(P2pNetworkPubsubAction::Broadcast { message });
-                }
                 Ok(())
             }
         }