From a9149fb92ede2ca22677c0c3bc40b0d35e97f604 Mon Sep 17 00:00:00 2001
From: Sukka <isukkaw@gmail.com>
Date: Sun, 16 Jun 2024 18:25:33 +0800
Subject: [PATCH] feat: add a wrapper around sockette w/ error retry (#1219)

* feat: add a wrapper around sockette w/ error retry

* chore: use import path alias

* perf: reduce retry
---
 src/components/layout/layout-traffic.tsx | 47 +++++++++---------------
 src/utils/websocket.ts                   | 39 ++++++++++++++++++++
 2 files changed, 57 insertions(+), 29 deletions(-)
 create mode 100644 src/utils/websocket.ts

diff --git a/src/components/layout/layout-traffic.tsx b/src/components/layout/layout-traffic.tsx
index 6b8dd329..c6a25c55 100644
--- a/src/components/layout/layout-traffic.tsx
+++ b/src/components/layout/layout-traffic.tsx
@@ -12,7 +12,7 @@ import { useLogSetup } from "./use-log-setup";
 import { useVisibility } from "@/hooks/use-visibility";
 import parseTraffic from "@/utils/parse-traffic";
 import useSWRSubscription from "swr/subscription";
-import Sockette from "sockette";
+import { createSockette } from "@/utils/websocket";
 
 interface MemoryUsage {
   inuse: number;
@@ -42,24 +42,17 @@ export const LayoutTraffic = () => {
     (_key, { next }) => {
       const { server = "", secret = "" } = clashInfo!;
 
-      let errorCount = 10;
-
-      const s = new Sockette(
+      const s = createSockette(
         `ws://${server}/traffic?token=${encodeURIComponent(secret)}`,
         {
           onmessage(event) {
-            errorCount = 0; // reset counter
             const data = JSON.parse(event.data) as ITrafficItem;
             trafficRef.current?.appendData(data);
             next(null, data);
           },
           onerror(event) {
-            errorCount -= 1;
-
-            if (errorCount <= 0) {
-              this.close();
-              next(event, { up: 0, down: 0 });
-            }
+            this.close();
+            next(event, { up: 0, down: 0 });
           },
         }
       );
@@ -86,27 +79,23 @@ export const LayoutTraffic = () => {
     clashInfo && pageVisible && displayMemory ? "getRealtimeMemory" : null,
     (_key, { next }) => {
       const { server = "", secret = "" } = clashInfo!;
-      const ws = new WebSocket(
-        `ws://${server}/memory?token=${encodeURIComponent(secret)}`
+
+      const s = createSockette(
+        `ws://${server}/memory?token=${encodeURIComponent(secret)}`,
+        {
+          onmessage(event) {
+            const data = JSON.parse(event.data) as MemoryUsage;
+            next(null, data);
+          },
+          onerror(event) {
+            this.close();
+            next(event, { inuse: 0 });
+          },
+        }
       );
 
-      let errorCount = 10;
-
-      ws.addEventListener("message", (event) => {
-        errorCount = 0; // reset counter
-        next(null, JSON.parse(event.data));
-      });
-      ws.addEventListener("error", (event) => {
-        errorCount -= 1;
-
-        if (errorCount <= 0) {
-          ws.close();
-          next(event, { inuse: 0 });
-        }
-      });
-
       return () => {
-        ws.close();
+        s.close();
       };
     },
     {
diff --git a/src/utils/websocket.ts b/src/utils/websocket.ts
new file mode 100644
index 00000000..9721341d
--- /dev/null
+++ b/src/utils/websocket.ts
@@ -0,0 +1,39 @@
+import Sockette, { type SocketteOptions } from "sockette";
+
+/**
+ * A wrapper of Sockette that will automatically reconnect up to `maxError` before emitting an error event.
+ */
+export const createSockette = (
+  url: string,
+  opt: SocketteOptions,
+  maxError = 10
+) => {
+  let remainRetryCount = maxError;
+
+  return new Sockette(url, {
+    ...opt,
+    // Sockette has a built-in reconnect when ECONNREFUSED feature
+    // Use maxError if opt.maxAttempts is not specified
+    maxAttempts: opt.maxAttempts ?? maxError,
+    onmessage(this: Sockette, ev) {
+      remainRetryCount = maxError; // reset counter
+      opt.onmessage?.call(this, ev);
+    },
+    onerror(this: Sockette, ev) {
+      remainRetryCount -= 1;
+
+      if (remainRetryCount >= 0) {
+        this.close();
+        this.reconnect();
+      } else {
+        opt.onerror?.call(this, ev);
+      }
+    },
+    onmaximum(this: Sockette, ev) {
+      opt.onmaximum?.call(this, ev);
+      // onmaximum will be fired when Sockette reaches built-in reconnect limit,
+      // We will also set remainRetryCount to 0 to prevent further reconnect.
+      remainRetryCount = 0;
+    },
+  });
+};