use helper method to terminate the receive loop on forceful websocket shutdown

Author: LucHeart

Common/Websocket/WebsockBaseController.cs

@@ -34,7 +34,7 @@ public abstract class WebsocketBaseController<T> : OpenShockControllerBase, IAsy
     private CancellationTokenSource? _linkedSource;
 
     protected CancellationToken LinkedToken;
-
+    
     /// <summary>
     /// Channel for multithreading thread safety of the websocket, MessageLoop is the only reader for this channel
     /// </summary>
@@ -53,7 +53,6 @@ protected WebsocketBaseController(ILogger<WebsocketBaseController<T>> logger)
         Logger = logger;
     }
 
-
     /// <inheritdoc />
     [NonAction]
     public ValueTask QueueMessage(T data) => Channel.Writer.WriteAsync(data, LinkedToken);
@@ -200,14 +199,18 @@ protected virtual Task SendWebSocketMessage(T message, WebSocket websocket, Canc
 
     #endregion
 
+    private CancellationTokenSource _receiveCancellationTokenSource = new();
+    
     /// <summary>
     /// Main receiver logic for the websocket
     /// </summary>
     /// <returns></returns>
     [NonAction]
     private async Task Logic()
     {
-        while (!LinkedToken.IsCancellationRequested)
+        using var linkedReceiverToken = CancellationTokenSource.CreateLinkedTokenSource(LinkedToken, _receiveCancellationTokenSource.Token);
+        
+        while (!linkedReceiverToken.IsCancellationRequested)
         {
             try
             {
@@ -231,7 +234,7 @@ private async Task Logic()
                     return;
                 }
 
-                if (!await HandleReceive())
+                if (!await HandleReceive(linkedReceiverToken.Token))
                 {
                     // HandleReceive returned false, we will close the connection after this
                     Logger.LogDebug("HandleReceive returned false, closing connection");
@@ -263,7 +266,18 @@ private async Task Logic()
     /// </summary>
     /// <returns>True if you want to continue the receiver loop, false if you want to terminate</returns>
     [NonAction]
-    protected abstract Task<bool> HandleReceive();
+    protected abstract Task<bool> HandleReceive(CancellationToken cancellationToken);
+    
+    [NonAction]
+    protected async Task ForceClose(WebSocketCloseStatus closeStatus, string? statusDescription)
+    {
+        await _receiveCancellationTokenSource.CancelAsync();
+
+        if (WebSocket is { State: WebSocketState.CloseReceived or WebSocketState.Open })
+        {
+            await WebSocket.CloseOutputAsync(closeStatus, statusDescription, LinkedToken);
+        }
+    }
 
     /// <summary>
     /// Send initial data to the client

LiveControlGateway/Controllers/HubControllerBase.cs

@@ -102,13 +102,7 @@ ILogger<FlatbuffersWebsocketBaseController<TIn, TOut>> logger
             try
             {
                 Logger.LogInformation("Keep alive timeout reached, closing websocket connection");
-
-                if (WebSocket is { State: WebSocketState.Open })
-                {
-                    await WebSocket!.CloseOutputAsync(WebSocketCloseStatus.ProtocolError, "Keep alive timeout reached",
-                        LinkedToken);
-                }
-                
+                await ForceClose(WebSocketCloseStatus.ProtocolError, "Keep alive timeout reached");
                 WebSocket?.Abort();
             }
             catch (Exception ex)
@@ -187,14 +181,7 @@ public async Task DisconnectOld()
         if (WebSocket == null)
             return;
 
-        if (WebSocket is { State: WebSocketState.Open })
-        {
-            await WebSocket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure,
-                "Hub is connecting from a different location",
-                LinkedToken);
-        }
-        
-        WebSocket?.Abort();
+        await ForceClose(WebSocketCloseStatus.NormalClosure, "Hub is connecting from a different location");
     }
 
     private static DateTimeOffset? GetBootedAtFromUptimeMs(ulong uptimeMs)

LiveControlGateway/Controllers/LiveControlController.cs

@@ -255,7 +255,7 @@ await QueueMessage(new LiveControlResponse<LiveResponseType>
     }
 
     /// <inheritdoc />
-    protected override async Task<bool> HandleReceive()
+    protected override async Task<bool> HandleReceive(CancellationToken cancellationToken)
     {
         var message =
             await JsonWebSocketUtils.ReceiveFullMessageAsyncNonAlloc<BaseRequest<LiveRequestType>>(WebSocket!,
@@ -266,21 +266,18 @@ await JsonWebSocketUtils.ReceiveFullMessageAsyncNonAlloc<BaseRequest<LiveRequest
                 if (request?.Data == null)
                 {
                     Logger.LogWarning("Received null data from client");
-                    await WebSocket!.CloseOutputAsync(WebSocketCloseStatus.InvalidPayloadData,
-                        "Invalid json message received", LinkedToken);
+                    await ForceClose(WebSocketCloseStatus.InvalidPayloadData, "Invalid json message received");
                     return false;
                 }
 
-#pragma warning disable CS4014
-                OsTask.Run(() => ProcessResult(request));
-#pragma warning restore CS4014
+                await ProcessResult(request);
+
                 return true;
             },
             async failed =>
             {
                 Logger.LogWarning(failed.Exception, "Deserialization failed for websocket message");
-                await WebSocket!.CloseOutputAsync(WebSocketCloseStatus.InvalidPayloadData,
-                    "Invalid json message received", LinkedToken);
+                await ForceClose(WebSocketCloseStatus.InvalidPayloadData, "Invalid json message received");
                 return false;
             }, closure =>
             {
@@ -551,8 +548,8 @@ await SendWebSocketMessage(new LiveControlResponse<LiveResponseType>
                 ResponseType = LiveResponseType.DeviceNotConnected,
             }, WebSocket!, LinkedToken);
 
-            await WebSocket!.CloseOutputAsync(WebSocketCloseStatus.NormalClosure,
-                "Hub is disconnected", LinkedToken);
+            
+            await ForceClose(WebSocketCloseStatus.NormalClosure, "Hub is disconnected");
         }
         catch (Exception e)
         {

LiveControlGateway/Websocket/FlatbuffersWebsocketBaseController.cs

@@ -48,18 +48,17 @@ protected override Task
     protected abstract Task<bool> Handle(TIn data);
 
     /// <inheritdoc />
-    protected override async Task<bool> HandleReceive()
+    protected override async Task<bool> HandleReceive(CancellationToken cancellationToken)
     {
         var message =
             await FlatbufferWebSocketUtils.ReceiveFullMessageAsyncNonAlloc(WebSocket!,
-                _incomingSerializer, LinkedToken);
+                _incomingSerializer, cancellationToken);
 
         var continueLoop = await message.Match(
             Handle,
             async _ =>
             {
-                await WebSocket!.CloseAsync(WebSocketCloseStatus.InvalidPayloadData, "Invalid flatbuffers message",
-                    LinkedToken);
+                await ForceClose(WebSocketCloseStatus.InvalidPayloadData, "Invalid flatbuffers message");
                 return false;
             },
             _ =>