From 5a4eb54693fb54668f29354454e640e217fda193 Mon Sep 17 00:00:00 2001
From: Matt <matthew.morley.ca@gmail.com>
Date: Thu, 2 Nov 2023 19:34:31 -0400
Subject: [PATCH] Check if WS is closing before sending message (#993)

---
 .../org/photonvision/server/DataSocketHandler.java     | 10 +++++-----
 1 file changed, 5 insertions(+), 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
index f7f0345723..9fe6189f45 100644
--- 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)