Improve ConsoleCancellationManager with process termination timeout

src/Aspire.Cli/Commands/BaseCommand.cs

@@ -58,7 +58,7 @@ protected BaseCommand(string name, string description, IFeatures features, ICliU
             catch (NonInteractiveException)
             {
                 // Error messages have already been displayed by the interaction service.
-                result = CommandResult.Failure((int)CliExitCodes.MissingRequiredArgument);
+                result = CommandResult.Failure(CliExitCodes.MissingRequiredArgument);
             }
 
             var isErrorExitCode = result.ExitCode != CliExitCodes.Success;

src/Aspire.Cli/ConsoleCancellationManager.cs

@@ -6,55 +6,125 @@
 namespace Aspire.Cli;
 
 /// <summary>
-/// Manages Ctrl+C and SIGTERM signal handling with a shared CancellationTokenSource.
+/// Manages Ctrl+C, SIGINT, and SIGTERM signal handling with a shared CancellationTokenSource.
+/// After cancellation is requested, waits up to <c>processTerminationTimeout</c> for the running
+/// handler to complete before signaling forced termination via <see cref="ProcessTerminationCompletionSource"/>.
 /// Disposing this instance unregisters all signal handlers and disposes the token source.
 /// </summary>
 internal sealed class ConsoleCancellationManager : IDisposable
 {
+    private const int SigIntExitCode = 130;
+    private const int SigTermExitCode = 143;
+
     private readonly CancellationTokenSource _cts = new();
-    private readonly CancellationToken _token;
-    private readonly ConsoleCancelEventHandler _cancelKeyPressHandler;
+    private readonly TimeSpan _processTerminationTimeout;
+    private readonly PosixSignalRegistration? _sigIntRegistration;
     private readonly PosixSignalRegistration? _sigTermRegistration;
+    private readonly CancellationToken _token;
+    private Task<int>? _startedHandler;
+    private int _cancelCalled;
+
+    internal readonly TaskCompletionSource<int> _processTerminationCompletionSource = new(TaskCreationOptions.RunContinuationsAsynchronously);
 
-    public ConsoleCancellationManager()
+    /// <summary>
+    /// A completion source that is signaled with a native exit code when the running handler
+    /// does not complete within the configured timeout after a termination signal.
+    /// </summary>
+    internal TaskCompletionSource<int> ProcessTerminationCompletionSource => _processTerminationCompletionSource;
+
+    /// <summary>
+    /// Sets the handler task that represents the currently executing command. When a termination
+    /// signal arrives, the manager will wait for this task to complete within the configured timeout.
+    /// </summary>
+    internal void SetStartedHandler(Task<int> handler) => Volatile.Write(ref _startedHandler, handler);
+
+    public ConsoleCancellationManager(TimeSpan processTerminationTimeout)
     {
+        _processTerminationTimeout = processTerminationTimeout;
+
+        // Set to a field so getting the token doesn't error after dispose.
         _token = _cts.Token;
-        _cancelKeyPressHandler = (sender, eventArgs) =>
+
+        // Prefer PosixSignalRegistration for both SIGINT and SIGTERM as it handles
+        // both signals uniformly and allows cancelling SIGTERM (which Console.CancelKeyPress cannot).
+        // Despite the name, PosixSignalRegistration is supported on Windows: the runtime maps
+        // SIGINT to CTRL_C_EVENT and SIGTERM to CTRL_CLOSE_EVENT/CTRL_SHUTDOWN_EVENT.
+        if (!OperatingSystem.IsAndroid()
+            && !OperatingSystem.IsIOS()
+            && !OperatingSystem.IsTvOS()
+            && !OperatingSystem.IsBrowser())
         {
-            TryCancel();
-            eventArgs.Cancel = true;
-        };
-        Console.CancelKeyPress += _cancelKeyPressHandler;
-
-        _sigTermRegistration = OperatingSystem.IsWindows()
-            ? null
-            : PosixSignalRegistration.Create(PosixSignal.SIGTERM, context =>
-            {
-                TryCancel();
-                context.Cancel = true;
-            });
+            _sigIntRegistration = PosixSignalRegistration.Create(PosixSignal.SIGINT, OnPosixSignal);
+            _sigTermRegistration = PosixSignalRegistration.Create(PosixSignal.SIGTERM, OnPosixSignal);
+        }
+
+        Console.CancelKeyPress += OnCancelKeyPress;
+        AppDomain.CurrentDomain.ProcessExit += OnProcessExit;
     }
 
     public CancellationToken Token => _token;
 
     public bool IsCancellationRequested => _cts.IsCancellationRequested;
 
-    private void TryCancel()
+    private void OnPosixSignal(PosixSignalContext context)
+    {
+        context.Cancel = true;
+        Cancel(context.Signal == PosixSignal.SIGINT ? SigIntExitCode : SigTermExitCode);
+    }
+
+    private void OnCancelKeyPress(object? sender, ConsoleCancelEventArgs e)
     {
+        e.Cancel = true;
+        Cancel(SigIntExitCode);
+    }
+
+    private void OnProcessExit(object? sender, EventArgs e) => Cancel(SigTermExitCode);
+
+    private void Cancel(int forcedTerminationExitCode)
+    {
+        // Ensure only the first signal triggers cancellation logic; subsequent signals are no-ops.
+        if (Interlocked.CompareExchange(ref _cancelCalled, 1, 0) != 0)
+        {
+            return;
+        }
+
+        // Request cancellation so cooperative listeners can begin shutting down.
         try
         {
             _cts.Cancel();
         }
         catch (ObjectDisposedException)
         {
             // A signal can race with process shutdown after cancellation resources are disposed.
+            return;
+        }
+
+        try
+        {
+            var startedHandler = Volatile.Read(ref _startedHandler);
+
+            // Wait for the configured interval to allow graceful shutdown.
+            if (startedHandler is null || !startedHandler.Wait(_processTerminationTimeout))
+            {
+                // If the handler does not finish within configured time, use the completion
+                // source to signal forced completion (preserving native exit code).
+                _processTerminationCompletionSource.TrySetResult(forcedTerminationExitCode);
+            }
+        }
+        catch (AggregateException)
+        {
+            // The task was cancelled or an exception was thrown during task execution.
         }
     }
 
     public void Dispose()
     {
-        Console.CancelKeyPress -= _cancelKeyPressHandler;
+        _sigIntRegistration?.Dispose();
         _sigTermRegistration?.Dispose();
+
+        Console.CancelKeyPress -= OnCancelKeyPress;
+        AppDomain.CurrentDomain.ProcessExit -= OnProcessExit;
+
         _cts.Dispose();
     }
 }

src/Aspire.Cli/Program.cs

@@ -729,7 +729,7 @@ public static async Task<int> Main(string[] args)
         // Setup handling of CTRL-C and SIGTERM as early as possible so that if
         // we get a signal anywhere that is not handled by Spectre Console
         // already that we know to trigger cancellation.
-        using var cancellationManager = new ConsoleCancellationManager();
+        using var cancellationManager = new ConsoleCancellationManager(processTerminationTimeout: TimeSpan.FromSeconds(5));
 
         Console.OutputEncoding = Encoding.UTF8;
 
@@ -788,7 +788,9 @@ public static async Task<int> Main(string[] args)
         var invokeConfig = new InvocationConfiguration()
         {
             // Disable default exception handler so we can log exceptions to telemetry.
-            EnableDefaultExceptionHandler = false
+            EnableDefaultExceptionHandler = false,
+            // Set timeout to null so that System.Commandline doesn't manage cancellation tokens or timeouts for us.
+            ProcessTerminationTimeout = null
         };
 
         app.Services.GetRequiredService<CliExecutionContext>();
@@ -833,7 +835,21 @@ public static async Task<int> Main(string[] args)
                         profileCommandActivity = profilingTelemetry.StartCommand(commandName);
                     }
 
-                    exitCode = await parseResult.InvokeAsync(invokeConfig, cancellationManager.Token);
+                    // Parse commandline and invoke the handler.
+                    var handlerTask = parseResult.InvokeAsync(invokeConfig, cancellationManager.Token);
+                    cancellationManager.SetStartedHandler(handlerTask);
+
+                    // Wait for either the handler to complete or a termination signal to trigger cancellation and timeout.
+                    var firstCompletedTask = await Task.WhenAny(handlerTask, cancellationManager.ProcessTerminationCompletionSource.Task);
+                    if (firstCompletedTask != handlerTask)
+                    {
+                        // The termination signal triggered cancellation and the timeout has completed. Kill the process.
+                        logger.LogWarning("Timeout waiting for cancellation from termination signal.");
+                    }
+
+                    exitCode = await firstCompletedTask; // return the result or propagate the exception
+
+                    // Set telemetry tags based on how the command completed.
                     profileCommandActivity.SetProcessExitCode(exitCode);
                     if (exitCode != CliExitCodes.Success)
                     {
@@ -848,33 +864,24 @@ public static async Task<int> Main(string[] args)
                 // Log exit code for debugging
                 logger.LogInformation("Exit code: {ExitCode}", exitCode);
             }
+            catch (OperationCanceledException ex) when (cancellationManager.IsCancellationRequested || ex is ExtensionOperationCanceledException)
+            {
+                exitCode = CliExitCodes.Cancelled;
+            }
             catch (Exception ex)
             {
-                exitCode = 1;
-                // Catch block is used instead of System.Commandline's default handler behavior.
-                // Allows logging of exceptions to telemetry.
-
-                // Don't log or display cancellation exceptions.
-                // Check both Ctrl+C cancellation (cancellationManager.IsCancellationRequested) and
-                // extension prompt cancellation (ExtensionOperationCanceledException).
-                if (!(ex is OperationCanceledException && cancellationManager.IsCancellationRequested) && ex is not ExtensionOperationCanceledException)
-                {
-                    logger.LogError(ex, "An unexpected error occurred.");
-
-                    telemetry.RecordError("An unexpected error occurred.", ex);
-
-                    errorWriter.WriteLine(string.Format(CultureInfo.CurrentCulture, InteractionServiceStrings.UnexpectedErrorOccurred, ex.Message));
-                }
-
-                // Log exit code for debugging
-                logger.LogError("Exit code: {ExitCode} (exception)", exitCode);
+                exitCode = CliExitCodes.InvalidCommand;
 
+                logger.LogError(ex, "An unexpected error occurred.");
+                telemetry.RecordError("An unexpected error occurred.", ex);
+                errorWriter.WriteLine(string.Format(CultureInfo.CurrentCulture, InteractionServiceStrings.UnexpectedErrorOccurred, ex.Message));
+            }
+            finally
+            {
                 mainActivity?.SetTag(TelemetryConstants.Tags.ProcessExitCode, exitCode);
+                mainActivity?.Stop();
             }
 
-            mainActivity?.SetTag(TelemetryConstants.Tags.ProcessExitCode, exitCode);
-            mainActivity?.Stop();
-
             if (profileCaptureSession is not null)
             {
                 try