Skip to content

Remove partitioning from CancellationTokenSource #48251

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
stephentoub merged 1 commit into from
Feb 18, 2021
Merged

Remove partitioning from CancellationTokenSource #48251

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/libraries/System.Private.CoreLib/src/System/Threading/CancellationToken.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -288,7 +288,7 @@ private CancellationTokenRegistration Register(Delegate callback, object? state,

CancellationTokenSource? source = _source;
return source != null ?
source.InternalRegister(callback, state, useSynchronizationContext ? SynchronizationContext.Current : null, useExecutionContext ? ExecutionContext.Capture() : null) :
source.Register(callback, state, useSynchronizationContext ? SynchronizationContext.Current : null, useExecutionContext ? ExecutionContext.Capture() : null) :
default; // Nothing to do for tokens than can never reach the canceled state. Give back a dummy registration.
}

Expand Down
113 changes: 51 additions & 62 deletions src/libraries/System.Private.CoreLib/src/System/Threading/CancellationTokenRegistration.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,10 +31,31 @@ internal CancellationTokenRegistration(long id, CancellationTokenSource.Callback
/// </summary>
public void Dispose()
{
CancellationTokenSource.CallbackNode node = _node;
if (node != null && !node.Partition.Unregister(_id, node))
if (_node is CancellationTokenSource.CallbackNode node && !node.Registrations.Unregister(_id, node))
{
WaitForCallbackIfNecessary();
WaitForCallbackIfNecessary(_id, node);

static void WaitForCallbackIfNecessary(long id, CancellationTokenSource.CallbackNode node)
{
// We're a valid registration but we were unable to unregister, which means the callback wasn't in the list,
// which means either it already executed or it's currently executing. We guarantee that we will not return
// if the callback is being executed (assuming we are not currently called by the callback itself)
// We achieve this by the following rules:
// 1. If we are called in the context of an executing callback, no need to wait (determined by tracking callback-executor threadID)
// - if the currently executing callback is this CTR, then waiting would deadlock. (We choose to return rather than deadlock)
// - if not, then this CTR cannot be the one executing, hence no need to wait
// 2. If unregistration failed, and we are on a different thread, then the callback may be running under control of cts.Cancel()
// => poll until cts.ExecutingCallback is not the one we are trying to unregister.
CancellationTokenSource source = node.Registrations.Source;
if (source.IsCancellationRequested && // Running callbacks has commenced.
!source.IsCancellationCompleted && // Running callbacks hasn't finished.
node.Registrations.ThreadIDExecutingCallbacks != Environment.CurrentManagedThreadId) // The executing thread ID is not this thread's ID.
{
// Callback execution is in progress, the executing thread is different from this thread and has taken the callback for execution
// so observe and wait until this target callback is no longer the executing callback.
node.Registrations.WaitForCallbackToComplete(id);
}
}
}
}

Expand All @@ -47,10 +68,27 @@ public void Dispose()
/// </summary>
public ValueTask DisposeAsync()
{
CancellationTokenSource.CallbackNode node = _node;
return node != null && !node.Partition.Unregister(_id, node) ?
WaitForCallbackIfNecessaryAsync() :
return _node is CancellationTokenSource.CallbackNode node && !node.Registrations.Unregister(_id, node) ?
WaitForCallbackIfNecessaryAsync(_id, node) :
default;

static ValueTask WaitForCallbackIfNecessaryAsync(long id, CancellationTokenSource.CallbackNode node)
{
// Same as WaitForCallbackIfNecessary, except returning a task that'll be completed when callbacks complete.

CancellationTokenSource source = node.Registrations.Source;
if (source.IsCancellationRequested && // Running callbacks has commenced.
!source.IsCancellationCompleted && // Running callbacks hasn't finished.
node.Registrations.ThreadIDExecutingCallbacks != Environment.CurrentManagedThreadId) // The executing thread ID is not this thread's ID.
{
// Callback execution is in progress, the executing thread is different from this thread and has taken the callback for execution
// so get a task that'll complete when this target callback is no longer the executing callback.
return node.Registrations.WaitForCallbackToCompleteAsync(id);
}

// Callback is either already completed, won't execute, or the callback itself is calling this.
return default;
}
}

/// <summary>Gets the <see cref="CancellationToken"/> with which this registration is associated.</summary>
Expand All @@ -59,66 +97,17 @@ public ValueTask DisposeAsync()
/// to <see cref="CancellationToken.Register"/> on a token that already had cancellation requested),
/// this will return a default token.
/// </remarks>
public CancellationToken Token
{
get
{
CancellationTokenSource.CallbackNode node = _node;
return node != null ?
new CancellationToken(node.Partition.Source) : // avoid CTS.Token, which throws after disposal
default;
}
}
public CancellationToken Token =>
_node is CancellationTokenSource.CallbackNode node ?
new CancellationToken(node.Registrations.Source) : // avoid CTS.Token, which throws after disposal
default;

/// <summary>
/// Disposes of the registration and unregisters the target callback from the associated
/// <see cref="System.Threading.CancellationToken">CancellationToken</see>.
/// </summary>
public bool Unregister()
{
CancellationTokenSource.CallbackNode node = _node;
return node != null && node.Partition.Unregister(_id, node);
}

private void WaitForCallbackIfNecessary()
{
// We're a valid registration but we were unable to unregister, which means the callback wasn't in the list,
// which means either it already executed or it's currently executing. We guarantee that we will not return
// if the callback is being executed (assuming we are not currently called by the callback itself)
// We achieve this by the following rules:
// 1. If we are called in the context of an executing callback, no need to wait (determined by tracking callback-executor threadID)
// - if the currently executing callback is this CTR, then waiting would deadlock. (We choose to return rather than deadlock)
// - if not, then this CTR cannot be the one executing, hence no need to wait
// 2. If unregistration failed, and we are on a different thread, then the callback may be running under control of cts.Cancel()
// => poll until cts.ExecutingCallback is not the one we are trying to unregister.
CancellationTokenSource source = _node.Partition.Source;
if (source.IsCancellationRequested && // Running callbacks has commenced.
!source.IsCancellationCompleted && // Running callbacks hasn't finished.
source.ThreadIDExecutingCallbacks != Environment.CurrentManagedThreadId) // The executing thread ID is not this thread's ID.
{
// Callback execution is in progress, the executing thread is different from this thread and has taken the callback for execution
// so observe and wait until this target callback is no longer the executing callback.
source.WaitForCallbackToComplete(_id);
}
}

private ValueTask WaitForCallbackIfNecessaryAsync()
{
// Same as WaitForCallbackIfNecessary, except returning a task that'll be completed when callbacks complete.

CancellationTokenSource source = _node.Partition.Source;
if (source.IsCancellationRequested && // Running callbacks has commenced.
!source.IsCancellationCompleted && // Running callbacks hasn't finished.
source.ThreadIDExecutingCallbacks != Environment.CurrentManagedThreadId) // The executing thread ID is not this thread's ID.
{
// Callback execution is in progress, the executing thread is different from this thread and has taken the callback for execution
// so get a task that'll complete when this target callback is no longer the executing callback.
return source.WaitForCallbackToCompleteAsync(_id);
}

// Callback is either already completed, won't execute, or the callback itself is calling this.
return default;
}
public bool Unregister() =>
_node is CancellationTokenSource.CallbackNode node && node.Registrations.Unregister(_id, node);

/// <summary>
/// Determines whether two <see
Expand Down Expand Up @@ -148,7 +137,7 @@ private ValueTask WaitForCallbackIfNecessaryAsync()
/// they both refer to the output of a single call to the same Register method of a
/// <see cref="System.Threading.CancellationToken">CancellationToken</see>.
/// </returns>
public override bool Equals([NotNullWhen(true)] object? obj) => obj is CancellationTokenRegistration && Equals((CancellationTokenRegistration)obj);
public override bool Equals([NotNullWhen(true)] object? obj) => obj is CancellationTokenRegistration other && Equals(other);

/// <summary>
/// Determines whether the current <see cref="System.Threading.CancellationToken">CancellationToken</see> instance is equal to the
Expand Down
Loading