修改websocket代码

ad-platform-pojo/src/main/java/com/baiye/model/dto/TreeUserDTO.java

@@ -13,5 +13,7 @@ public class TreeUserDTO extends UserDto {
 
     private static final long serialVersionUID = -6936022647922476487L;
 
+    private Long templateId;
+
     private List<TreeUserDTO> children;
 }

ad-platform-pojo/src/main/java/com/baiye/model/enums/ResponseCode.java

@@ -109,6 +109,12 @@ public enum ResponseCode {
      */
     INSUFFICIENT_ACCOUNT_BALANCE("1019", "账户余额不足"),
 
+    /**
+     *
+     * 账户余额不足
+     */
+    CONNECTION_SUCCEEDED("1027", "连接成功"),
+
     /**
      *
      * 非超级管理员不能修改数据

ad-platform-pojo/src/main/java/com/baiye/model/enums/WebSocketEnums.java

@@ -11,6 +11,7 @@ import lombok.Getter;
 @Getter
 @AllArgsConstructor
 public enum WebSocketEnums {
+
     /**
      * 置顶
      */
@@ -34,7 +35,13 @@ public enum WebSocketEnums {
     /**
      * 新建
      */
-    CREATE(5,"create");
+    CREATE(5,"create"),
+
+
+    /**
+     * 心跳
+     */
+    HEART_BEAT(6,"heart_beat");
 
     private final Integer value;
     private final String type;

manage/ad-platform-management/src/main/java/com/baiye/modules/system/service/impl/CompanyServiceImpl.java

@@ -147,6 +147,7 @@ public class CompanyServiceImpl implements CompanyService {
     }
 
     @Override
+    @CacheEvict(cacheNames = "companyCache", key = "#companyId")
     public void updateCompanyTemplateId(Long templateId, Long companyId) {
          companyRepository.updateCompanyTemplateIdById(templateId,companyId);
     }

manage/ad-platform-management/src/main/java/com/baiye/modules/system/service/impl/UserServiceImpl.java

@@ -135,6 +135,8 @@ public class UserServiceImpl implements UserService {
                 for (RoleSmallDto role : roles) {
                     if (role.getId() == DefaultNumberConstants.EIGHT_NUMBER) {
                         user.setIsManager(Boolean.TRUE);
+                        CompanyDto companyById = companyService.findCompanyById(user.getCompanyId());
+                        user.setTemplateId(companyById.getTemplateId());
                     }
                     if (role.getId() == DefaultNumberConstants.ONE_NUMBER) {
                         user.setIsAdmin(Boolean.TRUE);

manage/ad-platform-management/src/main/java/com/baiye/socket/WebSocketServer.java

@@ -6,6 +6,7 @@ import com.baiye.constant.DefaultNumberConstants;
 import com.baiye.http.WebSocketResponse;
 import com.baiye.manager.UserTokenManager;
 import com.baiye.model.dto.SendWebSocketDTO;
+import com.baiye.model.enums.ResponseCode;
 import com.baiye.model.enums.WebSocketEnums;
 import com.baiye.modules.system.service.MessageNotificationService;
 import com.baiye.modules.system.service.UserMessageService;
@@ -40,6 +41,7 @@ public class WebSocketServer {
     }
 
     private static final AtomicInteger ONLINE_COUNT = new AtomicInteger(0);
+
     /**
      * concurrent包的线程安全Set，用来存放每个客户端对应的Session对象。
      */
@@ -69,10 +71,16 @@ public class WebSocketServer {
      */
     @OnClose
     public void onClose(Session session) {
-        if (onlineId != null) {
-            SESSIONS.remove(onlineId);
+        Long userId = getUserId(session);
+        try {
+            if (userId != null) {
+                SESSIONS.remove(userId);
                 int cnt = ONLINE_COUNT.decrementAndGet();
-            log.info("有连接关闭，当前连接数为：{}", cnt);
+                session.close();
+                log.info("有连接关闭，关闭id为{},当前连接数为：{}", userId, cnt);
+            }
+        } catch (IOException e) {
+            throw new RuntimeException("关闭失败");
         }
     }
 
@@ -83,7 +91,6 @@ public class WebSocketServer {
      */
     @OnMessage
     public void onMessage(String message, Session session) {
-        log.info("来自客户端的消息：{}", message);
         JSONObject jsonObject = JSONUtil.parseObj(message);
         String type = "type";
         UserMessageService userMessageService =
@@ -144,6 +151,15 @@ public class WebSocketServer {
                                             (jsonObject.getStr(type))));
                 }
                 break;
+            case DefaultNumberConstants.SIX_NUMBER:
+                userId = getUserId(session);
+                if (userId != null) {
+                    sendMessage(session, JSONUtil.toJsonStr
+                            (WebSocketResponse.createBySuccess
+                                    (jsonObject.getStr(type),
+                                            ResponseCode.CONNECTION_SUCCEEDED)));
+                }
+                break;
             default:
                 break;
         }
@@ -203,8 +219,9 @@ public class WebSocketServer {
         }
         if (session != null) {
             sendMessage(session, message);
+            log.info("=============== 发送消息成功 ===============");
         } else {
-            log.warn("没有找到你指定ID的会话：{}", sessionId);
+            log.error("发送消息失败：{}", sessionId);
         }
     }
 
@@ -225,7 +242,7 @@ public class WebSocketServer {
                 data.setCode(DefaultNumberConstants.TWO_HUNDRED);
                 session.getBasicRemote().sendText(JSONUtil.toJsonStr(data));
             } else {
-                log.warn("没有找到你指定ID的会话：{}", id);
+                log.error("没有找到你指定ID的会话：{}", id);
             }
         }
     }
@@ -238,7 +255,7 @@ public class WebSocketServer {
                     session.getBasicRemote().sendText(JSONUtil.toJsonStr
                             (WebSocketResponse.createBySuccess("userMessage", hashMap.get(id))));
                 } else {
-                    log.warn("没有找到你指定ID的会话：{}", id);
+                    log.error("没有找到你指定ID的会话：{}", id);
                 }
             }
         } catch (IOException e) {