feat: implement reconnecting

2025-05-15 00:02:40 -05:00
parent 16aa4b5775
commit 499b1e3d4b
5 changed files with 319 additions and 26 deletions
@@ -128,7 +128,7 @@ csharp_style_prefer_top_level_statements = true:silent

 # Expression-level preferences
 csharp_prefer_simple_default_expression = true:suggestion
-csharp_style_deconstructed_variable_declaration = true:suggestion
+csharp_style_deconstructed_variable_declaration = false:silent
 csharp_style_implicit_object_creation_when_type_is_apparent = true:suggestion
 csharp_style_inlined_variable_declaration = true:suggestion
 csharp_style_prefer_index_operator = true:suggestion
@@ -6,7 +6,6 @@ public sealed class DiscordGatewayClientTests : IDisposable
  private readonly Mock<IWebSocketFactory> _mockWebSocketFactory = new();
  private readonly Mock<ILogger<DiscordGatewayClient>> _mockLogger = new();
  private readonly DiscordClientOptions _options = new();
-# pragma warning disable CA2213
  private readonly DiscordGatewayClient _discordGatewayClient;

  public DiscordGatewayClientTests()
@@ -19,6 +18,58 @@ public sealed class DiscordGatewayClientTests : IDisposable
    );
  }

+  [Fact]
+  public void Constructor_WhenCalledAndOptionsIsNull_ItShouldThrowArgumentNullException()
+  {
+    var act = () => new DiscordGatewayClient(
+      null!,
+      _mockWebSocketFactory.Object,
+      _mockLogger.Object,
+      _mockDiscordRestClient.Object
+    );
+
+    act.Should().Throw<ArgumentNullException>();
+  }
+
+  [Fact]
+  public void Constructor_WhenCalledAndWebSocketFactoryIsNull_ItShouldThrowArgumentNullException()
+  {
+    var act = () => new DiscordGatewayClient(
+      _options,
+      null!,
+      _mockLogger.Object,
+      _mockDiscordRestClient.Object
+    );
+
+    act.Should().Throw<ArgumentNullException>();
+  }
+
+  [Fact]
+  public void Constructor_WhenCalledAndLoggerIsNull_ItShouldThrowArgumentNullException()
+  {
+    var act = () => new DiscordGatewayClient(
+      _options,
+      _mockWebSocketFactory.Object,
+      null!,
+      _mockDiscordRestClient.Object
+    );
+
+    act.Should().Throw<ArgumentNullException>();
+  }
+
+  [Fact]
+  public void Constructor_WhenCalledAndDiscordRestClientIsNull_ItShouldThrowArgumentNullException()
+  {
+    var act = () => new DiscordGatewayClient(
+      _options,
+      _mockWebSocketFactory.Object,
+      _mockLogger.Object,
+      null!
+    );
+
+    act.Should().Throw<ArgumentNullException>();
+  }
+
  [Fact]
  public async Task ConnectAsync_WhenCalled_ItShouldConnectToGateway()
  {
@@ -107,7 +158,7 @@ public sealed class DiscordGatewayClientTests : IDisposable

    var messageQueue = new Queue<(WebSocketReceiveResult, byte[])>();

-    var heartbeatInterval = 1000;
+    var heartbeatInterval = 100;

    var helloEventPayload = CreateEventPayload(new
    {
@@ -148,6 +199,145 @@ public sealed class DiscordGatewayClientTests : IDisposable
    );
  }

+  [Fact]
+  public async Task ConnectAsync_WhenCalledAndHeartbeatIsNotAcknowledged_ItShouldDisconnectAndAttemptToResume()
+  {
+    _mockDiscordRestClient
+      .Setup(static x => x.GetGatewayUrlAsync(It.IsAny<CancellationToken>()))
+      .ReturnsAsync("wss://gateway.discord.gg");
+
+    var mockWebSocket = new Mock<IWebSocket>();
+
+    mockWebSocket
+      .Setup(static x => x.State)
+      .Returns(WebSocketState.Open);
+
+    var messageQueue = new Queue<(WebSocketReceiveResult, byte[])>();
+
+    var heartbeatInterval = 100;
+
+    var helloEventPayload = CreateEventPayload(new
+    {
+      op = 10,
+      d = new
+      {
+        heartbeat_interval = heartbeatInterval,
+      }
+    });
+
+    messageQueue.Enqueue((
+      new(helloEventPayload.Bytes.Length, WebSocketMessageType.Text, true),
+      helloEventPayload.Bytes
+    ));
+
+    SetupReceiveMessageSequence(mockWebSocket, messageQueue);
+
+    _mockWebSocketFactory
+      .Setup(static x => x.Create())
+      .Returns(mockWebSocket.Object);
+
+    var cts = new CancellationTokenSource();
+    await _discordGatewayClient.ConnectAsync(cts.Token);
+
+    await Task.Delay((int)(heartbeatInterval * 2.5));
+    await cts.CancelAsync();
+
+    var expectedHeartbeatPayload = CreateEventPayload(new HeartbeatDiscordEvent(null));
+
+    mockWebSocket.Verify(
+      x => x.SendAsync(
+        It.Is<ArraySegment<byte>>(b => expectedHeartbeatPayload.Bytes.SequenceEqual(b)),
+        It.IsAny<WebSocketMessageType>(),
+        It.IsAny<bool>(),
+        It.IsAny<CancellationToken>()
+      ),
+      Times.Once
+    );
+
+    mockWebSocket.Verify(
+      static x => x.CloseAsync(
+        It.Is<WebSocketCloseStatus>(
+          x => x != WebSocketCloseStatus.NormalClosure && x != WebSocketCloseStatus.EndpointUnavailable
+        ),
+        It.IsAny<string>(),
+        It.IsAny<CancellationToken>()
+      ),
+      Times.Once
+    );
+  }
+
+  [Fact]
+  public async Task ConnectAsync_WhenCalledAndHeartbeatNotAcknowlegedAndAlreadyClosed_ItShouldNotDisconnect()
+  {
+    _mockDiscordRestClient
+      .Setup(static x => x.GetGatewayUrlAsync(It.IsAny<CancellationToken>()))
+      .ReturnsAsync("wss://gateway.discord.gg");
+
+    var mockWebSocket = new Mock<IWebSocket>();
+
+    mockWebSocket
+      .SetupSequence(static x => x.State)
+      .Returns(WebSocketState.Open)
+      .Returns(WebSocketState.Open)
+      .Returns(WebSocketState.Open)
+      .Returns(WebSocketState.Open)
+      .Returns(WebSocketState.Open)
+      .Returns(WebSocketState.Open)
+      .Returns(WebSocketState.Open)
+      .Returns(WebSocketState.Closed);
+
+    var messageQueue = new Queue<(WebSocketReceiveResult, byte[])>();
+
+    var heartbeatInterval = 100;
+
+    var helloEventPayload = CreateEventPayload(new
+    {
+      op = 10,
+      d = new
+      {
+        heartbeat_interval = heartbeatInterval,
+      }
+    });
+
+    messageQueue.Enqueue((
+      new(helloEventPayload.Bytes.Length, WebSocketMessageType.Text, true),
+      helloEventPayload.Bytes
+    ));
+
+    SetupReceiveMessageSequence(mockWebSocket, messageQueue);
+
+    _mockWebSocketFactory
+      .Setup(static x => x.Create())
+      .Returns(mockWebSocket.Object);
+
+    var cts = new CancellationTokenSource();
+    await _discordGatewayClient.ConnectAsync(cts.Token);
+
+    await Task.Delay((int)(heartbeatInterval * 2.5));
+    await cts.CancelAsync();
+
+    var expectedHeartbeatPayload = CreateEventPayload(new HeartbeatDiscordEvent(null));
+
+    mockWebSocket.Verify(
+      x => x.SendAsync(
+        It.Is<ArraySegment<byte>>(b => expectedHeartbeatPayload.Bytes.SequenceEqual(b)),
+        It.IsAny<WebSocketMessageType>(),
+        It.IsAny<bool>(),
+        It.IsAny<CancellationToken>()
+      ),
+      Times.Once
+    );
+
+    mockWebSocket.Verify(
+      static x => x.CloseAsync(
+        It.IsAny<WebSocketCloseStatus>(),
+        It.IsAny<string>(),
+        It.IsAny<CancellationToken>()
+      ),
+      Times.Never
+    );
+  }
+
  private static void SetupReceiveMessageSequence(
    Mock<IWebSocket> mockWebSocket,
    Queue<(WebSocketReceiveResult, byte[])> messageQueue
@@ -21,14 +21,15 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient

  private string _gatewayUrl = string.Empty;
  private IWebSocket? _webSocket;
-  private Task? _receiveTask;
+  private int? _lastSequence;
+  private int _heartbeatInterval;
  private DateTime _timeLastHeartbeatSent = DateTime.MinValue;
  private DateTime _timeLastHeartbeatAcknowledged = DateTime.MinValue;
  private CancellationTokenSource? _heartbeatCts;
  private CancellationTokenSource? _linkedCts;
-  private Task? _heartbeatTask;
-  // private string _sessionId = string.Empty;
-  // private string _resumeGatewayUrl = string.Empty;
+  private bool _canResume;
+  private string _sessionId = string.Empty;
+  private string _resumeGatewayUrl = string.Empty;

  public DiscordGatewayClient(
    DiscordClientOptions options,
@@ -59,7 +60,7 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient

      _logger.LogInformation("Connected to Discord Gateway at {GatewayUrl}", _gatewayUrl);

-      _receiveTask = ReceiveMessagesAsync(cancellationToken);
+      _ = ReceiveMessagesAsync(cancellationToken);
    }
  }

@@ -156,9 +157,19 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient

  private async Task HandleEventAsync(DiscordEvent e, CancellationToken cancellationToken)
  {
+    using (await _lock.LockAsync(cancellationToken))
+    {
+      _lastSequence = e.Sequence;
+    }
+
    if (e is HelloDiscordEvent he)
    {
-      await StartHeartbeatAsync(he, cancellationToken);
+      using (await _lock.LockAsync(cancellationToken))
+      {
+        _heartbeatInterval = he.Data.HeartbeatInterval;
+      }
+
+      await StartHeartbeatAsync(cancellationToken);
      return;
    }

@@ -174,7 +185,7 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
    }
  }

-  private async Task StartHeartbeatAsync(HelloDiscordEvent helloEvent, CancellationToken cancellationToken)
+  private async Task StartHeartbeatAsync(CancellationToken cancellationToken)
  {
    using (await _lock.LockAsync(cancellationToken))
    {
@@ -187,7 +198,7 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
      _heartbeatCts = new CancellationTokenSource();
      _linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, _heartbeatCts.Token);

-      _heartbeatTask = Task.Run(async () =>
+      _ = Task.Run(async () =>
      {
        _logger.LogInformation("Starting heartbeat task.");

@@ -207,7 +218,14 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
 # pragma warning disable CA5394 // Do not use insecure randomness
            var jitter = Random.Shared.NextDouble();
 # pragma warning restore CA5394 // Do not use insecure randomness
-            await Task.Delay((int)(helloEvent.Data.HeartbeatInterval + jitter), _linkedCts.Token);
+            int heartbeatInterval;
+
+            using (await _lock.LockAsync(_linkedCts.Token))
+            {
+              heartbeatInterval = _heartbeatInterval + (int)(jitter * _heartbeatInterval);
+            }
+
+            await Task.Delay(heartbeatInterval, _linkedCts.Token);

            using (await _lock.LockAsync(_linkedCts.Token))
            {
@@ -216,12 +234,22 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
                if (_webSocket?.State is WebSocketState.Open)
                {
                  _logger.LogWarning("Heartbeat not acknowledged. Closing WebSocket.");
-                  await _webSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Heartbeat not acknowledged", CancellationToken.None);
+
+                  await _webSocket.CloseAsync(
+                    WebSocketCloseStatus.ProtocolError,
+                    "Heartbeat not acknowledged",
+                    _linkedCts.Token
+                  );
+
+                  _canResume = true;
                }
+
+                await ReconnectAsync(_linkedCts.Token);
+
                break;
              }

-              _timeLastHeartbeatSent = await SendHeartbeatAsync(helloEvent.Sequence, _linkedCts.Token);
+              _timeLastHeartbeatSent = await SendHeartbeatAsync(_linkedCts.Token);
            }

            _logger.LogInformation("Heartbeat sent at {Time}", _timeLastHeartbeatSent);
@@ -242,9 +270,59 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
    }
  }

-  private async Task<DateTime> SendHeartbeatAsync(int? sequence, CancellationToken cancellationToken)
+  private async Task ReconnectAsync(CancellationToken token)
  {
-    var heartbeat = new HeartbeatDiscordEvent(sequence);
+    _webSocket?.Dispose();
+
+    if (_heartbeatCts is not null)
+    {
+      await _heartbeatCts.CancelAsync();
+      _heartbeatCts.Dispose();
+    }
+
+    if (_linkedCts is not null)
+    {
+      await _linkedCts.CancelAsync();
+      _linkedCts.Dispose();
+    }
+
+    if (_canResume)
+    {
+      _logger.LogInformation("Resuming connection to Discord Gateway.");
+
+      _canResume = false;
+      _webSocket = _webSocketFactory.Create();
+      var uri = new Uri(_resumeGatewayUrl);
+      await _webSocket.ConnectAsync(uri, token);
+      await SendResumeAsync(token);
+      _ = ReceiveMessagesAsync(token);
+      _ = StartHeartbeatAsync(token);
+      return;
+    }
+
+    _logger.LogInformation("Reconnecting to Discord Gateway.");
+    await ConnectAsync(token);
+  }
+
+  private async Task SendResumeAsync(CancellationToken cancellationToken)
+  {
+    if (_lastSequence is null)
+    {
+      throw new InvalidOperationException("Cannot resume without a sequence number.");
+    }
+
+    var resume = new ResumeDiscordEvent(
+      _options.AppToken,
+      _sessionId,
+      _lastSequence.Value
+    );
+
+    await SendJsonAsync(resume, cancellationToken);
+  }
+
+  private async Task<DateTime> SendHeartbeatAsync(CancellationToken cancellationToken)
+  {
+    var heartbeat = new HeartbeatDiscordEvent(_lastSequence);
    await SendJsonAsync(heartbeat, cancellationToken);
    return DateTime.UtcNow;
  }
@@ -280,16 +358,6 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
    _heartbeatCts?.Cancel();
    _linkedCts?.Cancel();

-    if (_heartbeatTask is not null && _heartbeatTask.IsCompleted)
-    {
-      _heartbeatTask.Dispose();
-    }
-
-    if (_receiveTask is not null && _receiveTask.IsCompleted)
-    {
-      _receiveTask.Dispose();
-    }
-
    _heartbeatCts?.Dispose();
    _linkedCts?.Dispose();
    _webSocket?.Dispose();
@@ -4,6 +4,7 @@ internal static class DiscordOpCodes
 {
  public const int Dispatch = 0;
  public const int Heartbeat = 1;
+  public const int Resume = 6;
  public const int HeartbeatAck = 11;
  public const int Hello = 10;
 }
@@ -0,0 +1,34 @@
+namespace StevesBot.Worker.Discord.Events;
+
+internal sealed record ResumeDiscordEvent : DiscordEvent
+{
+  [JsonPropertyName("d")]
+  public new ResumeData Data { get; init; } = new();
+
+  public ResumeDiscordEvent(
+    string token,
+    string sessionId,
+    int sequence
+  )
+  {
+    OpCode = DiscordOpCodes.Resume;
+    Data = new ResumeData
+    {
+      Token = token,
+      SessionId = sessionId,
+      Sequence = sequence
+    };
+  }
+}
+
+internal sealed record ResumeData
+{
+  [JsonPropertyName("token")]
+  public string Token { get; init; } = string.Empty;
+
+  [JsonPropertyName("session_id")]
+  public string SessionId { get; init; } = string.Empty;
+
+  [JsonPropertyName("seq")]
+  public int Sequence { get; init; }
+}