Skip to main content

reth_node_ethstats/
ethstats.rs

1use crate::{
2    connection::ConnWrapper,
3    credentials::EthstatsCredentials,
4    error::EthStatsError,
5    events::{
6        AuthMsg, BlockMsg, BlockStats, HistoryMsg, LatencyMsg, NodeInfo, NodeStats, PayloadMsg,
7        PayloadStats, PendingMsg, PendingStats, PingMsg, StatsMsg, TxStats, UncleStats,
8    },
9};
10use alloy_consensus::{BlockHeader, Sealable};
11use alloy_primitives::U256;
12use reth_chain_state::{CanonStateNotification, CanonStateSubscriptions};
13use reth_network_api::{NetworkInfo, Peers};
14use reth_primitives_traits::{Block, BlockBody};
15use reth_storage_api::{BlockReader, BlockReaderIdExt};
16use reth_transaction_pool::TransactionPool;
17
18use chrono::Local;
19use futures_util::FutureExt;
20use serde_json::Value;
21use std::{
22    str::FromStr,
23    sync::Arc,
24    time::{Duration, Instant},
25};
26use tokio::{
27    sync::{mpsc, Mutex, RwLock},
28    time::{interval, sleep, timeout},
29};
30use tokio_stream::StreamExt;
31use tokio_tungstenite::connect_async;
32use tracing::{debug, info};
33use url::Url;
34
35/// Number of historical blocks to include in a history update sent to the `EthStats` server
36const HISTORY_UPDATE_RANGE: u64 = 50;
37/// Duration to wait before attempting to reconnect to the `EthStats` server
38const RECONNECT_INTERVAL: Duration = Duration::from_secs(5);
39/// Maximum time to wait for a ping response from the server
40const PING_TIMEOUT: Duration = Duration::from_secs(5);
41/// Interval between regular stats reports to the server
42const REPORT_INTERVAL: Duration = Duration::from_secs(15);
43/// Maximum time to wait for initial connection establishment
44const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
45/// Maximum time to wait for reading messages from the server
46const READ_TIMEOUT: Duration = Duration::from_secs(30);
47
48/// Main service for interacting with an `EthStats` server
49///
50/// This service handles all communication with the `EthStats` server including
51/// authentication, stats reporting, block notifications, and connection management.
52/// It maintains a persistent `WebSocket` connection and automatically reconnects
53/// when the connection is lost.
54#[derive(Debug, Clone)]
55pub struct EthStatsService<Network, Provider, Pool> {
56    /// Authentication credentials for the `EthStats` server
57    credentials: EthstatsCredentials,
58    /// `WebSocket` connection wrapper, wrapped in `Arc<RwLock>` for shared access
59    conn: Arc<RwLock<Option<ConnWrapper>>>,
60    /// Timestamp of the last ping sent to the server
61    last_ping: Arc<Mutex<Option<Instant>>>,
62    /// Network interface for getting peer and sync information
63    network: Network,
64    /// Blockchain provider for reading block data and state
65    provider: Provider,
66    /// Transaction pool for getting pending transaction statistics
67    pool: Pool,
68}
69
70impl<Network, Provider, Pool> EthStatsService<Network, Provider, Pool>
71where
72    Network: NetworkInfo + Peers,
73    Provider: BlockReaderIdExt + CanonStateSubscriptions,
74    Pool: TransactionPool,
75{
76    /// Create a new `EthStats` service and establish initial connection
77    ///
78    /// # Arguments
79    /// * `url` - Connection string in format "`node_id:secret@host`"
80    /// * `network` - Network interface implementation
81    /// * `provider` - Blockchain provider implementation
82    /// * `pool` - Transaction pool implementation
83    pub async fn new(
84        url: &str,
85        network: Network,
86        provider: Provider,
87        pool: Pool,
88    ) -> Result<Self, EthStatsError> {
89        let credentials = EthstatsCredentials::from_str(url)?;
90        let service = Self {
91            credentials,
92            conn: Arc::new(RwLock::new(None)),
93            last_ping: Arc::new(Mutex::new(None)),
94            network,
95            provider,
96            pool,
97        };
98        service.connect().await?;
99
100        Ok(service)
101    }
102
103    /// Establish `WebSocket` connection to the `EthStats` server
104    ///
105    /// Attempts to connect to the server using the credentials and handles
106    /// connection timeouts and errors. Uses either `ws://` or `wss://` based
107    /// on the credentials configuration.
108    async fn connect(&self) -> Result<(), EthStatsError> {
109        debug!(
110            target: "ethstats",
111            "Attempting to connect to EthStats server at {}", self.credentials.host
112        );
113        let protocol = if self.credentials.use_tls { "wss" } else { "ws" };
114        let full_url = format!("{}://{}/api", protocol, self.credentials.host);
115        let url = Url::parse(&full_url).map_err(EthStatsError::Url)?;
116
117        match timeout(CONNECT_TIMEOUT, connect_async(url.as_str())).await {
118            Ok(Ok((ws_stream, _))) => {
119                debug!(
120                    target: "ethstats",
121                    "Successfully connected to EthStats server at {}", self.credentials.host
122                );
123                let conn: ConnWrapper = ConnWrapper::new(ws_stream);
124                *self.conn.write().await = Some(conn);
125                self.login().await?;
126                Ok(())
127            }
128            Ok(Err(e)) => Err(EthStatsError::WebSocket(e)),
129            Err(_) => {
130                debug!(target: "ethstats", "Connection to EthStats server timed out");
131                Err(EthStatsError::Timeout)
132            }
133        }
134    }
135
136    /// Authenticate with the `EthStats` server
137    ///
138    /// Sends authentication credentials and node information to the server
139    /// and waits for a successful acknowledgment.
140    async fn login(&self) -> Result<(), EthStatsError> {
141        debug!(
142            target: "ethstats",
143            "Attempting to login to EthStats server as node_id {}", self.credentials.node_id
144        );
145        let conn = self.conn.read().await;
146        let conn = conn.as_ref().ok_or(EthStatsError::NotConnected)?;
147
148        let network_status = self
149            .network
150            .network_status()
151            .await
152            .map_err(|e| EthStatsError::AuthError(e.to_string()))?;
153        let id = &self.credentials.node_id;
154        let secret = &self.credentials.secret;
155        let protocol = network_status
156            .capabilities
157            .iter()
158            .map(|cap| format!("{}/{}", cap.name, cap.version))
159            .collect::<Vec<_>>()
160            .join(", ");
161        let port = self.network.local_addr().port() as u64;
162
163        let auth = AuthMsg {
164            id: id.clone(),
165            secret: secret.clone(),
166            info: NodeInfo {
167                name: id.clone(),
168                node: network_status.client_version.clone(),
169                port,
170                network: self.network.chain_id().to_string(),
171                protocol,
172                api: "No".to_string(),
173                os: std::env::consts::OS.into(),
174                os_ver: std::env::consts::ARCH.into(),
175                client: "0.1.1".to_string(),
176                history: true,
177            },
178        };
179
180        let message = auth.generate_login_message();
181        conn.write_json(&message).await?;
182
183        let response =
184            timeout(READ_TIMEOUT, conn.read_json()).await.map_err(|_| EthStatsError::Timeout)??;
185
186        if let Some(ack) = response.get("emit") &&
187            ack.get(0) == Some(&Value::String("ready".to_string()))
188        {
189            info!(
190                target: "ethstats",
191                "Login successful to EthStats server as node_id {}", self.credentials.node_id
192            );
193            return Ok(());
194        }
195
196        debug!(target: "ethstats", "Login failed: Unauthorized or unexpected login response");
197        Err(EthStatsError::AuthError("Unauthorized or unexpected login response".into()))
198    }
199
200    /// Report current node statistics to the `EthStats` server
201    ///
202    /// Sends information about the node's current state including sync status,
203    /// peer count, and uptime.
204    async fn report_stats(&self) -> Result<(), EthStatsError> {
205        let conn = self.conn.read().await;
206        let conn = conn.as_ref().ok_or(EthStatsError::NotConnected)?;
207
208        let stats_msg = StatsMsg {
209            id: self.credentials.node_id.clone(),
210            stats: NodeStats {
211                active: true,
212                syncing: self.network.is_syncing(),
213                peers: self.network.num_connected_peers() as u64,
214                gas_price: self.pool.block_info().pending_basefee,
215                uptime: 100,
216            },
217        };
218
219        let message = stats_msg.generate_stats_message();
220        conn.write_json(&message).await?;
221
222        Ok(())
223    }
224
225    /// Send a ping message to the `EthStats` server
226    ///
227    /// Records the ping time and starts a timeout task to detect if the server
228    /// doesn't respond within the expected timeframe.
229    async fn send_ping(&self) -> Result<(), EthStatsError> {
230        let conn = self.conn.read().await;
231        let conn = conn.as_ref().ok_or(EthStatsError::NotConnected)?;
232
233        let ping_time = Instant::now();
234        *self.last_ping.lock().await = Some(ping_time);
235
236        let client_time = Local::now().format("%Y-%m-%d %H:%M:%S%.f %:z %Z").to_string();
237        let ping_msg = PingMsg { id: self.credentials.node_id.clone(), client_time };
238
239        let message = ping_msg.generate_ping_message();
240        conn.write_json(&message).await?;
241
242        // Start ping timeout
243        let active_ping = self.last_ping.clone();
244        let conn_ref = self.conn.clone();
245        tokio::spawn(
246            async move {
247                sleep(PING_TIMEOUT).await;
248                let timed_out = {
249                    let mut active = active_ping.lock().await;
250                    let timed_out = active.is_some();
251                    if timed_out {
252                        *active = None;
253                    }
254                    timed_out
255                };
256
257                if timed_out {
258                    debug!(target: "ethstats", "Ping timeout");
259                    // Clear connection to trigger reconnect
260                    if let Some(conn) = conn_ref.write().await.take() {
261                        let _ = conn.close().await;
262                    }
263                }
264            }
265            .boxed(),
266        );
267
268        Ok(())
269    }
270
271    /// Report latency measurement to the `EthStats` server
272    ///
273    /// Calculates the round-trip time from the last ping and sends it to
274    /// the server. This is called when a pong response is received.
275    async fn report_latency(&self) -> Result<(), EthStatsError> {
276        let start = {
277            let mut active = self.last_ping.lock().await;
278            active.take()
279        };
280
281        if let Some(start) = start {
282            let latency = start.elapsed().as_millis() as u64 / 2;
283
284            debug!(target: "ethstats", "Reporting latency: {}ms", latency);
285
286            let latency_msg = LatencyMsg { id: self.credentials.node_id.clone(), latency };
287
288            let message = latency_msg.generate_latency_message();
289            let conn = self.conn.read().await;
290            let conn = conn.as_ref().ok_or(EthStatsError::NotConnected)?;
291            conn.write_json(&message).await?;
292        }
293
294        Ok(())
295    }
296
297    /// Report pending transaction count to the `EthStats` server
298    ///
299    /// Gets the current number of pending transactions from the pool and
300    /// sends this information to the server.
301    async fn report_pending(&self) -> Result<(), EthStatsError> {
302        let conn = self.conn.read().await;
303        let conn = conn.as_ref().ok_or(EthStatsError::NotConnected)?;
304        let pending = self.pool.pool_size().pending as u64;
305
306        debug!(target: "ethstats", "Reporting pending txs: {}", pending);
307
308        let pending_msg =
309            PendingMsg { id: self.credentials.node_id.clone(), stats: PendingStats { pending } };
310
311        let message = pending_msg.generate_pending_message();
312        conn.write_json(&message).await?;
313
314        Ok(())
315    }
316
317    /// Report block information to the `EthStats` server
318    ///
319    /// Fetches block data either from a canonical state notification or
320    /// the current best block, converts it to stats format, and sends
321    /// it to the server.
322    ///
323    /// # Arguments
324    /// * `head` - Optional canonical state notification containing new block info
325    async fn report_block(
326        &self,
327        head: Option<CanonStateNotification<<Provider as CanonStateSubscriptions>::Primitives>>,
328    ) -> Result<(), EthStatsError> {
329        let conn = self.conn.read().await;
330        let conn = conn.as_ref().ok_or(EthStatsError::NotConnected)?;
331
332        let block_number = if let Some(head) = head {
333            head.tip().header().number()
334        } else {
335            self.provider
336                .best_block_number()
337                .map_err(|e| EthStatsError::DataFetchError(e.to_string()))?
338        };
339
340        match self.provider.block_by_id(block_number.into()) {
341            Ok(Some(block)) => {
342                let block_msg = BlockMsg {
343                    id: self.credentials.node_id.clone(),
344                    block: self.block_to_stats(&block)?,
345                };
346
347                debug!(target: "ethstats", "Reporting block: {}", block_number);
348
349                let message = block_msg.generate_block_message();
350                conn.write_json(&message).await?;
351            }
352            Ok(None) => {
353                // Block not found, stop fetching
354                debug!(target: "ethstats", "Block {} not found", block_number);
355                return Err(EthStatsError::BlockNotFound(block_number));
356            }
357            Err(e) => {
358                debug!(target: "ethstats", "Error fetching block {}: {}", block_number, e);
359                return Err(EthStatsError::DataFetchError(e.to_string()));
360            }
361        };
362
363        Ok(())
364    }
365
366    /// Report new payload information to the `EthStats` server
367    ///
368    /// Sends information about payload processing time and block details
369    /// to the server for monitoring purposes.
370    pub async fn report_new_payload(
371        &self,
372        block_hash: alloy_primitives::B256,
373        block_number: u64,
374        processing_time: Duration,
375    ) -> Result<(), EthStatsError> {
376        let conn = self.conn.read().await;
377        let conn = conn.as_ref().ok_or(EthStatsError::NotConnected)?;
378
379        let payload_stats = PayloadStats {
380            number: U256::from(block_number),
381            hash: block_hash,
382            processing_time: processing_time.as_millis() as u64,
383        };
384
385        let payload_msg =
386            PayloadMsg { id: self.credentials.node_id.clone(), payload: payload_stats };
387
388        debug!(
389            target: "ethstats",
390            "Reporting new payload: block={}, hash={:?}, processing_time={}ms",
391            block_number,
392            block_hash,
393            processing_time.as_millis()
394        );
395
396        let message = payload_msg.generate_new_payload_message();
397        conn.write_json(&message).await?;
398
399        Ok(())
400    }
401
402    /// Convert a block to `EthStats` block statistics format
403    ///
404    /// Extracts relevant information from a block and formats it according
405    /// to the `EthStats` protocol specification.
406    ///
407    /// # Arguments
408    /// * `block` - The block to convert
409    fn block_to_stats(
410        &self,
411        block: &<Provider as BlockReader>::Block,
412    ) -> Result<BlockStats, EthStatsError> {
413        let body = block.body();
414        let header = block.header();
415
416        let txs = body.transaction_hashes_iter().copied().map(|hash| TxStats { hash }).collect();
417
418        Ok(BlockStats {
419            number: U256::from(header.number()),
420            hash: header.hash_slow(),
421            parent_hash: header.parent_hash(),
422            timestamp: U256::from(header.timestamp()),
423            miner: header.beneficiary(),
424            gas_used: header.gas_used(),
425            gas_limit: header.gas_limit(),
426            diff: header.difficulty().to_string(),
427            total_diff: "0".into(),
428            txs,
429            tx_root: header.transactions_root(),
430            root: header.state_root(),
431            uncles: UncleStats(vec![]),
432        })
433    }
434
435    /// Report historical block data to the `EthStats` server
436    ///
437    /// Fetches multiple blocks by their numbers and sends their statistics
438    /// to the server. This is typically called in response to a history
439    /// request from the server.
440    ///
441    /// # Arguments
442    /// * `list` - Vector of block numbers to fetch and report
443    async fn report_history(&self, list: Option<&Vec<u64>>) -> Result<(), EthStatsError> {
444        let conn = self.conn.read().await;
445        let conn = conn.as_ref().ok_or(EthStatsError::NotConnected)?;
446
447        let indexes = if let Some(list) = list {
448            list
449        } else {
450            let best_block_number = self
451                .provider
452                .best_block_number()
453                .map_err(|e| EthStatsError::DataFetchError(e.to_string()))?;
454
455            let start = best_block_number.saturating_sub(HISTORY_UPDATE_RANGE);
456
457            &(start..=best_block_number).collect()
458        };
459
460        let mut blocks = Vec::with_capacity(indexes.len());
461        for &block_number in indexes {
462            match self.provider.block_by_id(block_number.into()) {
463                Ok(Some(block)) => {
464                    blocks.push(block);
465                }
466                Ok(None) => {
467                    // Block not found, stop fetching
468                    debug!(target: "ethstats", "Block {} not found", block_number);
469                    break;
470                }
471                Err(e) => {
472                    debug!(target: "ethstats", "Error fetching block {}: {}", block_number, e);
473                    break;
474                }
475            }
476        }
477
478        let history: Vec<BlockStats> =
479            blocks.iter().map(|block| self.block_to_stats(block)).collect::<Result<_, _>>()?;
480
481        if history.is_empty() {
482            debug!(target: "ethstats", "No history to send to stats server");
483        } else {
484            debug!(
485                target: "ethstats",
486                "Sending historical blocks to ethstats, first: {}, last: {}",
487                history.first().unwrap().number,
488                history.last().unwrap().number
489            );
490        }
491
492        let history_msg = HistoryMsg { id: self.credentials.node_id.clone(), history };
493
494        let message = history_msg.generate_history_message();
495        conn.write_json(&message).await?;
496
497        Ok(())
498    }
499
500    /// Send a complete status report to the `EthStats` server
501    ///
502    /// Performs all regular reporting tasks: ping, block info, pending
503    /// transactions, and general statistics.
504    async fn report(&self) -> Result<(), EthStatsError> {
505        self.send_ping().await?;
506        self.report_block(None).await?;
507        self.report_pending().await?;
508        self.report_stats().await?;
509
510        Ok(())
511    }
512
513    /// Handle incoming messages from the `EthStats` server
514    ///
515    /// # Expected Message Variants
516    ///
517    /// This function expects messages in the following format:
518    ///
519    /// ```json
520    /// { "emit": [<command: String>, <payload: Object>] }
521    /// ```
522    ///
523    /// ## Supported Commands:
524    ///
525    /// - `"node-pong"`: Indicates a pong response to a previously sent ping. The payload is
526    ///   ignored. Triggers a latency report to the server.
527    ///   - Example: ```json { "emit": [ "node-pong", { "clientTime": "2025-07-10 12:00:00.123
528    ///     +00:00 UTC", "serverTime": "2025-07-10 12:00:01.456 +00:00 UTC" } ] } ```
529    ///
530    /// - `"history"`: Requests historical block data. The payload may contain a `list` field with
531    ///   block numbers to fetch. If `list` is not present, the default range is used.
532    ///   - Example with list: `{ "emit": ["history", {"list": [1, 2, 3], "min": 1, "max": 3}] }`
533    ///   - Example without list: `{ "emit": ["history", {}] }`
534    ///
535    /// ## Other Commands:
536    ///
537    /// Any other command is logged as unhandled and ignored.
538    async fn handle_message(&self, msg: Value) -> Result<(), EthStatsError> {
539        let emit = match msg.get("emit") {
540            Some(emit) => emit,
541            None => {
542                debug!(target: "ethstats", "Stats server sent non-broadcast, msg {}", msg);
543                return Err(EthStatsError::InvalidRequest);
544            }
545        };
546
547        let command = match emit.get(0) {
548            Some(Value::String(command)) => command.as_str(),
549            _ => {
550                debug!(target: "ethstats", "Invalid stats server message type, msg {}", msg);
551                return Err(EthStatsError::InvalidRequest);
552            }
553        };
554
555        match command {
556            "node-pong" => {
557                self.report_latency().await?;
558            }
559            "history" => {
560                let block_numbers = emit
561                    .get(1)
562                    .and_then(|v| v.as_object())
563                    .and_then(|obj| obj.get("list"))
564                    .and_then(|v| v.as_array());
565
566                if block_numbers.is_none() {
567                    self.report_history(None).await?;
568
569                    return Ok(());
570                }
571
572                let block_numbers = block_numbers
573                    .unwrap()
574                    .iter()
575                    .take(HISTORY_UPDATE_RANGE as usize + 1)
576                    .map(|val| {
577                        val.as_u64().ok_or_else(|| {
578                            debug!(
579                                target: "ethstats",
580                                "Invalid stats history block number, msg {}", msg
581                            );
582                            EthStatsError::InvalidRequest
583                        })
584                    })
585                    .collect::<Result<_, _>>()?;
586
587                self.report_history(Some(&block_numbers)).await?;
588            }
589            other => debug!(target: "ethstats", "Unhandled command: {}", other),
590        }
591
592        Ok(())
593    }
594
595    /// Main service loop that handles all `EthStats` communication
596    ///
597    /// This method runs the main event loop that:
598    /// - Maintains the `WebSocket` connection
599    /// - Handles incoming messages from the server
600    /// - Reports statistics at regular intervals
601    /// - Processes new block notifications
602    /// - Automatically reconnects when the connection is lost
603    ///
604    /// The service runs until explicitly shut down or an unrecoverable
605    /// error occurs.
606    pub async fn run(self) {
607        // Create channels for internal communication
608        let (shutdown_tx, mut shutdown_rx) = mpsc::channel(1);
609        let (message_tx, mut message_rx) = mpsc::channel(32);
610        let (head_tx, mut head_rx) = mpsc::channel(10);
611
612        // Start the read loop in a separate task
613        let read_handle = {
614            let conn_arc = self.conn.clone();
615            let message_tx = message_tx.clone();
616            let shutdown_tx = shutdown_tx.clone();
617
618            tokio::spawn(async move {
619                loop {
620                    let conn_clone = {
621                        let conn_guard = conn_arc.read().await;
622                        conn_guard.as_ref().cloned()
623                    };
624                    if let Some(conn) = conn_clone {
625                        match conn.read_json().await {
626                            Ok(msg) => {
627                                if message_tx.send(msg).await.is_err() {
628                                    break;
629                                }
630                            }
631                            Err(e) => match e {
632                                crate::error::ConnectionError::Serialization(err) => {
633                                    debug!(target: "ethstats", "JSON parse error from stats server: {}", err);
634                                }
635                                other => {
636                                    debug!(target: "ethstats", "Read error: {}", other);
637                                    if let Some(conn) = conn_arc.write().await.take() {
638                                        let _ = conn.close().await;
639                                    }
640                                }
641                            },
642                        }
643                    } else {
644                        sleep(RECONNECT_INTERVAL).await;
645                    }
646                }
647
648                let _ = shutdown_tx.send(()).await;
649            }
650            .boxed())
651        };
652
653        let canonical_stream_handle = {
654            let mut canonical_stream = self.provider.canonical_state_stream();
655            let head_tx = head_tx.clone();
656            let shutdown_tx = shutdown_tx.clone();
657
658            tokio::spawn(
659                async move {
660                    while let Some(head) = canonical_stream.next().await {
661                        if head_tx.send(head).await.is_err() {
662                            break;
663                        }
664                    }
665
666                    let _ = shutdown_tx.send(()).await;
667                }
668                .boxed(),
669            )
670        };
671
672        let mut pending_tx_receiver = self.pool.pending_transactions_listener();
673
674        // Set up intervals
675        let mut report_interval = interval(REPORT_INTERVAL);
676        let mut reconnect_interval = interval(RECONNECT_INTERVAL);
677
678        // Main event loop using select!
679        loop {
680            tokio::select! {
681                // Handle shutdown signal
682                _ = shutdown_rx.recv() => {
683                    info!(target: "ethstats", "Shutting down ethstats service");
684                    break;
685                }
686
687                // Handle messages from the read loop
688                Some(msg) = message_rx.recv() => {
689                    if let Err(e) = self.handle_message(msg).await {
690                        debug!(target: "ethstats", "Error handling message: {}", e);
691                        self.disconnect().await;
692                    }
693                }
694
695                // Handle new block
696                Some(head) = head_rx.recv() => {
697                    if let Err(e) = self.report_block(Some(head)).await {
698                        debug!(target: "ethstats", "Failed to report block: {}", e);
699                        self.disconnect().await;
700                    }
701
702                    if let Err(e) = self.report_pending().await {
703                        debug!(target: "ethstats", "Failed to report pending: {}", e);
704                        self.disconnect().await;
705                    }
706                }
707
708                // Handle new pending tx
709                _= pending_tx_receiver.recv() => {
710                    if let Err(e) = self.report_pending().await {
711                        debug!(target: "ethstats", "Failed to report pending: {}", e);
712                        self.disconnect().await;
713                    }
714                }
715
716                // Handle stats reporting
717                _ = report_interval.tick() => {
718                    if let Err(e) = self.report().await {
719                        debug!(target: "ethstats", "Failed to report: {}", e);
720                        self.disconnect().await;
721                    }
722                }
723
724                // Handle reconnection
725                _ = reconnect_interval.tick() => {
726                    if self.conn.read().await.is_none() {
727                        match self.connect().await {
728                            Ok(_) => info!(target: "ethstats", "Reconnected successfully"),
729                            Err(e) => debug!(target: "ethstats", "Reconnect failed: {}", e),
730                        }
731                    }
732                }
733            }
734        }
735
736        // Cleanup
737        self.disconnect().await;
738
739        // Cancel background tasks
740        read_handle.abort();
741        canonical_stream_handle.abort();
742    }
743
744    /// Gracefully close the `WebSocket` connection
745    ///
746    /// Attempts to close the connection cleanly and logs any errors
747    /// that occur during the process.
748    async fn disconnect(&self) {
749        if let Some(conn) = self.conn.write().await.take() &&
750            let Err(e) = conn.close().await
751        {
752            debug!(target: "ethstats", "Error closing connection: {}", e);
753        }
754    }
755
756    /// Test helper to check connection status
757    #[cfg(test)]
758    pub async fn is_connected(&self) -> bool {
759        self.conn.read().await.is_some()
760    }
761}
762
763#[cfg(test)]
764mod tests {
765    use super::*;
766    use futures_util::{SinkExt, StreamExt};
767    use reth_network_api::noop::NoopNetwork;
768    use reth_storage_api::noop::NoopProvider;
769    use reth_transaction_pool::noop::NoopTransactionPool;
770    use serde_json::json;
771    use tokio::net::TcpListener;
772    use tokio_tungstenite::tungstenite::protocol::{frame::Utf8Bytes, Message};
773
774    const TEST_HOST: &str = "127.0.0.1";
775    const TEST_PORT: u16 = 0; // Let OS choose port
776
777    async fn setup_mock_server() -> (String, tokio::task::JoinHandle<()>) {
778        let listener = TcpListener::bind((TEST_HOST, TEST_PORT)).await.unwrap();
779        let addr = listener.local_addr().unwrap();
780
781        let handle = tokio::spawn(async move {
782            let (stream, _) = listener.accept().await.unwrap();
783            let mut ws_stream = tokio_tungstenite::accept_async(stream).await.unwrap();
784
785            // Handle login
786            if let Some(Ok(Message::Text(text))) = ws_stream.next().await {
787                let value: serde_json::Value = serde_json::from_str(&text).unwrap();
788                if value["emit"][0] == "hello" {
789                    let response = json!({
790                        "emit": ["ready", []]
791                    });
792                    ws_stream
793                        .send(Message::Text(Utf8Bytes::from(response.to_string())))
794                        .await
795                        .unwrap();
796                }
797            }
798
799            // Handle ping
800            while let Some(Ok(msg)) = ws_stream.next().await {
801                if let Message::Text(text) = msg &&
802                    text.contains("node-ping")
803                {
804                    let pong = json!({
805                        "emit": ["node-pong", {"id": "test-node"}]
806                    });
807                    ws_stream.send(Message::Text(Utf8Bytes::from(pong.to_string()))).await.unwrap();
808                }
809            }
810        });
811
812        (addr.to_string(), handle)
813    }
814
815    #[tokio::test]
816    async fn test_connection_and_login() {
817        let (server_url, server_handle) = setup_mock_server().await;
818        let ethstats_url = format!("test-node:test-secret@{server_url}");
819
820        let network = NoopNetwork::default();
821        let provider = NoopProvider::default();
822        let pool = NoopTransactionPool::default();
823
824        let service = EthStatsService::new(&ethstats_url, network, provider, pool)
825            .await
826            .expect("Service should connect");
827
828        // Verify connection was established
829        assert!(service.is_connected().await, "Service should be connected");
830
831        // Clean up server
832        server_handle.abort();
833    }
834
835    #[tokio::test]
836    async fn test_history_command_handling() {
837        let (server_url, server_handle) = setup_mock_server().await;
838        let ethstats_url = format!("test-node:test-secret@{server_url}");
839
840        let network = NoopNetwork::default();
841        let provider = NoopProvider::default();
842        let pool = NoopTransactionPool::default();
843
844        let service = EthStatsService::new(&ethstats_url, network, provider, pool)
845            .await
846            .expect("Service should connect");
847
848        // Simulate receiving a history command
849        let history_cmd = json!({
850            "emit": ["history", {"list": [1, 2, 3]}]
851        });
852
853        service.handle_message(history_cmd).await.expect("History command should be handled");
854
855        // Clean up server
856        server_handle.abort();
857    }
858
859    #[tokio::test]
860    async fn test_invalid_url_handling() {
861        let network = NoopNetwork::default();
862        let provider = NoopProvider::default();
863        let pool = NoopTransactionPool::default();
864
865        // Test missing secret
866        let result = EthStatsService::new(
867            "test-node@localhost",
868            network.clone(),
869            provider.clone(),
870            pool.clone(),
871        )
872        .await;
873        assert!(
874            matches!(result, Err(EthStatsError::InvalidUrl(_))),
875            "Should detect invalid URL format"
876        );
877
878        // Test invalid URL format
879        let result = EthStatsService::new("invalid-url", network, provider, pool).await;
880        assert!(
881            matches!(result, Err(EthStatsError::InvalidUrl(_))),
882            "Should detect invalid URL format"
883        );
884    }
885
886    #[tokio::test(flavor = "current_thread")]
887    async fn report_latency_lock_order_regression() {
888        // Simulate a live connection so a pong handler (report_latency) can grab conn.read().
889        let (server_url, server_handle) = setup_mock_server().await;
890        let ethstats_url = format!("test-node:test-secret@{server_url}");
891
892        let network = NoopNetwork::default();
893        let provider = NoopProvider::default();
894        let pool = NoopTransactionPool::default();
895
896        let service = EthStatsService::new(&ethstats_url, network, provider, pool)
897            .await
898            .expect("Service should connect");
899
900        // Keep last_ping set to mimic an outstanding ping while a pong is being handled.
901        let mut last_ping_guard = service.last_ping.lock().await;
902        *last_ping_guard = Some(Instant::now());
903
904        let mut report_latency = Box::pin(service.report_latency());
905        assert!(futures_util::poll!(report_latency.as_mut()).is_pending());
906
907        let write_guard = service
908            .conn
909            .try_write()
910            .expect("conn write lock should not be held while report_latency waits on last_ping");
911
912        drop(write_guard);
913        drop(last_ping_guard);
914
915        let _ = report_latency.await;
916        server_handle.abort();
917    }
918}