kabu_defi_health_monitor/
state_health_monitor.rs

1use kabu_core_components::Component;
2use reth_tasks::TaskExecutor;
3use std::collections::HashMap;
4use std::sync::Arc;
5use tokio::sync::{RwLock, broadcast};
6
7use alloy_eips::BlockNumberOrTag;
8use alloy_network::Ethereum;
9use alloy_provider::Provider;
10use chrono::{DateTime, Duration, Local};
11use eyre::Result;
12use tokio::sync::broadcast::Receiver;
13use tokio::sync::broadcast::error::RecvError;
14use tracing::{error, info, warn};
15
16use kabu_core_blockchain::{Blockchain, BlockchainState};
17use kabu_evm_db::DatabaseKabuExt;
18use kabu_types_events::{MarketEvents, MessageTxCompose, TxComposeMessageType};
19use kabu_types_market::{MarketState, PoolId};
20use reth_node_types::NodePrimitives;
21use revm::DatabaseRef;
22
23async fn verify_pool_state_task<P: Provider<Ethereum> + 'static, DB: DatabaseKabuExt>(
24    client: P,
25    address: PoolId,
26    market_state: Arc<RwLock<MarketState<DB>>>,
27) -> Result<()> {
28    info!("Verifying state {address:?}");
29    let address = match address {
30        PoolId::Address(addr) => addr,
31        PoolId::B256(_) => return Err(eyre::eyre!("B256 pool ID not supported for verification")),
32    };
33    let account = market_state.write().await.state_db.load_account(address).cloned()?;
34    let read_only_cell_hash_set = market_state.read().await.config.read_only_cells.get(&address).cloned().unwrap_or_default();
35
36    for (cell, current_value) in account.storage.iter() {
37        if read_only_cell_hash_set.contains(cell) {
38            continue;
39        }
40        match client.get_storage_at(address, *cell).block_id(BlockNumberOrTag::Latest.into()).await {
41            Ok(actual_value) => {
42                if actual_value.is_zero() {
43                    continue;
44                }
45                if *current_value != actual_value {
46                    warn!(
47                        "verify : account storage is different : {address:?} {cell:?} {current_value:#32x} -> {actual_value:#32x} storage size : {}",
48                        account.storage.len()
49                    );
50                    if let Err(e) = market_state.write().await.state_db.insert_account_storage(address, *cell, actual_value) {
51                        error!("{e}");
52                    }
53                }
54            }
55            Err(e) => {
56                error!("Cannot read storage {:?} {:?} : {}", account, cell, e)
57            }
58        }
59    }
60
61    Ok(())
62}
63
64pub async fn state_health_monitor_worker<
65    P: Provider<Ethereum> + Clone + 'static,
66    DB: DatabaseRef + DatabaseKabuExt + Send + Sync + Clone + 'static,
67>(
68    client: P,
69    market_state: Arc<RwLock<MarketState<DB>>>,
70    tx_compose_channel_rx: broadcast::Sender<MessageTxCompose>,
71    market_events_rx: broadcast::Sender<MarketEvents>,
72) -> Result<()> {
73    let mut tx_compose_channel_rx: Receiver<MessageTxCompose> = tx_compose_channel_rx.subscribe();
74    let mut market_events_rx: Receiver<MarketEvents> = market_events_rx.subscribe();
75
76    let mut check_time_map: HashMap<PoolId, DateTime<Local>> = HashMap::new();
77    let mut pool_address_to_verify_vec: Vec<PoolId> = Vec::new();
78
79    loop {
80        tokio::select! {
81            msg = market_events_rx.recv() => {
82                let market_event_msg : Result<MarketEvents, RecvError> = msg;
83                match market_event_msg {
84                    Ok(market_event)=>{
85                        if matches!(market_event, MarketEvents::BlockStateUpdate{..}) {
86                            for pool_address in pool_address_to_verify_vec {
87                                tokio::task::spawn(
88                                    verify_pool_state_task(
89                                        client.clone(),
90                                        pool_address,
91                                        market_state.clone()
92                                    )
93                                );
94                            }
95                            pool_address_to_verify_vec = Vec::new();
96                        }
97                    }
98                    Err(e)=>{error!("market_event_rx error : {e}")}
99                }
100            },
101
102            msg = tx_compose_channel_rx.recv() => {
103                let tx_compose_update : Result<MessageTxCompose, RecvError>  = msg;
104                match tx_compose_update {
105                    Ok(tx_compose_msg)=>{
106                        if let TxComposeMessageType::Sign(sign_request_data)= tx_compose_msg.inner
107                            && let Some(swap) = sign_request_data.swap {
108                                let pool_address_vec =  swap.get_pool_address_vec();
109                                let now = chrono::Local::now();
110                                for pool_address in pool_address_vec {
111                                    if now - check_time_map.get(&pool_address).cloned().unwrap_or(DateTime::<Local>::MIN_UTC.into()) > Duration::seconds(60) {
112                                        check_time_map.insert(pool_address, Local::now());
113                                        if !pool_address_to_verify_vec.contains(&pool_address){
114                                            pool_address_to_verify_vec.push(pool_address)
115                                        }
116                                    }
117                                }
118                        }
119
120                    }
121                    Err(e)=>{
122                        error!("tx_compose_channel_rx : {e}")
123                    }
124                }
125
126            }
127
128        }
129    }
130}
131
132pub struct StateHealthMonitorActor<P, DB: Clone + Send + Sync + 'static> {
133    client: P,
134
135    market_state: Option<Arc<RwLock<MarketState<DB>>>>,
136
137    tx_compose_channel_rx: Option<broadcast::Sender<MessageTxCompose>>,
138
139    market_events_rx: Option<broadcast::Sender<MarketEvents>>,
140}
141
142impl<P, DB> StateHealthMonitorActor<P, DB>
143where
144    P: Provider<Ethereum> + Send + Sync + Clone + 'static,
145    DB: DatabaseRef + DatabaseKabuExt + Send + Sync + Clone + Default + 'static,
146{
147    pub fn new(client: P) -> Self {
148        StateHealthMonitorActor { client, market_state: None, tx_compose_channel_rx: None, market_events_rx: None }
149    }
150
151    pub fn on_bc<NP: NodePrimitives>(self, bc: &Blockchain, state: &BlockchainState<DB, NP>) -> Self {
152        Self {
153            market_state: Some(state.market_state()),
154            tx_compose_channel_rx: Some(bc.tx_compose_channel()),
155            market_events_rx: Some(bc.market_events_channel()),
156            ..self
157        }
158    }
159}
160
161impl<P, DB> Component for StateHealthMonitorActor<P, DB>
162where
163    P: Provider<Ethereum> + Send + Sync + Clone + 'static,
164    DB: DatabaseRef + DatabaseKabuExt + Send + Sync + Clone + 'static,
165{
166    fn spawn(self, executor: TaskExecutor) -> Result<()> {
167        let market_state = self.market_state.clone().ok_or_else(|| eyre::eyre!("market_state not set"))?;
168        let tx_compose_channel_rx = self.tx_compose_channel_rx.clone().ok_or_else(|| eyre::eyre!("tx_compose_channel_rx not set"))?;
169        let market_events_rx = self.market_events_rx.clone().ok_or_else(|| eyre::eyre!("market_events_rx not set"))?;
170
171        executor.spawn(async move {
172            if let Err(e) = state_health_monitor_worker(self.client.clone(), market_state, tx_compose_channel_rx, market_events_rx).await {
173                tracing::error!("State health monitor worker failed: {}", e);
174            }
175        });
176        Ok(())
177    }
178    fn name(&self) -> &'static str {
179        "StateHealthMonitorActor"
180    }
181}
kabu_defi_health_monitor/state_health_monitor.rs

kabu_defi_health_monitor/
state_health_monitor.rs
