Fix race condition during rawconsumer creation

The processCheck now waits the subscription is completed before send the credits

Signed-off-by: Gabriele Santomaggio <[email protected]>

Author: Gsantomaggio

RabbitMQ.Stream.Client/Client.cs

@@ -713,23 +713,27 @@ internal async Task<CloseResponse> MaybeClose(string reason, string stream, Conn
             await _poolSemaphore.WaitAsync().ConfigureAwait(false);
             try
             {
-                if (!string.IsNullOrEmpty(ClientId))
-                {
-                    _logger.LogInformation("Releasing ids for the client id {ClientId}", ClientId);
-                    pool.Release(ClientId, stream);
-                }
-
                 if (!HasEntities())
                 {
                     if (!string.IsNullOrEmpty(ClientId))
                     {
                         _logger.LogInformation("Close connection for the {ClientId}", ClientId);
-                        // pool.remove(ClientId)  is a duplicate call here but it is ok
                         // the client can be closed in an unexpected way so we need to remove it from the pool
                         // so you will find pool.remove(ClientId) also to the disconnect event
+                        // pool.remove(ClientId) is a duplicate call here but it is ok. The pool is idempotent 
+                        pool.Remove(ClientId);
                         await Close(reason).ConfigureAwait(false);
                     }
                 }
+                else
+                {
+                    // we remove an id reference from the client
+                    // in case there are still active ids from the client and the stream 
+                    if (!string.IsNullOrEmpty(ClientId))
+                    {
+                        pool.Release(ClientId, stream);
+                    }
+                }
 
                 var result = new CloseResponse(0, ResponseCode.Ok);
                 return result;

RabbitMQ.Stream.Client/Connection.cs

@@ -145,14 +145,14 @@ private async Task ProcessIncomingFrames()
                     var buffer = result.Buffer;
                     if (buffer.Length == 0)
                     {
-                        Debug.WriteLine("TCP Connection Closed");
+                        Debug.WriteLine("TCP Connection Closed!");
                         // We're not going to receive any more bytes from the connection.
                         break;
                     }
 
                     // Let's try to read some frames!
 
-                    while (TryReadFrame(ref buffer, out var frame))
+                    while (TryReadFrame(ref buffer, out var frame) && !isClosed)
                     {
                         // Let's rent some memory to copy the frame from the network stream. This memory will be reclaimed once the frame has been handled.
 

RabbitMQ.Stream.Client/ConnectionsPool.cs

@@ -96,17 +96,25 @@ internal static byte FindNextValidId(List<byte> ids)
                 return 0;
             }
 
-            // we start with the recycle when we reach the max value
-            // in this way we can avoid to recycle the same ids in a short time
+            // // we start with the recycle when we reach the max value
+            // // in this way we can avoid to recycle the same ids in a short time
             ids.Sort();
-            if (ids[^1] == byte.MaxValue)
+
+/* Unmerged change from project 'RabbitMQ.Stream.Client(net7.0)'
+Before:
+            if (ids[^1] != byte.MaxValue) return (byte)(ids[^1] + 1);
+After:
+            if (ids[^1] != byte.MaxValue)
+                return (byte)(ids[^1] + 1);
+*/
+            if (ids[^1] != byte.MaxValue)
+                return (byte)(ids[^1] + 1);
+
+            for (var i = 0; i < ids.Count - 1; i++)
             {
-                for (var i = 0; i < ids.Count - 1; i++)
+                if (ids[i + 1] - ids[i] > 1)
                 {
-                    if (ids[i + 1] - ids[i] > 1)
-                    {
-                        return (byte)(ids[i] + 1);
-                    }
+                    return (byte)(ids[i] + 1);
                 }
             }
 

RabbitMQ.Stream.Client/RawConsumer.cs

@@ -120,6 +120,12 @@ public class RawConsumer : AbstractEntity, IConsumer, IDisposable
         private readonly ILogger _logger;
         private readonly Channel<Chunk> _chunksBuffer;
         private readonly ushort _initialCredits;
+        
+        // _completeSubscription is used to notify the ProcessChunks task
+        // that the subscription is completed and so it can start to process the chunks
+        // this is needed because the socket starts to receive the chunks before the subscription_id is 
+        // assigned. 
+        private readonly TaskCompletionSource _completeSubscription = new ();
 
         private string ConsumerInfo()
         {
@@ -192,6 +198,7 @@ public static async Task<IConsumer> Create(
             var client = await RoutingHelper<Routing>
                 .LookupRandomConnection(clientParameters, metaStreamInfo, config.Pool, logger)
                 .ConfigureAwait(false);
+            logger?.LogInformation("Creating consumer {ConsumerInfo}", client.ClientId);
             var consumer = new RawConsumer((Client)client, config, logger);
             await consumer.Init().ConfigureAwait(false);
             return consumer;
@@ -304,6 +311,14 @@ await _config.MessageHandler(this,
                     }
                     catch (Exception e)
                     {
+                        if (Token.IsCancellationRequested)
+                        {
+                            _logger?.LogDebug(
+                                "Dispatching  {ConsumerInfo}, Cancellation Requested, the consumer is closing. ",
+                                ConsumerInfo());
+                            return;
+                        }
+
                         _logger?.LogError(e,
                             "Error while Dispatching message, ChunkId : {ChunkId} {ConsumerInfo}",
                             chunk.ChunkId, ConsumerInfo());
@@ -373,6 +388,7 @@ private void ProcessChunks()
         {
             Task.Run(async () =>
             {
+                await _completeSubscription.Task.ConfigureAwait(false);
                 try
                 {
                     while (!Token.IsCancellationRequested &&
@@ -384,6 +400,7 @@ await _chunksBuffer.Reader.WaitToReadAsync(Token).ConfigureAwait(false)) //
                             // we request the credit before process the check to keep the network busy
                             try
                             {
+
                                 await _client.Credit(_subscriberId, 1).ConfigureAwait(false);
                             }
                             catch (InvalidOperationException)
@@ -488,7 +505,7 @@ private async Task Init()
             var chunkConsumed = 0;
             // this the default value for the consumer.
             _config.StoredOffsetSpec = _config.OffsetSpec;
-            var (consumerId, response) = await _client.Subscribe(
+            (_subscriberId, var response) = await _client.Subscribe(
                 _config,
                 _initialCredits,
                 consumerProperties,
@@ -553,7 +570,7 @@ private async Task Init()
             ).ConfigureAwait(false);
             if (response.ResponseCode == ResponseCode.Ok)
             {
-                _subscriberId = consumerId;
+                _completeSubscription.SetResult();
                 _status = EntityStatus.Open;
                 return;
             }
@@ -567,12 +584,12 @@ public override async Task<ResponseCode> Close()
             // this unlock the consumer if it is waiting for a message
             // see DispatchMessage method where the token is used
             MaybeCancelToken();
-
             if (!IsOpen())
             {
                 return ResponseCode.Ok;
             }
 
+            _status = EntityStatus.Closed;
             var result = ResponseCode.Ok;
             try
             {