Merge pull request #953 from openmina/fix/pubsub-tweaks

fix(p2p): Pubsub message handling tweaks

Author: tizoc

node/src/action_kind.rs

@@ -404,6 +404,7 @@ pub enum ActionKind {
     P2pNetworkPubsubGraft,
     P2pNetworkPubsubIncomingData,
     P2pNetworkPubsubIncomingMessage,
+    P2pNetworkPubsubIncomingMessageCleanup,
     P2pNetworkPubsubNewStream,
     P2pNetworkPubsubOutgoingData,
     P2pNetworkPubsubOutgoingMessage,
@@ -687,7 +688,7 @@ pub enum ActionKind {
 }
 
 impl ActionKind {
-    pub const COUNT: u16 = 578;
+    pub const COUNT: u16 = 579;
 }
 
 impl std::fmt::Display for ActionKind {
@@ -1892,6 +1893,9 @@ impl ActionKindGet for P2pNetworkPubsubAction {
             Self::NewStream { .. } => ActionKind::P2pNetworkPubsubNewStream,
             Self::IncomingData { .. } => ActionKind::P2pNetworkPubsubIncomingData,
             Self::IncomingMessage { .. } => ActionKind::P2pNetworkPubsubIncomingMessage,
+            Self::IncomingMessageCleanup { .. } => {
+                ActionKind::P2pNetworkPubsubIncomingMessageCleanup
+            }
             Self::Graft { .. } => ActionKind::P2pNetworkPubsubGraft,
             Self::Prune { .. } => ActionKind::P2pNetworkPubsubPrune,
             Self::Broadcast { .. } => ActionKind::P2pNetworkPubsubBroadcast,

p2p/src/network/pubsub/p2p_network_pubsub_actions.rs

@@ -25,6 +25,9 @@ pub enum P2pNetworkPubsubAction {
         message: pb::Message,
         seen_limit: usize,
     },
+    IncomingMessageCleanup {
+        peer_id: PeerId,
+    },
     Graft {
         peer_id: PeerId,
         topic_id: String,

p2p/src/network/pubsub/p2p_network_pubsub_reducer.rs

@@ -160,51 +160,79 @@ impl P2pNetworkPubsubState {
                 message,
                 seen_limit,
             } => {
-                pubsub_state.reduce_incoming_message(peer_id, message, seen_limit)?;
+                // Check result later to ensure we always dispatch the cleanup action
+                let reduce_incoming_result =
+                    pubsub_state.reduce_incoming_message(peer_id, message, seen_limit);
 
                 let (dispatcher, global_state) = state_context.into_dispatcher_and_state();
+
+                dispatcher.push(P2pNetworkPubsubAction::IncomingMessageCleanup { peer_id });
+
+                reduce_incoming_result?;
+
                 let state: &Self = global_state.substate()?;
                 let config: &P2pConfig = global_state.substate()?;
 
-                let incoming_block = state.incoming_block.as_ref().cloned();
-                let incoming_transactions = state.incoming_transactions.clone();
-                let incoming_snarks = state.incoming_snarks.clone();
-                let topics = state.topics.clone();
-
-                for (topic_id, map) in topics {
+                for (topic_id, map) in &state.topics {
                     let mesh_size = map.values().filter(|s| s.on_mesh()).count();
                     let could_accept = mesh_size < config.meshsub.outbound_degree_high;
 
                     if !could_accept {
                         if let Some(topic_state) = map.get(&peer_id) {
                             if topic_state.on_mesh() {
+                                let topic_id = topic_id.clone();
                                 dispatcher.push(P2pNetworkPubsubAction::Prune { peer_id, topic_id })
                             }
                         }
                     }
                 }
 
-                broadcast(dispatcher, global_state)?;
-                if let Some((_, block)) = incoming_block {
-                    let best_tip = BlockWithHash::try_new(block)?;
+                if let Err(error) = broadcast(dispatcher, global_state) {
+                    bug_condition!(
+                        "Failure when trying to broadcast incoming pubsub message: {error}"
+                    );
+                };
+
+                if let Some((_, block)) = state.incoming_block.as_ref() {
+                    let best_tip = BlockWithHash::try_new(block.clone())?;
                     dispatcher.push(P2pPeerAction::BestTipUpdate { peer_id, best_tip });
                 }
-                for (transaction, nonce) in incoming_transactions {
+                for (transaction, nonce) in &state.incoming_transactions {
                     dispatcher.push(P2pChannelsTransactionAction::Libp2pReceived {
                         peer_id,
-                        transaction: Box::new(transaction),
-                        nonce,
+                        transaction: Box::new(transaction.clone()),
+                        nonce: *nonce,
                     });
                 }
-                for (snark, nonce) in incoming_snarks {
+                for (snark, nonce) in &state.incoming_snarks {
                     dispatcher.push(P2pChannelsSnarkAction::Libp2pReceived {
                         peer_id,
-                        snark: Box::new(snark),
-                        nonce,
+                        snark: Box::new(snark.clone()),
+                        nonce: *nonce,
                     });
                 }
                 Ok(())
             }
+            P2pNetworkPubsubAction::IncomingMessageCleanup { peer_id } => {
+                pubsub_state.incoming_transactions.clear();
+                pubsub_state.incoming_snarks.clear();
+
+                pubsub_state.incoming_transactions.shrink_to(0x20);
+                pubsub_state.incoming_snarks.shrink_to(0x20);
+
+                pubsub_state.incoming_block = None;
+
+                let Some(client_state) = pubsub_state.clients.get_mut(&peer_id) else {
+                    bug_condition!(
+                        "State not found for action P2pNetworkPubsubAction::IncomingMessageCleanup"
+                    );
+                    return Ok(());
+                };
+                client_state.incoming_messages.clear();
+                client_state.incoming_messages.shrink_to(0x20);
+
+                Ok(())
+            }
             // we want to add peer to our mesh
             P2pNetworkPubsubAction::Graft { peer_id, topic_id } => {
                 let Some(state) = pubsub_state
@@ -440,21 +468,6 @@ impl P2pNetworkPubsubState {
         message: Message,
         seen_limit: usize,
     ) -> Result<(), String> {
-        self.incoming_transactions.clear();
-        self.incoming_snarks.clear();
-
-        self.incoming_transactions.shrink_to(0x20);
-        self.incoming_snarks.shrink_to(0x20);
-
-        let Some(state) = self.clients.get_mut(&peer_id) else {
-            bug_condition!("State not found for action P2pNetworkPubsubAction::IncomingMessage");
-            return Ok(());
-        };
-        state.incoming_messages.clear();
-        state.incoming_messages.shrink_to(0x20);
-
-        let message_id = self.mcache.put(message.clone());
-
         let topic = self.topics.entry(message.topic.clone()).or_default();
 
         if let Some(signature) = &message.signature {
@@ -470,27 +483,6 @@ impl P2pNetworkPubsubState {
             }
         }
 
-        self.clients
-            .iter_mut()
-            .filter(|(c, _)| {
-                // don't send back to who sent this
-                *c != &peer_id
-            })
-            .for_each(|(c, state)| {
-                let Some(topic_state) = topic.get(c) else {
-                    return;
-                };
-                if topic_state.on_mesh() {
-                    state.publish(&message)
-                } else {
-                    let ctr = state.message.control.get_or_insert_with(Default::default);
-                    ctr.ihave.push(pb::ControlIHave {
-                        topic_id: Some(message.topic.clone()),
-                        message_ids: message_id.clone().into_iter().collect(),
-                    })
-                }
-            });
-
         if let Some(data) = &message.data {
             if data.len() > 8 {
                 let mut slice = &data[8..];
@@ -517,6 +509,31 @@ impl P2pNetworkPubsubState {
             }
         }
 
+        let message_id = self.mcache.put(message.clone());
+
+        // TODO: this should only happen after the contents have been validated.
+        // The only validation that has happened so far is that the message can be parsed.
+        self.clients
+            .iter_mut()
+            .filter(|(c, _)| {
+                // don't send back to who sent this
+                *c != &peer_id
+            })
+            .for_each(|(c, state)| {
+                let Some(topic_state) = topic.get(c) else {
+                    return;
+                };
+                if topic_state.on_mesh() {
+                    state.publish(&message)
+                } else {
+                    let ctr = state.message.control.get_or_insert_with(Default::default);
+                    ctr.ihave.push(pb::ControlIHave {
+                        topic_id: Some(message.topic.clone()),
+                        message_ids: message_id.clone().into_iter().collect(),
+                    })
+                }
+            });
+
         Ok(())
     }