Add the semaphore to sync the

producers and consumers requests.
Remove the way to decide the nextid for producers and consumers now it uses the client lists
Make the consumer and producer close idempotent.
Add EntityStatus to handle the connection status for producers and consumers

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

Author: Gsantomaggio

RabbitMQ.Stream.Client/AbstractEntity.cs

@@ -6,20 +6,34 @@
 
 namespace RabbitMQ.Stream.Client
 {
+
+    internal enum EntityStatus
+    {
+        Open,
+        Closed,
+        Disposed
+    }
     public abstract class AbstractEntity
     {
         private readonly CancellationTokenSource _cancelTokenSource = new();
         protected CancellationToken Token => _cancelTokenSource.Token;
 
+        internal EntityStatus _status = EntityStatus.Closed;
         // here the _cancelTokenSource is disposed and the token is cancelled
         // in producer is used to cancel the send task
         // in consumer is used to cancel the receive task
         protected void MaybeCancelToken()
         {
+
             if (!_cancelTokenSource.IsCancellationRequested)
                 _cancelTokenSource.Cancel();
         }
 
+        public bool IsOpen()
+        {
+            return _status == EntityStatus.Open;
+        }
+
         protected Client _client;
 
     }

RabbitMQ.Stream.Client/Client.cs

@@ -109,17 +109,14 @@ public class Client : IClient
 
         private Connection connection;
 
-        private readonly IDictionary<byte, (Action<ReadOnlyMemory<ulong>>, Action<(ulong, ResponseCode)[]>)>
-            publishers =
-                new ConcurrentDictionary<byte, (Action<ReadOnlyMemory<ulong>>, Action<(ulong, ResponseCode)[]>)>();
-
         private readonly ConcurrentDictionary<uint, IValueTaskSource> requests = new();
 
         private readonly TaskCompletionSource<TuneResponse> tuneReceived =
             new TaskCompletionSource<TuneResponse>(TaskCreationOptions.RunContinuationsAsynchronously);
 
-        private readonly List<byte> subscriptionIds = new();
-        private readonly List<byte> publisherIds = new();
+        internal readonly IDictionary<byte, (Action<ReadOnlyMemory<ulong>>, Action<(ulong, ResponseCode)[]>)>
+            publishers =
+                new ConcurrentDictionary<byte, (Action<ReadOnlyMemory<ulong>>, Action<(ulong, ResponseCode)[]>)>();
 
         internal readonly IDictionary<byte, ConsumerEvents> consumers =
             new ConcurrentDictionary<byte, ConsumerEvents>();
@@ -157,38 +154,11 @@ private Client(ClientParameters parameters, ILogger logger = null)
             IsClosed = false;
             _logger = logger ?? NullLogger.Instance;
             ClientId = Guid.NewGuid().ToString();
-        }
-
-        private byte GetNextSubscriptionId()
-        {
-            byte result;
-            lock (Obj)
-            {
-                result = ConnectionsPool.FindMissingConsecutive(subscriptionIds);
-                subscriptionIds.Add(result);
-            }
-
-            return result;
-        }
-
-        private byte GetNextPublisherId()
-        {
-            byte result;
-            lock (Obj)
-            {
-                result = ConnectionsPool.FindMissingConsecutive(publisherIds);
-                publisherIds.Add(result);
-            }
-
-            return result;
-        }
-
-        internal void RemoveSubscriptionId(byte id)
-        {
-            lock (Obj)
+            AppDomain.CurrentDomain.UnhandledException += (sender, args) =>
             {
-                subscriptionIds.Remove(id);
-            }
+                _logger.LogError(args.ExceptionObject as Exception, "Unhandled exception");
+                Parameters.UnhandledExceptionHandler(args.ExceptionObject as Exception);
+            };
         }
 
         public bool IsClosed
@@ -327,14 +297,23 @@ public ValueTask<bool> Publish<T>(T msg) where T : struct, ICommand
             Action<ReadOnlyMemory<ulong>> confirmCallback,
             Action<(ulong, ResponseCode)[]> errorCallback)
         {
-            var publisherId = GetNextPublisherId();
-            publishers.Add(publisherId, (confirmCallback, errorCallback));
-            return (publisherId, await Request<DeclarePublisherRequest, DeclarePublisherResponse>(corr =>
-                new DeclarePublisherRequest(corr, publisherId, publisherRef, stream)).ConfigureAwait(false));
+            await _poolSemaphore.WaitAsync().ConfigureAwait(false);
+            try
+            {
+                var publisherId = ConnectionsPool.FindMissingConsecutive(publishers.Keys.ToList());
+                publishers.Add(publisherId, (confirmCallback, errorCallback));
+                return (publisherId, await Request<DeclarePublisherRequest, DeclarePublisherResponse>(corr =>
+                    new DeclarePublisherRequest(corr, publisherId, publisherRef, stream)).ConfigureAwait(false));
+            }
+            finally
+            {
+                _poolSemaphore.Release();
+            }
         }
 
         public async Task<DeletePublisherResponse> DeletePublisher(byte publisherId)
         {
+            await _poolSemaphore.WaitAsync().ConfigureAwait(false);
             try
             {
                 var result =
@@ -346,6 +325,7 @@ await Request<DeletePublisherRequest, DeletePublisherResponse>(corr =>
             finally
             {
                 publishers.Remove(publisherId);
+                _poolSemaphore.Release();
             }
         }
 
@@ -366,21 +346,29 @@ await Request<DeletePublisherRequest, DeletePublisherResponse>(corr =>
             Dictionary<string, string> properties, Func<Deliver, Task> deliverHandler,
             Func<bool, Task<IOffsetType>> consumerUpdateHandler)
         {
-            var subscriptionId = GetNextSubscriptionId();
-
-            consumers.Add(subscriptionId,
-                new ConsumerEvents(
-                    deliverHandler,
-                    consumerUpdateHandler));
+            await _poolSemaphore.WaitAsync().ConfigureAwait(false);
+            try
+            {
+                var subscriptionId = ConnectionsPool.FindMissingConsecutive(consumers.Keys.ToList());
+                consumers.Add(subscriptionId,
+                    new ConsumerEvents(
+                        deliverHandler,
+                        consumerUpdateHandler));
 
-            return (subscriptionId,
-                await Request<SubscribeRequest, SubscribeResponse>(corr =>
-                    new SubscribeRequest(corr, subscriptionId, config.Stream, config.OffsetSpec, initialCredit,
-                        properties)).ConfigureAwait(false));
+                return (subscriptionId,
+                    await Request<SubscribeRequest, SubscribeResponse>(corr =>
+                        new SubscribeRequest(corr, subscriptionId, config.Stream, config.OffsetSpec, initialCredit,
+                            properties)).ConfigureAwait(false));
+            }
+            finally
+            {
+                _poolSemaphore.Release();
+            }
         }
 
         public async Task<UnsubscribeResponse> Unsubscribe(byte subscriptionId)
         {
+            await _poolSemaphore.WaitAsync().ConfigureAwait(false);
             try
             {
                 // here we reduce a bit the timeout to avoid waiting too much
@@ -398,6 +386,7 @@ await Request<UnsubscribeRequest, UnsubscribeResponse>(corr =>
                 _logger.LogDebug("Unsubscribe: {SubscriptionId}", subscriptionId);
                 // remove consumer after RPC returns, this should avoid uncorrelated data being sent
                 consumers.Remove(subscriptionId);
+                _poolSemaphore.Release();
             }
         }
 
@@ -681,6 +670,7 @@ public async Task<CloseResponse> Close(string reason)
                 return new CloseResponse(0, ResponseCode.Ok);
             }
 
+            InternalClose();
             try
             {
                 var result =
@@ -700,41 +690,54 @@ public async Task<CloseResponse> Close(string reason)
             }
             finally
             {
-                // even if the close fails we need to close the connection
-                InternalClose();
                 connection.Dispose();
             }
 
             return new CloseResponse(0, ResponseCode.Ok);
         }
 
+        // _poolSemaphore is introduced here: https://github.com/rabbitmq/rabbitmq-stream-dotnet-client/pull/328 
+        // the MaybeClose can be called in different threads so we need to protect the pool
+        // the pool itself is thread safe but we need to protect the flow to be sure that the 
+        // connection is released only once
+        private readonly SemaphoreSlim _poolSemaphore = new(1, 1);
+
         // Safe close 
         // the client can be closed only if HasEntities is false
         // if the client has entities (publishers or consumers) it will be released from the pool
         // Release will decrement the active ids for the connection
         // if the active ids are 0 the connection will be closed
+
         internal async Task<CloseResponse> MaybeClose(string reason, string stream, ConnectionsPool pool)
         {
-            if (!string.IsNullOrEmpty(ClientId))
-            {
-                _logger.LogInformation("Releasing connection {Connection}", ClientId);
-                pool.Release(ClientId, stream);
-            }
-
-            if (!HasEntities())
+            await _poolSemaphore.WaitAsync().ConfigureAwait(false);
+            try
             {
                 if (!string.IsNullOrEmpty(ClientId))
                 {
-                    _logger.LogInformation("Close connection {Connection}", 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
-                    await Close(reason).ConfigureAwait(false);
+                    _logger.LogInformation("Releasing connection {Connection}", ClientId);
+                    pool.Release(ClientId, stream);
+                }
+
+                if (!HasEntities())
+                {
+                    if (!string.IsNullOrEmpty(ClientId))
+                    {
+                        _logger.LogInformation("Close connection {Connection}", 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
+                        await Close(reason).ConfigureAwait(false);
+                    }
                 }
-            }
 
-            var result = new CloseResponse(0, ResponseCode.Ok);
-            return result;
+                var result = new CloseResponse(0, ResponseCode.Ok);
+                return result;
+            }
+            finally
+            {
+                _poolSemaphore.Release();
+            }
         }
 
         public string ClientId { get; init; }

RabbitMQ.Stream.Client/ConnectionsPool.cs

@@ -84,23 +84,28 @@ public bool Available
 /// </summary>
 public class ConnectionsPool
 {
+    private static readonly object s_lock = new();
+
     internal static byte FindMissingConsecutive(List<byte> ids)
     {
-        if (ids.Count == 0)
+        lock (s_lock)
         {
-            return 0;
-        }
+            if (ids.Count == 0)
+            {
+                return 0;
+            }
 
-        ids.Sort();
-        for (var i = 0; i < ids.Count - 1; i++)
-        {
-            if (ids[i + 1] - ids[i] > 1)
+            ids.Sort();
+            for (var i = 0; i < ids.Count - 1; i++)
             {
-                return (byte)(ids[i] + 1);
+                if (ids[i + 1] - ids[i] > 1)
+                {
+                    return (byte)(ids[i] + 1);
+                }
             }
-        }
 
-        return (byte)(ids[^1] + 1);
+            return (byte)(ids[^1] + 1);
+        }
     }
 
     private readonly int _maxConnections;
@@ -233,9 +238,16 @@ public void Remove(string clientId)
 
     public int ActiveIdsCount => Connections.Values.Sum(x => x.StreamIds.Values.Sum(y => y.Count));
 
-    public int ActiveIdsCountForStream(string stream) => Connections.Values.Sum(x => x.StreamIds.TryGetValue(stream, out var streamIds) ? streamIds.Count : 0);
+    public int ActiveIdsCountForStream(string stream) => Connections.Values.Sum(x =>
+        x.StreamIds.TryGetValue(stream, out var streamIds) ? streamIds.Count : 0);
 
-    public int ActiveIdsCountForClient(string clientId) => Connections.TryGetValue(clientId, out var connectionItem) ? connectionItem.StreamIds.Values.Sum(y => y.Count) : 0;
+    public int ActiveIdsCountForClient(string clientId) => Connections.TryGetValue(clientId, out var connectionItem)
+        ? connectionItem.StreamIds.Values.Sum(y => y.Count)
+        : 0;
 
-    public int ActiveIdsCountForClientAndStream(string clientId, string stream) => Connections.TryGetValue(clientId, out var connectionItem) && connectionItem.StreamIds.TryGetValue(stream, out var streamIds) ? streamIds.Count : 0;
+    public int ActiveIdsCountForClientAndStream(string clientId, string stream) =>
+        Connections.TryGetValue(clientId, out var connectionItem) &&
+        connectionItem.StreamIds.TryGetValue(stream, out var streamIds)
+            ? streamIds.Count
+            : 0;
 }

RabbitMQ.Stream.Client/PublicAPI.Shipped.txt

@@ -547,7 +547,6 @@ RabbitMQ.Stream.Client.RawProducer.ConfirmFrames.get -> int
 RabbitMQ.Stream.Client.RawProducer.Dispose() -> void
 RabbitMQ.Stream.Client.RawProducer.GetLastPublishingId() -> System.Threading.Tasks.Task<ulong>
 RabbitMQ.Stream.Client.RawProducer.IncomingFrames.get -> int
-RabbitMQ.Stream.Client.RawProducer.IsOpen() -> bool
 RabbitMQ.Stream.Client.RawProducer.MessagesSent.get -> int
 RabbitMQ.Stream.Client.RawProducer.PendingCount.get -> int
 RabbitMQ.Stream.Client.RawProducer.PublishCommandsSent.get -> int

RabbitMQ.Stream.Client/PublicAPI.Unshipped.txt

@@ -1,6 +1,7 @@
 const RabbitMQ.Stream.Client.RouteQueryResponse.Key = 24 -> ushort
 const RabbitMQ.Stream.Client.StreamStatsResponse.Key = 28 -> ushort
 override RabbitMQ.Stream.Client.Broker.ToString() -> string
+RabbitMQ.Stream.Client.AbstractEntity.IsOpen() -> bool
 RabbitMQ.Stream.Client.AbstractEntity.MaybeCancelToken() -> void
 RabbitMQ.Stream.Client.AbstractEntity.Token.get -> System.Threading.CancellationToken
 RabbitMQ.Stream.Client.AuthMechanism

RabbitMQ.Stream.Client/RawConsumer.cs

@@ -115,7 +115,6 @@ internal void Validate()
 
     public class RawConsumer : AbstractEntity, IConsumer, IDisposable
     {
-        private bool _disposed;
         private readonly RawConsumerConfig _config;
         private byte _subscriberId;
         private readonly ILogger _logger;
@@ -555,6 +554,7 @@ private async Task Init()
             if (response.ResponseCode == ResponseCode.Ok)
             {
                 _subscriberId = consumerId;
+                _status = EntityStatus.Open;
                 return;
             }
 
@@ -568,7 +568,7 @@ public async Task<ResponseCode> Close()
             // see DispatchMessage method where the token is used
             MaybeCancelToken();
 
-            if (_client.IsClosed)
+            if (!IsOpen())
             {
                 return ResponseCode.Ok;
             }
@@ -596,8 +596,6 @@ public async Task<ResponseCode> Close()
                     _subscriberId, ConsumerInfo());
             }
 
-            _client.RemoveSubscriptionId(_subscriberId);
-
             var closed = await _client.MaybeClose($"_client-close-subscriber: {_subscriberId}",
                     _config.Stream, _config.Pool)
                 .ConfigureAwait(false);
@@ -614,7 +612,7 @@ private void Dispose(bool disposing)
                 return;
             }
 
-            if (_disposed)
+            if (_status == EntityStatus.Disposed)
             {
                 return;
             }
@@ -632,7 +630,7 @@ private void Dispose(bool disposing)
             }
             finally
             {
-                _disposed = true;
+                _status = EntityStatus.Disposed;
             }
         }