refactor: trying to keep operations that need locking isolated and atomic

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

@@ -5,16 +5,22 @@ 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 Mock<TimeProvider> _mockTimeProvider = new();
   private readonly DiscordClientOptions _options = new();
   private readonly DiscordGatewayClient _discordGatewayClient;
 
   public DiscordGatewayClientTests()
   {
+    _mockTimeProvider
+      .Setup(static x => x.GetUtcNow())
+      .Returns(DateTimeOffset.UtcNow);
+
     _discordGatewayClient = new DiscordGatewayClient(
       _options,
       _mockWebSocketFactory.Object,
       _mockLogger.Object,
-      _mockDiscordRestClient.Object
+      _mockDiscordRestClient.Object,
+      _mockTimeProvider.Object
     );
   }
 
@@ -25,7 +31,8 @@ public sealed class DiscordGatewayClientTests : IDisposable
       null!,
       _mockWebSocketFactory.Object,
       _mockLogger.Object,
-      _mockDiscordRestClient.Object
+      _mockDiscordRestClient.Object,
+      _mockTimeProvider.Object
     );
 
     act.Should().Throw<ArgumentNullException>();
@@ -38,7 +45,8 @@ public sealed class DiscordGatewayClientTests : IDisposable
       _options,
       null!,
       _mockLogger.Object,
-      _mockDiscordRestClient.Object
+      _mockDiscordRestClient.Object,
+      _mockTimeProvider.Object
     );
 
     act.Should().Throw<ArgumentNullException>();
@@ -51,7 +59,8 @@ public sealed class DiscordGatewayClientTests : IDisposable
       _options,
       _mockWebSocketFactory.Object,
       null!,
-      _mockDiscordRestClient.Object
+      _mockDiscordRestClient.Object,
+      _mockTimeProvider.Object
     );
 
     act.Should().Throw<ArgumentNullException>();
@@ -64,6 +73,21 @@ public sealed class DiscordGatewayClientTests : IDisposable
       _options,
       _mockWebSocketFactory.Object,
       _mockLogger.Object,
+      null!,
+      _mockTimeProvider.Object
+    );
+
+    act.Should().Throw<ArgumentNullException>();
+  }
+
+  [Fact]
+  public void Constructor_WhenCalledAndTimeProviderIsNull_ItShouldThrowArgumentNullException()
+  {
+    var act = () => new DiscordGatewayClient(
+      _options,
+      _mockWebSocketFactory.Object,
+      _mockLogger.Object,
+      _mockDiscordRestClient.Object,
       null!
     );
 

src/StevesBot.Worker/Discord/DiscordGatewayClient.cs

@@ -8,6 +8,7 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
   private readonly IWebSocketFactory _webSocketFactory;
   private readonly ILogger<DiscordGatewayClient> _logger;
   private readonly IDiscordRestClient _discordRestClient;
+  private readonly TimeProvider _timeProvider;
   private readonly JsonSerializerOptions _jsonSerializerOptions = new()
   {
     PropertyNamingPolicy = JsonNamingPolicy.CamelCase,
@@ -23,8 +24,8 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
   private IWebSocket? _webSocket;
   private int? _lastSequence;
   private int _heartbeatInterval;
-  private DateTime _timeLastHeartbeatSent = DateTime.MinValue;
-  private DateTime _timeLastHeartbeatAcknowledged = DateTime.MinValue;
+  private DateTimeOffset _timeLastHeartbeatSent = DateTimeOffset.MinValue;
+  private DateTimeOffset _timeLastHeartbeatAcknowledged = DateTimeOffset.MinValue;
   private CancellationTokenSource? _heartbeatCts;
   private CancellationTokenSource? _linkedCts;
   private bool _canResume;
@@ -35,34 +36,34 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
     DiscordClientOptions options,
     IWebSocketFactory webSocketFactory,
     ILogger<DiscordGatewayClient> logger,
-    IDiscordRestClient discordRestClient
+    IDiscordRestClient discordRestClient,
+    TimeProvider timeProvider
   )
   {
     _options = options ?? throw new ArgumentNullException(nameof(options));
     _webSocketFactory = webSocketFactory ?? throw new ArgumentNullException(nameof(webSocketFactory));
     _logger = logger ?? throw new ArgumentNullException(nameof(logger));
     _discordRestClient = discordRestClient ?? throw new ArgumentNullException(nameof(discordRestClient));
+    _timeProvider = timeProvider ?? throw new ArgumentNullException(nameof(timeProvider));
   }
 
   public async Task ConnectAsync(CancellationToken cancellationToken)
   {
-    using (await _lock.LockAsync(cancellationToken))
+    await SetCanResumeAsync(false, cancellationToken);
+
+    if (await IsGatewayUrlSetAsync(cancellationToken) is false)
     {
-      if (string.IsNullOrEmpty(_gatewayUrl))
-      {
-        _gatewayUrl = await _discordRestClient.GetGatewayUrlAsync(cancellationToken);
+      var gatewayUrl = await _discordRestClient.GetGatewayUrlAsync(cancellationToken);
+      await SetGatewayUrlAsync(gatewayUrl, cancellationToken);
     }
 
-      _webSocket = _webSocketFactory.Create();
-
-      var uri = new Uri(_gatewayUrl);
-      await _webSocket.ConnectAsync(uri, cancellationToken);
+    await SetWebSocketAsync(_webSocketFactory.Create(), cancellationToken);
+    await ConnectToGatewayAsync(cancellationToken);
 
     _logger.LogInformation("Connected to Discord Gateway at {GatewayUrl}", _gatewayUrl);
 
     _ = ReceiveMessagesAsync(cancellationToken);
   }
-  }
 
   private async Task ReceiveMessagesAsync(CancellationToken cancellationToken)
   {
@@ -72,21 +73,6 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
 
       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;
-          }
-        }
-
-        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
@@ -97,19 +83,14 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
 
         do
         {
-          result = await _webSocket.ReceiveAsync(new(messageBuffer), cancellationToken);
+          result = await ReceiveMessageAsync(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;
-            }
-
+            var canResume = result.CloseStatus is not WebSocketCloseStatus.NormalClosure or WebSocketCloseStatus.EndpointUnavailable;
+            await SetCanResumeAsync(canResume, cancellationToken);
+            await CloseAsync(result.CloseStatus ?? WebSocketCloseStatus.NormalClosure, result.CloseStatusDescription, cancellationToken);
+            await ReconnectAsync(cancellationToken);
             return;
           }
 
@@ -123,7 +104,6 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
 
         memoryStream.Seek(0, SeekOrigin.Begin);
 
-        var message = Encoding.UTF8.GetString(messageBuffer, 0, result.Count);
         var e = await JsonSerializer.DeserializeAsync<DiscordEvent>(
           memoryStream,
           _jsonSerializerOptions,
@@ -139,48 +119,60 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
         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);
     }
+# 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, "Unexpected error while receiving messages");
-      throw new DiscordGatewayClientException("Unexpected error while receiving messages.", ex);
+
+      using (await _lock.LockAsync(cancellationToken))
+      {
+        if (_webSocket?.State is WebSocketState.Open)
+        {
+          await _webSocket.CloseAsync(
+            WebSocketCloseStatus.NormalClosure,
+            "WebSocket error. Closing connection and invalidating session.",
+            cancellationToken
+          );
+        }
+
+        _canResume = false;
+        await ReconnectAsync(cancellationToken);
+      }
     }
   }
 
   private async Task HandleEventAsync(DiscordEvent e, CancellationToken cancellationToken)
   {
-    using (await _lock.LockAsync(cancellationToken))
-    {
-      _lastSequence = e.Sequence;
-    }
+    await SetSequenceAsync(e.Sequence, cancellationToken);
 
     if (e is HelloDiscordEvent he)
     {
-      using (await _lock.LockAsync(cancellationToken))
-      {
-        _heartbeatInterval = he.Data.HeartbeatInterval;
-      }
-
+      await SetHeartbeatIntervalAsync(he.Data.HeartbeatInterval, cancellationToken);
       await StartHeartbeatAsync(cancellationToken);
+      await IdentifyAsync(cancellationToken);
+
+      _logger.LogInformation("Hello event received. Heartbeat interval: {Interval}", _heartbeatInterval);
       return;
     }
 
     if (e is HeartbeatAckDiscordEvent hae)
     {
-      using (await _lock.LockAsync(cancellationToken))
-      {
-        _timeLastHeartbeatAcknowledged = DateTime.UtcNow;
+      await SetHeartbeatAcknowledgedAsync(_timeProvider.GetUtcNow(), cancellationToken);
+      _logger.LogInformation("Heartbeat acknowledged at {Time}", _timeLastHeartbeatAcknowledged);
+      return;
     }
 
-      _logger.LogInformation("Heartbeat acknowledged at {Time}", _timeLastHeartbeatAcknowledged);
+    if (e is ReadyDiscordEvent re)
+    {
+      await SetSessionIdAsync(re.Data.SessionId, cancellationToken);
+      await SetResumeGatewayUrlAsync(re.Data.ResumeGatewayUrl, cancellationToken);
+
+      _logger.LogInformation("Ready event received. Session ID: {SessionId}", _sessionId);
       return;
     }
   }
@@ -244,12 +236,12 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
                   _canResume = true;
                 }
 
-                await ReconnectAsync(_linkedCts.Token);
+                await ReconnectAsync(cancellationToken);
 
                 break;
               }
 
-              _timeLastHeartbeatSent = await SendHeartbeatAsync(_linkedCts.Token);
+              await SendHeartbeatAsync(_linkedCts.Token);
             }
 
             _logger.LogInformation("Heartbeat sent at {Time}", _timeLastHeartbeatSent);
@@ -263,7 +255,6 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
 # 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);
@@ -290,7 +281,6 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
     {
       _logger.LogInformation("Resuming connection to Discord Gateway.");
 
-      _canResume = false;
       _webSocket = _webSocketFactory.Create();
       var uri = new Uri(_resumeGatewayUrl);
       await _webSocket.ConnectAsync(uri, token);
@@ -320,11 +310,21 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
     await SendJsonAsync(resume, cancellationToken);
   }
 
-  private async Task<DateTime> SendHeartbeatAsync(CancellationToken cancellationToken)
+  private async Task IdentifyAsync(CancellationToken cancellationToken)
+  {
+    var identify = new IdentifyDiscordEvent(
+      _options.AppToken,
+      _options.Intents
+    );
+
+    await SendJsonAsync(identify, cancellationToken);
+  }
+
+  private async Task SendHeartbeatAsync(CancellationToken cancellationToken)
   {
     var heartbeat = new HeartbeatDiscordEvent(_lastSequence);
     await SendJsonAsync(heartbeat, cancellationToken);
-    return DateTime.UtcNow;
+    await SetHeartbeatSentAsync(_timeProvider.GetUtcNow(), cancellationToken);
   }
 
   private async Task SendJsonAsync(object data, CancellationToken cancellationToken)
@@ -353,6 +353,151 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
     }
   }
 
+  private async Task CloseAsync(WebSocketCloseStatus closeStatus, string? statusDescription, CancellationToken cancellationToken)
+  {
+    using (await _lock.LockAsync(cancellationToken))
+    {
+      if (_webSocket is null)
+      {
+        throw new InvalidOperationException("WebSocket is not set. Cannot close.");
+      }
+
+      if (_webSocket.State is not WebSocketState.Open)
+      {
+        throw new InvalidOperationException("WebSocket is not open. Cannot close.");
+      }
+
+      await _webSocket.CloseAsync(closeStatus, statusDescription, cancellationToken);
+    }
+  }
+
+  private async Task<WebSocketReceiveResult> ReceiveMessageAsync(
+    ArraySegment<byte> segment,
+    CancellationToken cancellationToken
+  )
+  {
+    using (await _lock.LockAsync(cancellationToken))
+    {
+      if (_webSocket is null)
+      {
+        throw new InvalidOperationException("WebSocket is not set. Cannot receive message.");
+      }
+
+      if (_webSocket.State is not WebSocketState.Open)
+      {
+        throw new InvalidOperationException("WebSocket is not open. Cannot receive message.");
+      }
+      var result = await _webSocket.ReceiveAsync(segment, cancellationToken);
+      return result;
+    }
+  }
+
+  private async Task ConnectToGatewayAsync(CancellationToken cancellationToken)
+  {
+    using (await _lock.LockAsync(cancellationToken))
+    {
+      if (_webSocket is null)
+      {
+        throw new InvalidOperationException("WebSocket is not set. Cannot connect.");
+      }
+
+      if (_webSocket.State is WebSocketState.Open)
+      {
+        throw new InvalidOperationException("WebSocket is already open. Cannot connect.");
+      }
+
+      await _webSocket.ConnectAsync(new Uri(_gatewayUrl), cancellationToken);
+    }
+  }
+
+  private async Task<bool> IsGatewayUrlSetAsync(CancellationToken cancellationToken)
+  {
+    using (await _lock.LockAsync(cancellationToken))
+    {
+      return !string.IsNullOrEmpty(_gatewayUrl);
+    }
+  }
+
+  private async Task SetGatewayUrlAsync(string url, CancellationToken cancellationToken)
+  {
+    using (await _lock.LockAsync(cancellationToken))
+    {
+      _gatewayUrl = url;
+    }
+  }
+
+  private async Task SetSequenceAsync(int? sequence, CancellationToken cancellationToken)
+  {
+    using (await _lock.LockAsync(cancellationToken))
+    {
+      _lastSequence = sequence;
+    }
+  }
+
+  private async Task SetHeartbeatIntervalAsync(int interval, CancellationToken cancellationToken)
+  {
+    using (await _lock.LockAsync(cancellationToken))
+    {
+      _heartbeatInterval = interval;
+    }
+  }
+
+  private async Task SetHeartbeatSentAsync(DateTimeOffset time, CancellationToken cancellationToken)
+  {
+    using (await _lock.LockAsync(cancellationToken))
+    {
+      _timeLastHeartbeatSent = time;
+    }
+  }
+
+  private async Task SetHeartbeatAcknowledgedAsync(DateTimeOffset time, CancellationToken cancellationToken)
+  {
+    using (await _lock.LockAsync(cancellationToken))
+    {
+      _timeLastHeartbeatAcknowledged = time;
+    }
+  }
+
+  private async Task SetSessionIdAsync(string sessionId, CancellationToken cancellationToken)
+  {
+    using (await _lock.LockAsync(cancellationToken))
+    {
+      _sessionId = sessionId;
+    }
+  }
+
+  private async Task SetResumeGatewayUrlAsync(string url, CancellationToken cancellationToken)
+  {
+    using (await _lock.LockAsync(cancellationToken))
+    {
+      _resumeGatewayUrl = url;
+    }
+  }
+
+  private async Task<bool> IsWebSocketOpenAsync(CancellationToken cancellationToken)
+  {
+    using (await _lock.LockAsync(cancellationToken))
+    {
+      return _webSocket?.State is WebSocketState.Open;
+    }
+  }
+
+  private async Task SetWebSocketAsync(IWebSocket webSocket, CancellationToken cancellationToken)
+  {
+    using (await _lock.LockAsync(cancellationToken))
+    {
+      _webSocket = webSocket;
+    }
+  }
+
+  private async Task SetCanResumeAsync(bool canResume, CancellationToken cancellationToken)
+  {
+    using (await _lock.LockAsync(cancellationToken))
+    {
+      _canResume = canResume;
+    }
+  }
+
   public void Dispose()
   {
     _heartbeatCts?.Cancel();

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

@@ -0,0 +1,40 @@
+namespace StevesBot.Worker.Discord.Events;
+
+internal sealed record IdentifyDiscordEvent : DiscordEvent
+{
+  [JsonPropertyName("d")]
+  public new IdentifyData Data { get; init; } = new IdentifyData();
+
+  public IdentifyDiscordEvent(string token, long intents)
+  {
+    Data = new IdentifyData
+    {
+      Token = token,
+      Intents = intents,
+    };
+  }
+}
+
+internal record IdentifyData
+{
+  [JsonPropertyName("token")]
+  public string Token { get; init; } = string.Empty;
+
+  [JsonPropertyName("properties")]
+  public IdentifyProperties Properties { get; init; } = new IdentifyProperties();
+
+  [JsonPropertyName("intents")]
+  public long Intents { get; init; }
+}
+
+internal record IdentifyProperties
+{
+  [JsonPropertyName("os")]
+  public string Os { get; init; } = Environment.OSVersion.ToString();
+
+  [JsonPropertyName("browser")]
+  public string Browser { get; init; } = Assembly.GetExecutingAssembly().GetName().FullName;
+
+  [JsonPropertyName("device")]
+  public string Device { get; init; } = Assembly.GetExecutingAssembly().GetName().FullName;
+}

src/StevesBot.Worker/Usings.cs

@@ -1,5 +1,6 @@
 global using System.Net.Http.Json;
 global using System.Net.WebSockets;
+global using System.Reflection;
 global using System.Text.Json;
 global using System.Text.Json.Serialization;