Improve hub disconnect and closure handling

Author: LucHeart

Common/Utils/JsonWebSocketUtils.cs

@@ -28,8 +28,6 @@ public static class JsonWebSocketUtils
                 bytes += result.Count;
                 if (result.MessageType == WebSocketMessageType.Close)
                 {
-                    await socket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Closure during message read",
-                        cancellationToken);
                     return new WebsocketClosure();
                 }
 

Common/Websocket/WebsockBaseController.cs

@@ -43,7 +43,7 @@ public abstract class WebsocketBaseController<T> : OpenShockControllerBase, IAsy
     /// <summary>
     /// Channel for multithreading thread safety of the websocket, MessageLoop is the only reader for this channel
     /// </summary>
-    private readonly Channel<T> _channel = Channel.CreateUnbounded<T>();
+    protected readonly Channel<T> Channel = System.Threading.Channels.Channel.CreateUnbounded<T>();
 
 #pragma warning disable IDISP008
     protected WebSocket? WebSocket;
@@ -64,7 +64,7 @@ public WebsocketBaseController(ILogger<WebsocketBaseController<T>> logger, IHost
 
     /// <inheritdoc />
     [NonAction]
-    public ValueTask QueueMessage(T data) => _channel.Writer.WriteAsync(data, LinkedToken);
+    public ValueTask QueueMessage(T data) => Channel.Writer.WriteAsync(data, LinkedToken);
 
     private bool _disposed;
 
@@ -88,7 +88,7 @@ public virtual async ValueTask DisposeAsync()
         await DisposeControllerAsync();
         await UnregisterConnection();
 
-        _channel.Writer.Complete();
+        Channel.Writer.TryComplete();
         await Close.CancelAsync();
         WebSocket?.Dispose();
         LinkedSource.Dispose();
@@ -167,7 +167,7 @@ await HttpContext.Response.WriteAsJsonAsync(response, jsonOptions.Value.Serializ
     [NonAction]
     private async Task MessageLoop()
     {
-        await foreach (var msg in _channel.Reader.ReadAllAsync(LinkedToken))
+        await foreach (var msg in Channel.Reader.ReadAllAsync(LinkedToken))
         {
             try
             {

LiveControlGateway/Controllers/LiveControlController.cs

@@ -280,23 +280,21 @@ await JsonWebSocketUtils.ReceiveFullMessageAsyncNonAlloc<BaseRequest<LiveRequest
 
                 if (message.IsT2)
                 {
-                    if (WebSocket.State != WebSocketState.Open)
-                    {
-                        Logger.LogTrace("Client sent closure, but connection state is not open");
-                        break;
-                    }
-
                     try
                     {
-                        await WebSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Normal close",
-                            LinkedToken);
+                        if (WebSocket.State is WebSocketState.Open or WebSocketState.CloseReceived
+                            or WebSocketState.CloseSent)
+                        {
+                            await WebSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Normal close",
+                                LinkedToken);
+                        }
                     }
                     catch (OperationCanceledException e)
                     {
                         Logger.LogError(e, "Error during close handshake");
                     }
 
-                    Logger.LogInformation("Closing websocket connection");
+                    Logger.LogTrace("Closing websocket connection");
                     break;
                 }
 
@@ -573,18 +571,26 @@ public async Task HubDisconnected()
         _unregistered = true;
 
         Logger.LogTrace("Hub disconnected, disposing controller");
+
+        Channel.Writer.TryComplete(); // Complete the channel so we stop sending messages
+        
         try
         {
+            var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10));
+            
             await SendWebSocketMessage(new LiveControlResponse<LiveResponseType>
             {
                 ResponseType = LiveResponseType.DeviceNotConnected,
-            }, WebSocket!, LinkedToken);
+            }, WebSocket!, cts.Token);
+            
+            await WebSocket!.CloseAsync(WebSocketCloseStatus.NormalClosure, "Hub is connecting from a different location", cts.Token);
         }
         catch (Exception e)
         {
             // We don't really care if this fails
-            Logger.LogDebug(e, "Error while sending disconnect message");
+            Logger.LogDebug(e, "Error while sending disconnect message or closing websocket");
         }
+        
         await DisposeAsync();
     }