Add support for timeout in process termination handling

src/System.CommandLine.ApiCompatibility.Tests/ApiCompatibilityApprovalTests.System_CommandLine_api_is_not_changed.approved.txt

@@ -240,7 +240,7 @@ System.CommandLine.Builder
   public static class CommandLineBuilderExtensions
     public static CommandLineBuilder AddMiddleware(this CommandLineBuilder builder, System.CommandLine.Invocation.InvocationMiddleware middleware, System.CommandLine.Invocation.MiddlewareOrder order = Default)
     public static CommandLineBuilder AddMiddleware(this CommandLineBuilder builder, System.Action<System.CommandLine.Invocation.InvocationContext> onInvoke, System.CommandLine.Invocation.MiddlewareOrder order = Default)
-    public static CommandLineBuilder CancelOnProcessTermination(this CommandLineBuilder builder)
+    public static CommandLineBuilder CancelOnProcessTermination(this CommandLineBuilder builder, System.Nullable<System.TimeSpan> timeout = null)
     public static CommandLineBuilder EnableDirectives(this CommandLineBuilder builder, System.Boolean value = True)
     public static CommandLineBuilder EnableLegacyDoubleDashBehavior(this CommandLineBuilder builder, System.Boolean value = True)
     public static CommandLineBuilder EnablePosixBundling(this CommandLineBuilder builder, System.Boolean value = True)

src/System.CommandLine.Tests/Invocation/CancelOnProcessTerminationTests.cs

@@ -20,9 +20,9 @@ public class CancelOnProcessTerminationTests
         private const int SIGTERM = 15;
 
         [LinuxOnlyTheory]
-        [InlineData(SIGINT, Skip = "https://github.com/dotnet/command-line-api/issues/1206")]  // Console.CancelKeyPress
+        [InlineData(SIGINT/*, Skip = "https://github.com/dotnet/command-line-api/issues/1206"*/)]  // Console.CancelKeyPress
         [InlineData(SIGTERM)] // AppDomain.CurrentDomain.ProcessExit
-        public async Task CancelOnProcessTermination_cancels_on_process_termination(int signo)
+        public async Task CancelOnProcessTermination_provides_CancellationToken_that_signals_termination_when_no_timeout_is_specified(int signo)
         {
             const string ChildProcessWaiting = "Waiting for the command to be cancelled";
             const int CancelledExitCode = 42;
@@ -91,6 +91,170 @@ public async Task CancelOnProcessTermination_cancels_on_process_termination(int
             process.ExitCode.Should().Be(CancelledExitCode);
         }
 
+        [LinuxOnlyTheory]
+        [InlineData(SIGINT)]
+        [InlineData(SIGTERM)]
+        public async Task CancelOnProcessTermination_provides_CancellationToken_that_signals_termination_when_null_timeout_is_specified(int signo)
+        {
+            const string ChildProcessWaiting = "Waiting for the command to be cancelled";
+            const int CancelledExitCode = 42;
+
+            Func<string[], Task<int>> childProgram = (string[] args) =>
+            {
+                var command = new Command("the-command");
+
+                command.SetHandler(async context =>
+                {
+                    var cancellationToken = context.GetCancellationToken();
+
+                    try
+                    {
+                        context.Console.WriteLine(ChildProcessWaiting);
+                        await Task.Delay(int.MaxValue, cancellationToken);
+                        context.ExitCode = 1;
+                    }
+                    catch (OperationCanceledException)
+                    {
+                        // For Process.Exit handling the event must remain blocked as long as the
+                        // command is executed.
+                        // We are currently blocking that event because CancellationTokenSource.Cancel
+                        // is called from the event handler.
+                        // We'll do an async Yield now. This means the Cancel call will return
+                        // and we're no longer actively blocking the event.
+                        // The event handler is responsible to continue blocking until the command
+                        // has finished executing. If it doesn't we won't get the CancelledExitCode.
+                        await Task.Yield();
+
+                        // Exit code gets set here - but then execution continues and is let run till code voluntarily returns
+                        //  hence exit code gets overwritten below
+                        context.ExitCode = 123;
+                    }
+
+                    // This is an example of bad pattern and reason why we need a timeout on termination processing
+                    await Task.Delay(TimeSpan.FromMilliseconds(200));
+
+                    context.ExitCode = CancelledExitCode;
+                });
+
+                return new CommandLineBuilder(new RootCommand
+                       {
+                           command
+                       })
+                       // Unfortunately we cannot use test parameter here - RemoteExecutor currently doesn't capture the closure
+                       .CancelOnProcessTermination(null)
+                       .Build()
+                       .InvokeAsync("the-command");
+            };
+
+            using RemoteExecution program = RemoteExecutor.Execute(childProgram, psi: new ProcessStartInfo { RedirectStandardOutput = true });
+
+            Process process = program.Process;
+
+            // Wait for the child to be in the command handler.
+            string childState = await process.StandardOutput.ReadLineAsync();
+            childState.Should().Be(ChildProcessWaiting);
+
+            // Request termination
+            kill(process.Id, signo).Should().Be(0);
+
+            // Verify the process terminates timely
+            bool processExited = process.WaitForExit(10000);
+            if (!processExited)
+            {
+                process.Kill();
+                process.WaitForExit();
+            }
+            processExited.Should().Be(true);
+
+            // Verify the process exit code
+            process.ExitCode.Should().Be(CancelledExitCode);
+        }
+
+        [LinuxOnlyTheory]
+        [InlineData(SIGINT)]
+        [InlineData(SIGTERM)]
+        public async Task CancelOnProcessTermination_provides_CancellationToken_that_signals_termination_and_execution_is_terminated_at_the_specified_timeout(int signo)
+        {
+            const string ChildProcessWaiting = "Waiting for the command to be cancelled";
+            const int CancelledExitCode = 42;
+            const int ForceTerminationCode = 130;
+
+            Func<string[], Task<int>> childProgram = (string[] args) =>
+            {
+                var command = new Command("the-command");
+
+                command.SetHandler(async context =>
+                {
+                    var cancellationToken = context.GetCancellationToken();
+
+                    try
+                    {
+                        context.Console.WriteLine(ChildProcessWaiting);
+                        await Task.Delay(int.MaxValue, cancellationToken);
+                        context.ExitCode = 1;
+                    }
+                    catch (OperationCanceledException)
+                    {
+                        // For Process.Exit handling the event must remain blocked as long as the
+                        // command is executed.
+                        // We are currently blocking that event because CancellationTokenSource.Cancel
+                        // is called from the event handler.
+                        // We'll do an async Yield now. This means the Cancel call will return
+                        // and we're no longer actively blocking the event.
+                        // The event handler is responsible to continue blocking until the command
+                        // has finished executing. If it doesn't we won't get the CancelledExitCode.
+                        await Task.Yield();
+
+                        context.ExitCode = CancelledExitCode;
+
+                        // This is an example of bad pattern and reason why we need a timeout on termination processing
+                        await Task.Delay(TimeSpan.FromMilliseconds(1000));
+
+                        // Execution should newer get here as termination processing has a timeout of 100ms
+                        Environment.Exit(123);
+                    }
+
+                    // This is an example of bad pattern and reason why we need a timeout on termination processing
+                    await Task.Delay(TimeSpan.FromMilliseconds(1000));
+
+                    // Execution should newer get here as termination processing has a timeout of 100ms
+                    Environment.Exit(123);
+                });
+
+                return new CommandLineBuilder(new RootCommand
+                       {
+                           command
+                       })
+                        // Unfortunately we cannot use test parameter here - RemoteExecutor currently doesn't capture the closure
+                       .CancelOnProcessTermination(TimeSpan.FromMilliseconds(100))
+                       .Build()
+                       .InvokeAsync("the-command");
+            };
+
+            using RemoteExecution program = RemoteExecutor.Execute(childProgram, psi: new ProcessStartInfo { RedirectStandardOutput = true });
+
+            Process process = program.Process;
+
+            // Wait for the child to be in the command handler.
+            string childState = await process.StandardOutput.ReadLineAsync();
+            childState.Should().Be(ChildProcessWaiting);
+
+            // Request termination
+            kill(process.Id, signo).Should().Be(0);
+
+            // Verify the process terminates timely
+            bool processExited = process.WaitForExit(10000);
+            if (!processExited)
+            {
+                process.Kill();
+                process.WaitForExit();
+            }
+            processExited.Should().Be(true);
+
+            // Verify the process exit code
+            process.ExitCode.Should().Be(ForceTerminationCode);
+        }
+
         [DllImport("libc", SetLastError = true)]
         private static extern int kill(int pid, int sig);
     }

src/System.CommandLine/Builder/CommandLineBuilderExtensions.cs

@@ -10,6 +10,7 @@
 using System.Linq;
 using System.Reflection;
 using System.Threading;
+using System.Threading.Tasks;
 using static System.Environment;
 using Process = System.CommandLine.Invocation.Process;
 
@@ -42,9 +43,25 @@ public static class CommandLineBuilderExtensions
         /// Enables signaling and handling of process termination via a <see cref="CancellationToken"/> that can be passed to a <see cref="ICommandHandler"/> during invocation.
         /// </summary>
         /// <param name="builder">A command line builder.</param>
+        /// <param name="timeout">
+        /// Optional timeout for the command to process the exit cancellation.
+        /// If not passed, or passed null or non-positive timeout (including <see cref="Timeout.InfiniteTimeSpan"/>), no timeout is enforced.
+        /// If positive value is passed - command is forcefully terminated after the timeout with exit code 130 (as if <see cref="CancelOnProcessTermination"/> was not called).
+        /// Host enforced timeout for ProcessExit event cannot be extended - default is 2 seconds: https://docs.microsoft.com/en-us/dotnet/api/system.appdomain.processexit?view=net-6.0.
+        /// </param>
         /// <returns>The same instance of <see cref="CommandLineBuilder"/>.</returns>
-        public static CommandLineBuilder CancelOnProcessTermination(this CommandLineBuilder builder)
+        public static CommandLineBuilder CancelOnProcessTermination(
+            this CommandLineBuilder builder, 
+            TimeSpan? timeout = null)
         {
+            // https://tldp.org/LDP/abs/html/exitcodes.html - 130 - script terminated by ctrl-c
+            const int SIGINT_EXIT_CODE = 130;
+
+            if (timeout == null || timeout.Value < TimeSpan.Zero)
+            {
+                timeout = Timeout.InfiniteTimeSpan;
+            }
+
             builder.AddMiddleware(async (context, next) =>
             {
                 bool cancellationHandlingAdded = false;
@@ -56,24 +73,58 @@ public static CommandLineBuilder CancelOnProcessTermination(this CommandLineBuil
                 {
                     blockProcessExit = new ManualResetEventSlim(initialState: false);
                     cancellationHandlingAdded = true;
-                    consoleHandler = (_, args) =>
-                    {
-                        cts.Cancel();
-                        // Stop the process from terminating.
-                        // Since the context was cancelled, the invocation should
-                        // finish and Main will return.
-                        args.Cancel = true;
-                    };
+                    // Default limit for ProcesExit handler is 2 seconds
+                    //  https://docs.microsoft.com/en-us/dotnet/api/system.appdomain.processexit?view=net-6.0
                     processExitHandler = (_, _) =>
                     {
-                        cts.Cancel();
+                        // Cancel asynchronously not to block the handler (as then the process might possibly run longer then what was the requested timeout)
+                        Task timeoutTask = Task.Delay(timeout.Value);
+                        Task cancelTask = Task.Factory.StartNew(cts.Cancel);
+
                         // The process exits as soon as the event handler returns.
                         // We provide a return value using Environment.ExitCode
                         // because Main will not finish executing.
                         // Wait for the invocation to finish.
-                        blockProcessExit.Wait();
+                        if (!blockProcessExit.Wait(timeout > TimeSpan.Zero
+                                ? timeout.Value
+                                : Timeout.InfiniteTimeSpan))
+                        {
+                            context.ExitCode = SIGINT_EXIT_CODE;
+                        }
+                        // Let's block here (to prevent process bailing out) for the rest of the timeout (if any), for cancellation to finish (if it hasn't yet)
+                        else if (Task.WaitAny(timeoutTask, cancelTask) == 0)
+                        {
+                            // The async cancellation didn't finish in timely manner
+                            context.ExitCode = SIGINT_EXIT_CODE;
+                        }
                         ExitCode = context.ExitCode;
                     };
+                    consoleHandler = (_, args) =>
+                    {
+                        // Stop the process from terminating.
+                        // Since the context was cancelled, the invocation should
+                        // finish and Main will return.
+                        args.Cancel = true;
+
+                        // If timeout was requested - make sure cancellation processing (or any other activity within the current process)
+                        //  doesn't keep the process running after the timeout
+                        if (timeout! > TimeSpan.Zero)
+                        {
+                            Task
+                                .Delay(timeout.Value, default)
+                                .ContinueWith(t =>
+                                {
+                                    // Prevent our ProcessExit from intervene and block the exit
+                                    AppDomain.CurrentDomain.ProcessExit -= processExitHandler;
+                                    Environment.Exit(SIGINT_EXIT_CODE);
+                                }, (CancellationToken)default);
+                        }
+
+                        // Cancel synchronously here - no need to perform it asynchronously as the timeout is already running (and would kill the process if needed),
+                        //  plus we cannot wait only on the cancellation (e.g. via `Task.Factory.StartNew(cts.Cancel).Wait(cancelationProcessingTimeout.Value)`)
+                        //  as we need to abort any other possible execution within the process - even outside the context of cancellation processing
+                        cts.Cancel();
+                    };
                     Console.CancelKeyPress += consoleHandler;
                     AppDomain.CurrentDomain.ProcessExit += processExitHandler;
                 };