From c85cc721b294c38cd52f5ac364a71d54efa58b68 Mon Sep 17 00:00:00 2001 From: Stephen Toub Date: Fri, 18 Oct 2019 17:02:04 -0400 Subject: [PATCH] Enable pooling for `async ValueTask/ValueTask` methods Today `async ValueTask/ValueTask` methods use builders that special-case the synchronously completing case (to just return a `default(ValueTask)` or `new ValueTask(result))` but defer to the equivalent of `async Task/Task` for when the operation completes asynchronously. This, however, doesn't take advantage of the fact that value tasks can wrap arbitrary `IValueTaskSource/IValueTaskSource` implementations. This commit gives `AsyncValueTaskMethodBuilder` and `AsyncValueTaskMethodBuilder` the ability to use pooled `IValueTaskSource/IValueTaskSource` instances, such that calls to an `async ValueTask/ValueTask` method incur 0 allocations (ammortized) as long as there's a cached object available. Currently the pooling is behind a feature flag, requiring opt-in via the DOTNET_SYSTEM_THREADING_POOLASYNCVALUETASKS environment variable (setting it to "true" or "1"). This is done for a few reasons: - There's a breaking change here, in that while we say/document that `ValueTask/ValueTask`s are more limited in what they can be used for, nothing in the implementation actually stops a `ValueTask` that was wrapping a `Task` from being used as permissively as `Task`, e.g. if you were to `await` such a `ValueTask` multiple times, it would happen to work, even though we say "never do that". This change means that anyone who was relying on such undocumented behaviors will now be broken. I think this is a reasonable thing to do in a major release, but I also want feedback and a lot of runway on it. - Policy around pooling. Pooling is always a difficult thing to tune. Right now I've chosen a policy that limits the number of pooled objects per state machine type to an arbitrary multiple of the processor count, and that tries to strike a balance between contention and garbage by using a spin lock and if there's any contention while trying to get or return a pooled object, the cache is ignored. We will need to think hard about what policy to use here. It's also possible it could be tuned per state machine, e.g. by having an attribute that's looked up via reflection when creating the cache for a state machine, but that'll add a lot of first-access overhead to any `async ValueTask/ValueTask` method. For now, it's tunable via the `DOTNET_SYSTEM_THREADING_POOLASYNCVALUETASKSLIMIT` environment variable, which may be set to the maximum number of pooled objects desired per state machine. - Performance validation. Preliminary numbers show that this accomplishes its goal, having on-par throughput with the current implementation but with significantly less allocation. That needs to be validated at scale and across a variety of workloads. - Diagnostics. There are several diagnostics-related abilities available for `async Task/Task` methods that are not included here when using `async ValueTask/ValueTask` when pooling. We need to evaluate these (e.g. tracing) and determine which pieces need to be enabled and which we're fine omitting. Before shipping .NET 5, we could choose to flip the polarity of the switch (making it opt-out rather than opt-in), remove the fallback altogether and just have it be always on, or revert this change, all based on experimentation and data we receive between now and then. --- .../AsyncMethodBuilderCore.cs | 2 +- .../CompilerServices/AsyncTaskCache.cs | 24 +- .../AsyncTaskMethodBuilder.cs | 54 +- .../AsyncTaskMethodBuilderT.cs | 121 ++--- .../AsyncValueTaskMethodBuilder.cs | 123 ++++- .../AsyncValueTaskMethodBuilderT.cs | 468 ++++++++++++++++-- .../shared/System/Threading/Tasks/Task.cs | 1 - 7 files changed, 654 insertions(+), 139 deletions(-) diff --git a/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncMethodBuilderCore.cs b/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncMethodBuilderCore.cs index 214ee72c9b5b..f8aad2f01788 100644 --- a/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncMethodBuilderCore.cs +++ b/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncMethodBuilderCore.cs @@ -62,7 +62,7 @@ public static void Start(ref TStateMachine stateMachine) where TS } } - public static void SetStateMachine(IAsyncStateMachine stateMachine, Task task) + public static void SetStateMachine(IAsyncStateMachine stateMachine, Task? task) { if (stateMachine == null) { diff --git a/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncTaskCache.cs b/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncTaskCache.cs index 9da2fc0d18de..3f164a02c437 100644 --- a/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncTaskCache.cs +++ b/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncTaskCache.cs @@ -12,9 +12,9 @@ namespace System.Runtime.CompilerServices internal static class AsyncTaskCache { /// A cached Task{Boolean}.Result == true. - internal static readonly Task s_trueTask = CreateCacheableTask(true); + internal static readonly Task s_trueTask = CreateCacheableTask(result: true); /// A cached Task{Boolean}.Result == false. - internal static readonly Task s_falseTask = CreateCacheableTask(false); + internal static readonly Task s_falseTask = CreateCacheableTask(result: false); /// The cache of Task{Int32}. internal static readonly Task[] s_int32Tasks = CreateInt32Tasks(); /// The minimum value, inclusive, for which we want a cached task. @@ -22,6 +22,26 @@ internal static class AsyncTaskCache /// The maximum value, exclusive, for which we want a cached task. internal const int ExclusiveInt32Max = 9; + /// true if we should use reusable boxes for async completions of ValueTask methods; false if we should use tasks. + /// + /// We rely on tiered compilation turning this into a const and doing dead code elimination to make checks on this efficient. + /// It's also required for safety that this value never changes once observed, as Unsafe.As casts are employed based on its value. + /// + internal static readonly bool s_valueTaskPoolingEnabled = GetPoolAsyncValueTasksSwitch(); + /// Maximum number of boxes that are allowed to be cached per state machine type. + internal static readonly int s_valueTaskPoolingCacheSize = GetPoolAsyncValueTasksLimitValue(); + + private static bool GetPoolAsyncValueTasksSwitch() + { + string? value = Environment.GetEnvironmentVariable("DOTNET_SYSTEM_THREADING_POOLASYNCVALUETASKS"); + return value != null && (bool.IsTrueStringIgnoreCase(value) || value.Equals("1")); + } + + private static int GetPoolAsyncValueTasksLimitValue() => + int.TryParse(Environment.GetEnvironmentVariable("DOTNET_SYSTEM_THREADING_POOLASYNCVALUETASKSLIMIT"), out int result) && result > 0 ? + result : + Environment.ProcessorCount * 4; // arbitrary default value + /// Creates a non-disposable task. /// Specifies the result type. /// The result for the task. diff --git a/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncTaskMethodBuilder.cs b/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncTaskMethodBuilder.cs index 83f00fe2578c..e13ea03ac07c 100644 --- a/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncTaskMethodBuilder.cs +++ b/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncTaskMethodBuilder.cs @@ -21,15 +21,12 @@ public struct AsyncTaskMethodBuilder /// A cached VoidTaskResult task used for builders that complete synchronously. private static readonly Task s_cachedCompleted = AsyncTaskMethodBuilder.s_defaultResultTask; - /// The generic builder object to which this non-generic instance delegates. - private AsyncTaskMethodBuilder m_builder; // mutable struct: must not be readonly. Debugger depends on the exact name of this field. + /// The lazily-initialized built task. + private Task? m_task; // Debugger depends on the exact name of this field. /// Initializes a new . /// The initialized . - public static AsyncTaskMethodBuilder Create() => - // m_builder should be initialized to AsyncTaskMethodBuilder.Create(), but on coreclr - // that Create() is a nop, so we can just return the default here. - default; + public static AsyncTaskMethodBuilder Create() => default; /// Initiates the builder's execution with the associated state machine. /// Specifies the type of the state machine. @@ -44,7 +41,7 @@ public void Start(ref TStateMachine stateMachine) where TStateMac /// The argument was null (Nothing in Visual Basic). /// The builder is incorrectly initialized. public void SetStateMachine(IAsyncStateMachine stateMachine) => - m_builder.SetStateMachine(stateMachine); + AsyncMethodBuilderCore.SetStateMachine(stateMachine, task: null); /// /// Schedules the specified state machine to be pushed forward when the specified awaiter completes. @@ -57,7 +54,7 @@ public void AwaitOnCompleted( ref TAwaiter awaiter, ref TStateMachine stateMachine) where TAwaiter : INotifyCompletion where TStateMachine : IAsyncStateMachine => - m_builder.AwaitOnCompleted(ref awaiter, ref stateMachine); + AsyncTaskMethodBuilder.AwaitOnCompleted(ref awaiter, ref stateMachine, ref m_task); /// /// Schedules the specified state machine to be pushed forward when the specified awaiter completes. @@ -66,11 +63,12 @@ public void AwaitOnCompleted( /// Specifies the type of the state machine. /// The awaiter. /// The state machine. + [MethodImpl(MethodImplOptions.AggressiveInlining)] public void AwaitUnsafeOnCompleted( ref TAwaiter awaiter, ref TStateMachine stateMachine) where TAwaiter : ICriticalNotifyCompletion where TStateMachine : IAsyncStateMachine => - m_builder.AwaitUnsafeOnCompleted(ref awaiter, ref stateMachine); + AsyncTaskMethodBuilder.AwaitUnsafeOnCompleted(ref awaiter, ref stateMachine, ref m_task); /// Gets the for this builder. /// The representing the builder's asynchronous operation. @@ -78,7 +76,19 @@ public void AwaitUnsafeOnCompleted( public Task Task { [MethodImpl(MethodImplOptions.AggressiveInlining)] - get => m_builder.Task; + get => m_task ?? InitializeTaskAsPromise(); + } + + /// + /// Initializes the task, which must not yet be initialized. Used only when the Task is being forced into + /// existence when no state machine is needed, e.g. when the builder is being synchronously completed with + /// an exception, when the builder is being used out of the context of an async method, etc. + /// + [MethodImpl(MethodImplOptions.NoInlining)] + private Task InitializeTaskAsPromise() + { + Debug.Assert(m_task == null); + return m_task = new Task(); } /// @@ -87,7 +97,20 @@ public Task Task /// /// The builder is not initialized. /// The task has already completed. - public void SetResult() => m_builder.SetResult(s_cachedCompleted); // Using s_cachedCompleted is faster than using s_defaultResultTask. + public void SetResult() + { + // Get the currently stored task, which will be non-null if get_Task has already been accessed. + // If there isn't one, store the supplied completed task. + if (m_task is null) + { + m_task = s_cachedCompleted; + } + else + { + // Otherwise, complete the task that's there. + AsyncTaskMethodBuilder.SetExistingTaskResult(m_task, default!); + } + } /// /// Completes the in the @@ -97,7 +120,8 @@ public Task Task /// The argument is null (Nothing in Visual Basic). /// The builder is not initialized. /// The task has already completed. - public void SetException(Exception exception) => m_builder.SetException(exception); + public void SetException(Exception exception) => + AsyncTaskMethodBuilder.SetException(exception, ref m_task); /// /// Called by the debugger to request notification when the first wait operation @@ -106,7 +130,8 @@ public Task Task /// /// true to enable notification; false to disable a previously set notification. /// - internal void SetNotificationForWaitCompletion(bool enabled) => m_builder.SetNotificationForWaitCompletion(enabled); + internal void SetNotificationForWaitCompletion(bool enabled) => + AsyncTaskMethodBuilder.SetNotificationForWaitCompletion(enabled, ref m_task); /// /// Gets an object that may be used to uniquely identify this builder to the debugger. @@ -116,6 +141,7 @@ public Task Task /// It must only be used by the debugger and tracing purposes, and only in a single-threaded manner /// when no other threads are in the middle of accessing this property or this.Task. /// - internal object ObjectIdForDebugger => m_builder.ObjectIdForDebugger; + internal object ObjectIdForDebugger => + m_task ??= AsyncTaskMethodBuilder.CreateWeaklyTypedStateMachineBox(); } } diff --git a/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncTaskMethodBuilderT.cs b/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncTaskMethodBuilderT.cs index 918358a06e19..e577b26c2a75 100644 --- a/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncTaskMethodBuilderT.cs +++ b/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncTaskMethodBuilderT.cs @@ -3,10 +3,10 @@ // See the LICENSE file in the project root for more information. using System.Diagnostics; +using System.Diagnostics.CodeAnalysis; using System.Threading; using System.Threading.Tasks; using Internal.Runtime.CompilerServices; -using System.Diagnostics.CodeAnalysis; namespace System.Runtime.CompilerServices { @@ -25,17 +25,11 @@ public struct AsyncTaskMethodBuilder internal static readonly Task s_defaultResultTask = AsyncTaskCache.CreateCacheableTask(default); /// The lazily-initialized built task. - private Task m_task; // lazily-initialized: must not be readonly. Debugger depends on the exact name of this field. + private Task? m_task; // Debugger depends on the exact name of this field. /// Initializes a new . /// The initialized . - public static AsyncTaskMethodBuilder Create() - { - // NOTE: If this method is ever updated to perform more initialization, - // other Create methods like AsyncTaskMethodBuilder.Create and - // AsyncValueTaskMethodBuilder.Create must be updated to call this. - return default; - } + public static AsyncTaskMethodBuilder Create() => default; /// Initiates the builder's execution with the associated state machine. /// Specifies the type of the state machine. @@ -49,8 +43,8 @@ public void Start(ref TStateMachine stateMachine) where TStateMac /// The heap-allocated state machine object. /// The argument was null (Nothing in Visual Basic). /// The builder is incorrectly initialized. - public void SetStateMachine(IAsyncStateMachine stateMachine) - => AsyncMethodBuilderCore.SetStateMachine(stateMachine, m_task); + public void SetStateMachine(IAsyncStateMachine stateMachine) => + AsyncMethodBuilderCore.SetStateMachine(stateMachine, m_task); /// /// Schedules the specified state machine to be pushed forward when the specified awaiter completes. @@ -62,11 +56,17 @@ public void SetStateMachine(IAsyncStateMachine stateMachine) public void AwaitOnCompleted( ref TAwaiter awaiter, ref TStateMachine stateMachine) where TAwaiter : INotifyCompletion + where TStateMachine : IAsyncStateMachine => + AwaitOnCompleted(ref awaiter, ref stateMachine, ref m_task); + + internal static void AwaitOnCompleted( + ref TAwaiter awaiter, ref TStateMachine stateMachine, [NotNull] ref Task? taskField) + where TAwaiter : INotifyCompletion where TStateMachine : IAsyncStateMachine { try { - awaiter.OnCompleted(GetStateMachineBox(ref stateMachine).MoveNextAction); + awaiter.OnCompleted(GetStateMachineBox(ref stateMachine, ref taskField).MoveNextAction); } catch (Exception e) { @@ -81,15 +81,28 @@ public void AwaitOnCompleted( /// Specifies the type of the state machine. /// The awaiter. /// The state machine. - // AggressiveOptimization to workaround boxing allocations in Tier0 until: https://github.com/dotnet/coreclr/issues/14474 - [MethodImpl(MethodImplOptions.AggressiveOptimization)] + [MethodImpl(MethodImplOptions.AggressiveInlining)] public void AwaitUnsafeOnCompleted( ref TAwaiter awaiter, ref TStateMachine stateMachine) where TAwaiter : ICriticalNotifyCompletion + where TStateMachine : IAsyncStateMachine => + AwaitUnsafeOnCompleted(ref awaiter, ref stateMachine, ref m_task); + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + internal static void AwaitUnsafeOnCompleted( + ref TAwaiter awaiter, ref TStateMachine stateMachine, [NotNull] ref Task? taskField) + where TAwaiter : ICriticalNotifyCompletion where TStateMachine : IAsyncStateMachine { - IAsyncStateMachineBox box = GetStateMachineBox(ref stateMachine); + IAsyncStateMachineBox box = GetStateMachineBox(ref stateMachine, ref taskField); + AwaitUnsafeOnCompleted(ref awaiter, box); + } + [MethodImpl(MethodImplOptions.AggressiveOptimization)] // workaround boxing allocations in Tier0: https://github.com/dotnet/coreclr/issues/14474 + internal static void AwaitUnsafeOnCompleted( + ref TAwaiter awaiter, IAsyncStateMachineBox box) + where TAwaiter : ICriticalNotifyCompletion + { // The null tests here ensure that the jit can optimize away the interface // tests when TAwaiter is a ref type. @@ -139,8 +152,9 @@ public void AwaitUnsafeOnCompleted( /// Specifies the type of the async state machine. /// The state machine. /// The "boxed" state machine. - private IAsyncStateMachineBox GetStateMachineBox( - ref TStateMachine stateMachine) + private static IAsyncStateMachineBox GetStateMachineBox( + ref TStateMachine stateMachine, + [NotNull] ref Task? taskField) where TStateMachine : IAsyncStateMachine { ExecutionContext? currentContext = ExecutionContext.Capture(); @@ -150,7 +164,7 @@ private IAsyncStateMachineBox GetStateMachineBox( // a strongly-typed manner into an AsyncStateMachineBox. It will already contain // the state machine as well as a MoveNextDelegate and a context. The only thing // we might need to do is update the context if that's changed since it was stored. - if (m_task is AsyncStateMachineBox stronglyTypedBox) + if (taskField is AsyncStateMachineBox stronglyTypedBox) { if (stronglyTypedBox.Context != currentContext) { @@ -168,7 +182,7 @@ private IAsyncStateMachineBox GetStateMachineBox( // result in a boxing allocation when storing the TStateMachine if it's a struct, but // this only happens in active debugging scenarios where such performance impact doesn't // matter. - if (m_task is AsyncStateMachineBox weaklyTypedBox) + if (taskField is AsyncStateMachineBox weaklyTypedBox) { // If this is the first await, we won't yet have a state machine, so store it. if (weaklyTypedBox.StateMachine == null) @@ -191,7 +205,7 @@ private IAsyncStateMachineBox GetStateMachineBox( // multiple callbacks registered to push the state machine, which could result in bad behavior. Debugger.NotifyOfCrossThreadDependency(); - // At this point, m_task should really be null, in which case we want to create the box. + // At this point, taskField should really be null, in which case we want to create the box. // However, in a variety of debugger-related (erroneous) situations, it might be non-null, // e.g. if the Task property is examined in a Watch window, forcing it to be lazily-intialized // as a Task rather than as an AsyncStateMachineBox. The worst that happens in such @@ -209,7 +223,7 @@ private IAsyncStateMachineBox GetStateMachineBox( CreateDebugFinalizableAsyncStateMachineBox() : new AsyncStateMachineBox(); #endif - m_task = box; // important: this must be done before storing stateMachine into box.StateMachine! + taskField = box; // important: this must be done before storing stateMachine into box.StateMachine! box.StateMachine = stateMachine; box.Context = currentContext; @@ -374,24 +388,16 @@ private Task InitializeTaskAsPromise() return m_task = new Task(); } - /// - /// Initializes the task, which must not yet be initialized. Used only when the Task is being forced into - /// existence due to the debugger trying to enable step-out/step-over/etc. prior to the first await yielding - /// in an async method. In that case, we don't know the actual TStateMachine type, so we're forced to - /// use IAsyncStateMachine instead. - /// - [MethodImpl(MethodImplOptions.NoInlining)] - private Task InitializeTaskAsStateMachineBox() + internal static Task CreateWeaklyTypedStateMachineBox() { - Debug.Assert(m_task == null); #if CORERT // DebugFinalizableAsyncStateMachineBox looks like a small type, but it actually is not because // it will have a copy of all the slots from its parent. It will add another hundred(s) bytes // per each async method in CoreRT / ProjectN binaries without adding much value. Avoid // generating this extra code until a better solution is implemented. - return (m_task = new AsyncStateMachineBox()); + return new AsyncStateMachineBox()); #else - return m_task = AsyncMethodBuilderCore.TrackAsyncMethodCompletion ? + return AsyncMethodBuilderCore.TrackAsyncMethodCompletion ? CreateDebugFinalizableAsyncStateMachineBox() : new AsyncStateMachineBox(); #endif @@ -407,7 +413,7 @@ public void SetResult(TResult result) { // Get the currently stored task, which will be non-null if get_Task has already been accessed. // If there isn't one, get a task and store it. - if (m_task == null) + if (m_task is null) { m_task = GetTaskForResult(result); Debug.Assert(m_task != null, $"{nameof(GetTaskForResult)} should never return null"); @@ -415,51 +421,27 @@ public void SetResult(TResult result) else { // Slow path: complete the existing task. - SetExistingTaskResult(result); + SetExistingTaskResult(m_task, result); } } /// Completes the already initialized task with the specified result. /// The result to use to complete the task. - private void SetExistingTaskResult([AllowNull] TResult result) + internal static void SetExistingTaskResult(Task taskField, [AllowNull] TResult result) { - Debug.Assert(m_task != null, "Expected non-null task"); + Debug.Assert(taskField != null, "Expected non-null task"); if (AsyncCausalityTracer.LoggingOn) { - AsyncCausalityTracer.TraceOperationCompletion(m_task, AsyncCausalityStatus.Completed); + AsyncCausalityTracer.TraceOperationCompletion(taskField, AsyncCausalityStatus.Completed); } - if (!m_task.TrySetResult(result)) + if (!taskField.TrySetResult(result)) { ThrowHelper.ThrowInvalidOperationException(ExceptionResource.TaskT_TransitionToFinal_AlreadyCompleted); } } - /// - /// Completes the builder by using either the supplied completed task, or by completing - /// the builder's previously accessed task using default(TResult). - /// - /// A task already completed with the value default(TResult). - /// The task has already completed. - internal void SetResult(Task completedTask) - { - Debug.Assert(completedTask != null, "Expected non-null task"); - Debug.Assert(completedTask.IsCompletedSuccessfully, "Expected a successfully completed task"); - - // Get the currently stored task, which will be non-null if get_Task has already been accessed. - // If there isn't one, store the supplied completed task. - if (m_task == null) - { - m_task = completedTask; - } - else - { - // Otherwise, complete the task that's there. - SetExistingTaskResult(default!); // Remove ! when nullable attributes are respected - } - } - /// /// Completes the in the /// Faulted state with the specified exception. @@ -467,7 +449,9 @@ internal void SetResult(Task completedTask) /// The to use to fault the task. /// The argument is null (Nothing in Visual Basic). /// The task has already completed. - public void SetException(Exception exception) + public void SetException(Exception exception) => SetException(exception, ref m_task); + + internal static void SetException(Exception exception, ref Task? taskField) { if (exception == null) { @@ -475,7 +459,7 @@ public void SetException(Exception exception) } // Get the task, forcing initialization if it hasn't already been initialized. - Task task = this.Task; + Task task = (taskField ??= new Task()); // If the exception represents cancellation, cancel the task. Otherwise, fault the task. bool successfullySet = exception is OperationCanceledException oce ? @@ -505,10 +489,13 @@ public void SetException(Exception exception) /// This should only be invoked from within an asynchronous method, /// and only by the debugger. /// - internal void SetNotificationForWaitCompletion(bool enabled) + internal void SetNotificationForWaitCompletion(bool enabled) => + SetNotificationForWaitCompletion(enabled, ref m_task); + + internal static void SetNotificationForWaitCompletion(bool enabled, [NotNull] ref Task? taskField) { // Get the task (forcing initialization if not already initialized), and set debug notification - (m_task ?? InitializeTaskAsStateMachineBox()).SetNotificationForWaitCompletion(enabled); + (taskField ??= CreateWeaklyTypedStateMachineBox()).SetNotificationForWaitCompletion(enabled); // NOTE: It's important that the debugger use builder.SetNotificationForWaitCompletion // rather than builder.Task.SetNotificationForWaitCompletion. Even though the latter will @@ -529,7 +516,7 @@ internal void SetNotificationForWaitCompletion(bool enabled) /// It must only be used by the debugger and tracing purposes, and only in a single-threaded manner /// when no other threads are in the middle of accessing this or other members that lazily initialize the task. /// - internal object ObjectIdForDebugger => m_task ?? InitializeTaskAsStateMachineBox(); + internal object ObjectIdForDebugger => m_task ??= CreateWeaklyTypedStateMachineBox(); /// /// Gets a task for the specified result. This will either diff --git a/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncValueTaskMethodBuilder.cs b/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncValueTaskMethodBuilder.cs index f5ea29bb00c7..bb5bd28393c1 100644 --- a/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncValueTaskMethodBuilder.cs +++ b/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncValueTaskMethodBuilder.cs @@ -4,6 +4,9 @@ using System.Runtime.InteropServices; using System.Threading.Tasks; +using Internal.Runtime.CompilerServices; + +using StateMachineBox = System.Runtime.CompilerServices.AsyncValueTaskMethodBuilder.StateMachineBox; namespace System.Runtime.CompilerServices { @@ -11,62 +14,98 @@ namespace System.Runtime.CompilerServices [StructLayout(LayoutKind.Auto)] public struct AsyncValueTaskMethodBuilder { - /// The to which most operations are delegated. - private AsyncTaskMethodBuilder _methodBuilder; // mutable struct; do not make it readonly - /// true if completed synchronously and successfully; otherwise, false. - private bool _haveResult; - /// true if the builder should be used for setting/getting the result; otherwise, false. - private bool _useBuilder; + /// Sentinel object used to indicate that the builder completed synchronously and successfully. + private static readonly object s_syncSuccessSentinel = AsyncValueTaskMethodBuilder.s_syncSuccessSentinel; + + /// The wrapped state machine box or task, based on the value of . + /// + /// If the operation completed synchronously and successfully, this will be . + /// + private object? m_task; // Debugger depends on the exact name of this field. /// Creates an instance of the struct. /// The initialized instance. - public static AsyncValueTaskMethodBuilder Create() => - // _methodBuilder should be initialized to AsyncTaskMethodBuilder.Create(), but on coreclr - // that Create() is a nop, so we can just return the default here. - default; + public static AsyncValueTaskMethodBuilder Create() => default; /// Begins running the builder with the associated state machine. /// The type of the state machine. /// The state machine instance, passed by reference. [MethodImpl(MethodImplOptions.AggressiveInlining)] - public void Start(ref TStateMachine stateMachine) where TStateMachine : IAsyncStateMachine => - // will provide the right ExecutionContext semantics + public void Start(ref TStateMachine stateMachine) + where TStateMachine : IAsyncStateMachine => AsyncMethodBuilderCore.Start(ref stateMachine); /// Associates the builder with the specified state machine. /// The state machine instance to associate with the builder. - public void SetStateMachine(IAsyncStateMachine stateMachine) => _methodBuilder.SetStateMachine(stateMachine); + public void SetStateMachine(IAsyncStateMachine stateMachine) => + AsyncMethodBuilderCore.SetStateMachine(stateMachine, task: null); /// Marks the task as successfully completed. public void SetResult() { - if (_useBuilder) + if (m_task is null) { - _methodBuilder.SetResult(); + m_task = s_syncSuccessSentinel; + } + else if (AsyncTaskCache.s_valueTaskPoolingEnabled) + { + Unsafe.As(m_task).SetResult(default); } else { - _haveResult = true; + AsyncTaskMethodBuilder.SetExistingTaskResult(Unsafe.As>(m_task), default); } } /// Marks the task as failed and binds the specified exception to the task. /// The exception to bind to the task. - public void SetException(Exception exception) => _methodBuilder.SetException(exception); + public void SetException(Exception exception) + { + if (AsyncTaskCache.s_valueTaskPoolingEnabled) + { + AsyncValueTaskMethodBuilder.SetException(exception, ref Unsafe.As(ref m_task)); + } + else + { + AsyncTaskMethodBuilder.SetException(exception, ref Unsafe.As?>(ref m_task)); + } + } /// Gets the task for this builder. public ValueTask Task { get { - if (_haveResult) + if (m_task == s_syncSuccessSentinel) { return default; } + + // With normal access paterns, m_task should always be non-null here: the async method should have + // either completed synchronously, in which case SetResult would have set m_task to a non-null object, + // or it should be completing asynchronously, in which case AwaitUnsafeOnCompleted would have similarly + // initialized m_task to a state machine object. However, if the type is used manually (not via + // compiler-generated code) and accesses Task directly, we force it to be initialized. Things will then + // "work" but in a degraded mode, as we don't know the TStateMachine type here, and thus we use a box around + // the interface instead. + + if (AsyncTaskCache.s_valueTaskPoolingEnabled) + { + var box = Unsafe.As(m_task); + if (box is null) + { + m_task = box = AsyncValueTaskMethodBuilder.CreateWeaklyTypedStateMachineBox(); + } + return new ValueTask(box, box.Version); + } else { - _useBuilder = true; - return new ValueTask(_methodBuilder.Task); + var task = Unsafe.As?>(m_task); + if (task is null) + { + m_task = task = new Task(); // base task used rather than box to minimize size when used as manual promise + } + return new ValueTask(task); } } } @@ -80,8 +119,14 @@ public void AwaitOnCompleted(ref TAwaiter awaiter, ref where TAwaiter : INotifyCompletion where TStateMachine : IAsyncStateMachine { - _useBuilder = true; - _methodBuilder.AwaitOnCompleted(ref awaiter, ref stateMachine); + if (AsyncTaskCache.s_valueTaskPoolingEnabled) + { + AsyncValueTaskMethodBuilder.AwaitOnCompleted(ref awaiter, ref stateMachine, ref Unsafe.As(ref m_task)); + } + else + { + AsyncTaskMethodBuilder.AwaitOnCompleted(ref awaiter, ref stateMachine, ref Unsafe.As?>(ref m_task)); + } } /// Schedules the state machine to proceed to the next action when the specified awaiter completes. @@ -89,12 +134,42 @@ public void AwaitOnCompleted(ref TAwaiter awaiter, ref /// The type of the state machine. /// The awaiter. /// The state machine. + [MethodImpl(MethodImplOptions.AggressiveInlining)] public void AwaitUnsafeOnCompleted(ref TAwaiter awaiter, ref TStateMachine stateMachine) where TAwaiter : ICriticalNotifyCompletion where TStateMachine : IAsyncStateMachine { - _useBuilder = true; - _methodBuilder.AwaitUnsafeOnCompleted(ref awaiter, ref stateMachine); + if (AsyncTaskCache.s_valueTaskPoolingEnabled) + { + AsyncValueTaskMethodBuilder.AwaitUnsafeOnCompleted(ref awaiter, ref stateMachine, ref Unsafe.As(ref m_task)); + } + else + { + AsyncTaskMethodBuilder.AwaitUnsafeOnCompleted(ref awaiter, ref stateMachine, ref Unsafe.As?>(ref m_task)); + } + } + + /// + /// Gets an object that may be used to uniquely identify this builder to the debugger. + /// + /// + /// This property lazily instantiates the ID in a non-thread-safe manner. + /// It must only be used by the debugger and tracing purposes, and only in a single-threaded manner + /// when no other threads are in the middle of accessing this or other members that lazily initialize the box. + /// + internal object ObjectIdForDebugger + { + get + { + if (m_task is null) + { + m_task = AsyncTaskCache.s_valueTaskPoolingEnabled ? (object) + AsyncValueTaskMethodBuilder.CreateWeaklyTypedStateMachineBox() : + AsyncTaskMethodBuilder.CreateWeaklyTypedStateMachineBox(); + } + + return m_task; + } } } } diff --git a/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncValueTaskMethodBuilderT.cs b/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncValueTaskMethodBuilderT.cs index fc5ce80737b2..d2bd7ec111a4 100644 --- a/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncValueTaskMethodBuilderT.cs +++ b/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncValueTaskMethodBuilderT.cs @@ -2,8 +2,13 @@ // The .NET Foundation licenses this file to you under the MIT license. // See the LICENSE file in the project root for more information. +using System.Diagnostics; +using System.Diagnostics.CodeAnalysis; using System.Runtime.InteropServices; +using System.Threading; using System.Threading.Tasks; +using System.Threading.Tasks.Sources; +using Internal.Runtime.CompilerServices; namespace System.Runtime.CompilerServices { @@ -12,66 +17,115 @@ namespace System.Runtime.CompilerServices [StructLayout(LayoutKind.Auto)] public struct AsyncValueTaskMethodBuilder { - /// The to which most operations are delegated. - private AsyncTaskMethodBuilder _methodBuilder; // mutable struct; do not make it readonly - /// The result for this builder, if it's completed before any awaits occur. + /// Sentinel object used to indicate that the builder completed synchronously and successfully. + /// + /// To avoid memory safety issues even in the face of invalid race conditions, we ensure that the type of this object + /// is valid for the mode in which we're operating. As such, it's cached on the generic builder per TResult + /// rather than having one sentinel instance for all types. + /// + internal static readonly object s_syncSuccessSentinel = AsyncTaskCache.s_valueTaskPoolingEnabled ? (object) + new SyncSuccessSentinelStateMachineBox() : + new Task(default(TResult)!); + + /// The wrapped state machine or task. If the operation completed synchronously and successfully, this will be a sentinel object compared by reference identity. + private object? m_task; // Debugger depends on the exact name of this field. + /// The result for this builder if it's completed synchronously, in which case will be . private TResult _result; - /// true if contains the synchronous result for the async method; otherwise, false. - private bool _haveResult; - /// true if the builder should be used for setting/getting the result; otherwise, false. - private bool _useBuilder; /// Creates an instance of the struct. /// The initialized instance. - public static AsyncValueTaskMethodBuilder Create() => - // _methodBuilder should be initialized to AsyncTaskMethodBuilder.Create(), but on coreclr - // that Create() is a nop, so we can just return the default here. - default; + public static AsyncValueTaskMethodBuilder Create() => default; /// Begins running the builder with the associated state machine. /// The type of the state machine. /// The state machine instance, passed by reference. [MethodImpl(MethodImplOptions.AggressiveInlining)] public void Start(ref TStateMachine stateMachine) where TStateMachine : IAsyncStateMachine => - // will provide the right ExecutionContext semantics AsyncMethodBuilderCore.Start(ref stateMachine); /// Associates the builder with the specified state machine. /// The state machine instance to associate with the builder. - public void SetStateMachine(IAsyncStateMachine stateMachine) => _methodBuilder.SetStateMachine(stateMachine); + public void SetStateMachine(IAsyncStateMachine stateMachine) => + AsyncMethodBuilderCore.SetStateMachine(stateMachine, task: null); - /// Marks the task as successfully completed. - /// The result to use to complete the task. + /// Marks the value task as successfully completed. + /// The result to use to complete the value task. public void SetResult(TResult result) { - if (_useBuilder) + if (m_task is null) { - _methodBuilder.SetResult(result); + _result = result; + m_task = s_syncSuccessSentinel; + } + else if (AsyncTaskCache.s_valueTaskPoolingEnabled) + { + Unsafe.As(m_task).SetResult(result); } else { - _result = result; - _haveResult = true; + AsyncTaskMethodBuilder.SetExistingTaskResult(Unsafe.As>(m_task), result); } } - /// Marks the task as failed and binds the specified exception to the task. - /// The exception to bind to the task. - public void SetException(Exception exception) => _methodBuilder.SetException(exception); + /// Marks the value task as failed and binds the specified exception to the value task. + /// The exception to bind to the value task. + public void SetException(Exception exception) + { + if (AsyncTaskCache.s_valueTaskPoolingEnabled) + { + SetException(exception, ref Unsafe.As(ref m_task)); + } + else + { + AsyncTaskMethodBuilder.SetException(exception, ref Unsafe.As?>(ref m_task)); + } + } - /// Gets the task for this builder. + internal static void SetException(Exception exception, [NotNull] ref StateMachineBox? boxFieldRef) + { + if (exception is null) + { + ThrowHelper.ThrowArgumentNullException(ExceptionArgument.exception); + } + + (boxFieldRef ??= CreateWeaklyTypedStateMachineBox()).SetException(exception); + } + + /// Gets the value task for this builder. public ValueTask Task { get { - if (_haveResult) + if (m_task == s_syncSuccessSentinel) { return new ValueTask(_result); } + + // With normal access paterns, m_task should always be non-null here: the async method should have + // either completed synchronously, in which case SetResult would have set m_task to a non-null object, + // or it should be completing asynchronously, in which case AwaitUnsafeOnCompleted would have similarly + // initialized m_task to a state machine object. However, if the type is used manually (not via + // compiler-generated code) and accesses Task directly, we force it to be initialized. Things will then + // "work" but in a degraded mode, as we don't know the TStateMachine type here, and thus we use a box around + // the interface instead. + + if (AsyncTaskCache.s_valueTaskPoolingEnabled) + { + var box = Unsafe.As(m_task); + if (box is null) + { + m_task = box = CreateWeaklyTypedStateMachineBox(); + } + return new ValueTask(box, box.Version); + } else { - _useBuilder = true; - return new ValueTask(_methodBuilder.Task); + var task = Unsafe.As?>(m_task); + if (task is null) + { + m_task = task = new Task(); // base task used rather than box to minimize size when used as manual promise + } + return new ValueTask(task); } } } @@ -85,8 +139,29 @@ public void AwaitOnCompleted(ref TAwaiter awaiter, ref where TAwaiter : INotifyCompletion where TStateMachine : IAsyncStateMachine { - _useBuilder = true; - _methodBuilder.AwaitOnCompleted(ref awaiter, ref stateMachine); + if (AsyncTaskCache.s_valueTaskPoolingEnabled) + { + AwaitOnCompleted(ref awaiter, ref stateMachine, ref Unsafe.As(ref m_task)); + } + else + { + AsyncTaskMethodBuilder.AwaitOnCompleted(ref awaiter, ref stateMachine, ref Unsafe.As?>(ref m_task)); + } + } + + internal static void AwaitOnCompleted( + ref TAwaiter awaiter, ref TStateMachine stateMachine, [NotNull] ref StateMachineBox? box) + where TAwaiter : INotifyCompletion + where TStateMachine : IAsyncStateMachine + { + try + { + awaiter.OnCompleted(GetStateMachineBox(ref stateMachine, ref box).MoveNextAction); + } + catch (Exception e) + { + System.Threading.Tasks.Task.ThrowAsync(e, targetContext: null); + } } /// Schedules the state machine to proceed to the next action when the specified awaiter completes. @@ -94,12 +169,345 @@ public void AwaitOnCompleted(ref TAwaiter awaiter, ref /// The type of the state machine. /// the awaiter /// The state machine. + [MethodImpl(MethodImplOptions.AggressiveInlining)] public void AwaitUnsafeOnCompleted(ref TAwaiter awaiter, ref TStateMachine stateMachine) where TAwaiter : ICriticalNotifyCompletion where TStateMachine : IAsyncStateMachine { - _useBuilder = true; - _methodBuilder.AwaitUnsafeOnCompleted(ref awaiter, ref stateMachine); + if (AsyncTaskCache.s_valueTaskPoolingEnabled) + { + AwaitUnsafeOnCompleted(ref awaiter, ref stateMachine, ref Unsafe.As(ref m_task)); + } + else + { + AsyncTaskMethodBuilder.AwaitUnsafeOnCompleted(ref awaiter, ref stateMachine, ref Unsafe.As?>(ref m_task)); + } + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + internal static void AwaitUnsafeOnCompleted( + ref TAwaiter awaiter, ref TStateMachine stateMachine, [NotNull] ref StateMachineBox? boxRef) + where TAwaiter : ICriticalNotifyCompletion + where TStateMachine : IAsyncStateMachine + { + IAsyncStateMachineBox box = GetStateMachineBox(ref stateMachine, ref boxRef); + AsyncTaskMethodBuilder.AwaitUnsafeOnCompleted(ref awaiter, box); + } + + /// Gets the "boxed" state machine object. + /// Specifies the type of the async state machine. + /// The state machine. + /// A reference to the field containing the initialized state machine box. + /// The "boxed" state machine. + private static IAsyncStateMachineBox GetStateMachineBox( + ref TStateMachine stateMachine, + [NotNull] ref StateMachineBox? boxFieldRef) + where TStateMachine : IAsyncStateMachine + { + ExecutionContext? currentContext = ExecutionContext.Capture(); + + // Check first for the most common case: not the first yield in an async method. + // In this case, the first yield will have already "boxed" the state machine in + // a strongly-typed manner into an AsyncStateMachineBox. It will already contain + // the state machine as well as a MoveNextDelegate and a context. The only thing + // we might need to do is update the context if that's changed since it was stored. + if (boxFieldRef is StateMachineBox stronglyTypedBox) + { + if (stronglyTypedBox.Context != currentContext) + { + stronglyTypedBox.Context = currentContext; + } + + return stronglyTypedBox; + } + + // The least common case: we have a weakly-typed boxed. This results if the debugger + // or some other use of reflection accesses a property like ObjectIdForDebugger. In + // such situations, we need to get an object to represent the builder, but we don't yet + // know the type of the state machine, and thus can't use TStateMachine. Instead, we + // use the IAsyncStateMachine interface, which all TStateMachines implement. This will + // result in a boxing allocation when storing the TStateMachine if it's a struct, but + // this only happens in active debugging scenarios where such performance impact doesn't + // matter. + if (boxFieldRef is StateMachineBox weaklyTypedBox) + { + // If this is the first await, we won't yet have a state machine, so store it. + if (weaklyTypedBox.StateMachine is null) + { + Debugger.NotifyOfCrossThreadDependency(); // same explanation as with usage below + weaklyTypedBox.StateMachine = stateMachine; + } + + // Update the context. This only happens with a debugger, so no need to spend + // extra IL checking for equality before doing the assignment. + weaklyTypedBox.Context = currentContext; + return weaklyTypedBox; + } + + // Alert a listening debugger that we can't make forward progress unless it slips threads. + // If we don't do this, and a method that uses "await foo;" is invoked through funceval, + // we could end up hooking up a callback to push forward the async method's state machine, + // the debugger would then abort the funceval after it takes too long, and then continuing + // execution could result in another callback being hooked up. At that point we have + // multiple callbacks registered to push the state machine, which could result in bad behavior. + Debugger.NotifyOfCrossThreadDependency(); + + // At this point, m_task should really be null, in which case we want to create the box. + // However, in a variety of debugger-related (erroneous) situations, it might be non-null, + // e.g. if the Task property is examined in a Watch window, forcing it to be lazily-intialized + // as a Task rather than as an ValueTaskStateMachineBox. The worst that happens in such + // cases is we lose the ability to properly step in the debugger, as the debugger uses that + // object's identity to track this specific builder/state machine. As such, we proceed to + // overwrite whatever's there anyway, even if it's non-null. + var box = StateMachineBox.GetOrCreateBox(); + boxFieldRef = box; // important: this must be done before storing stateMachine into box.StateMachine! + box.StateMachine = stateMachine; + box.Context = currentContext; + + return box; + } + + /// + /// Creates a box object for use when a non-standard access pattern is employed, e.g. when Task + /// is evaluated in the debugger prior to the async method yielding for the first time. + /// + internal static StateMachineBox CreateWeaklyTypedStateMachineBox() => new StateMachineBox(); + + /// + /// Gets an object that may be used to uniquely identify this builder to the debugger. + /// + /// + /// This property lazily instantiates the ID in a non-thread-safe manner. + /// It must only be used by the debugger and tracing purposes, and only in a single-threaded manner + /// when no other threads are in the middle of accessing this or other members that lazily initialize the box. + /// + internal object ObjectIdForDebugger + { + get + { + if (m_task is null) + { + m_task = AsyncTaskCache.s_valueTaskPoolingEnabled ? (object) + CreateWeaklyTypedStateMachineBox() : + AsyncTaskMethodBuilder.CreateWeaklyTypedStateMachineBox(); + } + + return m_task; + } + } + + /// The base type for all value task box reusable box objects, regardless of state machine type. + internal abstract class StateMachineBox : + IValueTaskSource, IValueTaskSource + { + /// A delegate to the MoveNext method. + protected Action? _moveNextAction; + /// Captured ExecutionContext with which to invoke MoveNext. + public ExecutionContext? Context; + /// Implementation for IValueTaskSource interfaces. + protected ManualResetValueTaskSourceCore _valueTaskSource; + + /// Completes the box with a result. + /// The result. + public void SetResult(TResult result) => + _valueTaskSource.SetResult(result); + + /// Completes the box with an error. + /// The exception. + public void SetException(Exception error) => + _valueTaskSource.SetException(error); + + /// Gets the status of the box. + public ValueTaskSourceStatus GetStatus(short token) => _valueTaskSource.GetStatus(token); + + /// Schedules the continuation action for this box. + public void OnCompleted(Action continuation, object? state, short token, ValueTaskSourceOnCompletedFlags flags) => + _valueTaskSource.OnCompleted(continuation, state, token, flags); + + /// Gets the current version number of the box. + public short Version => _valueTaskSource.Version; + + /// Implemented by derived type. + TResult IValueTaskSource.GetResult(short token) => throw NotImplemented.ByDesign; + + /// Implemented by derived type. + void IValueTaskSource.GetResult(short token) => throw NotImplemented.ByDesign; + } + + private sealed class SyncSuccessSentinelStateMachineBox : StateMachineBox + { + public SyncSuccessSentinelStateMachineBox() => SetResult(default!); + } + + /// Provides a strongly-typed box object based on the specific state machine type in use. + private sealed class StateMachineBox : + StateMachineBox, + IValueTaskSource, IValueTaskSource, IAsyncStateMachineBox, IThreadPoolWorkItem + where TStateMachine : IAsyncStateMachine + { + /// Delegate used to invoke on an ExecutionContext when passed an instance of this box type. + private static readonly ContextCallback s_callback = ExecutionContextCallback; + /// Lock used to protected the shared cache of boxes. + /// The code that uses this assumes a runtime without thread aborts. + private static int s_cacheLock; + /// Singly-linked list cache of boxes. + private static StateMachineBox? s_cache; + /// The number of items stored in . + private static int s_cacheSize; + + // TODO: + // AsyncTaskMethodBuilder logs about the state machine box lifecycle; AsyncValueTaskMethodBuilder currently + // does not when it employs these pooled boxes. That logging is based on Task IDs, which we lack here. + // We could use the box's Version, but that is very likely to conflict with the IDs of other tasks in the system. + // For now, we don't log, but should we choose to we'll probably want to store an int ID on the state machine box, + // and initialize it an ID from Task's generator. + + /// If this box is stored in the cache, the next box in the cache. + private StateMachineBox? _next; + /// The state machine itself. + [AllowNull, MaybeNull] + public TStateMachine StateMachine = default; + + /// Gets a box object to use for an operation. This may be a reused, pooled object, or it may be new. + [MethodImpl(MethodImplOptions.AggressiveInlining)] // only one caller + internal static StateMachineBox GetOrCreateBox() + { + // Try to acquire the lock to access the cache. If there's any contention, don't use the cache. + if (Interlocked.CompareExchange(ref s_cacheLock, 1, 0) == 0) + { + // If there are any instances cached, take one from the cache stack and use it. + StateMachineBox? box = s_cache; + if (!(box is null)) + { + s_cache = box._next; + box._next = null; + s_cacheSize--; + Debug.Assert(s_cacheSize >= 0, "Expected the cache size to be non-negative."); + + // Release the lock and return the box. + Volatile.Write(ref s_cacheLock, 0); + return box; + } + + // No objects were cached. We'll just create a new instance. + Debug.Assert(s_cacheSize == 0, "Expected cache size to be 0."); + + // Release the lock. + Volatile.Write(ref s_cacheLock, 0); + } + + // Couldn't quickly get a cached instance, so create a new instance. + return new StateMachineBox(); + } + + private void ReturnOrDropBox() + { + Debug.Assert(_next is null, "Expected box to not be part of cached list."); + + // Clear out the state machine and associated context to avoid keeping arbitrary state referenced by + // lifted locals. We want to do this regardless of whether we end up caching the box or not, in case + // the caller keeps the box alive for an arbitrary period of time. + StateMachine = default; + Context = default; + + // Reset the MRVTSC. We can either do this here, in which case we may be paying the (small) overhead + // to reset the box even if we're going to drop it, or we could do it while holding the lock, in which + // case we'll only reset it if necessary but causing the lock to be held for longer, thereby causing + // more contention. For now at least, we do it outside of the lock. (This must not be done after + // the lock is released, since at that point the instance could already be in use elsewhere.) + // We also want to increment the version number even if we're going to drop it, to maximize the chances + // that incorrectly double-awaiting a ValueTask will produce an error. + _valueTaskSource.Reset(); + + // If reusing the object would result in potentially wrapping around its version number, just throw it away. + // This provides a modicum of additional safety when ValueTasks are misused (helping to avoid the case where + // a ValueTask is illegally re-awaited and happens to do so at exactly 2^16 uses later on this exact same instance), + // at the expense of potentially incurring an additional allocation every 65K uses. + if ((ushort)_valueTaskSource.Version == ushort.MaxValue) + { + return; + } + + // Try to acquire the cache lock. If there's any contention, or if the cache is full, we just throw away the object. + if (Interlocked.CompareExchange(ref s_cacheLock, 1, 0) == 0) + { + if (s_cacheSize < AsyncTaskCache.s_valueTaskPoolingCacheSize) + { + // Push the box onto the cache stack for subsequent reuse. + _next = s_cache; + s_cache = this; + s_cacheSize++; + Debug.Assert(s_cacheSize > 0 && s_cacheSize <= AsyncTaskCache.s_valueTaskPoolingCacheSize, "Expected cache size to be within bounds."); + } + + // Release the lock. + Volatile.Write(ref s_cacheLock, 0); + } + } + + /// + /// Used to initialize s_callback above. We don't use a lambda for this on purpose: a lambda would + /// introduce a new generic type behind the scenes that comes with a hefty size penalty in AOT builds. + /// + private static void ExecutionContextCallback(object? s) + { + // Only used privately to pass directly to EC.Run + Debug.Assert(s is StateMachineBox); + Unsafe.As>(s).StateMachine!.MoveNext(); + } + + /// A delegate to the method. + public Action MoveNextAction => _moveNextAction ??= new Action(MoveNext); + + /// Invoked to run MoveNext when this instance is executed from the thread pool. + void IThreadPoolWorkItem.Execute() => MoveNext(); + + /// Calls MoveNext on + public void MoveNext() + { + ExecutionContext? context = Context; + + if (context is null) + { + Debug.Assert(!(StateMachine is null)); + StateMachine.MoveNext(); + } + else + { + ExecutionContext.RunInternal(context, s_callback, this); + } + } + + /// Get the result of the operation. + TResult IValueTaskSource.GetResult(short token) + { + try + { + return _valueTaskSource.GetResult(token); + } + finally + { + // Reuse this instance if possible, otherwise clear and drop it. + ReturnOrDropBox(); + } + } + + /// Get the result of the operation. + void IValueTaskSource.GetResult(short token) + { + try + { + _valueTaskSource.GetResult(token); + } + finally + { + // Reuse this instance if possible, otherwise clear and drop it. + ReturnOrDropBox(); + } + } + + /// Gets the state machine as a boxed object. This should only be used for debugging purposes. + IAsyncStateMachine IAsyncStateMachineBox.GetStateMachineObject() => StateMachine!; // likely boxes, only use for debugging } } } diff --git a/src/System.Private.CoreLib/shared/System/Threading/Tasks/Task.cs b/src/System.Private.CoreLib/shared/System/Threading/Tasks/Task.cs index 9d2c4f172379..521b62589041 100644 --- a/src/System.Private.CoreLib/shared/System/Threading/Tasks/Task.cs +++ b/src/System.Private.CoreLib/shared/System/Threading/Tasks/Task.cs @@ -2550,7 +2550,6 @@ internal void UnsafeSetContinuationForAwait(IAsyncStateMachineBox stateMachineBo // If we're unable to because the task has already completed, queue it. if (!AddTaskContinuation(stateMachineBox, addBeforeOthers: false)) { - Debug.Assert(stateMachineBox is Task, "Every state machine box should derive from Task"); ThreadPool.UnsafeQueueUserWorkItemInternal(stateMachineBox, preferLocal: true); } }