From 4da32cd409996302b6e522c66cedc6b8a8f952af Mon Sep 17 00:00:00 2001 From: Stephen Toub Date: Thu, 2 Mar 2023 11:41:25 -0500 Subject: [PATCH] Combine ExecutionContext/CapturedContext fields in ManualResetValueTaskSourceCore (#82181) There's currently an _executionContext field and a _capturedContext field. Both of these are for non-fast-path. This not only means the struct is larger than is necessary for the common case, it also means when we reset the instance between operations we need to clear an extra field, and it means we have an extra branch on some paths to check both fields. We can instead combine them into a single field. We allocate an extra tuple object if there's both an ExecutionContext and a scheduler, but this is exceedingly rare: when used as part of awaits, there will never be an ExecutionContext, so this will only happen in situations where someone is directly using the awaiter's OnCompleted (not UnsafeOnCompleted) method and there's a scheduler and ConfigureAwait(false) wasn't used. Such a situation not only is rare, it also already has additional overheads. This also cleans up a bit about how exceptions are handled and moves more logic out of the generic type to avoid code bloat with generic instantiations. --- .../Sources/ManualResetValueTaskSourceCore.cs | 312 ++++++++++-------- .../src/System/ThrowHelper.cs | 3 + 2 files changed, 174 insertions(+), 141 deletions(-) diff --git a/src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/Sources/ManualResetValueTaskSourceCore.cs b/src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/Sources/ManualResetValueTaskSourceCore.cs index c2bd5ab8f3f90..c29c5e759f2c1 100644 --- a/src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/Sources/ManualResetValueTaskSourceCore.cs +++ b/src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/Sources/ManualResetValueTaskSourceCore.cs @@ -8,7 +8,7 @@ namespace System.Threading.Tasks.Sources { /// Provides the core logic for implementing a manual-reset or . - /// + /// Specifies the type of results of the operation represented by this instance. [StructLayout(LayoutKind.Auto)] public struct ManualResetValueTaskSourceCore { @@ -20,38 +20,44 @@ public struct ManualResetValueTaskSourceCore private Action? _continuation; /// State to pass to . private object? _continuationState; - /// to flow to the callback, or null if no flowing is required. - private ExecutionContext? _executionContext; /// - /// A "captured" or with which to invoke the callback, - /// or null if no special context is required. + /// Null if no special context was found. + /// ExecutionContext if one was captured due to needing to be flowed. + /// A scheduler (TaskScheduler or SynchronizationContext) if one was captured and needs to be used for callback scheduling. + /// Or a CapturedContext if there's both an ExecutionContext and a scheduler. + /// The most common and the fast path case to optimize for is null. /// private object? _capturedContext; - /// Whether the current operation has completed. - private bool _completed; - /// The result with which the operation succeeded, or the default value if it hasn't yet completed or failed. - private TResult? _result; /// The exception with which the operation failed, or null if it hasn't yet completed or completed successfully. private ExceptionDispatchInfo? _error; + /// The result with which the operation succeeded, or the default value if it hasn't yet completed or failed. + private TResult? _result; /// The current version of this value, used to help prevent misuse. private short _version; + /// Whether the current operation has completed. + private bool _completed; + /// Whether to force continuations to run asynchronously. + private bool _runContinuationsAsynchronously; /// Gets or sets whether to force continuations to run asynchronously. /// Continuations may run asynchronously if this is false, but they'll never run synchronously if this is true. - public bool RunContinuationsAsynchronously { get; set; } + public bool RunContinuationsAsynchronously + { + get => _runContinuationsAsynchronously; + set => _runContinuationsAsynchronously = value; + } /// Resets to prepare for the next operation. public void Reset() { // Reset/update state for the next use/await of this instance. _version++; - _completed = false; - _result = default; - _error = null; - _executionContext = null; - _capturedContext = null; _continuation = null; _continuationState = null; + _capturedContext = null; + _error = null; + _result = default; + _completed = false; } /// Completes with a successful result. @@ -79,8 +85,8 @@ public ValueTaskSourceStatus GetStatus(short token) { ValidateToken(token); return - Volatile.Read(ref _continuation) == null || !_completed ? ValueTaskSourceStatus.Pending : - _error == null ? ValueTaskSourceStatus.Succeeded : + Volatile.Read(ref _continuation) is null || !_completed ? ValueTaskSourceStatus.Pending : + _error is null ? ValueTaskSourceStatus.Succeeded : _error.SourceException is OperationCanceledException ? ValueTaskSourceStatus.Canceled : ValueTaskSourceStatus.Faulted; } @@ -92,22 +98,18 @@ public TResult GetResult(short token) { if (token != _version || !_completed || _error is not null) { - ThrowForFailedGetResult(token); + ThrowForFailedGetResult(); } return _result!; } + /// Throws an exception in response to a failed . [StackTraceHidden] - private void ThrowForFailedGetResult(short token) + private void ThrowForFailedGetResult() { - if (token != _version || !_completed) - { - ThrowHelper.ThrowInvalidOperationException(); - } - _error?.Throw(); - Debug.Fail($"{nameof(ThrowForFailedGetResult)} should never get here"); + throw new InvalidOperationException(); // not using ThrowHelper.ThrowInvalidOperationException so that the JIT sees ThrowForFailedGetResult as always throwing } /// Schedules the continuation action for this operation. @@ -117,28 +119,34 @@ private void ThrowForFailedGetResult(short token) /// The flags describing the behavior of the continuation. public void OnCompleted(Action continuation, object? state, short token, ValueTaskSourceOnCompletedFlags flags) { - ArgumentNullException.ThrowIfNull(continuation); - + if (continuation is null) + { + ThrowHelper.ThrowArgumentNullException(ExceptionArgument.continuation); + } ValidateToken(token); if ((flags & ValueTaskSourceOnCompletedFlags.FlowExecutionContext) != 0) { - _executionContext = ExecutionContext.Capture(); + _capturedContext = ExecutionContext.Capture(); } if ((flags & ValueTaskSourceOnCompletedFlags.UseSchedulingContext) != 0) { - SynchronizationContext? sc = SynchronizationContext.Current; - if (sc != null && sc.GetType() != typeof(SynchronizationContext)) + if (SynchronizationContext.Current is SynchronizationContext sc && + sc.GetType() != typeof(SynchronizationContext)) { - _capturedContext = sc; + _capturedContext = _capturedContext is null ? + sc : + new CapturedSchedulerAndExecutionContext(sc, (ExecutionContext)_capturedContext); } else { TaskScheduler ts = TaskScheduler.Current; if (ts != TaskScheduler.Default) { - _capturedContext = ts; + _capturedContext = _capturedContext is null ? + ts : + new CapturedSchedulerAndExecutionContext(ts, (ExecutionContext)_capturedContext); } } } @@ -150,47 +158,41 @@ public void OnCompleted(Action continuation, object? state, short token // 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? oldContinuation = _continuation; - if (oldContinuation == null) + object? storedContinuation = _continuation; + if (storedContinuation is null) { _continuationState = state; - oldContinuation = Interlocked.CompareExchange(ref _continuation, continuation, null); + 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; + } } - if (oldContinuation != null) + // 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)) { - // Operation already completed, so we need to queue the supplied callback. - if (!ReferenceEquals(oldContinuation, ManualResetValueTaskSourceCoreShared.s_sentinel)) - { - ThrowHelper.ThrowInvalidOperationException(); - } + ThrowHelper.ThrowInvalidOperationException(); + } - switch (_capturedContext) - { - case null: - if (_executionContext != null) - { - ThreadPool.QueueUserWorkItem(continuation, state, preferLocal: true); - } - else - { - ThreadPool.UnsafeQueueUserWorkItem(continuation, state, preferLocal: true); - } - break; - - case SynchronizationContext sc: - sc.Post(static s => - { - var tuple = (TupleSlim, object?>)s!; - tuple.Item1(tuple.Item2); - }, new TupleSlim, object?>(continuation, state)); - break; - - case TaskScheduler ts: - Task.Factory.StartNew(continuation, state, CancellationToken.None, TaskCreationOptions.DenyChildAttach, ts); - break; - } + object? capturedContext = _capturedContext; + switch (capturedContext) + { + case null: + ThreadPool.UnsafeQueueUserWorkItem(continuation, state, preferLocal: true); + break; + + case ExecutionContext: + ThreadPool.QueueUserWorkItem(continuation, state, preferLocal: true); + break; + + default: + ManualResetValueTaskSourceCoreShared.ScheduleCapturedContext(capturedContext, continuation, state); + break; } } @@ -213,62 +215,125 @@ private void SignalCompletion() } _completed = true; - if (Volatile.Read(ref _continuation) is null && Interlocked.CompareExchange(ref _continuation, ManualResetValueTaskSourceCoreShared.s_sentinel, null) is null) - { - return; - } + Action? continuation = + Volatile.Read(ref _continuation) ?? + Interlocked.CompareExchange(ref _continuation, ManualResetValueTaskSourceCoreShared.s_sentinel, null); - Debug.Assert(_continuation is not null); - - if (_executionContext is null) + if (continuation is not null) { - if (_capturedContext is null) + Debug.Assert(continuation is not null, $"{nameof(continuation)} is null"); + + object? context = _capturedContext; + if (context is null) { - if (RunContinuationsAsynchronously) + if (_runContinuationsAsynchronously) { - ThreadPool.UnsafeQueueUserWorkItem(_continuation, _continuationState, preferLocal: true); + ThreadPool.UnsafeQueueUserWorkItem(continuation, _continuationState, preferLocal: true); } else { - _continuation(_continuationState); + continuation(_continuationState); } } + else if (context is ExecutionContext or CapturedSchedulerAndExecutionContext) + { + ManualResetValueTaskSourceCoreShared.InvokeContinuationWithContext(context, continuation, _continuationState, _runContinuationsAsynchronously); + } else { - InvokeSchedulerContinuation(); + Debug.Assert(context is TaskScheduler or SynchronizationContext, $"context is {context}"); + ManualResetValueTaskSourceCoreShared.ScheduleCapturedContext(context, continuation, _continuationState); } } - else + } + } + + /// A tuple of both a non-null scheduler and a non-null ExecutionContext. + internal sealed class CapturedSchedulerAndExecutionContext + { + internal readonly object _scheduler; + internal readonly ExecutionContext _executionContext; + + public CapturedSchedulerAndExecutionContext(object scheduler, ExecutionContext executionContext) + { + Debug.Assert(scheduler is SynchronizationContext or TaskScheduler, $"{nameof(scheduler)} is {scheduler}"); + Debug.Assert(executionContext is not null, $"{nameof(executionContext)} is null"); + + _scheduler = scheduler; + _executionContext = executionContext; + } + } + + internal static class ManualResetValueTaskSourceCoreShared // separated out of generic to avoid unnecessary duplication + { + internal static readonly Action s_sentinel = CompletionSentinel; + + private static void CompletionSentinel(object? _) // named method to aid debugging + { + Debug.Fail("The sentinel delegate should never be invoked."); + ThrowHelper.ThrowInvalidOperationException(); + } + + internal static void ScheduleCapturedContext(object context, Action continuation, object? state) + { + Debug.Assert( + context is SynchronizationContext or TaskScheduler or CapturedSchedulerAndExecutionContext, + $"{nameof(context)} is {context}"); + + switch (context) { - InvokeContinuationWithContext(); + case SynchronizationContext sc: + ScheduleSynchronizationContext(sc, continuation, state); + break; + + case TaskScheduler ts: + ScheduleTaskScheduler(ts, continuation, state); + break; + + default: + CapturedSchedulerAndExecutionContext cc = (CapturedSchedulerAndExecutionContext)context; + if (cc._scheduler is SynchronizationContext ccsc) + { + ScheduleSynchronizationContext(ccsc, continuation, state); + } + else + { + Debug.Assert(cc._scheduler is TaskScheduler, $"{nameof(cc._scheduler)} is {cc._scheduler}"); + ScheduleTaskScheduler((TaskScheduler)cc._scheduler, continuation, state); + } + break; } + + static void ScheduleSynchronizationContext(SynchronizationContext sc, Action continuation, object? state) => + sc.Post(continuation.Invoke, state); + + static void ScheduleTaskScheduler(TaskScheduler scheduler, Action continuation, object? state) => + Task.Factory.StartNew(continuation, state, CancellationToken.None, TaskCreationOptions.DenyChildAttach, scheduler); } - private void InvokeContinuationWithContext() + internal static void InvokeContinuationWithContext(object capturedContext, Action continuation, object? continuationState, bool runContinuationsAsynchronously) { // This is in a helper as the error handling causes the generated asm // for the surrounding code to become less efficient (stack spills etc) // and it is an uncommon path. + Debug.Assert(continuation is not null, $"{nameof(continuation)} is null"); + Debug.Assert(capturedContext is ExecutionContext or CapturedSchedulerAndExecutionContext, $"{nameof(capturedContext)} is {capturedContext}"); - Debug.Assert(_continuation != null, $"Null {nameof(_continuation)}"); - Debug.Assert(_executionContext != null, $"Null {nameof(_executionContext)}"); - + // Capture the current EC. We'll switch over to the target EC and then restore back to this one. ExecutionContext? currentContext = ExecutionContext.CaptureForRestore(); - // Restore the captured ExecutionContext before executing anything. - ExecutionContext.Restore(_executionContext); - if (_capturedContext is null) + if (capturedContext is ExecutionContext ec) { - if (RunContinuationsAsynchronously) + ExecutionContext.RestoreInternal(ec); // Restore the captured ExecutionContext before executing anything. + if (runContinuationsAsynchronously) { try { - ThreadPool.QueueUserWorkItem(_continuation, _continuationState, preferLocal: true); + ThreadPool.QueueUserWorkItem(continuation, continuationState, preferLocal: true); } finally { - // Restore the current ExecutionContext. - ExecutionContext.RestoreInternal(currentContext); + ExecutionContext.RestoreInternal(currentContext); // Restore the current ExecutionContext. } } else @@ -279,7 +344,7 @@ private void InvokeContinuationWithContext() SynchronizationContext? syncContext = SynchronizationContext.Current; try { - _continuation(_continuationState); + continuation(continuationState); } catch (Exception ex) { @@ -290,64 +355,29 @@ private void InvokeContinuationWithContext() } finally { - // Set sync context back to what it was prior to coming in + // Set sync context back to what it was prior to coming in. + // Then restore the current ExecutionContext. SynchronizationContext.SetSynchronizationContext(syncContext); - // Restore the current ExecutionContext. ExecutionContext.RestoreInternal(currentContext); } // Now rethrow the exception; if there is one. edi?.Throw(); } - - return; - } - - try - { - InvokeSchedulerContinuation(); } - finally - { - // Restore the current ExecutionContext. - ExecutionContext.RestoreInternal(currentContext); - } - } - - /// - /// Invokes the continuation with the appropriate scheduler. - /// This assumes that if is not null we're already - /// running within that . - /// - private void InvokeSchedulerContinuation() - { - Debug.Assert(_capturedContext != null, $"Null {nameof(_capturedContext)}"); - Debug.Assert(_continuation != null, $"Null {nameof(_continuation)}"); - - switch (_capturedContext) + else { - case SynchronizationContext sc: - sc.Post(static s => - { - var state = (TupleSlim, object?>)s!; - state.Item1(state.Item2); - }, new TupleSlim, object?>(_continuation, _continuationState)); - break; - - case TaskScheduler ts: - Task.Factory.StartNew(_continuation, _continuationState, CancellationToken.None, TaskCreationOptions.DenyChildAttach, ts); - break; + CapturedSchedulerAndExecutionContext cc = (CapturedSchedulerAndExecutionContext)capturedContext; + ExecutionContext.Restore(cc._executionContext); // Restore the captured ExecutionContext before executing anything. + try + { + ScheduleCapturedContext(capturedContext, continuation, continuationState); + } + finally + { + ExecutionContext.RestoreInternal(currentContext); // Restore the current ExecutionContext. + } } } } - - internal static class ManualResetValueTaskSourceCoreShared // separated out of generic to avoid unnecessary duplication - { - internal static readonly Action s_sentinel = CompletionSentinel; - private static void CompletionSentinel(object? _) // named method to aid debugging - { - Debug.Fail("The sentinel delegate should never be invoked."); - ThrowHelper.ThrowInvalidOperationException(); - } - } } diff --git a/src/libraries/System.Private.CoreLib/src/System/ThrowHelper.cs b/src/libraries/System.Private.CoreLib/src/System/ThrowHelper.cs index 59539d2eb06eb..3382fa9571d71 100644 --- a/src/libraries/System.Private.CoreLib/src/System/ThrowHelper.cs +++ b/src/libraries/System.Private.CoreLib/src/System/ThrowHelper.cs @@ -818,6 +818,8 @@ private static string GetArgumentName(ExceptionArgument argument) return "function"; case ExceptionArgument.scheduler: return "scheduler"; + case ExceptionArgument.continuation: + return "continuation"; case ExceptionArgument.continuationAction: return "continuationAction"; case ExceptionArgument.continuationFunction: @@ -1142,6 +1144,7 @@ internal enum ExceptionArgument creationOptions, function, scheduler, + continuation, continuationAction, continuationFunction, tasks,