diff --git a/.editorconfig b/.editorconfig index 568eeea..f99bdb0 100644 --- a/.editorconfig +++ b/.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 diff --git a/src/StevesBot.Worker.Tests/Unit/DiscordGatewayClientTests.cs b/src/StevesBot.Worker.Tests/Unit/DiscordGatewayClientTests.cs index 5fb39ef..9f9feba 100644 --- a/src/StevesBot.Worker.Tests/Unit/DiscordGatewayClientTests.cs +++ b/src/StevesBot.Worker.Tests/Unit/DiscordGatewayClientTests.cs @@ -6,7 +6,6 @@ public sealed class DiscordGatewayClientTests : IDisposable private readonly Mock _mockWebSocketFactory = new(); private readonly Mock> _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(); + } + + [Fact] + public void Constructor_WhenCalledAndWebSocketFactoryIsNull_ItShouldThrowArgumentNullException() + { + var act = () => new DiscordGatewayClient( + _options, + null!, + _mockLogger.Object, + _mockDiscordRestClient.Object + ); + + act.Should().Throw(); + } + + [Fact] + public void Constructor_WhenCalledAndLoggerIsNull_ItShouldThrowArgumentNullException() + { + var act = () => new DiscordGatewayClient( + _options, + _mockWebSocketFactory.Object, + null!, + _mockDiscordRestClient.Object + ); + + act.Should().Throw(); + } + + [Fact] + public void Constructor_WhenCalledAndDiscordRestClientIsNull_ItShouldThrowArgumentNullException() + { + var act = () => new DiscordGatewayClient( + _options, + _mockWebSocketFactory.Object, + _mockLogger.Object, + null! + ); + + act.Should().Throw(); + } + [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())) + .ReturnsAsync("wss://gateway.discord.gg"); + + var mockWebSocket = new Mock(); + + 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>(b => expectedHeartbeatPayload.Bytes.SequenceEqual(b)), + It.IsAny(), + It.IsAny(), + It.IsAny() + ), + Times.Once + ); + + mockWebSocket.Verify( + static x => x.CloseAsync( + It.Is( + x => x != WebSocketCloseStatus.NormalClosure && x != WebSocketCloseStatus.EndpointUnavailable + ), + It.IsAny(), + It.IsAny() + ), + Times.Once + ); + } + + [Fact] + public async Task ConnectAsync_WhenCalledAndHeartbeatNotAcknowlegedAndAlreadyClosed_ItShouldNotDisconnect() + { + _mockDiscordRestClient + .Setup(static x => x.GetGatewayUrlAsync(It.IsAny())) + .ReturnsAsync("wss://gateway.discord.gg"); + + var mockWebSocket = new Mock(); + + 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>(b => expectedHeartbeatPayload.Bytes.SequenceEqual(b)), + It.IsAny(), + It.IsAny(), + It.IsAny() + ), + Times.Once + ); + + mockWebSocket.Verify( + static x => x.CloseAsync( + It.IsAny(), + It.IsAny(), + It.IsAny() + ), + Times.Never + ); + } + private static void SetupReceiveMessageSequence( Mock mockWebSocket, Queue<(WebSocketReceiveResult, byte[])> messageQueue diff --git a/src/StevesBot.Worker/Discord/DiscordGatewayClient.cs b/src/StevesBot.Worker/Discord/DiscordGatewayClient.cs index cfec6d5..32bf9b1 100644 --- a/src/StevesBot.Worker/Discord/DiscordGatewayClient.cs +++ b/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 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 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 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(); diff --git a/src/StevesBot.Worker/Discord/Events/DiscordOpCodes.cs b/src/StevesBot.Worker/Discord/Events/DiscordOpCodes.cs index 1ddcaa5..b2a0e1c 100644 --- a/src/StevesBot.Worker/Discord/Events/DiscordOpCodes.cs +++ b/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; } \ No newline at end of file diff --git a/src/StevesBot.Worker/Discord/Events/ResumeDiscordEvent.cs b/src/StevesBot.Worker/Discord/Events/ResumeDiscordEvent.cs new file mode 100644 index 0000000..00bd78c --- /dev/null +++ b/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; } +} \ No newline at end of file