feat: Enhance WebSocket handling with improved ping/pong logic and session management

Signed-off-by: Mario Serrano <mario@dynamiasoluciones.com>

Author: marioserrano09

zk/src/main/java/tools/dynamia/zk/websocket/WebSocketGlobalCommandHandler.java

@@ -24,7 +24,7 @@
 import org.zkoss.zk.ui.Desktop;
 import tools.dynamia.commons.logger.Loggable;
 
-import java.util.Collection;
+import java.io.IOException;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
@@ -62,8 +62,8 @@
  */
 public class WebSocketGlobalCommandHandler extends TextWebSocketHandler implements Loggable {
 
-    private final Map<String, String> desktops = new ConcurrentHashMap<>();
-    private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();
+    // Map of active WebSocket sessions by session ID
+    private final Map<String, DeskstopWebSocketSession> sessions = new ConcurrentHashMap<>();
 
     /**
      * Handles incoming text messages from WebSocket clients.
@@ -91,17 +91,18 @@ protected void handleTextMessage(WebSocketSession session, TextMessage message)
             return;
         }
 
+        if ("PONG".equals(payload)) {
+            // Ignore PONG messages
+            return;
+        }
+
         // Handle desktop ID registration
-        String desktopId = payload;
-        String oldSessionId = desktops.get(desktopId);
-        if (oldSessionId != null) {
-            WebSocketSession oldSession = sessions.get(oldSessionId);
-            if (oldSession != null) {
-                oldSession.close(CloseStatus.NORMAL);
-            }
+        var oldSession = sessions.get(payload);
+        if (oldSession != null) {
+            oldSession.close(CloseStatus.NORMAL);
         }
-        log("Associating desktop " + desktopId + " with ws session " + session.getId());
-        desktops.put(desktopId, session.getId());
+        log("Associating desktop " + payload + " with ws session " + session.getId());
+        sessions.put(session.getId(), new DeskstopWebSocketSession(payload, session));
     }
 
     /**
@@ -114,8 +115,7 @@ protected void handleTextMessage(WebSocketSession session, TextMessage message)
      */
     @Override
     public void afterConnectionEstablished(WebSocketSession session) {
-        log("WebSocket connection established: " + session.getId());
-        sessions.put(session.getId(), session);
+        log("New webSocket connection established: " + session.getId() + " waiting for desktop ID...");
     }
 
     /**
@@ -125,16 +125,11 @@ public void afterConnectionEstablished(WebSocketSession session) {
      * the internal maps. This ensures that closed connections don't accumulate in memory.</p>
      *
      * @param session the closed WebSocket session
-     * @param status the status code indicating why the connection was closed
+     * @param status  the status code indicating why the connection was closed
      */
     @Override
     public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
         log("WebSocket connection closed: " + session.getId() + " with status " + status);
-
-        String desktopId = desktops.entrySet().stream().filter(e -> e.getValue().equals(session.getId())).map(Map.Entry::getKey).findFirst().orElse(null);
-        if (desktopId != null) {
-            desktops.remove(desktopId);
-        }
         sessions.remove(session.getId());
     }
 
@@ -147,15 +142,11 @@ public void afterConnectionClosed(WebSocketSession session, CloseStatus status)
      * @param desktop the ZK Desktop to find the session for
      * @return the associated WebSocket session, or {@code null} if not found or desktop is null
      */
-    WebSocketSession findSession(Desktop desktop) {
-        WebSocketSession session = null;
-        if (desktop != null && desktop.getId() != null) {
-            String sessionID = desktops.get(desktop.getId());
-            if (sessionID != null) {
-                session = sessions.get(sessionID);
-            }
-        }
-        return session;
+    public DeskstopWebSocketSession findSession(Desktop desktop) {
+        return sessions.values().stream()
+                .filter(ds -> ds.matchesDesktop(desktop))
+                .findFirst()
+                .orElse(null);
     }
 
     /**
@@ -166,7 +157,61 @@ WebSocketSession findSession(Desktop desktop) {
      *
      * @return a collection of all active WebSocket sessions
      */
-    Collection<WebSocketSession> getAllSessions() {
-        return sessions.values();
+    public Map<String, DeskstopWebSocketSession> getSessions() {
+        return sessions;
+    }
+
+    /**
+     * Closes the WebSocket session associated with a specific ZK Desktop.
+     *
+     * <p>This method safely closes the session and removes it from the internal
+     * session map to prevent resource leaks.</p>
+     *
+     * @param desktop the ZK Desktop whose session should be closed
+     */
+    public void closeSession(Desktop desktop) {
+        var session = findSession(desktop);
+        if (session != null) {
+            try {
+                session.close(CloseStatus.NORMAL);
+            } finally {
+                sessions.remove(session.session.getId());
+            }
+        } else {
+            log("No websocket session found for desktop " + desktop.getId());
+        }
+    }
+
+
+    /**
+     * Represents a WebSocket session associated with a specific ZK Desktop.
+     *
+     * <p>This record encapsulates the desktop ID and the WebSocket session,
+     * providing utility methods for session management and message sending.</p>
+     */
+    public record DeskstopWebSocketSession(String desktopId, WebSocketSession session) {
+        boolean isOpen() {
+            return session.isOpen();
+        }
+
+        boolean matchesDesktop(Desktop desktop) {
+            return desktop != null && desktopId.equals(desktop.getId());
+        }
+
+        boolean matchesSession(WebSocketSession otherSession) {
+            return session.getId().equals(otherSession.getId());
+        }
+
+        void close(CloseStatus closeStatus) {
+            try {
+                session.close(closeStatus);
+            } catch (Exception e) {
+                //ignore
+            }
+        }
+
+        void sendMessage(String message) throws IOException {
+            session.sendMessage(new TextMessage(message));
+        }
     }
 }

zk/src/main/java/tools/dynamia/zk/websocket/WebSocketPushSender.java

@@ -17,6 +17,7 @@
 
 package tools.dynamia.zk.websocket;
 
+import org.springframework.web.socket.CloseStatus;
 import org.springframework.web.socket.TextMessage;
 import org.springframework.web.socket.WebSocketSession;
 import org.zkoss.zk.ui.Desktop;
@@ -31,6 +32,7 @@
 import java.io.IOException;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
 
 /**
  * Helper class to send push command to desktop client using web socket
@@ -66,21 +68,23 @@ public static boolean sendPushCommand(Desktop desktop, String command, Map<Strin
                     data.putAll(payload);
                 }
                 data.put("command", command);
-                WebSocketSession session = handler.findSession(desktop);
+                var session = handler.findSession(desktop);
                 if (session != null) {
                     String textData = StringPojoParser.convertMapToJson(data);
-                    session.sendMessage(new TextMessage(textData));
+                    session.sendMessage(textData);
                     return true;
                 } else {
                     LOGGER.warn("No websocket session found for desktop " + desktop);
                 }
+            } catch (IOException e) {
+                LOGGER.error("IO Error sending push command '" + command + "' to Dekstop: " + desktop, e);
+                handler.closeSession(desktop);
             } catch (Exception e) {
                 LOGGER.error("Error sending push command '" + command + "' to Dekstop: " + desktop, e);
             }
         } else {
             LOGGER.warn("No websocket handler found");
         }
-
         return false;
     }
 
@@ -101,17 +105,31 @@ public static boolean sendPushCommand(String command) {
     public static void broadcastCommand(String command) {
         WebSocketGlobalCommandHandler handler = getHandler();
         if (handler != null) {
-            handler.getAllSessions().forEach(s -> {
+            AtomicInteger count = new AtomicInteger();
+            handler.getSessions().values().forEach(s -> {
                 try {
-                    s.sendMessage(new TextMessage(command));
+                    s.sendMessage(command);
+                    count.incrementAndGet();
                 } catch (IOException e) {
-                    LOGGER.error("Error sending command " + command + " to WS Session: " + s);
+                    LOGGER.error("IO Error sending command " + command + " to WS Session: " + s, e);
+                    s.close(CloseStatus.NORMAL);
+                } catch (Exception e) {
+                    LOGGER.error("Error sending command " + command + " to WS Session: " + s, e);
                 }
             });
+            LOGGER.info("Broadcasted command '" + command + "' to " + count.get() + " WS sessions.");
         }
     }
 
-    private static WebSocketGlobalCommandHandler getHandler() {
+    /**
+     * Broadcast a heartbeat PING command to all connected sessions
+     */
+    public static void broadcastHeartbeat() {
+        broadcastCommand("PING");
+    }
+
+
+    public static WebSocketGlobalCommandHandler getHandler() {
         return Containers.get().findObject(WebSocketGlobalCommandHandler.class);
     }
 

zk/src/main/resources/static/dynamia-tools/js/dynamia-tools-ws.js

@@ -118,7 +118,7 @@
     // Keep-alive configuration
     var keepAliveConfig = {
         enabled: true,
-        interval: 30000 // Send ping every 30 seconds
+        interval: 60000 // Send ping every 60 seconds
     };
 
     var wsManager = {
@@ -301,6 +301,17 @@
                     return;
                 }
 
+                if (e.data === 'PING') {
+                    // Respond to server ping
+                    try {
+                        wsManager.socket.send('PONG');
+                        console.debug('DynamiaTools WebSocket: pong sent in response to server ping');
+                    } catch (error) {
+                        console.warn('DynamiaTools WebSocket: error sending pong', error);
+                    }
+                    return;
+                }
+
                 console.debug('DynamiaTools command received:', e.data);
 
                 try {