Stevan/steves-bot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: begin working on connection management

Browse Source
This commit is contained in:
Stevan Freeborn
2025-05-13 23:51:14 -05:00
parent df9b3cd7fc
commit 35d6ee5503
7 changed files with 376 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/StevesBot.Worker.Tests/Unit/DiscordGatewayClientOptionsTests.cs → src/StevesBot.Worker.Tests/Unit/DiscordClientOptionsTests.cs
+5 -5
View File
@@ -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
+98
View File
@@ -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
+34
View File
@@ -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
+4 -12
View File
@@ -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
+6
View File
@@ -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/DiscordGatewayClientOptions.cs → src/StevesBot.Worker/Discord/DiscordClientOptions.cs
+1 -1
View File
@@ -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
+228 -20
View File
@@ -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()
// {
// PropertyNamingPolicy = JsonNamingPolicy.CamelCase,
// ReferenceHandler = ReferenceHandler.IgnoreCycles,
// Converters =
// {
// new DiscordEventConverter(),
// },
// };
private readonly JsonSerializerOptions _jsonSerializerOptions = new()
{
PropertyNamingPolicy = JsonNamingPolicy.CamelCase,
ReferenceHandler = ReferenceHandler.IgnoreCycles,
Converters =
{
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 _timeLastHeartbeatAcknowledged = DateTime.MinValue;
// private CancellationTokenSource? _heartbeatCts;
// private Task? _heartbeatTask;
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;
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?.Dispose();
// _heartbeatTask?.Dispose();
_heartbeatCts?.Cancel();
_heartbeatCts?.Dispose();
_linkedCts?.Dispose();
_heartbeatTask?.Dispose();
_receiveTask?.Dispose();
_webSocket?.Dispose();
_lock.Dispose();