refactor: unify Mihomo WS subscription with shared hook (#5719)

* refactor: unify Mihomo WS subscription with shared hook

* refactor: relocate clash log hook and streamline services

* docs: Changelog.md

Changelog.md

@@ -56,6 +56,7 @@
 - 优化前端数据刷新
 - 优化流量采样和数据处理
 - 优化应用重启/退出时的资源清理性能, 大幅缩短执行时间
+- 优化 WebSocket 连接机制
 
 </details>
 

src/components/setting/setting-clash.tsx

@@ -9,10 +9,10 @@ import { updateGeo } from "tauri-plugin-mihomo-api";
 import { DialogRef, Switch } from "@/components/base";
 import { TooltipIcon } from "@/components/base/base-tooltip-icon";
 import { useClash } from "@/hooks/use-clash";
+import { useClashLog } from "@/hooks/use-clash-log";
 import { useVerge } from "@/hooks/use-verge";
 import { invoke_uwp_tool } from "@/services/cmds";
 import { showNotice } from "@/services/noticeService";
-import { useClashLog } from "@/services/states";
 import getSystem from "@/utils/get-system";
 
 import { ClashCoreViewer } from "./mods/clash-core-viewer";

src/hooks/use-clash-log.ts

@@ -0,0 +1,14 @@
+import { useLocalStorage } from "foxact/use-local-storage";
+
+const defaultClashLog: IClashLog = {
+  enable: true,
+  logLevel: "info",
+  logFilter: "all",
+  logOrder: "asc",
+};
+
+export const useClashLog = () =>
+  useLocalStorage<IClashLog>("clash-log", defaultClashLog, {
+    serializer: JSON.stringify,
+    deserializer: JSON.parse,
+  });

src/hooks/use-connection-data.ts

@@ -1,9 +1,10 @@
-import { useLocalStorage } from "foxact/use-local-storage";
-import { useEffect, useRef } from "react";
 import { mutate } from "swr";
-import useSWRSubscription from "swr/subscription";
 import { MihomoWebSocket } from "tauri-plugin-mihomo-api";
 
+import { useMihomoWsSubscription } from "./use-mihomo-ws-subscription";
+
+const MAX_CLOSED_CONNS_NUM = 500;
+
 export const initConnData: ConnectionMonitorData = {
   uploadTotal: 0,
   downloadTotal: 0,
@@ -18,129 +19,91 @@ export interface ConnectionMonitorData {
   closedConnections: IConnectionsItem[];
 }
 
-const MAX_CLOSED_CONNS_NUM = 500;
+const trimClosedConnections = (
+  closedConnections: IConnectionsItem[],
+): IConnectionsItem[] =>
+  closedConnections.length > MAX_CLOSED_CONNS_NUM
+    ? closedConnections.slice(-MAX_CLOSED_CONNS_NUM)
+    : closedConnections;
+
+const mergeConnectionSnapshot = (
+  payload: IConnections,
+  previous: ConnectionMonitorData = initConnData,
+): ConnectionMonitorData => {
+  const nextConnections = payload.connections ?? [];
+  const previousActive = previous.activeConnections ?? [];
+  const nextById = new Map(nextConnections.map((conn) => [conn.id, conn]));
+  const newIds = new Set(nextConnections.map((conn) => conn.id));
+
+  // Keep surviving connections in their previous relative order to reduce row reshuffle,
+  // but constrain the array to the incoming snapshot length.
+  const carried = previousActive
+    .map((prev) => {
+      const next = nextById.get(prev.id);
+      if (!next) return null;
+
+      nextById.delete(prev.id);
+      return {
+        ...next,
+        curUpload: next.upload - prev.upload,
+        curDownload: next.download - prev.download,
+      } as IConnectionsItem;
+    })
+    .filter(Boolean) as IConnectionsItem[];
+
+  const newcomers = nextConnections
+    .filter((conn) => nextById.has(conn.id))
+    .map((conn) => ({
+      ...conn,
+      curUpload: 0,
+      curDownload: 0,
+    }));
+
+  const activeConnections = [...carried, ...newcomers];
+
+  const closedConnections = trimClosedConnections([
+    ...(previous.closedConnections ?? []),
+    ...previousActive.filter((conn) => !newIds.has(conn.id)),
+  ]);
+
+  return {
+    uploadTotal: payload.uploadTotal ?? 0,
+    downloadTotal: payload.downloadTotal ?? 0,
+    activeConnections,
+    closedConnections,
+  };
+};
 
 export const useConnectionData = () => {
-  const [date, setDate] = useLocalStorage("mihomo_connection_date", Date.now());
-  const subscriptKey = `getClashConnection-${date}`;
-
-  const ws = useRef<MihomoWebSocket | null>(null);
-  const wsFirstConnection = useRef<boolean>(true);
-  const timeoutRef = useRef<ReturnType<typeof setTimeout>>(null);
-
-  const response = useSWRSubscription<
-    ConnectionMonitorData,
-    any,
-    string | null
-  >(
-    subscriptKey,
-    (_key, { next }) => {
-      const reconnect = async () => {
-        await ws.current?.close();
-        ws.current = null;
-        timeoutRef.current = setTimeout(async () => await connect(), 500);
-      };
-
-      const connect = () =>
-        MihomoWebSocket.connect_connections()
-          .then((ws_) => {
-            ws.current = ws_;
-            if (timeoutRef.current) clearTimeout(timeoutRef.current);
-
-            ws_.addListener(async (msg) => {
-              if (msg.type === "Text") {
-                if (msg.data.startsWith("Websocket error")) {
-                  next(msg.data);
-                  await reconnect();
-                } else {
-                  const data = JSON.parse(msg.data) as IConnections;
-                  next(null, (old = initConnData) => {
-                    const oldConn = old.activeConnections;
-                    const maxLen = data.connections?.length;
-                    const activeConns: IConnectionsItem[] = [];
-                    const rest = (data.connections || []).filter((each) => {
-                      const index = oldConn.findIndex((o) => o.id === each.id);
-                      if (index >= 0 && index < maxLen) {
-                        const old = oldConn[index];
-                        each.curUpload = each.upload - old.upload;
-                        each.curDownload = each.download - old.download;
-                        activeConns[index] = each;
-                        return false;
-                      }
-                      return true;
-                    });
-                    for (let i = 0; i < maxLen; ++i) {
-                      if (!activeConns[i] && rest.length > 0) {
-                        activeConns[i] = rest.shift()!;
-                        activeConns[i].curUpload = 0;
-                        activeConns[i].curDownload = 0;
-                      }
-                    }
-                    const currentClosedConns = oldConn.filter((each) => {
-                      const index = activeConns.findIndex(
-                        (o) => o.id === each.id,
-                      );
-                      return index < 0;
-                    });
-                    let closedConns =
-                      old.closedConnections.concat(currentClosedConns);
-                    if (closedConns.length > 500) {
-                      closedConns = closedConns.slice(-MAX_CLOSED_CONNS_NUM);
-                    }
-                    return {
-                      uploadTotal: data.uploadTotal,
-                      downloadTotal: data.downloadTotal,
-                      activeConnections: activeConns,
-                      closedConnections: closedConns,
-                    };
-                  });
-                }
-              }
-            });
-          })
-          .catch((_) => {
-            if (!ws.current) {
-              timeoutRef.current = setTimeout(async () => await connect(), 500);
-            }
-          });
-
-      if (
-        wsFirstConnection.current ||
-        (ws.current && !wsFirstConnection.current)
-      ) {
-        wsFirstConnection.current = false;
-        if (ws.current) {
-          ws.current.close();
-          ws.current = null;
-        }
-        connect();
-      }
-
-      return () => {
-        if (timeoutRef.current) {
-          clearTimeout(timeoutRef.current);
-          timeoutRef.current = null;
-        }
-        ws.current?.close();
-        ws.current = null;
-      };
-    },
-    {
+  const { response, refresh, subscriptionCacheKey } =
+    useMihomoWsSubscription<ConnectionMonitorData>({
+      storageKey: "mihomo_connection_date",
+      buildSubscriptKey: (date) => `getClashConnection-${date}`,
       fallbackData: initConnData,
-      keepPreviousData: true,
-    },
-  );
+      connect: () => MihomoWebSocket.connect_connections(),
+      setupHandlers: ({ next, scheduleReconnect }) => ({
+        handleMessage: (data) => {
+          if (data.startsWith("Websocket error")) {
+            next(data);
+            void scheduleReconnect();
+            return;
+          }
 
-  useEffect(() => {
-    mutate(`$sub$${subscriptKey}`);
-  }, [date, subscriptKey]);
-
-  const refreshGetClashConnection = () => {
-    setDate(Date.now());
-  };
+          try {
+            const parsed = JSON.parse(data) as IConnections;
+            next(null, (old = initConnData) =>
+              mergeConnectionSnapshot(parsed, old),
+            );
+          } catch (error) {
+            next(error);
+          }
+        },
+      }),
+    });
 
   const clearClosedConnections = () => {
-    mutate(`$sub$${subscriptKey}`, {
+    if (!subscriptionCacheKey) return;
+    mutate(subscriptionCacheKey, {
       uploadTotal: response.data?.uploadTotal ?? 0,
       downloadTotal: response.data?.downloadTotal ?? 0,
       activeConnections: response.data?.activeConnections ?? [],
@@ -148,5 +111,9 @@ export const useConnectionData = () => {
     });
   };
 
-  return { response, refreshGetClashConnection, clearClosedConnections };
+  return {
+    response,
+    refreshGetClashConnection: refresh,
+    clearClosedConnections,
+  };
 };

src/hooks/use-connection-setting.ts

@@ -0,0 +1,13 @@
+import { useLocalStorage } from "foxact/use-local-storage";
+
+const defaultConnectionSetting: IConnectionSetting = { layout: "table" };
+
+export const useConnectionSetting = () =>
+  useLocalStorage<IConnectionSetting>(
+    "connections-setting",
+    defaultConnectionSetting,
+    {
+      serializer: JSON.stringify,
+      deserializer: JSON.parse,
+    },
+  );

src/hooks/use-log-data.ts

@@ -1,142 +1,113 @@
 import dayjs from "dayjs";
-import { useLocalStorage } from "foxact/use-local-storage";
 import { useEffect, useRef } from "react";
 import { mutate } from "swr";
-import useSWRSubscription from "swr/subscription";
-import { MihomoWebSocket } from "tauri-plugin-mihomo-api";
+import { MihomoWebSocket, type LogLevel } from "tauri-plugin-mihomo-api";
 
 import { getClashLogs } from "@/services/cmds";
-import { useClashLog } from "@/services/states";
+
+import { useClashLog } from "./use-clash-log";
+import { useMihomoWsSubscription } from "./use-mihomo-ws-subscription";
 
 const MAX_LOG_NUM = 1000;
+const FLUSH_DELAY_MS = 50;
+type LogType = ILogItem["type"];
+
+const DEFAULT_LOG_TYPES: LogType[] = ["debug", "info", "warning", "error"];
+const LOG_LEVEL_FILTERS: Record<LogLevel, LogType[]> = {
+  debug: DEFAULT_LOG_TYPES,
+  info: ["info", "warning", "error"],
+  warning: ["warning", "error"],
+  error: ["error"],
+  silent: [],
+};
+
+const clampLogs = (logs: ILogItem[]): ILogItem[] =>
+  logs.length > MAX_LOG_NUM ? logs.slice(-MAX_LOG_NUM) : logs;
+
+const filterLogsByLevel = (
+  logs: ILogItem[],
+  allowedTypes: LogType[],
+): ILogItem[] => {
+  if (allowedTypes.length === 0) return [];
+  if (allowedTypes.length === DEFAULT_LOG_TYPES.length) return logs;
+  return logs.filter((log) => allowedTypes.includes(log.type));
+};
+
+const appendLogs = (
+  current: ILogItem[] | undefined,
+  incoming: ILogItem[],
+): ILogItem[] => clampLogs([...(current ?? []), ...incoming]);
 
 export const useLogData = () => {
   const [clashLog] = useClashLog();
   const enableLog = clashLog.enable;
   const logLevel = clashLog.logLevel;
+  const allowedTypes = LOG_LEVEL_FILTERS[logLevel] ?? DEFAULT_LOG_TYPES;
 
-  const [date, setDate] = useLocalStorage("mihomo_logs_date", Date.now());
-  const subscriptKey = enableLog ? `getClashLog-${date}` : null;
+  const { response, refresh, subscriptionCacheKey } = useMihomoWsSubscription<
+    ILogItem[]
+  >({
+    storageKey: "mihomo_logs_date",
+    buildSubscriptKey: (date) => (enableLog ? `getClashLog-${date}` : null),
+    fallbackData: [],
+    keepPreviousData: true,
+    connect: () => MihomoWebSocket.connect_logs(logLevel),
+    setupHandlers: ({ next, scheduleReconnect, isMounted }) => {
+      let flushTimer: ReturnType<typeof setTimeout> | null = null;
+      const buffer: ILogItem[] = [];
 
-  const ws = useRef<MihomoWebSocket | null>(null);
-  const wsFirstConnection = useRef<boolean>(true);
-  const timeoutRef = useRef<ReturnType<typeof setTimeout>>(null);
-
-  const response = useSWRSubscription<ILogItem[], any, string | null>(
-    subscriptKey,
-    (_key, { next }) => {
-      const reconnect = async () => {
-        await ws.current?.close();
-        ws.current = null;
-        timeoutRef.current = setTimeout(async () => await connect(), 500);
+      const clearFlushTimer = () => {
+        if (flushTimer) {
+          clearTimeout(flushTimer);
+          flushTimer = null;
+        }
       };
 
-      const connect = () =>
-        MihomoWebSocket.connect_logs(logLevel)
-          .then(async (ws_) => {
-            ws.current = ws_;
-            if (timeoutRef.current) clearTimeout(timeoutRef.current);
-
-            const logs = await getClashLogs();
-            let filterLogs: ILogItem[] = [];
-            switch (logLevel) {
-              case "debug":
-                filterLogs = logs.filter((i) =>
-                  ["debug", "info", "warning", "error"].includes(i.type),
-                );
-                break;
-              case "info":
-                filterLogs = logs.filter((i) =>
-                  ["info", "warning", "error"].includes(i.type),
-                );
-                break;
-              case "warning":
-                filterLogs = logs.filter((i) =>
-                  ["warning", "error"].includes(i.type),
-                );
-                break;
-              case "error":
-                filterLogs = logs.filter((i) => i.type === "error");
-                break;
-              case "silent":
-                filterLogs = [];
-                break;
-              default:
-                filterLogs = logs;
-                break;
-            }
-            next(null, filterLogs);
-
-            const buffer: ILogItem[] = [];
-            let flushTimer: ReturnType<typeof setTimeout> | null = null;
-            const flush = () => {
-              if (buffer.length > 0) {
-                next(null, (l) => {
-                  let newList = [...(l ?? []), ...buffer.splice(0)];
-                  if (newList.length > MAX_LOG_NUM) {
-                    newList = newList.slice(
-                      -Math.min(MAX_LOG_NUM, newList.length),
-                    );
-                  }
-                  return newList;
-                });
-              }
-              flushTimer = null;
-            };
-            ws_.addListener(async (msg) => {
-              if (msg.type === "Text") {
-                if (msg.data.startsWith("Websocket error")) {
-                  next(msg.data);
-                  await reconnect();
-                } else {
-                  const data = JSON.parse(msg.data) as ILogItem;
-                  data.time = dayjs().format("MM-DD HH:mm:ss");
-                  buffer.push(data);
-
-                  // flush data
-                  if (!flushTimer) {
-                    flushTimer = setTimeout(flush, 50);
-                  }
-                }
-              }
-            });
-          })
-          .catch((_) => {
-            if (!ws.current) {
-              timeoutRef.current = setTimeout(async () => await connect(), 500);
-            }
-          });
-
-      if (
-        wsFirstConnection.current ||
-        (ws.current && !wsFirstConnection.current)
-      ) {
-        wsFirstConnection.current = false;
-        if (ws.current) {
-          ws.current.close();
-          ws.current = null;
+      const flush = () => {
+        if (!buffer.length || !isMounted()) {
+          flushTimer = null;
+          return;
         }
-        connect();
-      }
+        const pendingLogs = buffer.splice(0, buffer.length);
+        next(null, (current) => appendLogs(current, pendingLogs));
+        flushTimer = null;
+      };
 
-      return () => {
-        if (timeoutRef.current) {
-          clearTimeout(timeoutRef.current);
-          timeoutRef.current = null;
-        }
-        ws.current?.close();
-        ws.current = null;
+      return {
+        handleMessage: (data) => {
+          if (data.startsWith("Websocket error")) {
+            next(data);
+            void scheduleReconnect();
+            return;
+          }
+
+          try {
+            const parsed = JSON.parse(data) as ILogItem;
+            if (
+              allowedTypes.length > 0 &&
+              !allowedTypes.includes(parsed.type)
+            ) {
+              return;
+            }
+            parsed.time = dayjs().format("MM-DD HH:mm:ss");
+            buffer.push(parsed);
+            if (!flushTimer) {
+              flushTimer = setTimeout(flush, FLUSH_DELAY_MS);
+            }
+          } catch (error) {
+            next(error);
+          }
+        },
+        async onConnected() {
+          const logs = await getClashLogs();
+          if (isMounted()) {
+            next(null, clampLogs(filterLogsByLevel(logs, allowedTypes)));
+          }
+        },
+        cleanup: clearFlushTimer,
       };
     },
-    {
-      fallbackData: [],
-      keepPreviousData: true,
-    },
-  );
-
-  useEffect(() => {
-    mutate(`$sub$${subscriptKey}`);
-  }, [date, subscriptKey]);
+  });
 
   const previousLogLevel = useRef<string | undefined>(undefined);
 
@@ -151,15 +122,16 @@ export const useLogData = () => {
     }
 
     previousLogLevel.current = logLevel;
-    ws.current?.close();
-    setDate(Date.now());
-  }, [logLevel, setDate]);
+    refresh();
+  }, [logLevel, refresh]);
 
   const refreshGetClashLog = (clear = false) => {
     if (clear) {
-      mutate(`$sub$${subscriptKey}`, []);
+      if (subscriptionCacheKey) {
+        mutate(subscriptionCacheKey, []);
+      }
     } else {
-      setDate(Date.now());
+      refresh();
     }
   };
 

src/hooks/use-memory-data.ts

@@ -1,89 +1,37 @@
-import { useLocalStorage } from "foxact/use-local-storage";
-import { useEffect, useRef } from "react";
-import { mutate } from "swr";
-import useSWRSubscription from "swr/subscription";
 import { MihomoWebSocket } from "tauri-plugin-mihomo-api";
 
+import { useMihomoWsSubscription } from "./use-mihomo-ws-subscription";
+
 export interface IMemoryUsageItem {
   inuse: number;
   oslimit?: number;
 }
 
+const FALLBACK_MEMORY_USAGE: IMemoryUsageItem = { inuse: 0 };
+
 export const useMemoryData = () => {
-  const [date, setDate] = useLocalStorage("mihomo_memory_date", Date.now());
-  const subscriptKey = `getClashMemory-${date}`;
-
-  const ws = useRef<MihomoWebSocket | null>(null);
-  const wsFirstConnection = useRef<boolean>(true);
-  const timeoutRef = useRef<ReturnType<typeof setTimeout>>(null);
-
-  const response = useSWRSubscription<IMemoryUsageItem, any, string | null>(
-    subscriptKey,
-    (_key, { next }) => {
-      const reconnect = async () => {
-        await ws.current?.close();
-        ws.current = null;
-        timeoutRef.current = setTimeout(async () => await connect(), 500);
-      };
-
-      const connect = () =>
-        MihomoWebSocket.connect_memory()
-          .then((ws_) => {
-            ws.current = ws_;
-            if (timeoutRef.current) clearTimeout(timeoutRef.current);
-
-            ws_.addListener(async (msg) => {
-              if (msg.type === "Text") {
-                if (msg.data.startsWith("Websocket error")) {
-                  next(msg.data, { inuse: 0 });
-                  await reconnect();
-                } else {
-                  const data = JSON.parse(msg.data) as IMemoryUsageItem;
-                  next(null, data);
-                }
-              }
-            });
-          })
-          .catch((_) => {
-            if (!ws.current) {
-              timeoutRef.current = setTimeout(async () => await connect(), 500);
-            }
-          });
-
-      if (
-        wsFirstConnection.current ||
-        (ws.current && !wsFirstConnection.current)
-      ) {
-        wsFirstConnection.current = false;
-        if (ws.current) {
-          ws.current.close();
-          ws.current = null;
+  const { response, refresh } = useMihomoWsSubscription<IMemoryUsageItem>({
+    storageKey: "mihomo_memory_date",
+    buildSubscriptKey: (date) => `getClashMemory-${date}`,
+    fallbackData: FALLBACK_MEMORY_USAGE,
+    connect: () => MihomoWebSocket.connect_memory(),
+    setupHandlers: ({ next, scheduleReconnect }) => ({
+      handleMessage: (data) => {
+        if (data.startsWith("Websocket error")) {
+          next(data, FALLBACK_MEMORY_USAGE);
+          void scheduleReconnect();
+          return;
         }
-        connect();
-      }
 
-      return () => {
-        if (timeoutRef.current) {
-          clearTimeout(timeoutRef.current);
-          timeoutRef.current = null;
+        try {
+          const parsed = JSON.parse(data) as IMemoryUsageItem;
+          next(null, parsed);
+        } catch (error) {
+          next(error, FALLBACK_MEMORY_USAGE);
         }
-        ws.current?.close();
-        ws.current = null;
-      };
-    },
-    {
-      fallbackData: { inuse: 0 },
-      keepPreviousData: true,
-    },
-  );
+      },
+    }),
+  });
 
-  useEffect(() => {
-    mutate(`$sub$${subscriptKey}`);
-  }, [date, subscriptKey]);
-
-  const refreshGetClashMemory = () => {
-    setDate(Date.now());
-  };
-
-  return { response, refreshGetClashMemory };
+  return { response, refreshGetClashMemory: refresh };
 };

src/hooks/use-mihomo-ws-subscription.ts

@@ -0,0 +1,156 @@
+import { useLocalStorage } from "foxact/use-local-storage";
+import { useCallback, useEffect, useRef } from "react";
+import { mutate, type MutatorCallback } from "swr";
+import useSWRSubscription from "swr/subscription";
+import { type Message, type MihomoWebSocket } from "tauri-plugin-mihomo-api";
+
+export const RECONNECT_DELAY_MS = 500;
+
+type NextFn<T> = (error?: any, data?: T | MutatorCallback<T>) => void;
+
+interface HandlerContext<T> {
+  next: NextFn<T>;
+  scheduleReconnect: () => Promise<void>;
+  isMounted: () => boolean;
+}
+
+interface HandlerResult {
+  handleMessage: (data: string) => void;
+  onConnected?: (ws: MihomoWebSocket) => Promise<void> | void;
+  cleanup?: () => void;
+}
+
+interface UseMihomoWsSubscriptionOptions<T> {
+  storageKey: string;
+  buildSubscriptKey: (date: number) => string | null;
+  fallbackData: T;
+  connect: () => Promise<MihomoWebSocket>;
+  keepPreviousData?: boolean;
+  setupHandlers: (ctx: HandlerContext<T>) => HandlerResult;
+}
+
+export const useMihomoWsSubscription = <T>(
+  options: UseMihomoWsSubscriptionOptions<T>,
+) => {
+  const {
+    storageKey,
+    buildSubscriptKey,
+    fallbackData,
+    connect,
+    keepPreviousData = true,
+    setupHandlers,
+  } = options;
+
+  const [date, setDate] = useLocalStorage(storageKey, Date.now());
+  const subscriptKey = buildSubscriptKey(date);
+  const subscriptionCacheKey = subscriptKey ? `$sub$${subscriptKey}` : null;
+
+  const wsRef = useRef<MihomoWebSocket | null>(null);
+  const wsFirstConnection = useRef<boolean>(true);
+  const timeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null);
+
+  const response = useSWRSubscription<T, any, string | null>(
+    subscriptKey,
+    (_key, { next }) => {
+      let isMounted = true;
+
+      const clearReconnectTimer = () => {
+        if (timeoutRef.current) {
+          clearTimeout(timeoutRef.current);
+          timeoutRef.current = null;
+        }
+      };
+
+      const closeSocket = async () => {
+        if (wsRef.current) {
+          await wsRef.current.close();
+          wsRef.current = null;
+        }
+      };
+
+      const scheduleReconnect = async () => {
+        if (!isMounted) return;
+        clearReconnectTimer();
+        await closeSocket();
+        if (!isMounted) return;
+        timeoutRef.current = setTimeout(connectWs, RECONNECT_DELAY_MS);
+      };
+
+      const {
+        handleMessage: handleTextMessage,
+        onConnected,
+        cleanup,
+      } = setupHandlers({
+        next,
+        scheduleReconnect,
+        isMounted: () => isMounted,
+      });
+
+      const cleanupAll = () => {
+        clearReconnectTimer();
+        cleanup?.();
+        void closeSocket();
+      };
+
+      const handleMessage = (msg: Message) => {
+        if (msg.type !== "Text") return;
+        handleTextMessage(msg.data);
+      };
+
+      async function connectWs() {
+        try {
+          const ws_ = await connect();
+          if (!isMounted) {
+            await ws_.close();
+            return;
+          }
+
+          wsRef.current = ws_;
+          clearReconnectTimer();
+
+          if (onConnected) {
+            await onConnected(ws_);
+            if (!isMounted) {
+              await ws_.close();
+              return;
+            }
+          }
+
+          ws_.addListener(handleMessage);
+        } catch (ignoreError) {
+          if (!wsRef.current && isMounted) {
+            timeoutRef.current = setTimeout(connectWs, RECONNECT_DELAY_MS);
+          }
+        }
+      }
+
+      if (wsFirstConnection.current || !wsRef.current) {
+        wsFirstConnection.current = false;
+        cleanupAll();
+        void connectWs();
+      }
+
+      return () => {
+        isMounted = false;
+        wsFirstConnection.current = true;
+        cleanupAll();
+      };
+    },
+    {
+      fallbackData,
+      keepPreviousData,
+    },
+  );
+
+  useEffect(() => {
+    if (subscriptionCacheKey) {
+      mutate(subscriptionCacheKey);
+    }
+  }, [subscriptionCacheKey]);
+
+  const refresh = useCallback(() => {
+    setDate(Date.now());
+  }, [setDate]);
+
+  return { response, refresh, subscriptionCacheKey, wsRef };
+};

src/hooks/use-traffic-data.ts

@@ -1,95 +1,37 @@
-import { useLocalStorage } from "foxact/use-local-storage";
-import { useEffect, useRef } from "react";
-import { mutate } from "swr";
-import useSWRSubscription from "swr/subscription";
 import { MihomoWebSocket, Traffic } from "tauri-plugin-mihomo-api";
 
-import { TrafficRef } from "@/components/layout/traffic-graph";
-
+import { useMihomoWsSubscription } from "./use-mihomo-ws-subscription";
 import { useTrafficMonitorEnhanced } from "./use-traffic-monitor";
 
-export const useTrafficData = () => {
-  const [date, setDate] = useLocalStorage("mihomo_traffic_date", Date.now());
-  const subscriptKey = `getClashTraffic-${date}`;
+const FALLBACK_TRAFFIC: Traffic = { up: 0, down: 0 };
 
-  const trafficRef = useRef<TrafficRef>(null);
+export const useTrafficData = () => {
   const {
     graphData: { appendData },
   } = useTrafficMonitorEnhanced({ subscribe: false });
-  const ws = useRef<MihomoWebSocket | null>(null);
-  const wsFirstConnection = useRef<boolean>(true);
-  const timeoutRef = useRef<ReturnType<typeof setTimeout>>(null);
-
-  const response = useSWRSubscription<ITrafficItem, any, string | null>(
-    subscriptKey,
-    (_key, { next }) => {
-      const reconnect = async () => {
-        await ws.current?.close();
-        ws.current = null;
-        timeoutRef.current = setTimeout(async () => await connect(), 500);
-      };
-
-      const connect = async () => {
-        MihomoWebSocket.connect_traffic()
-          .then(async (ws_) => {
-            ws.current = ws_;
-            if (timeoutRef.current) clearTimeout(timeoutRef.current);
-
-            ws_.addListener(async (msg) => {
-              if (msg.type === "Text") {
-                if (msg.data.startsWith("Websocket error")) {
-                  next(msg.data, { up: 0, down: 0 });
-                  await reconnect();
-                } else {
-                  const data = JSON.parse(msg.data) as Traffic;
-                  trafficRef.current?.appendData(data);
-                  appendData(data);
-                  next(null, data);
-                }
-              }
-            });
-          })
-          .catch((_) => {
-            if (!ws.current) {
-              timeoutRef.current = setTimeout(async () => await connect(), 500);
-            }
-          });
-      };
-
-      if (
-        wsFirstConnection.current ||
-        (ws.current && !wsFirstConnection.current)
-      ) {
-        wsFirstConnection.current = false;
-        if (ws.current) {
-          ws.current.close();
-          ws.current = null;
+  const { response, refresh } = useMihomoWsSubscription<ITrafficItem>({
+    storageKey: "mihomo_traffic_date",
+    buildSubscriptKey: (date) => `getClashTraffic-${date}`,
+    fallbackData: FALLBACK_TRAFFIC,
+    connect: () => MihomoWebSocket.connect_traffic(),
+    setupHandlers: ({ next, scheduleReconnect }) => ({
+      handleMessage: (data) => {
+        if (data.startsWith("Websocket error")) {
+          next(data, FALLBACK_TRAFFIC);
+          void scheduleReconnect();
+          return;
         }
-        connect();
-      }
 
-      return () => {
-        if (timeoutRef.current) {
-          clearTimeout(timeoutRef.current);
-          timeoutRef.current = null;
+        try {
+          const parsed = JSON.parse(data) as Traffic;
+          appendData(parsed);
+          next(null, parsed);
+        } catch (error) {
+          next(error, FALLBACK_TRAFFIC);
         }
-        ws.current?.close();
-        ws.current = null;
-      };
-    },
-    {
-      fallbackData: { up: 0, down: 0 },
-      keepPreviousData: true,
-    },
-  );
+      },
+    }),
+  });
 
-  useEffect(() => {
-    mutate(`$sub$${subscriptKey}`);
-  }, [date, subscriptKey]);
-
-  const refreshGetClashTraffic = () => {
-    setDate(Date.now());
-  };
-
-  return { response, refreshGetClashTraffic };
+  return { response, refreshGetClashTraffic: refresh };
 };

src/pages/connections.tsx

@@ -28,7 +28,7 @@ import {
 import { ConnectionItem } from "@/components/connection/connection-item";
 import { ConnectionTable } from "@/components/connection/connection-table";
 import { useConnectionData } from "@/hooks/use-connection-data";
-import { useConnectionSetting } from "@/services/states";
+import { useConnectionSetting } from "@/hooks/use-connection-setting";
 import parseTraffic from "@/utils/parse-traffic";
 
 type OrderFunc = (list: IConnectionsItem[]) => IConnectionsItem[];

src/pages/logs.tsx

@@ -13,8 +13,8 @@ import { BaseSearchBox } from "@/components/base/base-search-box";
 import { SearchState } from "@/components/base/base-search-box";
 import { BaseStyledSelect } from "@/components/base/base-styled-select";
 import LogItem from "@/components/log/log-item";
+import { useClashLog } from "@/hooks/use-clash-log";
 import { useLogData } from "@/hooks/use-log-data";
-import { LogFilter, useClashLog } from "@/services/states";
 
 const LogPage = () => {
   const { t } = useTranslation();

src/services/states.ts

@@ -1,50 +1,9 @@
 import { createContextState } from "foxact/create-context-state";
-import { useLocalStorage } from "foxact/use-local-storage";
-import { LogLevel } from "tauri-plugin-mihomo-api";
 
 const [ThemeModeProvider, useThemeMode, useSetThemeMode] = createContextState<
   "light" | "dark"
 >();
 
-export type LogFilter = "all" | "debug" | "info" | "warn" | "err";
-export type LogOrder = "asc" | "desc";
-
-interface IClashLog {
-  enable: boolean;
-  logLevel: LogLevel;
-  logFilter: LogFilter;
-  logOrder: LogOrder;
-}
-const defaultClashLog: IClashLog = {
-  enable: true,
-  logLevel: "info",
-  logFilter: "all",
-  logOrder: "asc",
-};
-export const useClashLog = () =>
-  useLocalStorage<IClashLog>("clash-log", defaultClashLog, {
-    serializer: JSON.stringify,
-    deserializer: JSON.parse,
-  });
-
-// export const useEnableLog = () => useLocalStorage("enable-log", false);
-
-interface IConnectionSetting {
-  layout: "table" | "list";
-}
-
-const defaultConnectionSetting: IConnectionSetting = { layout: "table" };
-
-export const useConnectionSetting = () =>
-  useLocalStorage<IConnectionSetting>(
-    "connections-setting",
-    defaultConnectionSetting,
-    {
-      serializer: JSON.stringify,
-      deserializer: JSON.parse,
-    },
-  );
-
 // save the state of each profile item loading
 const [LoadingCacheProvider, useLoadingCache, useSetLoadingCache] =
   createContextState<Record<string, boolean>>({});

src/types/types.d.ts

@@ -197,6 +197,17 @@ interface ILogItem {
   payload: string;
 }
 
+type LogLevel = import("tauri-plugin-mihomo-api").LogLevel;
+type LogFilter = "all" | "debug" | "info" | "warn" | "err";
+type LogOrder = "asc" | "desc";
+
+interface IClashLog {
+  enable: boolean;
+  logLevel: LogLevel;
+  logFilter: LogFilter;
+  logOrder: LogOrder;
+}
+
 interface IConnectionsItem {
   id: string;
   metadata: {
@@ -227,6 +238,10 @@ interface IConnections {
   connections: IConnectionsItem[];
 }
 
+interface IConnectionSetting {
+  layout: "table" | "list";
+}
+
 /**
  * Some interface for command
  */