Different more sane approach

LucHeart · LucHeart · commit 58dd8da905b4 · 2025-01-31T18:56:46.000+01:00
diff --git a/LiveControlGateway/Controllers/HubControllerBase.cs b/LiveControlGateway/Controllers/HubControllerBase.cs
@@ -146,7 +146,6 @@ protected override async Task<bool> TryRegisterConnection()
     /// </summary>
     protected override async Task UnregisterConnection()
     {
-        if (_newConnection) return; // We dont want to call this here, as it would lead to a deadlock, this is already taken care of in the manager
         await _hubLifetimeManager.RemoveDeviceConnection(this);
     }
 
@@ -159,21 +158,6 @@ protected override async Task UnregisterConnection()
     /// <inheritdoc />
     public abstract ValueTask OtaInstall(SemVersion version);
 
-    /// <summary>
-    /// When we are disposing the controller because there is a new connection already
-    /// </summary>
-    private bool _newConnection = false;
-    
-    /// <summary>
-    /// Called by the hub lifetime manager to dispose the connection because there is a new connection
-    /// This is a direct replacement for DisposeAsync
-    /// </summary>
-    /// <returns></returns>
-    public ValueTask DisposeForNewConnection()
-    {
-        _newConnection = true;
-        return DisposeAsync();
-    }
 
     /// <summary>
     /// Keep the hub online
diff --git a/LiveControlGateway/Controllers/IHubController.cs b/LiveControlGateway/Controllers/IHubController.cs
@@ -34,10 +34,4 @@ public interface IHubController : IAsyncDisposable
     /// <param name="version"></param>
     /// <returns></returns>
     public ValueTask OtaInstall(SemVersion version);
-
-    /// <summary>
-    /// Dispose the current hub controller because a new connection has been established
-    /// </summary>
-    /// <returns></returns>
-    public ValueTask DisposeForNewConnection();
 }
diff --git a/LiveControlGateway/LifetimeManager/HubLifetime.cs b/LiveControlGateway/LifetimeManager/HubLifetime.cs
@@ -26,13 +26,24 @@ namespace OpenShock.LiveControlGateway.LifetimeManager;
 /// </summary>
 public sealed class HubLifetime : IAsyncDisposable
 {
+    
+    public enum HubLifetimeState
+    {
+        Idle,
+        SettingUp,
+        Swapping,
+        Removing
+    }
+
+    private volatile HubLifetimeState _state = HubLifetimeState.SettingUp;
+    public IHubController HubController { get; private set; }
+    
     private readonly TimeSpan _waitBetweenTicks;
     private readonly ushort _commandDuration;
 
     private Dictionary<Guid, ShockerState> _shockerStates = new();
     private readonly byte _tps;
-    public IHubController HubController { get; }
-    private readonly CancellationToken _cancellationToken;
+    private readonly CancellationTokenSource _cancellationToken;
 
     private readonly IDbContextFactory<OpenShockContext> _dbContextFactory;
     private readonly IRedisConnectionProvider _redisConnectionProvider;
@@ -49,17 +60,15 @@ public sealed class HubLifetime : IAsyncDisposable
     /// <param name="redisConnectionProvider"></param>
     /// <param name="redisPubService"></param>
     /// <param name="logger"></param>
-    /// <param name="cancellationToken"></param>
     public HubLifetime([Range(1, 10)] byte tps, IHubController hubController,
         IDbContextFactory<OpenShockContext> dbContextFactory,
         IRedisConnectionProvider redisConnectionProvider,
         IRedisPubService redisPubService,
-        ILogger<HubLifetime> logger,
-        CancellationToken cancellationToken = default)
+        ILogger<HubLifetime> logger)
     {
         _tps = tps;
         HubController = hubController;
-        _cancellationToken = cancellationToken;
+        _cancellationToken = new CancellationTokenSource();
         _dbContextFactory = dbContextFactory;
         _redisConnectionProvider = redisConnectionProvider;
         _redisPubService = redisPubService;
@@ -72,21 +81,62 @@ public HubLifetime([Range(1, 10)] byte tps, IHubController hubController,
     /// <summary>
     /// Call on creation to setup shockers for the first time
     /// </summary>
-    /// <param name="db"></param>
-    public async Task InitAsync(OpenShockContext db)
+    public async Task InitAsync()
     {
+        await using var db = await _dbContextFactory.CreateDbContextAsync(_cancellationToken.Token);
         await UpdateShockers(db);
 #pragma warning disable CS4014
         LucTask.Run(UpdateLoop);
 #pragma warning restore CS4014
+        
+        _state = HubLifetimeState.Idle; // We are fully setup, we can go back to idle state
+    }
+    
+    /// <summary>
+    /// Swap to a new underlying controller
+    /// </summary>
+    /// <param name="newController"></param>
+    public async Task Swap(IHubController newController)
+    {
+        var oldController = HubController;
+
+        await oldController.DisposeAsync();
+        
+        HubController = newController;
+        await UpdateDevice();
+
+        _state = HubLifetimeState.Idle; // Swap is done, return to (~~monke~~) idle
+    }
+
+    /// <summary>
+    /// Try to mark the lifetime as swapping
+    /// This needs external synchronization
+    /// </summary>
+    /// <returns>true if the lifetime is not busy, false if it is. Consider rejecting the connection</returns>
+    public bool TryMarkSwapping()
+    {
+        if (_state != HubLifetimeState.Idle) return false;
+        _state = HubLifetimeState.Swapping;
+        return true;
+    }
+
+    /// <summary>
+    /// Mark the lifetime as removing
+    /// This needs external synchronization
+    /// </summary>
+    /// <returns>true if the lifetime is not swapping or already removing</returns>
+    public bool TryMarkRemoving()
+    {
+        if (_state is HubLifetimeState.Swapping or HubLifetimeState.Removing) return false;
+        _state = HubLifetimeState.Removing;
+        return true;
     }
 
     private async Task UpdateLoop()
     {
-        var stopwatch = Stopwatch.StartNew();
         while (!_cancellationToken.IsCancellationRequested)
         {
-            stopwatch.Restart();
+            var startingTime = Stopwatch.GetTimestamp();
 
             try
             {
@@ -97,15 +147,16 @@ private async Task UpdateLoop()
                 _logger.LogError(e, "Error in Update()");
             }
 
-            var elapsed = stopwatch.Elapsed;
+
+            var elapsed = Stopwatch.GetElapsedTime(startingTime);
             var waitTime = _waitBetweenTicks - elapsed;
             if (waitTime.TotalMilliseconds < 1)
             {
                 _logger.LogWarning("Update loop running behind for device [{DeviceId}]", HubController.Id);
                 continue;
             }
 
-            await Task.Delay(waitTime, _cancellationToken);
+            await Task.Delay(waitTime, _cancellationToken.Token);
         }
     }
 
@@ -138,7 +189,7 @@ private async Task Update()
     /// </summary>
     public async Task UpdateDevice()
     {
-        await using var db = await _dbContextFactory.CreateDbContextAsync(_cancellationToken);
+        await using var db = await _dbContextFactory.CreateDbContextAsync(_cancellationToken.Token);
         await UpdateShockers(db);
 
         foreach (var websocketController in
@@ -158,7 +209,7 @@ private async Task UpdateShockers(OpenShockContext db)
             Id = x.Id,
             Model = x.Model,
             RfId = x.RfId
-        }).ToDictionaryAsync(x => x.Id, x => x, cancellationToken: _cancellationToken);
+        }).ToDictionaryAsync(x => x.Id, x => x, cancellationToken: _cancellationToken.Token);
     }
 
     /// <summary>
@@ -294,11 +345,17 @@ await deviceOnline.InsertAsync(new DeviceOnline
         return new Success();
     }
 
+    private bool _disposed = false;
+    
     /// <inheritdoc />
-    public ValueTask DisposeAsync() => HubController.DisposeAsync();
-
-    /// <inheritdoc />
-    public ValueTask DisposeForNewConnection() => HubController.DisposeForNewConnection();
+    public async ValueTask DisposeAsync()
+    {
+        if(_disposed) return;
+        _disposed = true;
+        
+        await _cancellationToken.CancelAsync();
+    }
+    
 }
 
 /// <summary>
diff --git a/LiveControlGateway/LifetimeManager/HubLifetimeManager.cs b/LiveControlGateway/LifetimeManager/HubLifetimeManager.cs
@@ -24,11 +24,9 @@ public sealed class HubLifetimeManager
     private readonly IRedisPubService _redisPubService;
     private readonly ILoggerFactory _loggerFactory;
     private readonly ILogger<HubLifetimeManager> _logger;
-    private readonly ConcurrentDictionary<Guid, HubLifetime> _lifetimes = new();
-    private readonly ConcurrentDictionary<Guid, SemaphoreSlim> _hubLocks = new();
-    
-    private readonly HashSet<Guid> _pendingHubs = [];
-    private readonly SemaphoreSlim _pendingHubsLock = new(1);
+    private readonly Dictionary<Guid, HubLifetime> _lifetimes = new();
+
+    private readonly SemaphoreSlim _lifetimesLock = new(1);
 
     /// <summary>
     /// DI constructor
@@ -62,100 +60,99 @@ ILoggerFactory loggerFactory
     public async Task<bool> TryAddDeviceConnection(byte tps, IHubController hubController,
         CancellationToken cancellationToken)
     {
-        // Try finally for _pendingHubs Add <-> Remove scope
-        try
+        var isSwapping = false;
+        HubLifetime? hubLifetime;
+
+        using (await _lifetimesLock.LockAsyncScoped(cancellationToken))
         {
-            using (await _pendingHubsLock.LockAsyncScoped(cancellationToken))
-            {
-                if (!_pendingHubs.Add(hubController.Id)) return false; // Another hub is already queued here
-            }
-            
-            using (await _hubLocks.GetOrAdd(hubController.Id, new SemaphoreSlim(1)).LockAsyncScoped(cancellationToken)) // Any only one thread is allowed here, per hub, and that is what we want.
+            if (_lifetimes.TryGetValue(hubController.Id, out hubLifetime))
             {
-                if (_lifetimes.TryGetValue(hubController.Id, out var oldControllerLifetime))
+                // There already is a hub lifetime, lets swap!
+                if (!hubLifetime.TryMarkSwapping())
                 {
-                    _logger.LogDebug("Disposing old hub controller");
-                    await oldControllerLifetime
-                        .DisposeForNewConnection(); // Use this to not call the remove connection method from the controller
-
-                    foreach (var websocketController in WebsocketManager.LiveControlUsers.GetConnections(hubController
-                                 .Id))
-                        await websocketController.UpdateConnectedState(false);
+                    return
+                        false; // Tell the controller that we are busy right now TODO: Tell the connecting client why it failed
                 }
 
-                var hubLifetime = GetLifetime(tps, hubController, cancellationToken);
+                isSwapping = true;
+            }
+            else
+            {
+                // This is a fresh connection with no existing lifetime, create one!
+                hubLifetime = CreateNewLifetime(tps, hubController);
                 _lifetimes[hubController.Id] = hubLifetime;
+            }
+        }
 
-                await using var db = await _dbContextFactory.CreateDbContextAsync(cancellationToken);
-
-                await hubLifetime.InitAsync(db);
-
-                foreach (var websocketController in WebsocketManager.LiveControlUsers.GetConnections(hubController.Id))
-                    await websocketController.UpdateConnectedState(true);
 
-                return true;
-            }
+        if (isSwapping)
+        {
+            await hubLifetime.Swap(hubController);
         }
-        finally
+        else
         {
-            using (await _pendingHubsLock.LockAsyncScoped(cancellationToken))
-            {
-                _pendingHubs.Remove(hubController.Id);
-            }
+            await hubLifetime.InitAsync();
+
+            foreach (var websocketController in WebsocketManager.LiveControlUsers.GetConnections(hubLifetime
+                         .HubController.Id))
+                await websocketController.UpdateConnectedState(true);
         }
+
+        return true;
     }
 
-    private HubLifetime GetLifetime(byte tps, IHubController hubController, CancellationToken cancellationToken)
+    private HubLifetime CreateNewLifetime(byte tps, IHubController hubController)
     {
-        _logger.LogInformation("New device connected, creating lifetime [{DeviceId}]", hubController.Id);
+        _logger.LogInformation("New hub connected, creating lifetime [{DeviceId}]", hubController.Id);
 
         var deviceLifetime = new HubLifetime(
             tps,
             hubController,
             _dbContextFactory,
             _redisConnectionProvider,
             _redisPubService,
-            _loggerFactory.CreateLogger<HubLifetime>(),
-            cancellationToken);
+            _loggerFactory.CreateLogger<HubLifetime>());
 
         return deviceLifetime;
     }
 
     /// <summary>
-    /// Remove device from Lifetime Manager, called on dispose of device controller, this is the actual end of life of the hub
+    /// Remove device from Lifetime Manager, called on dispose of device controller,
+    /// this is the actual end of life of the hub
     /// </summary>
     /// <param name="hubController"></param>
     public async Task RemoveDeviceConnection(IHubController hubController)
     {
-        if (!_hubLocks.TryGetValue(hubController.Id, out var hubLock))
-        {
-            // Our lock doesnt exist, this shouldn't happen
-            _logger.LogWarning("Hub lock not found for hub [{HubId}]", hubController.Id);
-            return;
-        }
-
-        using (await hubLock.LockAsyncScoped())
+        HubLifetime? hubLifetime;
+        
+        using (await _lifetimesLock.LockAsyncScoped())
         {
-            try
+            if (!_lifetimes.TryGetValue(hubController.Id, out hubLifetime))
             {
-                if(!_lifetimes.TryGetValue(hubController.Id, out var oldControllerLifetime)) return;
-            
-                if(oldControllerLifetime.HubController != hubController) return;
-            
-                _lifetimes.TryRemove(hubController.Id, out _);
+                _logger.LogError("Hub lifetime not found for hub [{HubId}]", hubController.Id);
+                return;
+            }
             
-                foreach (var websocketController in WebsocketManager.LiveControlUsers.GetConnections(hubController.Id))
-                    await websocketController.UpdateConnectedState(false);
+            // Dont remove a hub lifetime that has a different hub controller,
+            // this might happen when remove is called after a swap has been fully done
+            if(hubLifetime.HubController != hubController) return;
+
+            if (hubLifetime.TryMarkRemoving())
+            {
+                return;
             }
-            finally
+        }
+        
+        foreach (var websocketController in WebsocketManager.LiveControlUsers.GetConnections(hubController.Id))
+            await websocketController.UpdateConnectedState(false);
+        
+        await hubLifetime.DisposeAsync();
+        
+        using (await _lifetimesLock.LockAsyncScoped())
+        {
+            if (!_lifetimes.Remove(hubController.Id))
             {
-                using (await _pendingHubsLock.LockAsyncScoped())
-                {
-                    if (!_pendingHubs.Contains(hubController.Id))
-                    {
-                        _hubLocks.TryRemove(hubController.Id, out _);
-                    }
-                }
+                _logger.LogError("Failed to remove hub lifetime [{HubId}], this shouldnt happen WTF?!", hubController.Id);
             }
         }
     }
