1use kabu_core_components::Component;
2use reth_tasks::TaskExecutor;
3use std::collections::BTreeMap;
4use std::collections::btree_map::Entry;
5use std::marker::PhantomData;
6use std::sync::Arc;
7use tokio::sync::RwLock;
8
9use alloy_eips::{BlockId, BlockNumberOrTag};
10use alloy_network::Network;
11use alloy_primitives::{Address, Bytes, U256};
12use alloy_provider::Provider;
13use alloy_rpc_types_trace::geth::AccountState;
14use eyre::{Result, eyre};
15
16use kabu_core_blockchain::{Blockchain, BlockchainState};
17use kabu_defi_address_book::TokenAddressEth;
18use kabu_evm_utils::{BalanceCheater, NWETH};
19use kabu_types_blockchain::GethStateUpdate;
20use kabu_types_entities::{AccountNonceAndBalanceState, TxSigners};
21use kabu_types_market::MarketState;
22use reth_node_types::NodePrimitives;
23use revm::{Database, DatabaseCommit, DatabaseRef};
24use tracing::{debug, error, trace};
25
26async fn fetch_account_state<P, N>(client: P, address: Address) -> Result<AccountState>
27where
28 N: Network,
29 P: Provider<N> + Send + Sync + Clone + 'static,
30{
31 let code = client.get_code_at(address).block_id(BlockId::Number(BlockNumberOrTag::Latest)).await.ok();
32 let balance = client.get_balance(address).block_id(BlockId::Number(BlockNumberOrTag::Latest)).await.ok();
33 let nonce = client.get_transaction_count(address).block_id(BlockId::Number(BlockNumberOrTag::Latest)).await.ok();
34
35 Ok(AccountState { balance, code, nonce, storage: BTreeMap::new() })
36}
37
38async fn set_monitor_token_balance(
39 account_nonce_balance_state: Option<Arc<RwLock<AccountNonceAndBalanceState>>>,
40 owner: Address,
41 token: Address,
42 balance: U256,
43) {
44 if let Some(account_nonce_balance) = account_nonce_balance_state {
45 debug!("set_monitor_balance {} {} {}", owner, token, balance);
46 let mut account_nonce_balance_guard = account_nonce_balance.write().await;
47 let entry = account_nonce_balance_guard.get_entry_or_default(owner);
48 debug!("set_monitor_balance {:?}", entry);
49
50 entry.add_balance(token, balance);
51 }
52}
53
54async fn set_monitor_nonce(account_nonce_balance_state: Option<Arc<RwLock<AccountNonceAndBalanceState>>>, owner: Address, nonce: u64) {
55 if let Some(account_nonce_balance) = account_nonce_balance_state {
56 debug!("set_monitor_nonce {} {}", owner, nonce);
57 let mut account_nonce_balance_guard = account_nonce_balance.write().await;
58 let entry = account_nonce_balance_guard.get_entry_or_default(owner);
59 debug!("set_monitor_nonce {:?}", entry);
60 entry.set_nonce(nonce);
61 }
62}
63
64pub async fn preload_market_state<P, N, DB>(
65 client: P,
66 copied_accounts_vec: Vec<Address>,
67 new_accounts_vec: Vec<(Address, u64, U256, Option<Bytes>)>,
68 token_balances_vec: Vec<(Address, Address, U256)>,
69 market_state: Arc<RwLock<MarketState<DB>>>,
70 account_nonce_balance_state: Option<Arc<RwLock<AccountNonceAndBalanceState>>>,
71) -> Result<()>
72where
73 N: Network,
74 P: Provider<N> + Send + Sync + Clone + 'static,
75 DB: DatabaseRef + Database + DatabaseCommit + Send + Sync + Clone + 'static,
76{
77 let mut market_state_guard = market_state.write().await;
78
79 let mut state: GethStateUpdate = BTreeMap::new();
80
81 for address in copied_accounts_vec {
82 trace!("Loading address : {address}");
83 let acc_state = fetch_account_state(client.clone(), address).await?;
84
85 set_monitor_token_balance(account_nonce_balance_state.clone(), address, Address::ZERO, acc_state.balance.unwrap_or_default()).await;
86
87 set_monitor_nonce(account_nonce_balance_state.clone(), address, acc_state.nonce.unwrap_or_default()).await;
88 trace!("Loaded address : {address} {:?}", acc_state);
89
90 state.insert(address, acc_state);
91 }
92
93 for (address, nonce, balance, code) in new_accounts_vec {
94 trace!("new_accounts added {} {} {}", address, nonce, balance);
95 set_monitor_token_balance(account_nonce_balance_state.clone(), address, NWETH::NATIVE_ADDRESS, balance).await;
96 state.insert(address, AccountState { balance: Some(balance), code, nonce: Some(nonce), storage: BTreeMap::new() });
97 }
98
99 for (token, owner, balance) in token_balances_vec {
100 if token == TokenAddressEth::ETH_NATIVE {
101 match state.entry(owner) {
102 Entry::Vacant(e) => {
103 e.insert(AccountState { balance: Some(balance), nonce: Some(0), code: None, storage: BTreeMap::new() });
104 }
105 Entry::Occupied(mut e) => {
106 e.get_mut().balance = Some(balance);
107 }
108 }
109 } else {
110 match state.entry(token) {
111 Entry::Vacant(e) => {
112 let mut acc_state = fetch_account_state(client.clone(), token).await?;
113 acc_state.storage.insert(BalanceCheater::get_balance_cell(token, owner)?.into(), balance.into());
114 e.insert(acc_state);
115 }
116 Entry::Occupied(mut e) => {
117 e.get_mut().storage.insert(BalanceCheater::get_balance_cell(token, owner)?.into(), balance.into());
118 }
119 }
120 }
121
122 set_monitor_token_balance(account_nonce_balance_state.clone(), owner, token, balance).await;
123 }
124 market_state_guard.apply_geth_update(state);
125
126 Ok(())
127}
128
129#[allow(dead_code)]
130#[derive(Clone)]
131pub struct MarketStatePreloadedOneShotComponent<P, N, DB> {
132 name: &'static str,
133 client: P,
134 copied_accounts: Vec<Address>,
135 new_accounts: Vec<(Address, u64, U256, Option<Bytes>)>,
136 token_balances: Vec<(Address, Address, U256)>,
137
138 market_state: Option<Arc<RwLock<MarketState<DB>>>>,
139
140 account_nonce_balance_state: Option<Arc<RwLock<AccountNonceAndBalanceState>>>,
141 _n: PhantomData<N>,
142}
143
144#[allow(dead_code)]
145impl<P, N, DB> MarketStatePreloadedOneShotComponent<P, N, DB>
146where
147 N: Network,
148 P: Provider<N> + Send + Sync + Clone + 'static,
149 DB: DatabaseRef + DatabaseCommit + Send + Sync + Clone + 'static,
150{
151 fn name(&self) -> &'static str {
152 self.name
153 }
154
155 pub fn new(client: P) -> Self {
156 Self {
157 name: "MarketStatePreloadedOneShotComponent",
158 client,
159 copied_accounts: Vec::new(),
160 new_accounts: Vec::new(),
161 token_balances: Vec::new(),
162 market_state: None,
163 account_nonce_balance_state: None,
164 _n: PhantomData,
165 }
166 }
167
168 pub fn with_name(self, name: &'static str) -> Self {
169 Self { name, ..self }
170 }
171
172 pub fn on_bc<NP: NodePrimitives>(self, bc: &Blockchain<NP>, state: &BlockchainState<DB, NP>) -> Self {
173 Self { account_nonce_balance_state: Some(bc.nonce_and_balance()), market_state: Some(state.market_state_commit()), ..self }
174 }
175
176 pub fn with_signers(self, tx_signers: Arc<RwLock<TxSigners>>) -> Self {
177 match tx_signers.try_read() {
178 Ok(signers) => {
179 let mut addresses = self.copied_accounts;
180 addresses.extend(signers.get_address_vec());
181 Self { copied_accounts: addresses, ..self }
182 }
183 Err(e) => {
184 error!("tx_signers.try_read() {}", e);
185 self
186 }
187 }
188 }
189
190 pub fn with_copied_account(self, address: Address) -> Self {
191 let mut copied_accounts = self.copied_accounts;
192 copied_accounts.push(address);
193 Self { copied_accounts, ..self }
194 }
195
196 pub fn with_copied_accounts<T: Into<Address>>(self, address_vec: Vec<T>) -> Self {
197 let mut copied_accounts = self.copied_accounts;
198 let more: Vec<Address> = address_vec.into_iter().map(Into::<Address>::into).collect();
199 copied_accounts.extend(more);
200 Self { copied_accounts, ..self }
201 }
202
203 pub fn with_new_account(self, address: Address, nonce: u64, balance: U256, code: Option<Bytes>) -> Self {
204 let mut new_accounts = self.new_accounts;
205 new_accounts.push((address, nonce, balance, code));
206 Self { new_accounts, ..self }
207 }
208
209 pub fn with_market_state(self, market_state: Arc<RwLock<MarketState<DB>>>) -> Self {
210 Self { market_state: Some(market_state), ..self }
211 }
212
213 pub fn with_token_balance(self, token: Address, owner: Address, balance: U256) -> Self {
214 let mut token_balances = self.token_balances;
215 token_balances.push((token, owner, balance));
216 Self { token_balances, ..self }
217 }
218}
219
220impl<P, N, DB> Component for MarketStatePreloadedOneShotComponent<P, N, DB>
221where
222 N: Network,
223 P: Provider<N> + Send + Sync + Clone + 'static,
224 DB: DatabaseRef + Database + DatabaseCommit + Send + Sync + Clone + 'static,
225{
226 fn spawn(self, executor: TaskExecutor) -> Result<()> {
227 let market_state = self.market_state.clone().ok_or_else(|| eyre!("market_state not set"))?;
228 executor.spawn(async move {
229 if let Err(e) = preload_market_state(
230 self.client.clone(),
231 self.copied_accounts.clone(),
232 self.new_accounts.clone(),
233 self.token_balances.clone(),
234 market_state,
235 self.account_nonce_balance_state.clone(),
236 )
237 .await
238 {
239 tracing::error!("Preload market state failed: {}", e);
240 }
241 });
242 Ok(())
243 }
244 fn name(&self) -> &'static str {
245 "MarketStatePreloadedOneShotComponent"
246 }
247}