kabu_types_blockchain/
mempool.rs

1use crate::{AccountNonceAndTransactions, FetchState, GethStateUpdate, MempoolTx};
2use alloy_consensus::Transaction;
3use alloy_consensus::transaction::{SignerRecoverable, TxHashRef};
4use alloy_primitives::BlockNumber;
5use alloy_primitives::map::HashMap;
6use alloy_primitives::{Address, TxHash};
7use alloy_rpc_types_eth::Log;
8use chrono::{DateTime, Utc};
9use eyre::{Result, eyre};
10use reth_node_types::NodePrimitives;
11use reth_primitives_traits::SignedTransaction;
12use std::collections::hash_map::Entry;
13
14#[derive(Debug, Default)]
15pub struct Mempool<N: NodePrimitives> {
16    pub txs: HashMap<TxHash, MempoolTx<N>>,
17    accounts: HashMap<Address, AccountNonceAndTransactions>,
18}
19
20impl<N: NodePrimitives> Mempool<N> {
21    pub fn new() -> Mempool<N> {
22        Mempool { txs: HashMap::default(), accounts: HashMap::default() }
23    }
24
25    pub fn len(&self) -> usize {
26        self.txs.len()
27    }
28
29    pub fn is_empty(&self) -> bool {
30        self.txs.is_empty()
31    }
32
33    pub fn is_tx(&self, tx_hash: &TxHash) -> bool {
34        self.txs.contains_key(tx_hash)
35    }
36
37    pub fn add_tx(&mut self, tx: N::SignedTx) -> &mut Self {
38        let entry = self.txs.entry(*tx.tx_hash()).or_default();
39        entry.tx = Some(tx);
40        self
41    }
42
43    pub fn add_tx_logs(&mut self, tx_hash: TxHash, logs: Vec<Log>) -> &mut Self {
44        let entry = self.txs.entry(tx_hash).or_default();
45        entry.logs = Some(logs);
46        self
47    }
48
49    pub fn add_tx_state_change(&mut self, tx_hash: TxHash, state_update: GethStateUpdate) -> &mut Self {
50        let entry = self.txs.entry(tx_hash).or_default();
51        entry.state_update = Some(state_update);
52        self
53    }
54
55    pub fn filter_by_gas_price(&self, gas_price: u128) -> Vec<&MempoolTx<N>>
56    where
57        N::SignedTx: SignedTransaction,
58    {
59        self.txs
60            .values()
61            .filter(|&item| item.mined.is_none() && item.tx.as_ref().is_some_and(|tx| tx.max_fee_per_gas() >= gas_price))
62            .collect()
63    }
64
65    pub fn filter_ok_by_gas_price(&self, gas_price: u128) -> Vec<&MempoolTx<N>>
66    where
67        N::SignedTx: SignedTransaction,
68    {
69        self.txs
70            .values()
71            .filter(|&item| {
72                item.mined.is_none()
73                    && !item.failed.unwrap_or(false)
74                    && item.tx.as_ref().is_some_and(|tx| tx.max_fee_per_gas() >= gas_price)
75            })
76            .collect()
77    }
78
79    pub fn filter_on_block(&self, block_number: BlockNumber) -> Vec<&MempoolTx<N>> {
80        self.txs.values().filter(|&item| item.mined == Some(block_number)).collect()
81    }
82
83    pub fn is_mined(&self, tx_hash: &TxHash) -> bool {
84        match self.txs.get(tx_hash) {
85            Some(tx) => tx.mined.is_some(),
86            None => false,
87        }
88    }
89
90    pub fn is_failed(&self, tx_hash: &TxHash) -> bool {
91        match self.txs.get(tx_hash) {
92            Some(e) => e.failed.unwrap_or(false),
93            None => false,
94        }
95    }
96
97    pub fn clean(&mut self) {
98        self.txs = Default::default();
99        self.accounts = Default::default();
100    }
101
102    pub fn clean_txs(&mut self, max_block_number: BlockNumber, max_time: DateTime<Utc>) {
103        self.txs = self
104            .txs
105            .clone()
106            .into_iter()
107            .filter(|(_, v)| v.mined.unwrap_or(max_block_number + 1) > max_block_number && v.time > max_time)
108            .collect();
109    }
110
111    pub fn set_mined(&mut self, tx_hash: TxHash, block_number: BlockNumber) -> &mut Self {
112        let entry = self.txs.entry(tx_hash).or_default();
113        entry.mined = Some(block_number);
114        self
115    }
116
117    pub fn set_failed(&mut self, tx_hash: TxHash) {
118        if let Entry::Occupied(mut e) = self.txs.entry(tx_hash) {
119            let value = e.get_mut();
120            value.failed = Some(true)
121        }
122    }
123
124    pub fn set_nonce(&mut self, account: Address, nonce: u64) -> &mut Self {
125        let entry = self.accounts.entry(account).or_default();
126        entry.set_nonce(Some(nonce));
127        self
128    }
129
130    pub fn is_valid_tx(&self, tx: &N::SignedTx) -> bool
131    where
132        N::SignedTx: SignedTransaction,
133    {
134        let from = match tx.recover_signer() {
135            Ok(addr) => addr,
136            Err(_) => return false,
137        };
138        self.accounts.get(&from).is_none_or(|acc| acc.nonce.is_none_or(|nonce| tx.nonce() == nonce + 1))
139    }
140
141    pub fn get_tx_by_hash(&self, tx_hash: &TxHash) -> Option<&MempoolTx<N>> {
142        self.txs.get(tx_hash)
143    }
144
145    pub fn get_or_fetch_pre_state(&mut self, _tx_hash: &TxHash) -> Result<FetchState<GethStateUpdate>> {
146        Err(eyre!("NOT_IMPLEMENTED"))
147    }
148
149    pub fn remove_tx(&mut self, tx_hash: &TxHash) -> Option<MempoolTx<N>> {
150        self.txs.remove(tx_hash)
151    }
152}
kabu_types_blockchain/mempool.rs

kabu_types_blockchain/
mempool.rs
