package com.ard.sample.component.websocket.handler; import com.alibaba.fastjson2.JSON; import com.ard.sample.component.websocket.model.Payload; import com.ard.sample.component.websocket.model.WebSocketMessage; 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; @Component public class DockWebSocketHandler extends TextWebSocketHandler { // 存储用户会话: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<>(); /** * 连接建立后回调 */ @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()); System.out.println("用户[" + 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 (!Arrays.asList("status", "control", "progress", "state").contains(type)) { session.sendMessage(new TextMessage("不支持的消息类型:" + type)); return; } // 处理订阅/取消订阅 handleSubscription(userId, type, payload); } catch (Exception e) { System.err.println("消息处理错误:" + e.getMessage()); session.sendMessage(new TextMessage("消息格式错误")); } } /** * 处理订阅/取消订阅逻辑 */ 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); System.out.println("用户[" + userId + "]订阅了 设备[" + sn + "]的[" + type + "]类型消息"); } else { // 取消订阅操作 if (subscriptions.containsKey(subscriptionKey)) { subscriptions.get(subscriptionKey).remove(userId); if (subscriptions.get(subscriptionKey).isEmpty()) { subscriptions.remove(subscriptionKey); } System.out.println("用户[" + userId + "]取消订阅了 设备[" + sn + "]的[" + type + "]类型消息"); } } } /** * 连接关闭后回调 */ @Override public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception { String userId = (String) session.getAttributes().get("userId"); if (userId != null) { userSessions.remove(userId); sessionLocks.remove(session); // 移除会话锁 // 清除该用户的所有订阅 removeUserAllSubscriptions(userId); System.out.println("用户[" + userId + "]已断开连接,当前在线人数:" + userSessions.size()); } } /** * 清除用户的所有订阅 */ private void removeUserAllSubscriptions(String userId) { for (Map.Entry> entry : subscriptions.entrySet()) { entry.getValue().remove(userId); if (entry.getValue().isEmpty()) { subscriptions.remove(entry.getKey()); } } } /** * 向订阅了特定设备特定类型的用户发送消息 */ public void sendMessageToSubscribers(String sn, String type, String content) { String subscriptionKey = getSubscriptionKey(sn, type); Set subscriberIds = subscriptions.getOrDefault(subscriptionKey, new HashSet<>()); if (subscriberIds.isEmpty()) { //System.out.println("没有用户订阅 设备[" + sn + "]的[" + type + "]类型消息"); return; } try { // 构建发送的消息 // WebSocketMessage message = new WebSocketMessage(); // message.setType(type); // // Payload payload = new Payload(); // payload.setSn(sn); // payload.setSubscribe(true); // 仅作为标识,无实际意义 // message.setPayload(payload); // // String jsonMessage = JSON.toJSONString(message); // 向所有订阅者发送消息 for (String userId : subscriberIds) { WebSocketSession session = userSessions.get(userId); // 1. 检查会话是否有效 if (session == null || !session.isOpen()) { // 移除无效会话 userSessions.remove(userId); sessionLocks.remove(session); continue; } // 2. 获取当前会话的锁对象,确保同步写操作 Object lock = sessionLocks.get(session); if (lock == null) { continue; // 理论上不会出现,防止NPE } // 3. 同步发送消息,避免并发冲突 synchronized (lock) { session.sendMessage(new TextMessage(content)); System.out.println("向用户[" + userId + "]发送 设备[" + sn + "]的[" + type + "]类型消息:" + content); } // if (session != null && session.isOpen()) { // session.sendMessage(new TextMessage(content)); // System.out.println("向用户[" + userId + "]发送 设备[" + sn + "]的[" + type + "]类型消息"); // System.out.println(content); // } } } catch (IOException e) { System.err.println("发送消息错误:" + e.getMessage()); } } /** * 生成订阅标识(sn+type的组合) */ private String getSubscriptionKey(String sn, String type) { return sn + ":" + type; } }