feat: it's working!!!

src/StevesBot.Worker/Discord/DiscordGatewayClient.cs

@@ -26,8 +26,10 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
   private int _heartbeatInterval;
   private DateTimeOffset _timeLastHeartbeatSent = DateTimeOffset.MinValue;
   private DateTimeOffset _timeLastHeartbeatAcknowledged = DateTimeOffset.MinValue;
+  private CancellationTokenSource? _receiveMessageCts;
+  private CancellationTokenSource? _linkedReceiveMessageCts;
   private CancellationTokenSource? _heartbeatCts;
-  private CancellationTokenSource? _linkedCts;
+  private CancellationTokenSource? _linkedHeartbeatCts;
   private bool _canResume;
   private string _sessionId = string.Empty;
   private string _resumeGatewayUrl = string.Empty;
@@ -51,27 +53,46 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
   {
     await SetCanResumeAsync(false, cancellationToken);
 
-    if (await IsGatewayUrlSetAsync(cancellationToken) is false)
+    if (string.IsNullOrWhiteSpace(_gatewayUrl))
     {
       var gatewayUrl = await _discordRestClient.GetGatewayUrlAsync(cancellationToken);
       await SetGatewayUrlAsync(gatewayUrl, cancellationToken);
     }
 
     await SetWebSocketAsync(_webSocketFactory.Create(), cancellationToken);
-    await ConnectToGatewayAsync(cancellationToken);
+    await ConnectWithGatewayUrlAsync(cancellationToken);
 
     _logger.LogInformation("Connected to Discord Gateway at {GatewayUrl}", _gatewayUrl);
 
-    _ = ReceiveMessagesAsync(cancellationToken);
+    await StartReceiveMessagesAsync(cancellationToken);
   }
 
-  private async Task ReceiveMessagesAsync(CancellationToken cancellationToken)
+  public async Task DisconnectAsync(CancellationToken cancellationToken)
   {
+    await CancelReceiveMessagesTaskAsync(cancellationToken);
+    await CancelHeartbeatTaskAsync(cancellationToken);
+    await CloseIfOpenAsync(WebSocketCloseStatus.NormalClosure, "Disconnecting", cancellationToken);
+  }
+
+  private async Task StartReceiveMessagesAsync(CancellationToken cancellationToken)
+  {
+    await CancelReceiveMessagesTaskAsync(cancellationToken);
+
+    var newReceiveCts = new CancellationTokenSource();
+    var newReceiveLinkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, newReceiveCts.Token);
+
+    await SetReceiveCtsAsync(newReceiveCts, cancellationToken);
+    await SetLinkedReceiveCtsAsync(newReceiveLinkedCts, cancellationToken);
+
+    _ = Task.Run(async () =>
+    {
+      _logger.LogInformation("Starting receive messages task.");
+
       try
       {
         var messageBuffer = new byte[8192];
 
-      while (cancellationToken.IsCancellationRequested is false)
+        while (_linkedReceiveMessageCts?.IsCancellationRequested is false)
         {
           // websocket message might be larger than the
           // size of the buffer so we need to loop until
@@ -83,21 +104,27 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
 
           do
           {
-          result = await ReceiveMessageAsync(new(messageBuffer), cancellationToken);
+            result = await ReceiveMessageAsync(new(messageBuffer), _linkedReceiveMessageCts.Token);
 
             if (result.MessageType is WebSocketMessageType.Close)
             {
-            var canResume = result.CloseStatus is not WebSocketCloseStatus.NormalClosure or WebSocketCloseStatus.EndpointUnavailable;
+              // TODO: Handle all possible close codes
-            await SetCanResumeAsync(canResume, cancellationToken);
+              // reconnect accordingly to documentation
-            await CloseAsync(result.CloseStatus ?? WebSocketCloseStatus.NormalClosure, result.CloseStatusDescription, cancellationToken);
+              var canResume = result.CloseStatus is not WebSocketCloseStatus.NormalClosure or WebSocketCloseStatus.EndpointUnavailable
+                && string.IsNullOrWhiteSpace(_resumeGatewayUrl) is false
+                && string.IsNullOrWhiteSpace(_sessionId) is false;
+
+              await SetCanResumeAsync(canResume, _linkedReceiveMessageCts.Token);
+              await CloseIfOpenAsync(result.CloseStatus ?? WebSocketCloseStatus.NormalClosure, result.CloseStatusDescription, _linkedReceiveMessageCts.Token);
+              _logger.LogInformation("Reconnecting because of close message: {CloseStatus} - {CloseStatusDescription}", result.CloseStatus, result.CloseStatusDescription);
               await ReconnectAsync(cancellationToken);
               return;
             }
 
             if (result.MessageType is WebSocketMessageType.Text)
             {
-            await memoryStream.WriteAsync(messageBuffer.AsMemory(0, result.Count), cancellationToken);
+              await memoryStream.WriteAsync(messageBuffer.AsMemory(0, result.Count), _linkedReceiveMessageCts.Token);
-            await memoryStream.FlushAsync(cancellationToken);
+              await memoryStream.FlushAsync(_linkedReceiveMessageCts.Token);
             }
 
           } while (result.EndOfMessage is false);
@@ -107,7 +134,7 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
           var e = await JsonSerializer.DeserializeAsync<DiscordEvent>(
             memoryStream,
             _jsonSerializerOptions,
-          cancellationToken
+            _linkedReceiveMessageCts.Token
           );
 
           if (e is null)
@@ -116,7 +143,7 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
             continue;
           }
 
-        await HandleEventAsync(e, cancellationToken);
+          await HandleEventAsync(e, _linkedReceiveMessageCts.Token);
         }
       }
       catch (OperationCanceledException ex)
@@ -129,122 +156,94 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
       {
         _logger.LogError(ex, "Unexpected error while receiving messages");
 
-      using (await _lock.LockAsync(cancellationToken))
+        _logger.LogWarning("Closing connection and invalidating session.");
-      {
+
-        if (_webSocket?.State is WebSocketState.Open)
+        await CloseIfOpenAsync(
-        {
-          await _webSocket.CloseAsync(
           WebSocketCloseStatus.NormalClosure,
           "WebSocket error. Closing connection and invalidating session.",
-            cancellationToken
+          newReceiveLinkedCts.Token
         );
-        }
 
-        _canResume = false;
+        await SetCanResumeAsync(false, newReceiveLinkedCts.Token);
+        _logger.LogInformation("Reconnecting because of error in receive message task");
         await ReconnectAsync(cancellationToken);
       }
-    }
+    }, newReceiveLinkedCts.Token);
   }
 
   private async Task HandleEventAsync(DiscordEvent e, CancellationToken cancellationToken)
   {
+    // TODO: Handle other events
     await SetSequenceAsync(e.Sequence, cancellationToken);
 
-    if (e is HelloDiscordEvent he)
+    switch (e)
     {
+      case HelloDiscordEvent he:
         await SetHeartbeatIntervalAsync(he.Data.HeartbeatInterval, cancellationToken);
         await StartHeartbeatAsync(cancellationToken);
         await IdentifyAsync(cancellationToken);
-
+        _logger.LogInformation("Hello event received. Heartbeat interval: {Interval}", he.Data.HeartbeatInterval);
-      _logger.LogInformation("Hello event received. Heartbeat interval: {Interval}", _heartbeatInterval);
+        break;
-      return;
+      case HeartbeatAckDiscordEvent:
-    }
-
-    if (e is HeartbeatAckDiscordEvent hae)
-    {
         await SetHeartbeatAcknowledgedAsync(_timeProvider.GetUtcNow(), cancellationToken);
-      _logger.LogInformation("Heartbeat acknowledged at {Time}", _timeLastHeartbeatAcknowledged);
+        _logger.LogInformation("Heartbeat acknowledged");
-      return;
+        break;
-    }
+      case ReadyDiscordEvent re:
-
-    if (e is ReadyDiscordEvent re)
-    {
         await SetSessionIdAsync(re.Data.SessionId, cancellationToken);
         await SetResumeGatewayUrlAsync(re.Data.ResumeGatewayUrl, cancellationToken);
-
+        _logger.LogInformation("Ready event received");
-      _logger.LogInformation("Ready event received. Session ID: {SessionId}", _sessionId);
+        break;
-      return;
+      default:
+        _logger.LogInformation("Received event: {Event}", e.GetType().Name);
+        break;
     }
   }
 
   private async Task StartHeartbeatAsync(CancellationToken cancellationToken)
   {
-    using (await _lock.LockAsync(cancellationToken))
+    await CancelHeartbeatTaskAsync(cancellationToken);
-    {
-      if (_heartbeatCts is not null)
-      {
-        await _heartbeatCts.CancelAsync();
-        _heartbeatCts.Dispose();
-      }
 
-      _heartbeatCts = new CancellationTokenSource();
+    var newHeartbeatCts = new CancellationTokenSource();
-      _linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, _heartbeatCts.Token);
+    var newLinkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, newHeartbeatCts.Token);
+
+    await SetHeartbeatCtsAsync(newHeartbeatCts, cancellationToken);
+    await SetLinkedHeartbeatCtsAsync(newLinkedCts, cancellationToken);
 
     _ = Task.Run(async () =>
     {
       _logger.LogInformation("Starting heartbeat task.");
 
-        while (_linkedCts.Token.IsCancellationRequested is false)
+      while (_linkedHeartbeatCts?.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 heartbeatInterval = CalculateHeartbeatInterval();
-            var jitter = Random.Shared.NextDouble();
-# pragma warning restore CA5394 // Do not use insecure randomness
-            int heartbeatInterval;
 
-            using (await _lock.LockAsync(_linkedCts.Token))
+          await Task.Delay(heartbeatInterval, _linkedHeartbeatCts.Token);
-            {
-              heartbeatInterval = _heartbeatInterval + (int)(jitter * _heartbeatInterval);
-            }
 
-            await Task.Delay(heartbeatInterval, _linkedCts.Token);
+          if (IsHeartbeatAcknowledged() is false)
-
-            using (await _lock.LockAsync(_linkedCts.Token))
           {
-              if (_timeLastHeartbeatAcknowledged < _timeLastHeartbeatSent)
+            if (IsWebSocketOpen())
-              {
-                if (_webSocket?.State is WebSocketState.Open)
             {
               _logger.LogWarning("Heartbeat not acknowledged. Closing WebSocket.");
 
-                  await _webSocket.CloseAsync(
+              await CloseAsync(
                 WebSocketCloseStatus.ProtocolError,
                 "Heartbeat not acknowledged",
-                    _linkedCts.Token
+                _linkedHeartbeatCts.Token
               );
 
-                  _canResume = true;
+              await SetCanResumeAsync(true, _linkedHeartbeatCts.Token);
             }
 
+            _logger.LogWarning("Heartbeat not acknowledged. Reconnecting.");
             await ReconnectAsync(cancellationToken);
-
             break;
           }
 
-              await SendHeartbeatAsync(_linkedCts.Token);
+          await SendHeartbeatAsync(_linkedHeartbeatCts.Token);
-            }
 
-            _logger.LogInformation("Heartbeat sent at {Time}", _timeLastHeartbeatSent);
+          _logger.LogInformation("Heartbeat sent");
         }
         catch (OperationCanceledException ex)
         {
@@ -255,57 +254,59 @@ 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);
+
+          _logger.LogWarning("Closing connection and invalidating session.");
+
+          await CloseIfOpenAsync(
+            WebSocketCloseStatus.NormalClosure,
+            "Heartbeat error. Closing connection and invalidating session.",
+            newLinkedCts.Token
+          );
+
+          _logger.LogInformation("Reconnecting because of error in heartbeat task");
+          await ReconnectAsync(cancellationToken);
         }
       }
-      }, _linkedCts.Token);
+    }, newLinkedCts.Token);
-    }
   }
 
-  private async Task ReconnectAsync(CancellationToken token)
+  private async Task ReconnectAsync(CancellationToken cancellationToken)
   {
     _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.");
 
-      _webSocket = _webSocketFactory.Create();
+      await SetWebSocketAsync(_webSocketFactory.Create(), cancellationToken);
-      var uri = new Uri(_resumeGatewayUrl);
+      await ConnectWithResumeUrlAsync(cancellationToken);
-      await _webSocket.ConnectAsync(uri, token);
+      await SendResumeAsync(cancellationToken);
-      await SendResumeAsync(token);
+      await StartReceiveMessagesAsync(cancellationToken);
-      _ = ReceiveMessagesAsync(token);
+      await StartHeartbeatAsync(cancellationToken);
-      _ = StartHeartbeatAsync(token);
       return;
     }
 
     _logger.LogInformation("Reconnecting to Discord Gateway.");
-    await ConnectAsync(token);
+    await ConnectAsync(cancellationToken);
   }
 
   private async Task SendResumeAsync(CancellationToken cancellationToken)
+  {
+    ResumeDiscordEvent resume;
+
+    using (await _lock.LockAsync(cancellationToken))
     {
       if (_lastSequence is null)
       {
-      throw new InvalidOperationException("Cannot resume without a sequence number.");
+        throw new DiscordGatewayClientException("Cannot resume without a sequence number.");
       }
 
-    var resume = new ResumeDiscordEvent(
+      resume = new ResumeDiscordEvent(
         _options.AppToken,
         _sessionId,
         _lastSequence.Value
       );
+    }
 
     await SendJsonAsync(resume, cancellationToken);
   }
@@ -329,10 +330,16 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
 
   private async Task SendJsonAsync(object data, CancellationToken cancellationToken)
   {
-    if (_webSocket?.State is not WebSocketState.Open)
+    using var _ = await _lock.LockAsync(cancellationToken);
+
+    if (_webSocket is null)
     {
-      _logger.LogWarning("WebSocket is not open. Cannot send message.");
+      throw new DiscordGatewayClientException("WebSocket is not set. Cannot send heartbeat.");
-      return;
+    }
+
+    if (_webSocket.State is not WebSocketState.Open)
+    {
+      throw new DiscordGatewayClientException("WebSocket is not open. Cannot send heartbeat.");
     }
 
     try
@@ -353,158 +360,230 @@ internal sealed class DiscordGatewayClient : IDiscordGatewayClient
     }
   }
 
-  private async Task CloseAsync(WebSocketCloseStatus closeStatus, string? statusDescription, CancellationToken cancellationToken)
+  private bool IsWebSocketOpen()
   {
-    using (await _lock.LockAsync(cancellationToken))
+    return _webSocket?.State is WebSocketState.Open;
-    {
-      if (_webSocket is null)
-      {
-        throw new InvalidOperationException("WebSocket is not set. Cannot close.");
   }
 
-      if (_webSocket.State is not WebSocketState.Open)
+  private bool IsHeartbeatAcknowledged()
   {
-        throw new InvalidOperationException("WebSocket is not open. Cannot close.");
+    return _timeLastHeartbeatAcknowledged >= _timeLastHeartbeatSent;
+  }
+
+  private int CalculateHeartbeatInterval()
+  {
+#pragma warning disable CA5394 // Do not use insecure randomness
+    var jitter = Random.Shared.NextDouble();
+#pragma warning restore CA5394 // Do not use insecure randomness
+    return (int)(jitter * _heartbeatInterval);
+  }
+
+  private async Task CancelReceiveMessagesTaskAsync(CancellationToken cancellationToken)
+  {
+    if (_receiveMessageCts is not null)
+    {
+      await _receiveMessageCts.CancelAsync();
+      _receiveMessageCts.Dispose();
+      await SetReceiveCtsAsync(null, cancellationToken);
+    }
+
+    if (_linkedReceiveMessageCts is not null)
+    {
+      await _linkedReceiveMessageCts.CancelAsync();
+      _linkedReceiveMessageCts.Dispose();
+      await SetLinkedReceiveCtsAsync(null, cancellationToken);
+    }
+  }
+
+  private async Task SetLinkedReceiveCtsAsync(CancellationTokenSource? cts, CancellationToken cancellationToken)
+  {
+    using var _ = await _lock.LockAsync(cancellationToken);
+    _linkedReceiveMessageCts = cts;
+  }
+
+  private async Task SetReceiveCtsAsync(CancellationTokenSource? cts, CancellationToken cancellationToken)
+  {
+    using var _ = await _lock.LockAsync(cancellationToken);
+    _receiveMessageCts = cts;
+  }
+
+
+  private async Task CancelHeartbeatTaskAsync(CancellationToken cancellationToken)
+  {
+    if (_heartbeatCts is not null)
+    {
+      await _heartbeatCts.CancelAsync();
+      _heartbeatCts.Dispose();
+      await SetHeartbeatCtsAsync(null, cancellationToken);
+    }
+
+    if (_linkedHeartbeatCts is not null)
+    {
+      await _linkedHeartbeatCts.CancelAsync();
+      _linkedHeartbeatCts.Dispose();
+      await SetLinkedHeartbeatCtsAsync(null, cancellationToken);
+    }
+  }
+
+  private async Task SetLinkedHeartbeatCtsAsync(CancellationTokenSource? cts, CancellationToken cancellationToken)
+  {
+    using var _ = await _lock.LockAsync(cancellationToken);
+    _linkedHeartbeatCts = cts;
+  }
+
+  private async Task SetHeartbeatCtsAsync(CancellationTokenSource? cts, CancellationToken cancellationToken)
+  {
+    using var _ = await _lock.LockAsync(cancellationToken);
+    _heartbeatCts = cts;
+  }
+
+
+  private async Task CloseIfOpenAsync(
+    WebSocketCloseStatus closeStatus,
+    string? statusDescription,
+    CancellationToken cancellationToken
+  )
+  {
+    if (_webSocket?.State is not WebSocketState.Open)
+    {
+      return;
     }
 
     await _webSocket.CloseAsync(closeStatus, statusDescription, cancellationToken);
   }
+
+  private async Task CloseAsync(WebSocketCloseStatus closeStatus, string? statusDescription, CancellationToken cancellationToken)
+  {
+    if (_webSocket is null)
+    {
+      throw new DiscordGatewayClientException("WebSocket is not set. Cannot close.");
+    }
+
+    if (_webSocket.State is not WebSocketState.Open)
+    {
+      throw new DiscordGatewayClientException("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.");
+      throw new DiscordGatewayClientException("WebSocket is not set. Cannot receive message.");
     }
 
     if (_webSocket.State is not WebSocketState.Open)
     {
-        throw new InvalidOperationException("WebSocket is not open. Cannot receive message.");
+      throw new DiscordGatewayClientException("WebSocket is not open. Cannot receive message.");
     }
+
     var result = await _webSocket.ReceiveAsync(segment, cancellationToken);
     return result;
   }
-  }
 
-  private async Task ConnectToGatewayAsync(CancellationToken cancellationToken)
+  private async Task ConnectWithGatewayUrlAsync(CancellationToken cancellationToken)
-  {
-    using (await _lock.LockAsync(cancellationToken))
   {
     if (_webSocket is null)
     {
-        throw new InvalidOperationException("WebSocket is not set. Cannot connect.");
+      throw new DiscordGatewayClientException("WebSocket is not set. Cannot connect.");
     }
 
     if (_webSocket.State is WebSocketState.Open)
     {
-        throw new InvalidOperationException("WebSocket is already open. Cannot connect.");
+      throw new DiscordGatewayClientException("WebSocket is already open. Cannot connect.");
     }
 
     await _webSocket.ConnectAsync(new Uri(_gatewayUrl), cancellationToken);
   }
+
+  private async Task ConnectWithResumeUrlAsync(CancellationToken cancellationToken)
+  {
+    if (_webSocket is null)
+    {
+      throw new DiscordGatewayClientException("WebSocket is not set. Cannot connect.");
     }
 
-  private async Task<bool> IsGatewayUrlSetAsync(CancellationToken cancellationToken)
+    if (_webSocket.State is WebSocketState.Open)
     {
-    using (await _lock.LockAsync(cancellationToken))
+      throw new DiscordGatewayClientException("WebSocket is already open. Cannot connect.");
-    {
-      return !string.IsNullOrEmpty(_gatewayUrl);
     }
+
+    await _webSocket.ConnectAsync(new Uri(_resumeGatewayUrl), cancellationToken);
   }
 
   private async Task SetGatewayUrlAsync(string url, CancellationToken cancellationToken)
   {
-    using (await _lock.LockAsync(cancellationToken))
+    using var _ = await _lock.LockAsync(cancellationToken);
-    {
     _gatewayUrl = url;
   }
-  }
 
   private async Task SetSequenceAsync(int? sequence, CancellationToken cancellationToken)
   {
-    using (await _lock.LockAsync(cancellationToken))
+    using var _ = await _lock.LockAsync(cancellationToken);
-    {
     _lastSequence = sequence;
   }
-  }
 
   private async Task SetHeartbeatIntervalAsync(int interval, CancellationToken cancellationToken)
   {
-    using (await _lock.LockAsync(cancellationToken))
+    using var _ = await _lock.LockAsync(cancellationToken);
-    {
     _heartbeatInterval = interval;
   }
-  }
 
   private async Task SetHeartbeatSentAsync(DateTimeOffset time, CancellationToken cancellationToken)
   {
-    using (await _lock.LockAsync(cancellationToken))
+    using var _ = await _lock.LockAsync(cancellationToken);
-    {
     _timeLastHeartbeatSent = time;
   }
-  }
 
   private async Task SetHeartbeatAcknowledgedAsync(DateTimeOffset time, CancellationToken cancellationToken)
   {
-    using (await _lock.LockAsync(cancellationToken))
+    using var _ = await _lock.LockAsync(cancellationToken);
-    {
     _timeLastHeartbeatAcknowledged = time;
   }
-  }
 
   private async Task SetSessionIdAsync(string sessionId, CancellationToken cancellationToken)
   {
-    using (await _lock.LockAsync(cancellationToken))
+    using var _ = await _lock.LockAsync(cancellationToken);
-    {
     _sessionId = sessionId;
   }
-  }
 
   private async Task SetResumeGatewayUrlAsync(string url, CancellationToken cancellationToken)
   {
-    using (await _lock.LockAsync(cancellationToken))
+    using var _ = 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))
+    using var _ = await _lock.LockAsync(cancellationToken);
-    {
     _webSocket = webSocket;
   }
-  }
 
   private async Task SetCanResumeAsync(bool canResume, CancellationToken cancellationToken)
   {
-    using (await _lock.LockAsync(cancellationToken))
+    using var _ = await _lock.LockAsync(cancellationToken);
-    {
     _canResume = canResume;
   }
-  }
 
   public void Dispose()
   {
     _heartbeatCts?.Cancel();
-    _linkedCts?.Cancel();
+    _linkedHeartbeatCts?.Cancel();
+
+    _receiveMessageCts?.Cancel();
+    _linkedReceiveMessageCts?.Cancel();
 
     _heartbeatCts?.Dispose();
-    _linkedCts?.Dispose();
+    _linkedHeartbeatCts?.Dispose();
+
+    _receiveMessageCts?.Dispose();
+    _linkedReceiveMessageCts?.Dispose();
+
     _webSocket?.Dispose();
     _lock.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 Identify = 2;
   public const int Resume = 6;
   public const int HeartbeatAck = 11;
   public const int Hello = 10;

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

@@ -7,6 +7,7 @@ internal sealed record IdentifyDiscordEvent : DiscordEvent
 
   public IdentifyDiscordEvent(string token, long intents)
   {
+    OpCode = DiscordOpCodes.Identify;
     Data = new IdentifyData
     {
       Token = token,

src/StevesBot.Worker/Discord/IDiscordGatewayClient.cs

@@ -3,4 +3,5 @@ namespace StevesBot.Worker.Discord;
 internal interface IDiscordGatewayClient : IDisposable
 {
   Task ConnectAsync(CancellationToken cancellationToken);
+  Task DisconnectAsync(CancellationToken cancellationToken);
 }

src/StevesBot.Worker/Program.cs

@@ -1,6 +1,31 @@
-using StevesBot.Worker;
+using Microsoft.Extensions.Options;
 
 var builder = Host.CreateApplicationBuilder(args);
+
+builder.Services.Configure<HostOptions>(static options => options.ShutdownTimeout = TimeSpan.FromSeconds(30));
+
+builder.Services.AddOptions<DiscordClientOptions>()
+  .BindConfiguration(nameof(DiscordClientOptions));
+
+builder.Services.AddSingleton(static sp =>
+{
+  var discordOptions = sp.GetRequiredService<IOptions<DiscordClientOptions>>().Value;
+  return discordOptions;
+});
+
+builder.Services.AddSingleton<IWebSocketFactory, WebSocketFactory>();
+builder.Services.AddSingleton(TimeProvider.System);
+
+builder.Services
+  .AddHttpClient<IDiscordRestClient, DiscordRestClient>(static (sp, c) =>
+  {
+    var discordOptions = sp.GetRequiredService<DiscordClientOptions>();
+    c.BaseAddress = new Uri(discordOptions.ApiUrl);
+    c.DefaultRequestHeaders.Authorization = new("Bot", discordOptions.AppToken);
+  })
+  .AddStandardResilienceHandler();
+
+builder.Services.AddSingleton<IDiscordGatewayClient, DiscordGatewayClient>();
 builder.Services.AddHostedService<Worker>();
 
 var host = builder.Build();

src/StevesBot.Worker/StevesBot.Worker.csproj

@@ -6,6 +6,8 @@
 
   <ItemGroup>
     <PackageReference Include="Microsoft.Extensions.Hosting" Version="9.0.4" />
+    <PackageReference Include="Microsoft.Extensions.Http" Version="9.0.5" />
+    <PackageReference Include="Microsoft.Extensions.Http.Resilience" Version="9.5.0" />
   </ItemGroup>
 
   <ItemGroup>

src/StevesBot.Worker/Threading/LockReleaser.cs

@@ -19,7 +19,16 @@ internal sealed class LockReleaser : IDisposable
       return;
     }
 
+    try
+    {
       _semaphore.Release();
+    }
+    catch (ObjectDisposedException)
+    {
+    }
+    finally
+    {
       _released = true;
     }
   }
+}

src/StevesBot.Worker/Usings.cs

@@ -7,3 +7,6 @@ global using System.Text.Json.Serialization;
 global using StevesBot.Worker.Discord.Events;
 global using StevesBot.Worker.Threading;
 global using StevesBot.Worker.WebSockets;
+
+global using StevesBot.Worker;
+global using StevesBot.Worker.Discord;

src/StevesBot.Worker/Worker.cs

@@ -1,24 +1,26 @@
+
 namespace StevesBot.Worker;
 
-internal class Worker : BackgroundService
+internal class Worker : IHostedService
 {
   private readonly ILogger<Worker> _logger;
+  private readonly IDiscordGatewayClient _discordGatewayClient;
 
-  public Worker(ILogger<Worker> logger)
+  public Worker(ILogger<Worker> logger, IDiscordGatewayClient discordGatewayClient)
   {
     _logger = logger;
+    _discordGatewayClient = discordGatewayClient;
   }
 
-  protected override async Task ExecuteAsync(CancellationToken stoppingToken)
+  public Task StartAsync(CancellationToken cancellationToken)
   {
-    while (!stoppingToken.IsCancellationRequested)
+    _logger.LogInformation("Connecting Discord Gateway Client");
-    {
+    return _discordGatewayClient.ConnectAsync(cancellationToken);
-      if (_logger.IsEnabled(LogLevel.Information))
-      {
-        _logger.LogInformation("Worker running at: {Time}", DateTimeOffset.Now);
   }
 
-      await Task.Delay(1000, stoppingToken);
+  public Task StopAsync(CancellationToken cancellationToken)
-    }
+  {
+    _logger.LogInformation("Disconnecting Discord Gateway Client");
+    return _discordGatewayClient.DisconnectAsync(cancellationToken);
   }
 }

src/StevesBot.Worker/appsettings.Example.json

@@ -5,7 +5,7 @@
       "Microsoft.Hosting.Lifetime": "Information"
     }
   },
-  "DiscordGatewayClientOptions": {
+  "DiscordClientOptions": {
     "ApiUrl": "ApiUrl",
     "AppToken": "AppToken",
     "Intents": 0