webSocket优化

hx-socket/src/main/java/com/fjhx/socket/service/WebSocketServer.java

@@ -5,7 +5,9 @@ import com.alibaba.fastjson.JSON;
 import com.baomidou.mybatisplus.core.toolkit.IdWorker;
 import com.fjhx.socket.entity.MessageEntity;
 import com.fjhx.socket.event.*;
+import com.ruoyi.common.core.domain.entity.SysUser;
 import com.ruoyi.common.core.domain.model.LoginUser;
+import com.ruoyi.common.utils.SecurityUtils;
 import com.ruoyi.framework.web.service.TokenService;
 import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
@@ -19,6 +21,8 @@ import javax.websocket.server.ServerEndpoint;
 import java.io.IOException;
 import java.util.Date;
 import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 
 @ServerEndpoint("/webStock/{token}")
@@ -34,20 +38,28 @@ public class WebSocketServer {
      */
     private static final Map<Long, Map<Long, WebSocketServer>> webSocketMap = new ConcurrentHashMap<>();
     private static final TokenService tokenService = SpringUtil.getBean(TokenService.class);
+
     /**
      * 与某个客户端的连接会话，需要通过它来给客户端发送数据
      */
     private Session session;
+
     /**
      * 接收userId
      */
     private Long userId;
+
     /**
      * sessionId
      */
     private Long sessionId;
 
     /**
+     * 租户id
+     */
+    private String tenantId;
+
+    /**
      * 发送自定义消息
      *
      * @param userId  发送用户
@@ -55,6 +67,9 @@ public class WebSocketServer {
      * @param message 消息内容
      */
     public static void sendInfo(Long userId, Integer type, Object message) {
+
+        String tenantId = SecurityUtils.getTenantId();
+
         Map<Long, WebSocketServer> sessionIdWebSocketServerMap = webSocketMap.get(userId);
 
         if (sessionIdWebSocketServerMap == null) {
@@ -63,7 +78,13 @@ public class WebSocketServer {
         }
 
         for (WebSocketServer webSocketServer : sessionIdWebSocketServerMap.values()) {
-            webSocketServer.sendMessage(type, message);
+
+            String itemTenantId = webSocketServer.getTenantId();
+
+            if (Objects.equals(tenantId, itemTenantId)) {
+                webSocketServer.sendMessage(type, message);
+            }
+
         }
 
     }
@@ -100,7 +121,8 @@ public class WebSocketServer {
      */
     public static void sendInfoGroup(Integer type, Object message) {
         if (webSocketMap.size() > 0) {
-            for (Long userId : WebSocketServer.webSocketMap.keySet()) {
+            Set<Long> keySet = WebSocketServer.webSocketMap.keySet();
+            for (Long userId : keySet) {
                 sendInfo(userId, type, message);
             }
         }
@@ -119,9 +141,6 @@ public class WebSocketServer {
     @OnOpen
     public void onOpen(@RequestBody Session session, @PathParam("token") String token) {
 
-        this.session = session;
-        this.sessionId = IdWorker.getId();
-
         // 解析token
         LoginUser loginUser = tokenService.getLoginUser(token);
 
@@ -135,7 +154,12 @@ public class WebSocketServer {
             return;
         }
 
-        this.userId = loginUser.getUserId();
+        SysUser user = loginUser.getUser();
+
+        this.session = session;
+        this.sessionId = IdWorker.getId();
+        this.userId = user.getUserId();
+        this.tenantId = user.getTenantId();
 
         synchronized (this) {
             Map<Long, WebSocketServer> sessionIdWebSocketServerMap = webSocketMap.computeIfAbsent(userId, k -> new ConcurrentHashMap<>());
@@ -195,6 +219,7 @@ public class WebSocketServer {
      * @param message 消息内容
      */
     public void sendMessage(Integer type, Object message) {
+
         MessageEntity sendEntity = new MessageEntity();
         sendEntity.setUserId(userId);
         sendEntity.setSessionId(sessionId);
@@ -208,6 +233,7 @@ public class WebSocketServer {
         } catch (IOException e) {
             applicationContext.publishEvent(new WebSocketSendFailEvent(this, sendEntity));
         }
+
     }
 
 }