Allow continuations to have no ExecutionContext (#128323)

`ValueTaskContinuation` does not need to save/restore ExecutionContext,
and also in the future we expect to allow continuations to skip this
save/restore when the JIT proved that it is not going to be looked at.

Author: jakobbotsch

src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs

@@ -99,13 +99,20 @@ public bool HasException()
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        public unsafe ExecutionContext? GetExecutionContext()
+        public unsafe bool TryGetExecutionContext(out ExecutionContext? execContext)
         {
             const uint mask = (1u << (int)ContinuationFlags.ExecutionContextIndexNumBits) - 1;
             uint index = ((uint)Flags >> (int)ContinuationFlags.ExecutionContextIndexFirstBit) & mask;
+            if (index == 0)
+            {
+                execContext = null;
+                return false;
+            }
+
             Debug.Assert(index != 0);
             ref byte data = ref RuntimeHelpers.GetRawData(this);
-            return Unsafe.As<byte, ExecutionContext?>(ref Unsafe.Add(ref data, (DataOffset - PointerSize) + index * PointerSize));
+            execContext = Unsafe.As<byte, ExecutionContext?>(ref Unsafe.Add(ref data, (DataOffset - PointerSize) + index * PointerSize));
+            return true;
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
@@ -128,16 +135,6 @@ public ref byte GetResultStorageOrNull()
             ref byte data = ref RuntimeHelpers.GetRawData(this);
             return ref Unsafe.Add(ref data, (DataOffset - PointerSize) + index * PointerSize);
         }
-
-        protected void EncodeFieldOffsetInFlags(ref byte field, ContinuationFlags firstBit, ContinuationFlags numBits)
-        {
-            int offset = (int)Unsafe.ByteOffset(ref RuntimeHelpers.GetRawData(this), ref field);
-            offset -= DataOffset;
-            Debug.Assert(offset >= 0 && offset % PointerSize == 0);
-            uint index = 1 + (uint)offset / PointerSize;
-            Debug.Assert(index < (1 << (int)numBits));
-            Flags |= (ContinuationFlags)((uint)index << (int)firstBit);
-        }
     }
 
     [StructLayout(LayoutKind.Explicit)]
@@ -373,7 +370,6 @@ private static unsafe void TransparentAwaitValueTask(ValueTask valueTask)
 
             Debug.Assert(valueTask._obj != null);
             vtsCont.Initialize(valueTask._obj, valueTask._token);
-            vtsCont.ExecutionContext = ExecutionContext.CaptureForSuspension(state.CurrentThread!);
 
             sentinelContinuation.Next = vtsCont;
             state.StackState->ValueTaskContinuation = vtsCont;
@@ -401,7 +397,6 @@ private static unsafe void TransparentAwaitValueTaskOfT<T>(ValueTask<T?> valueTa
 
             Debug.Assert(valueTask._obj != null);
             vtsCont.Initialize<T>(valueTask._obj, valueTask._token);
-            vtsCont.ExecutionContext = ExecutionContext.CaptureForSuspension(state.CurrentThread!);
 
             sentinelContinuation.Next = vtsCont;
             state.StackState->ValueTaskContinuation = vtsCont;
@@ -660,7 +655,11 @@ private unsafe void DispatchContinuations()
                         asyncDispatcherInfo.NextContinuation = nextContinuation;
 
                         Debug.Assert(awaitState.CurrentThread != null);
-                        RestoreExecutionContext(awaitState.CurrentThread, curContinuation.GetExecutionContext());
+                        if (curContinuation.TryGetExecutionContext(out ExecutionContext? execContext))
+                        {
+                            RestoreExecutionContext(awaitState.CurrentThread, execContext);
+                        }
+
                         ref byte resultLoc = ref nextContinuation != null ? ref nextContinuation.GetResultStorageOrNull() : ref GetResultStorage();
 
                         Continuation? newContinuation = curContinuation.ResumeInfo->Resume(curContinuation, ref resultLoc);
@@ -766,7 +765,11 @@ private unsafe void InstrumentedDispatchContinuations(AsyncInstrumentation.Flags
                         asyncDispatcherInfo.NextContinuation = nextContinuation;
 
                         Debug.Assert(awaitState.CurrentThread != null);
-                        RestoreExecutionContext(awaitState.CurrentThread, curContinuation.GetExecutionContext());
+                        if (curContinuation.TryGetExecutionContext(out ExecutionContext? execContext))
+                        {
+                            RestoreExecutionContext(awaitState.CurrentThread, execContext);
+                        }
+
                         ref byte resultLoc = ref nextContinuation != null ? ref nextContinuation.GetResultStorageOrNull() : ref GetResultStorage();
 
                         RuntimeAsyncInstrumentationHelpers.ResumeRuntimeAsyncMethod(ref asyncDispatcherInfo, flags, curContinuation);

src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.ValueTaskContinuation.cs

@@ -12,9 +12,6 @@ public static partial class AsyncHelpers
     {
         private sealed unsafe class ValueTaskContinuation : Continuation
         {
-            // Currently all continuations are expected to capture and restore
-            // ExecutionContext, even though we do not actually need it here.
-            public ExecutionContext? ExecutionContext;
             internal object? Source;
             internal short Token;
             internal delegate*<object, Action<object?>, object?, short, ValueTaskSourceOnCompletedFlags, void> OnCompletedValueTaskSource;
@@ -23,11 +20,6 @@ private sealed unsafe class ValueTaskContinuation : Continuation
             public ValueTaskContinuation()
             {
                 ResumeInfo = (ResumeInfo*)Unsafe.AsPointer(in ValueTaskContinuationResume.ResumeInfo);
-
-                EncodeFieldOffsetInFlags(
-                    ref Unsafe.As<ExecutionContext?, byte>(ref ExecutionContext),
-                    ContinuationFlags.ExecutionContextIndexFirstBit,
-                    ContinuationFlags.ExecutionContextIndexNumBits);
             }
 
             public void GetResult(ref byte returnValue)
@@ -113,7 +105,6 @@ private static class ValueTaskContinuationResume
                 {
                     var vtsCont = (ValueTaskContinuation)cont;
                     vtsCont.Next = null;
-                    vtsCont.ExecutionContext = null;
                     t_runtimeAsyncAwaitState.CachedValueTaskContinuation = vtsCont;
 
                     vtsCont.GetResult(ref result);