Merge pull request #4187 from YevhenBondarenko/fix/ws

added ping for WS

application/src/main/java/org/thingsboard/server/controller/plugin/TbWebSocketHandler.java

@@ -41,9 +41,13 @@ import org.thingsboard.server.service.telemetry.TelemetryWebSocketMsgEndpoint;
 import org.thingsboard.server.service.telemetry.TelemetryWebSocketService;
 import org.thingsboard.server.service.telemetry.TelemetryWebSocketSessionRef;
 
-import javax.websocket.*;
+import javax.websocket.RemoteEndpoint;
+import javax.websocket.SendHandler;
+import javax.websocket.SendResult;
+import javax.websocket.Session;
 import java.io.IOException;
 import java.net.URI;
+import java.nio.ByteBuffer;
 import java.security.InvalidParameterException;
 import java.util.Queue;
 import java.util.Set;
@@ -79,6 +83,9 @@ public class TbWebSocketHandler extends TextWebSocketHandler implements Telemetr
     @Value("${server.ws.limits.max_updates_per_session:}")
     private String perSessionUpdatesConfiguration;
 
+    @Value("${server.ws.ping_timeout:30000}")
+    private long pingTimeout;
+
     private ConcurrentMap<String, TelemetryWebSocketSessionRef> blacklistedSessions = new ConcurrentHashMap<>();
     private ConcurrentMap<String, TbRateLimits> perSessionUpdateLimits = new ConcurrentHashMap<>();
 
@@ -120,6 +127,7 @@ public class TbWebSocketHandler extends TextWebSocketHandler implements Telemetr
                 return;
             }
             internalSessionMap.put(internalSessionId, new SessionMetaData(session, sessionRef, maxMsgQueuePerSession));
+
             externalSessionMap.put(externalSessionId, internalSessionId);
             processInWebSocketService(sessionRef, SessionEvent.onEstablished());
             log.info("[{}][{}][{}] Session is opened", sessionRef.getSecurityCtx().getTenantId(), externalSessionId, session.getId());
@@ -189,6 +197,8 @@ public class TbWebSocketHandler extends TextWebSocketHandler implements Telemetr
         private volatile boolean isSending = false;
         private final Queue<String> msgQueue;
 
+        private volatile long lastActivityTime;
+
         SessionMetaData(WebSocketSession session, TelemetryWebSocketSessionRef sessionRef, int maxMsgQueuePerSession) {
             super();
             this.session = session;
@@ -196,6 +206,23 @@ public class TbWebSocketHandler extends TextWebSocketHandler implements Telemetr
             this.asyncRemote = nativeSession.getAsyncRemote();
             this.sessionRef = sessionRef;
             this.msgQueue = new LinkedBlockingQueue<>(maxMsgQueuePerSession);
+            this.lastActivityTime = System.currentTimeMillis();
+        }
+
+        synchronized void sendPing(long currentTime) {
+            try {
+                if (currentTime - lastActivityTime >= pingTimeout) {
+                    this.asyncRemote.sendPing(ByteBuffer.wrap(new byte[]{}));
+                    lastActivityTime = currentTime;
+                }
+            } catch (Exception e) {
+                log.trace("[{}] Failed to send ping msg", session.getId(), e);
+                try {
+                    close(this.sessionRef, CloseStatus.SESSION_NOT_RELIABLE);
+                } catch (IOException ioe) {
+                    log.trace("[{}] Session transport error", session.getId(), ioe);
+                }
+            }
         }
 
         synchronized void sendMsg(String msg) {
@@ -243,6 +270,7 @@ public class TbWebSocketHandler extends TextWebSocketHandler implements Telemetr
                     log.trace("[{}] Session transport error", session.getId(), ioe);
                 }
             } else {
+                lastActivityTime = System.currentTimeMillis();
                 String msg = msgQueue.poll();
                 if (msg != null) {
                     sendMsgInternal(msg);
@@ -284,6 +312,22 @@ public class TbWebSocketHandler extends TextWebSocketHandler implements Telemetr
         }
     }
 
+    @Override
+    public void sendPing(TelemetryWebSocketSessionRef sessionRef, long currentTime) throws IOException {
+        String externalId = sessionRef.getSessionId();
+        String internalId = externalSessionMap.get(externalId);
+        if (internalId != null) {
+            SessionMetaData sessionMd = internalSessionMap.get(internalId);
+            if (sessionMd != null) {
+                sessionMd.sendPing(currentTime);
+            } else {
+                log.warn("[{}][{}] Failed to find session by internal id", externalId, internalId);
+            }
+        } else {
+            log.warn("[{}] Failed to find session by external id", externalId);
+        }
+    }
+
     @Override
     public void close(TelemetryWebSocketSessionRef sessionRef, CloseStatus reason) throws IOException {
         String externalId = sessionRef.getSessionId();

application/src/main/java/org/thingsboard/server/service/telemetry/DefaultTelemetryWebSocketService.java

@@ -64,7 +64,6 @@ import org.thingsboard.server.service.telemetry.cmd.v1.TelemetryPluginCmd;
 import org.thingsboard.server.service.telemetry.cmd.v1.TimeseriesSubscriptionCmd;
 import org.thingsboard.server.service.telemetry.cmd.v2.AlarmDataCmd;
 import org.thingsboard.server.service.telemetry.cmd.v2.CmdUpdate;
-import org.thingsboard.server.service.telemetry.cmd.v2.DataUpdate;
 import org.thingsboard.server.service.telemetry.cmd.v2.EntityCountCmd;
 import org.thingsboard.server.service.telemetry.cmd.v2.EntityDataCmd;
 import org.thingsboard.server.service.telemetry.cmd.v2.EntityDataUpdate;
@@ -89,6 +88,8 @@ import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
@@ -151,14 +152,23 @@ public class DefaultTelemetryWebSocketService implements TelemetryWebSocketServi
     private ExecutorService executor;
     private String serviceId;
 
+    private ScheduledExecutorService pingExecutor;
+
     @PostConstruct
     public void initExecutor() {
         serviceId = serviceInfoProvider.getServiceId();
         executor = Executors.newWorkStealingPool(50);
+
+        pingExecutor = Executors.newSingleThreadScheduledExecutor();
+        pingExecutor.scheduleWithFixedDelay(this::sendPing, 10000, 10000, TimeUnit.MILLISECONDS);
     }
 
     @PreDestroy
     public void shutdownExecutor() {
+        if (pingExecutor != null) {
+            pingExecutor.shutdownNow();
+        }
+
         if (executor != null) {
             executor.shutdownNow();
         }
@@ -774,6 +784,17 @@ public class DefaultTelemetryWebSocketService implements TelemetryWebSocketServi
         }
     }
 
+    private void sendPing() {
+        long currentTime = System.currentTimeMillis();
+        wsSessionsMap.values().forEach(md ->
+                executor.submit(() -> {
+                    try {
+                        msgEndpoint.sendPing(md.getSessionRef(), currentTime);
+                    } catch (IOException e) {
+                        log.warn("[{}] Failed to send ping: {}", md.getSessionRef().getSessionId(), e);
+                    }
+                }));
+    }
 
     private static Optional<Set<String>> getKeys(TelemetryPluginCmd cmd) {
         if (!StringUtils.isEmpty(cmd.getKeys())) {

application/src/main/java/org/thingsboard/server/service/telemetry/TelemetryWebSocketMsgEndpoint.java

@@ -26,5 +26,7 @@ public interface TelemetryWebSocketMsgEndpoint {
 
     void send(TelemetryWebSocketSessionRef sessionRef, int subscriptionId, String msg) throws IOException;
 
+    void sendPing(TelemetryWebSocketSessionRef sessionRef, long currentTime) throws IOException;
+
     void close(TelemetryWebSocketSessionRef sessionRef, CloseStatus withReason) throws IOException;
 }

application/src/main/resources/thingsboard.yml

@@ -34,6 +34,7 @@ server:
   log_controller_error_stack_trace: "${HTTP_LOG_CONTROLLER_ERROR_STACK_TRACE:false}"
   ws:
     send_timeout: "${TB_SERVER_WS_SEND_TIMEOUT:5000}"
+    ping_timeout: "${TB_SERVER_WS_PING_TIMEOUT:30000}"
     limits:
       # Limit the amount of sessions and subscriptions available on each server. Put values to zero to disable particular limitation
       max_sessions_per_tenant: "${TB_SERVER_WS_TENANT_RATE_LIMITS_MAX_SESSIONS_PER_TENANT:0}"