kabu_core_mempool/
mempool_component.rs

1use eyre::Result;
2use kabu_core_components::Component;
3use kabu_types_blockchain::{ChainParameters, Mempool};
4use kabu_types_events::{MarketEvents, MempoolEvents, MessageMempoolDataUpdate};
5use reth_ethereum_primitives::EthPrimitives;
6use reth_node_types::NodePrimitives;
7use reth_tasks::TaskExecutor;
8use std::sync::Arc;
9use tokio::sync::{RwLock, broadcast};
10use tracing::{error, info};
11
12pub struct MempoolComponent<N: NodePrimitives = EthPrimitives> {
13    chain_parameters: ChainParameters,
14    mempool: Arc<RwLock<Mempool<N>>>,
15    mempool_update_rx: broadcast::Receiver<MessageMempoolDataUpdate<N>>,
16    market_events_rx: broadcast::Receiver<MarketEvents>,
17    mempool_events_tx: broadcast::Sender<MempoolEvents>,
18    influxdb_tx: Option<broadcast::Sender<influxdb::WriteQuery>>,
19}
20
21impl<N: NodePrimitives> MempoolComponent<N> {
22    pub fn new(
23        chain_parameters: ChainParameters,
24        mempool: Arc<RwLock<Mempool<N>>>,
25        mempool_update_rx: broadcast::Receiver<MessageMempoolDataUpdate<N>>,
26        market_events_rx: broadcast::Receiver<MarketEvents>,
27        mempool_events_tx: broadcast::Sender<MempoolEvents>,
28        influxdb_tx: Option<broadcast::Sender<influxdb::WriteQuery>>,
29    ) -> Self {
30        Self { chain_parameters, mempool, mempool_update_rx, market_events_rx, mempool_events_tx, influxdb_tx }
31    }
32}
33
34impl<N: NodePrimitives + 'static> Component for MempoolComponent<N> {
35    fn spawn(self, executor: TaskExecutor) -> Result<()> {
36        let name = self.name();
37
38        let mut mempool_update_rx = self.mempool_update_rx;
39        let mut market_events_rx = self.market_events_rx;
40        let mempool = self.mempool;
41        let _mempool_events_tx = self.mempool_events_tx;
42        let _influxdb_tx = self.influxdb_tx;
43        let _chain_parameters = self.chain_parameters;
44
45        // Spawn the main mempool task
46        executor.spawn_critical(name, async move {
47            info!("Starting mempool component");
48
49            loop {
50                tokio::select! {
51                    Ok(_msg) = mempool_update_rx.recv() => {
52                        // Process mempool update
53                        let mempool_guard = mempool.write().await;
54                        // Add transaction processing logic here
55                        drop(mempool_guard);
56                    }
57                    Ok(event) = market_events_rx.recv() => {
58                        match event {
59                            MarketEvents::BlockHeaderUpdate { .. } => {
60                                // Process new block header
61                                let mempool_guard = mempool.write().await;
62                                // Remove confirmed transactions based on block number
63                                drop(mempool_guard);
64                            }
65                            MarketEvents::BlockTxUpdate { .. } => {
66                                // Process block with transactions
67                                let mempool_guard = mempool.write().await;
68                                // Update mempool state based on block transactions
69                                drop(mempool_guard);
70                            }
71                            _ => {
72                                // Ignore other market events
73                            }
74                        }
75                    }
76                    else => {
77                        error!("All channels closed, stopping mempool component");
78                        break;
79                    }
80                }
81            }
82        });
83
84        Ok(())
85    }
86    fn name(&self) -> &'static str {
87        "MempoolComponent"
88    }
89}