SpringBoot与Vue实现WebSocket心跳机制

思路

前端每隔一段时间向后端发送一次字符串ping-${uid}，后端收到后返回pong响应

后端

后端配置

package org.example.config;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.server.standard.ServerEndpointExporter;

@Configuration
public class WebSocketConfig {
    @Bean
    public ServerEndpointExporter serverEndpointExporter() {
        return new ServerEndpointExporter();
    }
}

package org.example.controller;

import cn.hutool.json.JSONObject;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

@Component
@ServerEndpoint("/ws/{userId}")
public class WebSocketController {

    private Session session;
    // 当前连接的用户ID
    private String userId;
    // 存储所有的 WebSocket 连接
    private static final CopyOnWriteArraySet<WebSocketController> webSockets = new CopyOnWriteArraySet<>();
    // 存储用户ID和对应的会话，方便查找和管理
    private static final ConcurrentHashMap<String, Session> sessionPool = new ConcurrentHashMap<>();
    private static final ObjectMapper objectMapper = new ObjectMapper();


    @Autowired
    public void setService(MessageService messageService) {
        WebSocketController.messageService = messageService;
    }

    // 当新的 WebSocket 连接建立时调用此方法
    @OnOpen
    public void onOpen(Session session, @PathParam("userId") String userId) {
        this.session = session;
        this.userId = userId;
        webSockets.add(this);
        sessionPool.put(userId, session);
    }

    // 当 WebSocket 连接关闭时调用此方法
    @OnClose
    public void onClose() {
        webSockets.remove(this);
        sessionPool.remove(this.userId);
    }

    // 当收到消息时调用此方法
    @OnMessage
    public void onMessage(String message) {
        try {
            // 心跳机制
            if (message.startsWith("ping")) {
                // 对心跳消息进行解析，获取用户ID
                String[] parts = message.split("-");
                String uid = parts[1]; // message格式：ping-uid
                // 获取对应的会话
                Session session = sessionPool.get(uid);
                if (session!= null && session.isOpen()) {
                    // 发送心跳回复
                    session.getBasicRemote().sendText("pong");
                }
                return;
            }
            // 将接收到的消息反序列化为 Message 对象
            Message msg = objectMapper.readValue(message, Message.class);
            // msg 是对象，message 是源文本
            handleMessageSend(msg, message);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }


    // 当发生错误时调用此方法
    @OnError
    public void onError(Session session, Throwable error) {
        System.err.println("Error: " + error.getMessage());
    }


    // 处理消息发送的逻辑 messageService已另行实现
    private void handleMessageSend(Message msg, String msgText) {
        // 获取消息中的群组ID
        int groupId = msg.getGroupId();
        // 根据群组ID获取群组信息
        Group group = messageService.getGroupById(groupId);
        // 获取消息中的用户ID
        int Uid = msg.getUid();
        String nickname = messageService.getNicknameByUid(Uid);
        // 单人-------------
        if (group == null) {
            sendOneMessage(String.valueOf(groupId), msgText);
            return;
        }
        // 群聊-------------
        List<String> userIds = messageService.getUserIdsByGroupId(groupId);
        if (group.getMulti() == 1) { // Group chat
            // 发送者的用户ID
            int uid = msg.getUid();
            // 根据用户ID获取用户昵称、头像等信息
            nickname = userService.getUserByUid(uid).getNickname();
            // 在源文本消息中添加发送者信息
            JSONObject jsonObject = new JSONObject(msgText);
            jsonObject.append("nickname", nickname);
            String updatedMsgText = jsonObject.toString();
            // 给群组成员发送消息
            for (String userId : userIds) {
                if (userId.equals(this.userId)) continue;
                sendOneMessage(userId, updatedMsgText);
            }
        }
    }



    // 发送消息给单个用户
    public void sendOneMessage(String userId, String message) {
        Session session = sessionPool.get(userId);
        if (session!= null && session.isOpen()) {
            // 异步发送消息
            session.getAsyncRemote().sendText(message);
        }
    }


}

前端

import { watch } from 'vue';
import { defineStore, storeToRefs } from 'pinia';
import notificationSound from '@/assets/notification.mp3';

const wsurl = import.meta.env.VITE_WS_URL;
const myUid = localStorage.getItem('uid')

export const useWsStore = defineStore('ws', {
    state: () => ({
        ws: null,

    }),

    actions: {
        async wsConnect(uid) {
            return new Promise((resolve, reject) => {
                this.ws = new WebSocket(`${wsurl}/ws/${uid}`);
                // 提示音
                const audio = new Audio(notificationSound);
                audio.volume = 0;
                audio.play().then(() => {// 静音播放
                    audio.volume = 1; // 恢复音量
                }).catch((e) => {
                    console.log(e);
                });

                this.ws.onopen = () => {
                    this.startHeartBeat(); // 启动心跳
                    resolve();
                };
                // ws连接关闭
                this.ws.onclose = () => {
                    this.ws = new WebSocket(`${wsurl}/ws/${uid}`);
                };
                this.ws.onerror = (error) => {
                    reject(error);
                };
                this.ws.onmessage = (e) => {
                    let newMsg;
                    try {
                        newMsg = JSON.parse(e.data);
                    } catch {
                        newMsg = e.data; //pong消息，心跳回应
                        return;
                    }
                    audio.play().catch((e) => { console.log(e) });
                };
            });
        },
        

        startHeartBeat() {
            this.heartBeatTimer = setInterval(() => {
                if (this.ws.readyState === WebSocket.OPEN) {
                    this.ws.send(`ping-${myUid}`); // 发送ping消息
                } else {
                    this.reconnectWebSocket();
                }
            }, 20000); // 每20秒发送一次心跳
        },
        stopHeartBeat() {
            if (this.heartBeatTimer) {
                clearInterval(this.heartBeatTimer);
                this.heartBeatTimer = null;
            }
        },
        reconnectWebSocket() {
            this.stopHeartBeat(); // 停止心跳
            if (this.ws) {
                this.ws.close();
            }
            this.wsConnect(myUid); // 重新连接WebSocket
        },

        wsSend(data) {
            if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
                this.wsConnect(data.uid).then(() => {
                    // 确保发送的信息是字符串
                    this.ws.send(typeof data === "string" ? data : JSON.stringify(data));
                }).catch(error => {
                    console.error(error);
                });
            } else {
                this.ws.send(typeof data === "string" ? data : JSON.stringify(data));
            }

        },


        disconnectWs() {
            this.stopHeartBeat();
            if (this.ws) {
                this.ws.close();
                this.ws = null;
            }
        },
    },
});

前端使用

import { storeToRefs } from 'pinia';
import { useWsStore } from '@/store/wsStore'

const wsStore = useWsStore()
const useWs = storeToRefs(wsStore)

onMounted(() => {
    if (!useWs.ws) wsStore.connectWs(uid)
})

await wsStore.wsSend(newMsg);

const sendMsg = async (newMsg) => {
    await wsStore.wsSend(newMsg);
}