Spring Boot 整合原生 WebSocket

1、简述

1.1 什么是 WebSocket ?

WebSocket是一种在单个TCP连接上进行全双工、双向通信的网络协议，它实现了浏览器与服务器之间的持久性实时对话。与传统的HTTP请求-响应模式不同，WebSocket在初次握手后，双方可以随时主动发送数据。

1.2 WebSocket特点

特点维度	具体说明	带来的优势与考量
全双工通信	在单个TCP连接上，客户端与服务器可以同时、独立地发送和接收数据。	实现了真正的实时双向对话，是构建聊天室、在线游戏、协同编辑等强交互应用的理想协议。
独立协议	通过一次HTTP Upgrade握手，将连接升级为独立的 ws或wss（加密）协议进行通信。	握手后切换至专为实时交互设计的轻量级帧协议，摆脱了HTTP的无状态和头部冗余，通信效率更高。
持久化连接	连接一旦建立便会一直保持，直到显式关闭，期间可进行无数次数据交换。	彻底消除了HTTP的重复连接建立开销（如TCP握手、SSL协商），极大降低了延迟，并减少了服务器资源消耗。
低开销数据帧	数据传输采用自定义的帧格式，每个消息的协议头部额外开销极小。	特别适合高频、小数据量的通信场景（如心跳、实时坐标更新），网络利用率极高，延迟极低。
支持二进制与文本	可在同一连接中无缝传输文本数据和二进制数据（如ArrayBuffer, Blob）	功能全面强大，既能高效传输JSON等文本，也能直接处理文件、图片、音视频流，无需像SSE那样进行Base64编解码。
灵活但需自建的API	浏览器提供原生的 WebSocket API，包含 onopen, onmessage, send, close 等核心事件和方法。	为双向实时通信提供了标准化的底层构建块，控制力强。

2、Spring Boot 整合 WebSocket

2.1 引入依赖

bash 复制代码

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

2.2 WebSocket 配置类

java 复制代码

@Configuration
@EnableWebSocket
@Slf4j
public class WebSocketConfig implements WebSocketConfigurer {

    @Resource
    private MyWebSocketHandler myWebSocketHandler;

    @Resource
    private WebSocketAuthInterceptor authInterceptor;

    @Value("${server.port}")
    private String port;

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry
                // 注册 WebSocket 处理器并设置连接路劲
                .addHandler(myWebSocketHandler, "/ceshi")
                // 添加鉴权拦截器
                .addInterceptors(authInterceptor)
                // 允许跨域（生产需限制）
                .setAllowedOrigins("*");

        log.info("======================================");
        log.info("🚀 WebSocket 模块启动完成");
        log.info("📡 WebSocket 地址：ws://localhost:{}/ceshi", port);
        log.info("🔐 鉴权方式：HandshakeInterceptor (userId)");
        log.info("🌐 跨域策略：AllowedOrigins = *");
        log.info("======================================");
    }
}

2.3 WebSocket 处理器

java 复制代码

@Component
@Slf4j
public class MyWebSocketHandler extends TextWebSocketHandler {

    @Resource
    private WebSocketSessionManager sessionManager;

    /**
     * 建立连接成功后回调
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) {

        // 从拦截器中获取 userId
        String userId = (String) session.getAttributes().get("userId");

        // 保存连接
        sessionManager.add(userId, session);

        log.info("用户上线：{}", userId);
    }

    /**
     * 收到客户端消息时回调
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {

        String payload = message.getPayload();

        // 心跳消息
        if ("ping".equals(payload)) {
            session.sendMessage(new TextMessage("pong"));
            return;
        }

        log.info("收到客户端消息：{}", payload);
    }

    /**
     * 连接关闭时回调
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {

        String userId = (String) session.getAttributes().get("userId");

        // 移除连接
        sessionManager.remove(userId);

        log.info("用户离线：{}", userId);
    }
}

2.4 WebSocket 握手拦截器

java 复制代码

@Component
public class WebSocketAuthInterceptor implements HandshakeInterceptor {

    @Override
    public boolean beforeHandshake(@NonNull ServerHttpRequest request,
                                   @NonNull ServerHttpResponse response,
                                   @NonNull WebSocketHandler handler,
                                   @NonNull Map<String, Object> attributes) {

        // 获取 ws://xxx/ws?userId=1001 中的参数
        String query = request.getURI().getQuery();
        if (query == null) {
            return false;
        }

        // 简单解析 userId（生产建议用 token / JWT）
        String userId = Arrays.stream(query.split("&"))
                .filter(s -> s.startsWith("userId="))
                .map(s -> s.replace("userId=", ""))
                .findFirst()
                .orElse(null);

        // 校验失败，拒绝建立连接
        if (userId == null) {
            return false;
        }

        // 将 userId 保存到 session 属性中
        attributes.put("userId", userId);
        return true;
    }

    @Override
    public void afterHandshake(@NonNull ServerHttpRequest request,
                               @NonNull ServerHttpResponse response,
                               @NonNull WebSocketHandler handler,
                               Exception ex) {
        // 握手完成后无需处理
    }
}

2.5 WebSocketMessage 会话管理实体类

java 复制代码

@Data
@AllArgsConstructor
@NoArgsConstructor
public class WebSocketMessage {

    /**
     * 消息类型
     * HEARTBEAT：心跳
     * NOTICE：通知类消息
     * DATA：业务数据
     * ERROR：错误提示
     */
    private String type;

    /**
     * 业务标识
     * 如：workOrder、device、alarm
     * 前端可根据 biz 分模块处理
     */
    private String biz;

    /**
     * 实际业务数据
     * 可放 Map / DTO / VO
     */
    private Object data;
}

2.6 WebSocket 连接会话管理

java 复制代码

@Component
public class WebSocketSessionManager {

    /**
     * 保存在线 WebSocket 连接
     * key：userId
     * value：WebSocketSession
     */
    public static final Map<String, WebSocketSession> SESSION_MAP = new ConcurrentHashMap<>();

    /**
     * 新用户上线时调用
     */
    public void add(String userId, WebSocketSession session) {
        SESSION_MAP.put(userId, session);
    }

    /**
     * 用户下线 / 断开连接时调用
     */
    public void remove(String userId) {
        SESSION_MAP.remove(userId);
    }

    /**
     * 获取指定用户的连接
     */
    public WebSocketSession get(String userId) {
        return SESSION_MAP.get(userId);
    }

    /**
     * 获取所有在线连接
     */
    public Collection<WebSocketSession> all() {
        return SESSION_MAP.values();
    }

    /**
     * 给指定用户推送消息
     */
    public void sendToUser(String userId, String msg) throws IOException {
        WebSocketSession session = SESSION_MAP.get(userId);
        if (session != null && session.isOpen()) {
            session.sendMessage(new TextMessage(msg));
        }
    }

    /**
     * 广播消息（慎用）
     */
    public void sendToAll(String msg) throws IOException {
        for (WebSocketSession session : SESSION_MAP.values()) {
            if (session.isOpen()) {
                session.sendMessage(new TextMessage(msg));
            }
        }
    }
}

2.7 WebSocket 服务类

java 复制代码

public interface WebSocketService {

    /**
     * 获取指定的连接信息
     */
    WebSocketSession get(String userId);

    /**
     * 获取所有在线连接
     */
    Collection<WebSocketSession> all();

    /**
     * 给指定用户推送消息
     */
    Boolean sendToUser(String userId, String msg);

    /**
     * 广播消息（慎用）
     */
    Integer sendToAll(String msg);
}

2.7 WebSocket 服务实现类

java 复制代码

@Service
@Slf4j
public class WebSocketServiceImpl implements WebSocketService {

    @Resource
    WebSocketSessionManager webSocketSessionManager;

    @Override
    public WebSocketSession get(String userId) {
        return webSocketSessionManager.get(userId);
    }

    @Override
    public Collection<WebSocketSession> all() {
        return webSocketSessionManager.all();
    }

    @Override
    public Boolean sendToUser(String userId, String msg) {
        WebSocketSession session = SESSION_MAP.get(userId);
        if (session == null || !session.isOpen()) {
            return false;
        }

        try {
            webSocketSessionManager.sendToUser(userId, msg);
            return true;
        } catch (IOException e) {
            log.error("WebSocket 推送失败 userId={}", userId, e);
            return false;
        }
    }

    @Override
    public Integer sendToAll(String msg) {
        int count = 0;
        for (WebSocketSession session : SESSION_MAP.values()) {
            if (session.isOpen()) {
                try {
                    webSocketSessionManager.sendToAll(msg);
                    count++;
                } catch (IOException e) {
                    log.warn("广播发送失败", e);
                }
            }
        }
        return count;
    }
}

2.8 Controller 对外接口

java 复制代码

@Api(tags = "WebSocket API")
@RestController
@RequestMapping("/websocket")
public class WebSocketController {

    @Resource
    private WebSocketService webSocketService;

    @ApiOperation("获取指定用户的 WebSocket 连接信息")
    @GetMapping("/connection")
    public AjaxResult getConnection(@Parameter(name = "userId", description = "用户ID", required = true, example = "1001") String userId) {

        return AjaxResult.success(webSocketService.get(userId));
    }

    @ApiOperation("获取所有在线 WebSocket 连接")
    @GetMapping("/connections")
    public AjaxResult getAllConnections() {
        return AjaxResult.success(webSocketService.all());
    }

    @ApiOperation("给指定用户推送 WebSocket 消息")
    @PostMapping("/send/user")
    public AjaxResult sendToUser(
            @Parameter(name = "userId", description = "用户ID", required = true, example = "1001") String userId,
            @Parameter(name = "msg", description = "消息内容", required = true, example = "这是一条通知") String msg) {

        return webSocketService.sendToUser(userId, msg) ? AjaxResult.success("发送成功") : AjaxResult.error("用户不在线或发送失败");
    }

    @ApiOperation("广播 WebSocket 消息（慎用）")
    @PostMapping("/send/all")
    public AjaxResult sendToAll(
            @Parameter(name = "msg", description = "消息内容", required = true, example = "系统广播消息") String msg) {

        int count = webSocketService.sendToAll(msg);
        return AjaxResult.success("广播完成，发送给 " + count + " 个在线用户");
    }
}

3、测试

我这里使用的是 ApiPost 软件进行测试，记住新建 WebSocket 进行连接。

连接成功如图：

进行心跳监测，发送消息：ping，会收到回复信息：pong 。

新建普通接口进行消息发送：

此时连接会显示收到的消息：