diff --git a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs index 90f74e4b867457..228458d68ebfc4 100644 --- a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs +++ b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs @@ -435,64 +435,30 @@ private static unsafe T Suspend(Task task, ConfigureAwaitOptions options) return default!; } - /// - /// Used by internal thunks that implement awaiting on ValueTask. - /// A ValueTask may wrap: - /// - Completed result (we never await this) - /// - Task - /// - ValueTaskSource - /// Therefore, when we are awaiting a ValueTask completion we are really - /// awaiting a completion of an underlying Task or ValueTaskSource. - /// - /// ValueTask whose completion we are awaiting. [Intrinsic] [BypassReadyToRun] [MethodImpl(MethodImplOptions.NoInlining | MethodImplOptions.Async)] - private static unsafe void TransparentSuspend(ValueTask valueTask) + private static unsafe void TransparentSuspend(IValueTaskSource source, short token) { ref RuntimeAsyncAwaitState state = ref t_runtimeAsyncAwaitState; Continuation? sentinelContinuation = state.SentinelContinuation ??= new Continuation(); - Continuation nextCont; - object? obj = valueTask._obj; - if (obj is Task t) + ValueTaskSourceContinuation? vtsCont = state.CachedValueTaskSourceContinuation; + if (vtsCont != null) { - RuntimeAsyncTaskContinuation? taskCont = state.CachedTaskContinuation; - if (taskCont != null) - { - state.CachedTaskContinuation = null; - } - else - { - taskCont = new RuntimeAsyncTaskContinuation(); - } - - taskCont.Initialize(t); - state.StackState->TaskContinuation = taskCont; - nextCont = taskCont; + state.CachedValueTaskSourceContinuation = null; } else { - ValueTaskSourceContinuation? vtsCont = state.CachedValueTaskSourceContinuation; - if (vtsCont != null) - { - state.CachedValueTaskSourceContinuation = null; - } - else - { - vtsCont = new ValueTaskSourceContinuation(); - } - - Debug.Assert(obj is IValueTaskSource); - vtsCont.Initialize(Unsafe.As(ref obj), valueTask._token); - state.StackState->ValueTaskSourceContinuation = vtsCont; - nextCont = vtsCont; + vtsCont = new ValueTaskSourceContinuation(); } - sentinelContinuation.Next = nextCont; + vtsCont.Initialize(source, token); + sentinelContinuation.Next = vtsCont; + state.StackState->ValueTaskSourceContinuation = vtsCont; state.CaptureContexts(); - AsyncSuspend(nextCont); + AsyncSuspend(vtsCont); } [Intrinsic] @@ -536,50 +502,27 @@ private static unsafe void Suspend(IValueTaskSource source, short token, bool co [Intrinsic] [BypassReadyToRun] [MethodImpl(MethodImplOptions.NoInlining | MethodImplOptions.Async)] - private static unsafe T TransparentSuspend(ValueTask valueTask) + private static unsafe T TransparentSuspend(IValueTaskSource source, short token) { ref RuntimeAsyncAwaitState state = ref t_runtimeAsyncAwaitState; Continuation? sentinelContinuation = state.SentinelContinuation ??= new Continuation(); - Continuation nextCont; - object? obj = valueTask._obj; - if (obj is Task t) + ValueTaskSourceContinuation? vtsCont = state.CachedValueTaskSourceContinuation; + if (vtsCont != null) { - RuntimeAsyncTaskContinuation? taskCont = state.CachedTaskContinuation; - if (taskCont != null) - { - state.CachedTaskContinuation = null; - } - else - { - taskCont = new RuntimeAsyncTaskContinuation(); - } - - taskCont.Initialize(t); - state.StackState->TaskContinuation = taskCont; - nextCont = taskCont; + state.CachedValueTaskSourceContinuation = null; } else { - ValueTaskSourceContinuation? vtsCont = state.CachedValueTaskSourceContinuation; - if (vtsCont != null) - { - state.CachedValueTaskSourceContinuation = null; - } - else - { - vtsCont = new ValueTaskSourceContinuation(); - } - - Debug.Assert(obj is IValueTaskSource); - vtsCont.Initialize(Unsafe.As>(ref obj), valueTask._token); - state.StackState->ValueTaskSourceContinuation = vtsCont; - nextCont = vtsCont; + vtsCont = new ValueTaskSourceContinuation(); } - sentinelContinuation.Next = nextCont; + vtsCont.Initialize(source, token); + + sentinelContinuation.Next = vtsCont; + state.StackState->ValueTaskSourceContinuation = vtsCont; state.CaptureContexts(); - AsyncSuspend(nextCont); + AsyncSuspend(vtsCont); return default!; } @@ -706,14 +649,35 @@ private static void TransparentAwait(Task task) [MethodImpl(MethodImplOptions.Async)] private static void TransparentAwait(ValueTask task) { - if (!task.IsCompleted) + object? obj = task._obj; + if (obj == null) + { + return; + } + + if (obj is Task t) + { + if (!t.IsCompleted) + { + TailAwait(); + TransparentSuspend(t); + return; + } + + TaskAwaiter.ValidateEnd(t); + return; + } + + Debug.Assert(obj is IValueTaskSource); + IValueTaskSource vts = Unsafe.As(ref obj); + if (vts.GetStatus(task._token) == ValueTaskSourceStatus.Pending) { TailAwait(); - TransparentSuspend(task); + TransparentSuspend(vts, task._token); return; } - task.ThrowIfCompletedUnsuccessfully(); + vts.GetResult(task._token); } [BypassReadyToRun] @@ -734,13 +698,33 @@ private static T TransparentAwait(Task task) [MethodImpl(MethodImplOptions.Async)] private static T TransparentAwait(ValueTask task) { - if (!task.IsCompleted) + object? obj = task._obj; + if (obj == null) + { + return task._result!; + } + + if (obj is Task t) + { + if (!t.IsCompleted) + { + TailAwait(); + return TransparentSuspend(t); + } + + TaskAwaiter.ValidateEnd(t); + return t.ResultOnSuccess; + } + + Debug.Assert(obj is IValueTaskSource); + IValueTaskSource vts = Unsafe.As>(ref obj); + if (vts.GetStatus(task._token) == ValueTaskSourceStatus.Pending) { TailAwait(); - return TransparentSuspend(task); + return TransparentSuspend(vts, task._token); } - return task.Result; + return vts.GetResult(task._token); } // Represents execution of a chain of suspended and resuming runtime