From efe5cd7415174198cbf0f534c0dbfa454a55c586 Mon Sep 17 00:00:00 2001
From: voidash <ashish.thapa477@gmail.com>
Date: Sun, 5 Jan 2025 16:46:27 +0545
Subject: [PATCH] bridge-exec: remove working, notworking state from WsClient

---
 .../src/modes/operator/task_manager.rs        |  3 +-
 crates/bridge-exec/src/handler.rs             |  4 +-
 crates/bridge-exec/src/ws_client.rs           | 67 +++++--------------
 3 files changed, 18 insertions(+), 56 deletions(-)

diff --git a/bin/bridge-client/src/modes/operator/task_manager.rs b/bin/bridge-client/src/modes/operator/task_manager.rs
index 9fe40dd45..ad18c5142 100644
--- a/bin/bridge-client/src/modes/operator/task_manager.rs
+++ b/bin/bridge-client/src/modes/operator/task_manager.rs
@@ -115,8 +115,7 @@ where
 
         let l2_rpc_client = self
             .exec_handler
-            .l2_rpc_client_pool
-            .get()
+            .get_ready_rpc_client()
             .await
             .map_err(|_| PollDutyError::WsPool)?;
 
diff --git a/crates/bridge-exec/src/handler.rs b/crates/bridge-exec/src/handler.rs
index 571419644..90738cac7 100644
--- a/crates/bridge-exec/src/handler.rs
+++ b/crates/bridge-exec/src/handler.rs
@@ -22,7 +22,7 @@ use tracing::{debug, info, warn};
 
 use crate::{
     errors::{ExecError, ExecResult},
-    ws_client::{WsClient, WsClientManager},
+    ws_client::WsClientManager,
 };
 
 /// Websocket client pool
@@ -320,7 +320,7 @@ where
     }
 
     /// Retrieves a ready-to-use RPC client from the client pool.
-    async fn get_ready_rpc_client(&self) -> Result<Object<WsClientManager>, ExecError> {
+    pub async fn get_ready_rpc_client(&self) -> Result<Object<WsClientManager>, ExecError> {
         self.l2_rpc_client_pool
             .get()
             .await
diff --git a/crates/bridge-exec/src/ws_client.rs b/crates/bridge-exec/src/ws_client.rs
index 8e8c79862..4fce783eb 100644
--- a/crates/bridge-exec/src/ws_client.rs
+++ b/crates/bridge-exec/src/ws_client.rs
@@ -10,7 +10,7 @@ use jsonrpsee::{
         client::{BatchResponse, ClientT},
         params::BatchRequestBuilder,
         traits::ToRpcParams,
-        BoxError, ClientError, DeserializeOwned,
+        ClientError, DeserializeOwned,
     },
     ws_client::{WsClient as WebsocketClient, WsClientBuilder},
 };
@@ -35,25 +35,14 @@ pub struct WsClientManager {
     pub config: WsClientConfig,
 }
 
-/// Represents the state of a WebSocket client.
-///
-/// - `Working`: The client is connected and operational.
-/// - `NotWorking`: The client is not connected or has encountered an error.
-#[derive(Debug)]
-pub enum WsClientState {
-    /// The WebSocket client is connected and operational.
-    Working(WebsocketClient),
-
-    /// The WebSocket client is not connected or is in a failed state.
-    NotWorking,
-}
-
 /// Wrapper for the WebSocket client state.
 ///
 /// This struct encapsulates the `WsClientState`, enabling unified management of
 /// both connected and failed client states.
 #[derive(Debug)]
-pub struct WsClient(WsClientState);
+pub struct WsClient {
+    inner: WebsocketClient,
+}
 
 /// Implements the `Manager` trait for managing WebSocket clients.
 ///
@@ -71,13 +60,9 @@ impl Manager for WsClientManager {
     async fn create(&self) -> Result<Self::Type, Self::Error> {
         let client = WsClientBuilder::default()
             .build(self.config.url.clone())
-            .await;
+            .await?;
 
-        let bl = match client {
-            Ok(cl) => WsClientState::Working(cl),
-            Err(_) => WsClientState::NotWorking,
-        };
-        Ok(WsClient(bl))
+        Ok(WsClient { inner: client })
     }
 
     /// Recycles an existing WebSocket client.
@@ -90,29 +75,16 @@ impl Manager for WsClientManager {
         obj: &mut Self::Type,
         _metrics: &managed::Metrics,
     ) -> RecycleResult<Self::Error> {
-        match &obj.0 {
-            WsClientState::Working(cl) => {
-                if cl.is_connected() {
-                    Ok(())
-                } else {
-                    Err(RecycleError::Message(
-                        "Connection lost, recreate client".to_string().into(),
-                    ))
-                }
-            }
-            WsClientState::NotWorking => Err(RecycleError::Message(
-                "Connection still not found, recreate client"
-                    .to_string()
-                    .into(),
-            )),
+        if obj.inner.is_connected() {
+            Ok(())
+        } else {
+            Err(RecycleError::Message(
+                "Connection lost, recreate client".to_string().into(),
+            ))
         }
     }
 }
 
-fn make_not_working_error() -> ClientError {
-    ClientError::Transport(BoxError::from("Client is Not Working".to_string()))
-}
-
 /// Implements the `ClientT` trait for `WsClient`.
 ///
 /// This implementation allows `WsClient` to perform JSON-RPC operations,
@@ -124,10 +96,7 @@ impl ClientT for WsClient {
     where
         Params: ToRpcParams + Send,
     {
-        match &self.0 {
-            WsClientState::Working(inner) => inner.notification(method, params).await,
-            WsClientState::NotWorking => Err(make_not_working_error()),
-        }
+        self.inner.notification(method, params).await
     }
 
     /// Sends a method call request.
@@ -136,10 +105,7 @@ impl ClientT for WsClient {
         R: DeserializeOwned,
         Params: ToRpcParams + Send,
     {
-        match &self.0 {
-            WsClientState::Working(inner) => inner.request(method, params).await,
-            WsClientState::NotWorking => Err(make_not_working_error()),
-        }
+        self.inner.request(method, params).await
     }
 
     /// Sends a batch request.
@@ -150,9 +116,6 @@ impl ClientT for WsClient {
     where
         R: DeserializeOwned + fmt::Debug + 'a,
     {
-        match &self.0 {
-            WsClientState::Working(inner) => inner.batch_request(batch).await,
-            WsClientState::NotWorking => Err(make_not_working_error()),
-        }
+        self.inner.batch_request(batch).await
     }
 }