Fix use of AddTo/RemoveFromActiveTasks in async methods (#25915)

When we rewrote the async methods implementation with AsyncStateMachineBox, we neglected to call AddToActiveTasks if the debugger is paying attention to tasks.  This hasn't mattered as the debugger's support for Tasks hasn't worked for other reasons in .NET Core, but there is now a renewed focus on it, and this is preventing some of that support from working.

This change is a minimal fix to ensure that we're adding the state machine box task when it's created and removing it when it completes.  Post-3.0, we should look at overhauling this, e.g. to clean up a lot of this logging and tracking that's done, to use a weak table in order to avoid keeping task objects alive artificially if they're dropped without completing by the developer code, etc.

This only affects when the debugger is attached, as the s_asyncDebuggingEnabled field is only ever set by the debugger.

1 files changed, 15 insertions, 20 deletions

diff --git a/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncMethodBuilder.cs b/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncMethodBuilder.cs
index 3c9ecc3a12..97c0535e3c 100644
--- a/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncMethodBuilder.cs
+++ b/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/AsyncMethodBuilder.cs
@@ -517,12 +517,18 @@ namespace System.Runtime.CompilerServices
             box.StateMachine = stateMachine;
             box.Context = currentContext;
 
-            // Finally, log the creation of the state machine box object / task for this async method.
+            // Log the creation of the state machine box object / task for this async method.
             if (AsyncCausalityTracer.LoggingOn)
             {
                 AsyncCausalityTracer.TraceOperationCreation(box, "Async: " + stateMachine.GetType().Name);
             }
 
+            // And if async debugging is enabled, track the task.
+            if (System.Threading.Tasks.Task.s_asyncDebuggingEnabled)
+            {
+                System.Threading.Tasks.Task.AddToActiveTasks(box);
+            }
+
             return box;
         }
 
@@ -618,6 +624,12 @@ namespace System.Runtime.CompilerServices
 
                 if (IsCompleted)
                 {
+                    // If async debugging is enabled, remove the task from tracking.
+                    if (System.Threading.Tasks.Task.s_asyncDebuggingEnabled)
+                    {
+                        System.Threading.Tasks.Task.RemoveFromActiveTasks(this);
+                    }
+
                     // Clear out state now that the async method has completed.
                     // This avoids keeping arbitrary state referenced by lifted locals
                     // if this Task / state machine box is held onto.
@@ -716,31 +728,14 @@ namespace System.Runtime.CompilerServices
         {
             Debug.Assert(m_task != null, "Expected non-null task");
 
-            if (AsyncCausalityTracer.LoggingOn || System.Threading.Tasks.Task.s_asyncDebuggingEnabled)
-            {
-                LogExistingTaskCompletion();
-            }
-
-            if (!m_task.TrySetResult(result))
-            {
-                ThrowHelper.ThrowInvalidOperationException(ExceptionResource.TaskT_TransitionToFinal_AlreadyCompleted);
-            }
-        }
-
-        /// <summary>Handles logging for the successful completion of an operation.</summary>
-        private void LogExistingTaskCompletion()
-        {
-            Debug.Assert(m_task != null);
-
             if (AsyncCausalityTracer.LoggingOn)
             {
                 AsyncCausalityTracer.TraceOperationCompletion(m_task, AsyncCausalityStatus.Completed);
             }
 
-            // only log if we have a real task that was previously created
-            if (System.Threading.Tasks.Task.s_asyncDebuggingEnabled)
+            if (!m_task.TrySetResult(result))
             {
-                System.Threading.Tasks.Task.RemoveFromActiveTasks(m_task);
+                ThrowHelper.ThrowInvalidOperationException(ExceptionResource.TaskT_TransitionToFinal_AlreadyCompleted);
             }
         }