diff --git a/src/JsxCore/Hosting/JsxCoreServiceCollectionExtensions.cs b/src/JsxCore/Hosting/JsxCoreServiceCollectionExtensions.cs index 1fffe0c..ce633c0 100644 --- a/src/JsxCore/Hosting/JsxCoreServiceCollectionExtensions.cs +++ b/src/JsxCore/Hosting/JsxCoreServiceCollectionExtensions.cs @@ -217,6 +217,7 @@ minifier is null options.HotReload ?? false, provider.GetRequiredService(), provider.GetRequiredService(), + provider.GetRequiredService(), provider.GetRequiredService>())); services.TryAddSingleton(provider => provider.GetRequiredService()); diff --git a/src/JsxCore/Hosting/JsxHotReloadService.cs b/src/JsxCore/Hosting/JsxHotReloadService.cs index aa379e7..42621c8 100644 --- a/src/JsxCore/Hosting/JsxHotReloadService.cs +++ b/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,6 +17,7 @@ public sealed class JsxHotReloadService( bool enabled, JsxCompilationService compilation, JsxServerRendererReset reset, + IHostApplicationLifetime lifetime, ILogger logger) : IJsxHotReloadState, IDisposable { @@ -23,9 +25,14 @@ public sealed class JsxHotReloadService( private readonly ILogger _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 sockets; + lock (_clientGate) + { + _stopped = true; + sockets = [.. _clients.Values]; + _clients.Clear(); + } + + foreach (var socket in sockets) + { + socket.Abort(); + } + } + + public void Dispose() + { + Stop(); + _stoppingRegistration.Dispose(); } } diff --git a/test/JsxCore.Tests/Component/HotReload/HotReloadTests.cs b/test/JsxCore.Tests/Component/HotReload/HotReloadTests.cs index 3c2ab05..6dd94c3 100644 --- a/test/JsxCore.Tests/Component/HotReload/HotReloadTests.cs +++ b/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(); + SpinWait.SpinUntil(() => hotReload.ClientCount == 1, TimeSpan.FromSeconds(5)).ShouldBeTrue(); + + host.Server.Services.GetRequiredService().StopApplication(); + + SpinWait.SpinUntil(() => hotReload.ClientCount == 0, TimeSpan.FromSeconds(5)).ShouldBeTrue(); + } }