diff --git a/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncStateMachineDispatcher.cs b/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncStateMachineDispatcher.cs index 5a1f6ef5179e02..cef3d9e422f1c9 100644 --- a/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncStateMachineDispatcher.cs +++ b/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncStateMachineDispatcher.cs @@ -177,7 +177,6 @@ internal static unsafe void CompleteAsyncMethod(object completingBox, AsyncInstr internal sealed class AsyncStateMachineDispatcher : Task, IAsyncStateMachineBox { private IAsyncStateMachineBox? _inner; - private Action? _moveNextAction; internal IAsyncStateMachineBox? LastContinuation; @@ -235,7 +234,7 @@ public unsafe void MoveNext() } } - public Action MoveNextAction => _moveNextAction ??= MoveNext; + public Action MoveNextAction => (Action)(m_action ??= new Action(MoveNext)); public IAsyncStateMachine GetStateMachineObject() { diff --git a/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncTaskMethodBuilderT.cs b/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncTaskMethodBuilderT.cs index 513de50677ec00..bc5df6b275b987 100644 --- a/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncTaskMethodBuilderT.cs +++ b/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncTaskMethodBuilderT.cs @@ -61,7 +61,16 @@ internal static void AwaitOnCompleted( { try { - awaiter.OnCompleted(GetStateMachineBox(ref stateMachine, ref taskField).MoveNextAction); + IAsyncStateMachineBox box = GetStateMachineBox(ref stateMachine, ref taskField); + if (AsyncInstrumentation.IsActive && AsyncInstrumentation.LoadFlags(out AsyncInstrumentation.Flags flags)) + { + if (AsyncInstrumentation.IsEnabled.AsyncProfiler(flags)) + { + box = AsyncStateMachineDispatcherInfo.CreateDispatcher(box, flags); + } + } + + awaiter.OnCompleted(box.MoveNextAction); } catch (Exception e) { @@ -136,6 +145,14 @@ internal static void AwaitUnsafeOnCompleted( // The awaiter isn't specially known. Fall back to doing a normal await. try { + if (AsyncInstrumentation.IsActive && AsyncInstrumentation.LoadFlags(out AsyncInstrumentation.Flags flags)) + { + if (AsyncInstrumentation.IsEnabled.AsyncProfiler(flags)) + { + box = AsyncStateMachineDispatcherInfo.CreateDispatcher(box, flags); + } + } + awaiter.UnsafeOnCompleted(box.MoveNextAction); } catch (Exception e) diff --git a/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/PoolingAsyncValueTaskMethodBuilderT.cs b/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/PoolingAsyncValueTaskMethodBuilderT.cs index 49db22cb6e6064..5e4c41c576a4a2 100644 --- a/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/PoolingAsyncValueTaskMethodBuilderT.cs +++ b/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/PoolingAsyncValueTaskMethodBuilderT.cs @@ -115,7 +115,16 @@ internal static void AwaitOnCompleted( { try { - awaiter.OnCompleted(GetStateMachineBox(ref stateMachine, ref box).MoveNextAction); + IAsyncStateMachineBox ibox = GetStateMachineBox(ref stateMachine, ref box); + if (AsyncInstrumentation.IsActive && AsyncInstrumentation.LoadFlags(out AsyncInstrumentation.Flags flags)) + { + if (AsyncInstrumentation.IsEnabled.AsyncProfiler(flags)) + { + ibox = AsyncStateMachineDispatcherInfo.CreateDispatcher(ibox, flags); + } + } + + awaiter.OnCompleted(ibox.MoveNextAction); } catch (Exception e) { diff --git a/src/libraries/System.Runtime/tests/System.Threading.Tasks.Tests/System.Runtime.CompilerServices/AsyncProfilerV1Tests.cs b/src/libraries/System.Runtime/tests/System.Threading.Tasks.Tests/System.Runtime.CompilerServices/AsyncProfilerV1Tests.cs index 31a5eb5a8dfee4..2c8b661e290a40 100644 --- a/src/libraries/System.Runtime/tests/System.Threading.Tasks.Tests/System.Runtime.CompilerServices/AsyncProfilerV1Tests.cs +++ b/src/libraries/System.Runtime/tests/System.Threading.Tasks.Tests/System.Runtime.CompilerServices/AsyncProfilerV1Tests.cs @@ -1109,6 +1109,118 @@ public void StateMachineAsync_CustomTaskScheduler_EmitsContextEventsAndCallstack AssertTrue(stream, completeIdx > resumeIdx, "Expected CompleteStateMachineAsyncContext after Resume"); } + // Custom awaiters that post their continuation directly to the thread pool, bypassing the + // Task/ValueTask/Yield "known awaiter" fast paths (ITaskAwaiter / IConfiguredTaskAwaiter / + // IStateMachineBoxAwareAwaiter). They exercise the builder's fallback completion paths: + // the ICriticalNotifyCompletion awaiter routes through AwaitUnsafeOnCompleted's else-branch, + // the INotifyCompletion-only awaiter routes through AwaitOnCompleted. Both must still create a + // dispatcher so the await is represented in the V1 event stream. + private sealed class DirectPostCriticalAwaitable + { + public DirectPostCriticalAwaiter GetAwaiter() => default; + } + + private readonly struct DirectPostCriticalAwaiter : ICriticalNotifyCompletion + { + public bool IsCompleted => false; + public void GetResult() { } + public void OnCompleted(Action continuation) => UnsafeOnCompleted(continuation); + public void UnsafeOnCompleted(Action continuation) => + ThreadPool.QueueUserWorkItem(static c => c(), continuation, preferLocal: false); + } + + private sealed class DirectPostNotifyAwaitable + { + public DirectPostNotifyAwaiter GetAwaiter() => default; + } + + private readonly struct DirectPostNotifyAwaiter : INotifyCompletion + { + public bool IsCompleted => false; + public void GetResult() { } + public void OnCompleted(Action continuation) => + ThreadPool.QueueUserWorkItem(static c => c(), continuation, preferLocal: false); + } + + [RuntimeAsyncMethodGeneration(false)] + [MethodImpl(MethodImplOptions.NoInlining)] + private static async Task StateMachineAsync_CustomAwaiter_EmitsCreateResumeComplete_Critical_Marker() + { + await new DirectPostCriticalAwaitable(); + } + + [RuntimeAsyncMethodGeneration(false)] + [MethodImpl(MethodImplOptions.NoInlining)] + private static async Task StateMachineAsync_CustomAwaiter_EmitsCreateResumeComplete_Notify_Marker() + { + await new DirectPostNotifyAwaitable(); + } + + [RuntimeAsyncMethodGeneration(false)] + [MethodImpl(MethodImplOptions.NoInlining)] + [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] + private static async ValueTask StateMachineAsync_CustomAwaiter_EmitsCreateResumeComplete_PoolingCritical_Marker() + { + await new DirectPostCriticalAwaitable(); + } + + [RuntimeAsyncMethodGeneration(false)] + [MethodImpl(MethodImplOptions.NoInlining)] + [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] + private static async ValueTask StateMachineAsync_CustomAwaiter_EmitsCreateResumeComplete_PoolingNotify_Marker() + { + await new DirectPostNotifyAwaitable(); + } + + // Covers all four builder fallback sites that wrap the box for an opaque custom awaiter: + // - Task + ICriticalNotifyCompletion -> AsyncTaskMethodBuilderT.AwaitUnsafeOnCompleted else-branch + // - Task + INotifyCompletion -> AsyncTaskMethodBuilderT.AwaitOnCompleted + // - pooling + ICriticalNotifyCompletion -> PoolingAsyncValueTaskMethodBuilderT.AwaitUnsafeOnCompleted (delegates to the shared else-branch) + // - pooling + INotifyCompletion -> PoolingAsyncValueTaskMethodBuilderT.AwaitOnCompleted (its own create site) + [ConditionalTheory(typeof(AsyncProfilerTests), nameof(IsStateMachineAsyncAndThreadingSupported))] + [InlineData(false, true)] + [InlineData(false, false)] + [InlineData(true, true)] + [InlineData(true, false)] + public void StateMachineAsync_CustomAwaiter_EmitsCreateResumeComplete(bool pooling, bool criticalNotifyCompletion) + { + (string markerName, Func scenario) = (pooling, criticalNotifyCompletion) switch + { + (false, true) => (nameof(StateMachineAsync_CustomAwaiter_EmitsCreateResumeComplete_Critical_Marker), (Func)StateMachineAsync_CustomAwaiter_EmitsCreateResumeComplete_Critical_Marker), + (false, false) => (nameof(StateMachineAsync_CustomAwaiter_EmitsCreateResumeComplete_Notify_Marker), StateMachineAsync_CustomAwaiter_EmitsCreateResumeComplete_Notify_Marker), + (true, true) => (nameof(StateMachineAsync_CustomAwaiter_EmitsCreateResumeComplete_PoolingCritical_Marker), () => StateMachineAsync_CustomAwaiter_EmitsCreateResumeComplete_PoolingCritical_Marker().AsTask()), + _ => (nameof(StateMachineAsync_CustomAwaiter_EmitsCreateResumeComplete_PoolingNotify_Marker), () => StateMachineAsync_CustomAwaiter_EmitsCreateResumeComplete_PoolingNotify_Marker().AsTask()), + }; + + var events = CollectEvents(ResumeStateMachineAsyncCallstackKeyword | StateMachineAsyncCoreKeywords, () => + { + RunScenarioAndFlush(scenario); + }); + + // DumpAllEvents(events); + + var stream = ParseAllEvents(events); + + // A custom awaiter that posts its continuation directly (bypassing the Task/ValueTask/Yield + // fast paths) must still produce a dispatcher: the marker frame appears in a Resume callstack + // and the chain emits the standard Create -> Resume -> Complete sequence. Without the fallback + // dispatcher creation the marker is absent from the stream and this assertion fails. + var markerCallstacks = stream.CallstacksWithMarker(AsyncEventID.ResumeStateMachineAsyncCallstack, markerName); + AssertNotEmpty(stream, markerCallstacks); + + ulong dispatcherId = markerCallstacks[0].DispatcherId; + var ids = stream.ChainEventsFromDispatcher(dispatcherId).Select(e => e.EventId).ToList(); + + int createIdx = ids.IndexOf(AsyncEventID.CreateStateMachineAsyncContext); + AssertTrue(stream, createIdx >= 0, "Expected CreateStateMachineAsyncContext for the custom awaiter scenario"); + + int resumeIdx = ids.IndexOf(AsyncEventID.ResumeStateMachineAsyncContext, createIdx + 1); + AssertTrue(stream, resumeIdx > createIdx, "Expected ResumeStateMachineAsyncContext after Create"); + + int completeIdx = ids.IndexOf(AsyncEventID.CompleteStateMachineAsyncContext, resumeIdx + 1); + AssertTrue(stream, completeIdx > resumeIdx, "Expected CompleteStateMachineAsyncContext after Resume"); + } + [RuntimeAsyncMethodGeneration(false)] [MethodImpl(MethodImplOptions.NoInlining)] private static async Task StateMachineAsync_NoEventsWhenDisabled_Marker()