Merge pull request #180 from angelobreuer/degradedStateDetection

Implements a degraded state detection for LavalinkClusterNode
angelobreuer · Sep 12, 2024 · c8be66b · c8be66b
2 parents a14dc91 + 1db2661
commit c8be66b
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 3 deletions.
diff --git a/src/Lavalink4NET.Cluster/Nodes/LavalinkClusterNode.cs b/src/Lavalink4NET.Cluster/Nodes/LavalinkClusterNode.cs
@@ -1,4 +1,6 @@
-namespace Lavalink4NET.Cluster;
+using Lavalink4NET.Events;
+
+namespace Lavalink4NET.Cluster;
 
 using System.Collections.Immutable;
 using Lavalink4NET.Clients;
@@ -145,6 +147,9 @@ private async Task RunInternalAsync(CancellationToken shutdownCancellationToken
             var clientInformation = await _readyTask
                 .WaitAsync(shutdownCancellationToken)
                 .ConfigureAwait(false);
+
+            _node.SocketConnectionClosed += OnNodeSocketConnectionClosed;
+            _node.SocketConnectionReady += OnNodeSocketConnectionReady;
 
             var nodeTask = _node.RunAsync(clientInformation, _shutdownCancellationToken).AsTask();
 
@@ -158,11 +163,24 @@ await _node
         }
         finally
         {
+            _node.SocketConnectionClosed -= OnNodeSocketConnectionClosed;
+            _node.SocketConnectionReady -= OnNodeSocketConnectionReady;
+
             var exitStatus = _shutdownCancellationToken.IsCancellationRequested
                 ? LavalinkNodeStatus.OnDemand
                 : LavalinkNodeStatus.Unavailable;
 
             await UpdateStatusAsync(exitStatus, shutdownCancellationToken);
         }
     }
+
+    private async Task OnNodeSocketConnectionReady(object sender, EventArgs eventargs)
+    {
+        await UpdateStatusAsync(LavalinkNodeStatus.Available, _shutdownCancellationToken);
+    }
+
+    private async Task OnNodeSocketConnectionClosed(object sender, ConnectionClosedEventArgs eventargs)
+    {
+        await UpdateStatusAsync(LavalinkNodeStatus.Degraded, _shutdownCancellationToken);
+    }
 }
diff --git a/src/Lavalink4NET/LavalinkNode.cs b/src/Lavalink4NET/LavalinkNode.cs
@@ -39,6 +39,9 @@ internal sealed class LavalinkNode : IAsyncDisposable
     private TaskCompletionSource<string> _readyTaskCompletionSource;
     private Task? _executeTask;
     private bool _disposed;
+
+    internal event AsyncEventHandler<ConnectionClosedEventArgs>? SocketConnectionClosed;
+    internal event AsyncEventHandler<EventArgs>? SocketConnectionReady;
 
     public LavalinkNode(
         LavalinkNodeServiceContext serviceContext,
@@ -167,6 +170,7 @@ await _apiClient
             }
 
             _logger.Ready(Label, SessionId);
+            await SocketConnectionReady.InvokeAsync(this, EventArgs.Empty);
         }
 
         if (SessionId is null)
@@ -491,12 +495,13 @@ private async Task ReceiveInternalAsync(ClientInformation clientInformation, Can
         }
     }
 
-    private Task InvokeConnectionClosedAsync(object sender, ConnectionClosedEventArgs eventArgs)
+    private async Task InvokeConnectionClosedAsync(object sender, ConnectionClosedEventArgs eventArgs)
     {
         ArgumentNullException.ThrowIfNull(sender);
         ArgumentNullException.ThrowIfNull(eventArgs);
 
-        return _serviceContext.NodeListener.OnConnectionClosedAsync(eventArgs).AsTask();
+        await SocketConnectionClosed.InvokeAsync(this, eventArgs);
+        await _serviceContext.NodeListener.OnConnectionClosedAsync(eventArgs).AsTask();
     }
 
     private async ValueTask ReceiveInternalAsync(ILavalinkSocket socket, CancellationToken cancellationToken)