Skip to content

Fixes some theoretical races and retentions in AsyncOperation and ManualResetValueTaskSource #127256

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

Open
VSadov wants to merge 5 commits into
base: main
Choose a base branch
from
+111 −74
Open

Fixes some theoretical races and retentions in AsyncOperation and ManualResetValueTaskSource #127256

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
06ad1e7
reset _continuation on completion
VSadov Apr 21, 2026
246579a
some theoretical races and retentions
VSadov Apr 22, 2026
4ac167b
Apply suggestions from code review
VSadov Apr 22, 2026
38bf20c
PR feedback
VSadov Apr 22, 2026
20ec443
PR feedback
VSadov Apr 28, 2026
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
76 changes: 44 additions & 32 deletions ...stem.Private.CoreLib/src/System/Threading/Tasks/Sources/ManualResetValueTaskSourceCore.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,8 +34,6 @@ public struct ManualResetValueTaskSourceCore<TResult>
private TResult? _result;
/// <summary>The current version of this value, used to help prevent misuse.</summary>
private short _version;
/// <summary>Whether the current operation has completed.</summary>
private bool _completed;
/// <summary>Whether to force continuations to run asynchronously.</summary>
private bool _runContinuationsAsynchronously;

Expand All @@ -51,13 +49,15 @@ public bool RunContinuationsAsynchronously
public void Reset()
{
// Reset/update state for the next use/await of this instance.
// Order of assignments is unimportant here.
// The outer user always ensures that the state is not accessed across
// the reset point when implementing Rent/Return operations.
_version++;
_continuation = null;
_continuationState = null;
Comment thread
VSadov marked this conversation as resolved.
_capturedContext = null;
_error = null;
_result = default;
_completed = false;
}

/// <summary>Completes with a successful result.</summary>
Expand All @@ -79,13 +79,16 @@ public void SetException(Exception error)
/// <summary>Gets the operation version.</summary>
public short Version => _version;

/// <summary>Gets whether the operation has completed.</summary>
internal bool IsCompleted => ReferenceEquals(Volatile.Read(ref _continuation), ManualResetValueTaskSourceCoreShared.s_sentinel);

/// <summary>Gets the status of the operation.</summary>
/// <param name="token">Opaque value that was provided to the <see cref="ValueTask"/>'s constructor.</param>
public ValueTaskSourceStatus GetStatus(short token)
{
ValidateToken(token);
return
Volatile.Read(ref _continuation) is null || !_completed ? ValueTaskSourceStatus.Pending :
!IsCompleted ? ValueTaskSourceStatus.Pending :
Comment thread
VSadov marked this conversation as resolved.
_error is null ? ValueTaskSourceStatus.Succeeded :
_error.SourceException is OperationCanceledException ? ValueTaskSourceStatus.Canceled :
ValueTaskSourceStatus.Faulted;
Expand All @@ -96,7 +99,7 @@ public ValueTaskSourceStatus GetStatus(short token)
[StackTraceHidden]
public TResult GetResult(short token)
{
if (token != _version || !_completed || _error is not null)
if (token != _version || !IsCompleted || _error is not null)
{
ThrowForFailedGetResult();
}
Expand Down Expand Up @@ -125,6 +128,17 @@ public void OnCompleted(Action<object?> continuation, object? state, short token
}
ValidateToken(token);

// We need to store the state before the CompareExchange, so that if it completes immediately
// after the CompareExchange, it'll find the state already stored. If someone misuses this
// and schedules multiple continuations erroneously, we could end up using the wrong state.
// Make a best-effort attempt to catch such misuse.
if (_continuationState is not null)
{
ThrowHelper.ThrowInvalidOperationException();
}
_continuationState = state;

Debug.Assert(_capturedContext is null);
if ((flags & ValueTaskSourceOnCompletedFlags.FlowExecutionContext) != 0)
{
_capturedContext = ExecutionContext.Capture();
Expand All @@ -151,35 +165,36 @@ public void OnCompleted(Action<object?> continuation, object? state, short token
}
}

// We need to set the continuation state before we swap in the delegate, so that
// if there's a race between this and SetResult/Exception and SetResult/Exception
// sees the _continuation as non-null, it'll be able to invoke it with the state
// stored here. However, this also means that if this is used incorrectly (e.g.
// awaited twice concurrently), _continuationState might get erroneously overwritten.
// To minimize the chances of that, we check preemptively whether _continuation
// is already set to something other than the completion sentinel.
object? storedContinuation = _continuation;
if (storedContinuation is null)
// Try to set the provided continuation into _continuation. If this succeeds, that means the operation
// has not yet completed, and the completer will be responsible for invoking the callback. If this fails,
// that means the operation has already completed, and we must invoke the callback, but because we're still
// inside the awaiter's OnCompleted method and we want to avoid possible stack dives, we must invoke
// the continuation asynchronously rather than synchronously.
Action<object?>? prevContinuation = Interlocked.CompareExchange(ref _continuation, continuation, null);
if (prevContinuation is null)
{
_continuationState = state;
storedContinuation = Interlocked.CompareExchange(ref _continuation, continuation, null);
if (storedContinuation is null)
{
// Operation hadn't already completed, so we're done. The continuation will be
// invoked when SetResult/Exception is called at some later point.
return;
}
// Operation hadn't already completed, so we're done. The continuation will be
// invoked when SetResult/Exception is called at some later point.
return;
}

// Operation already completed, so we need to queue the supplied callback.
// At this point the storedContinuation should be the sentinal; if it's not, the instance was misused.
Debug.Assert(storedContinuation is not null, $"{nameof(storedContinuation)} is null");
if (!ReferenceEquals(storedContinuation, ManualResetValueTaskSourceCoreShared.s_sentinel))
// Queue the continuation. We always queue here, even if !RunContinuationsAsynchronously, in order
// to avoid stack diving; this path happens in the rare race when we're setting up to await and the
// object is completed after the awaiter.IsCompleted but before the awaiter.OnCompleted.

// We no longer need the stored values as we will be passing the state when queuing directly
_continuationState = null;
object? capturedContext = _capturedContext;
_capturedContext = null;

// If the set failed because there's already a delegate in _continuation, but that delegate is
// something other than the completion sentinel, something went wrong, which should only happen if
// the instance was erroneously used, likely to hook up multiple continuations.
if (!ReferenceEquals(prevContinuation, ManualResetValueTaskSourceCoreShared.s_sentinel))
{
ThrowHelper.ThrowInvalidOperationException();
}

object? capturedContext = _capturedContext;
switch (capturedContext)
{
case null:
Expand Down Expand Up @@ -209,15 +224,12 @@ private void ValidateToken(short token)
/// <summary>Signals that the operation has completed. Invoked after the result or error has been set.</summary>
private void SignalCompletion()
{
if (_completed)
if (IsCompleted)
{
ThrowHelper.ThrowInvalidOperationException();
}
_completed = true;

Action<object?>? continuation =
Volatile.Read(ref _continuation) ??
Interlocked.CompareExchange(ref _continuation, ManualResetValueTaskSourceCoreShared.s_sentinel, null);
Action<object?>? continuation = Interlocked.Exchange(ref _continuation, ManualResetValueTaskSourceCoreShared.s_sentinel);

if (continuation is not null)
{
Expand Down
109 changes: 67 additions & 42 deletions src/libraries/System.Threading.Channels/src/System/Threading/Channels/AsyncOperation.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,7 @@ private CancellationToken CancellationToken
#endif

/// <summary>Gets whether the operation has completed.</summary>
internal bool IsCompleted => ReferenceEquals(_continuation, s_completedSentinel);
internal bool IsCompleted => ReferenceEquals(Volatile.Read(ref _continuation), s_completedSentinel);
Copy link
Copy Markdown
Member Author

@VSadov VSadov Apr 22, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Performing an ordinary read could allow prefetching _error and we could end up assuming that a failed/cancelled operation actually cleanly completed.


/// <summary>Completes the operation with a failed state and the specified error.</summary>
/// <param name="exception">The error.</param>
Expand Down Expand Up @@ -267,12 +267,13 @@ private void SetCompletionAndInvokeContinuation()
Debug.Assert(_continuation is not null);

object? ctx = _capturedContext;
_capturedContext = null;

ExecutionContext? ec =
ctx is null ? null :
ctx as ExecutionContext ??
(ctx as CapturedSchedulerAndExecutionContext)?._executionContext;

_capturedContext = null;
if (ec is null)
{
Action<object?> c = _continuation!;
Expand Down Expand Up @@ -359,43 +360,52 @@ public void OnCompleted(Action<object?> continuation, object? state, short token
// inside the awaiter's OnCompleted method and we want to avoid possible stack dives, we must invoke
// the continuation asynchronously rather than synchronously.
Action<object?>? prevContinuation = Interlocked.CompareExchange(ref _continuation, continuation, null);
if (prevContinuation is not null)
if (prevContinuation is null)
{
// If the set failed because there's already a delegate in _continuation, but that delegate is
// something other than s_completedSentinel, something went wrong, which should only happen if
// the instance was erroneously used, likely to hook up multiple continuations.
Debug.Assert(IsCompleted, $"Expected IsCompleted");
if (!ReferenceEquals(prevContinuation, s_completedSentinel))
{
Debug.Assert(prevContinuation != s_availableSentinel, "Continuation was the available sentinel.");
ThrowMultipleContinuations();
}
// Operation hadn't already completed, so we're done. The continuation will be
// invoked when SetResult/Exception is called at some later point.
return;
}

// Queue the continuation. We always queue here, even if !RunContinuationsAsynchronously, in order
// to avoid stack diving; this path happens in the rare race when we're setting up to await and the
// object is completed after the awaiter.IsCompleted but before the awaiter.OnCompleted.
if (_capturedContext is null)
{
ChannelUtilities.UnsafeQueueUserWorkItem(continuation, state);
}
else if (sc is not null)
{
sc.Post(static s =>
{
var t = (KeyValuePair<Action<object?>, object?>)s!;
t.Key(t.Value);
}, new KeyValuePair<Action<object?>, object?>(continuation, state));
}
else if (ts is not null)
{
Debug.Assert(ts is not null);
Task.Factory.StartNew(continuation, state, CancellationToken.None, TaskCreationOptions.DenyChildAttach, ts);
}
else
// Queue the continuation. We always queue here, even if !RunContinuationsAsynchronously, in order
// to avoid stack diving; this path happens in the rare race when we're setting up to await and the
// object is completed after the awaiter.IsCompleted but before the awaiter.OnCompleted.

// We no longer need the stored values as we will be passing the state when queuing directly
_continuationState = null;
object? capturedContext = _capturedContext;
_capturedContext = null;

// If the set failed because there's already a delegate in _continuation, but that delegate is
// something other than the completion sentinel, something went wrong, which should only happen if
// the instance was erroneously used, likely to hook up multiple continuations.
Debug.Assert(prevContinuation != s_availableSentinel, "Continuation was the available sentinel.");
Debug.Assert(IsCompleted, $"Expected IsCompleted");
if (!ReferenceEquals(prevContinuation, s_completedSentinel))
{
ThrowMultipleContinuations();
}

if (capturedContext is null)
{
ChannelUtilities.UnsafeQueueUserWorkItem(continuation, state);
}
else if (sc is not null)
{
sc.Post(static s =>
{
Debug.Assert(_capturedContext is ExecutionContext);
ChannelUtilities.QueueUserWorkItem(continuation, state);
}
var t = (KeyValuePair<Action<object?>, object?>)s!;
t.Key(t.Value);
}, new KeyValuePair<Action<object?>, object?>(continuation, state));
}
else if (ts is not null)
{
Task.Factory.StartNew(continuation, state, CancellationToken.None, TaskCreationOptions.DenyChildAttach, ts);
}
else
{
Debug.Assert(capturedContext is ExecutionContext);
ChannelUtilities.QueueUserWorkItem(continuation, state);
}
}

Expand Down Expand Up @@ -470,11 +480,22 @@ void IValueTaskSource.GetResult(short token)

if (_pooled)
{
Volatile.Write(ref _continuation, s_availableSentinel); // only after fetching all needed data
ClearRetainedState();
// enable reuse only after fetching all needed data
Volatile.Write(ref _continuation, s_availableSentinel);
}

error?.Throw();
}

internal virtual void ClearRetainedState()
{
_error = null;
Comment thread
VSadov marked this conversation as resolved.
Debug.Assert((object?)Next == null);
Debug.Assert((object?)Previous == null);
Debug.Assert(_continuationState == null);
Debug.Assert(_capturedContext == null);
}
Comment thread
VSadov marked this conversation as resolved.
}

/// <summary>Represents an asynchronous operation with a result on a channel.</summary>
Expand Down Expand Up @@ -515,24 +536,28 @@ public TResult GetResult(short token)

if (_pooled)
{
Volatile.Write(ref _continuation, s_availableSentinel); // only after fetching all needed data
ClearRetainedState();
// enable reuse only after fetching all needed data
Volatile.Write(ref _continuation, s_availableSentinel);
}

error?.Throw();
return result!;
}

internal override void ClearRetainedState()
{
base.ClearRetainedState();
_result = default;
}

/// <summary>Attempts to take ownership of the pooled instance.</summary>
/// <returns>true if the instance is now owned by the caller, in which case its state has been reset; otherwise, false.</returns>
public bool TryOwnAndReset()
{
Debug.Assert(_pooled, "Should only be used for pooled objects");
if (ReferenceEquals(Interlocked.CompareExchange(ref _continuation, null, s_availableSentinel), s_availableSentinel))
{
_continuationState = null;
_result = default;
_error = null;
_capturedContext = null;
return true;
Comment thread
VSadov marked this conversation as resolved.
}

Expand Down
Loading