Add PreAllocatedOverlapped.UnsafeCreate / UnsafeAllocateNativeOverlap…

…ped (#53196)

src/coreclr/System.Private.CoreLib/src/System/Threading/ClrThreadPoolBoundHandle.cs

@@ -124,14 +124,65 @@ public static ThreadPoolBoundHandle BindHandle(SafeHandle handle)
         ///     This method was called after the <see cref="ThreadPoolBoundHandle"/> was disposed.
         /// </exception>
         [CLSCompliant(false)]
-        public unsafe NativeOverlapped* AllocateNativeOverlapped(IOCompletionCallback callback, object? state, object? pinData)
+        public unsafe NativeOverlapped* AllocateNativeOverlapped(IOCompletionCallback callback, object? state, object? pinData) =>
+            AllocateNativeOverlapped(callback, state, pinData, flowExecutionContext: true);
+
+        /// <summary>
+        ///     Returns an unmanaged pointer to a <see cref="NativeOverlapped"/> structure, specifying
+        ///     a delegate that is invoked when the asynchronous I/O operation is complete, a user-provided
+        ///     object providing context, and managed objects that serve as buffers.
+        /// </summary>
+        /// <param name="callback">
+        ///     An <see cref="IOCompletionCallback"/> delegate that represents the callback method
+        ///     invoked when the asynchronous I/O operation completes.
+        /// </param>
+        /// <param name="state">
+        ///     A user-provided object that distinguishes this <see cref="NativeOverlapped"/> from other
+        ///     <see cref="NativeOverlapped"/> instances. Can be <see langword="null"/>.
+        /// </param>
+        /// <param name="pinData">
+        ///     An object or array of objects representing the input or output buffer for the operation. Each
+        ///     object represents a buffer, for example an array of bytes.  Can be <see langword="null"/>.
+        /// </param>
+        /// <returns>
+        ///     An unmanaged pointer to a <see cref="NativeOverlapped"/> structure.
+        /// </returns>
+        /// <remarks>
+        ///     <para>
+        ///         The unmanaged pointer returned by this method can be passed to the operating system in
+        ///         overlapped I/O operations. The <see cref="NativeOverlapped"/> structure is fixed in
+        ///         physical memory until <see cref="FreeNativeOverlapped(NativeOverlapped*)"/> is called.
+        ///     </para>
+        ///     <para>
+        ///         The buffer or buffers specified in <paramref name="pinData"/> must be the same as those passed
+        ///         to the unmanaged operating system function that performs the asynchronous I/O.
+        ///     </para>
+        ///     <para>
+        ///         <see cref="ExecutionContext"/> is not flowed to the invocation of the callback.
+        ///     </para>
+        ///     <note>
+        ///         The buffers specified in <paramref name="pinData"/> are pinned for the duration of
+        ///         the I/O operation.
+        ///     </note>
+        /// </remarks>
+        /// <exception cref="ArgumentNullException">
+        ///     <paramref name="callback"/> is <see langword="null"/>.
+        /// </exception>
+        /// <exception cref="ObjectDisposedException">
+        ///     This method was called after the <see cref="ThreadPoolBoundHandle"/> was disposed.
+        /// </exception>
+        [CLSCompliant(false)]
+        public unsafe NativeOverlapped* UnsafeAllocateNativeOverlapped(IOCompletionCallback callback, object? state, object? pinData) =>
+            AllocateNativeOverlapped(callback, state, pinData, flowExecutionContext: false);
+
+        private unsafe NativeOverlapped* AllocateNativeOverlapped(IOCompletionCallback callback, object? state, object? pinData, bool flowExecutionContext)
         {
             if (callback == null)
                 throw new ArgumentNullException(nameof(callback));
 
             EnsureNotDisposed();
 
-            ThreadPoolBoundHandleOverlapped overlapped = new ThreadPoolBoundHandleOverlapped(callback, state, pinData, preAllocated: null);
+            ThreadPoolBoundHandleOverlapped overlapped = new ThreadPoolBoundHandleOverlapped(callback, state, pinData, preAllocated: null, flowExecutionContext);
             overlapped._boundHandle = this;
             return overlapped._nativeOverlapped;
         }

src/coreclr/System.Private.CoreLib/src/System/Threading/ClrThreadPoolBoundHandleOverlapped.cs

@@ -6,29 +6,32 @@ namespace System.Threading
     /// <summary>
     /// Overlapped subclass adding data needed by ThreadPoolBoundHandle.
     /// </summary>
-    internal sealed class ThreadPoolBoundHandleOverlapped : Overlapped
+    internal unsafe sealed class ThreadPoolBoundHandleOverlapped : Overlapped
     {
-        private static readonly unsafe IOCompletionCallback s_completionCallback = CompletionCallback;
+        private static readonly IOCompletionCallback s_completionCallback = CompletionCallback;
 
         private readonly IOCompletionCallback _userCallback;
         internal readonly object? _userState;
-        internal PreAllocatedOverlapped? _preAllocated;
-        internal unsafe NativeOverlapped* _nativeOverlapped;
+        internal readonly PreAllocatedOverlapped? _preAllocated;
+
+        internal NativeOverlapped* _nativeOverlapped;
         internal ThreadPoolBoundHandle? _boundHandle;
         internal bool _completed;
 
-        public unsafe ThreadPoolBoundHandleOverlapped(IOCompletionCallback callback, object? state, object? pinData, PreAllocatedOverlapped? preAllocated)
+        public ThreadPoolBoundHandleOverlapped(IOCompletionCallback callback, object? state, object? pinData, PreAllocatedOverlapped? preAllocated, bool flowExecutionContext)
         {
             _userCallback = callback;
             _userState = state;
             _preAllocated = preAllocated;
 
-            _nativeOverlapped = Pack(s_completionCallback, pinData);
-            _nativeOverlapped->OffsetLow = 0;        // CLR reuses NativeOverlapped instances and does not reset these
+            _nativeOverlapped = flowExecutionContext ?
+                Pack(s_completionCallback, pinData) :
+                UnsafePack(s_completionCallback, pinData);
+            _nativeOverlapped->OffsetLow = 0; // CLR reuses NativeOverlapped instances and does not reset these
             _nativeOverlapped->OffsetHigh = 0;
         }
 
-        private static unsafe void CompletionCallback(uint errorCode, uint numBytes, NativeOverlapped* nativeOverlapped)
+        private static void CompletionCallback(uint errorCode, uint numBytes, NativeOverlapped* nativeOverlapped)
         {
             ThreadPoolBoundHandleOverlapped overlapped = (ThreadPoolBoundHandleOverlapped)Overlapped.Unpack(nativeOverlapped);
 

src/coreclr/System.Private.CoreLib/src/System/Threading/ClrThreadPoolPreAllocatedOverlapped.cs

@@ -47,12 +47,56 @@ public sealed class PreAllocatedOverlapped : IDisposable, IDeferredDisposable
         ///     This method was called after the <see cref="ThreadPoolBoundHandle"/> was disposed.
         /// </exception>
         [CLSCompliant(false)]
-        public PreAllocatedOverlapped(IOCompletionCallback callback, object? state, object? pinData)
+        public PreAllocatedOverlapped(IOCompletionCallback callback, object? state, object? pinData) :
+            this(callback, state, pinData, flowExecutionContext: true)
+        {
+        }
+
+        /// <summary>
+        ///     Initializes a new instance of the <see cref="PreAllocatedOverlapped"/> class, specifying
+        ///     a delegate that is invoked when each asynchronous I/O operation is complete, a user-provided
+        ///     object providing context, and managed objects that serve as buffers.
+        /// </summary>
+        /// <param name="callback">
+        ///     An <see cref="IOCompletionCallback"/> delegate that represents the callback method
+        ///     invoked when each asynchronous I/O operation completes.
+        /// </param>
+        /// <param name="state">
+        ///     A user-provided object that distinguishes <see cref="NativeOverlapped"/> instance produced from this
+        ///     object from other <see cref="NativeOverlapped"/> instances. Can be <see langword="null"/>.
+        /// </param>
+        /// <param name="pinData">
+        ///     An object or array of objects representing the input or output buffer for the operations. Each
+        ///     object represents a buffer, for example an array of bytes.  Can be <see langword="null"/>.
+        /// </param>
+        /// <remarks>
+        ///     The new <see cref="PreAllocatedOverlapped"/> instance can be passed to
+        ///     <see cref="ThreadPoolBoundHandle.AllocateNativeOverlapped(PreAllocatedOverlapped)"/>, to produce
+        ///     a <see cref="NativeOverlapped"/> instance that can be passed to the operating system in overlapped
+        ///     I/O operations.  A single <see cref="PreAllocatedOverlapped"/> instance can only be used for
+        ///     a single native I/O operation at a time.  However, the state stored in the <see cref="PreAllocatedOverlapped"/>
+        ///     instance can be reused for subsequent native operations. ExecutionContext is not flowed to the invocation
+        ///     of the callback.
+        ///     <note>
+        ///         The buffers specified in <paramref name="pinData"/> are pinned until <see cref="Dispose"/> is called.
+        ///     </note>
+        /// </remarks>
+        /// <exception cref="ArgumentNullException">
+        ///     <paramref name="callback"/> is <see langword="null"/>.
+        /// </exception>
+        /// <exception cref="ObjectDisposedException">
+        ///     This method was called after the <see cref="ThreadPoolBoundHandle"/> was disposed.
+        /// </exception>
+        [CLSCompliant(false)]
+        public static PreAllocatedOverlapped UnsafeCreate(IOCompletionCallback callback, object? state, object? pinData) =>
+            new PreAllocatedOverlapped(callback, state, pinData, flowExecutionContext: false);
+
+        private PreAllocatedOverlapped(IOCompletionCallback callback, object? state, object? pinData, bool flowExecutionContext)
         {
             if (callback == null)
                 throw new ArgumentNullException(nameof(callback));
 
-            _overlapped = new ThreadPoolBoundHandleOverlapped(callback, state, pinData, this);
+            _overlapped = new ThreadPoolBoundHandleOverlapped(callback, state, pinData, this, flowExecutionContext);
         }
 
         internal bool AddRef()

src/libraries/System.Net.Sockets/src/System/Net/Sockets/SocketAsyncEventArgs.Windows.cs

@@ -83,20 +83,9 @@ private enum PinState : byte { None = 0, MultipleBuffer, SendPackets }
         [MemberNotNull(nameof(_preAllocatedOverlapped))]
         private void InitializeInternals()
         {
-            // PreAllocatedOverlapped captures ExecutionContext, but SocketAsyncEventArgs ensures
-            // that context is properly flowed if necessary, and thus we don't need the overlapped
-            // infrastructure capturing and flowing as well.
-            bool suppressFlow = !ExecutionContext.IsFlowSuppressed();
-            try
-            {
-                Debug.Assert(OperatingSystem.IsWindows());
-                if (suppressFlow) ExecutionContext.SuppressFlow();
-                _preAllocatedOverlapped = new PreAllocatedOverlapped(s_completionPortCallback, _strongThisRef, null);
-            }
-            finally
-            {
-                if (suppressFlow) ExecutionContext.RestoreFlow();
-            }
+            Debug.Assert(OperatingSystem.IsWindows());
+
+            _preAllocatedOverlapped = PreAllocatedOverlapped.UnsafeCreate(s_completionPortCallback, _strongThisRef, null);
 
             if (NetEventSource.Log.IsEnabled()) NetEventSource.Info(this, $"new PreAllocatedOverlapped {_preAllocatedOverlapped}");
         }

src/libraries/System.Private.CoreLib/src/System/IO/Strategies/AsyncWindowsFileStreamStrategy.ValueTaskSource.cs

@@ -32,7 +32,7 @@ internal ValueTaskSource(AsyncWindowsFileStreamStrategy strategy)
             {
                 _strategy = strategy;
                 _source.RunContinuationsAsynchronously = true;
-                _preallocatedOverlapped = new PreAllocatedOverlapped(s_ioCallback, this, null);
+                _preallocatedOverlapped = PreAllocatedOverlapped.UnsafeCreate(s_ioCallback, this, null);
             }
 
             internal void Dispose()

src/libraries/System.Private.CoreLib/src/System/IO/Strategies/Net5CompatFileStreamStrategy.Windows.cs

@@ -544,7 +544,7 @@ partial void OnBufferAllocated()
             Debug.Assert(_preallocatedOverlapped == null);
 
             if (_useAsyncIO)
-                _preallocatedOverlapped = new PreAllocatedOverlapped(CompletionSource.s_ioCallback, this, _buffer);
+                _preallocatedOverlapped = PreAllocatedOverlapped.UnsafeCreate(CompletionSource.s_ioCallback, this, _buffer);
         }
 
         private CompletionSource? CompareExchangeCurrentOverlappedOwner(CompletionSource? newSource, CompletionSource? existingSource)

src/libraries/System.Private.CoreLib/src/System/Threading/ThreadPoolBoundHandle.PlatformNotSupported.cs

@@ -33,6 +33,10 @@ public static ThreadPoolBoundHandle BindHandle(SafeHandle handle)
             throw new PlatformNotSupportedException(SR.PlatformNotSupported_OverlappedIO);
         }
 
+        [CLSCompliant(false)]
+        public unsafe NativeOverlapped* UnsafeAllocateNativeOverlapped(IOCompletionCallback callback, object? state, object? pinData) =>
+            AllocateNativeOverlapped(callback, state, pinData);
+
         [CLSCompliant(false)]
         public unsafe NativeOverlapped* AllocateNativeOverlapped(PreAllocatedOverlapped preAllocated)
         {

src/libraries/System.Threading.Overlapped/ref/System.Threading.Overlapped.cs

@@ -49,6 +49,8 @@ public sealed partial class PreAllocatedOverlapped : System.IDisposable
         public PreAllocatedOverlapped(System.Threading.IOCompletionCallback callback, object? state, object? pinData) { }
         public void Dispose() { }
         ~PreAllocatedOverlapped() { }
+        [System.CLSCompliantAttribute(false)]
+        public static System.Threading.PreAllocatedOverlapped UnsafeCreate(System.Threading.IOCompletionCallback callback, object? state, object? pinData) { throw null; }
     }
     public sealed partial class ThreadPoolBoundHandle : System.IDisposable
     {
@@ -64,5 +66,7 @@ public void Dispose() { }
         public unsafe void FreeNativeOverlapped(System.Threading.NativeOverlapped* overlapped) { }
         [System.CLSCompliantAttribute(false)]
         public unsafe static object? GetNativeOverlappedState(System.Threading.NativeOverlapped* overlapped) { throw null; }
+        [System.CLSCompliantAttribute(false)]
+        public unsafe System.Threading.NativeOverlapped* UnsafeAllocateNativeOverlapped(System.Threading.IOCompletionCallback callback, object? state, object? pinData) { throw null; }
     }
 }