Multiple Producers and Consumers per connection (#328)

* Implement multiple producers and consumers per connection.
Introduce a new set of classes to handle the connection pool.
The pool is a dictionary with client_id and the connectioninfo.

* The client_id is a field to identify the connection uniquely.

* the pool increases the reference count for each request, and for each Close(), it decreases the reference count.

* client_poolSemaphore is needed to coordinate the Producer and Consumer creation and destruction since the pool is thread-safe.

* To configure the pool, there is ConnectionPoolConfig, where the user can define  Producers and/or Consumers per connection. 

* The client still contains consumers' and publishers' lists to avoid changing the code too much.
These two lists are the source of the truth.

* This PR intercepts the entities' Creation and Destruction to make the change as safe as possible 

---------

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

Author: Gsantomaggio

.ci/versions.json

@@ -1,4 +1,4 @@
 {
   "erlang": "26.1.2",
-  "rabbitmq": "3.13.0-beta.6"
+  "rabbitmq": "3.13.0-rc.2"
 }

.github/workflows/build-test.yaml

@@ -43,7 +43,7 @@ jobs:
     runs-on: ubuntu-latest
     services:
       rabbitmq:
-        image: rabbitmq:3.13.0-beta.6-management
+        image: rabbitmq:3.13-rc-management
         env:
           RABBITMQ_SERVER_ADDITIONAL_ERL_ARGS: -rabbitmq_stream advertised_host localhost -rabbit collect_statistics_interval 4
         ports:

RabbitMQ.Stream.Client/AbstractEntity.cs

@@ -2,25 +2,74 @@
 // 2.0, and the Mozilla Public License, version 2.0.
 // Copyright (c) 2007-2023 VMware, Inc.
 
+using System;
 using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Extensions.Logging;
 
 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();
         }
 
-        protected Client _client;
+        public abstract Task<ResponseCode> Close();
+
+        protected void Dispose(bool disposing, string entityInfo, ILogger logger)
+        {
+            if (!disposing)
+            {
+                return;
+            }
+
+            if (_status == EntityStatus.Disposed)
+            {
+                return;
+            }
+
+            try
+            {
+                var closeTask = Close();
+                if (!closeTask.Wait(Consts.MidWait))
+                {
+                    logger.LogWarning("Failed to close {EntityInfo} in time", entityInfo);
+                }
+            }
+            catch (Exception e)
+            {
+                logger?.LogWarning("Failed to close {EntityInfo}, error {Error} ", entityInfo, e.Message);
+            }
+            finally
+            {
+                _status = EntityStatus.Disposed;
+            }
+        }
+
+        public bool IsOpen()
+        {
+            return _status == EntityStatus.Open;
+        }
+
+        internal Client _client;
 
     }
 }

RabbitMQ.Stream.Client/Client.cs

@@ -107,20 +107,16 @@ public class Client : IClient
 
         private uint correlationId = 0; // allow for some pre-amble
 
-        private byte nextPublisherId = 0;
-
         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 byte nextSubscriptionId;
+        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,17 +153,12 @@ private Client(ClientParameters parameters, ILogger logger = null)
                 (int)parameters.Heartbeat.TotalSeconds);
             IsClosed = false;
             _logger = logger ?? NullLogger.Instance;
-        }
-
-        private byte GetNextSubscriptionId()
-        {
-            byte result;
-            lock (Obj)
+            ClientId = Guid.NewGuid().ToString();
+            AppDomain.CurrentDomain.UnhandledException += (sender, args) =>
             {
-                result = nextSubscriptionId++;
-            }
-
-            return result;
+                _logger.LogError(args.ExceptionObject as Exception, "Unhandled exception");
+                Parameters.UnhandledExceptionHandler(args.ExceptionObject as Exception);
+            };
         }
 
         public bool IsClosed
@@ -304,16 +295,36 @@ public ValueTask<bool> Publish<T>(T msg) where T : struct, ICommand
         public async Task<(byte, DeclarePublisherResponse)> DeclarePublisher(string publisherRef,
             string stream,
             Action<ReadOnlyMemory<ulong>> confirmCallback,
-            Action<(ulong, ResponseCode)[]> errorCallback)
+            Action<(ulong, ResponseCode)[]> errorCallback, ConnectionsPool pool = null)
         {
-            var publisherId = nextPublisherId++;
-            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);
+            var publisherId = ConnectionsPool.FindNextValidId(publishers.Keys.ToList());
+            DeclarePublisherResponse response;
+
+            try
+            {
+                publishers.Add(publisherId, (confirmCallback, errorCallback));
+                response = await Request<DeclarePublisherRequest, DeclarePublisherResponse>(corr =>
+                    new DeclarePublisherRequest(corr, publisherId, publisherRef, stream)).ConfigureAwait(false);
+            }
+            finally
+            {
+                _poolSemaphore.Release();
+            }
+
+            if (response.ResponseCode == ResponseCode.Ok || pool == null)
+                return (publisherId, response);
+
+            // if the response code is not ok we need to remove the subscription
+            // and close the connection if necessary. 
+            publishers.Remove(publisherId);
+            await MaybeClose("Create Publisher Exception", stream, pool).ConfigureAwait(false);
+            return (publisherId, response);
         }
 
         public async Task<DeletePublisherResponse> DeletePublisher(byte publisherId)
         {
+            await _poolSemaphore.WaitAsync().ConfigureAwait(false);
             try
             {
                 var result =
@@ -325,6 +336,7 @@ await Request<DeletePublisherRequest, DeletePublisherResponse>(corr =>
             finally
             {
                 publishers.Remove(publisherId);
+                _poolSemaphore.Release();
             }
         }
 
@@ -345,21 +357,38 @@ 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);
+            var subscriptionId = ConnectionsPool.FindNextValidId(consumers.Keys.ToList());
+            SubscribeResponse response;
+            try
+            {
+                consumers.Add(subscriptionId,
+                    new ConsumerEvents(
+                        deliverHandler,
+                        consumerUpdateHandler));
 
-            return (subscriptionId,
-                await Request<SubscribeRequest, SubscribeResponse>(corr =>
+                response = await Request<SubscribeRequest, SubscribeResponse>(corr =>
                     new SubscribeRequest(corr, subscriptionId, config.Stream, config.OffsetSpec, initialCredit,
-                        properties)).ConfigureAwait(false));
+                        properties)).ConfigureAwait(false);
+            }
+            finally
+            {
+                _poolSemaphore.Release();
+            }
+
+            if (response.ResponseCode == ResponseCode.Ok)
+                return (subscriptionId, response);
+
+            // if the response code is not ok we need to remove the subscription
+            // and close the connection if necessary. 
+            consumers.Remove(subscriptionId);
+            await MaybeClose("Create Consumer Exception", config.Stream, config.Pool).ConfigureAwait(false);
+            return (subscriptionId, response);
         }
 
         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
@@ -377,6 +406,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();
             }
         }
 
@@ -640,6 +670,14 @@ private void InternalClose()
             IsClosed = true;
         }
 
+        private bool HasEntities()
+        {
+            lock (Obj)
+            {
+                return publishers.Count > 0 || consumers.Count > 0;
+            }
+        }
+
         private async ValueTask<bool> ConsumerUpdateResponse(uint rCorrelationId, IOffsetType offsetSpecification)
         {
             return await Publish(new ConsumerUpdateRequest(rCorrelationId, offsetSpecification)).ConfigureAwait(false);
@@ -652,6 +690,7 @@ public async Task<CloseResponse> Close(string reason)
                 return new CloseResponse(0, ResponseCode.Ok);
             }
 
+            InternalClose();
             try
             {
                 var result =
@@ -671,28 +710,62 @@ 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 the publishers are == 0
-        // not a public method used internally by producers and consumers
-        internal async Task<CloseResponse> MaybeClose(string reason)
+        // 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 (publishers.Count == 0 && consumers.Count == 0)
+            await _poolSemaphore.WaitAsync().ConfigureAwait(false);
+            try
             {
-                await Close(reason).ConfigureAwait(false);
-            }
+                if (!HasEntities())
+                {
+                    if (!string.IsNullOrEmpty(ClientId))
+                    {
+                        _logger.LogInformation("Close connection for the {ClientId}", ClientId);
+                        // 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;
+                var result = new CloseResponse(0, ResponseCode.Ok);
+                return result;
+            }
+            finally
+            {
+                _poolSemaphore.Release();
+            }
         }
 
+        public string ClientId { get; init; }
+
         public async ValueTask<QueryPublisherResponse> QueryPublisherSequence(string publisherRef, string stream)
         {
             return await Request<QueryPublisherRequest, QueryPublisherResponse>(corr =>

RabbitMQ.Stream.Client/ClientExceptions.cs

@@ -125,4 +125,12 @@ public CrcException(string s)
         {
         }
     }
+
+    public class TooManyConnectionsException : Exception
+    {
+        public TooManyConnectionsException(string s)
+            : base(s)
+        {
+        }
+    }
 }

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.