kabu_rpc_handler/handler/
ws.rs

1use axum::extract::{ConnectInfo, State};
2use axum::{
3    extract::ws::{Message, WebSocket, WebSocketUpgrade},
4    response::IntoResponse,
5};
6
7use crate::dto::block::{BlockHeader, WebSocketMessage};
8use kabu_rpc_state::AppState;
9use kabu_types_events::MarketEvents;
10use revm::{DatabaseCommit, DatabaseRef};
11use std::net::SocketAddr;
12use tracing::{error, warn};
13
14/// Handle websocket upgrade
15pub async fn ws_handler<DB: DatabaseRef<Error = kabu_evm_db::KabuDBError> + DatabaseCommit + Send + Sync + Clone + 'static>(
16    ws: WebSocketUpgrade,
17    ConnectInfo(addr): ConnectInfo<SocketAddr>,
18    State(app_state): State<AppState<DB>>,
19) -> impl IntoResponse {
20    ws.on_failed_upgrade(move |e| {
21        warn!("ws upgrade error: {} with {}", e, addr);
22    })
23    .on_upgrade(move |socket| on_upgrade(socket, addr, app_state))
24}
25
26/// Actual websocket statemachine (one will be spawned per connection)
27async fn on_upgrade<DB: DatabaseRef + DatabaseCommit + Send + Sync + Clone + 'static>(
28    mut socket: WebSocket,
29    _who: SocketAddr,
30    app_state: AppState<DB>,
31) {
32    let mut receiver = app_state.bc.market_events_channel().subscribe();
33
34    while let Ok(event) = receiver.recv().await {
35        // Only process BlockHeaderUpdate events for WebSocket streaming
36        let (block_number, timestamp, base_fee, next_base_fee) = match event {
37            MarketEvents::BlockHeaderUpdate { block_number, timestamp, base_fee, next_base_fee, .. } => {
38                (block_number, timestamp, base_fee, next_base_fee)
39            }
40            _ => continue, // Skip non-block events
41        };
42
43        let ws_msg = WebSocketMessage::BlockHeader(BlockHeader {
44            number: block_number,
45            timestamp,
46            base_fee_per_gas: Some(base_fee),
47            next_block_base_fee: next_base_fee,
48        });
49        match serde_json::to_string(&ws_msg) {
50            Ok(json) => {
51                let _ = socket.send(Message::Text(json.into())).await;
52            }
53            Err(e) => {
54                error!("Failed to serialize block header: {}", e);
55            }
56        }
57    }
58}
kabu_rpc_handler/handler/ws.rs

kabu_rpc_handler/handler/
ws.rs
