基于 Spring Boot 的 WebSocket 实时通信服务实战

2. 消息处理器与会话管理

ChatWebSocketHandler 是核心组件。这里不仅处理单条消息，还实现了两个关键逻辑：分片重组和用户隔离路由。

• 分片处理：当消息超过缓冲区限制时，WebSocket 会将其拆分为多个片段。我们需要缓存这些片段直到收到最后一块再合并。
• 会话映射：使用 ConcurrentHashMap 维护 userId 到 Set<WebSocketSession> 的映射，确保消息能准确发送给特定用户或广播给同一用户的所有设备。

package com.example.websocket;

import java.io.ByteArrayOutputStream;
import java.net.URI;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

import org.springframework.web.socket.BinaryMessage;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.AbstractWebSocketHandler;

import com.fasterxml.jackson.databind.ObjectMapper;

public class ChatWebSocketHandler extends AbstractWebSocketHandler {

    private final ConcurrentHashMap<String, Set<WebSocketSession>> userSessions = new ConcurrentHashMap<>();
    private static final ObjectMapper MAPPER = new ObjectMapper();
    private final ConcurrentHashMap<String, StringBuilder> textFragments = new ConcurrentHashMap<>();
    private final ConcurrentHashMap<String, ByteArrayOutputStream> binaryFragments = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        // 验证用户 ID 的有效性，从 URL 查询参数中提取
        String uid = resolveUserId(session);
        if (uid == null || uid.isEmpty()) {
            session.close(CloseStatus.BAD_DATA);
            return;
        }
        session.getAttributes().put("userId", uid);
        // 多会话管理：一个用户可能对应多个连接
        userSessions.computeIfAbsent(uid, k -> ConcurrentHashMap.newKeySet()).add(session);
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String id = session.getId();
        // 分片处理：如果不是最后一片，先缓存
        if (!message.isLast()) {
            textFragments.computeIfAbsent(id, k -> new StringBuilder()).append(message.getPayload());
            return;
        }
        // 合并所有片段
        StringBuilder sb = textFragments.remove(id);
        String payload = sb != null ? sb.append(message.getPayload()).toString() : message.getPayload();
        routePayload(session, payload);
    }

    @Override
    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) throws Exception {
        String id = session.getId();
        ByteBuffer buf = message.getPayload();
        byte[] chunk = new byte[buf.remaining()];
        buf.get(chunk);
        ByteArrayOutputStream acc = binaryFragments.computeIfAbsent(id, k -> new ByteArrayOutputStream());
        acc.write(chunk);
        if (message.isLast()) {
            byte[] all = acc.toByteArray();
            binaryFragments.remove(id);
            String payload = new String(all, StandardCharsets.UTF_8);
            routePayload(session, payload);
        }
    }

    @Override
    public boolean supportsPartialMessages() {
        return true;
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        Object v = session.getAttributes().get("userId");
        if (v == null) return;
        String uid = String.valueOf(v);
        Set<WebSocketSession> set = userSessions.get(uid);
        if (set != null) {
            set.remove(session);
            if (set.isEmpty()) userSessions.remove(uid);
        }
    }

    /** 从 WebSocket 连接的 URL 查询参数中提取用户 ID */
    private String resolveUserId(WebSocketSession session) {
        URI uri = session.getUri();
        if (uri == null) return null;
        String q = uri.getQuery();
        if (q == null || q.isEmpty()) return null;
        String[] parts = q.split("&");
        for (String p : parts) {
            int i = p.indexOf('=');
            if (i > 0) {
                String k = p.substring(0, i);
                String val = p.substring(i + 1);
                if ("userId".equals(k)) return val;
            }
        }
        return null;
    }

    private void routePayload(WebSocketSession session, String payload) throws Exception {
        Object v = session.getAttributes().get("userId");
        if (v == null) return;
        String fromUid = String.valueOf(v);

        // 解析消息对象，注意 Message 类需自行定义 DTO
        Message message = new Message();
        message.setFromUserId(fromUid);
        try {
            Message receivedMsg = MAPPER.readValue(payload, Message.class);
            message.setToUserId(receivedMsg.getToUserId());
            message.setContent(receivedMsg.getContent());
            message.setType(receivedMsg.getType());
        } catch (Exception e) {
            // 如果解析失败，将整个 payload 作为 content
            message.setContent(payload);
        }

        String toUid = message.getToUserId();
        boolean isP2P = toUid != null && !toUid.isEmpty();
        Set<WebSocketSession> targets;

        if (isP2P) {
            targets = userSessions.get(toUid);
        } else {
            targets = userSessions.get(fromUid);
        }

        String outStr = MAPPER.writeValueAsString(message);
        TextMessage msg = new TextMessage(outStr);

        if (targets == null || targets.isEmpty()) {
            if (session.isOpen()) {
                session.sendMessage(msg);
            }
            return;
        }

        for (WebSocketSession s : targets) {
            if (s.isOpen()) {
                s.sendMessage(msg);
            }
        }
        // 如果是点对点且自己不在目标列表，也发给自己确认
        if (isP2P && session.isOpen()) {
            session.sendMessage(msg);
        }
    }
}

运行与测试

启动应用后，可以通过浏览器访问静态页面进行测试。连接地址通常为 ws://localhost:8080/ws?userId=xxx。通过不同的 userId 可以模拟多个用户同时在线并进行私聊或群聊测试。

在实际部署时，请注意防火墙策略是否放行了 WebSocket 端口，并确保 Nginx 或其他反向代理层正确配置了 Upgrade 头以支持长连接。