Removed all state machine creation for sync path in message reader.

Author: mattnischan

src/Bedrock.Framework/Protocols/WebSockets/WebSocketFrameReader.cs

@@ -13,6 +13,11 @@ namespace Bedrock.Framework.Protocols.WebSockets
     /// </summary>
     public class WebSocketFrameReader : IMessageReader<WebSocketReadFrame>
     {
+        /// <summary>
+        /// An instance of the WebSocketFrameReader.
+        /// </summary>
+        private WebSocketPayloadReader _payloadReader;
+
         /// <summary>
         /// Attempts to parse a message from a sequence.
         /// </summary>
@@ -30,7 +35,7 @@ public bool TryParseMessage(in ReadOnlySequence<byte> input, ref SequencePositio
                 return false;
             }
 
-            if (input.IsSingleSegment)
+            if (input.IsSingleSegment || input.FirstSpan.Length >= 14)
             {
                 if (TryParseSpan(input.FirstSpan, input.Length, out var bytesRead, out message))
                 {
@@ -123,8 +128,17 @@ private bool TryParseSpan(in ReadOnlySpan<byte> span, long inputLength, out int
             }
 
             var header = new WebSocketHeader(fin, opcode, masked, payloadLength, maskingKey);
-            message = new WebSocketReadFrame(header, new WebSocketPayloadReader(header));
 
+            if(_payloadReader == null)
+            {
+                _payloadReader = new WebSocketPayloadReader(header);
+            }
+            else
+            {
+                _payloadReader.Reset(header);
+            }
+
+            message = new WebSocketReadFrame(header, _payloadReader);
             bytesRead = 2 + extendedPayloadLengthSize + maskSize;
             return true;
         }

src/Bedrock.Framework/Protocols/WebSockets/WebSocketMessageReader.cs

@@ -101,62 +101,91 @@ public WebSocketMessageReader(PipeReader transport, IControlFrameHandler control
         /// </summary>
         /// <param name="cancellationToken">A cancellation token, if any.</param>
         /// <returns>A message read result.</returns>
-        public async ValueTask<MessageReadResult> ReadAsync(CancellationToken cancellationToken = default)
+        public ValueTask<MessageReadResult> ReadAsync(CancellationToken cancellationToken = default)
         {
             if (_awaitingHeader)
             {
-                var frame = await GetNextMessageFrameAsync(cancellationToken).ConfigureAwait(false);
-                ValidateHeader(frame.Header);
+                var readTask = GetNextMessageFrameAsync(cancellationToken);
+                if(readTask.IsCompletedSuccessfully)
+                {
+                    var frame = readTask.Result;
+                    ValidateHeader(frame.Header);
 
-                _header = frame.Header;
-                _payloadReader = frame.Payload;
+                    _header = frame.Header;
+                    _payloadReader = frame.Payload;
+                }
+                else
+                {
+                    return DoReadHeaderRequiredAsync(readTask, cancellationToken);
+                }
             }
 
+            return ReadPayloadAsync(cancellationToken);
+        }
+
+        /// <summary>
+        /// Completes an async read when reading a header is required.
+        /// </summary>
+        /// <param name="readTask">The active async read task from the ProtocolReader.</param>
+        /// <param name="cancellationToken">A cancellation token.</param>
+        /// <returns>A MessageReadResult.</returns>
+        private async ValueTask<MessageReadResult> DoReadHeaderRequiredAsync(ValueTask<WebSocketReadFrame> readTask, CancellationToken cancellationToken)
+        {
+            var frame = await readTask.ConfigureAwait(false);
+
+            ValidateHeader(frame.Header);
+
+            _header = frame.Header;
+            _payloadReader = frame.Payload;
+
+            return await ReadPayloadAsync(cancellationToken);
+        }
+
+        /// <summary>
+        /// Reads a portion of a message payload.
+        /// </summary>
+        /// <param name="cancellationToken">A cancellation token.</param>
+        /// <returns>A MessageReadResult.</returns>
+        private ValueTask<MessageReadResult> ReadPayloadAsync(CancellationToken cancellationToken)
+        {
             //Don't keep reading data into the buffer if we've hit a threshold
             //TODO: Is this even the right value to use in this context?
             if (_buffer.UnconsumedWrittenCount < _options.PauseWriterThreshold)
             {
                 var readTask = _protocolReader.ReadAsync(_payloadReader, cancellationToken);
-                ProtocolReadResult<ReadOnlySequence<byte>> payloadSequence;
-
                 if (readTask.IsCompletedSuccessfully)
                 {
-                    payloadSequence = readTask.Result;
+                    PopulateFromRead(readTask.Result);
                 }
                 else
                 {
-                    payloadSequence = await readTask;
-                }
-
-                if (payloadSequence.IsCanceled)
-                {
-                    throw new OperationCanceledException("Read canceled while attempting to read WebSocket payload.");
-                }
-
-                var sequence = payloadSequence.Message;
-
-                //If there is already data in the buffer, we'll need to add to it
-                if (_buffer.UnconsumedWrittenCount > 0)
-                {
-                    if (sequence.IsSingleSegment)
-                    {
-                        _buffer.Write(sequence.FirstSpan);
-                    }
-                    else
-                    {
-                        foreach (var segment in sequence)
-                        {
-                            _buffer.Write(segment.Span);
-                        }
-                    }
-                }
+                    return CreateMessageReadResultAsync(readTask, cancellationToken);
+                }   
+            }
 
-                _currentSequence = payloadSequence.Message;
-                _isCompleted = payloadSequence.IsCompleted;
-                _isCanceled = payloadSequence.IsCanceled;
+            var endOfMessage = _header.Fin && _payloadReader.BytesRemaining == 0;
 
-                _awaitingHeader = _payloadReader.BytesRemaining == 0;
+            //Serve back buffered data, if it exists, else give the direct sequence without buffering
+            if (_buffer.UnconsumedWrittenCount > 0)
+            {
+                return new ValueTask<MessageReadResult>(
+                    new MessageReadResult(new ReadOnlySequence<byte>(_buffer.WrittenMemory), endOfMessage, _isCanceled, _isCompleted));
             }
+            else
+            {
+                return new ValueTask<MessageReadResult>(new MessageReadResult(_currentSequence, endOfMessage, _isCanceled, _isCompleted));
+            }
+        }
+
+        /// <summary>
+        /// Creates a new MessageReadResult asynchronously.
+        /// </summary>
+        /// <param name="readTask">The active read task from the ProtocolReader.</param>
+        /// <param name="cancellationToken">A cancellation token.</param>
+        /// <returns>A new MessageReadResult.</returns>
+        private async ValueTask<MessageReadResult> CreateMessageReadResultAsync(ValueTask<ProtocolReadResult<ReadOnlySequence<byte>>> readTask, CancellationToken cancellationToken)
+        {
+            PopulateFromRead(await readTask);
 
             var endOfMessage = _header.Fin && _payloadReader.BytesRemaining == 0;
 
@@ -171,6 +200,42 @@ public async ValueTask<MessageReadResult> ReadAsync(CancellationToken cancellati
             }
         }
 
+        /// <summary>
+        /// Populates the message reader from a payload read result.
+        /// </summary>
+        /// <param name="readResult">The read result to populate the message reader from.</param>
+        private void PopulateFromRead(ProtocolReadResult<ReadOnlySequence<byte>> readResult)
+        {
+            if (readResult.IsCanceled)
+            {
+                throw new OperationCanceledException("Read canceled while attempting to read WebSocket payload.");
+            }
+
+            var sequence = readResult.Message;
+
+            //If there is already data in the buffer, we'll need to add to it
+            if (_buffer.UnconsumedWrittenCount > 0)
+            {
+                if (sequence.IsSingleSegment)
+                {
+                    _buffer.Write(sequence.FirstSpan);
+                }
+                else
+                {
+                    foreach (var segment in sequence)
+                    {
+                        _buffer.Write(segment.Span);
+                    }
+                }
+            }
+
+            _currentSequence = readResult.Message;
+            _isCompleted = readResult.IsCompleted;
+            _isCanceled = readResult.IsCanceled;
+
+            _awaitingHeader = _payloadReader.BytesRemaining == 0;
+        }
+
         /// <summary>
         /// Advances the reader to the provided position.
         /// </summary>
@@ -223,15 +288,40 @@ public void AdvanceTo(SequencePosition consumed, SequencePosition examined)
         /// </summary>
         /// <param name="cancellationToken">A cancellation token, if any.</param>
         /// <returns>True if the message is text, false otherwise.</returns>
-        public async ValueTask<bool> MoveNextMessageAsync(CancellationToken cancellationToken = default)
+        public ValueTask<bool> MoveNextMessageAsync(CancellationToken cancellationToken = default)
         {
             if (_payloadReader is object && _payloadReader.BytesRemaining != 0)
             {
                 throw new InvalidOperationException("MoveNextMessageAsync cannot be called while a message is still being read.");
             }
 
-            var frame = await GetNextMessageFrameAsync(cancellationToken);
+            var readTask = GetNextMessageFrameAsync(cancellationToken);
+            if (readTask.IsCompletedSuccessfully)
+            {
+                return new ValueTask<bool>(SetNextMessageAndGetIsText(readTask.Result));
+            }
 
+            return DoSetNextMessageAsync(readTask);
+        }
+
+        /// <summary>
+        /// Sets the next message frame asynchronously.
+        /// </summary>
+        /// <param name="readTask">The active ProtocolReader read task.</param>
+        /// <returns>True if the next message is a text message, false otherwise.</returns>
+        private async ValueTask<bool> DoSetNextMessageAsync(ValueTask<WebSocketReadFrame> readTask)
+        {
+            return SetNextMessageAndGetIsText(await readTask);
+        }
+
+        /// <summary>
+        /// Sets the message reader up with the next message frame data and determines if the message
+        /// is a text or binary message.
+        /// </summary>
+        /// <param name="frame">The read frame to set the message reader with.</param>
+        /// <returns>True if the next message is text, false otherwise.</returns>
+        private bool SetNextMessageAndGetIsText(WebSocketReadFrame frame)
+        {
             if (frame.Header.Opcode != WebSocketOpcode.Binary && frame.Header.Opcode != WebSocketOpcode.Text)
             {
                 ThrowBadProtocol($"Expected a start of message frame of Binary or Text but received {frame.Header.Opcode} instead.");

src/Bedrock.Framework/Protocols/WebSockets/WebSocketPayloadEncoder.cs

@@ -13,7 +13,7 @@ namespace Bedrock.Framework.Protocols.WebSockets
     /// Masks or unmasks a WebSocket payload according to the provided masking key, tracking the
     /// masking key index accross mask or unmasking requests.
     /// </summary>
-    internal struct WebSocketPayloadEncoder
+    internal class WebSocketPayloadEncoder
     {
         /// <summary>
         /// The masking key to use to mask or unmask the payload.
@@ -30,6 +30,15 @@ internal struct WebSocketPayloadEncoder
         /// </summary>
         /// <param name="maskingKey">The masking key to use to mask or unmask payloads.</param>
         public WebSocketPayloadEncoder(int maskingKey)
+        {
+            Reset(maskingKey);
+        }
+
+        /// <summary>
+        /// Resets the payload encoder.
+        /// </summary>
+        /// <param name="maskingKey">The masking key to use to mask or unmask payloads.</param>
+        public void Reset(int maskingKey)
         {
             _maskingKey = maskingKey;
             _currentMaskIndex = 0;

src/Bedrock.Framework/Protocols/WebSockets/WebSocketPayloadReader.cs

@@ -43,6 +43,18 @@ public WebSocketPayloadReader(WebSocketHeader header)
             _masked = header.Masked;
         }
 
+        /// <summary>
+        /// Resets the payload reader.
+        /// </summary>
+        /// <param name="header">The WebSocketHeader associated with this payload.</param>
+        public void Reset(WebSocketHeader header)
+        {
+            BytesRemaining = header.PayloadLength;
+            _masked = header.Masked;
+
+            _payloadEncoder.Reset(header.MaskingKey);
+        }
+
         /// <summary>
         /// Attempts to read the WebSocket payload from a sequence.
         /// </summary>

tests/Bedrock.Framework.Benchmarks/Program.cs

@@ -1,5 +1,6 @@
 using BenchmarkDotNet.Running;
 using System;
+using System.Threading.Tasks;
 
 namespace Bedrock.Framework.Benchmarks
 {

tests/Bedrock.Framework.Benchmarks/WebSocketProtocolBenchmarks.cs

@@ -1,5 +1,6 @@
 using Bedrock.Framework.Protocols.WebSockets;
 using BenchmarkDotNet.Attributes;
+using BenchmarkDotNet.Configs;
 using Microsoft.AspNetCore.Connections;
 using System;
 using System.Buffers;
@@ -13,21 +14,27 @@
 
 namespace Bedrock.Framework.Benchmarks
 {
+    [GroupBenchmarksBy(BenchmarkLogicalGroupRule.ByCategory)]
+    [CategoriesColumn]
     public class WebSocketProtocolBenchmarks
     {
-        private WebSocket _webSocket;
+        private WebSocket _webSocketServer;
 
-        private WebSocketProtocol _webSocketProtocol;
+        private WebSocket _webSocketClient;
 
-        private DefaultConnectionContext _connectionContext;
+        private WebSocketProtocol _webSocketProtocolServer;
 
-        private MemoryStream _stream;
+        private WebSocketProtocol _webSocketProtocolClient;
 
-        private byte[] _message;
+        private DefaultConnectionContext _serverConnectionContext;
 
-        private ArraySegment<byte> _arrayBuffer;
+        private DefaultConnectionContext _clientConnectionContext;
+
+        private MemoryStream _serverStream;
 
-        private ReadOnlyMemory<byte> _romBuffer;
+        private MemoryStream _clientStream;
+
+        private ArraySegment<byte> _arrayBuffer;
 
         private class DummyPipeReader : PipeReader
         {
@@ -64,50 +71,85 @@ private class DummyDuplexPipe : IDuplexPipe
 
         [GlobalSetup]
         public async ValueTask Setup()
+        {
+            var serverMessage = await GetMessageBytes(true, 4000);
+            var clientMessage = await GetMessageBytes(false, 4000);
+
+            (_serverConnectionContext, _serverStream) = CreateContextAndStream(serverMessage);
+            (_clientConnectionContext, _clientStream) = CreateContextAndStream(clientMessage);
+
+            _webSocketServer = WebSocket.CreateFromStream(_serverStream, true, null, TimeSpan.FromSeconds(30));
+            _webSocketProtocolServer = new WebSocketProtocol(_serverConnectionContext, WebSocketProtocolType.Server);
+
+            _webSocketClient = WebSocket.CreateFromStream(_clientStream, false, null, TimeSpan.FromSeconds(30));
+            _webSocketProtocolClient = new WebSocketProtocol(_clientConnectionContext, WebSocketProtocolType.Server);
+
+            _arrayBuffer = new ArraySegment<byte>(new byte[10000]);
+        }
+
+        private async ValueTask<byte[]> GetMessageBytes(bool isMasked, long size)
         {
             var writer = new WebSocketFrameWriter();
             var pipe = new Pipe();
 
-            _message = new byte[4000];
-
-            var header = WebSocketHeader.CreateMasked(true, WebSocketOpcode.Binary, 4000);
-            writer.WriteMessage(new WebSocketWriteFrame(header, new ReadOnlySequence<byte>(_message)), pipe.Writer);
+            var header = new WebSocketHeader(true, WebSocketOpcode.Binary, isMasked, (ulong)size, isMasked ? WebSocketHeader.GenerateMaskingKey() : default);
+            writer.WriteMessage(new WebSocketWriteFrame(header, new ReadOnlySequence<byte>(new byte[4000])), pipe.Writer);
 
             await pipe.Writer.FlushAsync();
 
             var result = await pipe.Reader.ReadAsync();
-            _message = result.Buffer.ToArray();
+            return result.Buffer.ToArray();
+        }
 
-            var dummyReader = new DummyPipeReader { Result = new ReadResult(new ReadOnlySequence<byte>(_message), false, false) };
-            var dummyDuplexPipe = new DummyDuplexPipe { DummyReader = dummyReader };
+        private (DefaultConnectionContext context, MemoryStream stream) CreateContextAndStream(byte[] message)
+        {
+            var reader = new DummyPipeReader { Result = new ReadResult(new ReadOnlySequence<byte>(message), false, false) };
+            var duplexPipe = new DummyDuplexPipe { DummyReader = reader };
 
-            _connectionContext = new DefaultConnectionContext { Transport = dummyDuplexPipe };
-            _stream = new MemoryStream(_message);
+            var stream = new MemoryStream(message);
+            var context = new DefaultConnectionContext { Transport = duplexPipe };
 
-            _webSocket = WebSocket.CreateFromStream(_stream, true, null, TimeSpan.FromSeconds(30));
-            _webSocketProtocol = new WebSocketProtocol(_connectionContext, WebSocketProtocolType.Server);
+            return (context, stream);
+        }
 
-            _arrayBuffer = new ArraySegment<byte>(new byte[10000]);
-            _romBuffer = new ReadOnlyMemory<byte>(_message);
+        [BenchmarkCategory("Masked"), Benchmark(Baseline = true)]
+        public async ValueTask WebSocketRead()
+        {
+            _clientStream.Seek(0, SeekOrigin.Begin);
+
+            var endOfMessage = false;
+            while (!endOfMessage)
+            {
+                var result = await _webSocketClient.ReceiveAsync(_arrayBuffer, CancellationToken.None);
+                endOfMessage = result.EndOfMessage;
+            }
         }
 
-        [Benchmark(Baseline = true)]
+        [BenchmarkCategory("Unmasked"), Benchmark(Baseline = true)]
         public async ValueTask WebSocketReadMasked()
         {
-            _stream.Seek(0, SeekOrigin.Begin);
+            _serverStream.Seek(0, SeekOrigin.Begin);
 
             var endOfMessage = false;
             while (!endOfMessage)
             {
-                var result = await _webSocket.ReceiveAsync(_arrayBuffer, CancellationToken.None);
+                var result = await _webSocketServer.ReceiveAsync(_arrayBuffer, CancellationToken.None);
                 endOfMessage = result.EndOfMessage;
             }
         }
 
-        [Benchmark]
+        [BenchmarkCategory("Masked"), Benchmark]
+        public async ValueTask WebSocketProtocolRead()
+        {
+            var message = await _webSocketProtocolClient.ReadAsync();
+            var data = await message.Reader.ReadAsync();
+            message.Reader.AdvanceTo(data.Data.End);
+        }
+
+        [BenchmarkCategory("Unmasked"), Benchmark]
         public async ValueTask WebSocketProtocolReadMasked()
         {
-            var message = await _webSocketProtocol.ReadAsync();
+            var message = await _webSocketProtocolServer.ReadAsync();
             var data = await message.Reader.ReadAsync();
             message.Reader.AdvanceTo(data.Data.End);
         }