Fix/hot reload shutdown

src/JsxCore/Hosting/JsxCoreServiceCollectionExtensions.cs

@@ -217,6 +217,7 @@ minifier is null
             options.HotReload ?? false,
             provider.GetRequiredService<JsxCompilationService>(),
             provider.GetRequiredService<JsxServerRendererReset>(),
+            provider.GetRequiredService<IHostApplicationLifetime>(),
             provider.GetRequiredService<ILogger<JsxHotReloadService>>()));
 
         services.TryAddSingleton<IJsxHotReloadState>(provider => provider.GetRequiredService<JsxHotReloadService>());

src/JsxCore/Hosting/JsxHotReloadService.cs

@@ -3,6 +3,7 @@
 using System.Text;
 using System.Text.Json;
 using JsxCore.Compilation;
+using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging;
 
 namespace JsxCore.Hosting;
@@ -16,16 +17,22 @@ public sealed class JsxHotReloadService(
     bool enabled,
     JsxCompilationService compilation,
     JsxServerRendererReset reset,
+    IHostApplicationLifetime lifetime,
     ILogger<JsxHotReloadService> logger)
     : IJsxHotReloadState, IDisposable
 {
     private readonly ConcurrentDictionary<Guid, WebSocket> _clients = new();
     private readonly ILogger<JsxHotReloadService> _logger = logger ?? throw new ArgumentNullException(nameof(logger));
     private readonly JsxCompilationService _compilation = compilation ?? throw new ArgumentNullException(nameof(compilation));
     private readonly JsxServerRendererReset _reset = reset ?? throw new ArgumentNullException(nameof(reset));
+    private readonly IHostApplicationLifetime _lifetime = lifetime ?? throw new ArgumentNullException(nameof(lifetime));
+    private readonly object _clientGate = new();
+    private CancellationTokenRegistration _stoppingRegistration;
     private bool _subscribed;
+    private bool _stopped;
 
     public bool Enabled { get; } = enabled;
+    internal int ClientCount => _clients.Count;
 
     public void Start()
     {
@@ -35,6 +42,7 @@ public void Start()
         }
         _subscribed = true;
         _compilation.BuildCompleted += OnBuildCompleted;
+        _stoppingRegistration = _lifetime.ApplicationStopping.Register(Stop);
     }
 
     private void OnBuildCompleted(BuildState state)
@@ -59,7 +67,21 @@ public async Task AcceptAsync(WebSocket socket, CancellationToken cancellationTo
         ArgumentNullException.ThrowIfNull(socket);
 
         var id = Guid.NewGuid();
-        _clients[id] = socket;
+        var accepted = false;
+        lock (_clientGate)
+        {
+            if (!_stopped)
+            {
+                _clients[id] = socket;
+                accepted = true;
+            }
+        }
+
+        if (!accepted)
+        {
+            socket.Abort();
+            return;
+        }
         _logger.LogDebug("JsxCore hot reload client {ClientId} connected.", id);
 
         try
@@ -120,14 +142,32 @@ await socket.SendAsync(payload, WebSocketMessageType.Text, endOfMessage: true, C
         }
     }
 
-    public void Dispose()
+    private void Stop()
     {
         if (_subscribed)
         {
             _compilation.BuildCompleted -= OnBuildCompleted;
             _subscribed = false;
         }
-        _clients.Clear();
+
+        List<WebSocket> sockets;
+        lock (_clientGate)
+        {
+            _stopped = true;
+            sockets = [.. _clients.Values];
+            _clients.Clear();
+        }
+
+        foreach (var socket in sockets)
+        {
+            socket.Abort();
+        }
+    }
+
+    public void Dispose()
+    {
+        Stop();
+        _stoppingRegistration.Dispose();
     }
 }
 

test/JsxCore.Tests/Component/HotReload/HotReloadTests.cs

@@ -1,8 +1,11 @@
 using System.Net;
+using JsxCore.Hosting;
 using JsxCore.Mvc;
 using Microsoft.AspNetCore.Builder;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Routing;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Hosting;
 using System.Text.RegularExpressions;
 using Shouldly;
 
@@ -75,4 +78,22 @@ public async Task HotReloadEndpoint_RequestIsNotAWebSocket_IsRejected()
 
         response.StatusCode.ShouldBe(HttpStatusCode.BadRequest);
     }
+
+    [Fact]
+    public async Task HotReloadService_ApplicationStopping_AbortsConnectedSockets()
+    {
+        using var project = HostedViews.Project();
+        await using var host = await JsxTestHost.StartAsync(project, o => o.HotReload = true);
+
+        var webSocketClient = host.Server.CreateWebSocketClient();
+        var uri = new Uri(host.Server.BaseAddress, "/_jsx/hmr");
+        using var socket = await webSocketClient.ConnectAsync(uri, CancellationToken.None);
+
+        var hotReload = host.Server.Services.GetRequiredService<JsxHotReloadService>();
+        SpinWait.SpinUntil(() => hotReload.ClientCount == 1, TimeSpan.FromSeconds(5)).ShouldBeTrue();
+
+        host.Server.Services.GetRequiredService<IHostApplicationLifetime>().StopApplication();
+
+        SpinWait.SpinUntil(() => hotReload.ClientCount == 0, TimeSpan.FromSeconds(5)).ShouldBeTrue();
+    }
 }