fix

zhuyifan
Commit c3b93742a067ec86ad43421f811eca175d126dc4 c3b93742 1 parent 1ebb8cf5
Showing 1 changed file with 64 additions and 56 deletions
src/main/java/com/infoloop/tianting/server/WebSocketServer.java
--- a/src/main/java/com/infoloop/tianting/server/WebSocketServer.java
View file @c3b9374
+++ b/src/main/java/com/infoloop/tianting/server/WebSocketServer.java
View file @c3b9374
@@ -7,14 +7,17 @@ import com.infoloop.tianting.server.session.UserSessionData;
 import com.infoloop.tianting.server.session.UserSessionKey;
 import com.infoloop.tianting.server.session.UserTypeEnum;
 import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.lang3.StringUtils;
 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 javax.annotation.Nullable;
 import java.io.IOException;
 import java.util.List;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
@@ -23,7 +26,7 @@ import java.util.stream.Collectors;
 @Slf4j
 public class WebSocketServer extends TextWebSocketHandler {
-    private static final ConcurrentHashMap<UserSessionKey, UserSessionData> userSessions = new ConcurrentHashMap<>();
+    private static final ConcurrentHashMap<UserSessionKey, List<UserSessionData>> userSessions = new ConcurrentHashMap<>();
     private static final long INACTIVITY_TIMEOUT = 60 * 60 * 1000;
     private static final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1);
@@ -33,38 +36,47 @@ public class WebSocketServer extends TextWebSocketHandler {
     private static void cleanInactiveSessions() {
         final var currentTime = System.currentTimeMillis();
-        final var iterator = userSessions.entrySet().iterator();
+        userSessions.forEach((key, sessionList) ->
-        log.info("cleanInactiveSessions; sessions:{}", userSessions.keySet());
+                userSessions.computeIfPresent(key, (k, sessions) -> {
-        while (iterator.hasNext()) {
+                    sessions.removeIf(userData -> {
-            final var entry = iterator.next();
+                        if (currentTime - userData.getLastActiveTime() > INACTIVITY_TIMEOUT) {
-            final var key = entry.getKey();
+                            try {
-            final var userData = entry.getValue();
+                                userData.closeSession();
-            final var lastActiveTime = userData.getLastActiveTime();
+                                log.info("closeSession, userId: {}, userType: {}", key.getUserId(), key.getUserType());
-            if (currentTime - lastActiveTime > INACTIVITY_TIMEOUT) {
+                            } catch (IOException e) {
-                try {
+                                log.error("closeSession error, userId: {}, userType: {}", key.getUserId(), key.getUserType(), e);
-                    userData.closeSession();
+                            }
-                    log.info("closeSession，userId : {}, userType : {}", key.getUserId(), key.getUserType());
+                            return true;
-                } catch (IOException e) {
+                        }
-                    log.error("closeSession error，userId : {}, userType : {}", key.getUserId(), key.getUserType(), e);
+                        return false;
-                }
+                    });
-                iterator.remove();
+                    return sessions.isEmpty() ? null : sessions; // 为空时移除
-            }
+                })
-        }
+        );
     }
+
     public static <T> void sendMessageToUser(UserSessionKey key, SocketMessage<T> socketMessage) throws IOException {
         sendMessageToUsers(List.of(key), socketMessage);
     }
     public static <T> void sendMessageToUsers(List<UserSessionKey> keys, SocketMessage<T> socketMessage) throws IOException {
         for (final var key : keys) {
-            final var userData = userSessions.get(key);
+            final var sessionList = userSessions.get(key);
-            if (userData != null && userData.getSession().isOpen()) {
+            if (sessionList != null) {
-                userData.getSession().sendMessage(new TextMessage(socketMessage.toJsonString()));
+                sessionList.forEach(userData -> {
-                userData.updateLastActiveTime();
+                    try {
-                log.info("send message to userId : {}, userType : {}", key.getUserId(), key.getUserType());
+                        if (userData.getSession().isOpen()) {
-            } else {
+                            userData.getSession().sendMessage(new TextMessage(socketMessage.toJsonString()));
-                log.info("userId : {}, userType : {} WebSocket connect closed; ", key.getUserId(), key.getUserType());
+                            userData.updateLastActiveTime();
+                            log.info("Send message to userId: {}, userType: {}", key.getUserId(), key.getUserType());
+                        } else {
+                            log.info("WebSocket closed; userId: {}, userType: {}", key.getUserId(), key.getUserType());
+                        }
+                    } catch (IOException e) {
+                        log.error("Failed to send message; userId: {}, userType: {}", key.getUserId(), key.getUserType(), e);
+                    }
+                });
             }
         }
     }
@@ -75,39 +87,40 @@ public class WebSocketServer extends TextWebSocketHandler {
     public static <T> void sendMessageByUserTypes(List<UserTypeEnum> userTypes, SocketMessage<T> socketMessage) throws IOException {
         final var keys = userSessions.keySet().stream()
-                .filter(userSessionData -> userTypes.contains(userSessionData.getUserType()))
+                .filter(userKey -> userTypes.contains(userKey.getUserType()))
                 .collect(Collectors.toList());
         sendMessageToUsers(keys, socketMessage);
     }
     public static <T> void sendMessageByUserTypes(List<UserTypeEnum> userTypes, List<UserAuthData> userAuthData, SocketMessage<T> socketMessage) throws IOException {
         final var keys = userSessions.keySet().stream()
-                .filter(userSessionData -> userTypes.contains(userSessionData.getUserType()))
+                .filter(userKey -> userTypes.contains(userKey.getUserType()))
-                .filter(userSessionData -> userAuthData.isEmpty() || (userAuthData.stream().anyMatch(data -> String.valueOf(data.getOperatorId()).equals(userSessionData.getUserId()) && data.getUserType().equals(userSessionData.getUserType()))))
+                .filter(userKey -> userAuthData.isEmpty()
+                        || userAuthData.stream().anyMatch(data -> String.valueOf(data.getOperatorId()).equals(userKey.getUserId()) && data.getUserType().equals(userKey.getUserType())))
                 .collect(Collectors.toList());
         sendMessageToUsers(keys, socketMessage);
     }
     public static void closeConnection(UserSessionKey key) {
-        final var userData = userSessions.get(key);
+        userSessions.computeIfPresent(key, (k, sessionList) -> {
-        if (userData != null) {
+            sessionList.forEach(userData -> {
-            try {
+                try {
-                userData.closeSession();
+                    userData.closeSession();
-                userSessions.remove(key);
+                } catch (IOException e) {
-                log.info("WebSocket close; userId = {}, userType = {}", key.getUserId(), key.getUserType());
+                    log.error("WebSocket close failed: userId: {}, userType: {}", key.getUserId(), key.getUserType(), e);
-            } catch (IOException e) {
+                }
-                log.error("WebSocket close failed: userId = {}, userType = {}", key.getUserId(), key.getUserType(), e);
+            });
-            }
+            return null; // 直接返回 null，移除 key
-        } else {
+        });
-            log.warn("Not Found userId : {}, userType : {} WebSocket Connect ", key.getUserId(), key.getUserType());
+        log.info("WebSocket closed; userId: {}, userType: {}", key.getUserId(), key.getUserType());
-        }
     }
+    @Nullable
     private UserSessionKey getUserSessionKey(WebSocketSession session) {
         final var userId = (String) session.getAttributes().get(CommonConstants.USER_ID);
         final var userTypeStr = (String) session.getAttributes().get(CommonConstants.USER_TYPE);
         final var userType = UserTypeEnum.fromString(userTypeStr);
-        if (userId != null && userType != null) {
+        if (StringUtils.isNotEmpty(userId) && userType != null) {
             return UserSessionKey.builder().userId(userId).userType(userType).build();
         }
         return null;
@@ -117,30 +130,25 @@ public class WebSocketServer extends TextWebSocketHandler {
     public void afterConnectionEstablished(WebSocketSession session) {
         final var key = getUserSessionKey(session);
         if (key != null) {
-            final var existingSession = userSessions.remove(key);
+            userSessions.computeIfAbsent(key, k -> new CopyOnWriteArrayList<>()).add(new UserSessionData(session));
-            if (existingSession != null) {
+            log.info("WebSocket connection established; userId: {}, userType: {}", key.getUserId(), key.getUserType());
-                try {
-                    existingSession.closeSession();
-                    log.info("WebSocket old connect closed，userId : {}, userType : {}", key.getUserId(), key.getUserType());
-                } catch (IOException e) {
-                    log.error("old WebSocket close failed，userId : {}, userType : {}", key.getUserId(), key.getUserType(), e);
-                }
-            }
-            userSessions.put(key, new UserSessionData(session));
-            log.info("WebSocket connect success, userId : {}, userType : {}", key.getUserId(), key.getUserType());
         } else {
-            log.warn("WebSocket connect failed，unable to get valid user information");
+            log.warn("WebSocket connection failed; unable to retrieve valid user information");
         }
     }
+
     @Override
     public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
         final var key = getUserSessionKey(session);
         if (key != null) {
-            userSessions.remove(key);
+            userSessions.computeIfPresent(key, (k, sessionList) -> {
-            log.info("WebSocket closed: userId : {}, userType : {}", key.getUserId(), key.getUserType());
+                sessionList.removeIf(userData -> userData.getSession().getId().equals(session.getId()));
+                return sessionList.isEmpty() ? null : sessionList;
+            });
+            log.info("WebSocket closed; userId: {}, userType: {}", key.getUserId(), key.getUserType());
         } else {
-            log.warn("WebSocket closed failed，unable to get valid user information");
+            log.warn("WebSocket closed; unable to retrieve valid user information");
         }
     }
 }
\ No newline at end of file