优化websocket服务

pom.xml

@@ -10,7 +10,7 @@
     </parent>
     <groupId>thyyxxk</groupId>
     <artifactId>wxservice-server</artifactId>
-    <version>5.4</version>
+    <version>5.5</version>
     <name>wxservice-server</name>
     <description>server for wxservice-web</description>
 

src/main/java/thyyxxk/wxservice_server/controller/WebSocketController.java

@@ -22,4 +22,15 @@ public class WebSocketController {
         return ResultVoUtil.success(data);
     }
 
+    @GetMapping("/systemBreak")
+    public ResultVo<String> systemBreak() {
+        int count = WebSocketServer.systemBreak();
+        return ResultVoUtil.success("操作成功，成功关闭连接数：" + count);
+    }
+
+    @GetMapping("/systemRecover")
+    public ResultVo<String> systemRecover() {
+        WebSocketServer.systemRecover();
+        return ResultVoUtil.success("操作成功，WebSocket服务已恢复。");
+    }
 }

src/main/java/thyyxxk/wxservice_server/websocket/WebSocketServer.java

@@ -2,15 +2,17 @@ package thyyxxk.wxservice_server.websocket;
 
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.stereotype.Component;
+import thyyxxk.wxservice_server.constant.Constants;
 
 import javax.websocket.*;
 import javax.websocket.server.PathParam;
 import javax.websocket.server.ServerEndpoint;
+import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicInteger;
 
 /**
- * @description: websoket config
+ * @description: websocket config
  * @author: DingJie
  * @create: 2021-05-18 11:20:43
  **/
@@ -19,26 +21,12 @@ import java.util.concurrent.atomic.AtomicInteger;
 @Component
 public class WebSocketServer {
     private static final AtomicInteger ONLINE_COUNT = new AtomicInteger();
-    private static final ConcurrentHashMap<String, Session> SESSION_MAP = new ConcurrentHashMap<>();
-
-    @OnOpen
-    public void onOpen(Session session, @PathParam("openId") String openId) {
-        SESSION_MAP.put(openId, session);
-        ONLINE_COUNT.incrementAndGet();
-    }
-
-    @OnClose
-    public void onClose(@PathParam("openId") String openId) {
-        SESSION_MAP.remove(openId);
-        ONLINE_COUNT.decrementAndGet();
-    }
-
-    @OnError
-    public void onError(Session session, Throwable error) {
-        log.error("WebSocket发生错误>>> session: {}，错误：{} ", session.getId(), error.getMessage());
-    }
+    private static ConcurrentHashMap<String, Session> SESSION_MAP = new ConcurrentHashMap<>(Constants.Capacity.DEFAULT);
 
     public static void sendMessage(String openId, String message) {
+        if (null == SESSION_MAP) {
+            return;
+        }
         Session session = SESSION_MAP.get(openId);
         if (null == session) {
             return;
@@ -54,4 +42,69 @@ public class WebSocketServer {
     public static int getOnlineCount() {
         return ONLINE_COUNT.get();
     }
+
+    public static int systemBreak() {
+        int count = 0;
+        if (null == SESSION_MAP) {
+            return count;
+        }
+        for (Map.Entry<String, Session> entry : SESSION_MAP.entrySet()) {
+            if (entry.getValue().isOpen()) {
+                try {
+                    CloseReason closeReason = new CloseReason(CloseReason.CloseCodes.NORMAL_CLOSURE, "服务更新。");
+                    entry.getValue().close(closeReason);
+                    count++;
+                } catch (Exception e) {
+                    log.error("关闭WebSocket发生错误>>> session: {}，错误：{}", entry.getValue().getId(), e.getMessage());
+                }
+            }
+        }
+        SESSION_MAP = null;
+        ONLINE_COUNT.set(0);
+        return count;
+    }
+
+    public static void systemRecover() {
+        if (null == SESSION_MAP) {
+            SESSION_MAP = new ConcurrentHashMap<>(Constants.Capacity.DEFAULT);
+        }
+    }
+
+    @OnOpen
+    public void onOpen(Session session, @PathParam("openId") String openId) {
+        if (null == SESSION_MAP) {
+            if (session.isOpen()) {
+                try {
+                    CloseReason closeReason = new CloseReason(CloseReason.CloseCodes.NORMAL_CLOSURE, "服务更新。");
+                    session.close(closeReason);
+                } catch (Exception e) {
+                    log.error("关闭WebSocket发生错误>>> session: {}，错误：{}", session.getId(), e.getMessage());
+                }
+            }
+        } else {
+            SESSION_MAP.put(openId, session);
+            ONLINE_COUNT.incrementAndGet();
+        }
+    }
+
+    @OnClose
+    public void onClose(@PathParam("openId") String openId) {
+        if (null != SESSION_MAP) {
+            SESSION_MAP.remove(openId);
+        }
+        ONLINE_COUNT.decrementAndGet();
+    }
+
+    @OnError
+    public void onError(Session session, Throwable error) {
+        log.error("WebSocket发生错误>>> session: {}，错误：{} ", session.getId(), error.getMessage());
+        if (session.isOpen()) {
+            try {
+                CloseReason closeReason = new CloseReason(CloseReason.CloseCodes.NORMAL_CLOSURE, "服务更新。");
+                session.close(closeReason);
+            } catch (Exception e) {
+                log.error("关闭WebSocket发生错误>>> session: {}，错误：{}", session.getId(), e.getMessage());
+            }
+        }
+    }
 }

src/main/resources/application-prod.yml

@@ -1,5 +1,5 @@
 server:
-  port: 8085
+  port: 8083
   servlet:
     context-path: /wxserver
 spring: