自定義SecuritySessionKeyGenerator類,為每個客戶端連接建立唯一的key

package com.pig4cloud.plugin.websocket.custom;import com.pig4cloud.plugin.websocket.holder.SessionKeyGenerator;
import org.springframework.web.socket.WebSocketSession;import java.util.UUID;
public class SecuritySessionKeyGenerator implements SessionKeyGenerator {@Overridepublic Object sessionKey(WebSocketSession webSocketSession) {Object user = webSocketSession.getAttributes().get("USER_KEY_ATTR_NAME");// 添加隨機后綴使每個連接鍵唯一return user != null ? user + "-" + UUID.randomUUID().toString() : null;}
}

重寫WebSocketSessionHolder類方法

package com.pig4cloud.plugin.websocket.holder;import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;import org.springframework.web.socket.WebSocketSession;public final class WebSocketSessionHolder {// 僅修改這一行：將Value改為Set類型private static final Map<String, Set<WebSocketSession>> USER_SESSION_MAP = new ConcurrentHashMap<>();private WebSocketSessionHolder() {}public static void addSession(Object sessionKey, WebSocketSession session) {// 修改為支持多會話的添加方式USER_SESSION_MAP.computeIfAbsent(sessionKey.toString(),k -> new CopyOnWriteArraySet<>()).add(session);}public static void removeSession(Object sessionKey) {// 移除時不再自動清除所有會話Set<WebSocketSession> sessions = USER_SESSION_MAP.get(sessionKey.toString());if (sessions != null) {sessions.removeIf(s -> !s.isOpen()); // 只移除已關閉的連接if (sessions.isEmpty()) {USER_SESSION_MAP.remove(sessionKey.toString());}}}// 保持原有方法簽名不變（兼容現有調用）public static WebSocketSession getSession(Object sessionKey) {Set<WebSocketSession> sessions = USER_SESSION_MAP.get(sessionKey.toString());return sessions != null ? sessions.stream().findFirst().orElse(null) : null;}// 新增方法：獲取用戶的所有會話public static Set<WebSocketSession> getSessions(Object sessionKey) {return USER_SESSION_MAP.getOrDefault(sessionKey.toString(), Collections.emptySet());}// 保持原有方法不變public static Collection<WebSocketSession> getSessions() {List<WebSocketSession> allSessions = new ArrayList<>();USER_SESSION_MAP.values().forEach(allSessions::addAll);return allSessions;}// 保持原有方法不變public static Set<String> getSessionKeys() {return USER_SESSION_MAP.keySet();}
}