From 8735a594ad59a8439c31e0858b68e90896e658a9 Mon Sep 17 00:00:00 2001 From: Kevin Bost Date: Thu, 30 Jun 2022 21:22:59 -0700 Subject: [PATCH 1/5] Adding in cancellation support for InvokeAsync Moving the storage of the cancellation token to the InvocationContext Adding InvocationContext tests Updating the approved public API Fixes after rebase Reworking the code so it doesn't register for events without a cancellation token --- .../CancelOnProcessTerminationTests.cs | 1 - .../Invocation/InvocationContextTests.cs | 80 +++++++++++++++++++ .../Invocation/InvocationExtensionsTests.cs | 25 ++++++ .../Invocation/InvocationPipelineTests.cs | 37 ++++++++- .../ParserTests.RootCommandAndArg0.cs | 11 +-- .../Binding/BindingContext.cs | 2 - .../Builder/CommandLineBuilderExtensions.cs | 31 +++---- src/System.CommandLine/CommandExtensions.cs | 13 ++- .../Invocation/InvocationContext.cs | 65 +++++++++------ .../Invocation/InvocationPipeline.cs | 11 +-- .../Invocation/ServiceProvider.cs | 2 - .../Parsing/ParseResultExtensions.cs | 7 +- .../Parsing/ParserExtensions.cs | 11 ++- 13 files changed, 232 insertions(+), 64 deletions(-) create mode 100644 src/System.CommandLine.Tests/Invocation/InvocationContextTests.cs diff --git a/src/System.CommandLine.Tests/Invocation/CancelOnProcessTerminationTests.cs b/src/System.CommandLine.Tests/Invocation/CancelOnProcessTerminationTests.cs index 22f6c05798..9fe62e3c5e 100644 --- a/src/System.CommandLine.Tests/Invocation/CancelOnProcessTerminationTests.cs +++ b/src/System.CommandLine.Tests/Invocation/CancelOnProcessTerminationTests.cs @@ -6,7 +6,6 @@ using System.Diagnostics; using System.Runtime.InteropServices; using System.Threading.Tasks; -using FluentAssertions; using Xunit; using Process = System.Diagnostics.Process; diff --git a/src/System.CommandLine.Tests/Invocation/InvocationContextTests.cs b/src/System.CommandLine.Tests/Invocation/InvocationContextTests.cs new file mode 100644 index 0000000000..8d5bffaeb7 --- /dev/null +++ b/src/System.CommandLine.Tests/Invocation/InvocationContextTests.cs @@ -0,0 +1,80 @@ +using FluentAssertions; +using System.CommandLine; +using System.CommandLine.Builder; +using System.CommandLine.Invocation; +using System.CommandLine.Parsing; +using System.Threading; +using Xunit; + +namespace System.CommandLine.Tests.Invocation +{ + public class InvocationContextTests + { + [Fact] + public void InvocationContext_with_cancellation_token_returns_it() + { + using CancellationTokenSource cts = new(); + var parseResult = new CommandLineBuilder(new RootCommand()) + .Build() + .Parse(""); + using InvocationContext context = new(parseResult, cancellationToken:cts.Token); + + var token = context.GetCancellationToken(); + + token.IsCancellationRequested.Should().BeFalse(); + cts.Cancel(); + token.IsCancellationRequested.Should().BeTrue(); + } + + [Fact] + public void InvocationContext_with_linked_cancellation_token_can_cancel_by_passed_token() + { + using CancellationTokenSource cts1 = new(); + using CancellationTokenSource cts2 = new(); + var parseResult = new CommandLineBuilder(new RootCommand()) + .Build() + .Parse(""); + using InvocationContext context = new(parseResult, cancellationToken: cts1.Token); + context.AddLinkedCancellationToken(() => cts2.Token); + + var token = context.GetCancellationToken(); + + token.IsCancellationRequested.Should().BeFalse(); + cts1.Cancel(); + token.IsCancellationRequested.Should().BeTrue(); + } + + [Fact] + public void InvocationContext_with_linked_cancellation_token_can_cancel_by_linked_token() + { + using CancellationTokenSource cts1 = new(); + using CancellationTokenSource cts2 = new(); + var parseResult = new CommandLineBuilder(new RootCommand()) + .Build() + .Parse(""); + using InvocationContext context = new(parseResult, cancellationToken: cts1.Token); + context.AddLinkedCancellationToken(() => cts2.Token); + + var token = context.GetCancellationToken(); + + token.IsCancellationRequested.Should().BeFalse(); + cts2.Cancel(); + token.IsCancellationRequested.Should().BeTrue(); + } + + [Fact] + public void InvocationContext_adding_additional_linked_token_after_token_has_been_built_throws() + { + using CancellationTokenSource cts = new(); + + var parseResult = new CommandLineBuilder(new RootCommand()) + .Build() + .Parse(""); + using InvocationContext context = new(parseResult, cancellationToken: cts.Token); + context.AddLinkedCancellationToken(() => cts.Token); + _ = context.GetCancellationToken(); + + Assert.Throws(() => context.AddLinkedCancellationToken(() => cts.Token)); + } + } +} diff --git a/src/System.CommandLine.Tests/Invocation/InvocationExtensionsTests.cs b/src/System.CommandLine.Tests/Invocation/InvocationExtensionsTests.cs index 947f94a7c3..4f6d2e7361 100644 --- a/src/System.CommandLine.Tests/Invocation/InvocationExtensionsTests.cs +++ b/src/System.CommandLine.Tests/Invocation/InvocationExtensionsTests.cs @@ -1,7 +1,9 @@ // Copyright (c) .NET Foundation and contributors. All rights reserved. // Licensed under the MIT license. See LICENSE file in the project root for full license information. +using System.CommandLine.Invocation; using System.CommandLine.IO; +using System.Threading; using System.Threading.Tasks; using FluentAssertions; using Xunit; @@ -149,5 +151,28 @@ public void RootCommand_Invoke_can_set_custom_result_code() resultCode.Should().Be(123); } + + [Fact] + public async Task Command_InvokeAsync_with_cancelation_token_invokes_command_handler() + { + CancellationTokenSource cts = new(); + var command = new Command("test"); + command.SetHandler((InvocationContext context) => + { + CancellationToken cancellationToken = context.GetCancellationToken(); + Assert.True(cancellationToken.CanBeCanceled); + if (cancellationToken.IsCancellationRequested) + { + return Task.FromResult(42); + } + + return Task.FromResult(0); + }); + + cts.Cancel(); + int rv = await command.InvokeAsync("test", cancellationToken: cts.Token); + + rv.Should().Be(42); + } } } diff --git a/src/System.CommandLine.Tests/Invocation/InvocationPipelineTests.cs b/src/System.CommandLine.Tests/Invocation/InvocationPipelineTests.cs index a707ca879b..ec71dc090b 100644 --- a/src/System.CommandLine.Tests/Invocation/InvocationPipelineTests.cs +++ b/src/System.CommandLine.Tests/Invocation/InvocationPipelineTests.cs @@ -5,6 +5,7 @@ using System.CommandLine.IO; using System.CommandLine.Parsing; using System.Linq; +using System.Threading; using System.Threading.Tasks; using FluentAssertions; using Xunit; @@ -44,7 +45,7 @@ public async Task InvokeAsync_chooses_the_appropriate_command() var parser = new CommandLineBuilder(new RootCommand { - first, + first, second }) .Build(); @@ -327,5 +328,39 @@ public async Task When_help_builder_factory_is_specified_it_is_used_to_create_th handlerWasCalled.Should().BeTrue(); factoryWasCalled.Should().BeTrue(); } + + [Fact] + public async Task Command_InvokeAsync_can_cancel_from_middleware() + { + var handlerWasCalled = false; + var isCancelRequested = false; + + var command = new Command("the-command"); + command.SetHandler((InvocationContext context) => + { + handlerWasCalled = true; + isCancelRequested = context.GetCancellationToken().IsCancellationRequested; + return Task.FromResult(0); + }); + + + using CancellationTokenSource cts = new(); + var parser = new CommandLineBuilder(new RootCommand + { + command + }) + .AddMiddleware(async (context, next) => + { + context.AddLinkedCancellationToken(() => cts.Token); + cts.Cancel(); + await next(context); + }) + .Build(); + + await parser.InvokeAsync("the-command"); + + handlerWasCalled.Should().BeTrue(); + isCancelRequested.Should().BeTrue(); + } } } diff --git a/src/System.CommandLine.Tests/ParserTests.RootCommandAndArg0.cs b/src/System.CommandLine.Tests/ParserTests.RootCommandAndArg0.cs index 993b8c2d10..1e2819c22c 100644 --- a/src/System.CommandLine.Tests/ParserTests.RootCommandAndArg0.cs +++ b/src/System.CommandLine.Tests/ParserTests.RootCommandAndArg0.cs @@ -42,10 +42,11 @@ public void When_parsing_a_string_array_input_then_a_full_path_to_an_executable_ command.Parse(Split("inner -x hello")).Errors.Should().BeEmpty(); - command.Parse(Split($"{RootCommand.ExecutablePath} inner -x hello")) - .Errors - .Should() - .ContainSingle(e => e.Message == $"{LocalizationResources.Instance.UnrecognizedCommandOrArgument(RootCommand.ExecutablePath)}"); + var parserResult = command.Parse(Split($"\"{RootCommand.ExecutablePath}\" inner -x hello")); + parserResult + .Errors + .Should() + .ContainSingle(e => e.Message == LocalizationResources.Instance.UnrecognizedCommandOrArgument(RootCommand.ExecutablePath)); } [Fact] @@ -76,7 +77,7 @@ public void When_parsing_an_unsplit_string_then_input_a_full_path_to_an_executab } }; - var result2 = command.Parse($"{RootCommand.ExecutablePath} inner -x hello"); + var result2 = command.Parse($"\"{RootCommand.ExecutablePath}\" inner -x hello"); result2.RootCommandResult.Token.Value.Should().Be(RootCommand.ExecutablePath); } diff --git a/src/System.CommandLine/Binding/BindingContext.cs b/src/System.CommandLine/Binding/BindingContext.cs index 79cbe88a5a..13d3bd6e4b 100644 --- a/src/System.CommandLine/Binding/BindingContext.cs +++ b/src/System.CommandLine/Binding/BindingContext.cs @@ -6,8 +6,6 @@ using System.Diagnostics.CodeAnalysis; using System.Linq; -#nullable enable - namespace System.CommandLine.Binding { /// diff --git a/src/System.CommandLine/Builder/CommandLineBuilderExtensions.cs b/src/System.CommandLine/Builder/CommandLineBuilderExtensions.cs index c20f08c841..e06e1496e6 100644 --- a/src/System.CommandLine/Builder/CommandLineBuilderExtensions.cs +++ b/src/System.CommandLine/Builder/CommandLineBuilderExtensions.cs @@ -65,17 +65,15 @@ public static CommandLineBuilder CancelOnProcessTermination( builder.AddMiddleware(async (context, next) => { - bool cancellationHandlingAdded = false; - ManualResetEventSlim? blockProcessExit = null; ConsoleCancelEventHandler? consoleHandler = null; EventHandler? processExitHandler = null; + ManualResetEventSlim? blockProcessExit = null; - context.CancellationHandlingAdded += (CancellationTokenSource cts) => + context.AddLinkedCancellationToken(() => { + //TODO: This CancellationTokenSource is never disposed... + CancellationTokenSource cts = new(); blockProcessExit = new ManualResetEventSlim(initialState: false); - cancellationHandlingAdded = 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 = (_, _) => { // Cancel asynchronously not to block the handler (as then the process might possibly run longer then what was the requested timeout) @@ -100,8 +98,11 @@ public static CommandLineBuilder CancelOnProcessTermination( } ExitCode = context.ExitCode; }; + // Default limit for ProcesExit handler is 2 seconds + // https://docs.microsoft.com/en-us/dotnet/api/system.appdomain.processexit?view=net-6.0 consoleHandler = (_, args) => { + cts.Cancel(); // Stop the process from terminating. // Since the context was cancelled, the invocation should // finish and Main will return. @@ -128,7 +129,10 @@ public static CommandLineBuilder CancelOnProcessTermination( }; Console.CancelKeyPress += consoleHandler; AppDomain.CurrentDomain.ProcessExit += processExitHandler; - }; + + return cts.Token; + }); + try { @@ -136,12 +140,9 @@ public static CommandLineBuilder CancelOnProcessTermination( } finally { - if (cancellationHandlingAdded) - { - Console.CancelKeyPress -= consoleHandler; - AppDomain.CurrentDomain.ProcessExit -= processExitHandler; - blockProcessExit!.Set(); - } + Console.CancelKeyPress -= consoleHandler; + AppDomain.CurrentDomain.ProcessExit -= processExitHandler; + blockProcessExit?.Set(); } }, MiddlewareOrderInternal.Startup); @@ -418,7 +419,7 @@ public static CommandLineBuilder UseHelp( int? maxWidth = null) { builder.CustomizeHelpLayout(customize); - + if (builder.HelpOption is null) { builder.UseHelp(new HelpOption(() => builder.LocalizationResources), maxWidth); @@ -599,7 +600,7 @@ public static CommandLineBuilder UseSuggestDirective( /// The maximum Levenshtein distance for suggestions based on detected typos in command line input. /// The same instance of . public static CommandLineBuilder UseTypoCorrections( - this CommandLineBuilder builder, + this CommandLineBuilder builder, int maxLevenshteinDistance = 3) { builder.AddMiddleware(async (context, next) => diff --git a/src/System.CommandLine/CommandExtensions.cs b/src/System.CommandLine/CommandExtensions.cs index 1a50cd0873..fb96cc19f8 100644 --- a/src/System.CommandLine/CommandExtensions.cs +++ b/src/System.CommandLine/CommandExtensions.cs @@ -4,6 +4,7 @@ using System.CommandLine.Invocation; using System.CommandLine.Parsing; using System.Linq; +using System.Threading; using System.Threading.Tasks; namespace System.CommandLine @@ -48,13 +49,15 @@ public static int Invoke( /// The command to invoke. /// The arguments to parse. /// The console to which output is written during invocation. + /// A token that can be used to cancel the invocation. /// The exit code for the invocation. public static async Task InvokeAsync( this Command command, string[] args, - IConsole? console = null) + IConsole? console = null, + CancellationToken cancellationToken = default) { - return await GetDefaultInvocationPipeline(command, args).InvokeAsync(console); + return await GetDefaultInvocationPipeline(command, args).InvokeAsync(console, cancellationToken); } /// @@ -64,12 +67,14 @@ public static async Task InvokeAsync( /// The command to invoke. /// The command line to parse. /// The console to which output is written during invocation. + /// A token that can be used to cancel the invocation. /// The exit code for the invocation. public static Task InvokeAsync( this Command command, string commandLine, - IConsole? console = null) => - command.InvokeAsync(CommandLineStringSplitter.Instance.Split(commandLine).ToArray(), console); + IConsole? console = null, + CancellationToken cancellationToken = default) => + command.InvokeAsync(CommandLineStringSplitter.Instance.Split(commandLine).ToArray(), console, cancellationToken); private static InvocationPipeline GetDefaultInvocationPipeline(Command command, string[] args) { diff --git a/src/System.CommandLine/Invocation/InvocationContext.cs b/src/System.CommandLine/Invocation/InvocationContext.cs index 5a14abef22..2fc5edeb0e 100644 --- a/src/System.CommandLine/Invocation/InvocationContext.cs +++ b/src/System.CommandLine/Invocation/InvocationContext.cs @@ -1,6 +1,7 @@ // Copyright (c) .NET Foundation and contributors. All rights reserved. // Licensed under the MIT license. See LICENSE file in the project root for full license information. +using System.Collections.Generic; using System.CommandLine.Binding; using System.CommandLine.Help; using System.CommandLine.IO; @@ -12,22 +13,27 @@ namespace System.CommandLine.Invocation /// /// Supports command invocation by providing access to parse results and other services. /// - public sealed class InvocationContext + public sealed class InvocationContext : IDisposable { - private CancellationTokenSource? _cts; - private Action? _cancellationHandlingAddedEvent; private HelpBuilder? _helpBuilder; private BindingContext? _bindingContext; private IConsole? _console; + private CancellationTokenSource? _linkedTokensSource; + private List> _cancellationTokens = new(3); + private Lazy _lazyCancellationToken; /// The result of the current parse operation. /// The console to which output is to be written. + /// A cancellation token that can be used to cancel and invocation. public InvocationContext( ParseResult parseResult, - IConsole? console = null) + IConsole? console = null, + CancellationToken cancellationToken = default) { ParseResult = parseResult; _console = console; + _cancellationTokens.Add(() => cancellationToken); + _lazyCancellationToken = new Lazy(BuildCancellationToken); } /// @@ -40,6 +46,8 @@ public BindingContext BindingContext if (_bindingContext is null) { _bindingContext = new BindingContext(this); + _bindingContext.ServiceProvider.AddService(_ => GetCancellationToken()); + _bindingContext.ServiceProvider.AddService(_ => this); } return _bindingContext; @@ -94,33 +102,44 @@ public IConsole Console /// As the is passed through the invocation pipeline to the associated with the invoked command, only the last value of this property will be the one applied. public IInvocationResult? InvocationResult { get; set; } - internal event Action CancellationHandlingAdded + /// + /// Gets a cancellation token that can be used to check if cancellation has been requested. + /// + public CancellationToken GetCancellationToken() => _lazyCancellationToken.Value; + + private CancellationToken BuildCancellationToken() { - add + switch(_cancellationTokens.Count) { - if (_cts is not null) - { - throw new InvalidOperationException("Handlers must be added before adding cancellation handling."); - } + 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; + }; + } - _cancellationHandlingAddedEvent += value; - } - remove => _cancellationHandlingAddedEvent -= value; + + /// + public void Dispose() + { + _linkedTokensSource?.Dispose(); + _linkedTokensSource = null; + (Console as IDisposable)?.Dispose(); } - /// - /// Gets token to implement cancellation handling. - /// - /// Token used by the caller to implement cancellation handling. - public CancellationToken GetCancellationToken() + public void AddLinkedCancellationToken(Func token) { - if (_cts is null) + if (_lazyCancellationToken.IsValueCreated) { - _cts = new CancellationTokenSource(); - _cancellationHandlingAddedEvent?.Invoke(_cts); + throw new InvalidOperationException($"Cannot add additional linked cancellation tokens once {nameof(InvocationContext)}.{nameof(CancellationToken)} has been invoked"); } - - return _cts.Token; + _cancellationTokens.Add(token); } } } diff --git a/src/System.CommandLine/Invocation/InvocationPipeline.cs b/src/System.CommandLine/Invocation/InvocationPipeline.cs index b428beba18..ba4c1d60e0 100644 --- a/src/System.CommandLine/Invocation/InvocationPipeline.cs +++ b/src/System.CommandLine/Invocation/InvocationPipeline.cs @@ -3,22 +3,23 @@ using System.Collections.Generic; using System.Linq; +using System.Threading; using System.Threading.Tasks; namespace System.CommandLine.Invocation { internal class InvocationPipeline { - private readonly ParseResult parseResult; + private readonly ParseResult _parseResult; public InvocationPipeline(ParseResult parseResult) { - this.parseResult = parseResult ?? throw new ArgumentNullException(nameof(parseResult)); + _parseResult = parseResult ?? throw new ArgumentNullException(nameof(parseResult)); } - public Task InvokeAsync(IConsole? console = null) + public Task InvokeAsync(IConsole? console = null, CancellationToken cancellationToken = default) { - var context = new InvocationContext(parseResult, console); + var context = new InvocationContext(_parseResult, console, cancellationToken); if (context.Parser.Configuration.Middleware.Count == 0 && context.ParseResult.CommandResult.Command.Handler is ICommandHandler handler) @@ -40,7 +41,7 @@ static async Task FullInvocationChainAsync(InvocationContext context) public int Invoke(IConsole? console = null) { - var context = new InvocationContext(parseResult, console); + var context = new InvocationContext(_parseResult, console); if (context.Parser.Configuration.Middleware.Count == 0 && context.ParseResult.CommandResult.Command.Handler is ICommandHandler handler) diff --git a/src/System.CommandLine/Invocation/ServiceProvider.cs b/src/System.CommandLine/Invocation/ServiceProvider.cs index acf3091947..852a1c9913 100644 --- a/src/System.CommandLine/Invocation/ServiceProvider.cs +++ b/src/System.CommandLine/Invocation/ServiceProvider.cs @@ -6,8 +6,6 @@ using System.CommandLine.Help; using System.Threading; -#nullable enable - namespace System.CommandLine.Invocation { internal class ServiceProvider : IServiceProvider diff --git a/src/System.CommandLine/Parsing/ParseResultExtensions.cs b/src/System.CommandLine/Parsing/ParseResultExtensions.cs index 0f233fe4f7..7cd3ebb8d6 100644 --- a/src/System.CommandLine/Parsing/ParseResultExtensions.cs +++ b/src/System.CommandLine/Parsing/ParseResultExtensions.cs @@ -6,6 +6,7 @@ using System.CommandLine.Invocation; using System.Linq; using System.Text; +using System.Threading; using System.Threading.Tasks; namespace System.CommandLine.Parsing @@ -20,11 +21,13 @@ public static class ParseResultExtensions /// /// A parse result on which the invocation is based. /// A console to which output can be written. By default, is used. + /// A token that can be used to cancel an invocation. /// A task whose result can be used as a process exit code. public static async Task InvokeAsync( this ParseResult parseResult, - IConsole? console = null) => - await new InvocationPipeline(parseResult).InvokeAsync(console); + IConsole? console = null, + CancellationToken cancellationToken = default) => + await new InvocationPipeline(parseResult).InvokeAsync(console, cancellationToken); /// /// Invokes the appropriate command handler for a parsed command line input. diff --git a/src/System.CommandLine/Parsing/ParserExtensions.cs b/src/System.CommandLine/Parsing/ParserExtensions.cs index 0f51f66ec7..4c90a1783b 100644 --- a/src/System.CommandLine/Parsing/ParserExtensions.cs +++ b/src/System.CommandLine/Parsing/ParserExtensions.cs @@ -2,6 +2,7 @@ // Licensed under the MIT license. See LICENSE file in the project root for full license information. using System.Linq; +using System.Threading; using System.Threading.Tasks; namespace System.CommandLine.Parsing @@ -40,8 +41,9 @@ public static int Invoke( public static Task InvokeAsync( this Parser parser, string commandLine, - IConsole? console = null) => - parser.InvokeAsync(CommandLineStringSplitter.Instance.Split(commandLine).ToArray(), console); + IConsole? console = null, + CancellationToken cancellationToken = default) => + parser.InvokeAsync(CommandLineStringSplitter.Instance.Split(commandLine).ToArray(), console, cancellationToken); /// /// Parses a command line string array and invokes the handler for the indicated command. @@ -50,8 +52,9 @@ public static Task InvokeAsync( public static async Task InvokeAsync( this Parser parser, string[] args, - IConsole? console = null) => - await parser.Parse(args).InvokeAsync(console); + IConsole? console = null, + CancellationToken cancellationToken = default) => + await parser.Parse(args).InvokeAsync(console, cancellationToken); /// /// Parses a command line string. From 56ec9b38d0101024c18fb9ce3becd4784dfdafa1 Mon Sep 17 00:00:00 2001 From: Kevin Bost Date: Thu, 30 Jun 2022 22:32:49 -0700 Subject: [PATCH 2/5] Handling TODO, to dispose of CancellationTokenSource --- .../Builder/CommandLineBuilderExtensions.cs | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/src/System.CommandLine/Builder/CommandLineBuilderExtensions.cs b/src/System.CommandLine/Builder/CommandLineBuilderExtensions.cs index e06e1496e6..7fd815c896 100644 --- a/src/System.CommandLine/Builder/CommandLineBuilderExtensions.cs +++ b/src/System.CommandLine/Builder/CommandLineBuilderExtensions.cs @@ -68,11 +68,11 @@ public static CommandLineBuilder CancelOnProcessTermination( ConsoleCancelEventHandler? consoleHandler = null; EventHandler? processExitHandler = null; ManualResetEventSlim? blockProcessExit = null; + CancellationTokenSource? cts = null; context.AddLinkedCancellationToken(() => { - //TODO: This CancellationTokenSource is never disposed... - CancellationTokenSource cts = new(); + cts = new CancellationTokenSource(); blockProcessExit = new ManualResetEventSlim(initialState: false); processExitHandler = (_, _) => { @@ -102,7 +102,6 @@ public static CommandLineBuilder CancelOnProcessTermination( // https://docs.microsoft.com/en-us/dotnet/api/system.appdomain.processexit?view=net-6.0 consoleHandler = (_, args) => { - cts.Cancel(); // Stop the process from terminating. // Since the context was cancelled, the invocation should // finish and Main will return. @@ -125,7 +124,7 @@ public static CommandLineBuilder CancelOnProcessTermination( // 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(); + cts?.Cancel(); }; Console.CancelKeyPress += consoleHandler; AppDomain.CurrentDomain.ProcessExit += processExitHandler; @@ -133,7 +132,6 @@ public static CommandLineBuilder CancelOnProcessTermination( return cts.Token; }); - try { await next(context); @@ -142,6 +140,7 @@ public static CommandLineBuilder CancelOnProcessTermination( { Console.CancelKeyPress -= consoleHandler; AppDomain.CurrentDomain.ProcessExit -= processExitHandler; + Interlocked.Exchange(ref cts, null)?.Dispose(); blockProcessExit?.Set(); } }, MiddlewareOrderInternal.Startup); From 9b6224a30f12425c12953f519f2d1e5230950a32 Mon Sep 17 00:00:00 2001 From: Kevin Bost Date: Thu, 7 Jul 2022 23:56:18 -0700 Subject: [PATCH 3/5] Update with suggestions to better handled linked token --- ...ommandLine_api_is_not_changed.approved.txt | 1 + .../Invocation/InvocationContextTests.cs | 21 +--- .../Invocation/InvocationPipelineTests.cs | 2 +- .../Builder/CommandLineBuilderExtensions.cs | 118 ++++++++---------- .../Invocation/InvocationContext.cs | 53 ++++---- 5 files changed, 82 insertions(+), 113 deletions(-) diff --git a/src/System.CommandLine.ApiCompatibility.Tests/ApiCompatibilityApprovalTests.System_CommandLine_api_is_not_changed.approved.txt b/src/System.CommandLine.ApiCompatibility.Tests/ApiCompatibilityApprovalTests.System_CommandLine_api_is_not_changed.approved.txt index 6b42819629..d8ed392d97 100644 --- a/src/System.CommandLine.ApiCompatibility.Tests/ApiCompatibilityApprovalTests.System_CommandLine_api_is_not_changed.approved.txt +++ b/src/System.CommandLine.ApiCompatibility.Tests/ApiCompatibilityApprovalTests.System_CommandLine_api_is_not_changed.approved.txt @@ -368,6 +368,7 @@ System.CommandLine.Invocation public System.CommandLine.Parsing.Parser Parser { get; } public System.CommandLine.ParseResult ParseResult { get; set; } public System.Threading.CancellationToken GetCancellationToken() + public System.Void LinkToken(System.Threading.CancellationToken token) public delegate InvocationMiddleware : System.MulticastDelegate, System.ICloneable, System.Runtime.Serialization.ISerializable .ctor(System.Object object, System.IntPtr method) public System.IAsyncResult BeginInvoke(InvocationContext context, System.Func next, System.AsyncCallback callback, System.Object object) diff --git a/src/System.CommandLine.Tests/Invocation/InvocationContextTests.cs b/src/System.CommandLine.Tests/Invocation/InvocationContextTests.cs index 8d5bffaeb7..c62027f46d 100644 --- a/src/System.CommandLine.Tests/Invocation/InvocationContextTests.cs +++ b/src/System.CommandLine.Tests/Invocation/InvocationContextTests.cs @@ -17,7 +17,7 @@ public void InvocationContext_with_cancellation_token_returns_it() var parseResult = new CommandLineBuilder(new RootCommand()) .Build() .Parse(""); - using InvocationContext context = new(parseResult, cancellationToken:cts.Token); + using InvocationContext context = new(parseResult, cancellationToken: cts.Token); var token = context.GetCancellationToken(); @@ -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.LinkToken(cts2.Token); var token = context.GetCancellationToken(); @@ -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.LinkToken(cts2.Token); var token = context.GetCancellationToken(); @@ -61,20 +61,5 @@ public void InvocationContext_with_linked_cancellation_token_can_cancel_by_linke cts2.Cancel(); token.IsCancellationRequested.Should().BeTrue(); } - - [Fact] - public void InvocationContext_adding_additional_linked_token_after_token_has_been_built_throws() - { - using CancellationTokenSource cts = new(); - - var parseResult = new CommandLineBuilder(new RootCommand()) - .Build() - .Parse(""); - using InvocationContext context = new(parseResult, cancellationToken: cts.Token); - context.AddLinkedCancellationToken(() => cts.Token); - _ = context.GetCancellationToken(); - - Assert.Throws(() => context.AddLinkedCancellationToken(() => cts.Token)); - } } } diff --git a/src/System.CommandLine.Tests/Invocation/InvocationPipelineTests.cs b/src/System.CommandLine.Tests/Invocation/InvocationPipelineTests.cs index ec71dc090b..eeeebeed97 100644 --- a/src/System.CommandLine.Tests/Invocation/InvocationPipelineTests.cs +++ b/src/System.CommandLine.Tests/Invocation/InvocationPipelineTests.cs @@ -351,7 +351,7 @@ public async Task Command_InvokeAsync_can_cancel_from_middleware() }) .AddMiddleware(async (context, next) => { - context.AddLinkedCancellationToken(() => cts.Token); + context.LinkToken(cts.Token); cts.Cancel(); await next(context); }) diff --git a/src/System.CommandLine/Builder/CommandLineBuilderExtensions.cs b/src/System.CommandLine/Builder/CommandLineBuilderExtensions.cs index 7fd815c896..38868cc915 100644 --- a/src/System.CommandLine/Builder/CommandLineBuilderExtensions.cs +++ b/src/System.CommandLine/Builder/CommandLineBuilderExtensions.cs @@ -52,7 +52,7 @@ public static class CommandLineBuilderExtensions /// /// The same instance of . public static CommandLineBuilder CancelOnProcessTermination( - this CommandLineBuilder builder, + this CommandLineBuilder builder, TimeSpan? timeout = null) { // https://tldp.org/LDP/abs/html/exitcodes.html - 130 - script terminated by ctrl-c @@ -67,70 +67,63 @@ public static CommandLineBuilder CancelOnProcessTermination( { ConsoleCancelEventHandler? consoleHandler = null; EventHandler? processExitHandler = null; - ManualResetEventSlim? blockProcessExit = null; - CancellationTokenSource? cts = null; + ManualResetEventSlim blockProcessExit = new(initialState: false); - context.AddLinkedCancellationToken(() => + processExitHandler = (_, _) => { - cts = new CancellationTokenSource(); - blockProcessExit = new ManualResetEventSlim(initialState: false); - processExitHandler = (_, _) => + // 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(context.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. + if (!blockProcessExit.Wait(timeout > TimeSpan.Zero + ? timeout.Value + : Timeout.InfiniteTimeSpan)) { - // 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. - 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; - }; - // Default limit for ProcesExit handler is 2 seconds - // https://docs.microsoft.com/en-us/dotnet/api/system.appdomain.processexit?view=net-6.0 - consoleHandler = (_, args) => + 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) { - // 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); - } + // The async cancellation didn't finish in timely manner + context.ExitCode = SIGINT_EXIT_CODE; + } + ExitCode = context.ExitCode; + }; + // Default limit for ProcesExit handler is 2 seconds + // https://docs.microsoft.com/en-us/dotnet/api/system.appdomain.processexit?view=net-6.0 + 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; + // 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 + context.Cancel(); + }; - return cts.Token; - }); + Console.CancelKeyPress += consoleHandler; + AppDomain.CurrentDomain.ProcessExit += processExitHandler; try { @@ -140,14 +133,13 @@ public static CommandLineBuilder CancelOnProcessTermination( { Console.CancelKeyPress -= consoleHandler; AppDomain.CurrentDomain.ProcessExit -= processExitHandler; - Interlocked.Exchange(ref cts, null)?.Dispose(); blockProcessExit?.Set(); } }, MiddlewareOrderInternal.Startup); return builder; } - + /// /// Enables the parser to recognize command line directives. /// @@ -206,7 +198,7 @@ public static CommandLineBuilder EnablePosixBundling( builder.EnablePosixBundling = value; return builder; } - + /// /// Ensures that the application is registered with the dotnet-suggest tool to enable command line completions. /// @@ -485,7 +477,7 @@ public static CommandLineBuilder AddMiddleware( return builder; } - + /// /// Adds a middleware delegate to the invocation pipeline called before a command handler is invoked. /// diff --git a/src/System.CommandLine/Invocation/InvocationContext.cs b/src/System.CommandLine/Invocation/InvocationContext.cs index 2fc5edeb0e..07e6e024f3 100644 --- a/src/System.CommandLine/Invocation/InvocationContext.cs +++ b/src/System.CommandLine/Invocation/InvocationContext.cs @@ -1,4 +1,4 @@ -// Copyright (c) .NET Foundation and contributors. All rights reserved. +// Copyright (c) .NET Foundation and contributors. All rights reserved. // Licensed under the MIT license. See LICENSE file in the project root for full license information. using System.Collections.Generic; @@ -18,9 +18,9 @@ public sealed class InvocationContext : IDisposable private HelpBuilder? _helpBuilder; private BindingContext? _bindingContext; private IConsole? _console; - private CancellationTokenSource? _linkedTokensSource; - private List> _cancellationTokens = new(3); - private Lazy _lazyCancellationToken; + private readonly CancellationToken _token; + private readonly LinkedList _registrations = new(); + private volatile CancellationTokenSource? _source; /// The result of the current parse operation. /// The console to which output is to be written. @@ -32,8 +32,13 @@ public InvocationContext( { ParseResult = parseResult; _console = console; - _cancellationTokens.Add(() => cancellationToken); - _lazyCancellationToken = new Lazy(BuildCancellationToken); + + _source = new CancellationTokenSource(); + _token = _source.Token; + if (cancellationToken.CanBeCanceled) + { + LinkToken(cancellationToken); + } } /// @@ -105,41 +110,27 @@ public IConsole Console /// /// Gets a cancellation token that can be used to check if cancellation has been requested. /// - public CancellationToken GetCancellationToken() => _lazyCancellationToken.Value; + public CancellationToken GetCancellationToken() => _token; - private CancellationToken BuildCancellationToken() + internal void Cancel() { - switch(_cancellationTokens.Count) - { - 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; - }; + using var source = Interlocked.Exchange(ref _source, null); + source?.Cancel(); } - - /// - public void Dispose() + public void LinkToken(CancellationToken token) { - _linkedTokensSource?.Dispose(); - _linkedTokensSource = null; - (Console as IDisposable)?.Dispose(); + _registrations.AddLast(token.Register(Cancel)); } - public void AddLinkedCancellationToken(Func token) + /// + void IDisposable.Dispose() { - if (_lazyCancellationToken.IsValueCreated) + Interlocked.Exchange(ref _source, null)?.Dispose(); + foreach (CancellationTokenRegistration registration in _registrations) { - throw new InvalidOperationException($"Cannot add additional linked cancellation tokens once {nameof(InvocationContext)}.{nameof(CancellationToken)} has been invoked"); + registration.Dispose(); } - _cancellationTokens.Add(token); } } } From 149f6b0f9b613752f9d2534afe58c2b3c14cb811 Mon Sep 17 00:00:00 2001 From: Kevin Bost Date: Fri, 8 Jul 2022 00:03:18 -0700 Subject: [PATCH 4/5] Rebased and fixed conflicts --- ...tem_CommandLine_api_is_not_changed.approved.txt | 14 +++++++------- .../Invocation/CancelOnProcessTerminationTests.cs | 2 ++ .../Invocation/InvocationContextTests.cs | 1 - 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/src/System.CommandLine.ApiCompatibility.Tests/ApiCompatibilityApprovalTests.System_CommandLine_api_is_not_changed.approved.txt b/src/System.CommandLine.ApiCompatibility.Tests/ApiCompatibilityApprovalTests.System_CommandLine_api_is_not_changed.approved.txt index d8ed392d97..01429ca129 100644 --- a/src/System.CommandLine.ApiCompatibility.Tests/ApiCompatibilityApprovalTests.System_CommandLine_api_is_not_changed.approved.txt +++ b/src/System.CommandLine.ApiCompatibility.Tests/ApiCompatibilityApprovalTests.System_CommandLine_api_is_not_changed.approved.txt @@ -66,8 +66,8 @@ System.CommandLine public static class CommandExtensions public static System.Int32 Invoke(this Command command, System.String[] args, IConsole console = null) public static System.Int32 Invoke(this Command command, System.String commandLine, IConsole console = null) - public static System.Threading.Tasks.Task InvokeAsync(this Command command, System.String[] args, IConsole console = null) - public static System.Threading.Tasks.Task InvokeAsync(this Command command, System.String commandLine, IConsole console = null) + public static System.Threading.Tasks.Task InvokeAsync(this Command command, System.String[] args, IConsole console = null, System.Threading.CancellationToken cancellationToken = null) + public static System.Threading.Tasks.Task InvokeAsync(this Command command, System.String commandLine, IConsole console = null, System.Threading.CancellationToken cancellationToken = null) public static ParseResult Parse(this Command command, System.String[] args) public static ParseResult Parse(this Command command, System.String commandLine) public class CommandLineBuilder @@ -357,8 +357,8 @@ System.CommandLine.Help System.CommandLine.Invocation public interface IInvocationResult public System.Void Apply(InvocationContext context) - public class InvocationContext - .ctor(System.CommandLine.ParseResult parseResult, System.CommandLine.IConsole console = null) + public class InvocationContext, System.IDisposable + .ctor(System.CommandLine.ParseResult parseResult, System.CommandLine.IConsole console = null, System.Threading.CancellationToken cancellationToken = null) public System.CommandLine.Binding.BindingContext BindingContext { get; } public System.CommandLine.IConsole Console { get; set; } public System.Int32 ExitCode { get; set; } @@ -451,12 +451,12 @@ System.CommandLine.Parsing public static System.String Diagram(this System.CommandLine.ParseResult parseResult) public static System.Boolean HasOption(this System.CommandLine.ParseResult parseResult, System.CommandLine.Option option) public static System.Int32 Invoke(this System.CommandLine.ParseResult parseResult, System.CommandLine.IConsole console = null) - public static System.Threading.Tasks.Task InvokeAsync(this System.CommandLine.ParseResult parseResult, System.CommandLine.IConsole console = null) + public static System.Threading.Tasks.Task InvokeAsync(this System.CommandLine.ParseResult parseResult, System.CommandLine.IConsole console = null, System.Threading.CancellationToken cancellationToken = null) public static class ParserExtensions public static System.Int32 Invoke(this Parser parser, System.String commandLine, System.CommandLine.IConsole console = null) public static System.Int32 Invoke(this Parser parser, System.String[] args, System.CommandLine.IConsole console = null) - public static System.Threading.Tasks.Task InvokeAsync(this Parser parser, System.String commandLine, System.CommandLine.IConsole console = null) - public static System.Threading.Tasks.Task InvokeAsync(this Parser parser, System.String[] args, System.CommandLine.IConsole console = null) + public static System.Threading.Tasks.Task InvokeAsync(this Parser parser, System.String commandLine, System.CommandLine.IConsole console = null, System.Threading.CancellationToken cancellationToken = null) + public static System.Threading.Tasks.Task InvokeAsync(this Parser parser, System.String[] args, System.CommandLine.IConsole console = null, System.Threading.CancellationToken cancellationToken = null) public static System.CommandLine.ParseResult Parse(this Parser parser, System.String commandLine) public abstract class SymbolResult public System.Collections.Generic.IReadOnlyList Children { get; } diff --git a/src/System.CommandLine.Tests/Invocation/CancelOnProcessTerminationTests.cs b/src/System.CommandLine.Tests/Invocation/CancelOnProcessTerminationTests.cs index 9fe62e3c5e..aaae8d2b2b 100644 --- a/src/System.CommandLine.Tests/Invocation/CancelOnProcessTerminationTests.cs +++ b/src/System.CommandLine.Tests/Invocation/CancelOnProcessTerminationTests.cs @@ -1,6 +1,8 @@ // Copyright (c) .NET Foundation and contributors. All rights reserved. // Licensed under the MIT license. See LICENSE file in the project root for full license information. +using FluentAssertions; +using System.CommandLine.Invocation; using System.CommandLine.Parsing; using System.CommandLine.Tests.Utility; using System.Diagnostics; diff --git a/src/System.CommandLine.Tests/Invocation/InvocationContextTests.cs b/src/System.CommandLine.Tests/Invocation/InvocationContextTests.cs index c62027f46d..5e168e7191 100644 --- a/src/System.CommandLine.Tests/Invocation/InvocationContextTests.cs +++ b/src/System.CommandLine.Tests/Invocation/InvocationContextTests.cs @@ -1,6 +1,5 @@ using FluentAssertions; using System.CommandLine; -using System.CommandLine.Builder; using System.CommandLine.Invocation; using System.CommandLine.Parsing; using System.Threading; From 6e0dba8be8aa032c00b4ae1e67da6d79452c2222 Mon Sep 17 00:00:00 2001 From: Kevin Bost Date: Fri, 8 Jul 2022 08:40:07 -0700 Subject: [PATCH 5/5] Missing using after rebase. --- .../Invocation/InvocationPipelineTests.cs | 1 + 1 file changed, 1 insertion(+) diff --git a/src/System.CommandLine.Tests/Invocation/InvocationPipelineTests.cs b/src/System.CommandLine.Tests/Invocation/InvocationPipelineTests.cs index eeeebeed97..1938f9548b 100644 --- a/src/System.CommandLine.Tests/Invocation/InvocationPipelineTests.cs +++ b/src/System.CommandLine.Tests/Invocation/InvocationPipelineTests.cs @@ -2,6 +2,7 @@ // Licensed under the MIT license. See LICENSE file in the project root for full license information. using System.CommandLine.Help; +using System.CommandLine.Invocation; using System.CommandLine.IO; using System.CommandLine.Parsing; using System.Linq;