kabu_broadcast_accounts/accounts_monitor/
account_monitor_component.rs

1use alloy_consensus::Transaction;
2use alloy_eips::BlockId;
3use eyre::Result;
4use std::collections::HashSet;
5use std::sync::Arc;
6use tokio::sync::RwLock;
7use tracing::{debug, error, info, warn};
8
9use alloy_network::Network;
10use alloy_primitives::Address;
11use alloy_primitives::utils::format_ether;
12use futures_util::StreamExt;
13use kabu_core_components::Component;
14use kabu_types_entities::{AccountNonceAndBalanceState, TxSigners};
15use reth_chain_state::{CanonStateNotification, CanonStateSubscriptions};
16use reth_ethereum_primitives::EthPrimitives;
17use reth_node_types::NodePrimitives;
18use reth_primitives::Block;
19use reth_primitives_traits::{RecoveredBlock, SignerRecoverable};
20use reth_provider::StateProviderFactory;
21use reth_revm::database::StateProviderDatabase;
22use reth_rpc_eth_types::cache::db::StateProviderTraitObjWrapper;
23use reth_tasks::TaskExecutor;
24use revm::DatabaseRef;
25use revm::database::CacheDB;
26
27// TODO: Does this provider has any positive effect on performance?
28
29/// Component that monitors account nonces and balances for managed accounts
30#[derive(Clone)]
31pub struct AccountMonitorComponent<R, N, P = EthPrimitives>
32where
33    P: NodePrimitives,
34{
35    /// Reth provider for canonical state subscriptions
36    reth_provider: R,
37    /// Shared state containing account nonces and balances
38    account_state: Arc<RwLock<AccountNonceAndBalanceState>>,
39    /// Signers to monitor accounts for
40    signers: Arc<RwLock<TxSigners<P>>>,
41    /// Phantom data for network type
42    _network: std::marker::PhantomData<N>,
43}
44
45impl<R, N, P> AccountMonitorComponent<R, N, P>
46where
47    R: CanonStateSubscriptions<Primitives = EthPrimitives> + StateProviderFactory + Send + Sync + Clone + 'static,
48    N: Network + 'static,
49    P: reth_node_types::NodePrimitives + 'static,
50{
51    pub fn new(reth_provider: R, account_state: Arc<RwLock<AccountNonceAndBalanceState>>, signers: Arc<RwLock<TxSigners<P>>>) -> Self {
52        Self { reth_provider, account_state, signers, _network: std::marker::PhantomData }
53    }
54
55    async fn run(self) -> Result<()> {
56        info!("Starting account monitor component");
57
58        // Get initial account list from signers
59        let monitored_accounts = self.get_monitored_accounts().await;
60        info!("Monitoring {} accounts for nonce and balance updates", monitored_accounts.len());
61
62        // Initialize account state
63        self.initialize_accounts(&monitored_accounts).await?;
64
65        // Subscribe to canonical state stream
66        let mut canon_stream = self.reth_provider.canonical_state_stream();
67
68        while let Some(canon_notification) = canon_stream.next().await {
69            if let Err(e) = self.handle_canon_state_notification(canon_notification).await {
70                error!("Error handling canonical state notification: {}", e);
71            }
72        }
73
74        Ok(())
75    }
76
77    async fn get_monitored_accounts(&self) -> HashSet<Address> {
78        let signers = self.signers.read().await;
79        signers.get_address_vec().into_iter().collect()
80    }
81
82    async fn initialize_accounts(&self, accounts: &HashSet<Address>) -> Result<()> {
83        let mut account_state = self.account_state.write().await;
84
85        for &account_address in accounts {
86            account_state.add_account(account_address);
87
88            let state = self.reth_provider.state_by_block_id(BlockId::latest())?;
89            let db = CacheDB::new(StateProviderDatabase::new(StateProviderTraitObjWrapper(&state)));
90            // Fetch initial nonce and balance
91            let account = db.basic_ref(account_address)?;
92
93            let Some(account) = account else {
94                warn!("Account {} does not exist on chain", account_address);
95                continue;
96            };
97
98            if let Some(account_data) = account_state.get_mut_account(&account_address) {
99                account_data.set_nonce(account.nonce);
100                account_data.set_balance(Address::ZERO, account.balance);
101                info!(
102                    "Initialized account {} with nonce {} and balance {} ETH",
103                    account_address,
104                    account.nonce,
105                    format_ether(account.balance)
106                );
107            }
108        }
109
110        Ok(())
111    }
112
113    async fn handle_canon_state_notification(&self, notification: CanonStateNotification<EthPrimitives>) -> Result<()> {
114        let new = match notification {
115            CanonStateNotification::Reorg { old: _, new } => {
116                debug!("Processing reorg with new chain");
117                new
118            }
119            CanonStateNotification::Commit { new } => {
120                debug!("Processing committed chain");
121                new
122            }
123        };
124
125        for block in new.blocks().values() {
126            update_nonces_from_block(&self.account_state, block).await?;
127        }
128
129        Ok(())
130    }
131}
132
133async fn update_nonces_from_block(account_state: &Arc<RwLock<AccountNonceAndBalanceState>>, block: &RecoveredBlock<Block>) -> Result<()> {
134    let mut state = account_state.write().await;
135
136    // Process transactions to update nonces
137    for tx in &block.body().transactions {
138        let from = tx.recover_signer()?;
139
140        if state.is_monitored(&from)
141            && let Some(account_data) = state.get_mut_account(&from)
142        {
143            let current_nonce = account_data.get_nonce();
144            let tx_nonce = tx.nonce();
145
146            // Update nonce if this transaction has a higher nonce
147            if tx_nonce >= current_nonce {
148                account_data.set_nonce(tx_nonce + 1);
149                debug!("Updated nonce for {} to {} (from tx nonce {})", from, tx_nonce + 1, tx_nonce);
150            }
151        }
152    }
153
154    Ok(())
155}
156
157impl<R, N, P> Component for AccountMonitorComponent<R, N, P>
158where
159    R: CanonStateSubscriptions<Primitives = EthPrimitives> + StateProviderFactory + Send + Sync + Clone + 'static,
160    N: Network + 'static,
161    P: reth_node_types::NodePrimitives + 'static,
162{
163    fn spawn(self, executor: TaskExecutor) -> Result<()> {
164        executor.spawn_critical(self.name(), async move {
165            if let Err(e) = self.run().await {
166                error!("Account monitor component failed: {}", e);
167            }
168        });
169        Ok(())
170    }
171
172    fn name(&self) -> &'static str {
173        "AccountMonitorComponent"
174    }
175}
176
177/// Builder for AccountMonitorComponent
178pub struct AccountMonitorComponentBuilder {}
179
180impl AccountMonitorComponentBuilder {
181    pub fn new() -> Self {
182        Self {}
183    }
184
185    pub fn build<R, N, P>(
186        self,
187        reth_provider: R,
188        account_state: Arc<RwLock<AccountNonceAndBalanceState>>,
189        signers: Arc<RwLock<TxSigners<P>>>,
190    ) -> AccountMonitorComponent<R, N, P>
191    where
192        R: CanonStateSubscriptions<Primitives = EthPrimitives> + StateProviderFactory + Send + Sync + Clone + 'static,
193        N: Network + 'static,
194        P: NodePrimitives + 'static,
195    {
196        AccountMonitorComponent::new(reth_provider, account_state, signers)
197    }
198}
199
200impl Default for AccountMonitorComponentBuilder {
201    fn default() -> Self {
202        Self::new()
203    }
204}
kabu_broadcast_accounts/accounts_monitor/account_monitor_component.rs

kabu_broadcast_accounts/accounts_monitor/
account_monitor_component.rs
