feat: begin working on connection management

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

@@ -1,14 +1,14 @@
 namespace StevesBot.Worker.Tests.Unit;
 
-public class DiscordGatewayClientOptionsTests
+public class DiscordClientOptionsTests
 {
   [Fact]
   public void Constructor_WhenCalledWithoutParameters_ItShouldCreateInstance()
   {
-    var options = new DiscordGatewayClientOptions();
+    var options = new DiscordClientOptions();
 
     options.Should().NotBeNull();
-    options.Should().BeOfType<DiscordGatewayClientOptions>();
+    options.Should().BeOfType<DiscordClientOptions>();
     options.ApiUrl.Should().Be(string.Empty);
     options.AppToken.Should().Be(string.Empty);
     options.Intents.Should().Be(0);
@@ -21,7 +21,7 @@ public class DiscordGatewayClientOptionsTests
     var appToken = "test-token";
     var intents = 123;
 
-    var options = new DiscordGatewayClientOptions
+    var options = new DiscordClientOptions
     {
       ApiUrl = apiUrl,
       AppToken = appToken,
@@ -29,7 +29,7 @@ public class DiscordGatewayClientOptionsTests
     };
 
     options.Should().NotBeNull();
-    options.Should().BeOfType<DiscordGatewayClientOptions>();
+    options.Should().BeOfType<DiscordClientOptions>();
     options.ApiUrl.Should().Be(apiUrl);
     options.AppToken.Should().Be(appToken);
     options.Intents.Should().Be(intents);

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

@@ -0,0 +1,98 @@
+namespace StevesBot.Worker.Tests.Unit;
+
+public sealed class DiscordGatewayClientTests : IDisposable
+{
+  private readonly Mock<IDiscordRestClient> _mockDiscordRestClient = new();
+  private readonly Mock<IWebSocketFactory> _mockWebSocketFactory = new();
+  private readonly Mock<ILogger<DiscordGatewayClient>> _mockLogger = new();
+  private readonly DiscordClientOptions _options = new();
+  private readonly DiscordGatewayClient _discordGatewayClient;
+
+  public DiscordGatewayClientTests()
+  {
+    _discordGatewayClient = new DiscordGatewayClient(
+      _options,
+      _mockWebSocketFactory.Object,
+      _mockLogger.Object,
+      _mockDiscordRestClient.Object
+    );
+  }
+
+  [Fact]
+  public async Task ConnectAsync_WhenCalled_ItShouldConnectToGateway()
+  {
+    _mockDiscordRestClient
+      .Setup(static x => x.GetGatewayUrlAsync(It.IsAny<CancellationToken>()))
+      .ReturnsAsync("wss://gateway.discord.gg");
+
+    var mockWebSocket = new Mock<IWebSocket>();
+
+    _mockWebSocketFactory
+      .Setup(static x => x.Create())
+      .Returns(mockWebSocket.Object);
+
+    await _discordGatewayClient.ConnectAsync(CancellationToken.None);
+
+    mockWebSocket
+      .Verify(
+        static x => x.ConnectAsync(It.IsAny<Uri>(), It.IsAny<CancellationToken>()),
+        Times.Once
+      );
+  }
+
+  [Fact]
+  public async Task ConnectAsync_WhenCalledAndWebSocketIsClosed_ItShouldStopReceivingMessages()
+  {
+    _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.Closed);
+
+    _mockWebSocketFactory
+      .Setup(static x => x.Create())
+      .Returns(mockWebSocket.Object);
+
+    await _discordGatewayClient.ConnectAsync(CancellationToken.None);
+
+    mockWebSocket.Verify(
+      static x => x.ReceiveAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<CancellationToken>()),
+      Times.Never
+    );
+  }
+
+  [Fact]
+  public async Task ConnectAsync_WhenCalledAndCancellationIsRequested_ItShouldStopReceivingMessages()
+  {
+    _mockDiscordRestClient
+      .Setup(static x => x.GetGatewayUrlAsync(It.IsAny<CancellationToken>()))
+      .ReturnsAsync("wss://gateway.discord.gg");
+
+    var mockWebSocket = new Mock<IWebSocket>();
+
+    _mockWebSocketFactory
+      .Setup(static x => x.Create())
+      .Returns(mockWebSocket.Object);
+
+    using var cts = new CancellationTokenSource();
+
+    await _discordGatewayClient.ConnectAsync(cts.Token);
+    await Task.Delay(100);
+    await cts.CancelAsync();
+
+    mockWebSocket.Verify(static x => x.State, Times.AtMostOnce);
+    mockWebSocket.Verify(
+      static x => x.ReceiveAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<CancellationToken>()),
+      Times.AtMostOnce
+    );
+  }
+
+  public void Dispose()
+  {
+    _discordGatewayClient.Dispose();
+  }
+}

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

@@ -0,0 +1,34 @@
+namespace StevesBot.Worker.Tests.Unit;
+
+public class DiscordRestClientExceptionTests
+{
+  [Fact]
+  public void Constructor_WhenCalledWithNoParameters_ItShouldCreateInstance()
+  {
+    var exception = new DiscordRestClientException();
+
+    exception.Should().NotBeNull();
+  }
+
+  [Fact]
+  public void Constructor_WhenCalledWithMessage_ItShouldCreateInstance()
+  {
+    var message = "Test message";
+    var exception = new DiscordRestClientException(message);
+
+    exception.Should().NotBeNull();
+    exception.Message.Should().Be(message);
+  }
+
+  [Fact]
+  public void Constructor_WhenCalledWithMessageAndInnerException_ItShouldCreateInstance()
+  {
+    var message = "Test message";
+    var innerException = new Exception("Inner exception");
+    var exception = new DiscordRestClientException(message, innerException);
+
+    exception.Should().NotBeNull();
+    exception.Message.Should().Be(message);
+    exception.InnerException.Should().Be(innerException);
+  }
+}

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

@@ -1,17 +1,9 @@
-using System.Net;
-
-using Microsoft.Extensions.Logging;
-
-using Moq;
-
-using RichardSzalay.MockHttp;
-
 namespace StevesBot.Worker.Tests.Unit;
 
 public sealed class DiscordRestClientTests : IDisposable
 {
   private const string BaseUrl = "https://discord.com/api/v10";
-  private static string GatewayUrl => $"{BaseUrl}/gateway";
+  private static string GatewayEndpoint => $"{BaseUrl}/gateway";
 
   private readonly Mock<ILogger<DiscordRestClient>> _loggerMock = new();
   private readonly MockHttpMessageHandler _mockHttpMessageHandler;
@@ -29,7 +21,7 @@ public sealed class DiscordRestClientTests : IDisposable
   public async Task GetGatewayUrlAsync_WhenRequestFails_ItShouldThrowException()
   {
     _mockHttpMessageHandler
-      .When(GatewayUrl)
+      .When(GatewayEndpoint)
       .Respond(HttpStatusCode.InternalServerError);
 
     var act = async () => await _discordRestClient.GetGatewayUrlAsync(CancellationToken.None);
@@ -41,7 +33,7 @@ public sealed class DiscordRestClientTests : IDisposable
   public async Task GetGatewayUrlAsync_WhenResponseIsNull_ItShouldThrowException()
   {
     _mockHttpMessageHandler
-      .When(GatewayUrl)
+      .When(GatewayEndpoint)
       .Respond(HttpStatusCode.OK, "application/json", "null");
 
     var act = async () => await _discordRestClient.GetGatewayUrlAsync(CancellationToken.None);
@@ -58,7 +50,7 @@ public sealed class DiscordRestClientTests : IDisposable
     }}";
 
     _mockHttpMessageHandler
-      .When(GatewayUrl)
+      .When(GatewayEndpoint)
       .Respond(HttpStatusCode.OK, "application/json", jsonResponse);
 
     var result = await _discordRestClient.GetGatewayUrlAsync(CancellationToken.None);

src/StevesBot.Worker.Tests/Usings.cs

@@ -1,3 +1,4 @@
+global using System.Net;
 global using System.Net.WebSockets;
 global using System.Text;
 global using System.Text.Json;
@@ -9,6 +10,11 @@ global using Microsoft.AspNetCore.Hosting.Server.Features;
 global using Microsoft.AspNetCore.Http.Features;
 global using Microsoft.Extensions.DependencyInjection;
 global using Microsoft.Extensions.Hosting;
+global using Microsoft.Extensions.Logging;
+
+global using Moq;
+
+global using RichardSzalay.MockHttp;
 
 global using StevesBot.Worker.Discord;
 global using StevesBot.Worker.Discord.Events;

src/StevesBot.Worker/Discord/DiscordClientOptions.cs

@@ -1,6 +1,6 @@
 namespace StevesBot.Worker.Discord;
 
-internal sealed class DiscordGatewayClientOptions
+internal sealed class DiscordClientOptions
 {
   public string ApiUrl { get; init; } = string.Empty;
   public string AppToken { get; init; } = string.Empty;

src/StevesBot.Worker/Discord/DiscordGatewayClient.cs

@@ -1,34 +1,37 @@
+using System.Text;
+
 namespace StevesBot.Worker.Discord;
 
 internal class DiscordGatewayClient : IDiscordGatewayClient
 {
-  private readonly DiscordGatewayClientOptions _options;
+  private readonly DiscordClientOptions _options;
   private readonly IWebSocketFactory _webSocketFactory;
   private readonly ILogger<DiscordGatewayClient> _logger;
   private readonly IDiscordRestClient _discordRestClient;
-  // private readonly JsonSerializerOptions _jsonSerializerOptions = new()
+  private readonly JsonSerializerOptions _jsonSerializerOptions = new()
-  // {
+  {
-  //   PropertyNamingPolicy = JsonNamingPolicy.CamelCase,
+    PropertyNamingPolicy = JsonNamingPolicy.CamelCase,
-  //   ReferenceHandler = ReferenceHandler.IgnoreCycles,
+    ReferenceHandler = ReferenceHandler.IgnoreCycles,
-  //   Converters =
+    Converters =
-  //   {
+    {
-  //     new DiscordEventConverter(),
+      new DiscordEventConverter(),
-  //   },
+    },
-  // };
+  };
   private readonly AsyncLock _lock = new();
 
   private string _gatewayUrl = string.Empty;
   private IWebSocket? _webSocket;
   private Task? _receiveTask;
-  // private DateTime _timeLastHeartbeatSent = DateTime.MinValue;
+  private DateTime _timeLastHeartbeatSent = DateTime.MinValue;
-  // private DateTime _timeLastHeartbeatAcknowledged = DateTime.MinValue;
+  private DateTime _timeLastHeartbeatAcknowledged = DateTime.MinValue;
-  // private CancellationTokenSource? _heartbeatCts;
+  private CancellationTokenSource? _heartbeatCts;
-  // private Task? _heartbeatTask;
+  private CancellationTokenSource? _linkedCts;
+  private Task? _heartbeatTask;
   // private string _sessionId = string.Empty;
   // private string _resumeGatewayUrl = string.Empty;
 
   public DiscordGatewayClient(
-    DiscordGatewayClientOptions options,
+    DiscordClientOptions options,
     IWebSocketFactory webSocketFactory,
     ILogger<DiscordGatewayClient> logger,
     IDiscordRestClient discordRestClient
@@ -60,16 +63,221 @@ internal class DiscordGatewayClient : IDiscordGatewayClient
     }
   }
 
-  private static Task ReceiveMessagesAsync(CancellationToken cancellationToken)
+  private async Task ReceiveMessagesAsync(CancellationToken cancellationToken)
   {
-    return Task.CompletedTask;
+    try
+    {
+      var messageBuffer = new byte[8192];
+
+      while (cancellationToken.IsCancellationRequested is false)
+      {
+        using (await _lock.LockAsync(cancellationToken))
+        {
+          if (_webSocket?.State is not WebSocketState.Open)
+          {
+            _logger.LogWarning("WebSocket is not open. Cannot receive messages.");
+            return;
+          }
+        }
+
+        // websocket message might be larger than the
+        // size of the buffer so we need to loop until
+        // we receive the end of the message and
+        // write the data we receive on each iteration
+        // to the memory stream
+        using var memoryStream = new MemoryStream();
+        WebSocketReceiveResult result;
+
+        do
+        {
+          result = await _webSocket.ReceiveAsync(new(messageBuffer), cancellationToken);
+
+          if (result.MessageType is WebSocketMessageType.Close)
+          {
+            // TODO: If close status is 1000 or 1001 we cannot resume.
+            // if 1000 or 1001 we should close the connection and reconnect
+            // else we should close the connection and attempt to resume
+
+            if (result.CloseStatus is WebSocketCloseStatus.NormalClosure or WebSocketCloseStatus.EndpointUnavailable)
+            {
+              return;
+            }
+
+            return;
+          }
+
+          if (result.MessageType is WebSocketMessageType.Text)
+          {
+# pragma warning disable CA1849
+            memoryStream.Write(messageBuffer, 0, result.Count);
+          }
+
+        } while (result.EndOfMessage is false);
+
+        memoryStream.Seek(0, SeekOrigin.Begin);
+
+        var message = Encoding.UTF8.GetString(messageBuffer, 0, result.Count);
+        var e = await JsonSerializer.DeserializeAsync<DiscordEvent>(
+          memoryStream,
+          _jsonSerializerOptions,
+          cancellationToken
+        );
+
+        if (e is null)
+        {
+          _logger.LogInformation("Received null event.");
+          continue;
+        }
+
+        await HandleEventAsync(e, cancellationToken);
+      }
+    }
+    catch (WebSocketException ex)
+    {
+      _logger.LogError(ex, "WebSocket error");
+      throw new DiscordGatewayClientException("WebSocket error.", ex);
+    }
+    catch (OperationCanceledException ex)
+    {
+      _logger.LogInformation(ex, "Receive messages operation canceled: {Message}", ex.Message);
+    }
+    catch (Exception ex)
+    {
+      _logger.LogError(ex, "Unexpected error while receiving messages");
+      throw new DiscordGatewayClientException("Unexpected error while receiving messages.", ex);
+    }
+  }
+
+  private async Task HandleEventAsync(DiscordEvent e, CancellationToken cancellationToken)
+  {
+    if (e is HelloDiscordEvent he)
+    {
+      await StartHeartbeatAsync(he, cancellationToken);
+      return;
+    }
+
+    if (e is HeartbeatAckDiscordEvent hae)
+    {
+      using (await _lock.LockAsync(cancellationToken))
+      {
+        _timeLastHeartbeatAcknowledged = DateTime.UtcNow;
+      }
+
+      _logger.LogInformation("Heartbeat acknowledged at {Time}", _timeLastHeartbeatAcknowledged);
+      return;
+    }
+  }
+
+  private async Task StartHeartbeatAsync(HelloDiscordEvent helloEvent, CancellationToken cancellationToken)
+  {
+    using (await _lock.LockAsync(cancellationToken))
+    {
+      if (_heartbeatCts is not null)
+      {
+        await _heartbeatCts.CancelAsync();
+        _heartbeatCts.Dispose();
+      }
+
+      _heartbeatCts = new CancellationTokenSource();
+      _linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, _heartbeatCts.Token);
+
+      _heartbeatTask = Task.Run(async () =>
+      {
+        _logger.LogInformation("Starting heartbeat task.");
+
+        while (_linkedCts.Token.IsCancellationRequested is false)
+        {
+          using (await _lock.LockAsync(_linkedCts.Token))
+          {
+            if (_webSocket?.State is not WebSocketState.Open)
+            {
+              _logger.LogWarning("WebSocket is not open. Cannot send heartbeat.");
+              break;
+            }
+          }
+
+          try
+          {
+# 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);
+
+            using (await _lock.LockAsync(_linkedCts.Token))
+            {
+              if (_timeLastHeartbeatAcknowledged < _timeLastHeartbeatSent)
+              {
+                if (_webSocket?.State is WebSocketState.Open)
+                {
+                  _logger.LogWarning("Heartbeat not acknowledged. Closing WebSocket.");
+                  await _webSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Heartbeat not acknowledged", CancellationToken.None);
+                }
+                break;
+              }
+
+              _timeLastHeartbeatSent = await SendHeartbeatAsync(helloEvent.Sequence, _linkedCts.Token);
+            }
+
+            _logger.LogInformation("Heartbeat sent at {Time}", _timeLastHeartbeatSent);
+          }
+          catch (OperationCanceledException ex)
+          {
+            _logger.LogInformation(ex, "Heartbeat task canceled");
+          }
+# pragma warning disable CA1031 // Do not catch general exception types
+          catch (Exception ex)
+# pragma warning restore CA1031 // Do not catch general exception types
+          {
+            _logger.LogError(ex, "Error in heartbeat task: {Message}", ex.Message);
+            // TODO: Attempt to reconnect
+          }
+        }
+      }, _linkedCts.Token);
+    }
+  }
+
+  private async Task<DateTime> SendHeartbeatAsync(int? sequence, CancellationToken cancellationToken)
+  {
+    var heartbeat = new HeartbeatDiscordEvent(sequence);
+    await SendJsonAsync(heartbeat, cancellationToken);
+    return DateTime.UtcNow;
+  }
+
+  private async Task SendJsonAsync(object data, CancellationToken cancellationToken)
+  {
+    using (await _lock.LockAsync(cancellationToken))
+    {
+      if (_webSocket?.State is not WebSocketState.Open)
+      {
+        _logger.LogWarning("WebSocket is not open. Cannot send message.");
+        return;
+      }
+    }
+
+    try
+    {
+      var json = JsonSerializer.Serialize(data, _jsonSerializerOptions);
+      var bytes = Encoding.UTF8.GetBytes(json);
+      var buffer = new ArraySegment<byte>(bytes);
+      await _webSocket.SendAsync(buffer, WebSocketMessageType.Text, true, cancellationToken);
+    }
+    catch (OperationCanceledException ex)
+    {
+      _logger.LogInformation(ex, "Send operation canceled");
+    }
+    catch (Exception ex)
+    {
+      _logger.LogError(ex, "Failed to send message: {Message}", ex.Message);
+      throw new DiscordGatewayClientException("Failed to send message.", ex);
+    }
   }
 
   public void Dispose()
   {
-    // _heartbeatCts?.Cancel();
+    _heartbeatCts?.Cancel();
-    // _heartbeatCts?.Dispose();
+    _heartbeatCts?.Dispose();
-    // _heartbeatTask?.Dispose();
+    _linkedCts?.Dispose();
+    _heartbeatTask?.Dispose();
     _receiveTask?.Dispose();
     _webSocket?.Dispose();
     _lock.Dispose();