Pool SocketAsyncEventArgs used to establish connections

SocketAsyncEventArgs is a heavy object to use once and throw away.  This change lets us pool and reuse them.

Author: stephentoub

src/Common/src/System/Collections/Concurrent/ConcurrentQueue_Segment.cs

@@ -2,14 +2,13 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
-using System.Collections.Generic;
 using System.Diagnostics;
 using System.Runtime.InteropServices;
 using System.Threading;
 
 namespace System.Collections.Concurrent
 {
-    public partial class ConcurrentQueue<T>
+    partial class ConcurrentQueue<T>
     {
         /// <summary>
         /// Provides a multi-producer, multi-consumer thread-safe bounded segment.  When the queue is full,
@@ -34,8 +33,10 @@ internal sealed class Segment
             internal bool _preservedForObservation;
             /// <summary>Indicates whether the segment has been marked such that no additional items may be enqueued.</summary>
             internal bool _frozenForEnqueues;
+#pragma warning disable 0649 // some builds don't assign to this field
             /// <summary>The segment following this one in the queue, or null if this segment is the last in the queue.</summary>
             internal Segment _nextSegment;
+#pragma warning restore 0649
 
             /// <summary>Creates the segment.</summary>
             /// <param name="boundedLength">

src/System.Net.Http/src/System.Net.Http.csproj

@@ -152,6 +152,9 @@
     <Compile Include="System\Net\Http\SocketsHttpHandler\SocketsHttpHandler.cs" />
     <Compile Include="System\Net\Http\SocketsHttpHandler\RawConnectionStream.cs" />
     <Compile Include="System\Net\Http\SocketsHttpHandler\RedirectHandler.cs" />
+    <Compile Include="$(CommonPath)\System\Collections\Concurrent\ConcurrentQueue_Segment.cs">
+      <Link>Common\System\Collections\Concurrent\ConcurrentQueue_Segment.cs</Link>
+    </Compile>
     <Compile Include="$(CommonPath)\System\Net\NTAuthentication.Common.cs">
       <Link>Common\System\Net\NTAuthentication.Common.cs</Link>
     </Compile>

src/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/ConnectHelper.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+using System.Collections.Concurrent;
 using System.Diagnostics;
 using System.IO;
 using System.Net.Security;
@@ -15,6 +16,11 @@ namespace System.Net.Http
 {
     internal static class ConnectHelper
     {
+        /// <summary>Pool of event args to use to establish connections.</summary>
+        private static readonly ConcurrentQueue<ConnectEventArgs>.Segment s_connectEventArgs =
+            new ConcurrentQueue<ConnectEventArgs>.Segment(
+                ConcurrentQueue<ConnectEventArgs>.Segment.RoundUpToPowerOf2(Environment.ProcessorCount));
+
         /// <summary>
         /// Helper type used by HttpClientHandler when wrapping SocketsHttpHandler to map its
         /// certificate validation callback to the one used by SslStream.
@@ -34,86 +40,100 @@ public CertificateCallbackMapper(Func<HttpRequestMessage, X509Certificate2, X509
 
         public static async ValueTask<(Socket, Stream)> ConnectAsync(string host, int port, CancellationToken cancellationToken)
         {
-            try
+            // Rather than creating a new Socket and calling ConnectAsync on it, we use the static
+            // Socket.ConnectAsync with a SocketAsyncEventArgs, as we can then use Socket.CancelConnectAsync
+            // to cancel it if needed. Rent or allocate one.
+            ConnectEventArgs saea;
+            if (!s_connectEventArgs.TryDequeue(out saea))
             {
-                // Rather than creating a new Socket and calling ConnectAsync on it, we use the static
-                // Socket.ConnectAsync with a SocketAsyncEventArgs, as we can then use Socket.CancelConnectAsync
-                // to cancel it if needed.
-                using (var saea = new BuilderAndCancellationTokenSocketAsyncEventArgs(cancellationToken))
-                {
-                    // Configure which server to which to connect.
-                    saea.RemoteEndPoint = IPAddress.TryParse(host, out IPAddress address) ?
-                        (EndPoint)new IPEndPoint(address, port) :
-                        new DnsEndPoint(host, port);
+                saea = new ConnectEventArgs();
+            }
+            saea.Initialize(cancellationToken);
 
-                    // Hook up a callback that'll complete the Task when the operation completes.
-                    saea.Completed += (s, e) =>
-                    {
-                        var csaea = (BuilderAndCancellationTokenSocketAsyncEventArgs)e;
-                        switch (e.SocketError)
-                        {
-                            case SocketError.Success:
-                                csaea.Builder.SetResult();
-                                break;
-                            case SocketError.OperationAborted:
-                            case SocketError.ConnectionAborted:
-                                if (csaea.CancellationToken.IsCancellationRequested)
-                                {
-                                    csaea.Builder.SetException(CancellationHelper.CreateOperationCanceledException(null, csaea.CancellationToken));
-                                    break;
-                                }
-                                goto default;
-                            default:
-                                csaea.Builder.SetException(new SocketException((int)e.SocketError));
-                                break;
-                        }
-                    };
+            try
+            {
+                // Configure which server to which to connect.
+                saea.RemoteEndPoint = IPAddress.TryParse(host, out IPAddress address) ?
+                    (EndPoint)new IPEndPoint(address, port) :
+                    new DnsEndPoint(host, port);
 
-                    // Initiate the connection.
-                    if (Socket.ConnectAsync(SocketType.Stream, ProtocolType.Tcp, saea))
-                    {
-                        // Connect completing asynchronously. Enable it to be canceled and wait for it.
-                        using (cancellationToken.Register(s => Socket.CancelConnectAsync((SocketAsyncEventArgs)s), saea))
-                        {
-                            await saea.Builder.Task.ConfigureAwait(false);
-                        }
-                    }
-                    else if (saea.SocketError != SocketError.Success)
+                // Initiate the connection.
+                if (Socket.ConnectAsync(SocketType.Stream, ProtocolType.Tcp, saea))
+                {
+                    // Connect completing asynchronously. Enable it to be canceled and wait for it.
+                    using (cancellationToken.Register(s => Socket.CancelConnectAsync((SocketAsyncEventArgs)s), saea))
                     {
-                        // Connect completed synchronously but unsuccessfully.
-                        throw new SocketException((int)saea.SocketError);
+                        await saea.Builder.Task.ConfigureAwait(false);
                     }
-
-                    Debug.Assert(saea.SocketError == SocketError.Success, $"Expected Success, got {saea.SocketError}.");
-                    Debug.Assert(saea.ConnectSocket != null, "Expected non-null socket");
-                   
-                    // Configure the socket and return a stream for it.
-                    Socket socket = saea.ConnectSocket;
-                    socket.NoDelay = true;
-                    return (socket, new NetworkStream(socket, ownsSocket: true));
                 }
+                else if (saea.SocketError != SocketError.Success)
+                {
+                    // Connect completed synchronously but unsuccessfully.
+                    throw new SocketException((int)saea.SocketError);
+                }
+
+                Debug.Assert(saea.SocketError == SocketError.Success, $"Expected Success, got {saea.SocketError}.");
+                Debug.Assert(saea.ConnectSocket != null, "Expected non-null socket");
+
+                // Configure the socket and return a stream for it.
+                Socket socket = saea.ConnectSocket;
+                socket.NoDelay = true;
+                return (socket, new NetworkStream(socket, ownsSocket: true));
             }
             catch (Exception error)
             {
                 throw CancellationHelper.ShouldWrapInOperationCanceledException(error, cancellationToken) ?
                     CancellationHelper.CreateOperationCanceledException(error, cancellationToken) :
                     new HttpRequestException(error.Message, error);
             }
+            finally
+            {
+                // Pool the event args, or if the pool is full, dispose of it.
+                saea.Clear();
+                if (!s_connectEventArgs.TryEnqueue(saea))
+                {
+                    saea.Dispose();
+                }
+            }
         }
 
         /// <summary>SocketAsyncEventArgs that carries with it additional state for a Task builder and a CancellationToken.</summary>
-        private sealed class BuilderAndCancellationTokenSocketAsyncEventArgs : SocketAsyncEventArgs
+        private sealed class ConnectEventArgs : SocketAsyncEventArgs
         {
-            public AsyncTaskMethodBuilder Builder { get; }
-            public CancellationToken CancellationToken { get; }
+            public AsyncTaskMethodBuilder Builder { get; private set; }
+            public CancellationToken CancellationToken { get; private set; }
 
-            public BuilderAndCancellationTokenSocketAsyncEventArgs(CancellationToken cancellationToken)
+            public void Initialize(CancellationToken cancellationToken)
             {
+                CancellationToken = cancellationToken;
                 var b = new AsyncTaskMethodBuilder();
                 var ignored = b.Task; // force initialization
                 Builder = b;
+            }
 
-                CancellationToken = cancellationToken;
+            public void Clear() => CancellationToken = default;
+
+            protected override void OnCompleted(SocketAsyncEventArgs _)
+            {
+                switch (SocketError)
+                {
+                    case SocketError.Success:
+                        Builder.SetResult();
+                        break;
+
+                    case SocketError.OperationAborted:
+                    case SocketError.ConnectionAborted:
+                        if (CancellationToken.IsCancellationRequested)
+                        {
+                            Builder.SetException(CancellationHelper.CreateOperationCanceledException(null, CancellationToken));
+                            break;
+                        }
+                        goto default;
+
+                    default:
+                        Builder.SetException(new SocketException((int)SocketError));
+                        break;
+                }
             }
         }