Async Profiler: Instrument custom awaiters in V1 dispatcher.

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<VoidTaskResult>, 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()
         {

src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncTaskMethodBuilderT.cs

@@ -61,7 +61,16 @@ internal static void AwaitOnCompleted<TAwaiter, TStateMachine>(
         {
             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<TAwaiter>(
                 // 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)

src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/PoolingAsyncValueTaskMethodBuilderT.cs

@@ -115,7 +115,16 @@ internal static void AwaitOnCompleted<TAwaiter, TStateMachine>(
         {
             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)
             {

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<Task> scenario) = (pooling, criticalNotifyCompletion) switch
+            {
+                (false, true) => (nameof(StateMachineAsync_CustomAwaiter_EmitsCreateResumeComplete_Critical_Marker), (Func<Task>)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()