feat(ws): retry mechanism in WsConnect (#2303)

* feat(`ws`): retry mechanism in WsConnect

* wasm

Author: yash-atreya

Diff for: crates/pubsub/src/handle.rs

@@ -1,8 +1,11 @@
 use alloy_json_rpc::PubSubItem;
 use serde_json::value::RawValue;
-use tokio::sync::{
-    mpsc,
-    oneshot::{self, error::TryRecvError},
+use tokio::{
+    sync::{
+        mpsc,
+        oneshot::{self, error::TryRecvError},
+    },
+    time::Duration,
 };
 
 /// A handle to a backend. Communicates to a `ConnectionInterface` on the
@@ -23,6 +26,13 @@ pub struct ConnectionHandle {
 
     /// Notify the backend of intentional shutdown.
     pub(crate) shutdown: oneshot::Sender<()>,
+
+    /// Max number of retries before failing and exiting the connection.
+    /// Default is 10.
+    pub(crate) max_retries: u32,
+    /// The interval between retries.
+    /// Default is 3 seconds.
+    pub(crate) retry_interval: Duration,
 }
 
 impl ConnectionHandle {
@@ -33,7 +43,14 @@ impl ConnectionHandle {
         let (error_tx, error_rx) = oneshot::channel();
         let (shutdown_tx, shutdown_rx) = oneshot::channel();
 
-        let handle = Self { to_socket, from_socket, error: error_rx, shutdown: shutdown_tx };
+        let handle = Self {
+            to_socket,
+            from_socket,
+            error: error_rx,
+            shutdown: shutdown_tx,
+            max_retries: 10,
+            retry_interval: Duration::from_secs(3),
+        };
         let interface = ConnectionInterface {
             from_frontend,
             to_frontend,
@@ -43,6 +60,19 @@ impl ConnectionHandle {
         (handle, interface)
     }
 
+    /// Set the max number of retries before failing and exiting the connection.
+    /// Default is 10.
+    pub fn with_max_retries(mut self, max_retries: u32) -> Self {
+        self.max_retries = max_retries;
+        self
+    }
+
+    /// Set the interval between retries.
+    pub fn with_retry_interval(mut self, retry_interval: Duration) -> Self {
+        self.retry_interval = retry_interval;
+        self
+    }
+
     /// Shutdown the backend.
     pub fn shutdown(self) {
         let _ = self.shutdown.send(());

Diff for: crates/pubsub/src/service.rs

@@ -37,11 +37,6 @@ pub(crate) struct PubSubService<T> {
 
     /// The request manager.
     pub(crate) in_flights: RequestManager,
-
-    /// Number of retries. Default is 10.
-    ///
-    /// Every retry is made at an interval of 3 seconds.
-    pub(crate) retries: u32,
 }
 
 impl<T: PubSubConnect> PubSubService<T> {
@@ -56,7 +51,6 @@ impl<T: PubSubConnect> PubSubService<T> {
             reqs,
             subs: SubscriptionManager::default(),
             in_flights: Default::default(),
-            retries: 10,
         };
         this.spawn();
         Ok(PubSubFrontend::new(tx))
@@ -205,7 +199,8 @@ impl<T: PubSubConnect> PubSubService<T> {
     /// Attempt to reconnect with retries
     async fn reconnect_with_retries(&mut self) -> TransportResult<()> {
         let mut retry_count = 0;
-        let max_retries = self.retries;
+        let max_retries = self.handle.max_retries;
+        let interval = self.handle.retry_interval;
         loop {
             match self.reconnect().await {
                 Ok(()) => break Ok(()),
@@ -215,13 +210,12 @@ impl<T: PubSubConnect> PubSubService<T> {
                         error!("Reconnect failed after {max_retries} attempts, shutting down: {e}");
                         break Err(e);
                     }
-                    let duration = std::time::Duration::from_secs(3);
                     warn!(
                         "Reconnection attempt {retry_count}/{max_retries} failed: {e}. \
-                         Retrying in {:.3}s...",
-                        duration.as_secs_f64(),
+                         Retrying in {:?}s...",
+                        interval.as_secs_f64(),
                     );
-                    sleep(duration).await;
+                    sleep(interval).await;
                 }
             }
         }

Diff for: crates/transport-ws/src/native.rs

@@ -25,12 +25,24 @@ pub struct WsConnect {
     pub auth: Option<Authorization>,
     /// The websocket config.
     pub config: Option<WebSocketConfig>,
+    /// Max number of retries before failing and exiting the connection.
+    /// Default is 10.
+    max_retries: u32,
+    /// The interval between retries.
+    /// Default is 3 seconds.
+    retry_interval: Duration,
 }
 
 impl WsConnect {
     /// Creates a new websocket connection configuration.
     pub fn new<S: Into<String>>(url: S) -> Self {
-        Self { url: url.into(), auth: None, config: None }
+        Self {
+            url: url.into(),
+            auth: None,
+            config: None,
+            max_retries: 10,
+            retry_interval: Duration::from_secs(3),
+        }
     }
 
     /// Sets the authorization header.
@@ -44,6 +56,20 @@ impl WsConnect {
         self.config = Some(config);
         self
     }
+
+    /// Sets the max number of retries before failing and exiting the connection.
+    /// Default is 10.
+    pub const fn with_max_retries(mut self, max_retries: u32) -> Self {
+        self.max_retries = max_retries;
+        self
+    }
+
+    /// Sets the interval between retries.
+    /// Default is 3 seconds.
+    pub const fn with_retry_interval(mut self, retry_interval: Duration) -> Self {
+        self.retry_interval = retry_interval;
+        self
+    }
 }
 
 impl IntoClientRequest for WsConnect {
@@ -77,7 +103,7 @@ impl PubSubConnect for WsConnect {
 
         backend.spawn();
 
-        Ok(handle)
+        Ok(handle.with_max_retries(self.max_retries).with_retry_interval(self.retry_interval))
     }
 }
 

Diff for: crates/transport-ws/src/wasm.rs

@@ -6,19 +6,40 @@ use futures::{
     stream::{Fuse, StreamExt},
 };
 use serde_json::value::RawValue;
+use std::time::Duration;
 use ws_stream_wasm::{WsErr, WsMessage, WsMeta, WsStream};
 
 /// Simple connection info for the websocket.
 #[derive(Clone, Debug)]
 pub struct WsConnect {
     /// The URL to connect to.
     pub url: String,
+    /// Max number of retries before failing and exiting the connection.
+    /// Default is 10.
+    max_retries: u32,
+    /// The interval between retries.
+    /// Default is 3 seconds.
+    retry_interval: Duration,
 }
 
 impl WsConnect {
     /// Creates a new websocket connection configuration.
     pub fn new<S: Into<String>>(url: S) -> Self {
-        Self { url: url.into() }
+        Self { url: url.into(), max_retries: 10, retry_interval: Duration::from_secs(3) }
+    }
+
+    /// Sets the max number of retries before failing and exiting the connection.
+    /// Default is 10.
+    pub const fn with_max_retries(mut self, max_retries: u32) -> Self {
+        self.max_retries = max_retries;
+        self
+    }
+
+    /// Sets the interval between retries.
+    /// Default is 3 seconds.
+    pub const fn with_retry_interval(mut self, retry_interval: Duration) -> Self {
+        self.retry_interval = retry_interval;
+        self
     }
 }
 
@@ -36,7 +57,7 @@ impl PubSubConnect for WsConnect {
 
         backend.spawn();
 
-        Ok(handle)
+        Ok(handle.with_max_retries(self.max_retries).with_retry_interval(self.retry_interval))
     }
 }