From d9d3726aa97d7a429cbaf5a76076828ad040c290 Mon Sep 17 00:00:00 2001
From: Declan Chidlow <git.rh92c@8shield.net>
Date: Sat, 13 Jul 2024 16:47:24 +0800
Subject: [PATCH] Better handle WS connection errors

---
 bot/src/bot/modules/api_communication.ts | 40 +++++++++++++++++-------
 1 file changed, 29 insertions(+), 11 deletions(-)

diff --git a/bot/src/bot/modules/api_communication.ts b/bot/src/bot/modules/api_communication.ts
index 163b04c..1efcd0b 100644
--- a/bot/src/bot/modules/api_communication.ts
+++ b/bot/src/bot/modules/api_communication.ts
@@ -15,13 +15,17 @@ const wsEvents = new EventEmitter();
 const { API_WS_URL, API_WS_TOKEN } = process.env;
 const wsQueue: { [key: string]: string }[] = [];
 let client: ws | undefined = undefined;
+let retryCount = 0;
+const MAX_RETRIES = 10;
+const INITIAL_RETRY_DELAY = 3000;
 
 type WSResponse = { success: false; error: string; statusCode?: number } | { success: true; [key: string]: any };
 
-if (!API_WS_URL || !API_WS_TOKEN) logger.info("$API_WS_URL or $API_WS_TOKEN not found.");
-else {
-	logger.info(`$API_WS_URL and $API_WS_TOKEN set; Connecting to ${API_WS_URL}`);
-	connect();
+if (!API_WS_URL || !API_WS_TOKEN) {
+  logger.error("$API_WS_URL or $API_WS_TOKEN not found. Please set these environment variables.");
+} else {
+  logger.info(`$API_WS_URL and $API_WS_TOKEN set; Connecting to ${API_WS_URL}`);
+  connect();
 }
 
 function connect() {
@@ -29,7 +33,8 @@ function connect() {
 	client = new ws(API_WS_URL!, { headers: { authorization: API_WS_TOKEN! } });
 
 	client.once("open", () => {
-		logger.debug("WS connected");
+		logger.info("WebSocket connected successfully");
+    retryCount = 0;
 		if (wsQueue.length > 0) {
 			logger.debug(`Attempting to send ${wsQueue.length} queued WS messages`);
 
@@ -44,14 +49,14 @@ function connect() {
 
 	client.once("close", () => {
 		client = undefined;
-		logger.warn(`WS closed, reconnecting in 3 seconds`);
-		setTimeout(connect, 3000);
+		retryConnection();
 	});
 
-	client.once("error", (err: Error) => {
-		client = undefined;
-		logger.warn(`WS: ${err}`);
-	});
+  client.once("error", (err: Error) => {
+    client = undefined;
+    logger.error(`WebSocket error: ${err.message}`);
+    retryConnection();
+  });
 
 	client.on("message", (msg: ws.Data) => {
 		logger.debug(`[WS] [<] ${msg.toString("utf8")}`);
@@ -89,6 +94,19 @@ function wsSend(data: { [key: string]: any }) {
 	}
 }
 
+function retryConnection() {
+  if (retryCount >= MAX_RETRIES) {
+    logger.error(`Failed to connect after ${MAX_RETRIES} attempts. Please check your network and API_WS_URL.`);
+    return;
+  }
+
+  const delay = INITIAL_RETRY_DELAY * Math.pow(2, retryCount);
+  retryCount++;
+
+  logger.warn(`WebSocket disconnected. Attempting to reconnect in ${delay / 1000} seconds (Attempt ${retryCount} of ${MAX_RETRIES})`);
+  setTimeout(connect, delay);
+}
+
 wsEvents.on("req:test", (data: any, res: (data: any) => void) => {
 	res({ received: data });
 });