Clear unnecessary state from completed Tasks (#20294)

When Tasks backed by delegates are created, an ExecutionContext is captured.  When the task completes, its delegate is being cleared, but its ExecutionContext is not, which means if the Task is subsequently kept alive (e.g. stored in a cache), so too is its ExecutionContext, which can capture an arbitrary amount of ambient state via async locals.  This commit augments the clearing of the delegate to similarly clear the ExecutionContext.

Related, async methods can also capture ExecutionContext when awaits yield, so this clears out that context as well.  And as long as we're doing that, we may as well also clear the state machine state, so that any hoisted locals in the state machine aren't kept alive if the resulting task is kept alive.  Not doing so previously was a conscious choice, in order to aid in debugging, but as we've heard of at least a couple of cases where it unexpectedly caused a leak, I'm going ahead and changing it.

diff --git a/src/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncMethodBuilder.cs b/src/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncMethodBuilder.cs
index 2fb284a..fc37ca7 100644 (file)
--- a/src/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncMethodBuilder.cs
+++ b/src/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncMethodBuilder.cs
@@ -541,6 +541,8 @@ namespace System.Runtime.CompilerServices
             /// <summary>Calls MoveNext on <see cref="StateMachine"/></summary>
             public void MoveNext()
             {
+                Debug.Assert(!IsCompleted);
+
                 bool loggingOn = AsyncCausalityTracer.LoggingOn;
                 if (loggingOn)
                 {
@@ -557,12 +559,21 @@ namespace System.Runtime.CompilerServices
                     ExecutionContext.RunInternal(context, s_callback, this);
                 }
 
-                // In case this is a state machine box with a finalizer, suppress its finalization
-                // if it's now complete.  We only need the finalizer to run if the box is collected
-                // without having been completed.
-                if (IsCompleted && AsyncMethodBuilderCore.TrackAsyncMethodCompletion)
+                if (IsCompleted)
                 {
-                    GC.SuppressFinalize(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.
+                    StateMachine = default;
+                    Context = default;
+
+                    // In case this is a state machine box with a finalizer, suppress its finalization
+                    // as it's now complete.  We only need the finalizer to run if the box is collected
+                    // without having been completed.
+                    if (AsyncMethodBuilderCore.TrackAsyncMethodCompletion)
+                    {
+                        GC.SuppressFinalize(this);
+                    }
                 }
 
                 if (loggingOn)

diff --git a/src/System.Private.CoreLib/src/System/Threading/Tasks/Task.cs b/src/System.Private.CoreLib/src/System/Threading/Tasks/Task.cs
index 161540a..3bce5fb 100644 (file)
--- a/src/System.Private.CoreLib/src/System/Threading/Tasks/Task.cs
+++ b/src/System.Private.CoreLib/src/System/Threading/Tasks/Task.cs
@@ -2178,9 +2178,15 @@ namespace System.Threading.Tasks
             // continuations hold onto the task, and therefore are keeping it alive.
             m_action = null;
 
-            // Notify parent if this was an attached task
-            if (m_contingentProperties != null)
+            ContingentProperties cp = m_contingentProperties;
+            if (cp != null)
             {
+                // Similarly, null out any ExecutionContext we may have captured,
+                // to avoid keeping state like async locals alive unnecessarily
+                // when the Task is kept alive.
+                cp.m_capturedContext = null;
+
+                // Notify parent if this was an attached task
                 NotifyParentIfPotentiallyAttachedTask();
             }