Skip to content

Commit 8c41aa3

Browse files
KebooKeboo
committed
Update with suggestions to better handled linked token
1 parent abe1b8d commit 8c41aa3

File tree

5 files changed

+82
-113
lines changed

5 files changed

+82
-113
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -368,6 +368,7 @@ System.CommandLine.Invocation
368368
public System.CommandLine.Parsing.Parser Parser { get; }
369369
public System.CommandLine.ParseResult ParseResult { get; set; }
370370
public System.Threading.CancellationToken GetCancellationToken()
371+
public System.Void LinkToken(System.Threading.CancellationToken token)
371372
public delegate InvocationMiddleware : System.MulticastDelegate, System.ICloneable, System.Runtime.Serialization.ISerializable
372373
.ctor(System.Object object, System.IntPtr method)
373374
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
@@ -353,7 +353,7 @@ public async Task Command_InvokeAsync_can_cancel_from_middleware()
353353
})
354354
.AddMiddleware(async (context, next) =>
355355
{
356-
context.AddLinkedCancellationToken(() => cts.Token);
356+
context.LinkToken(cts.Token);
357357
cts.Cancel();
358358
await next(context);
359359
})

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

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

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

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

132-
return cts.Token;
133-
});
125+
Console.CancelKeyPress += consoleHandler;
126+
AppDomain.CurrentDomain.ProcessExit += processExitHandler;
134127

135128
try
136129
{
@@ -140,14 +133,13 @@ public static CommandLineBuilder CancelOnProcessTermination(
140133
{
141134
Console.CancelKeyPress -= consoleHandler;
142135
AppDomain.CurrentDomain.ProcessExit -= processExitHandler;
143-
Interlocked.Exchange(ref cts, null)?.Dispose();
144136
blockProcessExit?.Set();
145137
}
146138
}, MiddlewareOrderInternal.Startup);
147139

148140
return builder;
149141
}
150-
142+
151143
/// <summary>
152144
/// Enables the parser to recognize command line directives.
153145
/// </summary>
@@ -206,7 +198,7 @@ public static CommandLineBuilder EnablePosixBundling(
206198
builder.EnablePosixBundling = value;
207199
return builder;
208200
}
209-
201+
210202
/// <summary>
211203
/// Ensures that the application is registered with the <c>dotnet-suggest</c> tool to enable command line completions.
212204
/// </summary>
@@ -485,7 +477,7 @@ public static CommandLineBuilder AddMiddleware(
485477

486478
return builder;
487479
}
488-
480+
489481
/// <summary>
490482
/// Adds a middleware delegate to the invocation pipeline called before a command handler is invoked.
491483
/// </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)