Reworking the code so it doesn't register for events without a cancellation token

Author: Keboo

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

@@ -374,7 +374,7 @@ System.CommandLine.Invocation
     public System.CommandLine.LocalizationResources LocalizationResources { get; }
     public System.CommandLine.Parsing.Parser Parser { get; }
     public System.CommandLine.Parsing.ParseResult ParseResult { get; set; }
-    public System.Void AddLinkedCancellationToken(System.Threading.CancellationToken token)
+    public System.Void AddLinkedCancellationToken(System.Func<System.Threading.CancellationToken> token)
     public System.Void Dispose()
   public delegate InvocationMiddleware : System.MulticastDelegate, System.ICloneable, System.Runtime.Serialization.ISerializable
     .ctor(System.Object object, System.IntPtr method)

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

@@ -35,7 +35,7 @@ public void InvocationContext_with_linked_cancellation_token_can_cancel_by_passe
                 .Build()
                 .Parse("");
             using InvocationContext context = new(parseResult, cancellationToken: cts1.Token);
-            context.AddLinkedCancellationToken(cts2.Token);
+            context.AddLinkedCancellationToken(() => cts2.Token);
 
             var token = context.CancellationToken;
 
@@ -53,7 +53,7 @@ public void InvocationContext_with_linked_cancellation_token_can_cancel_by_linke
                 .Build()
                 .Parse("");
             using InvocationContext context = new(parseResult, cancellationToken: cts1.Token);
-            context.AddLinkedCancellationToken(cts2.Token);
+            context.AddLinkedCancellationToken(() => cts2.Token);
 
             var token = context.CancellationToken;
 
@@ -71,10 +71,10 @@ public void InvocationContext_adding_additional_linked_token_after_token_has_bee
                 .Build()
                 .Parse("");
             using InvocationContext context = new(parseResult, cancellationToken: cts.Token);
-            context.AddLinkedCancellationToken(cts.Token);
+            context.AddLinkedCancellationToken(() => cts.Token);
             _ = context.CancellationToken;
 
-            Assert.Throws<InvalidOperationException>(() => context.AddLinkedCancellationToken(cts.Token));
+            Assert.Throws<InvalidOperationException>(() => context.AddLinkedCancellationToken(() => cts.Token));
         }
     }
 }

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

@@ -355,7 +355,7 @@ public async Task Command_InvokeAsync_can_cancel_from_middleware()
                          })
                          .AddMiddleware(async (context, next) =>
                          {
-                             context.AddLinkedCancellationToken(cts.Token);
+                             context.AddLinkedCancellationToken(() => cts.Token);
                              cts.Cancel();
                              await next(context);
                          })

src/System.CommandLine/Builder/CommandLineBuilderExtensions.cs

@@ -47,32 +47,39 @@ public static CommandLineBuilder CancelOnProcessTermination(this CommandLineBuil
         {
             builder.AddMiddleware(async (context, next) =>
             {
-                //TODO: This CancellationTokenSource is never disposed...
-                CancellationTokenSource cts = new();
+                ConsoleCancelEventHandler? consoleHandler = null;
+                EventHandler? processExitHandler = null;
+                ManualResetEventSlim? blockProcessExit = null;
 
-                context.AddLinkedCancellationToken(cts.Token);
-
-                ManualResetEventSlim blockProcessExit = new(initialState: false);
-                ConsoleCancelEventHandler 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;
-                };
-                EventHandler processExitHandler = (_, _) =>
+                context.AddLinkedCancellationToken(() => 
                 {
-                    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();
-                    ExitCode = context.ExitCode;
-                };
-                Console.CancelKeyPress += consoleHandler;
-                AppDomain.CurrentDomain.ProcessExit += processExitHandler;
+                    //TODO: This CancellationTokenSource is never disposed...
+                    CancellationTokenSource cts = new();
+                    blockProcessExit = new ManualResetEventSlim(initialState: false);
+                    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;
+                    };
+                    processExitHandler = (_, _) =>
+                    {
+                        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();
+                        ExitCode = context.ExitCode;
+                    };
+                    Console.CancelKeyPress += consoleHandler;
+                    AppDomain.CurrentDomain.ProcessExit += processExitHandler;
+
+                    return cts.Token;
+                });
+
 
                 try
                 {
@@ -82,7 +89,7 @@ public static CommandLineBuilder CancelOnProcessTermination(this CommandLineBuil
                 {
                     Console.CancelKeyPress -= consoleHandler;
                     AppDomain.CurrentDomain.ProcessExit -= processExitHandler;
-                    blockProcessExit.Set();
+                    blockProcessExit?.Set();
                 }
             }, MiddlewareOrderInternal.Startup);
 
@@ -439,7 +446,7 @@ public static CommandLineBuilder UseHelp(
             int? maxWidth = null)
         {
             builder.CustomizeHelpLayout(customize);
-            
+
             if (builder.HelpOption is null)
             {
                 builder.UseHelp(new HelpOption(() => builder.LocalizationResources), maxWidth);

src/System.CommandLine/Invocation/InvocationContext.cs

@@ -16,8 +16,7 @@ public sealed class InvocationContext : IDisposable
     {
         private HelpBuilder? _helpBuilder;
         private CancellationTokenSource? _linkedTokensSource;
-        private List<CancellationToken>? _linkedTokens;
-        private CancellationToken _cancellationToken;
+        private List<Func<CancellationToken>> _cancellationTokens = new(3);
         private Lazy<CancellationToken> _lazyCancellationToken;
 
         /// <param name="parseResult">The result of the current parse operation.</param>
@@ -32,7 +31,7 @@ public InvocationContext(
             BindingContext.ServiceProvider.AddService(_ => CancellationToken);
             BindingContext.ServiceProvider.AddService(_ => this);
 
-            _cancellationToken = cancellationToken;
+            _cancellationTokens.Add(() => cancellationToken);
             _lazyCancellationToken = new Lazy<CancellationToken>(BuildCancellationToken);
         }
 
@@ -88,15 +87,19 @@ public ParseResult ParseResult
 
         private CancellationToken BuildCancellationToken()
         {
-            if (_linkedTokens is not null)
+            switch(_cancellationTokens.Count)
             {
-                var tokens = new List<CancellationToken>(_linkedTokens.Count + 1);
-                tokens.Add(_cancellationToken);
-                tokens.AddRange(_linkedTokens);
-                _linkedTokensSource = CancellationTokenSource.CreateLinkedTokenSource(tokens.ToArray());
-                return _linkedTokensSource.Token;
-            }
-            return _cancellationToken;
+                case 0: return CancellationToken.None;
+                case 1: return _cancellationTokens[0]();
+                default:
+                    CancellationToken[] tokens = new CancellationToken[_cancellationTokens.Count];
+                    for(int i = 0; i < _cancellationTokens.Count; i++)
+                    {
+                        tokens[i] = _cancellationTokens[i]();
+                    }
+                    _linkedTokensSource = CancellationTokenSource.CreateLinkedTokenSource(tokens);
+                    return _linkedTokensSource.Token;
+            };
         }
 
 
@@ -108,13 +111,13 @@ public void Dispose()
             (Console as IDisposable)?.Dispose();
         }
 
-        public void AddLinkedCancellationToken(CancellationToken token)
+        public void AddLinkedCancellationToken(Func<CancellationToken> token)
         {
             if (_lazyCancellationToken.IsValueCreated)
             {
                 throw new InvalidOperationException($"Cannot add additional linked cancellation tokens once {nameof(InvocationContext)}.{nameof(CancellationToken)} has been invoked");
             }
-            (_linkedTokens ??= new()).Add(token);
+            _cancellationTokens.Add(token);
         }
     }
 }