fix: reconnect agent after ELS restarts

FeatBit.Agent.slnx

@@ -0,0 +1,9 @@
+<Solution>
+  <Folder Name="/src/">
+    <Project Path="src/Api/Api.csproj" />
+  </Folder>
+  <Folder Name="/tests/">
+    <Project Path="tests/Api.IntegrationTests/Api.IntegrationTests.csproj" />
+    <Project Path="tests/Api.UnitTests/Api.UnitTests.csproj" />
+  </Folder>
+</Solution>

src/Api/Properties/AssemblyInfo.cs

@@ -0,0 +1,4 @@
+using System.Runtime.CompilerServices;
+
+[assembly: InternalsVisibleTo("Api.UnitTests")]
+[assembly: InternalsVisibleTo("Api.IntegrationTests")]

src/Api/Transport/FbWebSocket.cs

@@ -5,6 +5,8 @@ namespace Api.Transport;
 
 internal sealed partial class FbWebSocket
 {
+    internal const WebSocketCloseStatus ServerRejectedCloseStatus = (WebSocketCloseStatus)4003;
+
     public event Func<Task>? OnConnected;
     public event Func<Task>? OnKeepAlive;
     public event Func<Exception?, Task>? OnReconnecting;
@@ -42,12 +44,15 @@ internal FbWebSocket(Uri websocketUri, ILoggerFactory loggerFactory)
     public async Task ConnectAsync(bool isReconnecting = false)
     {
         _websocket = new ClientWebSocket();
-        using var connectCts = new CancellationTokenSource(_connectTimeout);
+
+        using var connectCts = CancellationTokenSource.CreateLinkedTokenSource(_stopCts.Token);
+        connectCts.CancelAfter(_connectTimeout);
 
         try
         {
             Log.Starting(_logger, _websocketUri);
             await _websocket.ConnectAsync(_websocketUri, connectCts.Token);
+            _stopCts.Token.ThrowIfCancellationRequested();
 
             Log.Connected(_logger);
 
@@ -61,12 +66,14 @@ public async Task ConnectAsync(bool isReconnecting = false)
         {
             Log.ErrorStarting(_logger, ex);
 
-            if (ex is TaskCanceledException tce && tce.CancellationToken == connectCts.Token)
+            if (ex is OperationCanceledException && !_stopCts.IsCancellationRequested)
             {
                 Log.TimedOut(_logger, _connectTimeout.TotalSeconds);
             }
 
-            if (!isReconnecting)
+            _websocket.Dispose();
+
+            if (!isReconnecting && !_stopCts.IsCancellationRequested)
             {
                 _ = ReconnectAsync();
             }
@@ -102,7 +109,6 @@ public async Task CloseAsync(CancellationToken cancellationToken)
         }
 
         _stopCts.Cancel();
-        _stopCts.Dispose();
 
         _keepAliveTimer?.Dispose();
 
@@ -159,10 +165,24 @@ private void HandleConnectionClose()
         }
     }
 
-    private bool ShouldReconnect()
+    private bool ShouldReconnect() => ShouldReconnect(_websocket?.CloseStatus, _stopCts.IsCancellationRequested);
+
+    internal static bool ShouldReconnect(WebSocketCloseStatus? closeStatus, bool stopRequested)
     {
-        var closeStatus = _websocket?.CloseStatus;
-        return closeStatus != WebSocketCloseStatus.NormalClosure && closeStatus != (WebSocketCloseStatus)4003;
+        if (closeStatus == ServerRejectedCloseStatus)
+        {
+            // server rejected the connection explicitly, so we should not reconnect
+            return false;
+        }
+
+        if (stopRequested)
+        {
+            // connection was stopped manually, so we should not reconnect
+            return false;
+        }
+
+        // otherwise, we should reconnect
+        return true;
     }
 
     private async Task ReconnectAsync()
@@ -274,4 +294,4 @@ private async Task KeepAliveAsync(CancellationToken ct = default)
         Log.InvokingEventHandler(_logger, nameof(OnKeepAlive));
         _ = OnKeepAlive?.Invoke();
     }
-}
+}

tests/Api.IntegrationTests/Api.IntegrationTests.csproj

@@ -0,0 +1,25 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+    <PropertyGroup>
+        <TargetFramework>net10.0</TargetFramework>
+        <ImplicitUsings>enable</ImplicitUsings>
+        <Nullable>enable</Nullable>
+        <IsPackable>false</IsPackable>
+        <IsTestProject>true</IsTestProject>
+    </PropertyGroup>
+
+    <ItemGroup>
+        <PackageReference Include="Microsoft.NET.Test.Sdk" Version="18.5.1" />
+        <PackageReference Include="xunit" Version="2.9.3" />
+        <PackageReference Include="xunit.runner.visualstudio" Version="3.1.5">
+            <IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
+            <PrivateAssets>all</PrivateAssets>
+        </PackageReference>
+    </ItemGroup>
+
+    <ItemGroup>
+        <FrameworkReference Include="Microsoft.AspNetCore.App" />
+        <ProjectReference Include="..\..\src\Api\Api.csproj" />
+    </ItemGroup>
+
+</Project>

tests/Api.IntegrationTests/Transport/WebSocketReconnectTests.cs

@@ -0,0 +1,100 @@
+using System.Net.WebSockets;
+using Api.Transport;
+using Microsoft.Extensions.Logging.Abstractions;
+using Xunit;
+
+namespace Api.IntegrationTests.Transport;
+
+public class WebSocketReconnectTests
+{
+    [Fact]
+    public async Task NormalClosure_ReconnectsRepeatedly()
+    {
+        await using var server = await WebSocketStubServer.StartAsync();
+        var webSocket = new FbWebSocket(server.WebSocketUri, NullLoggerFactory.Instance);
+
+        try
+        {
+            await webSocket.ConnectAsync();
+            await server.WaitForConnectionsAsync(1);
+
+            await server.CloseAsync(1, WebSocketCloseStatus.NormalClosure);
+            await server.WaitForConnectionsAsync(2);
+
+            await server.CloseAsync(2, WebSocketCloseStatus.NormalClosure);
+            await server.WaitForConnectionsAsync(3);
+
+            Assert.Equal(3, server.ConnectionCount);
+        }
+        finally
+        {
+            server.AbortConnections();
+            await webSocket.CloseAsync(CancellationToken.None);
+        }
+    }
+
+    [Fact]
+    public async Task InvalidRequestClose_DoesNotReconnect()
+    {
+        await using var server = await WebSocketStubServer.StartAsync(
+            FbWebSocket.ServerRejectedCloseStatus
+        );
+        var webSocket = new FbWebSocket(server.WebSocketUri, NullLoggerFactory.Instance);
+        var closed = new TaskCompletionSource<WebSocketCloseStatus?>(
+            TaskCreationOptions.RunContinuationsAsynchronously
+        );
+        webSocket.OnClosed += (_, status, _) =>
+        {
+            closed.TrySetResult(status);
+            return Task.CompletedTask;
+        };
+
+        try
+        {
+            await webSocket.ConnectAsync();
+
+            var closeStatus = await closed.Task.WaitAsync(TimeSpan.FromSeconds(5));
+            await Task.Delay(TimeSpan.FromMilliseconds(250));
+
+            Assert.Equal(FbWebSocket.ServerRejectedCloseStatus, closeStatus);
+            Assert.Equal(1, server.ConnectionCount);
+        }
+        finally
+        {
+            server.AbortConnections();
+            await webSocket.CloseAsync(CancellationToken.None);
+        }
+    }
+
+    [Fact]
+    public async Task NormalClosure_WhenStoppedDuringReconnect_DoesNotReconnect()
+    {
+        await using var server = await WebSocketStubServer.StartAsync();
+        var webSocket = new FbWebSocket(server.WebSocketUri, NullLoggerFactory.Instance);
+        var reconnecting = new TaskCompletionSource(
+            TaskCreationOptions.RunContinuationsAsynchronously
+        );
+        webSocket.OnReconnecting += _ =>
+        {
+            reconnecting.TrySetResult();
+            return webSocket.CloseAsync(CancellationToken.None);
+        };
+
+        try
+        {
+            await webSocket.ConnectAsync();
+            await server.WaitForConnectionsAsync(1);
+
+            await server.CloseAsync(1, WebSocketCloseStatus.NormalClosure);
+            await reconnecting.Task.WaitAsync(TimeSpan.FromSeconds(5));
+            await Task.Delay(TimeSpan.FromMilliseconds(250));
+
+            Assert.Equal(1, server.ConnectionCount);
+        }
+        finally
+        {
+            server.AbortConnections();
+            await webSocket.CloseAsync(CancellationToken.None);
+        }
+    }
+}

tests/Api.IntegrationTests/Transport/WebSocketStubServer.cs

@@ -0,0 +1,142 @@
+using System.Collections.Concurrent;
+using System.Net;
+using System.Net.WebSockets;
+using Microsoft.AspNetCore.Builder;
+using Microsoft.AspNetCore.Hosting;
+using Microsoft.AspNetCore.Hosting.Server;
+using Microsoft.AspNetCore.Hosting.Server.Features;
+using Microsoft.AspNetCore.Http;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging;
+
+namespace Api.IntegrationTests.Transport;
+
+public sealed class WebSocketStubServer : IAsyncDisposable
+{
+    private readonly WebApplication _application;
+    private readonly WebSocketCloseStatus? _closeImmediatelyWith;
+    private readonly ConcurrentDictionary<int, WebSocket> _connections = new();
+    private readonly SemaphoreSlim _connectionSignal = new(0);
+    private int _connectionCount;
+
+    private WebSocketStubServer(WebApplication application, WebSocketCloseStatus? closeImmediatelyWith)
+    {
+        _application = application;
+        _closeImmediatelyWith = closeImmediatelyWith;
+    }
+
+    public int ConnectionCount => Volatile.Read(ref _connectionCount);
+
+    public Uri WebSocketUri { get; private set; } = null!;
+
+    public static async Task<WebSocketStubServer> StartAsync(WebSocketCloseStatus? closeImmediatelyWith = null)
+    {
+        var builder = WebApplication.CreateBuilder();
+        builder.Logging.ClearProviders();
+        builder.WebHost.ConfigureKestrel(options => options.Listen(IPAddress.Loopback, 0));
+
+        var application = builder.Build();
+        var server = new WebSocketStubServer(application, closeImmediatelyWith);
+
+        application.UseWebSockets();
+        application.Map("/streaming", server.HandleConnectionAsync);
+
+        await application.StartAsync();
+
+        var addresses = application.Services
+            .GetRequiredService<IServer>()
+            .Features
+            .Get<IServerAddressesFeature>()!;
+        var httpUri = new Uri(addresses.Addresses.Single());
+        server.WebSocketUri = new UriBuilder(new Uri(httpUri, "/streaming")) { Scheme = "ws" }.Uri;
+
+        return server;
+    }
+
+    public async Task WaitForConnectionsAsync(int expectedCount)
+    {
+        using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+        while (ConnectionCount < expectedCount)
+        {
+            await _connectionSignal.WaitAsync(timeout.Token);
+        }
+    }
+
+    public async Task CloseAsync(int connectionId, WebSocketCloseStatus closeStatus)
+    {
+        if (!_connections.TryGetValue(connectionId, out var connection))
+        {
+            throw new InvalidOperationException(
+                $"WebSocket connection {connectionId} is no longer active."
+            );
+        }
+
+        await connection.CloseOutputAsync(closeStatus, "closed by test server", CancellationToken.None);
+    }
+
+    public void AbortConnections()
+    {
+        foreach (var connection in _connections.Values)
+        {
+            connection.Abort();
+        }
+    }
+
+    private async Task HandleConnectionAsync(HttpContext context)
+    {
+        if (!context.WebSockets.IsWebSocketRequest)
+        {
+            context.Response.StatusCode = StatusCodes.Status400BadRequest;
+            return;
+        }
+
+        using var webSocket = await context.WebSockets.AcceptWebSocketAsync();
+        var connectionId = Interlocked.Increment(ref _connectionCount);
+        _connections[connectionId] = webSocket;
+        _connectionSignal.Release();
+
+        try
+        {
+            if (_closeImmediatelyWith.HasValue)
+            {
+                await webSocket.CloseOutputAsync(
+                    _closeImmediatelyWith.Value,
+                    "closed by test server",
+                    CancellationToken.None
+                );
+            }
+
+            var buffer = new byte[1024];
+            while (webSocket.State is WebSocketState.Open or WebSocketState.CloseSent)
+            {
+                var result = await webSocket.ReceiveAsync(buffer, context.RequestAborted);
+                if (result.MessageType == WebSocketMessageType.Close)
+                {
+                    break;
+                }
+            }
+        }
+        catch (OperationCanceledException)
+        {
+            // The test server is stopping.
+        }
+        catch (WebSocketException)
+        {
+            // The client disposed the connection after observing the close frame.
+        }
+        finally
+        {
+            _connections.TryRemove(connectionId, out _);
+        }
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        AbortConnections();
+
+        using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+        await _application.StopAsync(timeout.Token);
+        await _application.DisposeAsync();
+        _connectionSignal.Dispose();
+    }
+}