fix: ASW-148 novnc 实现超时关闭功能

2 files changed, 52 insertions, 4 deletions

diff --git a/src/main/java/net/geedge/asw/common/config/websocket/EnvironmentNovncWebSocketHandler.java b/src/main/java/net/geedge/asw/common/config/websocket/EnvironmentNovncWebSocketHandler.java
index 58554fa..d73b0c7 100644
--- a/src/main/java/net/geedge/asw/common/config/websocket/EnvironmentNovncWebSocketHandler.java
+++ b/src/main/java/net/geedge/asw/common/config/websocket/EnvironmentNovncWebSocketHandler.java
@@ -6,6 +6,7 @@ import cn.hutool.core.util.ObjectUtil;
 import cn.hutool.json.JSONObject;
 import cn.hutool.log.Log;
 import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
+import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
 import net.geedge.asw.common.util.Constants;
 import net.geedge.asw.common.util.T;
 import net.geedge.asw.module.environment.entity.EnvironmentEntity;
@@ -22,6 +23,9 @@ import java.net.http.HttpClient;
 import java.net.http.WebSocket;
 import java.nio.ByteBuffer;
 import java.util.concurrent.CompletionStage;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
 
 @Component
 public class EnvironmentNovncWebSocketHandler extends TextWebSocketHandler {
@@ -31,7 +35,6 @@ public class EnvironmentNovncWebSocketHandler extends TextWebSocketHandler {
      * env id
      */
     private String envId;
-
     /**
      * session
      */
@@ -42,14 +45,23 @@ public class EnvironmentNovncWebSocketHandler extends TextWebSocketHandler {
      */
     private String userId;
 
+    private long lastReceivedTime;
+
+    private Integer sessionTimeout;
 
     private IEnvironmentService environmentService;
 
     private IEnvironmentSessionService environmentSessionService;
 
-    public EnvironmentNovncWebSocketHandler(IEnvironmentService environmentService, IEnvironmentSessionService environmentSessionService) {
-        this.environmentService = environmentService;
+    private ScheduledExecutorService scheduler;
+
+    public EnvironmentNovncWebSocketHandler() {
+    }
+    public EnvironmentNovncWebSocketHandler(IEnvironmentService deviceService, IEnvironmentSessionService environmentSessionService, Integer sessionTimeout) {
+        this.environmentService = deviceService;
         this.environmentSessionService = environmentSessionService;
+        this.sessionTimeout = sessionTimeout;
+        this.scheduler = Executors.newSingleThreadScheduledExecutor();
     }
 
     private void initFieldVal(WebSocketSession session) {
@@ -57,8 +69,34 @@ public class EnvironmentNovncWebSocketHandler extends TextWebSocketHandler {
         this.sessionId = (String) session.getAttributes().get("sessionId");
         this.userId = (String) session.getAttributes().get("userId");
         Constants.ENV_NOVNC_WEBSOCKET_SESSION.put(sessionId, session);
+        lastReceivedTime = T.DateUtil.currentSeconds(); // 初始化接收时间
+        startConnectionMonitor(session);
     }
 
+    private void startConnectionMonitor(WebSocketSession session) {
+        scheduler.scheduleAtFixedRate(() -> {
+            if (System.currentTimeMillis() - lastReceivedTime > sessionTimeout) {
+                try {
+                    if (session.isOpen()) {
+                        log.info("current no connection info, clean no real use connection finshed, sessionId: {}", sessionId);
+
+                        // update session status
+                        environmentSessionService.update(new LambdaUpdateWrapper<EnvironmentSessionEntity>()
+                                .set(EnvironmentSessionEntity::getStatus, 2)
+                                .eq(EnvironmentSessionEntity::getId, sessionId));
+
+                        session.close(CloseStatus.NORMAL.withReason("current no connection info, clean no real use connection finshed."));
+
+                    }
+                } catch (IOException e) {
+                    log.error(e, "Error clean no real use connection");
+                }
+                scheduler.shutdown(); // 关闭调度器
+            }
+        }, 0, 1, TimeUnit.MINUTES); // 每分钟检查一次
+    }
+
+
     @Override
     public void afterConnectionEstablished(WebSocketSession session) throws Exception {
         super.afterConnectionEstablished(session);
@@ -139,6 +177,7 @@ public class EnvironmentNovncWebSocketHandler extends TextWebSocketHandler {
             WebSocket envSocket = (WebSocket) session.getAttributes().get("envWebsocket");
             if (envSocket != null) {
                 envSocket.sendBinary(message.getPayload(), true);
+                lastReceivedTime = System.currentTimeMillis(); // 更新接收时间
             }
         } catch (Exception e) {
             log.error(e, "[handleBinaryMessage] [error]");
@@ -153,6 +192,10 @@ public class EnvironmentNovncWebSocketHandler extends TextWebSocketHandler {
             envWebsocket.sendClose(WebSocket.NORMAL_CLOSURE, "Normal closure");
         }
         Constants.ENV_NOVNC_WEBSOCKET_SESSION.remove(sessionId);
+        if (scheduler != null && !scheduler.isShutdown()) {
+            scheduler.shutdownNow(); // 停止调度器
+            scheduler.close();
+        }
         super.afterConnectionClosed(session, status);
     }
 
diff --git a/src/main/java/net/geedge/asw/common/config/websocket/WebSocketConfig.java b/src/main/java/net/geedge/asw/common/config/websocket/WebSocketConfig.java
index 3010dea..1c893c2 100644
--- a/src/main/java/net/geedge/asw/common/config/websocket/WebSocketConfig.java
+++ b/src/main/java/net/geedge/asw/common/config/websocket/WebSocketConfig.java
@@ -3,6 +3,7 @@ package net.geedge.asw.common.config.websocket;
 import net.geedge.asw.module.environment.service.IEnvironmentService;
 import net.geedge.asw.module.environment.service.IEnvironmentSessionService;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.annotation.Value;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.web.socket.config.annotation.EnableWebSocket;
 import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
@@ -12,6 +13,10 @@ import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry
 @EnableWebSocket
 public class WebSocketConfig implements WebSocketConfigurer {
 
+
+    @Value("${session.timeout:1800000}")  // 默认为 30 分钟
+    private Integer sessionTimeout;
+
     @Autowired
     private IEnvironmentService deviceService;
 
@@ -20,7 +25,7 @@ public class WebSocketConfig implements WebSocketConfigurer {
 
     @Override
     public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
-        registry.addHandler(new EnvironmentNovncWebSocketHandler(deviceService, environmentSessionService), "/api/v1/env/{envId}/session/{sessionId}/novnc")
+        registry.addHandler(new EnvironmentNovncWebSocketHandler(deviceService, environmentSessionService, sessionTimeout), "/api/v1/env/{envId}/session/{sessionId}/novnc")
                 .addInterceptors(new EnvironmentWebSocketInterceptor())
                 .setAllowedOrigins("*");