Fix unconditional semaphore release in BasicPublishAsync when CancellationToken is cancelled

Author: EvheniyHlushko

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

@@ -54,20 +54,34 @@ internal partial class Channel : IChannel, IRecoverable
 
         private sealed class PublisherConfirmationInfo : IDisposable
         {
-            private TaskCompletionSource<bool>? _publisherConfirmationTcs;
+            private readonly SemaphoreSlim _semaphore;
+            private readonly TaskCompletionSource<bool>? _publisherConfirmationTcs;
             private readonly RateLimitLease? _lease;
+            private bool _semaphoreReleased;
 
-            internal PublisherConfirmationInfo(ulong publishSequenceNumber,
+            internal PublisherConfirmationInfo(
+                SemaphoreSlim semaphore,
+                ulong publishSequenceNumber,
                 TaskCompletionSource<bool>? publisherConfirmationTcs,
                 RateLimitLease? lease)
             {
+                _semaphore = semaphore;
                 PublishSequenceNumber = publishSequenceNumber;
                 _publisherConfirmationTcs = publisherConfirmationTcs;
                 _lease = lease;
             }
 
             internal ulong PublishSequenceNumber { get; }
 
+            internal void ReleaseSemaphore()
+            {
+                if (!_semaphoreReleased)
+                {
+                    _semaphoreReleased = true;
+                    _semaphore.Release();
+                }
+            }
+
             internal async Task MaybeWaitForConfirmationAsync(CancellationToken cancellationToken)
             {
                 if (_publisherConfirmationTcs is not null)
@@ -92,6 +106,7 @@ internal bool MaybeHandleException(Exception ex)
 
             public void Dispose()
             {
+                ReleaseSemaphore();
                 _lease?.Dispose();
             }
         }
@@ -291,47 +306,50 @@ await _confirmSemaphore.WaitAsync(reason.CancellationToken)
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         private async Task<PublisherConfirmationInfo?> MaybeStartPublisherConfirmationTrackingAsync(CancellationToken cancellationToken)
         {
-            if (_publisherConfirmationsEnabled)
+            if (!_publisherConfirmationsEnabled)
             {
-                RateLimitLease? lease = null;
-                if (_publisherConfirmationTrackingEnabled)
-                {
-                    if (_outstandingPublisherConfirmationsRateLimiter is not null)
-                    {
-                        lease = await _outstandingPublisherConfirmationsRateLimiter.AcquireAsync(
-                            cancellationToken: cancellationToken)
-                            .ConfigureAwait(false);
+                return null;
+            }
 
-                        if (!lease.IsAcquired)
-                        {
-                            throw new InvalidOperationException("Could not acquire a lease from the rate limiter.");
-                        }
-                    }
+            RateLimitLease? lease = null;
+            if (_publisherConfirmationTrackingEnabled && _outstandingPublisherConfirmationsRateLimiter is not null)
+            {
+                lease = await _outstandingPublisherConfirmationsRateLimiter.AcquireAsync(
+                    cancellationToken: cancellationToken)
+                    .ConfigureAwait(false);
+                if (!lease.IsAcquired)
+                {
+                    throw new InvalidOperationException("Could not acquire a lease from the rate limiter.");
                 }
+            }
 
-                await _confirmSemaphore.WaitAsync(cancellationToken)
-                    .ConfigureAwait(false);
+            try
+            {
+                await _confirmSemaphore.WaitAsync(cancellationToken).ConfigureAwait(false);
+            }
+            catch (OperationCanceledException)
+            {
+                lease?.Dispose();
+                throw;
+            }
 
-                ulong publishSequenceNumber = _nextPublishSeqNo;
+            ulong publishSequenceNumber = _nextPublishSeqNo;
 
-                TaskCompletionSource<bool>? publisherConfirmationTcs = null;
-                if (_publisherConfirmationTrackingEnabled)
+            TaskCompletionSource<bool>? publisherConfirmationTcs = null;
+            if (_publisherConfirmationTrackingEnabled)
+            {
+                publisherConfirmationTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+                if (!_confirmsTaskCompletionSources.TryAdd(publishSequenceNumber, publisherConfirmationTcs))
                 {
-                    publisherConfirmationTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
-                    if (!_confirmsTaskCompletionSources.TryAdd(publishSequenceNumber, publisherConfirmationTcs))
-                    {
-                        throw new InvalidOperationException($"Failed to track the publisher confirmation for sequence number '{publishSequenceNumber}' because it already exists.");
-                    }
+                    _confirmSemaphore.Release();
+                    lease?.Dispose();
+                    throw new InvalidOperationException($"Failed to track the publisher confirmation for sequence number '{publishSequenceNumber}' because it already exists.");
                 }
+            }
 
-                _nextPublishSeqNo++;
+            _nextPublishSeqNo++;
 
-                return new PublisherConfirmationInfo(publishSequenceNumber, publisherConfirmationTcs, lease);
-            }
-            else
-            {
-                return null;
-            }
+            return new PublisherConfirmationInfo(_confirmSemaphore, publishSequenceNumber, publisherConfirmationTcs, lease);
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
@@ -360,44 +378,26 @@ private bool MaybeHandleExceptionWithEnabledPublisherConfirmations(PublisherConf
         private async Task MaybeEndPublisherConfirmationTrackingAsync(PublisherConfirmationInfo? publisherConfirmationInfo,
             CancellationToken cancellationToken)
         {
-            if (_publisherConfirmationsEnabled)
+            if (publisherConfirmationInfo is null)
             {
-                try
-                {
-                    _confirmSemaphore.Release();
-                }
-                catch (SemaphoreFullException ex)
-                {
-                    /*
-                     * rabbitmq/rabbitmq-dotnet-client-1793
-                     * If MaybeStartPublisherConfirmationTracking throws an exception *prior* to acquiring
-                     * _confirmSemaphore, the above Release() call will throw SemaphoreFullException.
-                     * In "normal" cases, publisherConfirmationInfo will thus be null, but if not, throw
-                     * a "bug found" exception here.
-                     */
-                    if (publisherConfirmationInfo is not null)
-                    {
-                        throw new InvalidOperationException(InternalConstants.BugFound, ex);
-                    }
-                }
+                return;
+            }
 
-                if (publisherConfirmationInfo is not null)
-                {
-                    try
-                    {
-                        await publisherConfirmationInfo.MaybeWaitForConfirmationAsync(cancellationToken)
-                            .ConfigureAwait(false);
-                    }
-                    catch (OperationCanceledException)
-                    {
-                        _confirmsTaskCompletionSources.TryRemove(publisherConfirmationInfo.PublishSequenceNumber, out _);
-                        throw;
-                    }
-                    finally
-                    {
-                        publisherConfirmationInfo.Dispose();
-                    }
-                }
+            publisherConfirmationInfo.ReleaseSemaphore();
+
+            try
+            {
+                await publisherConfirmationInfo.MaybeWaitForConfirmationAsync(cancellationToken)
+                    .ConfigureAwait(false);
+            }
+            catch (OperationCanceledException)
+            {
+                _confirmsTaskCompletionSources.TryRemove(publisherConfirmationInfo.PublishSequenceNumber, out _);
+                throw;
+            }
+            finally
+            {
+                publisherConfirmationInfo.Dispose();
             }
         }