feat: implement reconnecting

.editorconfig

@@ -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

src/StevesBot.Worker.Tests/Unit/DiscordGatewayClientTests.cs

@@ -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

src/StevesBot.Worker/Discord/DiscordGatewayClient.cs

@@ -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 bool _canResume;
-  // private string _sessionId = string.Empty;
+  private string _sessionId = string.Empty;
-  // private string _resumeGatewayUrl = 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();

src/StevesBot.Worker/Discord/Events/DiscordOpCodes.cs

@@ -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;
 }

src/StevesBot.Worker/Discord/Events/ResumeDiscordEvent.cs

@@ -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; }
+}