Skip to content

Commit e31db44

Browse files
KebooKeboo
committed
Update with suggestions to better handled linked token
1 parent 45dd6eb commit e31db44

File tree

5 files changed

+82
-115
lines changed

5 files changed

+82
-115
lines changed

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -348,9 +348,8 @@ System.CommandLine.Invocation
348348
public System.CommandLine.LocalizationResources LocalizationResources { get; }
349349
public System.CommandLine.Parsing.Parser Parser { get; }
350350
public System.CommandLine.Parsing.ParseResult ParseResult { get; set; }
351-
public System.Void AddLinkedCancellationToken(System.Func<System.Threading.CancellationToken> token)
352-
public System.Void Dispose()
353351
public System.Threading.CancellationToken GetCancellationToken()
352+
public System.Void LinkToken(System.Threading.CancellationToken token)
354353
public delegate InvocationMiddleware : System.MulticastDelegate, System.ICloneable, System.Runtime.Serialization.ISerializable
355354
.ctor(System.Object object, System.IntPtr method)
356355
public System.IAsyncResult BeginInvoke(InvocationContext context, System.Func<InvocationContext,System.Threading.Tasks.Task> next, System.AsyncCallback callback, System.Object object)

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

Lines changed: 3 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ public void InvocationContext_with_cancellation_token_returns_it()
1717
var parseResult = new CommandLineBuilder(new RootCommand())
1818
.Build()
1919
.Parse("");
20-
using InvocationContext context = new(parseResult, cancellationToken:cts.Token);
20+
using InvocationContext context = new(parseResult, cancellationToken: cts.Token);
2121

2222
var token = context.GetCancellationToken();
2323

@@ -35,7 +35,7 @@ public void InvocationContext_with_linked_cancellation_token_can_cancel_by_passe
3535
.Build()
3636
.Parse("");
3737
using InvocationContext context = new(parseResult, cancellationToken: cts1.Token);
38-
context.AddLinkedCancellationToken(() => cts2.Token);
38+
context.LinkToken(cts2.Token);
3939

4040
var token = context.GetCancellationToken();
4141

@@ -53,28 +53,13 @@ public void InvocationContext_with_linked_cancellation_token_can_cancel_by_linke
5353
.Build()
5454
.Parse("");
5555
using InvocationContext context = new(parseResult, cancellationToken: cts1.Token);
56-
context.AddLinkedCancellationToken(() => cts2.Token);
56+
context.LinkToken(cts2.Token);
5757

5858
var token = context.GetCancellationToken();
5959

6060
token.IsCancellationRequested.Should().BeFalse();
6161
cts2.Cancel();
6262
token.IsCancellationRequested.Should().BeTrue();
6363
}
64-
65-
[Fact]
66-
public void InvocationContext_adding_additional_linked_token_after_token_has_been_built_throws()
67-
{
68-
using CancellationTokenSource cts = new();
69-
70-
var parseResult = new CommandLineBuilder(new RootCommand())
71-
.Build()
72-
.Parse("");
73-
using InvocationContext context = new(parseResult, cancellationToken: cts.Token);
74-
context.AddLinkedCancellationToken(() => cts.Token);
75-
_ = context.GetCancellationToken();
76-
77-
Assert.Throws<InvalidOperationException>(() => context.AddLinkedCancellationToken(() => cts.Token));
78-
}
7964
}
8065
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -354,7 +354,7 @@ public async Task Command_InvokeAsync_can_cancel_from_middleware()
354354
})
355355
.AddMiddleware(async (context, next) =>
356356
{
357-
context.AddLinkedCancellationToken(() => cts.Token);
357+
context.LinkToken(cts.Token);
358358
cts.Cancel();
359359
await next(context);
360360
})

‎src/System.CommandLine/Builder/CommandLineBuilderExtensions.cs‎

Lines changed: 55 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ public static class CommandLineBuilderExtensions
5151
/// </param>
5252
/// <returns>The same instance of <see cref="CommandLineBuilder"/>.</returns>
5353
public static CommandLineBuilder CancelOnProcessTermination(
54-
this CommandLineBuilder builder,
54+
this CommandLineBuilder builder,
5555
TimeSpan? timeout = null)
5656
{
5757
// https://tldp.org/LDP/abs/html/exitcodes.html - 130 - script terminated by ctrl-c
@@ -66,70 +66,63 @@ public static CommandLineBuilder CancelOnProcessTermination(
6666
{
6767
ConsoleCancelEventHandler? consoleHandler = null;
6868
EventHandler? processExitHandler = null;
69-
ManualResetEventSlim? blockProcessExit = null;
70-
CancellationTokenSource? cts = null;
69+
ManualResetEventSlim blockProcessExit = new(initialState: false);
7170

72-
context.AddLinkedCancellationToken(() =>
71+
processExitHandler = (_, _) =>
7372
{
74-
cts = new CancellationTokenSource();
75-
blockProcessExit = new ManualResetEventSlim(initialState: false);
76-
processExitHandler = (_, _) =>
73+
// Cancel asynchronously not to block the handler (as then the process might possibly run longer then what was the requested timeout)
74+
Task timeoutTask = Task.Delay(timeout.Value);
75+
Task cancelTask = Task.Factory.StartNew(context.Cancel);
76+
77+
// The process exits as soon as the event handler returns.
78+
// We provide a return value using Environment.ExitCode
79+
// because Main will not finish executing.
80+
// Wait for the invocation to finish.
81+
if (!blockProcessExit.Wait(timeout > TimeSpan.Zero
82+
? timeout.Value
83+
: Timeout.InfiniteTimeSpan))
7784
{
78-
// Cancel asynchronously not to block the handler (as then the process might possibly run longer then what was the requested timeout)
79-
Task timeoutTask = Task.Delay(timeout.Value);
80-
Task cancelTask = Task.Factory.StartNew(cts.Cancel);
81-
82-
// The process exits as soon as the event handler returns.
83-
// We provide a return value using Environment.ExitCode
84-
// because Main will not finish executing.
85-
// Wait for the invocation to finish.
86-
if (!blockProcessExit.Wait(timeout > TimeSpan.Zero
87-
? timeout.Value
88-
: Timeout.InfiniteTimeSpan))
89-
{
90-
context.ExitCode = SIGINT_EXIT_CODE;
91-
}
92-
// 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)
93-
else if (Task.WaitAny(timeoutTask, cancelTask) == 0)
94-
{
95-
// The async cancellation didn't finish in timely manner
96-
context.ExitCode = SIGINT_EXIT_CODE;
97-
}
98-
ExitCode = context.ExitCode;
99-
};
100-
// Default limit for ProcesExit handler is 2 seconds
101-
// https://docs.microsoft.com/en-us/dotnet/api/system.appdomain.processexit?view=net-6.0
102-
consoleHandler = (_, args) =>
85+
context.ExitCode = SIGINT_EXIT_CODE;
86+
}
87+
// 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)
88+
else if (Task.WaitAny(timeoutTask, cancelTask) == 0)
10389
{
104-
// Stop the process from terminating.
105-
// Since the context was cancelled, the invocation should
106-
// finish and Main will return.
107-
args.Cancel = true;
108-
109-
// If timeout was requested - make sure cancellation processing (or any other activity within the current process)
110-
// doesn't keep the process running after the timeout
111-
if (timeout! > TimeSpan.Zero)
112-
{
113-
Task
114-
.Delay(timeout.Value, default)
115-
.ContinueWith(t =>
116-
{
117-
// Prevent our ProcessExit from intervene and block the exit
118-
AppDomain.CurrentDomain.ProcessExit -= processExitHandler;
119-
Environment.Exit(SIGINT_EXIT_CODE);
120-
}, (CancellationToken)default);
121-
}
90+
// The async cancellation didn't finish in timely manner
91+
context.ExitCode = SIGINT_EXIT_CODE;
92+
}
93+
ExitCode = context.ExitCode;
94+
};
95+
// Default limit for ProcesExit handler is 2 seconds
96+
// https://docs.microsoft.com/en-us/dotnet/api/system.appdomain.processexit?view=net-6.0
97+
consoleHandler = (_, args) =>
98+
{
99+
// Stop the process from terminating.
100+
// Since the context was cancelled, the invocation should
101+
// finish and Main will return.
102+
args.Cancel = true;
103+
104+
// If timeout was requested - make sure cancellation processing (or any other activity within the current process)
105+
// doesn't keep the process running after the timeout
106+
if (timeout! > TimeSpan.Zero)
107+
{
108+
Task
109+
.Delay(timeout.Value, default)
110+
.ContinueWith(t =>
111+
{
112+
// Prevent our ProcessExit from intervene and block the exit
113+
AppDomain.CurrentDomain.ProcessExit -= processExitHandler;
114+
Environment.Exit(SIGINT_EXIT_CODE);
115+
}, (CancellationToken)default);
116+
}
122117

123-
// Cancel synchronously here - no need to perform it asynchronously as the timeout is already running (and would kill the process if needed),
124-
// plus we cannot wait only on the cancellation (e.g. via `Task.Factory.StartNew(cts.Cancel).Wait(cancelationProcessingTimeout.Value)`)
125-
// as we need to abort any other possible execution within the process - even outside the context of cancellation processing
126-
cts?.Cancel();
127-
};
128-
Console.CancelKeyPress += consoleHandler;
129-
AppDomain.CurrentDomain.ProcessExit += processExitHandler;
118+
// Cancel synchronously here - no need to perform it asynchronously as the timeout is already running (and would kill the process if needed),
119+
// plus we cannot wait only on the cancellation (e.g. via `Task.Factory.StartNew(cts.Cancel).Wait(cancelationProcessingTimeout.Value)`)
120+
// as we need to abort any other possible execution within the process - even outside the context of cancellation processing
121+
context.Cancel();
122+
};
130123

131-
return cts.Token;
132-
});
124+
Console.CancelKeyPress += consoleHandler;
125+
AppDomain.CurrentDomain.ProcessExit += processExitHandler;
133126

134127
try
135128
{
@@ -139,14 +132,13 @@ public static CommandLineBuilder CancelOnProcessTermination(
139132
{
140133
Console.CancelKeyPress -= consoleHandler;
141134
AppDomain.CurrentDomain.ProcessExit -= processExitHandler;
142-
Interlocked.Exchange(ref cts, null)?.Dispose();
143135
blockProcessExit?.Set();
144136
}
145137
}, MiddlewareOrderInternal.Startup);
146138

147139
return builder;
148140
}
149-
141+
150142
/// <summary>
151143
/// Enables the parser to recognize command line directives.
152144
/// </summary>
@@ -205,7 +197,7 @@ public static CommandLineBuilder EnablePosixBundling(
205197
builder.EnablePosixBundling = value;
206198
return builder;
207199
}
208-
200+
209201
/// <summary>
210202
/// Ensures that the application is registered with the <c>dotnet-suggest</c> tool to enable command line completions.
211203
/// </summary>
@@ -484,7 +476,7 @@ public static CommandLineBuilder AddMiddleware(
484476

485477
return builder;
486478
}
487-
479+
488480
/// <summary>
489481
/// Adds a middleware delegate to the invocation pipeline called before a command handler is invoked.
490482
/// </summary>

‎src/System.CommandLine/Invocation/InvocationContext.cs‎

Lines changed: 22 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright (c) .NET Foundation and contributors. All rights reserved.
1+
// Copyright (c) .NET Foundation and contributors. All rights reserved.
22
// Licensed under the MIT license. See LICENSE file in the project root for full license information.
33

44
using System.Collections.Generic;
@@ -18,9 +18,9 @@ public sealed class InvocationContext : IDisposable
1818
private HelpBuilder? _helpBuilder;
1919
private BindingContext? _bindingContext;
2020
private IConsole? _console;
21-
private CancellationTokenSource? _linkedTokensSource;
22-
private List<Func<CancellationToken>> _cancellationTokens = new(3);
23-
private Lazy<CancellationToken> _lazyCancellationToken;
21+
private readonly CancellationToken _token;
22+
private readonly LinkedList<CancellationTokenRegistration> _registrations = new();
23+
private volatile CancellationTokenSource? _source;
2424

2525
/// <param name="parseResult">The result of the current parse operation.</param>
2626
/// <param name="console">The console to which output is to be written.</param>
@@ -32,8 +32,13 @@ public InvocationContext(
3232
{
3333
ParseResult = parseResult;
3434
_console = console;
35-
_cancellationTokens.Add(() => cancellationToken);
36-
_lazyCancellationToken = new Lazy<CancellationToken>(BuildCancellationToken);
35+
36+
_source = new CancellationTokenSource();
37+
_token = _source.Token;
38+
if (cancellationToken.CanBeCanceled)
39+
{
40+
LinkToken(cancellationToken);
41+
}
3742
}
3843

3944
/// <summary>
@@ -105,41 +110,27 @@ public IConsole Console
105110
/// <summary>
106111
/// Gets a cancellation token that can be used to check if cancellation has been requested.
107112
/// </summary>
108-
public CancellationToken GetCancellationToken() => _lazyCancellationToken.Value;
113+
public CancellationToken GetCancellationToken() => _token;
109114

110-
private CancellationToken BuildCancellationToken()
115+
internal void Cancel()
111116
{
112-
switch(_cancellationTokens.Count)
113-
{
114-
case 0: return CancellationToken.None;
115-
case 1: return _cancellationTokens[0]();
116-
default:
117-
CancellationToken[] tokens = new CancellationToken[_cancellationTokens.Count];
118-
for(int i = 0; i < _cancellationTokens.Count; i++)
119-
{
120-
tokens[i] = _cancellationTokens[i]();
121-
}
122-
_linkedTokensSource = CancellationTokenSource.CreateLinkedTokenSource(tokens);
123-
return _linkedTokensSource.Token;
124-
};
117+
using var source = Interlocked.Exchange(ref _source, null);
118+
source?.Cancel();
125119
}
126120

127-
128-
/// <inheritdoc />
129-
public void Dispose()
121+
public void LinkToken(CancellationToken token)
130122
{
131-
_linkedTokensSource?.Dispose();
132-
_linkedTokensSource = null;
133-
(Console as IDisposable)?.Dispose();
123+
_registrations.AddLast(token.Register(Cancel));
134124
}
135125

136-
public void AddLinkedCancellationToken(Func<CancellationToken> token)
126+
/// <inheritdoc />
127+
void IDisposable.Dispose()
137128
{
138-
if (_lazyCancellationToken.IsValueCreated)
129+
Interlocked.Exchange(ref _source, null)?.Dispose();
130+
foreach (CancellationTokenRegistration registration in _registrations)
139131
{
140-
throw new InvalidOperationException($"Cannot add additional linked cancellation tokens once {nameof(InvocationContext)}.{nameof(CancellationToken)} has been invoked");
132+
registration.Dispose();
141133
}
142-
_cancellationTokens.Add(token);
143134
}
144135
}
145136
}

0 commit comments

Comments
 (0)