From cef755fd8251a1a869b7b892603adb772993546b Mon Sep 17 00:00:00 2001
From: Tom Deseyn <tom.deseyn@gmail.com>
Date: Wed, 18 Mar 2020 22:57:01 +0100
Subject: [PATCH] Add SocketTransportOption to enable/disable WaitForData
 (#19396)

---
 ...re.Server.Kestrel.Transport.Sockets.netcoreapp.cs |  1 +
 .../src/Client/SocketConnectionFactory.cs            |  3 ++-
 .../src/Internal/SocketConnection.cs                 | 12 +++++++++---
 .../src/SocketConnectionListener.cs                  |  3 ++-
 .../Transport.Sockets/src/SocketTransportOptions.cs  |  8 ++++++++
 5 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/src/Servers/Kestrel/Transport.Sockets/ref/Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.netcoreapp.cs b/src/Servers/Kestrel/Transport.Sockets/ref/Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.netcoreapp.cs
index 622010aa63..020e80124d 100644
--- a/src/Servers/Kestrel/Transport.Sockets/ref/Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.netcoreapp.cs
+++ b/src/Servers/Kestrel/Transport.Sockets/ref/Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.netcoreapp.cs
@@ -31,5 +31,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets
         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 { } }
     }
 }
diff --git a/src/Servers/Kestrel/Transport.Sockets/src/Client/SocketConnectionFactory.cs b/src/Servers/Kestrel/Transport.Sockets/src/Client/SocketConnectionFactory.cs
index dc860eb8f5..6f95d242a5 100644
--- a/src/Servers/Kestrel/Transport.Sockets/src/Client/SocketConnectionFactory.cs
+++ b/src/Servers/Kestrel/Transport.Sockets/src/Client/SocketConnectionFactory.cs
@@ -62,7 +62,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets
                 PipeScheduler.ThreadPool,
                 _trace,
                 _options.MaxReadBufferSize,
-                _options.MaxWriteBufferSize);
+                _options.MaxWriteBufferSize,
+                _options.WaitForDataBeforeAllocatingBuffer);
 
             socketConnection.Start();
             return socketConnection;
diff --git a/src/Servers/Kestrel/Transport.Sockets/src/Internal/SocketConnection.cs b/src/Servers/Kestrel/Transport.Sockets/src/Internal/SocketConnection.cs
index 8074739454..4088dab978 100644
--- a/src/Servers/Kestrel/Transport.Sockets/src/Internal/SocketConnection.cs
+++ b/src/Servers/Kestrel/Transport.Sockets/src/Internal/SocketConnection.cs
@@ -33,13 +33,15 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.Internal
         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 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.Internal
             _socket = socket;
             MemoryPool = memoryPool;
             _trace = trace;
+            _waitForData = waitForData;
 
             LocalEndPoint = _socket.LocalEndPoint;
             RemoteEndPoint = _socket.RemoteEndPoint;
@@ -186,8 +189,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.Internal
             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);
diff --git a/src/Servers/Kestrel/Transport.Sockets/src/SocketConnectionListener.cs b/src/Servers/Kestrel/Transport.Sockets/src/SocketConnectionListener.cs
index f829ae5591..4c49a90233 100644
--- a/src/Servers/Kestrel/Transport.Sockets/src/SocketConnectionListener.cs
+++ b/src/Servers/Kestrel/Transport.Sockets/src/SocketConnectionListener.cs
@@ -112,7 +112,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets
                         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();
 
diff --git a/src/Servers/Kestrel/Transport.Sockets/src/SocketTransportOptions.cs b/src/Servers/Kestrel/Transport.Sockets/src/SocketTransportOptions.cs
index 424a4375ae..957876ca59 100644
--- a/src/Servers/Kestrel/Transport.Sockets/src/SocketTransportOptions.cs
+++ b/src/Servers/Kestrel/Transport.Sockets/src/SocketTransportOptions.cs
@@ -16,6 +16,14 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets
         /// </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>