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

@@ -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 });
-            }
           },
         }
       );
@@ -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();
       };
     },
     {

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;
+    },
+  });
+};