package com.ard.sample.component.websocket.handler; import com.alibaba.fastjson2.JSON; import com.ard.sample.common.model.ResponseResult; import com.ard.sample.component.websocket.enums.MessageType; import com.ard.sample.component.websocket.model.Payload; import com.ard.sample.component.websocket.model.WebSocketMessage; import com.ard.sample.manage.service.IControlService; import lombok.extern.slf4j.Slf4j; import org.springframework.beans.factory.annotation.Autowired; 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.*; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; @Slf4j @Component public class DockWebSocketHandler extends TextWebSocketHandler { @Autowired private LiveStreamHandler _liveStreamService; @Autowired private IControlService _controlService; // ====================== 心跳相关常量 ====================== /** 心跳间隔(秒):客户端需每隔此时间发一次ping */ private static final long HEARTBEAT_INTERVAL = 30; /** 心跳超时阈值(秒):超过此时间没收到ping,判定连接失效 */ private static final long HEARTBEAT_TIMEOUT = 90; /** 心跳检测任务间隔(秒):定时检查超时连接 */ private static final long CHECK_INTERVAL = 10; // ====================== 核心存储 ====================== // 存储用户会话:userId -> WebSocketSession private final Map userSessions = new ConcurrentHashMap<>(); // 订阅关系:key=订阅标识(sn+type),value=订阅用户ID集合 private final Map> subscriptions = new ConcurrentHashMap<>(); // 关键:为每个WebSocketSession分配独立的锁对象,确保写操作同步 private final Map sessionLocks = new ConcurrentHashMap<>(); // 存储用户最后心跳时间:userId -> 最后心跳时间戳(毫秒) private final Map userLastHeartbeatTime = new ConcurrentHashMap<>(); // ====================== 线程池 ====================== // 延时任务线程池(单线程,守护线程),用于处理连接关闭后的延时清理 private final ScheduledExecutorService delayExecutor = Executors.newSingleThreadScheduledExecutor(runnable -> { Thread thread = new Thread(runnable, "websocket-关闭-延迟-线程"); thread.setDaemon(true); // 设为守护线程,不阻塞JVM退出 return thread; }); // 心跳检测线程池(单线程,守护线程),定时检查超时连接 private final ScheduledExecutorService heartbeatCheckExecutor = Executors.newSingleThreadScheduledExecutor(runnable -> { Thread thread = new Thread(runnable, "websocket-心跳检测-线程"); thread.setDaemon(true); return thread; }); // ====================== 初始化心跳检测任务 ====================== public DockWebSocketHandler() { // 项目启动后,每隔 CHECK_INTERVAL 秒执行一次心跳检测 heartbeatCheckExecutor.scheduleAtFixedRate(this::checkHeartbeatTimeout, CHECK_INTERVAL, CHECK_INTERVAL, TimeUnit.SECONDS); } /** * 连接建立后回调 */ @Override public void afterConnectionEstablished(WebSocketSession session) throws Exception { // 从会话属性中获取认证后的用户ID String userId = (String) session.getAttributes().get("userId"); if (userId == null) { session.close(CloseStatus.POLICY_VIOLATION.withReason("未认证")); return; } userSessions.put(userId, session); // 为新会话初始化锁对象 sessionLocks.put(session, new Object()); // 初始化用户最后心跳时间(连接建立即视为一次心跳) userLastHeartbeatTime.put(userId, System.currentTimeMillis()); log.info("用户[{}]已连接,当前在线人数:{}", userId, userSessions.size()); } /** * 处理客户端发送的消息(包含心跳消息) */ @Override protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception { String userId = (String) session.getAttributes().get("userId"); String jsonMessage = message.getPayload(); try { // 解析客户端消息 WebSocketMessage webSocketMessage = JSON.parseObject(jsonMessage, WebSocketMessage.class); String type = webSocketMessage.getType(); Payload payload = webSocketMessage.getPayload(); // ====================== 处理心跳消息 ====================== if (MessageType.PING.getType().equals(type)) { // 1. 更新用户最后心跳时间 userLastHeartbeatTime.put(userId, System.currentTimeMillis()); // 2. 回复pong消息 WebSocketMessage pongMessage = new WebSocketMessage(); pongMessage.setType(MessageType.PONG.getType()); pongMessage.setPayload(new Payload()); // 空payload即可 sendMessage(session, JSON.toJSONString(pongMessage)); log.debug("用户[{}]心跳检测成功,已回复pong", userId); System.out.println("用户["+userId+"]心跳检测成功,已回复pong"); return; // 心跳消息无需走后续订阅逻辑 } // 验证消息类型是否合法 if (!MessageType.fromString(type)) { sendMessage(session, "不支持的消息类型:" + type); return; } // 处理订阅/取消订阅 handleSubscription(userId, type, payload); } catch (Exception e) { log.error("用户[{}]消息处理错误", userId, e); sendMessage(session, "消息格式错误"); } } /** * 处理订阅/取消订阅逻辑 */ private void handleSubscription(String userId, String type, Payload payload) { String sn = payload.getSn(); boolean subscribe = payload.isSubscribe(); String subscriptionKey = getSubscriptionKey(sn, type); // 生成唯一订阅标识 if (subscribe) { // 订阅操作 subscriptions.computeIfAbsent(subscriptionKey, k -> ConcurrentHashMap.newKeySet()).add(userId); log.info("用户[{}]订阅了 设备[{}]的[{}]类型消息", userId, sn, type); if (MessageType.LIVESTREAM.getType().equals(type)) { // 开启直播 log.info("开启设备[{}]直播", sn); _liveStreamService.StartLiveStreaming(sn); } } else { // 取消订阅操作 if (subscriptions.containsKey(subscriptionKey)) { subscriptions.get(subscriptionKey).remove(userId); if (subscriptions.get(subscriptionKey).isEmpty()) { subscriptions.remove(subscriptionKey); // 若取消的是直播订阅且无其他订阅者,停止直播 if (MessageType.LIVESTREAM.getType().equals(type)) { log.info("设备[{}]无订阅者,停止直播", sn); _liveStreamService.StopLiveStreaming(sn); } } log.info("用户[{}]取消订阅了 设备[{}]的[{}]类型消息", userId, sn, type); } } } /** * 连接关闭后回调 */ @Override public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception { String userId = (String) session.getAttributes().get("userId"); if (userId != null) { log.info("用户[{}]连接关闭(状态码:{}),将延时5秒执行清理操作", userId, status.getCode()); // 延时5秒执行清理逻辑 delayExecutor.schedule(() -> { try { // 清理核心存储 userSessions.remove(userId); sessionLocks.remove(session); userLastHeartbeatTime.remove(userId); // 清理心跳时间 // 清除该用户的所有订阅 removeUserAllSubscriptions(userId); log.info("用户[{}]延时5秒清理完成,当前在线人数:{}", userId, userSessions.size()); } catch (Exception e) { log.error("用户[{}]延时清理操作异常", userId, e); } }, 5, TimeUnit.SECONDS); // 延时5秒执行 } else { // 无认证用户,直接清理 sessionLocks.remove(session); log.info("无认证用户连接关闭,直接清理会话锁"); } } /** * 心跳超时检测:断开超过HEARTBEAT_TIMEOUT未发心跳的用户 */ private void checkHeartbeatTimeout() { long currentTime = System.currentTimeMillis(); // 遍历所有在线用户,检查心跳超时 for (Map.Entry entry : userLastHeartbeatTime.entrySet()) { String userId = entry.getKey(); Long lastHeartbeat = entry.getValue(); WebSocketSession session = userSessions.get(userId); // 1. 计算超时时间差 long timeoutDiff = currentTime - lastHeartbeat; if (timeoutDiff > HEARTBEAT_TIMEOUT * 1000) { log.warn("用户[{}]心跳超时({}秒未发ping),将断开连接", userId, timeoutDiff / 1000); try { // 2. 主动关闭连接 if (session != null && session.isOpen()) { session.close(CloseStatus.SESSION_NOT_RELIABLE.withReason("心跳超时")); } } catch (IOException e) { log.error("关闭用户[{}]超时连接失败", userId, e); } finally { // 3. 清理该用户的所有数据 userSessions.remove(userId); sessionLocks.remove(session); userLastHeartbeatTime.remove(userId); removeUserAllSubscriptions(userId); } } } } /** * 清除用户的所有订阅 */ private void removeUserAllSubscriptions(String userId) { // 遍历所有订阅,移除该用户并检查是否需要停止直播 Iterator>> iterator = subscriptions.entrySet().iterator(); while (iterator.hasNext()) { Map.Entry> entry = iterator.next(); Set userIds = entry.getValue(); if (userIds.contains(userId)) { userIds.remove(userId); // 若订阅组为空,移除并停止直播(如果是直播类型) if (userIds.isEmpty()) { String subscriptionKey = entry.getKey(); if (subscriptionKey.contains(MessageType.LIVESTREAM.getType())) { String sn = subscriptionKey.split(":")[0]; _liveStreamService.StopLiveStreaming(sn); log.info("设备[{}]无订阅者,停止直播", sn); } iterator.remove(); // 安全移除空订阅组 } } } } /** * 向订阅了特定设备特定类型的用户发送消息 */ public void sendMessageToSubscribers(String sn, String type, String content) { String subscriptionKey = getSubscriptionKey(sn, type); Set subscriberIds = subscriptions.getOrDefault(subscriptionKey, new HashSet<>()); if (subscriberIds.isEmpty()) { return; } try { for (String userId : subscriberIds) { WebSocketSession session = userSessions.get(userId); if (session == null || !session.isOpen()) { userSessions.remove(userId); sessionLocks.remove(session); userLastHeartbeatTime.remove(userId); // 清理心跳时间 continue; } // 同步发送消息,避免并发冲突 Object lock = sessionLocks.get(session); if (lock == null) { continue; } synchronized (lock) { session.sendMessage(new TextMessage(content)); log.debug("向用户[{}]发送 设备[{}]的[{}]类型消息:{}", userId, sn, type, content); } } } catch (IOException e) { log.error("发送消息错误:{}", e.getMessage()); } } /** * 通用发送消息方法(封装同步逻辑) */ private void sendMessage(WebSocketSession session, String content) throws IOException { Object lock = sessionLocks.get(session); if (lock != null && session.isOpen()) { synchronized (lock) { session.sendMessage(new TextMessage(content)); } } } /** * 生成订阅标识(sn+type的组合) */ private String getSubscriptionKey(String sn, String type) { return sn + ":" + type; } // ====================== 销毁线程池(可选,防止内存泄漏) ====================== public void destroy() { delayExecutor.shutdown(); heartbeatCheckExecutor.shutdown(); try { if (!delayExecutor.awaitTermination(1, TimeUnit.SECONDS)) { delayExecutor.shutdownNow(); } if (!heartbeatCheckExecutor.awaitTermination(1, TimeUnit.SECONDS)) { heartbeatCheckExecutor.shutdownNow(); } } catch (InterruptedException e) { delayExecutor.shutdownNow(); heartbeatCheckExecutor.shutdownNow(); } log.info("WebSocket线程池已销毁"); } }