Changed action flow for transaction in pubsub

Author: 0xMimir

node/src/action_kind.rs

@@ -607,6 +607,7 @@ pub enum ActionKind {
     TransactionPoolCandidateFetchPending,
     TransactionPoolCandidateFetchSuccess,
     TransactionPoolCandidateInfoReceived,
+    TransactionPoolCandidateLibp2pTransactionsReceived,
     TransactionPoolCandidatePeerPrune,
     TransactionPoolCandidateVerifyError,
     TransactionPoolCandidateVerifyNext,
@@ -719,7 +720,7 @@ pub enum ActionKind {
 }
 
 impl ActionKind {
-    pub const COUNT: u16 = 609;
+    pub const COUNT: u16 = 610;
 }
 
 impl std::fmt::Display for ActionKind {
@@ -1553,6 +1554,9 @@ impl ActionKindGet for TransactionPoolCandidateAction {
             Self::FetchPending { .. } => ActionKind::TransactionPoolCandidateFetchPending,
             Self::FetchError { .. } => ActionKind::TransactionPoolCandidateFetchError,
             Self::FetchSuccess { .. } => ActionKind::TransactionPoolCandidateFetchSuccess,
+            Self::Libp2pTransactionsReceived { .. } => {
+                ActionKind::TransactionPoolCandidateLibp2pTransactionsReceived
+            }
             Self::VerifyNext => ActionKind::TransactionPoolCandidateVerifyNext,
             Self::VerifyPending { .. } => ActionKind::TransactionPoolCandidateVerifyPending,
             Self::VerifyError { .. } => ActionKind::TransactionPoolCandidateVerifyError,

node/src/snark_pool/candidate/snark_pool_candidate_actions.rs

@@ -112,7 +112,10 @@ impl redux::EnablingCondition<crate::State> for SnarkPoolCandidateAction {
                             }
                         })
             }
-            SnarkPoolCandidateAction::WorkVerifyNext => state.snark.work_verify.jobs.is_empty(),
+            SnarkPoolCandidateAction::WorkVerifyNext => {
+                state.snark.work_verify.jobs.is_empty()
+                    && state.transition_frontier.sync.is_synced()
+            }
             SnarkPoolCandidateAction::WorkVerifyPending {
                 peer_id, job_ids, ..
             } => {

node/src/state.rs

@@ -4,9 +4,7 @@ use std::time::Duration;
 use malloc_size_of_derive::MallocSizeOf;
 use mina_p2p_messages::v2;
 use openmina_core::constants::PROTOCOL_VERSION;
-use openmina_core::transaction::{
-    TransactionInfo, TransactionPoolMessageSource, TransactionWithHash,
-};
+use openmina_core::transaction::{TransactionInfo, TransactionWithHash};
 use p2p::P2pNetworkPubsubMessageCacheId;
 use rand::prelude::*;
 
@@ -59,7 +57,7 @@ pub use crate::transition_frontier::TransitionFrontierState;
 pub use crate::watched_accounts::WatchedAccountsState;
 pub use crate::Config;
 use crate::{config::GlobalConfig, SnarkPoolAction};
-use crate::{ActionWithMeta, RpcAction, TransactionPoolAction};
+use crate::{ActionWithMeta, RpcAction};
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
 pub struct State {
@@ -540,11 +538,12 @@ impl P2p {
                     }
                 }
             )),
-            on_p2p_channels_transaction_libp2p_received: Some(redux::callback!(
-                on_p2p_channels_transaction_libp2p_received((transactions: Vec<TransactionWithHash>, id: P2pNetworkPubsubMessageCacheId)) -> crate::Action {
-                    TransactionPoolAction::StartVerify {
-                        commands: transactions.into_iter().collect(),
-                        from_source: TransactionPoolMessageSource::pubsub(id),
+            on_p2p_channels_transactions_libp2p_received: Some(redux::callback!(
+                on_p2p_channels_transactions_libp2p_received((peer_id: PeerId, transactions: Vec<TransactionWithHash>, message_id: P2pNetworkPubsubMessageCacheId)) -> crate::Action {
+                    TransactionPoolCandidateAction::Libp2pTransactionsReceived {
+                        message_id,
+                        transactions,
+                        peer_id
                     }
                 }
             )),

node/src/transaction_pool/candidate/transaction_pool_candidate_actions.rs

@@ -1,5 +1,6 @@
 use openmina_core::transaction::{TransactionHash, TransactionInfo, TransactionWithHash};
 use openmina_core::ActionEvent;
+use p2p::P2pNetworkPubsubMessageCacheId;
 use serde::{Deserialize, Serialize};
 
 use crate::p2p::channels::rpc::P2pRpcId;
@@ -37,12 +38,19 @@ pub enum TransactionPoolCandidateAction {
         peer_id: PeerId,
         transaction: TransactionWithHash,
     },
+    /// Callback for transactions received over pubsub
+    Libp2pTransactionsReceived {
+        peer_id: PeerId,
+        transactions: Vec<TransactionWithHash>,
+        message_id: P2pNetworkPubsubMessageCacheId,
+    },
     #[action_event(level = trace)]
     VerifyNext,
     VerifyPending {
         peer_id: PeerId,
         transaction_hashes: Vec<TransactionHash>,
         verify_id: (),
+        from_source: Option<P2pNetworkPubsubMessageCacheId>,
     },
     VerifyError {
         peer_id: PeerId,
@@ -51,6 +59,7 @@ pub enum TransactionPoolCandidateAction {
     VerifySuccess {
         peer_id: PeerId,
         verify_id: (),
+        from_source: Option<P2pNetworkPubsubMessageCacheId>,
     },
     PeerPrune {
         peer_id: PeerId,
@@ -100,7 +109,10 @@ impl redux::EnablingCondition<crate::State> for TransactionPoolCandidateAction {
                 .candidates
                 .get(*peer_id, transaction.hash())
                 .is_some(),
-            TransactionPoolCandidateAction::VerifyNext => true,
+            TransactionPoolCandidateAction::Libp2pTransactionsReceived { .. } => true,
+            TransactionPoolCandidateAction::VerifyNext => {
+                state.transition_frontier.sync.is_synced()
+            }
             TransactionPoolCandidateAction::VerifyPending {
                 peer_id,
                 transaction_hashes,

node/src/transaction_pool/candidate/transaction_pool_candidate_reducer.rs

@@ -90,38 +90,55 @@ impl TransactionPoolCandidatesState {
             } => {
                 state.transaction_received(meta.time(), *peer_id, transaction.clone());
             }
+            TransactionPoolCandidateAction::Libp2pTransactionsReceived {
+                peer_id,
+                transactions,
+                message_id,
+            } => {
+                state.transactions_received(
+                    meta.time(),
+                    *peer_id,
+                    transactions.clone(),
+                    *message_id,
+                );
+            }
             TransactionPoolCandidateAction::VerifyNext => {
                 let (dispatcher, global_state) = state_context.into_dispatcher_and_state();
 
                 let batch = global_state
                     .transaction_pool
                     .candidates
                     .get_batch_to_verify();
-                let Some((peer_id, batch)) = batch else {
+                let Some((peer_id, batch, from_source)) = batch else {
                     return;
                 };
 
                 let transaction_hashes = batch.iter().map(|tx| tx.hash().clone()).collect();
                 dispatcher.push(TransactionPoolAction::StartVerify {
                     commands: batch.into_iter().collect(),
-                    from_source: TransactionPoolMessageSource::None,
+                    from_source: from_source
+                        .map(TransactionPoolMessageSource::pubsub)
+                        .unwrap_or_default(),
                 });
                 dispatcher.push(TransactionPoolCandidateAction::VerifyPending {
                     peer_id,
                     transaction_hashes,
                     verify_id: (),
+                    from_source,
                 });
             }
             TransactionPoolCandidateAction::VerifyPending {
                 peer_id,
                 transaction_hashes,
                 verify_id,
+                from_source,
             } => {
                 state.verify_pending(meta.time(), peer_id, *verify_id, transaction_hashes);
                 let dispatcher = state_context.into_dispatcher();
                 dispatcher.push(TransactionPoolCandidateAction::VerifySuccess {
                     peer_id: *peer_id,
                     verify_id: *verify_id,
+                    from_source: *from_source,
                 });
             }
             TransactionPoolCandidateAction::VerifyError {
@@ -139,8 +156,12 @@ impl TransactionPoolCandidatesState {
                 //     reason: P2pDisconnectionReason::TransactionPoolVerifyError,
                 // });
             }
-            TransactionPoolCandidateAction::VerifySuccess { peer_id, verify_id } => {
-                state.verify_result(meta.time(), peer_id, *verify_id, Ok(()));
+            TransactionPoolCandidateAction::VerifySuccess {
+                peer_id,
+                verify_id,
+                from_source,
+            } => {
+                state.verify_result(meta.time(), peer_id, *verify_id, from_source, Ok(()));
             }
             TransactionPoolCandidateAction::PeerPrune { peer_id } => {
                 state.peer_remove(*peer_id);

node/src/transaction_pool/candidate/transaction_pool_candidate_state.rs

@@ -3,6 +3,7 @@
 use std::collections::{BTreeMap, BTreeSet};
 
 use mina_p2p_messages::v2;
+use p2p::P2pNetworkPubsubMessageCacheId;
 use redux::Timestamp;
 use serde::{Deserialize, Serialize};
 
@@ -15,10 +16,17 @@ use crate::p2p::PeerId;
 static EMPTY_PEER_TX_CANDIDATES: BTreeMap<TransactionHash, TransactionPoolCandidateState> =
     BTreeMap::new();
 
+type NextBatch = Option<(
+    PeerId,
+    Vec<TransactionWithHash>,
+    Option<P2pNetworkPubsubMessageCacheId>,
+)>;
+
 #[derive(Serialize, Deserialize, Debug, Clone, Default)]
 pub struct TransactionPoolCandidatesState {
     by_peer: BTreeMap<PeerId, BTreeMap<TransactionHash, TransactionPoolCandidateState>>,
     by_hash: BTreeMap<TransactionHash, BTreeSet<PeerId>>,
+    by_message_id: BTreeMap<P2pNetworkPubsubMessageCacheId, (PeerId, Vec<TransactionHash>)>,
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -68,6 +76,10 @@ impl TransactionPoolCandidatesState {
         self.by_hash.contains_key(hash)
     }
 
+    pub fn message_id_contains(&self, message_id: &P2pNetworkPubsubMessageCacheId) -> bool {
+        self.by_message_id.contains_key(message_id)
+    }
+
     pub fn peer_contains(&self, peer_id: PeerId, hash: &TransactionHash) -> bool {
         self.by_peer
             .get(&peer_id)
@@ -203,10 +215,42 @@ impl TransactionPoolCandidatesState {
         self.by_peer.entry(peer_id).or_default().insert(hash, state);
     }
 
-    pub fn get_batch_to_verify(&self) -> Option<(PeerId, Vec<TransactionWithHash>)> {
-        for hash in self.by_hash.keys() {
+    pub fn transactions_received(
+        &mut self,
+        time: Timestamp,
+        peer_id: PeerId,
+        transactions: Vec<TransactionWithHash>,
+        message_id: P2pNetworkPubsubMessageCacheId,
+    ) {
+        let transaction_hashes = transactions
+            .iter()
+            .map(TransactionWithHash::hash)
+            .cloned()
+            .collect::<Vec<_>>();
+
+        self.by_message_id
+            .insert(message_id, (peer_id, transaction_hashes));
+
+        transactions.into_iter().for_each(|transaction| {
+            self.transaction_received(time, peer_id, transaction);
+        })
+    }
+
+    /// Get next batch of transactions to verify, 
+    /// first checks if there are any transactions to verify from pubsub
+    /// after that checks for transactions from peers
+    pub fn get_batch_to_verify(&self) -> NextBatch {
+        if let Some(batch) = self.next_batch_from_pubsub() {
+            return Some(batch);
+        }
+
+        self.next_batch_from_peers()
+    }
+
+    fn next_batch_from_peers(&self) -> NextBatch {
+        for (hash, peers) in self.by_hash.iter() {
             if let Some(res) = None.or_else(|| {
-                for peer_id in self.by_hash.get(hash)? {
+                for peer_id in peers {
                     let peer_transactions = self.by_peer.get(peer_id)?;
                     if peer_transactions.get(hash)?.transaction().is_some() {
                         let transactions = peer_transactions
@@ -219,17 +263,39 @@ impl TransactionPoolCandidatesState {
                             })
                             .cloned()
                             .collect();
-                        return Some((*peer_id, transactions));
+                        return Some((*peer_id, transactions, None));
                     }
                 }
                 None
             }) {
                 return Some(res);
             }
         }
+
         None
     }
 
+    fn next_batch_from_pubsub(&self) -> NextBatch {
+        let (message_id, (peer_id, transaction_hashes)) = self.by_message_id.iter().next()?;
+        let transactions = self
+            .by_peer
+            .get(peer_id)?
+            .iter()
+            .filter_map(|(hash, state)| {
+                let TransactionPoolCandidateState::Received { transaction, .. } = state else {
+                    return None;
+                };
+                if transaction_hashes.contains(hash) {
+                    Some(transaction)
+                } else {
+                    None
+                }
+            })
+            .cloned()
+            .collect();
+
+        return Some((*peer_id, transactions, Some(*message_id)));
+    }
     pub fn verify_pending(
         &mut self,
         time: Timestamp,
@@ -259,8 +325,21 @@ impl TransactionPoolCandidatesState {
         _time: Timestamp,
         peer_id: &PeerId,
         verify_id: (),
+        from_source: &Option<P2pNetworkPubsubMessageCacheId>,
         _result: Result<(), ()>,
     ) {
+        if let Some(from_source) = from_source {
+            let Some((_, transactions)) = self.by_message_id.remove(from_source) else {
+                return;
+            };
+
+            for hash in transactions {
+                self.transaction_remove(&hash);
+            }
+
+            return;
+        }
+
         if let Some(peer_transactions) = self.by_peer.get_mut(peer_id) {
             let txs_to_remove = peer_transactions
                 .iter()
@@ -344,6 +423,20 @@ impl TransactionPoolCandidatesState {
             !peers.is_empty()
         })
     }
+
+    pub fn remove_pubsub_transactions(&mut self, message_id: &P2pNetworkPubsubMessageCacheId) {
+        let Some((peer_id, transactions)) = self.by_message_id.remove(message_id) else {
+            return;
+        };
+
+        let Some(peer_transactions) = self.by_peer.get_mut(&peer_id) else {
+            return;
+        };
+
+        transactions.into_iter().for_each(|tx| {
+            peer_transactions.remove(&tx);
+        });
+    }
 }
 
 impl TransactionPoolCandidateState {

node/src/transaction_pool/transaction_pool_reducer.rs

@@ -165,7 +165,7 @@ impl TransactionPoolState {
                                         }),
                                         peer_id: None,
                                         reason: "Transaction diff rejected".to_owned(),
-                                    })
+                                    });
                                 }
                                 TransactionPoolMessageSource::None => {}
                             }

p2p/src/channels/transaction/p2p_channels_transaction_reducer.rs

@@ -215,22 +215,23 @@ impl P2pChannelsTransactionState {
             P2pChannelsTransactionAction::Libp2pReceived {
                 transactions,
                 message_id,
+                peer_id,
                 ..
             } => {
                 let (dispatcher, state) = state_context.into_dispatcher_and_state();
                 let p2p_state: &P2pState = state.substate()?;
 
                 if let Some(callback) = &p2p_state
                     .callbacks
-                    .on_p2p_channels_transaction_libp2p_received
+                    .on_p2p_channels_transactions_libp2p_received
                 {
                     let transactions = transactions
                         .into_iter()
                         .map(TransactionWithHash::try_new)
                         .filter_map(Result::ok)
                         .collect();
 
-                    dispatcher.push_callback(callback.clone(), (transactions, message_id));
+                    dispatcher.push_callback(callback.clone(), (peer_id, transactions, message_id));
                 }
 
                 Ok(())