SpringBoot使用WebSocket

mvn依赖

java 复制代码

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

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

后端配置文件

java 复制代码

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

WS服务

java 复制代码

import jakarta.websocket.OnClose;
import jakarta.websocket.OnError;
import jakarta.websocket.OnMessage;
import jakarta.websocket.OnOpen;
import jakarta.websocket.Session;
import jakarta.websocket.server.PathParam;
import jakarta.websocket.server.ServerEndpoint;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;

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

    private static final ConcurrentHashMap<String, Session> SESSION_MAP = new ConcurrentHashMap<>();

    private String userId;

    @OnOpen
    public void onOpen(Session session, @PathParam("userId") String userId) {
        this.userId = userId;
        SESSION_MAP.put(userId, session);
        log.info("用户连接成功，userId={}, 当前在线人数={}", userId, SESSION_MAP.size());
    }

    @OnClose
    public void onClose() {
        if (userId != null) {
            SESSION_MAP.remove(userId);
            log.info("用户连接关闭，userId={}, 当前在线人数={}", userId, SESSION_MAP.size());
        }
    }

    @OnMessage
    public void onMessage(String message, Session session) {
        log.info("收到用户{}的消息：{}", userId, message);

        String response = "服务端已收到：" + message;

        try {
            session.getBasicRemote().sendText(response);
        } catch (IOException e) {
            log.error("发送消息失败，userId={}", userId, e);
        }
    }

    @OnError
    public void onError(Session session, Throwable error) {
        log.error("WebSocket异常，userId={}", userId, error);
    }

    public static void sendMessageToAll(String message) {
        SESSION_MAP.forEach((uid, session) -> {
            try {
                session.getBasicRemote().sendText(message);
            } catch (IOException e) {
                log.error("广播消息失败，userId={}", uid, e);
            }
        });
    }

    public static void sendMessageToUser(String userId, String message) {
        Session session = SESSION_MAP.get(userId);
        if (session != null) {
            try {
                session.getBasicRemote().sendText(message);
            } catch (IOException e) {
                log.error("发送消息给指定用户失败，userId={}", userId, e);
            }
        }
    }
}

前端测试文件

javascript 复制代码

<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <meta charset="UTF-8">
    <title>WebSocket 测试</title>
</head>
<body>

<h3>WebSocket 测试</h3>

用户ID：<input id="userId" value="user1"/>
<button onclick="connect()">连接</button>
<button onclick="disconnect()">断开</button>

<br/><br/>

发送消息：<input id="message"/>
<button onclick="send()">发送</button>

<br/><br/>

<div>
    <strong>日志：</strong>
    <pre id="log"></pre>
</div>

<script>
    let ws;

    function log(text) {
        document.getElementById('log').innerText += text + '\n';
    }

    function connect() {
        const userId = document.getElementById('userId').value;
        ws = new WebSocket('ws://localhost:8080/ws/' + userId);

        ws.onopen = function () {
            log('连接成功');
        };

        ws.onmessage = function (event) {
            log('收到消息：' + event.data);
        };

        ws.onclose = function () {
            log('连接关闭');
        };

        ws.onerror = function (error) {
            log('连接异常');
            console.error(error);
        };
    }

    function disconnect() {
        if (ws) {
            ws.close();
        }
    }

    function send() {
        const message = document.getElementById('message').value;
        if (ws && ws.readyState === WebSocket.OPEN) {
            ws.send(message);
        } else {
            log('WebSocket 未连接');
        }
    }
</script>

</body>
</html>