Add SocketTransportOption to enable/disable WaitForData (#19396)

Author: tmds

src/Servers/Kestrel/Transport.Sockets/ref/Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.netcoreapp.cs

@@ -31,5 +31,6 @@ public SocketTransportOptions() { }
         public long? MaxReadBufferSize { [System.Runtime.CompilerServices.CompilerGeneratedAttribute] get { throw null; } [System.Runtime.CompilerServices.CompilerGeneratedAttribute] set { } }
         public long? MaxWriteBufferSize { [System.Runtime.CompilerServices.CompilerGeneratedAttribute] get { throw null; } [System.Runtime.CompilerServices.CompilerGeneratedAttribute] set { } }
         public bool NoDelay { [System.Runtime.CompilerServices.CompilerGeneratedAttribute] get { throw null; } [System.Runtime.CompilerServices.CompilerGeneratedAttribute] set { } }
+        public bool WaitForDataBeforeAllocatingBuffer { [System.Runtime.CompilerServices.CompilerGeneratedAttribute] get { throw null; } [System.Runtime.CompilerServices.CompilerGeneratedAttribute] set { } }
     }
 }

src/Servers/Kestrel/Transport.Sockets/src/Client/SocketConnectionFactory.cs

@@ -62,7 +62,8 @@ public async ValueTask<ConnectionContext> ConnectAsync(EndPoint endpoint, Cancel
                 PipeScheduler.ThreadPool,
                 _trace,
                 _options.MaxReadBufferSize,
-                _options.MaxWriteBufferSize);
+                _options.MaxWriteBufferSize,
+                _options.WaitForDataBeforeAllocatingBuffer);
 
             socketConnection.Start();
             return socketConnection;

src/Servers/Kestrel/Transport.Sockets/src/Internal/SocketConnection.cs

@@ -33,13 +33,15 @@ internal sealed class SocketConnection : TransportConnection
         private Task _processingTask;
         private readonly TaskCompletionSource<object> _waitForConnectionClosedTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
         private bool _connectionClosed;
+        private readonly bool _waitForData;
 
         internal SocketConnection(Socket socket,
                                   MemoryPool<byte> memoryPool,
                                   PipeScheduler scheduler,
                                   ISocketsTrace trace,
                                   long? maxReadBufferSize = null,
-                                  long? maxWriteBufferSize = null)
+                                  long? maxWriteBufferSize = null,
+                                  bool waitForData = true)
         {
             Debug.Assert(socket != null);
             Debug.Assert(memoryPool != null);
@@ -48,6 +50,7 @@ internal SocketConnection(Socket socket,
             _socket = socket;
             MemoryPool = memoryPool;
             _trace = trace;
+            _waitForData = waitForData;
 
             LocalEndPoint = _socket.LocalEndPoint;
             RemoteEndPoint = _socket.RemoteEndPoint;
@@ -186,8 +189,11 @@ private async Task ProcessReceives()
             var input = Input;
             while (true)
             {
-                // Wait for data before allocating a buffer.
-                await _receiver.WaitForDataAsync();
+                if (_waitForData)
+                {
+                    // Wait for data before allocating a buffer.
+                    await _receiver.WaitForDataAsync();
+                }
 
                 // Ensure we have some reasonable amount of buffer space
                 var buffer = input.GetMemory(MinAllocBufferSize);

src/Servers/Kestrel/Transport.Sockets/src/SocketConnectionListener.cs

@@ -112,7 +112,8 @@ public async ValueTask<ConnectionContext> AcceptAsync(CancellationToken cancella
                         acceptSocket.NoDelay = _options.NoDelay;
                     }
 
-                    var connection = new SocketConnection(acceptSocket, _memoryPool, _schedulers[_schedulerIndex], _trace, _options.MaxReadBufferSize, _options.MaxWriteBufferSize);
+                    var connection = new SocketConnection(acceptSocket, _memoryPool, _schedulers[_schedulerIndex], _trace,
+                        _options.MaxReadBufferSize, _options.MaxWriteBufferSize, _options.WaitForDataBeforeAllocatingBuffer);
 
                     connection.Start();
 

src/Servers/Kestrel/Transport.Sockets/src/SocketTransportOptions.cs

@@ -16,6 +16,14 @@ public class SocketTransportOptions
         /// </remarks>
         public int IOQueueCount { get; set; } = Math.Min(Environment.ProcessorCount, 16);
 
+        /// <summary>
+        /// Wait until there is data available to allocate a buffer. Setting this to false can increase throughput at the cost of increased memory usage.
+        /// </summary>
+        /// <remarks>
+        /// Defaults to true.
+        /// </remarks>
+        public bool WaitForDataBeforeAllocatingBuffer { get; set; } = true;
+
         /// <summary>
         /// Set to false to enable Nagle's algorithm for all connections.
         /// </summary>