如何使用Java搭建WebSocket服务端 Java构建实时推送系统

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

import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;

@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(new MyTextWebSocketHandler(), "/ws")
                .setAllowedOrigins("*"); // 允许所有来源，生产环境应限制特定域名
    }
}

import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

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

@Component
public class MyTextWebSocketHandler extends TextWebSocketHandler {

    // 用来存放所有在线的WebSocketSession
    private static final CopyOnWriteArraySet<WebSocketSession> sessions = new CopyOnWriteArraySet<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        sessions.add(session);
        System.out.println("新连接建立: " + session.getId() + ", 当前在线人数: " + sessions.size());
        session.sendMessage(new TextMessage("欢迎连接到WebSocket服务！"));
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        System.out.println("收到来自 " + session.getId() + " 的消息: " + message.getPayload());
        // 可以将消息广播给所有客户端
        broadcastMessage(message.getPayload());
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        sessions.remove(session);
        System.out.println("连接关闭: " + session.getId() + ", 状态: " + status.getCode() + ", 原因: " + status.getReason() + ", 当前在线人数: " + sessions.size());
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        System.err.println("传输错误: " + session.getId() + ", 错误信息: " + exception.getMessage());
        sessions.remove(session);
    }

    // 广播消息给所有在线客户端
    public static void broadcastMessage(String message) {
        TextMessage textMessage = new TextMessage(message);
        for (WebSocketSession session : sessions) {
            try {
                if (session.isOpen()) {
                    session.sendMessage(textMessage);
                }
            } catch (IOException e) {
                System.err.println("发送消息到 " + session.getId() + " 失败: " + e.getMessage());
                // 考虑移除发送失败的session，或者在afterConnectionClosed中统一处理
            }
        }
    }
}

public static void broadcastMessage(String message) {
    TextMessage textMessage = new TextMessage(message);
    for (WebSocketSession session : sessions) { // 遍历所有session
        try {
            if (session.isOpen()) { // 确保session仍然是打开状态
                session.sendMessage(textMessage);
            }
        } catch (IOException e) {
            // 发送失败通常意味着客户端已断开，但连接状态更新滞后，
            // 可以在这里记录日志，或者更健壮地处理（比如尝试移除这个session）
            System.err.println("发送消息到 " + session.getId() + " 失败: " + e.getMessage());
        }
    }
}

import java.util.concurrent.ConcurrentHashMap;
import java.util.Map;

// ... 其他import

@Component
public class MyTextWebSocketHandler extends TextWebSocketHandler {

    // 存储所有在线session，key可以是用户ID，value是WebSocketSession
    private static final Map<String, WebSocketSession> userSessions = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        // 假设你可以从session中获取用户ID，例如通过认证信息或者URL参数
        String userId = extractUserIdFromSession(session); // 这是一个你需要实现的方法
        if (userId != null) {
            userSessions.put(userId, session);
            System.out.println("用户 " + userId + " 连接建立: " + session.getId());
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        // 在连接关闭时移除对应的session
        userSessions.entrySet().removeIf(entry -> entry.getValue().equals(session));
        System.out.println("连接关闭: " + session.getId() + ", 当前在线用户数: " + userSessions.size());
    }

    // 定向发送消息给某个用户
    public static void sendMessageToUser(String userId, String message) {
        WebSocketSession session = userSessions.get(userId);
        if (session != null && session.isOpen()) {
            try {
                session.sendMessage(new TextMessage(message));
                System.out.println("向用户 " + userId + " 发送消息: " + message);
            } catch (IOException e) {
                System.err.println("发送消息到用户 " + userId + " 失败: " + e.getMessage());
            }
        } else {
            System.out.println("用户 " + userId + " 不在线或会话已关闭。");
        }
    }

    // 示例：如何从session中获取用户ID（实际情况可能通过认证令牌等）
    private String extractUserIdFromSession(WebSocketSession session) {
        // 实际应用中，这部分会涉及安全认证，例如从Spring Security的Principal中获取
        // 或者从WebSocket握手时的HTTP头、URL参数中解析
        // 简单示例：假设客户端连接时URL是 /ws?userId=xxx
        String query = session.getUri() != null ? session.getUri().getQuery() : null;
        if (query != null && query.contains("userId=")) {
            return query.split("userId=")[1].split("&")[0];
        }
        return null;
    }
}