Feature/rework websocketbase (#190)

* Rework logic in live controller

* Implement into base

* Fix hub swapping and other edge cases

* Remove unused class

* No need to call abort on a websocket we know is null

Author: LucHeart

Common/Websocket/WebsockBaseController.cs

@@ -28,17 +28,12 @@ public abstract class WebsocketBaseController<T> : OpenShockControllerBase, IAsy
     /// </summary>
     protected readonly ILogger<WebsocketBaseController<T>> Logger;
 
-    /// <summary>
-    /// Close cancellation token to be called manually when termination of the current websocket is requested. Called on Dispose as well.
-    /// </summary>
-    protected readonly CancellationTokenSource Close = new();
-
     /// <summary>
     /// When passing a cancellation token, pass this Linked token, it is a Link from ApplicationStopping and Close.
     /// </summary>
-    protected readonly CancellationTokenSource LinkedSource;
+    private CancellationTokenSource? _linkedSource;
 
-    protected readonly CancellationToken LinkedToken;
+    protected CancellationToken LinkedToken;
 
     /// <summary>
     /// Channel for multithreading thread safety of the websocket, MessageLoop is the only reader for this channel
@@ -53,12 +48,9 @@ public abstract class WebsocketBaseController<T> : OpenShockControllerBase, IAsy
     /// DI
     /// </summary>
     /// <param name="logger"></param>
-    /// <param name="lifetime"></param>
-    public WebsocketBaseController(ILogger<WebsocketBaseController<T>> logger, IHostApplicationLifetime lifetime)
+    protected WebsocketBaseController(ILogger<WebsocketBaseController<T>> logger)
     {
         Logger = logger;
-        LinkedSource = CancellationTokenSource.CreateLinkedTokenSource(Close.Token, lifetime.ApplicationStopping);
-        LinkedToken = LinkedSource.Token;
     }
 
 
@@ -89,10 +81,9 @@ public virtual async ValueTask DisposeAsync()
         await UnregisterConnection();
 
         Channel.Writer.TryComplete();
-        await Close.CancelAsync();
 
         WebSocket?.Dispose();
-        LinkedSource.Dispose();
+        _linkedSource?.Dispose();
 
         GC.SuppressFinalize(this);
         Logger.LogTrace("Disposed websocket controller");
@@ -110,18 +101,25 @@ public virtual async ValueTask DisposeAsync()
     /// </summary>
     [ApiExplorerSettings(IgnoreApi = true)]
     [HttpGet]
-    public async Task Get()
+    public async Task Get([FromServices] IHostApplicationLifetime lifetime, CancellationToken cancellationToken)
     {
+#pragma warning disable IDISP003
+        _linkedSource = CancellationTokenSource.CreateLinkedTokenSource(lifetime.ApplicationStopping, cancellationToken);
+#pragma warning restore IDISP003
+        LinkedToken = _linkedSource.Token;
+        
         if (!HttpContext.WebSockets.IsWebSocketRequest)
         {
             var jsonOptions = HttpContext.RequestServices.GetRequiredService<IOptions<JsonOptions>>();
             HttpContext.Response.StatusCode = StatusCodes.Status400BadRequest;
             var response = WebsocketError.NonWebsocketRequest;
             response.AddContext(HttpContext);
             // ReSharper disable once MethodSupportsCancellation
-            await HttpContext.Response.WriteAsJsonAsync(response, jsonOptions.Value.SerializerOptions,
-                contentType: MediaTypeNames.Application.ProblemJson);
-            await Close.CancelAsync();
+            await HttpContext.Response.WriteAsJsonAsync(
+                response,
+                jsonOptions.Value.SerializerOptions,
+                contentType: MediaTypeNames.Application.ProblemJson,
+                cancellationToken: cancellationToken);
             return;
         }
 
@@ -133,16 +131,19 @@ await HttpContext.Response.WriteAsJsonAsync(response, jsonOptions.Value.Serializ
             HttpContext.Response.StatusCode = response.Status ?? StatusCodes.Status400BadRequest;
             response.AddContext(HttpContext);
             // ReSharper disable once MethodSupportsCancellation
-            await HttpContext.Response.WriteAsJsonAsync(response, jsonOptions.Value.SerializerOptions,
-                contentType: MediaTypeNames.Application.ProblemJson);
-
-            await Close.CancelAsync();
+            await HttpContext.Response.WriteAsJsonAsync(
+                response,
+                jsonOptions.Value.SerializerOptions,
+                contentType: MediaTypeNames.Application.ProblemJson,
+                cancellationToken: cancellationToken);
             return;
         }
 
         Logger.LogInformation("Opening websocket connection");
-        WebSocket?.Dispose(); // This should never happen, suppresses warning
+        
+#pragma warning disable IDISP003
         WebSocket = await HttpContext.WebSockets.AcceptWebSocketAsync();
+#pragma warning restore IDISP003
 
 #pragma warning disable CS4014
         OsTask.Run(MessageLoop);
@@ -154,8 +155,14 @@ await HttpContext.Response.WriteAsJsonAsync(response, jsonOptions.Value.Serializ
         // Logic ended
 
         await UnregisterConnection();
-        
-        await Close.CancelAsync();
+
+        // Only send close if the socket is still open, this allows us to close the websocket from inside the logic
+        // We send close if the client sent a close message though
+        if (WebSocket is { State: WebSocketState.Open or WebSocketState.CloseReceived }) 
+        {
+            await WebSocket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, "Normal closure",
+                LinkedToken);
+        }
     }
 
     #region Send Loop
@@ -198,7 +205,54 @@ protected virtual Task SendWebSocketMessage(T message, WebSocket websocket, Canc
     /// </summary>
     /// <returns></returns>
     [NonAction]
-    protected abstract Task Logic();
+    private async Task Logic()
+    {
+        while (!LinkedToken.IsCancellationRequested)
+        {
+            try
+            {
+                if (WebSocket == null)
+                {
+                    Logger.LogWarning("WebSocket is null, aborting");
+                    return;
+                }
+
+                if (WebSocket.State is WebSocketState.CloseReceived or WebSocketState.CloseSent or WebSocketState.Closed)
+                {
+                    // Client or we sent close message or both, we will close the connection after this
+                    return; 
+                }
+
+                if (WebSocket!.State != WebSocketState.Open)
+                {
+                    Logger.LogWarning("WebSocket is not open [{State}], aborting", WebSocket.State);
+                    WebSocket?.Abort();
+                    return;
+                }
+
+                await HandleReceive();
+
+            }
+            catch (OperationCanceledException)
+            {
+                Logger.LogWarning("WebSocket connection terminated due to close or shutdown");
+                return;
+            }
+            catch (Exception ex)
+            {
+                Logger.LogError(ex, "Exception while processing websocket request");
+                WebSocket?.Abort();
+                return;
+            }
+        }
+    }
+
+    /// <summary>
+    /// 
+    /// </summary>
+    /// <returns>True if you want to continue the receiver loop, false if you want to terminate</returns>
+    [NonAction]
+    protected abstract Task<bool> HandleReceive();
 
     /// <summary>
     /// Send initial data to the client

LiveControlGateway/Controllers/HubControllerBase.cs

@@ -1,4 +1,5 @@
-using FlatSharp;
+using System.Net.WebSockets;
+using FlatSharp;
 using Microsoft.AspNetCore.Mvc;
 using Microsoft.AspNetCore.Mvc.Filters;
 using Microsoft.Extensions.Options;
@@ -78,30 +79,36 @@ public void OnActionExecuted(ActionExecutedContext context)
     /// <summary>
     /// Base for hub websocket controllers
     /// </summary>
-    /// <param name="lifetime"></param>
     /// <param name="incomingSerializer"></param>
     /// <param name="outgoingSerializer"></param>
     /// <param name="hubLifetimeManager"></param>
     /// <param name="serviceProvider"></param>
     /// <param name="options"></param>
     /// <param name="logger"></param>
     protected HubControllerBase(
-        IHostApplicationLifetime lifetime,
         ISerializer<TIn> incomingSerializer,
         ISerializer<TOut> outgoingSerializer,
         HubLifetimeManager hubLifetimeManager,
         IServiceProvider serviceProvider,
         IOptions<LcgOptions> options,
         ILogger<FlatbuffersWebsocketBaseController<TIn, TOut>> logger
-        ) : base(logger, lifetime, incomingSerializer, outgoingSerializer)
+        ) : base(logger, incomingSerializer, outgoingSerializer)
     {
         _hubLifetimeManager = hubLifetimeManager;
         ServiceProvider = serviceProvider;
         _options = options.Value;
         _keepAliveTimeoutTimer.Elapsed += async (_, _) =>
         {
-            Logger.LogInformation("Keep alive timeout reached, closing websocket connection");
-            await Close.CancelAsync();
+            try
+            {
+                Logger.LogInformation("Keep alive timeout reached, closing websocket connection");
+                await WebSocket!.CloseOutputAsync(WebSocketCloseStatus.ProtocolError, "Keep alive timeout reached",
+                    LinkedToken);
+            }
+            catch (Exception ex)
+            {
+                Logger.LogError(ex, "Error while closing websocket connection from keep alive timeout");
+            }
         };
         _keepAliveTimeoutTimer.Start();
     }
@@ -167,7 +174,17 @@ protected override async Task UnregisterConnection()
 
     /// <inheritdoc />
     public abstract ValueTask OtaInstall(SemVersion version);
-    
+
+    /// <inheritdoc />
+    public async Task DisconnectOld()
+    {
+        if (WebSocket == null)
+            return;
+        
+        await WebSocket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, "Hub is connecting from a different location",
+            LinkedToken);
+    }
+
     private static DateTimeOffset? GetBootedAtFromUptimeMs(ulong uptimeMs)
     {
         var uptime = TimeSpan.FromMilliseconds(uptimeMs);

LiveControlGateway/Controllers/HubV1Controller.cs

@@ -31,21 +31,19 @@ public sealed class HubV1Controller : HubControllerBase<HubToGatewayMessage, Gat
     /// <summary>
     /// DI
     /// </summary>
-    /// <param name="lifetime"></param>
     /// <param name="hubLifetimeManager"></param>
     /// <param name="userHubContext"></param>
     /// <param name="serviceProvider"></param>
     /// <param name="options"></param>
     /// <param name="logger"></param>
     public HubV1Controller(
-            IHostApplicationLifetime lifetime,
             HubLifetimeManager hubLifetimeManager,
             IHubContext<UserHub, IUserHub> userHubContext,
             IServiceProvider serviceProvider,
             IOptions<LcgOptions> options,
             ILogger<HubV1Controller> logger
         )
-        : base(lifetime, HubToGatewayMessage.Serializer, GatewayToHubMessage.Serializer, hubLifetimeManager, serviceProvider, options, logger)
+        : base(HubToGatewayMessage.Serializer, GatewayToHubMessage.Serializer, hubLifetimeManager, serviceProvider, options, logger)
     {
         _userHubContext = userHubContext;
     }

LiveControlGateway/Controllers/HubV2Controller.cs

@@ -36,21 +36,19 @@ public sealed class HubV2Controller : HubControllerBase<HubToGatewayMessage, Gat
     /// <summary>
     /// DI
     /// </summary>
-    /// <param name="lifetime"></param>
     /// <param name="hubLifetimeManager"></param>
     /// <param name="userHubContext"></param>
     /// <param name="serviceProvider"></param>
     /// <param name="options"></param>
     /// <param name="logger"></param>
     public HubV2Controller(
-        IHostApplicationLifetime lifetime,
         HubLifetimeManager hubLifetimeManager,
         IHubContext<UserHub, IUserHub> userHubContext,
         IServiceProvider serviceProvider,
         IOptions<LcgOptions> options,
         ILogger<HubV2Controller> logger
         )
-        : base(lifetime, HubToGatewayMessage.Serializer, GatewayToHubMessage.Serializer, hubLifetimeManager, serviceProvider, options, logger)
+        : base(HubToGatewayMessage.Serializer, GatewayToHubMessage.Serializer, hubLifetimeManager, serviceProvider, options, logger)
     {
         _userHubContext = userHubContext;
         _pingTimer = new Timer(PingTimerElapsed, null, Duration.DevicePingInitialDelay, Duration.DevicePingPeriod);

LiveControlGateway/Controllers/IHubController.cs

@@ -33,4 +33,10 @@ public interface IHubController : IAsyncDisposable
     /// <param name="version"></param>
     /// <returns></returns>
     public ValueTask OtaInstall(SemVersion version);
+
+    /// <summary>
+    /// Disconnect the old connection in favor of the new one
+    /// </summary>
+    /// <returns></returns>
+    public Task DisconnectOld();
 }