SpringBoot 集成WebSocket 的步骤

• 添加依赖（spring-boot-starter-websocket）
• 配置 WebSocket 端点（@EnableWebSocket 或 @EnableWebSocketMessageBroker）
• 实现 WebSocket 处理器（WebSocketHandler 或 TextWebSocketHandler）
• 定义消息模型（如JSON格式的消息体）

安全性与性能优化

• WebSocket 连接鉴权（HandshakeInterceptor 或 Spring Security 集成）
• 心跳机制与连接保活

总结

• WebSocket 在SpringBoot中的最佳实践
• 扩展方向（如结合MQTT、gRPC等协议） 也可结合AI模型相关接入Spring AI 实现流式对话

    <!-- WebSocket 支持 -->
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-websocket</artifactId>
        </dependency>

@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        // 注册独立的 WebSocket 处理器
        registry.addHandler(new WebSocketHandler(), "/ws/chat")
                .addInterceptors(new AuthHandshakeInterceptor())
                .setAllowedOriginPatterns("*");
    }
    
}

/**
 * WebSocket握手拦截器，用于验证用户身份
 */
@Component
@Slf4j
public class AuthHandshakeInterceptor extends HttpSessionHandshakeInterceptor {

    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                   WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        
        log.info("WebSocket握手开始: {}", request.getURI());
//        SysUserVO user = UserLoginHelper.getUser();
//
//        if (user == null) {
//            log.warn("无法从Token中提取用户信息");
//            response.setStatusCode(org.springframework.http.HttpStatus.UNAUTHORIZED);
//            return false;
//        }
        SysUserVO user = new SysUserVO();
        // 将用户信息存储到attributes中，可以在WebSocketHandler中获取
        attributes.put("user", user);
        attributes.put("token", StpUserUtil.getTokenValue());
        attributes.put("handshakeTime", System.currentTimeMillis());

        log.info("用户 {} (ID: {}) WebSocket握手成功", user.getPhone(), user.getId());
        return super.beforeHandshake(request, response, wsHandler, attributes);
    }



    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                               WebSocketHandler wsHandler, Exception exception) {
        if (exception == null) {
            log.info("WebSocket握手完成: {}", request.getURI());
        } else {
            log.error("WebSocket握手失败: {}", request.getURI(), exception);
        }
        super.afterHandshake(request, response, wsHandler, exception);
    }
}

@AllArgsConstructor
@Getter
public enum WebSocketType {
    /**
     * 启动
     */
    START("start"),

    SYSTEM("system"),
    PING("ping"),
    PONG("pong"),
    ERROR("error"),
    ;

    private String type;
}

/**
 * 优化版WebSocket处理器 - JDK21特性支持8万并发
 */
@Slf4j
@Component
public class WebSocketHandler extends TextWebSocketHandler {

    private final ConcurrentMap<String, String> sessionUserMap = new ConcurrentHashMap<>();
    // 新结构（多会话支持）
    private static final ConcurrentHashMap<String, CopyOnWriteArraySet<WebSocketSession>> userSessionsMap = new ConcurrentHashMap<>();


    // 使用虚拟线程优化 - JDK21特性
    private static final int PARTITION_COUNT = 128; // 更多分区减少竞争
    private static final Map<String, WebSocketSession>[] sessionPartitions =
            new ConcurrentHashMap[PARTITION_COUNT];

    // 连接数统计 - 使用原子操作
    private final AtomicInteger totalConnections = new AtomicInteger(0);
    private final AtomicInteger[] partitionCounts = new AtomicInteger[PARTITION_COUNT];

    private final ObjectMapper objectMapper = new ObjectMapper();

    // 使用虚拟线程执行器 - JDK21核心特性
    private final ScheduledExecutorService cleanupExecutor;
    private final ExecutorService messageExecutor;

    // 使用StampedLock提高读取性能
    private final StampedLock[] partitionLocks = new StampedLock[PARTITION_COUNT];

    public WebSocketHandler() {
        // 初始化虚拟线程执行器
        this.cleanupExecutor = Executors.newScheduledThreadPool(2, Thread.ofVirtual().factory());
        this.messageExecutor = Executors.newThreadPerTaskExecutor(Thread.ofVirtual().factory());

        // 初始化锁数组
        for (int i = 0; i < PARTITION_COUNT; i++) {
            partitionLocks[i] = new StampedLock();
            sessionPartitions[i] = new ConcurrentHashMap<>(400); // 更小的初始容量
            partitionCounts[i] = new AtomicInteger(0);
        }
    }

    @PostConstruct
    public void init() {
        // 更频繁的清理 - 使用虚拟线程
        cleanupExecutor.scheduleAtFixedRate(this::cleanupInvalidSessions, 1, 1, TimeUnit.MINUTES);

        // 内存监控任务
        cleanupExecutor.scheduleAtFixedRate(this::logMemoryStats, 30, 30, TimeUnit.SECONDS);
    }

    /**
     * 根据sessionId计算分区索引 - 优化哈希分布
     */
    private int getPartitionIndex(String sessionId) {
        // 使用更好的哈希算法减少冲突
        int hash = sessionId.hashCode();
        hash ^= (hash >>> 16); // 扩散哈希值
        return Math.abs(hash) % PARTITION_COUNT;
    }

    /**
     * 获取对应的分区Map - 使用乐观锁读取
     */
    private Map<String, WebSocketSession> getPartition(String sessionId) {
        int partitionIndex = getPartitionIndex(sessionId);
        long stamp = partitionLocks[partitionIndex].tryOptimisticRead();
        Map<String, WebSocketSession> partition = sessionPartitions[partitionIndex];

        if (!partitionLocks[partitionIndex].validate(stamp)) {
            // 乐观读失败，使用悲观读
            stamp = partitionLocks[partitionIndex].readLock();
            try {
                partition = sessionPartitions[partitionIndex];
            } finally {
                partitionLocks[partitionIndex].unlockRead(stamp);
            }
        }
        return partition;
    }

    /**
     * 连接建立时调用 - 优化版本
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String sessionId = session.getId();

        // 检查连接数限制 - 提高到8万上限
        if (totalConnections.get() >= 80000) {
            log.warn("连接数达到上限，拒绝新连接: {}", sessionId);
            session.close(CloseStatus.SESSION_NOT_RELIABLE);
            return;
        }

        String userId = extractUserId(session); // 从session中提取用户ID
        sessionUserMap.put(sessionId, userId);
        // 初始化用户会话集合
        CopyOnWriteArraySet<WebSocketSession> sessions = userSessionsMap.computeIfAbsent(userId, k -> new CopyOnWriteArraySet<>());
        sessions.add(session);

        int partitionIndex = getPartitionIndex(sessionId);
        Map<String, WebSocketSession> partition = sessionPartitions[partitionIndex];

        // 使用写锁添加会话
        long stamp = partitionLocks[partitionIndex].writeLock();
        try {
            partition.put(sessionId, session);
            partitionCounts[partitionIndex].incrementAndGet();
            totalConnections.incrementAndGet();
        } finally {
            partitionLocks[partitionIndex].unlockWrite(stamp);
        }

        log.info("新的WebSocket连接建立: {}, 分区: {}, 分区连接数: {}, 总连接数: {}",
                sessionId, partitionIndex, partitionCounts[partitionIndex].get(), totalConnections.get());


    }

    /**
     * 处理文本消息 - 使用虚拟线程异步处理
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String payload = message.getPayload();
        String sessionId = session.getId();

        // 使用虚拟线程异步处理消息，避免阻塞Netty线程
        messageExecutor.execute(() -> {
            try {
                processMessage(session, sessionId, payload);
            } catch (Exception e) {
                log.error("处理消息失败: {}", sessionId, e);
            }
        });
    }

    /**
     * 异步处理消息
     */
    private void processMessage(WebSocketSession session, String sessionId, String payload) {
        try {
            // 解析JSON消息
            Map<String, String> messageData = objectMapper.readValue(payload, Map.class);
            String type = messageData.get("type");
            String content = messageData.get("content");

            if (WebSocketType.START.getType().equals(type) && content != null && !content.trim().isEmpty()) {
                // 处理AI消息
//                handleChatMessage(session, content, sessionId);
            } else if (WebSocketType.PING.getType().equals(type)) {
                // 响应心跳
                session.sendMessage(new TextMessage(createMessage(WebSocketType.PONG.getType(), "pong", sessionId)));
            } else {
                session.sendMessage(new TextMessage(createMessage(WebSocketType.ERROR.getType(), "未知的消息类型", sessionId)));
            }
        } catch (Exception e) {
            log.error("消息处理失败: {}", payload, e);
            try {
                session.sendMessage(new TextMessage(createMessage(WebSocketType.ERROR.getType(), "消息处理失败", sessionId)));
            } catch (IOException ioException) {
                log.debug("发送错误消息失败: {}", sessionId);
            }
        }
    }

    /**
     * 处理传输错误
     */
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) {
        String sessionId = session.getId();
        log.debug("WebSocket传输错误: {}", sessionId, exception);

        removeSession(sessionId);

        try {
            session.close(CloseStatus.SERVER_ERROR);
        } catch (IOException e) {
            // 忽略关闭异常
        }
    }

    /**
     * 连接关闭时调用
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        String sessionId = session.getId();
        String userId = sessionUserMap.get(sessionId);
        removeSession(sessionId);
        sessionUserMap.remove(sessionId);

        if (userId != null) {
            userSessionsMap.computeIfPresent(userId, (k, v) -> {
                v.remove(session);
                return v.isEmpty() ? null : v;
            });
        }
        log.info("WebSocket连接关闭: {}, 状态: {}, 总连接数: {}",
                sessionId, status, totalConnections.get());
    }

    public static String extractUserId(WebSocketSession session) {
        // 实际从session属性中获取，例如：
        SysUserVO user = (SysUserVO) session.getAttributes().get("user");
        return user.getId();
    }

    /**
     * 移除会话 - 线程安全
     */
    private void removeSession(String sessionId) {
        int partitionIndex = getPartitionIndex(sessionId);
        Map<String, WebSocketSession> partition = sessionPartitions[partitionIndex];

        long stamp = partitionLocks[partitionIndex].writeLock();
        try {
            if (partition.remove(sessionId) != null) {
                partitionCounts[partitionIndex].decrementAndGet();
                totalConnections.decrementAndGet();
            }
        } finally {
            partitionLocks[partitionIndex].unlockWrite(stamp);
        }
    }


    /**
     * 定期清理无效连接 - 优化版本
     */
    private void cleanupInvalidSessions() {
        long startTime = System.currentTimeMillis();
        int cleanedCount = 0;

        // 使用并行流清理所有分区
        cleanedCount = Arrays.stream(sessionPartitions)
                .parallel()
                .mapToInt(this::cleanupPartition)
                .sum();

        if (cleanedCount > 0) {
            log.info("清理无效连接完成: 清理数量={}, 耗时={}ms, 当前总连接数={}",
                    cleanedCount, System.currentTimeMillis() - startTime, totalConnections.get());
        }
    }

    /**
     * 清理单个分区
     */
    private int cleanupPartition(Map<String, WebSocketSession> partition) {
        int cleaned = 0;
        Iterator<Map.Entry<String, WebSocketSession>> iterator = partition.entrySet().iterator();

        while (iterator.hasNext()) {
            Map.Entry<String, WebSocketSession> entry = iterator.next();
            WebSocketSession session = entry.getValue();

            if (!session.isOpen()) {
                iterator.remove();
                cleaned++;
                totalConnections.decrementAndGet();
            }
        }
        return cleaned;
    }

    /**
     * 内存状态监控
     */
    private void logMemoryStats() {
        Runtime runtime = Runtime.getRuntime();
        long maxMemory = runtime.maxMemory() / (1024 * 1024);
        long totalMemory = runtime.totalMemory() / (1024 * 1024);
        long freeMemory = runtime.freeMemory() / (1024 * 1024);
        long usedMemory = totalMemory - freeMemory;

        // 当已使用内存 ≥ 最大内存的80%时触发（整数运算避免浮点误差）
        if (usedMemory * 100 >= maxMemory * 80) {
            log.info("内存状态: 使用={}MB, 空闲={}MB, 总计={}MB, 最大={}MB, 连接数={}",
                    usedMemory, freeMemory, totalMemory, maxMemory, totalConnections.get());
        }
    }


    /**
     * 创建JSON格式的消息
     */
    private String createMessage(String type, String content, String sessionId) {
        try {
            return objectMapper.writeValueAsString(Map.of(
                    "type", type,
                    "content", content,
//                    "sessionId", sessionId,
                    "timestamp", System.currentTimeMillis()
            ));
        } catch (Exception e) {
            log.error("创建消息失败", e);
            return "{\"type\":\"error\", \"content\":\"消息创建失败\"}";
        }
    }

    /**
     * 获取当前连接数
     */
    public int getConnectionCount() {
        return totalConnections.get();
    }

    /**
     * 向特定会话发送消息 - 使用虚拟线程异步发送
     */
    public CompletableFuture<Boolean> sendToSessionAsync(String sessionId, SseEvent sseEvent, String message) {
        return CompletableFuture.supplyAsync(() -> {
            Map<String, WebSocketSession> partition = getPartition(sessionId);
            WebSocketSession session = partition.get(sessionId);

            if (session != null && session.isOpen()) {
                try {
                    synchronized (session) {
                        session.sendMessage(new TextMessage(createMessage(sseEvent.getValue(), message, sessionId)));
                    }
                    return true;
                } catch (Exception e) {
                    log.debug("向会话 {} 发送消息失败", sessionId, e);
                    removeSession(sessionId);
                }
            }
            return false;
        }, messageExecutor);
    }


    public CompletableFuture<Integer> sendToUserAsync(String userId, SseEvent sseEvent, String message) {
        return CompletableFuture.supplyAsync(() -> {
            CopyOnWriteArraySet<WebSocketSession> sessions = userSessionsMap.getOrDefault(userId, new CopyOnWriteArraySet<>());
            if (sessions.isEmpty()) {
                log.warn("用户[{}]没有活跃会话", userId);
                return 0;
            }

            // 创建并行发送任务
            List<CompletableFuture<Integer>> futures = sessions.stream()
                    .map(session -> sendToSessionAsync(session.getId(), sseEvent, message)
                            .exceptionally(e -> {
                                log.error("用户[{}]会话[{}]推送失败", userId, session.getId(), e);
                                return false;
                            })
                            .thenApply(success -> success ? 1 : 0)
                    ).collect(Collectors.toList());

            // 聚合所有结果（非阻塞式聚合）
            return CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]))
                    .thenApply(v -> futures.stream().mapToInt(CompletableFuture::join).sum())
                    .join(); // 此处join在supplyAsync内部，不会阻塞外部线程
        }, messageExecutor);
    }

    public CompletableFuture<Integer> sendToSessionAsyncById(String sessionId, SseEvent sseEvent, String message) {
        return CompletableFuture.supplyAsync(() -> {
            // 根据sessionId获取会话（需与您的存储结构匹配）
            WebSocketSession session = getSessionById(sessionId);

            if (session == null || !session.isOpen()) {
                log.warn("会话[{}]不存在或已关闭", sessionId);
                return 0;
            }

            // 创建单一发送任务
            CompletableFuture<Integer> future = sendToSessionAsync(sessionId, sseEvent, message)
                    .exceptionally(e -> {
                        log.error("会话[{}]推送失败", sessionId, e);
                        return false;
                    })
                    .thenApply(success -> success ? 1 : 0);

            // 直接返回单个结果（无需allOf聚合）
            return future.join();
        }, messageExecutor);
    }

    // 会话查询方法（需与您的存储结构适配）
    private WebSocketSession getSessionById(String sessionId) {
        // 示例实现：从分片存储中查找
        int partitionIndex = getPartitionIndex(sessionId);
        return sessionPartitions[partitionIndex].get(sessionId);
    }




    /**
     * 广播消息到所有客户端 - 高度优化的并行版本
     */
    public CompletableFuture<Void> broadcastToAllAsync(String message) {
        return CompletableFuture.runAsync(() -> {
            long startTime = System.currentTimeMillis();
            AtomicInteger successCount = new AtomicInteger(0);

            // 使用并行流处理所有分区
            Arrays.stream(sessionPartitions)
                    .parallel()
                    .forEach(partition -> {
                        partition.forEach((sessionId, session) -> {
                            if (session.isOpen()) {
                                try {
                                    synchronized (session) {
                                        session.sendMessage(new TextMessage(message));
                                    }
                                    successCount.incrementAndGet();
                                } catch (IOException e) {
                                    // 发送失败时不移除，等待清理任务处理
                                    log.debug("广播消息失败: {}", sessionId);
                                }
                            }
                        });
                    });

            log.info("广播完成: 成功发送={}, 总连接数={}, 耗时={}ms",
                    successCount.get(), totalConnections.get(), System.currentTimeMillis() - startTime);
        }, messageExecutor);
    }

    /**
     * 应用关闭时清理资源
     */
    public void destroy() {
        log.info("开始关闭WebSocket处理器, 当前连接数: {}", totalConnections.get());

        // 关闭线程池
        if (cleanupExecutor != null) {
            cleanupExecutor.shutdown();
        }

        if (messageExecutor != null) {
            messageExecutor.shutdown();
        }

        // 异步关闭所有连接
        messageExecutor.execute(() -> {
            Arrays.stream(sessionPartitions)
                    .parallel()
                    .forEach(partition -> {
                        partition.forEach((sessionId, session) -> {
                            try {
                                if (session.isOpen()) {
                                    session.close(CloseStatus.GOING_AWAY);
                                }
                            } catch (IOException e) {
                                // 忽略关闭异常
                            }
                        });
                        partition.clear();
                    });

            log.info("WebSocket处理器关闭完成");
        });
    }

//websoket连接后 消息推送格式
{
  "type"："ping"，
  "content":"心跳推送"
}