* Use SemaphoreSlim to block publishers when max outstanding is reached.

Author: lukebakken

projects/RabbitMQ.Client/Impl/Channel.BasicPublish.cs

@@ -59,9 +59,6 @@ await MaybeStartPublisherConfirmationTracking(cancellationToken)
                 await MaybeEnforceFlowControlAsync(cancellationToken)
                     .ConfigureAwait(false);
 
-                await MaybeEnforceOutstandingPublisherConfirmationsAsync(cancellationToken)
-                    .ConfigureAwait(false);
-
                 var cmd = new BasicPublish(exchange, routingKey, mandatory, default);
 
                 using Activity? sendActivity = RabbitMQActivitySource.PublisherHasListeners
@@ -117,9 +114,6 @@ await MaybeStartPublisherConfirmationTracking(cancellationToken)
                 await MaybeEnforceFlowControlAsync(cancellationToken)
                     .ConfigureAwait(false);
 
-                await MaybeEnforceOutstandingPublisherConfirmationsAsync(cancellationToken)
-                    .ConfigureAwait(false);
-
                 var cmd = new BasicPublishMemory(exchange.Bytes, routingKey.Bytes, mandatory, default);
 
                 using Activity? sendActivity = RabbitMQActivitySource.PublisherHasListeners

projects/RabbitMQ.Client/Impl/Channel.PublisherConfirms.cs

@@ -48,10 +48,10 @@ internal partial class Channel : IChannel, IRecoverable
         private bool _publisherConfirmationsEnabled = false;
         private bool _publisherConfirmationTrackingEnabled = false;
         private ushort? _maxOutstandingPublisherConfirmations = null;
+        private SemaphoreSlim? _maxOutstandingConfirmationsSemaphore;
         private ulong _nextPublishSeqNo = 0;
         private readonly SemaphoreSlim _confirmSemaphore = new(1, 1);
         private readonly ConcurrentDictionary<ulong, TaskCompletionSource<bool>> _confirmsTaskCompletionSources = new();
-        private readonly AsyncManualResetEvent _maxOutstandingPublisherConfirmsReached = new(true);
 
         private class PublisherConfirmationInfo
         {
@@ -124,6 +124,13 @@ private void ConfigurePublisherConfirmations(bool publisherConfirmationsEnabled,
             _publisherConfirmationsEnabled = publisherConfirmationsEnabled;
             _publisherConfirmationTrackingEnabled = publisherConfirmationTrackingEnabled;
             _maxOutstandingPublisherConfirmations = maxOutstandingPublisherConfirmations;
+
+            if (_maxOutstandingPublisherConfirmations is not null)
+            {
+                _maxOutstandingConfirmationsSemaphore = new SemaphoreSlim(
+                    (int)_maxOutstandingPublisherConfirmations,
+                    (int)_maxOutstandingPublisherConfirmations);
+            }
         }
 
         private async Task MaybeConfirmSelect(CancellationToken cancellationToken)
@@ -141,7 +148,6 @@ private async Task MaybeConfirmSelect(CancellationToken cancellationToken)
                         if (_publisherConfirmationTrackingEnabled)
                         {
                             _confirmsTaskCompletionSources.Clear();
-                            MaybeUnblockPublishers();
                         }
                         _nextPublishSeqNo = 1;
                     }
@@ -187,15 +193,13 @@ private void HandleAck(ulong deliveryTag, bool multiple)
                         {
                             pair.Value.SetResult(true);
                             _confirmsTaskCompletionSources.Remove(pair.Key, out _);
-                            MaybeUnblockPublishers();
                         }
                     }
                 }
                 else
                 {
                     if (_confirmsTaskCompletionSources.TryRemove(deliveryTag, out TaskCompletionSource<bool>? tcs))
                     {
-                        MaybeUnblockPublishers();
                         tcs.SetResult(true);
                     }
                 }
@@ -215,15 +219,13 @@ private void HandleNack(ulong deliveryTag, bool multiple, bool isReturn)
                         {
                             pair.Value.SetException(new PublishException(pair.Key, isReturn));
                             _confirmsTaskCompletionSources.Remove(pair.Key, out _);
-                            MaybeUnblockPublishers();
                         }
                     }
                 }
                 else
                 {
                     if (_confirmsTaskCompletionSources.Remove(deliveryTag, out TaskCompletionSource<bool>? tcs))
                     {
-                        MaybeUnblockPublishers();
                         tcs.SetException(new PublishException(deliveryTag, isReturn));
                     }
                 }
@@ -266,7 +268,6 @@ await _confirmSemaphore.WaitAsync(reason.CancellationToken)
                         }
 
                         _confirmsTaskCompletionSources.Clear();
-                        MaybeUnblockPublishers();
                     }
                 }
                 finally
@@ -291,7 +292,12 @@ await _confirmSemaphore.WaitAsync(cancellationToken)
                 {
                     publisherConfirmationTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
                     _confirmsTaskCompletionSources[publishSequenceNumber] = publisherConfirmationTcs;
-                    MaybeBlockPublishers();
+                }
+
+                if (_maxOutstandingConfirmationsSemaphore is not null)
+                {
+                    await _maxOutstandingConfirmationsSemaphore.WaitAsync(cancellationToken)
+                        .ConfigureAwait(false);
                 }
 
                 _nextPublishSeqNo++;
@@ -318,7 +324,6 @@ private bool MaybeHandleExceptionWithEnabledPublisherConfirmations(PublisherConf
                 if (_publisherConfirmationTrackingEnabled)
                 {
                     _confirmsTaskCompletionSources.TryRemove(publisherConfirmationInfo.PublishSequenceNumber, out _);
-                    MaybeUnblockPublishers();
                 }
 
                 exceptionWasHandled = publisherConfirmationInfo.MaybeHandleException(ex);
@@ -340,49 +345,10 @@ private async Task MaybeEndPublisherConfirmationTracking(PublisherConfirmationIn
                     await publisherConfirmationInfo.MaybeWaitForConfirmationAsync(cancellationToken)
                         .ConfigureAwait(false);
                 }
-            }
-        }
-
-        [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        private ValueTask MaybeEnforceOutstandingPublisherConfirmationsAsync(CancellationToken cancellationToken)
-        {
-            if (_publisherConfirmationTrackingEnabled)
-            {
-                if (_maxOutstandingPublisherConfirmsReached.IsSet)
-                {
-                    return default;
-                }
-                else
-                {
-                    return _maxOutstandingPublisherConfirmsReached.WaitAsync(cancellationToken);
-                }
-            }
-
-            return default;
-        }
-
-        [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        private void MaybeBlockPublishers()
-        {
-            if (_publisherConfirmationTrackingEnabled)
-            {
-                if (_maxOutstandingPublisherConfirmations is not null
-                    && _confirmsTaskCompletionSources.Count >= _maxOutstandingPublisherConfirmations)
-                {
-                    _maxOutstandingPublisherConfirmsReached.Reset();
-                }
-            }
-        }
 
-        [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        private void MaybeUnblockPublishers()
-        {
-            if (_publisherConfirmationTrackingEnabled)
-            {
-                if (_maxOutstandingPublisherConfirmations is not null
-                    && _confirmsTaskCompletionSources.Count < _maxOutstandingPublisherConfirmations)
+                if (_maxOutstandingConfirmationsSemaphore is not null)
                 {
-                    _maxOutstandingPublisherConfirmsReached.Set();
+                    _maxOutstandingConfirmationsSemaphore.Release();
                 }
             }
         }

projects/RabbitMQ.Client/Impl/Channel.cs

@@ -531,7 +531,8 @@ protected virtual void Dispose(bool disposing)
 
                 ConsumerDispatcher.Dispose();
                 _rpcSemaphore.Dispose();
-                _confirmSemaphore?.Dispose();
+                _confirmSemaphore.Dispose();
+                _maxOutstandingConfirmationsSemaphore?.Dispose();
             }
         }
 
@@ -552,7 +553,8 @@ protected virtual async ValueTask DisposeAsyncCore()
 
             ConsumerDispatcher.Dispose();
             _rpcSemaphore.Dispose();
-            _confirmSemaphore?.Dispose();
+            _confirmSemaphore.Dispose();
+            _maxOutstandingConfirmationsSemaphore?.Dispose();
         }
 
         public Task ConnectionTuneOkAsync(ushort channelMax, uint frameMax, ushort heartbeat, CancellationToken cancellationToken)