Feature/improve livecontrol handling

Common/Errors/WebsocketError.cs

@@ -10,4 +10,7 @@ public static class WebsocketError
     public static OpenShockProblem WebsocketHubFirmwareVersionInvalid => new("Websocket.Hub.FirmwareVersionInvalid", "Supplied firmware version header is not valid semver", HttpStatusCode.BadRequest);
     public static OpenShockProblem WebsocketLiveControlHubIdInvalid => new("Websocket.LiveControl.HubIdInvalid", "Hub (device) id was missing or invalid", HttpStatusCode.BadRequest);
     public static OpenShockProblem WebsocketLiveControlHubNotFound => new("Websocket.LiveControl.HubNotFound", "Hub was not found or you are missing access", HttpStatusCode.NotFound);
+
+    public static OpenShockProblem WebsocketLiveControlHubLifetimeBusy => new("Websocket.LiveControl.HubLifetimeBusy", "Hub Lifetime is currently busy, try again please", HttpStatusCode.PreconditionFailed);
+    public static OpenShockProblem WebsocketHubLifetimeBusy => new("Websocket.Hub.LifetimeBusy", "Hub Lifetime is currently busy, try again soon please", HttpStatusCode.PreconditionFailed);
 }

Common/Websocket/WebsockBaseController.cs

@@ -17,7 +17,8 @@ namespace OpenShock.Common.Websocket;
 /// Base for json serialized websocket controller, you can override the SendMessageMethod to implement a different serializer
 /// </summary>
 /// <typeparam name="T"></typeparam>
-public abstract class WebsocketBaseController<T> : OpenShockControllerBase, IAsyncDisposable, IDisposable, IWebsocketController<T> where T : class
+public abstract class WebsocketBaseController<T> : OpenShockControllerBase, IAsyncDisposable, IDisposable,
+    IWebsocketController<T> where T : class
 {
     /// <inheritdoc />
     public abstract Guid Id { get; }
@@ -36,6 +37,7 @@ public abstract class WebsocketBaseController<T> : OpenShockControllerBase, IAsy
     /// When passing a cancellation token, pass this Linked token, it is a Link from ApplicationStopping and Close.
     /// </summary>
     protected readonly CancellationTokenSource LinkedSource;
+
     protected readonly CancellationToken LinkedToken;
 
     /// <summary>
@@ -62,7 +64,7 @@ public WebsocketBaseController(ILogger<WebsocketBaseController<T>> logger, IHost
 
     /// <inheritdoc />
     [NonAction]
-    public ValueTask QueueMessage(T data) => _channel.Writer.WriteAsync(data);
+    public ValueTask QueueMessage(T data) => _channel.Writer.WriteAsync(data, LinkedToken);
 
     private bool _disposed;
 
@@ -73,28 +75,28 @@ public virtual void Dispose()
         // ReSharper disable once MethodSupportsCancellation
         DisposeAsync().AsTask().Wait();
     }
-    
+
     /// <inheritdoc />
     [NonAction]
     public virtual async ValueTask DisposeAsync()
     {
-        if(_disposed) return;
+        if (_disposed) return;
         _disposed = true;
-        
+
         Logger.LogTrace("Disposing websocket controller..");
-        
+
         await DisposeControllerAsync();
         await UnregisterConnection();
-        
+
         _channel.Writer.Complete();
         await Close.CancelAsync();
         WebSocket?.Dispose();
         LinkedSource.Dispose();
-        
+
         GC.SuppressFinalize(this);
         Logger.LogTrace("Disposed websocket controller");
     }
-    
+
     /// <summary>
     /// Dispose function for any inheriting controller
     /// </summary>
@@ -115,7 +117,9 @@ public async Task Get()
             HttpContext.Response.StatusCode = StatusCodes.Status400BadRequest;
             var response = WebsocketError.NonWebsocketRequest;
             response.AddContext(HttpContext);
-            await HttpContext.Response.WriteAsJsonAsync(response, jsonOptions.Value.SerializerOptions, contentType: MediaTypeNames.Application.ProblemJson);
+            // ReSharper disable once MethodSupportsCancellation
+            await HttpContext.Response.WriteAsJsonAsync(response, jsonOptions.Value.SerializerOptions,
+                contentType: MediaTypeNames.Application.ProblemJson);
             await Close.CancelAsync();
             return;
         }
@@ -127,35 +131,32 @@ public async Task Get()
             var response = connectionPrecondition.AsT1.Value;
             HttpContext.Response.StatusCode = response.Status ?? StatusCodes.Status400BadRequest;
             response.AddContext(HttpContext);
-            await HttpContext.Response.WriteAsJsonAsync(response, jsonOptions.Value.SerializerOptions, contentType: MediaTypeNames.Application.ProblemJson);
-
+            // ReSharper disable once MethodSupportsCancellation
+            await HttpContext.Response.WriteAsJsonAsync(response, jsonOptions.Value.SerializerOptions,
+                contentType: MediaTypeNames.Application.ProblemJson);
+
             await Close.CancelAsync();
             return;
         }
 
         Logger.LogInformation("Opening websocket connection");
-
-        WebSocket?.Dispose(); // This should never happen, but just in case
+        WebSocket?.Dispose(); // This should never happen, suppresses warning
         WebSocket = await HttpContext.WebSockets.AcceptWebSocketAsync();
 
-        if (await TryRegisterConnection())
-        {
 #pragma warning disable CS4014
-            OsTask.Run(MessageLoop);
+        OsTask.Run(MessageLoop);
 #pragma warning restore CS4014
 
-            await SendInitialData();
-
-            await Logic();
-
-
-            if(_disposed) return;
-
-            await UnregisterConnection();
-        }
+        await SendInitialData();
 
+        await Logic();
 
-        await Close.CancelAsync();
+
+        if (_disposed) return;
+
+        await UnregisterConnection();
+
+        await DisposeAsync();
     }
 
     #region Send Loop
@@ -191,7 +192,6 @@ private async Task MessageLoop()
     protected virtual Task SendWebSocketMessage(T message, WebSocket websocket, CancellationToken cancellationToken) =>
         JsonWebSocketUtils.SendFullMessage(message, websocket, cancellationToken);
 
-
     #endregion
 
     /// <summary>
@@ -200,20 +200,14 @@ protected virtual Task SendWebSocketMessage(T message, WebSocket websocket, Canc
     /// <returns></returns>
     [NonAction]
     protected abstract Task Logic();
-    
+
     /// <summary>
     /// Send initial data to the client
     /// </summary>
     /// <returns></returns>
     [NonAction]
     protected virtual Task SendInitialData() => Task.CompletedTask;
 
-    /// <summary>
-    /// Action when the websocket connection is created
-    /// </summary>
-    [NonAction]
-    protected virtual Task<bool> TryRegisterConnection() => Task.FromResult(true);
-
     /// <summary>
     /// Action when the websocket connection is finished or disposed
     /// </summary>

LiveControlGateway/Controllers/HubControllerBase.cs

@@ -45,7 +45,14 @@ public abstract class HubControllerBase<TIn, TOut> : FlatbuffersWebsocketBaseCon
     /// Service provider
     /// </summary>
     protected readonly IServiceProvider ServiceProvider;
-
+
+    private HubLifetime? _hubLifetime = null;
+
+    /// <summary>
+    /// Hub lifetime
+    /// </summary>
+    /// <exception cref="InvalidOperationException"></exception>
+    protected HubLifetime HubLifetime => _hubLifetime ?? throw new InvalidOperationException("Hub lifetime is null but was tried to access");
     private readonly LcgOptions _options;
 
     private readonly HubLifetimeManager _hubLifetimeManager;
@@ -113,7 +120,7 @@ ILogger<FlatbuffersWebsocketBaseController<TIn, TOut>> logger
     private SemVersion? _firmwareVersion;
 
     /// <inheritdoc />
-    protected override Task<OneOf<Success, Error<OpenShockProblem>>> ConnectionPrecondition()
+    protected override async Task<OneOf<Success, Error<OpenShockProblem>>> ConnectionPrecondition()
     {
         _connected = DateTimeOffset.UtcNow;
 
@@ -125,24 +132,29 @@ protected override Task<OneOf<Success, Error<OpenShockProblem>>> ConnectionPreco
         else
         {
             var err = new Error<OpenShockProblem>(WebsocketError.WebsocketHubFirmwareVersionInvalid);
-            return Task.FromResult(OneOf<Success, Error<OpenShockProblem>>.FromT1(err));
+            return err;
         }
 
         _userAgent = HttpContext.Request.Headers.UserAgent.ToString().Truncate(256);
+        var hubLifetimeResult = await _hubLifetimeManager.TryAddDeviceConnection(5, this, LinkedToken);
 
-        return Task.FromResult(OneOf<Success, Error<OpenShockProblem>>.FromT0(new Success()));
+        if (hubLifetimeResult.IsT1)
+        {
+            Logger.LogWarning("Hub lifetime busy, closing connection");
+            return new Error<OpenShockProblem>(ExceptionError.Exception);
+        }
+
+        if (hubLifetimeResult.IsT2)
+        {
+            Logger.LogError("Hub lifetime error, closing connection");
+            return new Error<OpenShockProblem>();
+        }
+
+        _hubLifetime = hubLifetimeResult.AsT0;
+
+        return new Success();
     }
 
-
-    /// <summary>
-    /// Register to the hub lifetime manager
-    /// </summary>
-    /// <returns></returns>
-    protected override async Task<bool> TryRegisterConnection()
-    {
-        return await _hubLifetimeManager.TryAddDeviceConnection(5, this, LinkedToken);
-    }
-
     private bool _unregistered;
 
     /// <summary>
@@ -177,7 +189,7 @@ protected async Task SelfOnline(DateTimeOffset bootedAt, ushort? latency = null,
         // Reset the keep alive timeout
         _keepAliveTimeoutTimer.Interval = Duration.DeviceKeepAliveTimeout.TotalMilliseconds;
 
-        var result = await _hubLifetimeManager.DeviceOnline(CurrentHub.Id, new SelfOnlineData()
+        await HubLifetime.Online(CurrentHub.Id, new SelfOnlineData()
         {
             Owner = CurrentHub.Owner,
             Gateway = _options.Fqdn,
@@ -188,15 +200,6 @@ protected async Task SelfOnline(DateTimeOffset bootedAt, ushort? latency = null,
             LatencyMs = latency,
             Rssi = rssi
         });
-
-        if (result.IsT1)
-        {
-            Logger.LogError("Error while updating hub online status [{HubId}], we dont exist in the managers list", CurrentHub.Id);
-            await Close.CancelAsync();
-            if (WebSocket != null)
-                await WebSocket.CloseAsync(WebSocketCloseStatus.InternalServerError, "Hub not found in manager",
-                    CancellationToken.None);
-        }
     }
 
     /// <inheritdoc />