Check if WS is closing before sending message (#993)

PhotonVision · Nov 2, 2023 · 5a4eb54 · 5a4eb54
1 parent 1277459
commit 5a4eb54
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/photon-server/src/main/java/org/photonvision/server/DataSocketHandler.java b/photon-server/src/main/java/org/photonvision/server/DataSocketHandler.java
@@ -53,8 +53,6 @@ public class DataSocketHandler {
     @SuppressWarnings("FieldCanBeLocal")
     private final UIOutboundSubscriber uiOutboundSubscriber = new UIOutboundSubscriber(this);
 
-    public static class UIMap extends HashMap<String, Object> {}
-
     private static class ThreadSafeSingleton {
         private static final DataSocketHandler INSTANCE = new DataSocketHandler();
     }
@@ -70,23 +68,23 @@ private DataSocketHandler() {
     }
 
     public void onConnect(WsConnectContext context) {
+        users.add(context);
         context.session.setIdleTimeout(
                 Duration.ofMillis(Long.MAX_VALUE)); // TODO: determine better value
         var remote = (InetSocketAddress) context.session.getRemoteAddress();
         var host = remote.getAddress().toString() + ":" + remote.getPort();
         logger.info("New websocket connection from " + host);
-        users.add(context);
         dcService.publishEvent(
                 new IncomingWebSocketEvent<>(
                         DataChangeDestination.DCD_GENSETTINGS, "userConnected", context));
     }
 
     protected void onClose(WsCloseContext context) {
+        users.remove(context);
         var remote = (InetSocketAddress) context.session.getRemoteAddress();
         var host = remote.getAddress().toString() + ":" + remote.getPort();
         var reason = context.reason() != null ? context.reason() : "Connection closed by client";
         logger.info("Closing websocket connection from " + host + " for reason: " + reason);
-        users.remove(context);
     }
 
     @SuppressWarnings({"unchecked"})
@@ -349,7 +347,9 @@ public void onBinaryMessage(WsBinaryMessageContext context) {
 
     private void sendMessage(Object message, WsContext user) throws JsonProcessingException {
         ByteBuffer b = ByteBuffer.wrap(objectMapper.writeValueAsBytes(message));
-        user.send(b);
+        if (user.session.isOpen()) {
+            user.send(b);
+        }
     }
 
     public void broadcastMessage(Object message, WsContext userToSkip)