Fix race conditions in ConnectionPool and SubProducer causing ThreadPool starvation

src/DotPulsar/Internal/ConnectionPool.cs

@@ -29,6 +29,7 @@ public sealed class ConnectionPool : IConnectionPool
     private readonly Connector _connector;
     private readonly EncryptionPolicy _encryptionPolicy;
     private readonly ConcurrentDictionary<PulsarUrl, Connection> _connections;
+    private readonly ConcurrentDictionary<PulsarUrl, SemaphoreSlim> _connectionGates;
     private readonly CancellationTokenSource _cancellationTokenSource;
     private readonly string? _listenerName;
     private readonly TimeSpan _closeInactiveConnectionsInterval;
@@ -51,6 +52,7 @@ public ConnectionPool(
         _encryptionPolicy = encryptionPolicy;
         _listenerName = listenerName;
         _connections = new ConcurrentDictionary<PulsarUrl, Connection>();
+        _connectionGates = new ConcurrentDictionary<PulsarUrl, SemaphoreSlim>();
         _cancellationTokenSource = new CancellationTokenSource();
         _closeInactiveConnectionsInterval = closeInactiveConnectionsInterval;
         _keepAliveInterval = keepAliveInterval;
@@ -65,6 +67,11 @@ public async ValueTask DisposeAsync()
         {
             await DisposeConnection(entry.Key, entry.Value).ConfigureAwait(false);
         }
+
+        foreach (var gate in _connectionGates.Values)
+        {
+            gate.Dispose();
+        }
     }
 
     public async ValueTask<IConnection> FindConnectionForTopic(string topic, CancellationToken cancellationToken)
@@ -143,7 +150,19 @@ private async ValueTask<Connection> GetConnection(PulsarUrl url, CancellationTok
         if (_connections.TryGetValue(url, out var connection) && connection is not null)
             return connection;
 
-        return await EstablishNewConnection(url, cancellationToken).ConfigureAwait(false);
+        var gate = _connectionGates.GetOrAdd(url, _ => new SemaphoreSlim(1, 1));
+        await gate.WaitAsync(cancellationToken).ConfigureAwait(false);
+        try
+        {
+            if (_connections.TryGetValue(url, out connection) && connection is not null)
+                return connection;
+
+            return await EstablishNewConnection(url, cancellationToken).ConfigureAwait(false);
+        }
+        finally
+        {
+            gate.Release();
+        }
     }
 
     private async Task<Connection> EstablishNewConnection(PulsarUrl url, CancellationToken cancellationToken)
@@ -165,7 +184,12 @@ private async Task<Connection> EstablishNewConnection(PulsarUrl url, Cancellatio
 
     private async ValueTask DisposeConnection(PulsarUrl serviceUrl, Connection connection)
     {
-        _connections.TryRemove(serviceUrl, out var _);
+        var pair = new KeyValuePair<PulsarUrl, Connection>(serviceUrl, connection);
+#if NETSTANDARD2_0 || NETSTANDARD2_1
+        ((ICollection<KeyValuePair<PulsarUrl, Connection>>)_connections).Remove(pair);
+#else
+        _connections.TryRemove(pair);
+#endif
         await connection.DisposeAsync().ConfigureAwait(false);
     }
 

src/DotPulsar/Internal/SubProducer.cs

@@ -219,8 +219,12 @@ public async Task EstablishNewChannel(CancellationToken cancellationToken)
         }
 
         _channel = await _executor.Execute(() => _factory.Create(cancellationToken), cancellationToken).ConfigureAwait(false);
-        _dispatcherCts = new CancellationTokenSource();
-        _dispatcherTask = Task.Run(async () => await MessageDispatcher(_channel, _dispatcherCts.Token), CancellationToken.None);
+
+        var cts = new CancellationTokenSource();
+        var token = cts.Token;
+        var channel = _channel;
+        _dispatcherCts = cts;
+        _dispatcherTask = Task.Run(() => MessageDispatcher(channel, token), CancellationToken.None);
     }
 
     public async ValueTask CloseChannel(CancellationToken cancellationToken)

tests/DotPulsar.Tests/IntegrationFixture.cs

@@ -73,6 +73,15 @@ public IntegrationFixture()
 
     public IAuthentication Authentication => AuthenticationFactory.Token(ct => ValueTask.FromResult(_token!));
 
+    public HttpClient CreateAdminClient() => new()
+    {
+        BaseAddress = AdminUrl,
+        DefaultRequestHeaders =
+        {
+            Authorization = AuthorizationHeader
+        }
+    };
+
     public async ValueTask DisposeAsync()
     {
         await _pulsarCluster.DisposeAsync();

tests/DotPulsar.Tests/Internal/ConnectionPoolTests.cs

@@ -0,0 +1,143 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+namespace DotPulsar.Tests.Internal;
+
+using DotPulsar.Abstractions;
+using DotPulsar.Extensions;
+using System.Text.Json;
+
+[Collection("Integration"), Trait("Category", "Integration")]
+public sealed class ConnectionPoolTests : IDisposable
+{
+    private const int ProducerCount = 20;
+
+    private readonly CancellationTokenSource _cts;
+    private readonly IntegrationFixture _fixture;
+    private readonly ITestOutputHelper _testOutputHelper;
+
+    public ConnectionPoolTests(IntegrationFixture fixture, ITestOutputHelper outputHelper)
+    {
+        _cts = new CancellationTokenSource(TimeSpan.FromMinutes(2));
+        _fixture = fixture;
+        _testOutputHelper = outputHelper;
+    }
+
+    [Fact]
+    public async Task Connectivity_WhenManyProducersReconnectSimultaneously_AllShouldShareOneConnection()
+    {
+        //Arrange
+        var topicName = await _fixture.CreateTopic(_cts.Token);
+        await using var client = CreateClient();
+        var producers = Enumerable.Range(0, ProducerCount).Select(_ => CreateProducer(client, topicName)).ToArray();
+
+        try
+        {
+            await Task.WhenAll(producers.Select(p => p.State.OnStateChangeTo(ProducerState.Connected, _cts.Token).AsTask()));
+
+            //Act
+            await using (await _fixture.DisableThePulsarConnection())
+            {
+                await Task.WhenAll(producers.Select(p => p.StateChangedTo(ProducerState.Disconnected, _cts.Token).AsTask()));
+            }
+
+            await Task.WhenAll(producers.Select(p => p.State.OnStateChangeTo(ProducerState.Connected, _cts.Token).AsTask()));
+
+            //Assert
+            using var adminClient = CreateAdminClient();
+            var uniqueConnections = await GetUniqueProducerConnectionCount(adminClient, topicName, _cts.Token);
+            uniqueConnections.ShouldBe(1);
+        }
+        finally
+        {
+            await Task.WhenAll(producers.Select(p => p.DisposeAsync().AsTask()));
+        }
+    }
+
+    [Fact]
+    public async Task Connectivity_WhenManyProducersReconnectTwice_AllShouldShareOneConnection()
+    {
+        //Arrange
+        var topicName = await _fixture.CreateTopic(_cts.Token);
+        await using var client = CreateClient();
+        var producers = Enumerable.Range(0, ProducerCount).Select(_ => CreateProducer(client, topicName)).ToArray();
+
+        try
+        {
+            await Task.WhenAll(producers.Select(p => p.State.OnStateChangeTo(ProducerState.Connected, _cts.Token).AsTask()));
+
+            //Act
+            for (var cycle = 0; cycle < 2; cycle++)
+            {
+                await using (await _fixture.DisableThePulsarConnection())
+                {
+                    await Task.WhenAll(producers.Select(p => p.StateChangedTo(ProducerState.Disconnected, _cts.Token).AsTask()));
+                }
+
+                await Task.WhenAll(producers.Select(p => p.State.OnStateChangeTo(ProducerState.Connected, _cts.Token).AsTask()));
+            }
+
+            //Assert
+            using var adminClient = CreateAdminClient();
+            var uniqueConnections = await GetUniqueProducerConnectionCount(adminClient, topicName, _cts.Token);
+            uniqueConnections.ShouldBe(1);
+        }
+        finally
+        {
+            await Task.WhenAll(producers.Select(p => p.DisposeAsync().AsTask()));
+        }
+    }
+
+    private static async ValueTask<int> GetUniqueProducerConnectionCount(HttpClient httpClient, string topicName, CancellationToken cancellationToken)
+    {
+        var topic = topicName.Replace("persistent://", string.Empty);
+        using var response = await httpClient.GetAsync($"/admin/v2/persistent/{topic}/stats", cancellationToken).ConfigureAwait(false);
+
+        if (!response.IsSuccessStatusCode)
+            return 0;
+
+        await using var stream = await response.Content.ReadAsStreamAsync(cancellationToken).ConfigureAwait(false);
+        var json = await JsonDocument.ParseAsync(stream, cancellationToken: cancellationToken).ConfigureAwait(false);
+
+        // to comply both Pulsar 3.x and 4.x
+        if (!json.RootElement.TryGetProperty("publishers", out var producers) &&
+            !json.RootElement.TryGetProperty("producers", out producers))
+            return 0;
+
+        return producers.EnumerateArray()
+            .Select(p => p.TryGetProperty("address", out var addr) ? addr.GetString() : null)
+            .Where(a => a is not null)
+            .Distinct()
+            .Count();
+    }
+
+    private IProducer<string> CreateProducer(IPulsarClient pulsarClient, string topicName)
+        => pulsarClient
+            .NewProducer(Schema.String)
+            .Topic(topicName)
+            .StateChangedHandler(_testOutputHelper.Log)
+            .Create();
+
+    private IPulsarClient CreateClient()
+        => PulsarClient
+            .Builder()
+            .Authentication(_fixture.Authentication)
+            .ExceptionHandler(_testOutputHelper.Log)
+            .ServiceUrl(_fixture.ServiceUrl)
+            .Build();
+
+    private HttpClient CreateAdminClient() => _fixture.CreateAdminClient();
+
+    public void Dispose() => _cts.Dispose();
+}

tests/DotPulsar.Tests/Internal/ConsumerTests.cs

@@ -591,14 +591,7 @@ private IPulsarClient CreateClient()
         .ServiceUrl(_fixture.ServiceUrl)
         .Build();
 
-    private HttpClient CreateAdminClient() => new()
-    {
-        BaseAddress = _fixture.AdminUrl,
-        DefaultRequestHeaders =
-        {
-            Authorization = _fixture.AuthorizationHeader
-        }
-    };
+    private HttpClient CreateAdminClient() => _fixture.CreateAdminClient();
 
     private static async ValueTask<long> GetPermits(HttpClient httpClient, string topic, string subscription, CancellationToken cancellationToken)
     {

tests/DotPulsar.Tests/Internal/SubProducerTests.cs

@@ -0,0 +1,115 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+namespace DotPulsar.Tests.Internal;
+
+using DotPulsar.Abstractions;
+using DotPulsar.Internal;
+using DotPulsar.Internal.Abstractions;
+
+[Trait("Category", "Unit")]
+public class SubProducerTests
+{
+    [Fact]
+    public async Task EstablishNewChannel_WhenCalledThreeTimesInSuccession_ShouldNotThrowObjectDisposedException()
+    {
+        //Arrange
+        var factory = Substitute.For<IProducerChannelFactory>();
+        factory.Create(Arg.Any<CancellationToken>())
+            .Returns(_ => Task.FromResult(Substitute.For<IProducerChannel>()));
+
+        await using var sut = CreateSubProducer(factory);
+
+        //Act
+        var exception = await Record.ExceptionAsync(async () =>
+        {
+            await sut.EstablishNewChannel(CancellationToken.None);
+            await sut.EstablishNewChannel(CancellationToken.None);
+            await sut.EstablishNewChannel(CancellationToken.None);
+        });
+
+        //Assert
+        exception.ShouldNotBeOfType<ObjectDisposedException>();
+    }
+
+    private static SubProducer CreateSubProducer(IProducerChannelFactory factory)
+        => new(
+            correlationId: Guid.NewGuid(),
+            registerEvent: Substitute.For<IRegisterEvent>(),
+            initialChannel: Substitute.For<IProducerChannel>(),
+            executor: new DirectExecutor(),
+            state: Substitute.For<IState<ProducerState>>(),
+            factory: factory,
+            partition: 0,
+            maxPendingMessages: 1000,
+            topic: "persistent://public/default/test");
+
+    private sealed class DirectExecutor : IExecute
+    {
+        public ValueTask Execute(Action action, CancellationToken cancellationToken = default)
+        {
+            action();
+            return ValueTask.CompletedTask;
+        }
+
+        public async ValueTask Execute(Func<Task> func, CancellationToken cancellationToken = default)
+            => await func().ConfigureAwait(false);
+
+        public async ValueTask Execute(Func<ValueTask> func, CancellationToken cancellationToken = default)
+            => await func().ConfigureAwait(false);
+
+        public ValueTask<TResult> Execute<TResult>(Func<TResult> func, CancellationToken cancellationToken = default)
+            => ValueTask.FromResult(func());
+
+        public async ValueTask<TResult> Execute<TResult>(Func<Task<TResult>> func, CancellationToken cancellationToken = default)
+            => await func().ConfigureAwait(false);
+
+        public async ValueTask<TResult> Execute<TResult>(Func<ValueTask<TResult>> func, CancellationToken cancellationToken = default)
+            => await func().ConfigureAwait(false);
+
+        public ValueTask<bool> TryExecuteOnce(Action action, CancellationToken cancellationToken = default)
+        {
+            action();
+            return ValueTask.FromResult(true);
+        }
+
+        public async ValueTask<bool> TryExecuteOnce(Func<Task> func, CancellationToken cancellationToken = default)
+        {
+            try
+            {
+                await func().ConfigureAwait(false);
+                return true;
+            }
+            catch (OperationCanceledException)
+            {
+                return false;
+            }
+        }
+
+        public async ValueTask<bool> TryExecuteOnce(Func<ValueTask> func, CancellationToken cancellationToken = default)
+        {
+            await func().ConfigureAwait(false);
+            return true;
+        }
+
+        public ValueTask<ExecutionResult<TResult>> TryExecuteOnce<TResult>(Func<TResult> func, CancellationToken cancellationToken = default)
+            => ValueTask.FromResult(new ExecutionResult<TResult>(true, func()));
+
+        public async ValueTask<ExecutionResult<TResult>> TryExecuteOnce<TResult>(Func<Task<TResult>> func, CancellationToken cancellationToken = default)
+            => new ExecutionResult<TResult>(true, await func().ConfigureAwait(false));
+
+        public async ValueTask<ExecutionResult<TResult>> TryExecuteOnce<TResult>(Func<ValueTask<TResult>> func, CancellationToken cancellationToken = default)
+            => new ExecutionResult<TResult>(true, await func().ConfigureAwait(false));
+    }
+}