Springboot 整合 WebSocket 实现聊天室功能

前言

WebSocket概述：
在日常的web应用开发中，常见的是前端向后端发起请求，有些时候会涉及到前后端互发消息，这时候就用到了WebSocket。

一、WebSocket原理

WebSocket是一种在单个TCP连接上进行全双工通信的协议。它通过一个简单的握手过程来建立连接，然后在连接上进行双向数据传输。与传统的HTTP请求不同，WebSocket连接一旦建立，就可以在客户端和服务器之间保持打开状态，直到被任何一方关闭。

核心特点包括：

全双工通信：客户端和服务器可以同时发送和接收消息。
持久连接：一旦建立连接，就可以持续进行数据交换，无需像HTTP那样频繁地建立新的连接。
低延迟：由于连接是持久的，数据可以几乎实时地发送和接收。
轻量级协议：WebSocket协议的头部信息非常简单，减少了数据传输的开销。

二、Spring Boot集成WebSocket

代码结构：
在这里插入图片描述

2.1. 引入依赖

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-websocket</artifactId>
</dependency>

2.2 配置类WebSocketConfig

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

/**
 * WebSocket配置类。
 * 用于启用Spring WebSocket支持，通过@Bean注解注册ServerEndpointExporter，
 * 从而允许使用@ServerEndpoint注解定义WebSocket端点。
 */
@Configuration
public class WebSocketConfig {

    /**
     * 注册ServerEndpointExporter Bean。
     * ServerEndpointExporter是Spring提供的一个工具类，
     * 它会扫描并注册所有使用@ServerEndpoint注解的类为WebSocket端点。
     *
     * @return ServerEndpointExporter实例
     */
    @Bean
    public ServerEndpointExporter serverEndpointExporter() {
        return new ServerEndpointExporter();
    }

    /**
     * 通信文本消息和二进制缓存区大小
     * 避免报文过大时，Websocket 1009 错误
     */

    @Bean
    public ServletServerContainerFactoryBean createWebSocketContainer() {
        ServletServerContainerFactoryBean container = new ServletServerContainerFactoryBean();
        // 文本/二进制消息最大缓冲区（10MB）
        container.setMaxTextMessageBufferSize(1024 * 1024 * 10);
        container.setMaxBinaryMessageBufferSize(1024 * 1024 * 10);
        // 最大会话空闲超时时间（1小时）
        container.setMaxSessionIdleTimeout(60 * 60 * 1000L);
        return container;
    }
}

2.3 WebSocketServer 类

WebSocketServer 类实现了 WebSocket 服务端的功能。负责处理 WebSocket 连接的建立、关闭、消息接收和发送等操作。

import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.Objects;
import java.util.concurrent.CopyOnWriteArraySet;

/**
 * WebSocketServer 类实现了 WebSocket 服务端的功能。
 * 它负责处理 WebSocket 连接的建立、关闭、消息接收和发送等操作。
 */
@Component
@Slf4j
@ServerEndpoint("/api/websocket/{sid}")
public class WebSocketServer {

    // 静态变量，用于记录当前在线连接数
    private static int onlineCount = 0;

    // 存储所有连接的 WebSocketServer 实例
    @Getter
    private static final CopyOnWriteArraySet<WebSocketServer> webSocketSet = new CopyOnWriteArraySet<>();

    // 当前连接的会话对象
    private Session session;

    // 客户端唯一标识符
    private String sid = "";

    /**
     * 连接建立成功时调用的方法。
     *
     * @param session 当前连接的会话对象
     * @param sid 客户端唯一标识符
     */
    @OnOpen
    public void onOpen(Session session, @PathParam("sid") String sid) {
        this.session = session;
        webSocketSet.add(this);     // 将当前实例加入集合
        this.sid = sid;
        addOnlineCount();           // 在线数加1
        try {
            sendMessage("WebSocket 连接成功");  // 发送连接成功的消息
            log.info("有新窗口开始监听:{},当前在线人数为:{}", sid, getOnlineCount());
        } catch (IOException e) {
            log.error("websocket IO Exception");
        }
    }

    /**
     * 连接关闭时调用的方法。
     */
    @OnClose
    public void onClose() {
        webSocketSet.remove(this);  // 从集合中移除当前实例
        subOnlineCount();           // 在线数减1
        log.info("释放的sid为：{}", sid);
        log.info("有一个连接关闭！当前在线人数为{}", getOnlineCount());
    }

    /**
     * 接收到客户端消息时调用的方法。
     *
     * @param message 客户端发送的消息
     * @param session 当前连接的会话对象
     */
    @OnMessage
    public void onMessage(String message, Session session) {
        log.info("收到来自窗口{}的信息:{}", sid, message);
        // 群发消息
        for (WebSocketServer item : webSocketSet) {
            if (Objects.equals(item.sid, this.sid)) {
                continue;
            }
            sendMessageToClient(item, message);
        }
    }

    /**
     * 实现服务器主动推送消息的方法，并统一处理异常。
     *
     * @param client 要推送的客户端实例
     * @param message 要推送的消息
     */
    private void sendMessageToClient(WebSocketServer client, String message) {
        try {
            client.sendMessage(message);
        } catch (IOException e) {
            log.error("向客户端 {} 发送消息时出错: {}", client.sid, message, e);
        }
    }

    /**
     * 群发自定义消息给指定的客户端。
     *
     * @param message 要发送的消息
     * @param sid 客户端唯一标识符，为 null 时发送给所有客户端
     * @throws IOException 如果发送消息时发生 I/O 错误
     */
    public static void sendInfo(String message, @PathParam("sid") String sid) throws IOException {
        log.info("推送消息到窗口" + sid + "，推送内容:" + message);

        for (WebSocketServer item : webSocketSet) {
            try {
                if (sid == null) {
                    // 如果 sid 为 null，则发送给所有客户端
                    item.sendMessage(message);
                } else if (item.sid.equals(sid)) {
                    // 如果 sid 匹配，则只发送给该客户端
                    item.sendMessage(message);
                }
            } catch (IOException e) {
                log.error("向客户端 {} 发送消息时出错: {}", item.sid, message, e);
            }
        }
    }

    /**
     * 发生错误时调用的方法。
     *
     * @param session 当前连接的会话对象
     * @param error 发生的错误
     */
    @OnError
    public void onError(Session session, Throwable error) {
        log.error("发生错误");
        error.printStackTrace();
    }

    /**
     * 实现服务器主动推送消息的方法。
     *
     * @param message 要推送的消息
     * @throws IOException 如果发送消息时发生 I/O 错误
     */
    public void sendMessage(String message) throws IOException {
        this.session.getBasicRemote().sendText(message);
    }

    /**
     * 获取当前在线连接数。
     *
     * @return 当前在线连接数
     */
    public static synchronized int getOnlineCount() {
        return onlineCount;
    }

    /**
     * 增加在线连接数。
     */
    public static synchronized void addOnlineCount() {
        WebSocketServer.onlineCount++;
    }

    /**
     * 减少在线连接数。
     */
    public static synchronized void subOnlineCount() {
        WebSocketServer.onlineCount--;
    }
}

2.4 前端代码 index.html

<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <meta charset="UTF-8">
    <title>WebSocket 聊天室</title>
    <script src="https://autherp.jd.com/js/jquery.js"></script>
    <style>
        .time {
            font-size: 0.8em;
            display: block;
            margin-bottom: 5px;
        }

        .user-msg {
            background-color: #90EE90;
            padding: 8px;
            border-radius: 8px;
            max-width: 70%;
            display: inline-block;
        }

        .system-msg {
            background-color: #D3D3D3;
            padding: 8px;
            border-radius: 8px;
            max-width: 70%;
            display: inline-block;
            font-size: 0.9em;
        }

        #message-box {
            height: 300px;
            overflow-y: auto;
            margin-bottom: 10px;
            border: 1px solid #ccc;
            padding: 10px;
            border-radius: 4px;
        }

        .message-right {
            text-align: right;
            margin: 10px 0;
        }

        .message-left {
            text-align: left;
            margin: 10px 0;
        }

        .message-center {
            text-align: center;
            margin: 10px 0;
        }

        .container {
            max-width: 600px;
            margin: 20px auto;
            padding: 20px;
            border: 1px solid #ddd;
            border-radius: 8px;
            box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
        }

        .input-group {
            display: flex;
            gap: 10px;
        }

        .input-group input[type="text"] {
            flex-grow: 1;
            padding: 8px;
            border: 1px solid #ccc;
            border-radius: 4px;
        }

        .btn-primary, .btn-danger {
            padding: 8px 16px;
            border: none;
            border-radius: 4px;
            cursor: pointer;
        }

        .btn-primary {
            background-color: #007BFF;
            color: white;
        }

        .btn-primary:hover {
            background-color: #0056b3;
        }

        .btn-danger {
            background-color: #DC3545;
            color: white;
        }

        .btn-danger:hover {
            background-color: #c82333;
        }

        /* 添加新样式让标题和按钮居中 */
        .container h2,
        .container .btn-danger {
            text-align: center;
            display: block;
            margin-left: auto;
            margin-right: auto;
        }

        /* 为按钮添加一些外边距，使其看起来更美观 */
        .container .btn-danger {
            margin-top: 10px;
        }
    </style>
</head>
<body>
<div class="container">
    <h2>WebSocket 聊天室</h2>
    <!-- 添加显示 sid 的元素 -->
    <p id="sid-display">当前用户 SID: <span id="sid-value"></span></p>

    <!-- 消息显示区域 -->
    <div id="message-box"></div>

    <!-- 输入框与发送按钮 -->
    <div class="input-group">
        <input type="text" id="text" placeholder="请输入消息..." />
        <button class="btn-primary" onclick="send()">发送</button>
    </div>

    <hr/>

    <!-- 关闭连接按钮 -->
    <button class="btn-danger" onclick="closeWebSocket()">关闭 WebSocket 连接</button>
</div>

<script type="text/javascript">
    let websocket = '';

    // 获取当前页面 URL 中的 sid 参数或随机生成一个
    function getSid() {
        const urlParams = new URLSearchParams(window.location.search);
        return urlParams.get('sid') || Math.floor(1000 + Math.random() * 9000); // 4位数字
    }

    const sid = getSid();
    const wsUrl = `ws://127.0.0.1:9999/api/websocket/${sid}`;

    // 页面加载完成后更新 sid 显示
    window.onload = function() {
        document.getElementById('sid-value').textContent = sid;
    };

    // 初始化 WebSocket
    if ('WebSocket' in window) {
        websocket = new WebSocket(wsUrl);
    } else {
        alert('当前浏览器不支持 WebSocket');
    }

    // 连接成功
    websocket.onopen = function () {
        console.log('WebSocket 连接成功');
    };

    // 接收消息
    websocket.onmessage = function (event) {
        addMessage(event.data);
    };

    // 错误处理
    websocket.onerror = function () {
        console.log('WebSocket 连接发生错误');
    };

    // 关闭连接
    websocket.onclose = function () {
        this.closeWebSocket();
        console.log('WebSocket 连接已关闭');
    };

    // 页面关闭前断开连接
    window.onbeforeunload = function () {
        this.closeWebSocket();
    };

    // 发送消息
    function send() {
        let message = document.getElementById('text').value.trim();
        if (!message) return;

        if (websocket && websocket.readyState === WebSocket.OPEN) {
            websocket.send(`{"msg":"${message}","sid":"${sid}", "time": "${new Date().toLocaleTimeString()}"}`);
            addMessage(`{"msg":"${message}","sid":"${sid}", "time": "${new Date().toLocaleTimeString()}"}`);
            document.getElementById('text').value = '';
        } else {
            addMessage("WebSocket 连接未建立，请稍后再试。");
        }
    }

    //关闭WebSocket连接
    function closeWebSocket() {
        if (websocket) {
            websocket.close();
        }
    }

    // 添加消息到聊天区
    function addMessage(content) {
        let msgBox = document.getElementById('message-box');
        const time = new Date().toLocaleTimeString();
        const div = document.createElement('div');

        let messageData;
        let isSystemMessage = false;

        try {
            messageData = JSON.parse(content);
        } catch (e) {
            isSystemMessage = true;
        }

        if (isSystemMessage) {
            div.className = 'message-center';
            div.innerHTML = `
            <span class="time">${time}</span>
            <span class="system-msg"> ${content}</span>
            `;
        } else {
            const isMyMessage = String(messageData.sid) === String(sid);
            div.className = isMyMessage ? 'message-right' : 'message-left';
            div.innerHTML = `
            <span class="time">${messageData.time}</span>
            <span class="${isMyMessage ? 'user-msg' : 'system-msg'}">
                ${isMyMessage ? '' : '用户#' + messageData.sid + ':'} ${messageData.msg}
            </span>
            `;
        }

        msgBox.appendChild(div);
        msgBox.scrollTop = msgBox.scrollHeight;
    }

</script>
</body>
</html>

2.5 Controller访问首页

import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.RequestMapping;

@Controller
public class TestController {

    @RequestMapping("/")
    public String index(){
        return "index.html";
    }
}