Improve shutdown sequence of Aspire service and fix race condition

src/Dotnet.Watch/.editorconfig

@@ -14,13 +14,15 @@ dotnet_diagnostic.CA2008.severity = none # Do not create tasks without passing a
 
 # CS - C# compiler warnings/errors
 dotnet_diagnostic.CS1591.severity = none # Missing XML comment for publicly visible type or member
-dotnet_diagnostic.CS1573.severity = none # Parameter 'sourceFile' has no matching param tag in the XML comment
+dotnet_diagnostic.CS1573.severity = none # Parameter has no matching param tag in the XML comment
+dotnet_diagnostic.CS1572.severity = warning # XML comment has a param tag for '...', but there is no parameter by that name
 
 # IDE - IDE/Style warnings
 dotnet_diagnostic.IDE0005.severity = none # Using directive is unnecessary
 dotnet_diagnostic.IDE0011.severity = none # Add braces
 dotnet_diagnostic.IDE0036.severity = none # Order modifiers
+dotnet_diagnostic.IDE0044.severity = none # Make field readonly
 dotnet_diagnostic.IDE0060.severity = none # Remove unused parameter
 dotnet_diagnostic.IDE0073.severity = none # File header does not match required text
 dotnet_diagnostic.IDE0161.severity = none # Convert to file-scoped namespace
-dotnet_diagnostic.IDE1006.severity = none # Naming rule violation
+dotnet_diagnostic.IDE1006.severity = none # Naming rule violation

src/Dotnet.Watch/AspireService/AspireServerService.cs

@@ -42,7 +42,12 @@ internal partial class AspireServerService : IAsyncDisposable
     private readonly string _currentSecret;
     private readonly string _displayName;
 
-    private readonly CancellationTokenSource _shutdownCancellationTokenSource = new();
+    /// <summary>
+    /// Triggered when the shutdown process has been initiated.
+    /// During this time all requests respond with <see cref="HttpStatusCode.ServiceUnavailable"/>.
+    /// </summary>
+    private readonly CancellationTokenSource _shutdownCancellationSource = new();
+
     private readonly int _port;
     private readonly X509Certificate2 _certificate;
     private readonly string _certificateEncodedBytes;
@@ -91,12 +96,20 @@ public AspireServerService(IAspireServerEvents aspireServerEvents, string displa
         _requestListener = StartListeningAsync();
     }
 
+    public void InitiateShutdown()
+    {
+        Log("Server shutdown initiated.");
+        _shutdownCancellationSource.Cancel();
+
+        // TODO: send message to DCP
+        // https://github.com/dotnet/aspire/issues/14987
+    }
+
     public async ValueTask DisposeAsync()
     {
-        // Shutdown the service:
-        _shutdownCancellationTokenSource.Cancel();
+        Log("Disposing server ...");
 
-        Log("Waiting for server to shutdown ...");
+        _shutdownCancellationSource.Cancel();
 
         try
         {
@@ -111,7 +124,7 @@ public async ValueTask DisposeAsync()
 
         _socketConnectionManager.Dispose();
         _certificate.Dispose();
-        _shutdownCancellationTokenSource.Dispose();
+        _shutdownCancellationSource.Dispose();
     }
 
     /// <inheritdoc/>
@@ -214,77 +227,63 @@ private Task StartListeningAsync()
         var app = builder.Build();
 
         app.MapGet("/", () => _displayName);
-        app.MapGet(InfoResponse.Url, GetInfoAsync);
+        app.MapGet(InfoResponse.Url, HandleInfoRequestAsync);
 
         // Set up the run session endpoints
         var runSessionApi = app.MapGroup(RunSessionRequest.Url);
 
-        runSessionApi.MapPut("/", RunSessionPutAsync);
-        runSessionApi.MapDelete("/{sessionId}", RunSessionDeleteAsync);
-        runSessionApi.Map(SessionNotification.Url, RunSessionNotifyAsync);
+        runSessionApi.MapPut("/", HandleStartSessionRequestAsync);
+        runSessionApi.MapDelete("/{sessionId}", HandleStopSessionRequestAsync);
+        runSessionApi.Map(SessionNotification.Url, HandleSessionNotifyRequestAsync);
 
         app.UseWebSockets(new WebSocketOptions
         {
              KeepAliveInterval = TimeSpan.FromSeconds(PingIntervalInSeconds)
         });
 
         // Run the application async. It will shutdown when the cancel token is signaled
-        return app.RunAsync(_shutdownCancellationTokenSource.Token);
+        return app.RunAsync(_shutdownCancellationSource.Token);
     }
 
-    private async Task RunSessionPutAsync(HttpContext context)
+    private bool TryAcceptRequest(HttpContext context)
     {
-        // Check the authentication header
-        if (!IsValidAuthentication(context))
-        {
-            Log("Authorization failure");
-            context.Response.StatusCode = (int)HttpStatusCode.Unauthorized;
-        }
-        else
+        if (_shutdownCancellationSource.IsCancellationRequested)
         {
-            await HandleStartSessionRequestAsync(context);
+            Log("Service unavailable -- shutdown in progress.");
+            context.Response.StatusCode = (int)HttpStatusCode.ServiceUnavailable;
+            return false;
         }
-    }
 
-    private async Task RunSessionDeleteAsync(HttpContext context, string sessionId)
-    {
         // Check the authentication header
         if (!IsValidAuthentication(context))
         {
             Log("Authorization failure");
             context.Response.StatusCode = (int)HttpStatusCode.Unauthorized;
+            return false;
         }
-        else
-        {
-            await HandleStopSessionRequestAsync(context, sessionId);
-        }
+
+        return true;
     }
 
-    private async Task GetInfoAsync(HttpContext context)
+    private async Task HandleInfoRequestAsync(HttpContext context)
     {
-        // Check the authentication header
-        if (!IsValidAuthentication(context))
+        if (!TryAcceptRequest(context))
         {
-            Log("Authorization failure");
-            context.Response.StatusCode = (int)HttpStatusCode.Unauthorized;
-        }
-        else
-        {
-            context.Response.StatusCode = (int)HttpStatusCode.OK;
-            await context.Response.WriteAsJsonAsync(InfoResponse.Instance, JsonSerializerOptions, _shutdownCancellationTokenSource.Token);
+            return;
         }
+
+        context.Response.StatusCode = (int)HttpStatusCode.OK;
+        await context.Response.WriteAsJsonAsync(InfoResponse.Instance, JsonSerializerOptions, _shutdownCancellationSource.Token);
     }
 
-    private async Task RunSessionNotifyAsync(HttpContext context)
+    private async Task HandleSessionNotifyRequestAsync(HttpContext context)
     {
-        // Check the authentication header
-        if (!IsValidAuthentication(context))
+        if (!TryAcceptRequest(context))
         {
-            Log("Authorization failure");
-            context.Response.StatusCode = (int)HttpStatusCode.Unauthorized;
             return;
         }
-        else if (!context.WebSockets.IsWebSocketRequest)
+
+        if (!context.WebSockets.IsWebSocketRequest)
         {
             context.Response.StatusCode = StatusCodes.Status400BadRequest;
             return;
@@ -323,6 +322,11 @@ private bool IsValidAuthentication(HttpContext context)
 
     private async Task HandleStartSessionRequestAsync(HttpContext context)
     {
+        if (!TryAcceptRequest(context))
+        {
+            return;
+        }
+
         string? projectPath = null;
 
         try
@@ -333,7 +337,7 @@ private async Task HandleStartSessionRequestAsync(HttpContext context)
             }
 
             // Get the project launch request data
-            var projectLaunchRequest = await context.GetProjectLaunchInformationAsync(_shutdownCancellationTokenSource.Token);
+            var projectLaunchRequest = await context.GetProjectLaunchInformationAsync(_shutdownCancellationSource.Token);
             if (projectLaunchRequest == null)
             {
                 // Unknown or unsupported version
@@ -343,7 +347,7 @@ private async Task HandleStartSessionRequestAsync(HttpContext context)
 
             projectPath = projectLaunchRequest.ProjectPath;
 
-            var sessionId = await _aspireServerEvents.StartProjectAsync(context.GetDcpId(), projectLaunchRequest, _shutdownCancellationTokenSource.Token);
+            var sessionId = await _aspireServerEvents.StartProjectAsync(context.GetDcpId(), projectLaunchRequest, _shutdownCancellationSource.Token);
 
             context.Response.StatusCode = (int)HttpStatusCode.Created;
             context.Response.Headers.Location = $"{context.Request.Scheme}://{context.Request.Host}{context.Request.Path}/{sessionId}";
@@ -370,14 +374,14 @@ private async Task WriteResponseTextAsync(HttpResponse response, Exception ex, b
                 Error = new ErrorDetail { ErrorCode = errorCode, Message = ex.GetMessageFromException() }
             };
 
-            await response.WriteAsJsonAsync(error, JsonSerializerOptions, _shutdownCancellationTokenSource.Token);
+            await response.WriteAsJsonAsync(error, JsonSerializerOptions, _shutdownCancellationSource.Token);
         }
         else
         {
             errorResponse = Encoding.UTF8.GetBytes(ex.GetMessageFromException());
             response.ContentType = "text/plain";
             response.ContentLength = errorResponse.Length;
-            await response.WriteAsync(ex.GetMessageFromException(), _shutdownCancellationTokenSource.Token);
+            await response.WriteAsync(ex.GetMessageFromException(), _shutdownCancellationSource.Token);
         }
     }
 
@@ -394,7 +398,7 @@ private async ValueTask<bool> SendMessageAsync(string dcpId, byte[] messageBytes
         try
         {
             using var cancelTokenSource = CancellationTokenSource.CreateLinkedTokenSource(
-                cancellationToken, _shutdownCancellationTokenSource.Token, connection.HttpRequestAborted);
+                cancellationToken, _shutdownCancellationSource.Token, connection.HttpRequestAborted);
 
             await _webSocketAccess.WaitAsync(cancelTokenSource.Token);
             await connection.Socket.SendAsync(new ArraySegment<byte>(messageBytes), WebSocketMessageType.Text, endOfMessage: true, cancelTokenSource.Token);
@@ -415,16 +419,21 @@ private async ValueTask<bool> SendMessageAsync(string dcpId, byte[] messageBytes
         return success;
     }
 
-    private async ValueTask HandleStopSessionRequestAsync(HttpContext context, string sessionId)
+    private async Task HandleStopSessionRequestAsync(HttpContext context, string sessionId)
     {
+        if (!TryAcceptRequest(context))
+        {
+            return;
+        }
+
         try
         {
             if (_isDisposed)
             {
                 throw new ObjectDisposedException(nameof(AspireServerService), "Received 'DELETE /run_session' request after the service has been disposed.");
             }
 
-            var sessionExists = await _aspireServerEvents.StopSessionAsync(context.GetDcpId(), sessionId, _shutdownCancellationTokenSource.Token);
+            var sessionExists = await _aspireServerEvents.StopSessionAsync(context.GetDcpId(), sessionId, _shutdownCancellationSource.Token);
             context.Response.StatusCode = (int)(sessionExists ? HttpStatusCode.OK : HttpStatusCode.NoContent);
         }
         catch (Exception e) when (e is not OperationCanceledException)

src/Dotnet.Watch/HotReloadClient/ClientTransport.cs

@@ -44,5 +44,10 @@ internal abstract class ClientTransport : IDisposable
     /// </summary>
     public abstract ValueTask<ClientTransportResponse?> ReadAsync(CancellationToken cancellationToken);
 
+    /// <summary>
+    /// True if <paramref name="exception"/> and <paramref name="cancellationToken"/> indicate an expected connection termination signal.
+    /// </summary>
+    public abstract bool IsExpectedConnectionTermination(Exception exception, CancellationToken cancellationToken);
+
     public abstract void Dispose();
 }

src/Dotnet.Watch/HotReloadClient/DefaultHotReloadClient.cs

@@ -82,7 +82,7 @@ async Task<ImmutableArray<string>> ConnectAsync()
                 {
                     // Don't report a warning when cancelled. The process has terminated or the host is shutting down in that case.
                     // Best effort: There is an inherent race condition due to time between the process exiting and the cancellation token triggering.
-                    if (!cancellationToken.IsCancellationRequested)
+                    if (!transport.IsExpectedConnectionTermination(e, cancellationToken))
                     {
                         Logger.LogError("Failed to read capabilities: {Message}", e.Message);
                     }
@@ -127,7 +127,7 @@ private async Task ListenForResponsesAsync(CancellationToken cancellationToken)
             }
             catch (Exception e) when (e is not OperationCanceledException)
             {
-                if (!cancellationToken.IsCancellationRequested)
+                if (!transport.IsExpectedConnectionTermination(e, cancellationToken))
                 {
                     Logger.LogError("Failed to read response: {Exception}", e.ToString());
                 }

src/Dotnet.Watch/HotReloadClient/NamedPipeClientTransport.cs

@@ -50,30 +50,14 @@ public override void ConfigureEnvironment(IDictionary<string, string> env)
     public override async Task WaitForConnectionAsync(CancellationToken cancellationToken)
     {
         _logger.LogDebug("Waiting for application to connect to pipe '{NamedPipeName}'.", _namedPipeName);
-
-        try
-        {
-            await _pipe.WaitForConnectionAsync(cancellationToken);
-        }
-        catch (Exception e) when (e is not OperationCanceledException)
-        {
-            // The process may die while we're waiting for the connection and the pipe may be disposed.
-            // Log and let subsequent ReadAsync return null gracefully.
-            if (IsExpectedPipeException(e, cancellationToken))
-            {
-                _logger.LogDebug("Pipe connection ended: {Message}", e.Message);
-                return;
-            }
-
-            throw;
-        }
+        await _pipe.WaitForConnectionAsync(cancellationToken);
     }
 
     /// <summary>
     /// Returns true if the exception is expected when the pipe is disposed or the process has terminated.
     /// On Unix named pipes can also throw SocketException with ErrorCode 125 (Operation canceled) when disposed.
     /// </summary>
-    private static bool IsExpectedPipeException(Exception e, CancellationToken cancellationToken)
+    public override bool IsExpectedConnectionTermination(Exception e, CancellationToken cancellationToken)
     {
         return e is ObjectDisposedException or EndOfStreamException or SocketException { ErrorCode: 125 }
             || cancellationToken.IsCancellationRequested;
@@ -93,16 +77,8 @@ public override async ValueTask WriteAsync(byte type, Func<Stream, CancellationT
 
     public override async ValueTask<ClientTransportResponse?> ReadAsync(CancellationToken cancellationToken)
     {
-        try
-        {
-            var type = (ResponseType)await _pipe.ReadByteAsync(cancellationToken);
-            return new ClientTransportResponse(type, _pipe, disposeStream: false);
-        }
-        catch (Exception e) when (e is not OperationCanceledException && IsExpectedPipeException(e, cancellationToken))
-        {
-            // Pipe has been disposed or the process has terminated.
-            return null;
-        }
+        var type = (ResponseType)await _pipe.ReadByteAsync(cancellationToken);
+        return new ClientTransportResponse(type, _pipe, disposeStream: false);
     }
 
     public override void Dispose()

src/Dotnet.Watch/HotReloadClient/Web/KestrelWebSocketServer.cs

@@ -39,10 +39,6 @@ public ImmutableArray<string> ServerUrls
     /// <summary>
     /// Starts the Kestrel WebSocket server.
     /// </summary>
-    /// <param name="hostName">Host name to bind to</param>
-    /// <param name="port">HTTP port to bind to (0 for auto-assign)</param>
-    /// <param name="securePort">HTTPS port to bind to in addition to HTTP port. Null to skip HTTPS.</param>
-    /// <param name="cancellationToken">Cancellation token</param>
     public static async ValueTask<KestrelWebSocketServer> StartServerAsync(WebSocketConfig config, RequestDelegate requestHandler, CancellationToken cancellationToken)
     {
         var host = new HostBuilder()

src/Dotnet.Watch/HotReloadClient/WebSocketClientTransport.cs

@@ -43,6 +43,9 @@ public override void Dispose()
         _handler.Dispose();
     }
 
+    public override bool IsExpectedConnectionTermination(Exception exception, CancellationToken cancellationToken)
+        => cancellationToken.IsCancellationRequested;
+
     /// <summary>
     /// Creates and starts a new <see cref="WebSocketClientTransport"/> instance.
     /// </summary>

src/Dotnet.Watch/Watch.Aspire/Server/AspireServerLauncher.cs

@@ -3,12 +3,6 @@
 
 using System.Collections.Immutable;
 using System.CommandLine;
-using System.Threading.Channels;
-using Microsoft.CodeAnalysis.Elfie.Diagnostics;
-using Microsoft.DotNet.HotReload;
-using Microsoft.Extensions.Logging;
-using Microsoft.Extensions.Logging.Abstractions;
-using static System.Runtime.InteropServices.JavaScript.JSType;
 
 namespace Microsoft.DotNet.Watch;
 

src/Dotnet.Watch/Watch.Aspire/Server/AspireWatcherLauncher.cs

@@ -2,7 +2,6 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Collections.Immutable;
-using Microsoft.CodeAnalysis.Elfie.Diagnostics;
 using Microsoft.Extensions.Logging;
 
 namespace Microsoft.DotNet.Watch;

src/Dotnet.Watch/Watch.Aspire/Server/ProcessLauncherFactory.cs

@@ -7,7 +7,6 @@
 using System.Runtime.CompilerServices;
 using System.Text.Json;
 using System.Threading.Channels;
-using Aspire.Tools.Service;
 using Microsoft.CodeAnalysis;
 using Microsoft.DotNet.HotReload;
 using Microsoft.Extensions.Logging;

src/Dotnet.Watch/Watch.Aspire/Server/StatusReportingLoggerFactory.cs

@@ -7,6 +7,7 @@ namespace Microsoft.DotNet.Watch;
 
 /// <summary>
 /// Intercepts select log messages reported by watch and forwards them to <see cref="WatchStatusWriter"/> to be sent to an external listener.
+/// Does not own (dispose) <paramref name="writer"/> and <paramref name="underlyingFactory"/>.
 /// </summary>
 internal sealed class StatusReportingLoggerFactory(WatchStatusWriter writer, LoggerFactory underlyingFactory) : ILoggerFactory
 {

src/Dotnet.Watch/Watch.Aspire/Server/WatchControlReader.cs

@@ -46,6 +46,8 @@ public async ValueTask DisposeAsync()
         {
             // Pipe may already be broken if the server disconnected
         }
+
+        _disposalCancellationSource.Dispose();
     }
 
     private async Task ListenAsync(CancellationToken cancellationToken)