Force secondary await continuations to run asynchronously (#25280)

* Force secondary await continuations to run asynchronously

For performance reasons, await continuations have been invoked synchronously, meaning they're invoked as part of the antecedent task's completion (as long as that task allows it, as long as there's sufficient stack space, etc.)  This generally works out well in the case where there's a single await continuation, which is far and away the common case.  However, it can cause problems in situations where there are multiple await continuations, as those continuations will end up being serialized, which can lead to slowdowns and deadlocks in niche situations.  To address that, this commit backs off a bit.  The first await continuation is still invoked synchronously, but subsequent await continuations are invoked asynchronously, such that they are not blocked by a previously registered await continuation.

* Fix nits

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 bf539a5..135f588 100644 (file)
--- a/src/System.Private.CoreLib/shared/System/Threading/Tasks/Task.cs
+++ b/src/System.Private.CoreLib/shared/System/Threading/Tasks/Task.cs
@@ -3296,11 +3296,10 @@ namespace System.Threading.Tasks
         {
             Debug.Assert(continuationObject != null);
 
-            TplEventSource log = TplEventSource.Log;
-            bool TplEventSourceLoggingEnabled = log.IsEnabled();
-            if (TplEventSourceLoggingEnabled)
+            TplEventSource? log = TplEventSource.Log;
+            if (!log.IsEnabled())
             {
-                log.RunningContinuation(Id, continuationObject);
+                log = null;
             }
 
             if (AsyncCausalityTracer.LoggingOn)
@@ -3351,26 +3350,64 @@ namespace System.Threading.Tasks
             lock (continuations) { }
             int continuationCount = continuations.Count;
 
-            // Fire the asynchronous continuations first ...
-            for (int i = 0; i < continuationCount; i++)
+            // Fire the asynchronous continuations first. However, if we're not able to run any continuations synchronously,
+            // then we can skip this first pass, since the second pass that tries to run everything synchronously will instead
+            // run everything asynchronously anyway.
+            if (canInlineContinuations)
             {
-                // Synchronous continuation tasks will have the ExecuteSynchronously option,
-                // and we're looking for asynchronous tasks...
-                if (continuations[i] is StandardTaskContinuation tc &&
-                    (tc.m_options & TaskContinuationOptions.ExecuteSynchronously) == 0)
+                bool forceContinuationsAsync = false;
+                for (int i = 0; i < continuationCount; i++)
                 {
-                    if (TplEventSourceLoggingEnabled)
+                    // For StandardTaskContinuations, we respect the TaskContinuationOptions.ExecuteSynchronously option,
+                    // as the developer needs to explicitly opt-into running the continuation synchronously, and if they do,
+                    // they get what they asked for. ITaskCompletionActions are only ever created by the runtime, and we always
+                    // try to execute them synchronously. For all other continuations (related to await), we only run it synchronously
+                    // if it's the first such continuation; otherwise, we force it to run asynchronously so as to not artificially
+                    // delay an await continuation behind other arbitrary user code created as a previous await continuation.
+
+                    object? currentContinuation = continuations[i];
+                    if (currentContinuation == null)
                     {
-                        log.RunningContinuationList(Id, i, tc);
+                        // The continuation was unregistered and null'd out, so just skip it.
+                        continue;
+                    }
+                    else if (currentContinuation is StandardTaskContinuation stc)
+                    {
+                        if ((stc.m_options & TaskContinuationOptions.ExecuteSynchronously) == 0)
+                        {
+                            continuations[i] = null; // so that we can skip this later
+                            log?.RunningContinuationList(Id, i, stc);
+                            stc.Run(this, canInlineContinuationTask: false);
+                        }
+                    }
+                    else if (!(currentContinuation is ITaskCompletionAction))
+                    {
+                        if (forceContinuationsAsync)
+                        {
+                            continuations[i] = null;
+                            log?.RunningContinuationList(Id, i, currentContinuation);
+                            switch (currentContinuation)
+                            {
+                                case IAsyncStateMachineBox stateMachineBox:
+                                    AwaitTaskContinuation.RunOrScheduleAction(stateMachineBox, allowInlining: false);
+                                    break;
+
+                                case Action action:
+                                    AwaitTaskContinuation.RunOrScheduleAction(action, allowInlining: false);
+                                    break;
+
+                                default:
+                                    Debug.Assert(currentContinuation is TaskContinuation);
+                                    ((TaskContinuation)currentContinuation).Run(this, canInlineContinuationTask: false);
+                                    break;
+                            }
+                        }
+                        forceContinuationsAsync = true;
                     }
-                    continuations[i] = null; // so that we can skip this later
-                    tc.Run(this, canInlineContinuations);
                 }
             }
 
             // ... and then fire the synchronous continuations (if there are any).
-            // This includes ITaskCompletionAction, AwaitTaskContinuations, IAsyncStateMachineBox,
-            // and Action delegates, which are all by default implicitly synchronous.
             for (int i = 0; i < continuationCount; i++)
             {
                 object? currentContinuation = continuations[i];
@@ -3379,10 +3416,7 @@ namespace System.Threading.Tasks
                     continue;
                 }
                 continuations[i] = null; // to enable free'ing up memory earlier
-                if (TplEventSourceLoggingEnabled)
-                {
-                    log.RunningContinuationList(Id, i, currentContinuation);
-                }
+                log?.RunningContinuationList(Id, i, currentContinuation);
 
                 switch (currentContinuation)
                 {