Fix errors on live control disconnect

Author: LucHeart

Common/Websocket/WebsockBaseController.cs

@@ -102,7 +102,7 @@ public virtual async ValueTask DisposeAsync()
     /// </summary>
     /// <returns></returns>
     [NonAction]
-    public virtual ValueTask DisposeControllerAsync() => ValueTask.CompletedTask;
+    protected virtual ValueTask DisposeControllerAsync() => ValueTask.CompletedTask;
 
     /// <summary>
     /// Initial get request to the websocket route - rewrite to websocket connection

LiveControlGateway/Controllers/HubControllerBase.cs

@@ -203,7 +203,7 @@ protected async Task SelfOnline(DateTimeOffset bootedAt, ushort? latency = null,
     }
 
     /// <inheritdoc />
-    public override ValueTask DisposeControllerAsync()
+    protected override ValueTask DisposeControllerAsync()
     {
         Logger.LogTrace("Disposing controller timer");
         _keepAliveTimeoutTimer.Dispose();

LiveControlGateway/Controllers/HubV2Controller.cs

@@ -226,7 +226,7 @@ public override ValueTask OtaInstall(SemVersion version)
 
 
     /// <inheritdoc />
-    public override async ValueTask DisposeControllerAsync()
+    protected override async ValueTask DisposeControllerAsync()
     {
         await _pingTimer.DisposeAsync();
         await base.DisposeControllerAsync();

LiveControlGateway/Controllers/LiveControlController.cs

@@ -88,7 +88,7 @@ public LiveControlController(
     }
 
 
-    private bool _unregistered;
+    private volatile bool _unregistered;
 
     /// <inheritdoc />
     protected override async Task UnregisterConnection()
@@ -566,13 +566,36 @@ await QueueMessage(new LiveControlResponse<LiveResponseType>
             }
         });
     }
+    
+    /// <summary>
+    /// Called by a hub lifetime when the hub is disconnected and this controller needs to die
+    /// </summary>
+    [NonAction]
+    public async Task HubDisconnected()
+    {
+        _unregistered = true;
+        
+        Logger.LogTrace("Hub disconnected, disposing controller");
+        try
+        {
+            await SendWebSocketMessage(new LiveControlResponse<LiveResponseType>
+            {
+                ResponseType = LiveResponseType.DeviceNotConnected,
+            }, WebSocket!, LinkedToken);
+        }
+        catch (Exception e)
+        {
+            // We don't really care if this fails
+            Logger.LogDebug(e, "Error while sending disconnect message");
+        }
+        await DisposeAsync();
+    }
 
     /// <inheritdoc />
-    public override async ValueTask DisposeControllerAsync()
+    protected override async ValueTask DisposeControllerAsync()
     {
         Logger.LogTrace("Disposing controller timer");
         _pingTimer.Dispose();
-        await UpdateConnectedState(false);
         await base.DisposeControllerAsync();
     }
 }

LiveControlGateway/LifetimeManager/HubLifetime.cs

@@ -90,12 +90,12 @@ public HubLifetime([Range(1, 10)] byte tps, IHubController hubController,
     {
         using (await _liveControlClientsLock.LockAsyncScoped())
         {
-            if (_liveControlClients.Contains(controller)) 
+            if (_liveControlClients.Contains(controller))
             {
                 _logger.LogWarning("Client already registered, not sure how this happened, probably a bug");
                 return null;
             }
-            
+
             _liveControlClients = _liveControlClients.Add(controller);
         }
 
@@ -120,21 +120,21 @@ public async Task<bool> RemoveLiveControlClient(LiveControlController controller
 
     private async Task DisposeLiveControlClients()
     {
-        foreach (var client in _liveControlClients)
+        ImmutableArray<LiveControlController> liveControlClients;
+        using (await _liveControlClientsLock.LockAsyncScoped())
         {
-            try
-            {
-                await client.DisposeAsync();
-            }
-            catch (Exception e)
-            {
-                _logger.LogError(e, "Error disposing live control client");
-            }
+            liveControlClients = _liveControlClients;
+            _liveControlClients = _liveControlClients.Clear(); // Returns just an empty lol
         }
 
-        using (await _liveControlClientsLock.LockAsyncScoped())
+        try
+        {
+            await Task.WhenAll(liveControlClients.Select(client => client.HubDisconnected()));
+        }
+        catch (Exception e)
         {
-            _liveControlClients = _liveControlClients.Clear();
+            // We dont expect any errors here, but if there is a bug then this might happen, and it would be cat
+            _logger.LogError(e, "Error disposing live control client");
         }
     }
 
@@ -160,7 +160,7 @@ public async Task<bool> InitAsync(CancellationToken cancellationToken)
 #pragma warning restore CS4014
 
         _state = HubLifetimeState.Idle; // We are fully setup, we can go back to idle state
-        
+
         return true;
     }
 
@@ -269,7 +269,7 @@ public async Task UpdateDevice()
     {
         await using var db = await _dbContextFactory.CreateDbContextAsync(_cancellationSource.Token);
         await UpdateShockers(db, _cancellationSource.Token);
-        
+
         foreach (var websocketController in _liveControlClients)
             await websocketController.UpdatePermissions(db);
     }
@@ -429,11 +429,10 @@ public async ValueTask DisposeAsync()
     {
         if (_disposed) return;
         _disposed = true;
-        
+
         await _cancellationSource.CancelAsync();
         await DisposeLiveControlClients();
     }
-    
 }
 
 /// <summary>