Remove another reference field from async state machines (#83737)

The async state machine Task-derived type currently adds three fields:
- The StateMachine
- An Action field for caching any delegate created to MoveNext
- The ExecutionContext to flow to the next MoveNext invocation

The other pending PR gets rid of the Action field by using the unused Action field from the base Task for that purpose.

This PR gets rid of the ExecutionContext field by using the unused state object field from the base Task for that purpose.  The field is exposed via the public AsyncState property, so this also uses a bit from the state flags field to prevent this state object from being returned from that property.

The combination of removing those two fields shaves 16 bytes off of every `async Task` state machine box on 64-bit.  The only remaining field added by the state machine type is for the state machine itself, which is required.

Author: stephentoub

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

@@ -3,7 +3,6 @@
 
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
-using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -290,12 +289,35 @@ private static void ExecutionContextCallback(object? s)
             private Action? _moveNextAction;
             /// <summary>The state machine itself.</summary>
             public TStateMachine? StateMachine; // mutable struct; do not make this readonly. SOS DumpAsync command depends on this name.
-            /// <summary>Captured ExecutionContext with which to invoke <see cref="MoveNextAction"/>; may be null.</summary>
-            public ExecutionContext? Context;
+
+            public AsyncStateMachineBox()
+            {
+                // The async state machine uses the base Task's state object field to store the captured execution context.
+                // Ensure that state object isn't published out for others to see.
+                Debug.Assert((m_stateFlags & (int)InternalTaskOptions.PromiseTask) != 0, "Expected state flags to already be configured.");
+                Debug.Assert(m_stateObject is null, "Expected to be able to use the state object field for ExecutionContext.");
+                m_stateFlags |= (int)InternalTaskOptions.HiddenState;
+            }
 
             /// <summary>A delegate to the <see cref="MoveNext()"/> method.</summary>
             public Action MoveNextAction => _moveNextAction ??= new Action(MoveNext);
 
+            /// <summary>Captured ExecutionContext with which to invoke <see cref="MoveNextAction"/>; may be null.</summary>
+            /// <remarks>
+            /// This uses the base Task.m_stateObject field to store the context, as that field is otherwise unused for state machine boxes.
+            /// This *must* not be set to anything other than null or an ExecutionContext, or it will result in a type safety hole.
+            /// We also don't want this ExecutionContext exposed out to consumers of the Task via Task.AsyncState, so
+            /// the ctor sets the HiddenState option to prevent this from leaking out.
+            /// </remarks>
+            public ref ExecutionContext? Context
+            {
+                get
+                {
+                    Debug.Assert(m_stateObject is null || m_stateObject is ExecutionContext, $"{nameof(m_stateObject)} must only be for ExecutionContext but contained {m_stateObject}.");
+                    return ref Unsafe.As<object?, ExecutionContext?>(ref m_stateObject);
+                }
+            }
+
             internal sealed override void ExecuteFromThreadPool(Thread threadPoolThread) => MoveNext(threadPoolThread);
 
             /// <summary>Calls MoveNext on <see cref="StateMachine"/></summary>

src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/Task.cs

@@ -9,7 +9,6 @@
 //
 // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
 
-using System.Buffers;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
@@ -120,7 +119,7 @@ public class Task : IAsyncResult, IDisposable
         [ThreadStatic]
         internal static Task? t_currentTask;  // The currently executing task.
 
-        internal static int s_taskIdCounter; // static counter used to generate unique task IDs
+        private static int s_taskIdCounter; // static counter used to generate unique task IDs
 
         private int m_taskId; // this task's unique ID. initialized only if it is ever requested
 
@@ -132,7 +131,7 @@ public class Task : IAsyncResult, IDisposable
         // the completion event which will be set when the Future class calls Finish().
         // But the event would now be signalled if Cancel() is called
 
-        internal object? m_stateObject; // A state object that can be optionally supplied, passed to action.
+        private protected object? m_stateObject; // A state object that can be optionally supplied, passed to action.
         internal TaskScheduler? m_taskScheduler; // The task scheduler this task runs under.
 
         internal volatile int m_stateFlags; // SOS DumpAsync command depends on this name
@@ -566,6 +565,7 @@ internal void TaskConstructorCore(Delegate? action, object? state, CancellationT
             int illegalInternalOptions =
                     (int)(internalOptions &
                             ~(InternalTaskOptions.PromiseTask |
+                              InternalTaskOptions.HiddenState |
                               InternalTaskOptions.ContinuationTask |
                               InternalTaskOptions.LazyCancellation |
                               InternalTaskOptions.QueuedByRuntime));
@@ -1446,7 +1446,7 @@ WaitHandle IAsyncResult.AsyncWaitHandle
         /// Gets the state object supplied when the <see cref="Task">Task</see> was created,
         /// or null if none was supplied.
         /// </summary>
-        public object? AsyncState => m_stateObject;
+        public object? AsyncState => (m_stateFlags & (int)InternalTaskOptions.HiddenState) == 0 ? m_stateObject : null;
 
         /// <summary>
         /// Gets an indication of whether the asynchronous operation completed synchronously.
@@ -6717,6 +6717,11 @@ internal enum InternalTaskOptions
         ContinuationTask = 0x0200,
         PromiseTask = 0x0400,
 
+        /// <summary>
+        /// The state object should not be returned from the AsyncState property.
+        /// </summary>
+        HiddenState = 0x0800,
+
         /// <summary>
         /// Store the presence of TaskContinuationOptions.LazyCancellation, since it does not directly
         /// translate into any TaskCreationOptions.

src/libraries/System.Threading.Tasks/tests/System.Runtime.CompilerServices/AsyncTaskMethodBuilderTests.cs

@@ -662,6 +662,19 @@ public void AsyncTaskMethodBuilder_Completed_RemovedFromTracking()
             }).Dispose();
         }
 
+        [Fact]
+        public void AsyncTaskMethodBuilder_NullStateEvenAfterSuspend()
+        {
+            Task t = AwaitSomething();
+            Assert.Null(t.AsyncState);
+
+            static async Task AwaitSomething()
+            {
+                Assert.NotNull(ExecutionContext.Capture());
+                await new TaskCompletionSource().Task;
+            }
+        }
+
         #region Helper Methods / Classes
 
         [MethodImpl(MethodImplOptions.NoInlining)]