sequencer(sequencer/providers/websocket): Refactor websocket connector to be used in provider and reorg tracker

Author: HristoStaykov

Cargo.lock

@@ -48,6 +48,7 @@ alloy-primitives = { workspace = true }
 alloy-u256-literal = { workspace = true }
 
 anyhow = { workspace = true }
+async-trait = "0.1"
 bridgetree = { git = "https://github.com/zcash/incrementalmerkletree", package = "bridgetree" }
 bytes = { workspace = true }
 chrono = { workspace = true }

apps/sequencer/Cargo.toml

@@ -331,12 +331,16 @@ pub async fn create_per_network_reorg_trackers(
         let net_clone = net.clone();
         let sequencer_state_providers_clone = sequencer_state.providers.clone();
         let sequencer_config = sequencer_state.sequencer_config.read().await;
-        let (reorg_tracker_config, websocket_url_opt) =
+        let (reorg_tracker_config, websocket_url_opt, websocket_reconnect_opt) =
             match sequencer_config.providers.get(net.as_str()) {
-                Some(c) => (c.reorg.clone(), c.websocket_url.clone()),
+                Some(c) => (
+                    c.reorg.clone(),
+                    c.websocket_url.clone(),
+                    c.websocket_reconnect.clone(),
+                ),
                 None => {
                     error!("No config for provider for network {net} will set to default!");
-                    (ReorgConfig::default(), None)
+                    (ReorgConfig::default(), None, None)
                 }
             };
         let relayer_send_channel = match sequencer_state
@@ -356,6 +360,7 @@ pub async fn create_per_network_reorg_trackers(
             sequencer_state_providers_clone,
             relayer_send_channel,
             websocket_url_opt,
+            websocket_reconnect_opt,
         );
         collected_futures.push(
             tokio::task::Builder::new()

apps/sequencer/src/providers/eth_send_utils.rs

@@ -2,3 +2,4 @@ pub mod eth_send_utils;
 pub mod inflight_observations;
 pub mod provider;
 pub mod reorg_tracking;
+pub mod ws;

apps/sequencer/src/providers/mod.rs

@@ -1,7 +1,5 @@
 use alloy::providers::Provider;
-use alloy::pubsub::{ConnectionHandle, PubSubConnect};
 use alloy::rpc::types::{TransactionInput, TransactionRequest};
-use alloy::transports::TransportResult;
 use alloy::{
     dyn_abi::DynSolValue,
     hex,
@@ -26,7 +24,7 @@ use incrementalmerkletree::{frontier::Frontier, Hashable, Level};
 use reqwest::Url; // TODO @ymadzhunkov include URL directly from url crate
 
 use blocksense_config::{
-    AllFeedsConfig, ContractConfig, PublishCriteria, SequencerConfig, WebsocketReconnectConfig,
+    AllFeedsConfig, ContractConfig, PublishCriteria, SequencerConfig,
     ADFS_ACCESS_CONTROL_CONTRACT_NAME, ADFS_CONTRACT_NAME,
 };
 use blocksense_data_feeds::feeds_processing::{
@@ -43,13 +41,15 @@ use std::collections::HashMap;
 use std::sync::Arc;
 use std::{fs, mem};
 use tokio::sync::{Mutex, RwLock};
-use tokio::time::{error::Elapsed, sleep, Duration};
+use tokio::time::{error::Elapsed, Duration};
 use tracing::{debug, error, info, warn};
 
 use crate::providers::eth_send_utils::{
     get_gas_limit, get_tx_retry_params, BatchOfUpdatesToProcess, GasFees,
 };
 use crate::providers::inflight_observations::InflightObservations;
+use crate::providers::ws::{ResilientWsConnect, WsReconnectMetrics, WsReconnectPolicy};
+use async_trait::async_trait;
 use std::time::Instant;
 
 pub type ProviderType =
@@ -108,149 +108,47 @@ impl Hashable for HashValue {
     }
 }
 
-#[derive(Debug, Clone)]
-struct WsReconnectPolicy {
-    initial: Duration,
-    max: Duration,
-    multiplier: f64,
-}
-
-impl WsReconnectPolicy {
-    fn from_config(config: Option<&WebsocketReconnectConfig>) -> Self {
-        let cfg = config.cloned().unwrap_or_default();
-        let initial = Duration::from_millis(cfg.initial_backoff_ms);
-        let max = Duration::from_millis(cfg.max_backoff_ms);
-        let multiplier = cfg.backoff_multiplier;
-
-        // Guard against misconfigured values just in case validation was skipped.
-        let clamped_initial = initial.min(max);
-
-        Self {
-            initial: clamped_initial,
-            max,
-            multiplier,
-        }
-    }
-
-    fn backoff_delay(&self, attempt: u64) -> Duration {
-        if attempt == 0 {
-            return Duration::ZERO;
-        }
-        let exponent = (attempt - 1) as f64;
-        let scaled = self.initial.mul_f64(self.multiplier.powf(exponent));
-        scaled.min(self.max)
-    }
-}
-
-impl Default for WsReconnectPolicy {
-    fn default() -> Self {
-        Self::from_config(None)
-    }
-}
-
-#[derive(Clone)]
-struct ResilientWsConnect {
-    inner: WsConnect,
-    policy: WsReconnectPolicy,
+struct ProviderWsRecorder {
     metrics: Arc<RwLock<ProviderMetrics>>,
-    network: Arc<String>,
+    network: String,
 }
 
-impl ResilientWsConnect {
-    fn new(
-        inner: WsConnect,
-        policy: WsReconnectPolicy,
-        metrics: Arc<RwLock<ProviderMetrics>>,
-        network: &str,
-    ) -> Self {
+impl ProviderWsRecorder {
+    fn new(metrics: Arc<RwLock<ProviderMetrics>>, network: &str) -> Self {
         Self {
-            inner,
-            policy,
             metrics,
-            network: Arc::new(network.to_owned()),
+            network: network.to_owned(),
         }
     }
 }
 
-impl PubSubConnect for ResilientWsConnect {
-    fn is_local(&self) -> bool {
-        self.inner.is_local()
+#[async_trait]
+impl WsReconnectMetrics for ProviderWsRecorder {
+    async fn on_disconnect(&self) {
+        self.metrics
+            .read()
+            .await
+            .ws_disconnects_detected
+            .with_label_values(&[self.network.as_str()])
+            .inc();
     }
 
-    async fn connect(&self) -> TransportResult<ConnectionHandle> {
-        let inner = self.inner.clone();
-        let handle = PubSubConnect::connect(&inner).await?;
-        Ok(handle
-            .with_max_retries(u32::MAX)
-            .with_retry_interval(Duration::from_secs(0)))
+    async fn on_attempt(&self) {
+        self.metrics
+            .read()
+            .await
+            .ws_reconnect_attempts
+            .with_label_values(&[self.network.as_str()])
+            .inc();
     }
 
-    async fn try_reconnect(&self) -> TransportResult<ConnectionHandle> {
-        let inner = self.inner.clone();
-        let policy = self.policy.clone();
-        let metrics = self.metrics.clone();
-        let network = self.network.clone();
-
-        {
-            let guard = metrics.read().await;
-            guard
-                .ws_disconnects_detected
-                .with_label_values(&[network.as_str()])
-                .inc();
-        }
-
-        warn!(
-            network = network.as_str(),
-            initial_backoff_ms = policy.initial.as_millis(),
-            max_backoff_ms = policy.max.as_millis(),
-            multiplier = policy.multiplier,
-            "WS transport disconnected; starting exponential reconnect attempts"
-        );
-
-        let mut attempt: u64 = 0;
-        loop {
-            attempt = attempt.saturating_add(1);
-            {
-                let guard = metrics.read().await;
-                guard
-                    .ws_reconnect_attempts
-                    .with_label_values(&[network.as_str()])
-                    .inc();
-            }
-
-            match PubSubConnect::connect(&inner).await {
-                Ok(handle) => {
-                    {
-                        let guard = metrics.read().await;
-                        guard
-                            .ws_reconnect_successes
-                            .with_label_values(&[network.as_str()])
-                            .inc();
-                    }
-
-                    info!(
-                        network = network.as_str(),
-                        attempt, "WS transport reconnected after {attempt} attempt(s)"
-                    );
-
-                    return Ok(handle
-                        .with_max_retries(u32::MAX)
-                        .with_retry_interval(Duration::from_secs(0)));
-                }
-                Err(err) => {
-                    let delay = policy.backoff_delay(attempt);
-                    warn!(
-                        network = network.as_str(),
-                        attempt,
-                        backoff_ms = delay.as_millis(),
-                        capped = delay == policy.max,
-                        error = %err,
-                        "WS reconnect attempt failed; will retry"
-                    );
-                    sleep(delay).await;
-                }
-            }
-        }
+    async fn on_success(&self) {
+        self.metrics
+            .read()
+            .await
+            .ws_reconnect_successes
+            .with_label_values(&[self.network.as_str()])
+            .inc();
     }
 }
 pub struct RpcProvider {
@@ -458,10 +356,14 @@ impl RpcProvider {
     ) -> RpcProvider {
         let provider = match rpc_url.scheme() {
             "ws" | "wss" => {
+                let metrics_recorder = Arc::new(ProviderWsRecorder::new(
+                    Arc::clone(provider_metrics),
+                    network,
+                ));
                 let resilient_connect = ResilientWsConnect::new(
                     WsConnect::new(rpc_url.as_str().to_owned()),
                     WsReconnectPolicy::from_config(p.websocket_reconnect.as_ref()),
-                    Arc::clone(provider_metrics),
+                    metrics_recorder,
                     network,
                 );