Reduce buffering in ManagedWebSocket.ReceiveAsync (#28577)

* Change SocketsHttpHandler's 101 response stream to buffer response reads

The connection already has a buffer; use it.

* Reduce buffering in ManagedWebSocket.ReceiveAsync

Today we use a big-ish receive buffer, reading all data into the receive buffer and then copying from there into the caller-provided destination buffer.  With this change, we assume that the underlying stream is providing the desired level of buffering, and thus only use a buffer small enough to handle message headers and control payloads (the latter for simplicity, as no caller-supplied buffer is available there).

* Change ClientWebSocket to avoid creating a large buffer

Author: stephentoub

src/Common/src/System/Net/WebSockets/ManagedWebSocket.cs

@@ -63,8 +63,6 @@ public static ManagedWebSocket CreateFromConnectedStream(
         private const int MaxControlPayloadLength = 125;
         /// <summary>Length of the mask XOR'd with the payload data.</summary>
         private const int MaskLength = 4;
-        /// <summary>Default length of a receive buffer to create when an invalid scratch buffer is provided.</summary>
-        private const int DefaultReceiveBufferSize = 0x1000;
 
         /// <summary>The stream used to communicate with the remote server.</summary>
         private readonly Stream _stream;
@@ -184,8 +182,11 @@ private ManagedWebSocket(Stream stream, bool isServer, string subprotocol, TimeS
             // socket rents a similarly sized buffer from the pool for its duration, we'll end up draining
             // the pool, such that other web sockets will allocate anyway, as will anyone else in the process using the
             // pool.  If someone wants to pool, they can do so by passing in the buffer they want to use, and they can
-            // get it from whatever pool they like.
-            _receiveBuffer = buffer.Length >= MaxMessageHeaderLength ? buffer : new byte[DefaultReceiveBufferSize];
+            // get it from whatever pool they like.  If we create our own buffer, it's small, large enough for message
+            // headers and control payloads, but data for other message payloads is read directly into the buffers
+            // passed into ReceiveAsync.
+            const int ReceiveBufferMinLength = MaxControlPayloadLength;
+            _receiveBuffer = buffer.Length >= ReceiveBufferMinLength ? buffer : new byte[ReceiveBufferMinLength];
 
             // Set up the abort source so that if it's triggered, we transition the instance appropriately.
             _abortSource.Token.Register(s =>
@@ -697,32 +698,56 @@ private async ValueTask<TWebSocketReceiveResult> ReceiveAsyncPrivate<TWebSocketR
                     }
 
                     // Otherwise, read as much of the payload as we can efficiently, and upate the header to reflect how much data
-                    // remains for future reads.
-                    int bytesToCopy = Math.Min(payloadBuffer.Length, (int)Math.Min(header.PayloadLength, _receiveBuffer.Length));
-                    Debug.Assert(bytesToCopy > 0, $"Expected {nameof(bytesToCopy)} > 0");
-                    if (_receiveBufferCount < bytesToCopy)
+                    // remains for future reads.  We first need to copy any data that may be lingering in the receive buffer
+                    // into the destination; then to minimize ReceiveAsync calls, we want to read as much as we can, stopping
+                    // only when we've either read the whole message or when we've filled the payload buffer.
+
+                    // First copy any data lingering in the receive buffer.
+                    int totalBytesReceived = 0;
+                    if (_receiveBufferCount > 0)
                     {
-                        await EnsureBufferContainsAsync(bytesToCopy, throwOnPrematureClosure: true).ConfigureAwait(false);
+                        int receiveBufferBytesToCopy = Math.Min(payloadBuffer.Length, (int)Math.Min(header.PayloadLength, _receiveBufferCount));
+                        Debug.Assert(receiveBufferBytesToCopy > 0);
+                        _receiveBuffer.Span.Slice(_receiveBufferOffset, receiveBufferBytesToCopy).CopyTo(payloadBuffer.Span);
+                        ConsumeFromBuffer(receiveBufferBytesToCopy);
+                        totalBytesReceived += receiveBufferBytesToCopy;
+                        Debug.Assert(
+                            _receiveBufferCount == 0 ||
+                            totalBytesReceived == payloadBuffer.Length ||
+                            totalBytesReceived == header.PayloadLength);
+                    }
+
+                    // Then read directly into the payload buffer until we've hit a limit.
+                    while (totalBytesReceived < payloadBuffer.Length &&
+                           totalBytesReceived < header.PayloadLength)
+                    {
+                        int numBytesRead = await _stream.ReadAsync(payloadBuffer.Slice(
+                            totalBytesReceived,
+                            (int)Math.Min(payloadBuffer.Length, header.PayloadLength) - totalBytesReceived)).ConfigureAwait(false);
+                        if (numBytesRead <= 0)
+                        {
+                            ThrowIfEOFUnexpected(throwOnPrematureClosure: true);
+                            break;
+                        }
+                        totalBytesReceived += numBytesRead;
                     }
 
                     if (_isServer)
                     {
-                        _receivedMaskOffsetOffset = ApplyMask(_receiveBuffer.Span.Slice(_receiveBufferOffset, bytesToCopy), header.Mask, _receivedMaskOffsetOffset);
+                        _receivedMaskOffsetOffset = ApplyMask(payloadBuffer.Span.Slice(0, totalBytesReceived), header.Mask, _receivedMaskOffsetOffset);
                     }
-                    _receiveBuffer.Span.Slice(_receiveBufferOffset, bytesToCopy).CopyTo(payloadBuffer.Span.Slice(0, bytesToCopy));
-                    ConsumeFromBuffer(bytesToCopy);
-                    header.PayloadLength -= bytesToCopy;
+                    header.PayloadLength -= totalBytesReceived;
 
                     // If this a text message, validate that it contains valid UTF8.
                     if (header.Opcode == MessageOpcode.Text &&
-                        !TryValidateUtf8(payloadBuffer.Span.Slice(0, bytesToCopy), header.Fin, _utf8TextState))
+                        !TryValidateUtf8(payloadBuffer.Span.Slice(0, totalBytesReceived), header.Fin, _utf8TextState))
                     {
                         await CloseWithReceiveErrorAndThrowAsync(WebSocketCloseStatus.InvalidPayloadData, WebSocketError.Faulted).ConfigureAwait(false);
                     }
 
                     _lastReceiveHeader = header;
                     return resultGetter.GetResult(
-                        bytesToCopy,
+                        totalBytesReceived,
                         header.Opcode == MessageOpcode.Text ? WebSocketMessageType.Text : WebSocketMessageType.Binary,
                         header.Fin && header.PayloadLength == 0,
                         null, null);
@@ -1165,27 +1190,32 @@ private async Task EnsureBufferContainsAsync(int minimumRequiredBytes, bool thro
                 {
                     int numRead = await _stream.ReadAsync(_receiveBuffer.Slice(_receiveBufferCount, _receiveBuffer.Length - _receiveBufferCount), default).ConfigureAwait(false);
                     Debug.Assert(numRead >= 0, $"Expected non-negative bytes read, got {numRead}");
-                    _receiveBufferCount += numRead;
-                    if (numRead == 0)
+                    if (numRead <= 0)
                     {
-                        // The connection closed before we were able to read everything we needed.
-                        // If it was due to use being disposed, fail.  If it was due to the connection
-                        // being closed and it wasn't expected, fail.  If it was due to the connection
-                        // being closed and that was expected, exit gracefully.
-                        if (_disposed)
-                        {
-                            throw new ObjectDisposedException(nameof(WebSocket));
-                        }
-                        else if (throwOnPrematureClosure)
-                        {
-                            throw new WebSocketException(WebSocketError.ConnectionClosedPrematurely);
-                        }
+                        ThrowIfEOFUnexpected(throwOnPrematureClosure);
                         break;
                     }
+                    _receiveBufferCount += numRead;
                 }
             }
         }
 
+        private void ThrowIfEOFUnexpected(bool throwOnPrematureClosure)
+        {
+            // The connection closed before we were able to read everything we needed.
+            // If it was due to us being disposed, fail.  If it was due to the connection
+            // being closed and it wasn't expected, fail.  If it was due to the connection
+            // being closed and that was expected, exit gracefully.
+            if (_disposed)
+            {
+                throw new ObjectDisposedException(nameof(WebSocket));
+            }
+            if (throwOnPrematureClosure)
+            {
+                throw new WebSocketException(WebSocketError.ConnectionClosedPrematurely);
+            }
+        }
+
         /// <summary>Gets a send buffer from the pool.</summary>
         private void AllocateSendBuffer(int minLength)
         {

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

@@ -1296,6 +1296,52 @@ private async ValueTask<int> ReadAsync(Memory<byte> destination)
             return count;
         }
 
+        private ValueTask<int> ReadBufferedAsync(Memory<byte> destination)
+        {
+            // If the caller provided buffer, and thus the amount of data desired to be read,
+            // is larger than the internal buffer, there's no point going through the internal
+            // buffer, so just do an unbuffered read.
+            return destination.Length >= _readBuffer.Length ?
+                ReadAsync(destination) :
+                ReadBufferedAsyncCore(destination);
+        }
+
+        private async ValueTask<int> ReadBufferedAsyncCore(Memory<byte> destination)
+        {
+            // This is called when reading the response body.
+
+            int remaining = _readLength - _readOffset;
+            if (remaining > 0)
+            {
+                // We have data in the read buffer.  Return it to the caller.
+                if (destination.Length <= remaining)
+                {
+                    ReadFromBuffer(destination.Span);
+                    return destination.Length;
+                }
+                else
+                {
+                    ReadFromBuffer(destination.Span.Slice(0, remaining));
+                    return remaining;
+                }
+            }
+
+            // No data in read buffer. 
+            _readOffset = _readLength = 0;
+
+            // Do a buffered read directly against the underlying stream.
+            Debug.Assert(_readAheadTask == null, "Read ahead task should have been consumed as part of the headers.");
+            int bytesRead = await _stream.ReadAsync(_readBuffer.AsMemory()).ConfigureAwait(false);
+            if (NetEventSource.IsEnabled) Trace($"Received {bytesRead} bytes.");
+            _readLength = bytesRead;
+
+            // Hand back as much data as we can fit.
+            int bytesToCopy = Math.Min(bytesRead, destination.Length);
+            _readBuffer.AsSpan(0, bytesToCopy).CopyTo(destination.Span);
+            _readOffset = bytesToCopy;
+            return bytesToCopy;
+        }
+
         private async Task CopyFromBufferAsync(Stream destination, int count, CancellationToken cancellationToken)
         {
             Debug.Assert(count <= _readLength - _readOffset);

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

@@ -26,7 +26,7 @@ public override async ValueTask<int> ReadAsync(Memory<byte> buffer, Cancellation
                     return 0;
                 }
 
-                ValueTask<int> readTask = _connection.ReadAsync(buffer);
+                ValueTask<int> readTask = _connection.ReadBufferedAsync(buffer);
                 int bytesRead;
                 if (readTask.IsCompletedSuccessfully)
                 {

src/System.Net.WebSockets.Client/src/System/Net/WebSockets/WebSocketHandle.Managed.cs

@@ -204,12 +204,12 @@ public async Task ConnectAsyncCore(Uri uri, CancellationToken cancellationToken,
                 }
 
                 // Get or create the buffer to use
-                const int MinBufferSize = 14; // from ManagedWebSocket.MaxMessageHeaderLength
+                const int MinBufferSize = 125; // from ManagedWebSocket.MaxControlPayloadLength
                 ArraySegment<byte> optionsBuffer = options.Buffer.GetValueOrDefault();
                 Memory<byte> buffer =
                     optionsBuffer.Count >= MinBufferSize ? optionsBuffer : // use the provided buffer if it's big enough
-                    options.ReceiveBufferSize >= MinBufferSize ? new byte[options.ReceiveBufferSize] : // or use the requested size if it's big enough
-                    Memory<byte>.Empty; // or let WebSocket.CreateFromStream use its default
+                    default; // or let WebSocket.CreateFromStream use its default
+                    // options.ReceiveBufferSize is ignored, as we rely on the buffer inside the SocketsHttpHandler stream
 
                 // Get the response stream and wrap it in a web socket.
                 Stream connectedStream = await response.Content.ReadAsStreamAsync().ConfigureAwait(false);