From 79ea96ab8a9f95db02afcca13599e5f4f99fe9af Mon Sep 17 00:00:00 2001
From: "Ray.Hao" <1490493387@qq.com>
Date: Thu, 13 Feb 2025 01:01:02 +0800
Subject: [PATCH] =?UTF-8?q?refactor:=20:recycle:=20websocket=20=E5=8A=9F?=
 =?UTF-8?q?=E8=83=BD=E9=87=8D=E6=9E=84?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .env.development                              |  2 +-
 src/hooks/useStomp.ts                         | 91 +++++++++----------
 .../NavBar/components/Notification.vue        | 37 +++++---
 src/views/demo/websocket.vue                  | 16 ++--
 4 files changed, 77 insertions(+), 69 deletions(-)

diff --git a/.env.development b/.env.development
index 10bd8d14..55b854a1 100644
--- a/.env.development
+++ b/.env.development
@@ -9,7 +9,7 @@ VITE_APP_API_URL=https://api.youlai.tech # 线上
 # VITE_APP_API_URL=http://localhost:8989    # 本地
 
 # WebSocket 端点（不配置则关闭），线上 ws://api.youlai.tech/ws ，本地 ws://localhost:8989/ws
-VITE_APP_WS_ENDPOINT=ws://localhost:8989/ws
+VITE_APP_WS_ENDPOINT=
 
 # 启用 Mock 服务
 VITE_MOCK_DEV_SERVER=false
diff --git a/src/hooks/useStomp.ts b/src/hooks/useStomp.ts
index b13d394a..af88f57a 100644
--- a/src/hooks/useStomp.ts
+++ b/src/hooks/useStomp.ts
@@ -1,4 +1,3 @@
-import { ref, watch, onMounted } from "vue";
 import { Client, IMessage, StompSubscription } from "@stomp/stompjs";
 import { getAccessToken } from "@/utils/auth";
 
@@ -14,26 +13,18 @@ export interface UseStompOptions {
 }
 
 export function useStomp(options: UseStompOptions = {}) {
-  // 默认值：brokerURL 从环境变量中获取，token 从 getAccessToken() 获取
   const defaultBrokerURL = import.meta.env.VITE_APP_WS_ENDPOINT || "";
   const defaultToken = getAccessToken();
 
-  // 将 brokerURL 定义为响应式 ref，便于动态修改
   const brokerURL = ref(options.brokerURL ?? defaultBrokerURL);
   const token = options.token ?? defaultToken;
 
-  // 连接状态标记
   const isConnected = ref(false);
-  // 存储所有订阅
   const subscriptions = new Map<string, StompSubscription>();
 
-  // 用于保存 STOMP 客户端的实例
-  let client: Client | null = null;
+  const client = ref<Client | null>(null);
 
-  /**
-   * 初始化 STOMP 客户端
-   * 只有在 brokerURL 非空时才会初始化客户端
-   */
+  // 初始化 STOMP 客户端
   const initializeClient = () => {
     if (!brokerURL.value) {
       console.warn(
@@ -42,8 +33,8 @@ export function useStomp(options: UseStompOptions = {}) {
       return;
     }
 
-    if (!client) {
-      client = new Client({
+    if (!client.value) {
+      client.value = new Client({
         brokerURL: brokerURL.value,
         reconnectDelay: options.reconnectDelay ?? 5000,
         debug: options.debug ? (msg) => console.log("[STOMP]", msg) : () => {},
@@ -54,17 +45,17 @@ export function useStomp(options: UseStompOptions = {}) {
         heartbeatOutgoing: 4000,
       });
 
-      client.onConnect = (frame) => {
+      client.value.onConnect = (frame) => {
         isConnected.value = true;
         console.log("STOMP connected", frame);
       };
 
-      client.onStompError = (frame) => {
+      client.value.onStompError = (frame) => {
         console.error("Broker reported error: " + frame.headers["message"]);
         console.error("Additional details: " + frame.body);
       };
 
-      client.onWebSocketClose = (evt) => {
+      client.value.onWebSocketClose = (evt) => {
         isConnected.value = false;
         console.warn("WebSocket closed", evt);
       };
@@ -75,9 +66,8 @@ export function useStomp(options: UseStompOptions = {}) {
   watch(brokerURL, (newURL, oldURL) => {
     if (newURL !== oldURL) {
       console.log(`brokerURL changed from ${oldURL} to ${newURL}`);
-      // 断开当前连接，重新激活客户端
-      if (client && client.connected) {
-        client.deactivate();
+      if (client.value && client.value.connected) {
+        client.value.deactivate();
       }
       brokerURL.value = newURL;
       initializeClient(); // 重新初始化客户端
@@ -89,36 +79,39 @@ export function useStomp(options: UseStompOptions = {}) {
     initializeClient();
   });
 
-  /**
-   * 激活连接（如果已经连接或正在激活则直接返回）
-   */
+  // 激活连接（如果已经连接或正在激活则直接返回）
   const connect = () => {
-    if (client && (client.connected || client.active)) {
+    if (client.value && (client.value.connected || client.value.active)) {
       console.log("Already connected or connecting, skipping connect() call.");
       return;
     }
-    client?.activate();
-  };
-
-  /**
-   * 订阅指定主题
-   * @param destination 目标主题地址
-   * @param callback 接收到消息时的回调函数
-   * @returns 返回订阅 id，用于后续取消订阅
-   */
-  const subscribe = (destination: string, callback: (message: IMessage) => void): string => {
-    if (client) {
-      const subscription = client.subscribe(destination, callback);
-      subscriptions.set(subscription.id, subscription);
-      return subscription.id;
+    if (client.value) {
+      client.value.activate();
+    } else {
+      console.warn("Client is not initialized.");
     }
-    return "";
   };
 
-  /**
-   * 取消指定订阅
-   * @param subscriptionId 要取消的订阅 id
-   */
+  // 订阅指定主题，连接成功后自动订阅
+  const subscribe = (destination: string, callback: (message: IMessage) => void): string => {
+    if (!client.value) {
+      console.error("STOMP client is not initialized.");
+      return "";
+    }
+
+    // 如果还没有连接，就先激活连接
+    if (!isConnected.value) {
+      console.log("Not connected yet. Connecting...");
+      connect();
+    }
+
+    // 连接成功后订阅主题
+    const subscription = client.value.subscribe(destination, callback);
+    subscriptions.set(subscription.id, subscription);
+    return subscription.id;
+  };
+
+  // 取消指定订阅
   const unsubscribe = (subscriptionId: string) => {
     const subscription = subscriptions.get(subscriptionId);
     if (subscription) {
@@ -127,15 +120,17 @@ export function useStomp(options: UseStompOptions = {}) {
     }
   };
 
-  /**
-   * 主动断开连接（如果未连接则不执行）
-   */
+  // 主动断开连接（如果未连接则不执行）
   const disconnect = () => {
-    if (client && !(client.connected || client.active)) {
+    if (client.value && !(client.value.connected || client.value.active)) {
       console.log("Already disconnected, skipping disconnect() call.");
       return;
     }
-    client?.deactivate();
+    if (client.value) {
+      client.value.deactivate();
+    } else {
+      console.warn("Client is not initialized.");
+    }
     isConnected.value = false;
   };
 
@@ -143,7 +138,7 @@ export function useStomp(options: UseStompOptions = {}) {
     client,
     isConnected,
     connect,
-    subscribe,
+    subscribe, // 订阅函数放到这里
     unsubscribe,
     disconnect,
     brokerURL, // 暴露 brokerURL 以便组件中动态修改
diff --git a/src/layout/components/NavBar/components/Notification.vue b/src/layout/components/NavBar/components/Notification.vue
index 06d55b2c..d29c0a3f 100644
--- a/src/layout/components/NavBar/components/Notification.vue
+++ b/src/layout/components/NavBar/components/Notification.vue
@@ -168,7 +168,6 @@
 <script setup lang="ts">
 import NoticeAPI, { NoticePageVO } from "@/api/system/notice";
 import router from "@/router";
-import { useStomp } from "@/hooks/useStomp";
 
 const activeTab = ref("notice");
 const notices = ref<NoticePageVO[]>([]);
@@ -176,20 +175,16 @@ const messages = ref<any[]>([]);
 const tasks = ref<any[]>([]);
 const noticeDetailRef = ref();
 
-// 初始化 useStomp hook（这里仅用于订阅通知消息），同时调用 connect 建立连接
-const { connect, subscribe, disconnect } = useStomp({
+import { useStomp } from "@/hooks/useStomp";
+
+const { subscribe, disconnect } = useStomp({
   debug: true,
 });
 
-// 获取未读消息列表并连接 WebSocket
-onMounted(() => {
-  NoticeAPI.getMyNoticePage({ pageNum: 1, pageSize: 5, isRead: 0 }).then((data) => {
-    notices.value = data.list;
-  });
-
-  // 建立连接
-  connect();
-
+/**
+ * 订阅通知消息
+ */
+function subscribeNotice() {
   subscribe("/user/queue/message", (message) => {
     console.log("收到消息：", message);
     const data = JSON.parse(message.body);
@@ -210,7 +205,16 @@ onMounted(() => {
       });
     }
   });
-});
+}
+
+/**
+ * 获取我的通知公告
+ */
+function featchMyNotice() {
+  NoticeAPI.getMyNoticePage({ pageNum: 1, pageSize: 5, isRead: 0 }).then((data) => {
+    notices.value = data.list;
+  });
+}
 
 // 阅读通知公告
 function handleReadNotice(id: string) {
@@ -233,8 +237,13 @@ function markAllAsRead() {
   });
 }
 
+// 获取未读消息列表并连接 WebSocket
+onMounted(() => {
+  featchMyNotice();
+  subscribeNotice();
+});
+
 onBeforeUnmount(() => {
-  // 如果需要取消订阅，可以在这里调用 disconnect 或 unsubscribe（本示例直接断开连接）
   disconnect();
 });
 </script>
diff --git a/src/views/demo/websocket.vue b/src/views/demo/websocket.vue
index 873c77ac..932294a7 100644
--- a/src/views/demo/websocket.vue
+++ b/src/views/demo/websocket.vue
@@ -34,6 +34,7 @@
             </el-col>
           </el-row>
         </el-card>
+
         <!-- 广播消息发送部分 -->
         <el-card class="mt-5">
           <el-form label-width="90px">
@@ -45,6 +46,7 @@
             </el-form-item>
           </el-form>
         </el-card>
+
         <!-- 点对点消息发送部分 -->
         <el-card class="mt-5">
           <el-form label-width="90px">
@@ -60,6 +62,7 @@
           </el-form>
         </el-card>
       </el-col>
+
       <!-- 消息接收显示部分 -->
       <el-col :span="12">
         <el-card>
@@ -96,10 +99,11 @@
 
 <script setup lang="ts">
 import { useStomp } from "@/hooks/useStomp";
-import { getAccessToken } from "@/utils/auth"; // 此处可与 hook 内 getToken 保持一致
-import { useUserStoreHook } from "@/store/modules/user";
+import { getAccessToken } from "@/utils/auth"; // 用于获取token
+import { useUserStoreHook } from "@/store/modules/user"; // 获取用户信息
 
 const userStore = useUserStoreHook();
+
 // 用于手动调整 WebSocket 地址
 const socketEndpoint = ref(import.meta.env.VITE_APP_WS_ENDPOINT);
 // 同步连接状态
@@ -175,8 +179,8 @@ function disconnectWebSocket() {
 
 // 发送广播消息
 function sendToAll() {
-  if (client && client.connected) {
-    client.publish({
+  if (client.value && isConnected.value) {
+    client.value.publish({
       destination: "/topic/notice",
       body: topicMessage.value,
     });
@@ -189,8 +193,8 @@ function sendToAll() {
 
 // 发送点对点消息
 function sendToUser() {
-  if (client && client.connected) {
-    client.publish({
+  if (client.value && isConnected.value) {
+    client.value.publish({
       destination: "/app/sendToUser/" + receiver.value,
       body: queneMessage.value,
     });