Stop artificially rooting SocketAsyncEventArgs (#28767)

SocketAsyncEventArgs creates an overlapped object that references the SAEA.  That overlapped object creates an async pinning handle in the runtime, which roots the SocketAsyncEventArgs.  The cycle from SAEA->Overlapped->handle->SAEA involving the root means that a dropped SAEA that's not Dispose'd will end up leaking.

This commit fixes that by adding a level of indirection between the handle and the SAEA.  Rather than wrapping the SAEA directly, the handle is given an intermediate object that references the SAEA, and the SAEA then stores itself as a reference into that object only while an active operation is in progress.  Once the operation completes, the reference in that object is nulled out, and the SAEA will no longer be kept alive by the pinning handle.

Author: stephentoub

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

@@ -5,6 +5,7 @@
 using System.Buffers;
 using System.Diagnostics;
 using System.IO;
+using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 using System.Threading;
 
@@ -40,6 +41,7 @@ private enum SingleBufferHandleState : byte { None, InProcess, Set }
 
         // Overlapped object related variables.
         private PreAllocatedOverlapped _preAllocatedOverlapped;
+        private readonly StrongBox<SocketAsyncEventArgs> _strongThisRef = new StrongBox<SocketAsyncEventArgs>(); // state for _preAllocatedOverlapped; .Value set to this while operations in flight
 
         private PinState _pinState;
         private enum PinState : byte { None = 0, MultipleBuffer, SendPackets }
@@ -53,7 +55,7 @@ private void InitializeInternals()
             try
             {
                 if (suppressFlow) ExecutionContext.SuppressFlow();
-                _preAllocatedOverlapped = new PreAllocatedOverlapped(s_completionPortCallback, this, null);
+                _preAllocatedOverlapped = new PreAllocatedOverlapped(s_completionPortCallback, _strongThisRef, null);
             }
             finally
             {
@@ -92,6 +94,15 @@ private unsafe void FreeNativeOverlapped(NativeOverlapped* overlapped)
             _currentSocket.SafeHandle.IOCPBoundHandle.FreeNativeOverlapped(overlapped);
         }
 
+        partial void StartOperationCommonCore()
+        {
+            // Store the reference to this instance so that it's kept alive by the preallocated
+            // overlapped during the asynchronous operation and so that it's available in the
+            // I/O completion callback.  Once the operation completes, we null this out so
+            // that the SocketAsyncEventArgs instance isn't kept alive unnecessarily.
+            _strongThisRef.Value = this;
+        }
+
         /// <summary>Handles the result of an IOCP operation.</summary>
         /// <param name="success">true if the operation completed synchronously and successfully; otherwise, false.</param>
         /// <param name="bytesTransferred">The number of bytes transferred, if the operation completed synchronously and successfully.</param>
@@ -1115,6 +1126,7 @@ private SocketError FinishOperationConnect()
 
         private void CompleteCore()
         {
+            _strongThisRef.Value = null; // null out this reference from the overlapped so this isn't kept alive artificially
             if (_singleBufferHandleState != SingleBufferHandleState.None)
             {
                 CompleteCoreSpin();
@@ -1173,7 +1185,10 @@ private void FinishOperationSendPackets()
 
         private static readonly unsafe IOCompletionCallback s_completionPortCallback = delegate (uint errorCode, uint numBytes, NativeOverlapped* nativeOverlapped)
         {
-            var saea = (SocketAsyncEventArgs)ThreadPoolBoundHandle.GetNativeOverlappedState(nativeOverlapped);
+            var saeaBox = (StrongBox<SocketAsyncEventArgs>)ThreadPoolBoundHandle.GetNativeOverlappedState(nativeOverlapped);
+            SocketAsyncEventArgs saea = saeaBox.Value;
+            Debug.Assert(saea != null);
+
             if ((SocketError)errorCode == SocketError.Success)
             {
                 saea.FreeNativeOverlapped(nativeOverlapped);

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

@@ -536,8 +536,12 @@ internal void StartOperationCommon(Socket socket, SocketAsyncOperation operation
             {
                 _context = ExecutionContext.Capture();
             }
+
+            StartOperationCommonCore();
         }
 
+        partial void StartOperationCommonCore();
+
         internal void StartOperationAccept()
         {
             // AcceptEx needs a single buffer that's the size of two native sockaddr buffers with 16

src/System.Net.Sockets/tests/FunctionalTests/SocketAsyncEventArgsTest.netcoreapp.cs

@@ -4,7 +4,10 @@
 
 using System.Buffers;
 using System.Collections.Generic;
+using System.Linq;
+using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
+using System.Threading;
 using Xunit;
 
 namespace System.Net.Sockets.Tests
@@ -192,5 +195,54 @@ public void SetBufferMemory_NonArray_BufferReturnsNull()
                 Assert.Null(saea.Buffer);
             }
         }
+
+        [OuterLoop("Involves GC and finalization")]
+        [Theory]
+        [InlineData(false)]
+        [InlineData(true)]
+        public void Finalizer_InvokedWhenNoLongerReferenced(bool afterAsyncOperation)
+        {
+            var cwt = new ConditionalWeakTable<object, object>();
+
+            for (int i = 0; i < 5; i++) // create several SAEA instances, stored into cwt
+            {
+                CreateSocketAsyncEventArgs();
+
+                void CreateSocketAsyncEventArgs() // separated out so that JIT doesn't extend lifetime of SAEA instances
+                {
+                    var saea = new SocketAsyncEventArgs();
+                    cwt.Add(saea, saea);
+
+                    if (afterAsyncOperation)
+                    {
+                        using (Socket listener = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp))
+                        {
+                            listener.Bind(new IPEndPoint(IPAddress.Loopback, 0));
+                            listener.Listen(1);
+
+                            using (Socket client = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp))
+                            {
+                                saea.RemoteEndPoint = listener.LocalEndPoint;
+                                using (var mres = new ManualResetEventSlim())
+                                {
+                                    saea.Completed += (s, e) => mres.Set();
+                                    if (client.ConnectAsync(saea))
+                                    {
+                                        mres.Wait();
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+
+            Assert.True(SpinWait.SpinUntil(() =>
+            {
+                GC.Collect();
+                GC.WaitForPendingFinalizers();
+                return cwt.Count() == 0; // validate that the cwt becomes empty
+            }, 30_000));
+        }
     }
 }