From 27e3616793f2e7ee139c93c378c29c9fa40f741f Mon Sep 17 00:00:00 2001
From: Yordan Atanasov <yordan.atanasov@limechain.tech>
Date: Thu, 19 Sep 2024 13:58:44 +0300
Subject: [PATCH] refactor: Address PR comments.

---
 .../com/limechain/rpc/WebsocketState.java     | 16 ++++++++++++
 .../com/limechain/rpc/WsRpcClientImpl.java    | 26 +++++++++++--------
 2 files changed, 31 insertions(+), 11 deletions(-)
 create mode 100644 src/main/java/com/limechain/rpc/WebsocketState.java

diff --git a/src/main/java/com/limechain/rpc/WebsocketState.java b/src/main/java/com/limechain/rpc/WebsocketState.java
new file mode 100644
index 000000000..e8e353f73
--- /dev/null
+++ b/src/main/java/com/limechain/rpc/WebsocketState.java
@@ -0,0 +1,16 @@
+package com.limechain.rpc;
+
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+
+@Getter
+@RequiredArgsConstructor
+public enum WebsocketState {
+
+    CONNECTING(0),
+    OPEN(1),
+    CLOSING(2),
+    CLOSED(3);
+
+    private final int intValue;
+}
diff --git a/src/main/java/com/limechain/rpc/WsRpcClientImpl.java b/src/main/java/com/limechain/rpc/WsRpcClientImpl.java
index 73410cc82..e4a1dfa90 100644
--- a/src/main/java/com/limechain/rpc/WsRpcClientImpl.java
+++ b/src/main/java/com/limechain/rpc/WsRpcClientImpl.java
@@ -2,6 +2,7 @@
 
 import com.limechain.constants.RpcConstants;
 import lombok.SneakyThrows;
+import lombok.extern.java.Log;
 import org.teavm.jso.browser.Window;
 import org.teavm.jso.core.JSString;
 import org.teavm.jso.websocket.WebSocket;
@@ -12,8 +13,11 @@
 /**
  * The implementation of {@link WsRpcClient}. Uses a native JS Websocket implementation.
  */
+@Log
 public class WsRpcClientImpl implements WsRpcClient {
 
+    private static final int WS_OPEN_WAIT_MS = 50;
+
     private WebSocket ws;
     private final Queue<String> responseQueue;
 
@@ -23,7 +27,7 @@ public WsRpcClientImpl() {
     }
 
     private void openWebsocketConnection() {
-        System.out.println("Initializing RPC websocket connection...");
+        log.info("Initializing RPC websocket connection...");
         //TODO change when configuring chain.
         ws = new WebSocket(RpcConstants.POLKADOT_WS_RPC);
         initHandlers();
@@ -31,17 +35,17 @@ private void openWebsocketConnection() {
 
     private void initHandlers() {
         ws.onClose(e -> {
-            System.out.println("RPC websocket connection was closed.");
-            System.out.println("Retrying connection...");
+            log.info("RPC websocket connection was closed.");
+            log.info("Retrying connection...");
             Window.setTimeout(this::openWebsocketConnection, 1000);
         });
 
         ws.onError(e -> {
-            System.out.println("There was an error in the RPC websocket connection. Closing connection...");
+            log.warning("There was an error in the RPC websocket connection. Closing connection...");
             ws.close();
         });
 
-        ws.onOpen(e -> System.out.println("Websocket connection is open."));
+        ws.onOpen(e -> log.info("Websocket connection is open."));
         ws.onMessage(e -> responseQueue.offer(e.getDataAsString()));
     }
 
@@ -62,16 +66,16 @@ public void send(JSString rpcString) {
      */
     @SneakyThrows
     private void handleSocketState() {
-        var startState = ws.getReadyState();
+        int startState = ws.getReadyState();
+        int openedState = WebsocketState.OPEN.getIntValue();
 
-        while (startState != 1) {
-            var currentState = ws.getReadyState();
-            if (currentState > 1) {
+        while (startState != openedState) {
+            if (startState > openedState) {
                 throw new Exception("Calling function of a closed websocket is prohibited.");
             }
 
-            startState = currentState;
-            Thread.sleep(50);
+            Thread.sleep(WS_OPEN_WAIT_MS);
+            startState = ws.getReadyState();
         }
     }