Run the Consumer MessageHandler in a Task (#250)

* Make the message handler async, avoiding blocking the socket thread. * reduce the initial credits to two * add the documentation * Use `Random.Shared` * Reduce the noise log when the consumer is removed --------- Signed-off-by: Gabriele Santomaggio <[email protected]> Co-authored-by: Luke Bakken <[email protected]>
rabbitmq · Mar 29, 2023 · c1b9961 · c1b9961
1 parent f966294
commit c1b9961
Show file tree

Hide file tree

Showing 15 changed files with 325 additions and 111 deletions.
diff --git a/RabbitMQ.Stream.Client/Client.cs b/RabbitMQ.Stream.Client/Client.cs
@@ -198,9 +198,8 @@ private async Task OnConnectionClosed(string reason)
         public static async Task<Client> Create(ClientParameters parameters, ILogger logger = null)
         {
             var client = new Client(parameters, logger);
-
             client.connection = await Connection
-                .Create(parameters.Endpoint, client.HandleIncoming, client.HandleClosed, parameters.Ssl)
+                .Create(parameters.Endpoint, client.HandleIncoming, client.HandleClosed, parameters.Ssl, logger)
                 .ConfigureAwait(false);
 
             // exchange properties
@@ -426,10 +425,11 @@ private async Task HandleIncoming(Memory<byte> frameMemory)
                         .ConfigureAwait(false);
                     if (off == null)
                     {
-                        _logger.LogWarning(
+                        _logger?.LogWarning(
                             "ConsumerUpdateHandler can't returned null, a default offsetType (OffsetTypeNext) will be used");
                         off = new OffsetTypeNext();
                     }
+
                     // event the consumer is not active, we need to send a ConsumerUpdateResponse
                     // by protocol definition. the offsetType can't be null so we use OffsetTypeNext as default
                     await ConsumerUpdateResponse(

diff --git a/RabbitMQ.Stream.Client/Connection.cs b/RabbitMQ.Stream.Client/Connection.cs
@@ -10,6 +10,7 @@
 using System.Net.Sockets;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.Extensions.Logging;
 
 namespace RabbitMQ.Stream.Client
 {
@@ -25,6 +26,7 @@ public class Connection : IDisposable
         private int numFrames;
         private bool isClosed = false;
         private bool _disposedValue;
+        private readonly ILogger _logger;
 
         internal int NumFrames => numFrames;
 
@@ -36,8 +38,9 @@ private static System.IO.Stream MaybeTcpUpgrade(NetworkStream networkStream, Ssl
         }
 
         private Connection(Socket socket, Func<Memory<byte>, Task> callback,
-            Func<string, Task> closedCallBack, SslOption sslOption)
+            Func<string, Task> closedCallBack, SslOption sslOption, ILogger logger)
         {
+            _logger = logger;
             this.socket = socket;
             commandCallback = callback;
             closedCallback = closedCallBack;
@@ -51,7 +54,7 @@ private Connection(Socket socket, Func<Memory<byte>, Task> callback,
         }
 
         public static async Task<Connection> Create(EndPoint endpoint, Func<Memory<byte>, Task> commandCallback,
-            Func<string, Task> closedCallBack, SslOption sslOption)
+            Func<string, Task> closedCallBack, SslOption sslOption, ILogger logger)
         {
             var socket = new Socket(SocketType.Stream, ProtocolType.Tcp);
             socket.NoDelay = true;
@@ -80,7 +83,7 @@ public static async Task<Connection> Create(EndPoint endpoint, Func<Memory<byte>
                 }
             }
 
-            return new Connection(socket, commandCallback, closedCallBack, sslOption);
+            return new Connection(socket, commandCallback, closedCallBack, sslOption, logger);
         }
 
         public async ValueTask<bool> Write<T>(T command) where T : struct, ICommand
@@ -95,6 +98,10 @@ public async ValueTask<bool> Write<T>(T command) where T : struct, ICommand
 
         private async Task WriteCommand<T>(T command) where T : struct, ICommand
         {
+            if (isClosed)
+            {
+                throw new InvalidOperationException("Connection is closed");
+            }
             // Only one thread should be able to write to the output pipeline at a time.
             await _writeLock.WaitAsync().ConfigureAwait(false);
             try
@@ -149,23 +156,36 @@ private async Task ProcessIncomingFrames()
                     reader.AdvanceTo(buffer.Start, buffer.End);
                 }
             }
+            catch (OperationCanceledException e)
+            {
+                caught = e;
+                // We don't need to log as error this exception
+                // It is raised when the socket is closed due of cancellation token 
+                // from the producer or consumer class
+                // we leave it as debug to avoid noise in the logs
+                _logger?.LogDebug("Operation Canceled Exception. TCP Connection Closed");
+            }
             catch (Exception e)
             {
                 caught = e;
                 // The exception is needed mostly to raise the 
                 // closedCallback event.
                 // It is useful to trace the error, but at this point
                 // the socket is closed maybe not in the correct way
-                Debug.WriteLine($"Error reading the socket, error: {e}");
+                if (!isClosed)
+                {
+                    _logger?.LogError(e, "Error reading the socket");
+                }
             }
             finally
             {
                 isClosed = true;
                 // Mark the PipeReader as complete
                 await reader.CompleteAsync(caught).ConfigureAwait(false);
                 var t = closedCallback?.Invoke("TCP Connection Closed")!;
-                await t.ConfigureAwait(false);
-                Debug.WriteLine("TCP Connection Closed");
+                if (t != null)
+                    await t.ConfigureAwait(false);
+                _logger?.LogDebug("TCP Connection Closed");
             }
         }
 
@@ -204,7 +224,7 @@ public void Dispose()
                     socket.Dispose();
                     if (!_incomingFramesTask.Wait(Consts.ShortWait))
                     {
-                        Debug.WriteLine($"frame reader task did not exit in {Consts.ShortWait}");
+                        _logger?.LogError("frame reader task did not exit in {ShortWait}", Consts.ShortWait);
                     }
                 }
                 finally

diff --git a/RabbitMQ.Stream.Client/Consts.cs b/RabbitMQ.Stream.Client/Consts.cs
@@ -14,5 +14,15 @@ internal static class Consts
         internal static readonly TimeSpan ShortWait = TimeSpan.FromSeconds(1);
         internal static readonly TimeSpan MidWait = TimeSpan.FromSeconds(3);
         internal static readonly TimeSpan LongWait = TimeSpan.FromSeconds(10);
+
+        internal static int RandomShort()
+        {
+            return Random.Shared.Next(500, 1500);
+        }
+
+        internal static int RandomMid()
+        {
+            return Random.Shared.Next(1000, 2500);
+        }
     }
 }
diff --git a/RabbitMQ.Stream.Client/Deliver.cs b/RabbitMQ.Stream.Client/Deliver.cs
@@ -51,7 +51,7 @@ internal readonly struct SubEntryChunk
         private SubEntryChunk(byte compress,
             ushort numRecordsInBatch,
             uint unCompressedDataSize, uint dataLen,
-            ReadOnlySequence<byte> data)
+            Memory<byte> data)
         {
             compressValue = compress;
             NumRecordsInBatch = numRecordsInBatch;
@@ -67,7 +67,18 @@ private SubEntryChunk(byte compress,
         public uint UnCompressedDataSize { get; }
 
         public uint DataLen { get; }
-        public ReadOnlySequence<byte> Data { get; }
+        public Memory<byte> Data { get; }
+
+        // This wrapper was added to be used in async methods
+        // where the SequenceReader is not available
+        // see RawConsumer:ParseChunk for more details
+        // at some point we could remove this wrapper
+        // and use system.io.pipeline instead of SequenceReader
+        internal static int Read(ref ReadOnlySequence<byte> seq, byte entryType, out SubEntryChunk subEntryChunk)
+        {
+            var reader = new SequenceReader<byte>(seq);
+            return Read(ref reader, entryType, out subEntryChunk);
+        }
 
         internal static int Read(ref SequenceReader<byte> reader, byte entryType, out SubEntryChunk subEntryChunk)
         {
@@ -77,14 +88,18 @@ internal static int Read(ref SequenceReader<byte> reader, byte entryType, out Su
             // Determinate what kind of the compression it is using
             // See Compress:CompressMode
             var compress = (byte)((byte)(entryType & 0x70) >> 4);
-            offset++;
 
             // Data contains the subEntryChunk information
             // We need to pass it to the subEntryChunk that will decode the information
+            var memory =
+                ArrayPool<byte>.Shared.Rent((int)dataLen).AsMemory(0, (int)dataLen);
             var data = reader.Sequence.Slice(reader.Consumed, dataLen);
+            data.CopyTo(memory.Span);
+
             subEntryChunk =
-                new SubEntryChunk(compress, numRecordsInBatch, unCompressedDataSize, dataLen, data);
-            offset += (int)dataLen;
+                new SubEntryChunk(compress, numRecordsInBatch, unCompressedDataSize, dataLen, memory);
+            offset += memory.Length;
+
             // Here we need to advance the reader to the datalen
             // Since Data is passed to the subEntryChunk.
             reader.Advance(dataLen);
@@ -101,7 +116,7 @@ private Chunk(byte magicVersion,
             ulong epoch,
             ulong chunkId,
             int crc,
-            ReadOnlySequence<byte> data)
+            Memory<byte> data)
         {
             MagicVersion = magicVersion;
             NumEntries = numEntries;
@@ -121,7 +136,7 @@ private Chunk(byte magicVersion,
         public ulong Epoch { get; }
         public ulong ChunkId { get; }
         public int Crc { get; }
-        public ReadOnlySequence<byte> Data { get; }
+        public Memory<byte> Data { get; }
 
         internal static int Read(ReadOnlySequence<byte> frame, out Chunk chunk)
         {
@@ -138,9 +153,11 @@ internal static int Read(ReadOnlySequence<byte> frame, out Chunk chunk)
             offset += WireFormatting.ReadUInt32(ref reader, out _);
             // offset += 4; // reserved
             offset += WireFormatting.ReadUInt32(ref reader, out _); // reserved
+            var memory =
+                ArrayPool<byte>.Shared.Rent((int)dataLen).AsMemory(0, (int)dataLen);
             var data = reader.Sequence.Slice(reader.Consumed, dataLen);
-            offset += (int)dataLen;
-            chunk = new Chunk(magicVersion, numEntries, numRecords, timestamp, epoch, chunkId, crc, data);
+            data.CopyTo(memory.Span);
+            chunk = new Chunk(magicVersion, numEntries, numRecords, timestamp, epoch, chunkId, crc, memory);
             return offset;
         }
     }

diff --git a/RabbitMQ.Stream.Client/Message.cs b/RabbitMQ.Stream.Client/Message.cs
@@ -71,6 +71,18 @@ public ReadOnlySequence<byte> Serialize()
             return new ReadOnlySequence<byte>(data);
         }
 
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        // This wrapper was added to be used in async methods
+        // where the SequenceReader is not available
+        // see RawConsumer:ParseChunk for more details
+        // at some point we could remove this wrapper
+        // and use system.io.pipeline instead of SequenceReader
+        public static Message From(ref ReadOnlySequence<byte> seq, uint len)
+        {
+            var reader = new SequenceReader<byte>(seq);
+            return From(ref reader, len);
+        }
+
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public static Message From(ref SequenceReader<byte> reader, uint len)
         {

diff --git a/RabbitMQ.Stream.Client/PublicAPI.Shipped.txt b/RabbitMQ.Stream.Client/PublicAPI.Shipped.txt
@@ -159,7 +159,6 @@ RabbitMQ.Stream.Client.Chunk
 RabbitMQ.Stream.Client.Chunk.Chunk() -> void
 RabbitMQ.Stream.Client.Chunk.ChunkId.get -> ulong
 RabbitMQ.Stream.Client.Chunk.Crc.get -> int
-RabbitMQ.Stream.Client.Chunk.Data.get -> System.Buffers.ReadOnlySequence<byte>
 RabbitMQ.Stream.Client.Chunk.Epoch.get -> ulong
 RabbitMQ.Stream.Client.Chunk.NumEntries.get -> ushort
 RabbitMQ.Stream.Client.Chunk.NumRecords.get -> uint
@@ -803,7 +802,6 @@ static RabbitMQ.Stream.Client.AMQP.DescribedFormatCode.Write(System.Span<byte> s
 static RabbitMQ.Stream.Client.Client.Create(RabbitMQ.Stream.Client.ClientParameters parameters, Microsoft.Extensions.Logging.ILogger logger = null) -> System.Threading.Tasks.Task<RabbitMQ.Stream.Client.Client>
 static RabbitMQ.Stream.Client.CompressionHelper.Compress(System.Collections.Generic.List<RabbitMQ.Stream.Client.Message> messages, RabbitMQ.Stream.Client.CompressionType compressionType) -> RabbitMQ.Stream.Client.ICompressionCodec
 static RabbitMQ.Stream.Client.CompressionHelper.UnCompress(RabbitMQ.Stream.Client.CompressionType compressionType, System.Buffers.ReadOnlySequence<byte> source, uint dataLen, uint unCompressedDataSize) -> System.Buffers.ReadOnlySequence<byte>
-static RabbitMQ.Stream.Client.Connection.Create(System.Net.EndPoint endpoint, System.Func<System.Memory<byte>, System.Threading.Tasks.Task> commandCallback, System.Func<string, System.Threading.Tasks.Task> closedCallBack, RabbitMQ.Stream.Client.SslOption sslOption) -> System.Threading.Tasks.Task<RabbitMQ.Stream.Client.Connection>
 static RabbitMQ.Stream.Client.LeaderLocator.ClientLocal.get -> RabbitMQ.Stream.Client.LeaderLocator
 static RabbitMQ.Stream.Client.LeaderLocator.LeastLeaders.get -> RabbitMQ.Stream.Client.LeaderLocator
 static RabbitMQ.Stream.Client.LeaderLocator.Random.get -> RabbitMQ.Stream.Client.LeaderLocator

diff --git a/RabbitMQ.Stream.Client/PublicAPI.Unshipped.txt b/RabbitMQ.Stream.Client/PublicAPI.Unshipped.txt
@@ -1,6 +1,7 @@
 const RabbitMQ.Stream.Client.StreamStatsResponse.Key = 28 -> ushort
 RabbitMQ.Stream.Client.AbstractEntity.MaybeCancelToken() -> void
 RabbitMQ.Stream.Client.AbstractEntity.Token.get -> System.Threading.CancellationToken
+RabbitMQ.Stream.Client.Chunk.Data.get -> System.Memory<byte>
 RabbitMQ.Stream.Client.Chunk.MagicVersion.get -> byte
 RabbitMQ.Stream.Client.Client.StreamStats(string stream) -> System.Threading.Tasks.ValueTask<RabbitMQ.Stream.Client.StreamStatsResponse>
 RabbitMQ.Stream.Client.Reliable.DeduplicatingProducer
@@ -24,4 +25,6 @@ RabbitMQ.Stream.Client.StreamStatsResponse.StreamStatsResponse() -> void
 RabbitMQ.Stream.Client.StreamStatsResponse.StreamStatsResponse(uint correlationId, RabbitMQ.Stream.Client.ResponseCode responseCode, System.Collections.Generic.IDictionary<string, long> statistic) -> void
 RabbitMQ.Stream.Client.StreamStatsResponse.Write(System.Span<byte> span) -> int
 RabbitMQ.Stream.Client.StreamSystem.StreamStats(string stream) -> System.Threading.Tasks.Task<RabbitMQ.Stream.Client.StreamStats>
-static RabbitMQ.Stream.Client.Reliable.DeduplicatingProducer.Create(RabbitMQ.Stream.Client.Reliable.DeduplicatingProducerConfig producerConfig, Microsoft.Extensions.Logging.ILogger<RabbitMQ.Stream.Client.Reliable.Producer> logger = null) -> System.Threading.Tasks.Task<RabbitMQ.Stream.Client.Reliable.DeduplicatingProducer>
+static RabbitMQ.Stream.Client.Connection.Create(System.Net.EndPoint endpoint, System.Func<System.Memory<byte>, System.Threading.Tasks.Task> commandCallback, System.Func<string, System.Threading.Tasks.Task> closedCallBack, RabbitMQ.Stream.Client.SslOption sslOption, Microsoft.Extensions.Logging.ILogger logger) -> System.Threading.Tasks.Task<RabbitMQ.Stream.Client.Connection>
+static RabbitMQ.Stream.Client.Message.From(ref System.Buffers.ReadOnlySequence<byte> seq, uint len) -> RabbitMQ.Stream.Client.Message
+static RabbitMQ.Stream.Client.Reliable.DeduplicatingProducer.Create(RabbitMQ.Stream.Client.Reliable.DeduplicatingProducerConfig producerConfig, Microsoft.Extensions.Logging.ILogger<RabbitMQ.Stream.Client.Reliable.Producer> logger = null) -> System.Threading.Tasks.Task<RabbitMQ.Stream.Client.Reliable.DeduplicatingProducer>