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<String, WebSocketSession> userSessions = new ConcurrentHashMap<>();
|
|
// 订阅关系：key=订阅标识(sn+type)，value=订阅用户ID集合
|
private final Map<String, Set<String>> subscriptions = new ConcurrentHashMap<>();
|
|
// 关键：为每个WebSocketSession分配独立的锁对象，确保写操作同步
|
private final Map<WebSocketSession, Object> 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<String, Set<String>> 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<String> 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;
|
}
|
}
|
