Further improvements

Author: Stefán J. Sigurðarson

projects/RabbitMQ.Client/client/impl/Connection.cs

@@ -30,6 +30,7 @@
 //---------------------------------------------------------------------------
 
 using System;
+using System.Buffers;
 using System.Collections.Generic;
 using System.IO;
 using System.Net;
@@ -70,7 +71,6 @@ internal sealed class Connection : IConnection
         private volatile bool _running = true;
         private readonly MainSession _session0;
         private SessionManager _sessionManager;
-        internal SemaphoreSlim _writeLock = new SemaphoreSlim(1, 1);
 
         //
         // Heartbeats
@@ -94,8 +94,6 @@ internal sealed class Connection : IConnection
         // errors, otherwise as read timeouts
         public ConsumerWorkService ConsumerWorkService { get; }
 
-        internal IFrameHandler FrameHandler => _frameHandler;
-
         public Connection(IConnectionFactory factory, bool insist, IFrameHandler frameHandler, string clientProvidedName = null)
         {
             ClientProvidedName = clientProvidedName;
@@ -346,7 +344,8 @@ public async ValueTask ClosingLoop()
                 // Wait for response/socket closure or timeout
                 while (!_closed)
                 {
-                    await MainLoopIteration().ConfigureAwait(false);
+                    await _frameHandler.FrameReader.WaitToReadAsync().ConfigureAwait(false);
+                    MainLoopIteration();
                 }
             }
             catch (ObjectDisposedException ode)
@@ -487,9 +486,15 @@ public async Task MainLoop()
                 {
                     while (_running)
                     {
+                        ValueTask<bool> waitToRead = _frameHandler.FrameReader.WaitToReadAsync();
+                        if (!waitToRead.IsCompleted)
+                        {
+                            await waitToRead.ConfigureAwait(false);
+                        }
+
                         try
                         {
-                            await MainLoopIteration().ConfigureAwait(false);
+                            MainLoopIteration();
                         }
                         catch (SoftProtocolException spe)
                         {
@@ -545,51 +550,53 @@ public async Task MainLoop()
             }
         }
 
-        public async ValueTask MainLoopIteration()
+        public void MainLoopIteration()
         {
-            InboundFrame frame = await _frameHandler.ReadFrame().ConfigureAwait(false);
-            NotifyHeartbeatListener();
-
-            bool shallReturn = true;
-            // We have received an actual frame.
-            if (frame.Type == FrameType.FrameHeartbeat)
-            {
-                // Ignore it: we've already just reset the heartbeat
-            }
-            else if (frame.Channel == 0)
-            {
-                // In theory, we could get non-connection.close-ok
-                // frames here while we're quiescing (m_closeReason !=
-                // null). In practice, there's a limited number of
-                // things the server can ask of us on channel 0 -
-                // essentially, just connection.close. That, combined
-                // with the restrictions on pipelining, mean that
-                // we're OK here to handle channel 0 traffic in a
-                // quiescing situation, even though technically we
-                // should be ignoring everything except
-                // connection.close-ok.
-                shallReturn = _session0.HandleFrame(in frame);
-            }
-            else
+            while (_frameHandler.FrameReader.TryRead(out InboundFrame frame))
             {
-                // If we're still m_running, but have a m_closeReason,
-                // then we must be quiescing, which means any inbound
-                // frames for non-zero channels (and any inbound
-                // commands on channel zero that aren't
-                // Connection.CloseOk) must be discarded.
-                if (_closeReason is null)
+                NotifyHeartbeatListener();
+
+                bool shallReturn = true;
+                // We have received an actual frame.
+                if (frame.Type == FrameType.FrameHeartbeat)
                 {
-                    // No close reason, not quiescing the
-                    // connection. Handle the frame. (Of course, the
-                    // Session itself may be quiescing this particular
-                    // channel, but that's none of our concern.)
-                    shallReturn = _sessionManager.Lookup(frame.Channel).HandleFrame(in frame);
+                    // Ignore it: we've already just reset the heartbeat
+                }
+                else if (frame.Channel == 0)
+                {
+                    // In theory, we could get non-connection.close-ok
+                    // frames here while we're quiescing (m_closeReason !=
+                    // null). In practice, there's a limited number of
+                    // things the server can ask of us on channel 0 -
+                    // essentially, just connection.close. That, combined
+                    // with the restrictions on pipelining, mean that
+                    // we're OK here to handle channel 0 traffic in a
+                    // quiescing situation, even though technically we
+                    // should be ignoring everything except
+                    // connection.close-ok.
+                    shallReturn = _session0.HandleFrame(in frame);
+                }
+                else
+                {
+                    // If we're still m_running, but have a m_closeReason,
+                    // then we must be quiescing, which means any inbound
+                    // frames for non-zero channels (and any inbound
+                    // commands on channel zero that aren't
+                    // Connection.CloseOk) must be discarded.
+                    if (_closeReason is null)
+                    {
+                        // No close reason, not quiescing the
+                        // connection. Handle the frame. (Of course, the
+                        // Session itself may be quiescing this particular
+                        // channel, but that's none of our concern.)
+                        shallReturn = _sessionManager.Lookup(frame.Channel).HandleFrame(in frame);
+                    }
                 }
-            }
 
-            if (shallReturn)
-            {
-                frame.ReturnPayload();
+                if (shallReturn)
+                {
+                    frame.ReturnPayload();
+                }
             }
         }
 
@@ -877,12 +884,7 @@ public void HeartbeatWriteTimerCallback(object state)
             {
                 if (!_closed)
                 {
-                    _writeLock.Wait();
-                    Span<byte> span = _frameHandler.PipeWriter.GetSpan(Client.Impl.Framing.Heartbeat.Payload.Length);
-                    Client.Impl.Framing.Heartbeat.Payload.CopyTo(span);
-                    _frameHandler.PipeWriter.Advance(Client.Impl.Framing.Heartbeat.Payload.Length);
-                    _frameHandler.Flush();
-                    _writeLock.Release();
+                    Write(OutgoingFrame.CreateHeartbeat());
                     _heartbeatWriteTimer?.Change((int)_heartbeatTimeSpan.TotalMilliseconds, Timeout.Infinite);
                 }
             }
@@ -919,6 +921,11 @@ public override string ToString()
             return string.Format("Connection({0},{1})", _id, Endpoint);
         }
 
+        public void Write(OutgoingFrame outgoingFrame)
+        {
+            _frameHandler.Write(outgoingFrame);
+        }
+
         public void UpdateSecret(string newSecret, string reason)
         {
             _model0.UpdateSecret(newSecret, reason);

projects/RabbitMQ.Client/client/impl/Frame.cs

@@ -31,12 +31,7 @@
 
 using System;
 using System.Buffers;
-using System.IO;
-using System.IO.Pipelines;
-using System.Net.Sockets;
 using System.Runtime.CompilerServices;
-using System.Runtime.ExceptionServices;
-using System.Threading.Tasks;
 
 using RabbitMQ.Client.Exceptions;
 using RabbitMQ.Util;
@@ -67,6 +62,16 @@ private static int WriteBaseFrame(Span<byte> span, FrameType type, ushort channe
             return StartPayload + 1 + payloadLength;
         }
 
+        internal static class Heartbeat
+        {
+            public const int FrameSize = BaseFrameSize;
+
+            public static int WriteTo(Span<byte> span)
+            {
+                return WriteBaseFrame(span, FrameType.FrameHeartbeat, 0, 0);
+            }
+        }
+
         internal static class Method
         {
             /* +----------+-----------+-----------+
@@ -130,23 +135,6 @@ public static int WriteTo(Span<byte> span, ushort channel, ReadOnlySpan<byte> bo
                 return WriteBaseFrame(span, FrameType.FrameBody, channel, StartBodyArgument - StartPayload + body.Length);
             }
         }
-
-        internal static class Heartbeat
-        {
-            /* Empty frame */
-            public const int FrameSize = BaseFrameSize;
-
-            /// <summary>
-            /// Compiler trick to directly refer to static data in the assembly, see here: https://github.com/dotnet/roslyn/pull/24621
-            /// </summary>
-            internal static ReadOnlySpan<byte> Payload => new byte[]
-            {
-                Constants.FrameHeartbeat,
-                0, 0, // channel
-                0, 0, 0, 0, // payload length
-                Constants.FrameEnd
-            };
-        }
     }
 
     internal readonly struct InboundFrame
@@ -164,96 +152,88 @@ private InboundFrame(FrameType type, int channel, ReadOnlyMemory<byte> payload,
             _rentedArray = rentedArray;
         }
 
-        private static async ValueTask ProcessProtocolHeader(PipeReader reader, ReadOnlyMemory<byte> protocolError)
+        private static void ProcessProtocolHeader(ReadOnlySpan<byte> protocolError)
         {
-            try
-            {
-                byte b1 = protocolError.Span[0];
-                byte b2 = protocolError.Span[1];
-                byte b3 = protocolError.Span[2];
-                if (b1 != 'M' || b2 != 'Q' || b3 != 'P')
-                {
-                    throw new MalformedFrameException("Invalid AMQP protocol header from server");
-                }
-
-                int transportHigh = protocolError.Span[3];
-                int transportLow = protocolError.Span[4];
-                int serverMajor = protocolError.Span[5];
-                int serverMinor = await reader.ReadByteAsync().ConfigureAwait(false); ;
-                throw new PacketNotRecognizedException(transportHigh, transportLow, serverMajor, serverMinor);
-            }
-            catch (EndOfStreamException)
+            byte b1 = protocolError[0];
+            byte b2 = protocolError[1];
+            byte b3 = protocolError[2];
+            if (b1 != 'M' || b2 != 'Q' || b3 != 'P')
             {
-                // Ideally we'd wrap the EndOfStreamException in the
-                // MalformedFrameException, but unfortunately the
-                // design of MalformedFrameException's superclass,
-                // ProtocolViolationException, doesn't permit
-                // this. Fortunately, the call stack in the
-                // EndOfStreamException is largely irrelevant at this
-                // point, so can safely be ignored.
                 throw new MalformedFrameException("Invalid AMQP protocol header from server");
             }
+
+            int transportHigh = protocolError[3];
+            int transportLow = protocolError[4];
+            int serverMajor = protocolError[5];
+            int serverMinor = protocolError[6];
+            throw new PacketNotRecognizedException(transportHigh, transportLow, serverMajor, serverMinor);
         }
 
-        internal static async ValueTask<InboundFrame> ReadFrom(PipeReader reader, Memory<byte> frameHeaderBuffer)
+        internal static bool TryReadFrame(ref ReadOnlySequence<byte> buffer, out InboundFrame frame)
         {
-            // We'll always need to read at least 7 bytes (type + channel + payloadSize) or (type + first 6 bytes of protocol error, see ProcessProtocolHeader).
-            int type;
-            try
+            // We'll always need to read at least 8 bytes (type (1) + channel (2) + payloadSize (4) + end marker (1)) or (8 bytes of protocol error, see ProcessProtocolHeader).
+            if (buffer.Length < 8)
             {
-                await reader.ReadAsync(frameHeaderBuffer.Slice(0, 7)).ConfigureAwait(false);
-
-                type = frameHeaderBuffer.Span[0];
+                frame = default;
+                return false;
             }
-            catch (IOException ioe)
+
+            if (buffer.First.Span[0] == 'A')
             {
-                // If it's a WSAETIMEDOUT SocketException, unwrap it.
-                // This might happen when the limit of half-open connections is
-                // reached.
-                if (ioe.InnerException is SocketException exception && exception.SocketErrorCode == SocketError.TimedOut)
+                // Probably an AMQP protocol header, otherwise meaningless
+                if (buffer.First.Length >= 8)
                 {
-                    ExceptionDispatchInfo.Capture(exception).Throw();
+                    ProcessProtocolHeader(buffer.First.Span.Slice(1, 7));
+                }
+                else
+                {
+                    Span<byte> protocolError = stackalloc byte[7];
+                    buffer.Slice(1, 7).CopyTo(protocolError);
+                    ProcessProtocolHeader(protocolError);
                 }
-
-                throw;
             }
 
-            switch (type)
+            int type = buffer.First.Span[0];
+            int channel;
+            int payloadSize;
+
+            if (buffer.First.Length >= 7)
             {
-                case -1:
-                    throw new EndOfStreamException("Reached the end of the stream. Possible authentication failure.");
-                case 'A':
-                    // Probably an AMQP protocol header, otherwise meaningless
-                    await ProcessProtocolHeader(reader, frameHeaderBuffer.Slice(1, 6)).ConfigureAwait(false);
-                    break;
+                channel = NetworkOrderDeserializer.ReadUInt16(buffer.First.Span.Slice(1, 2));
+                payloadSize = NetworkOrderDeserializer.ReadInt32(buffer.First.Span.Slice(3, 4)); // FIXME - throw exn on unreasonable value
+            }
+            else
+            {
+                Span<byte> headerBytes = stackalloc byte[6];
+                buffer.Slice(1, 6).CopyTo(headerBytes);
+                channel = NetworkOrderDeserializer.ReadUInt16(headerBytes.Slice(0, 2));
+                payloadSize = NetworkOrderDeserializer.ReadInt32(headerBytes.Slice(2, 4)); // FIXME - throw exn on unreasonable value
             }
-
-            int channel = NetworkOrderDeserializer.ReadUInt16(frameHeaderBuffer.Span.Slice(1, 2));
-            int payloadSize = NetworkOrderDeserializer.ReadInt32(frameHeaderBuffer.Span.Slice(3, 4)); // FIXME - throw exn on unreasonable value
 
             const int EndMarkerLength = 1;
-            // Is returned by InboundFrame.ReturnPayload in Connection.MainLoopIteration
             int readSize = payloadSize + EndMarkerLength;
-            byte[] payloadBytes = ArrayPool<byte>.Shared.Rent(readSize);
-            Memory<byte> payloadMemory = payloadBytes.AsMemory(0, readSize);
-            try
-            {
-                await reader.ReadAsync(payloadMemory).ConfigureAwait(false);
-            }
-            catch (Exception e)
+
+            // Do we have enough bytes to read an entire frame (type + channel + payloadSize + payload + end marker)
+            if (buffer.Length < (7 + readSize))
             {
-                // Early EOF.
-                ArrayPool<byte>.Shared.Return(payloadBytes);
-                throw new MalformedFrameException($"Short frame - expected to read {readSize} bytes: {e.Message}");
+                frame = default;
+                return false;
             }
 
+            // Is returned by InboundFrame.ReturnPayload in Connection.MainLoopIteration
+            byte[] payloadBytes = ArrayPool<byte>.Shared.Rent(readSize);
+            Memory<byte> payloadMemory = payloadBytes.AsMemory(0, readSize);
+            ReadOnlySequence<byte> payloadSlice = buffer.Slice(7, readSize);
+            payloadSlice.CopyTo(payloadMemory.Span);
             if (payloadBytes[payloadSize] != Constants.FrameEnd)
             {
                 ArrayPool<byte>.Shared.Return(payloadBytes);
                 throw new MalformedFrameException($"Bad frame end marker: {payloadBytes[payloadSize]}");
             }
 
-            return new InboundFrame((FrameType)type, channel, payloadMemory.Slice(0, payloadSize), payloadBytes);
+            buffer = buffer.Slice(payloadSlice.End);
+            frame = new InboundFrame((FrameType)type, channel, payloadMemory.Slice(0, payloadSize), payloadBytes);
+            return true;
         }
 
         public byte[] TakeoverPayload()