Fix missing ClearStateUponCompletion call in AsyncTaskMethodBuilder (#53962)

This was accidentally removed as part of a previous refactoring.  The lack of it means primarily that we're not always removing async method tasks from the s_currentActiveTasks dictionary when the debugger is attached (we didn't notice it because awaiting the resulting task also causes it to be removed).  It also means we're not clearing out some state we intended to in order to minimize extending lifetime of objects referenced by the state machine if the task is held onto for a prolonged period.

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 35c73e3..fdf9943 100644 (file)
--- a/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncTaskMethodBuilderT.cs
+++ b/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncTaskMethodBuilderT.cs
@@ -329,6 +329,11 @@ namespace System.Runtime.CompilerServices
                     }
                 }
 
+                if (IsCompleted)
+                {
+                    ClearStateUponCompletion();
+                }
+
                 if (loggingOn)
                 {
                     TplEventSource.Log.TraceSynchronousWorkEnd(CausalitySynchronousWork.Execution);

diff --git a/src/libraries/System.Threading.Tasks/tests/System.Runtime.CompilerServices/AsyncTaskMethodBuilderTests.cs b/src/libraries/System.Threading.Tasks/tests/System.Runtime.CompilerServices/AsyncTaskMethodBuilderTests.cs
index 1ce1e3f..1776960 100644 (file)
--- a/src/libraries/System.Threading.Tasks/tests/System.Runtime.CompilerServices/AsyncTaskMethodBuilderTests.cs
+++ b/src/libraries/System.Threading.Tasks/tests/System.Runtime.CompilerServices/AsyncTaskMethodBuilderTests.cs
@@ -5,6 +5,7 @@ using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Diagnostics.Tracing;
 using System.Linq;
+using System.Reflection;
 using System.Runtime.CompilerServices;
 using System.Text;
 using Microsoft.DotNet.RemoteExecutor;
@@ -644,6 +645,30 @@ namespace System.Threading.Tasks.Tests
             }).Dispose();
         }
 
+        [ConditionalFact(typeof(RemoteExecutor), nameof(RemoteExecutor.IsSupported))]
+        public void AsyncTaskMethodBuilder_Completed_RemovedFromTracking()
+        {
+            RemoteExecutor.Invoke(() =>
+            {
+                // NOTE: This depends on private implementation details generally only used by the debugger.
+                // If those ever change, this test will need to be updated as well.
+                typeof(Task).GetField("s_asyncDebuggingEnabled", BindingFlags.NonPublic | BindingFlags.Static).SetValue(null, true);
+
+                for (int i = 0; i < 1000; i++)
+                {
+                    static async Task YieldAsync(TaskCompletionSource tcs) => await tcs.Task;
+
+                    TaskCompletionSource tcs = new();
+                    Task t = YieldAsync(tcs);
+                    tcs.SetResult();
+                    t.Wait();
+                }
+
+                int activeCount = ((dynamic)typeof(Task).GetField("s_currentActiveTasks", BindingFlags.NonPublic | BindingFlags.Static).GetValue(null)).Count;
+                Assert.InRange(activeCount, 0, 10); // some other tasks may be created by the runtime, so this is just using a reasonably small upper bound
+            }).Dispose();
+        }
+
         #region Helper Methods / Classes
 
         [MethodImpl(MethodImplOptions.NoInlining)]