* Move publisher confirmation setup / teardown out of BasicPublishAsync.

Author: lukebakken

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

@@ -29,6 +29,7 @@
 //  Copyright (c) 2007-2024 Broadcom. All Rights Reserved.
 //---------------------------------------------------------------------------
 
+using System;
 using System.Buffers.Binary;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
@@ -52,6 +53,70 @@ internal partial class Channel : IChannel, IRecoverable
         private readonly SemaphoreSlim _confirmSemaphore = new(1, 1);
         private readonly ConcurrentDictionary<ulong, TaskCompletionSource<bool>> _confirmsTaskCompletionSources = new();
 
+        private class PublisherConfirmationInfo
+        {
+            private ulong _publishSequenceNumber;
+            private TaskCompletionSource<bool>? _publisherConfirmationTcs;
+
+            internal PublisherConfirmationInfo()
+            {
+                _publishSequenceNumber = 0;
+                _publisherConfirmationTcs = null;
+            }
+
+            internal PublisherConfirmationInfo(ulong publishSequenceNumber, TaskCompletionSource<bool>? publisherConfirmationTcs)
+            {
+                _publishSequenceNumber = publishSequenceNumber;
+                _publisherConfirmationTcs = publisherConfirmationTcs;
+            }
+
+            internal ulong PublishSequenceNumber => _publishSequenceNumber;
+
+            internal TaskCompletionSource<bool>? PublisherConfirmationTcs => _publisherConfirmationTcs;
+
+            internal async Task MaybeWaitForConfirmationAsync(CancellationToken cancellationToken)
+            {
+                if (_publisherConfirmationTcs is not null)
+                {
+                    await _publisherConfirmationTcs.Task.WaitAsync(cancellationToken)
+                        .ConfigureAwait(false);
+                }
+            }
+
+            internal bool MaybeHandleException(Exception ex)
+            {
+                bool exceptionWasHandled = false;
+
+                if (_publisherConfirmationTcs is not null)
+                {
+                    _publisherConfirmationTcs.SetException(ex);
+                    exceptionWasHandled = true;
+                }
+
+                return exceptionWasHandled;
+            }
+        }
+
+        public async ValueTask<ulong> GetNextPublishSequenceNumberAsync(CancellationToken cancellationToken = default)
+        {
+            if (_publisherConfirmationsEnabled)
+            {
+                await _confirmSemaphore.WaitAsync(cancellationToken).ConfigureAwait(false);
+                try
+                {
+                    return _nextPublishSeqNo;
+                }
+                finally
+                {
+                    _confirmSemaphore.Release();
+                }
+            }
+            else
+            {
+                return _nextPublishSeqNo;
+            }
+        }
+
         private void ConfigurePublisherConfirmations(bool publisherConfirmationsEnabled, bool publisherConfirmationTrackingEnabled)
         {
             _publisherConfirmationsEnabled = publisherConfirmationsEnabled;
@@ -98,10 +163,17 @@ await ModelSendAsync(in method, k.CancellationToken)
             }
         }
 
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private bool ShouldHandleAckOrNack(ulong deliveryTag)
+        {
+            return _publisherConfirmationsEnabled && _publisherConfirmationTrackingEnabled &&
+                deliveryTag > 0 && !_confirmsTaskCompletionSources.IsEmpty;
+        }
+
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         private void HandleAck(ulong deliveryTag, bool multiple)
         {
-            if (_publisherConfirmationsEnabled && _publisherConfirmationTrackingEnabled && deliveryTag > 0 && !_confirmsTaskCompletionSources.IsEmpty)
+            if (ShouldHandleAckOrNack(deliveryTag))
             {
                 if (multiple)
                 {
@@ -127,7 +199,7 @@ private void HandleAck(ulong deliveryTag, bool multiple)
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         private void HandleNack(ulong deliveryTag, bool multiple, bool isReturn)
         {
-            if (_publisherConfirmationsEnabled && _publisherConfirmationTrackingEnabled && deliveryTag > 0 && !_confirmsTaskCompletionSources.IsEmpty)
+            if (ShouldHandleAckOrNack(deliveryTag))
             {
                 if (multiple)
                 {
@@ -192,5 +264,67 @@ await _confirmSemaphore.WaitAsync(reason.CancellationToken)
                 }
             }
         }
+
+        private async Task<PublisherConfirmationInfo?> MaybeStartPublisherConfirmationTracking(CancellationToken cancellationToken)
+        {
+            if (_publisherConfirmationsEnabled)
+            {
+                await _confirmSemaphore.WaitAsync(cancellationToken)
+                    .ConfigureAwait(false);
+
+                ulong publishSequenceNumber = _nextPublishSeqNo;
+
+                TaskCompletionSource<bool>? publisherConfirmationTcs = null;
+                if (_publisherConfirmationTrackingEnabled)
+                {
+                    publisherConfirmationTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+                    _confirmsTaskCompletionSources[publishSequenceNumber] = publisherConfirmationTcs;
+                }
+
+                _nextPublishSeqNo++;
+
+                return new PublisherConfirmationInfo(publishSequenceNumber, publisherConfirmationTcs);
+            }
+            else
+            {
+                return null;
+            }
+        }
+
+        private bool MaybeHandleExceptionWithEnabledPublisherConfirmations(PublisherConfirmationInfo? publisherConfirmationInfo,
+            Exception ex)
+        {
+            bool exceptionWasHandled = false;
+
+            if (_publisherConfirmationsEnabled &&
+                publisherConfirmationInfo is not null)
+            {
+                _nextPublishSeqNo--;
+
+                if (_publisherConfirmationTrackingEnabled)
+                {
+                    _confirmsTaskCompletionSources.TryRemove(publisherConfirmationInfo.PublishSequenceNumber, out _);
+                }
+
+                exceptionWasHandled = publisherConfirmationInfo.MaybeHandleException(ex);
+            }
+
+            return exceptionWasHandled;
+        }
+
+        private async Task MaybeEndPublisherConfirmationTracking(PublisherConfirmationInfo? publisherConfirmationInfo,
+            CancellationToken cancellationToken)
+        {
+            if (_publisherConfirmationsEnabled)
+            {
+                _confirmSemaphore.Release();
+
+                if (publisherConfirmationInfo is not null)
+                {
+                    await publisherConfirmationInfo.MaybeWaitForConfirmationAsync(cancellationToken)
+                        .ConfigureAwait(false);
+                }
+            }
+        }
     }
 }

projects/RabbitMQ.Client/Impl/Channel.cs

@@ -807,26 +807,6 @@ await Session.Connection.HandleConnectionUnblockedAsync(cancellationToken)
             return true;
         }
 
-        public async ValueTask<ulong> GetNextPublishSequenceNumberAsync(CancellationToken cancellationToken = default)
-        {
-            if (_publisherConfirmationsEnabled)
-            {
-                await _confirmSemaphore.WaitAsync(cancellationToken).ConfigureAwait(false);
-                try
-                {
-                    return _nextPublishSeqNo;
-                }
-                finally
-                {
-                    _confirmSemaphore.Release();
-                }
-            }
-            else
-            {
-                return _nextPublishSeqNo;
-            }
-        }
-
         public virtual ValueTask BasicAckAsync(ulong deliveryTag, bool multiple,
             CancellationToken cancellationToken)
         {
@@ -969,26 +949,13 @@ public async ValueTask BasicPublishAsync<TProperties>(string exchange, string ro
             CancellationToken cancellationToken = default)
             where TProperties : IReadOnlyBasicProperties, IAmqpHeader
         {
-            TaskCompletionSource<bool>? publisherConfirmationTcs = null;
-            ulong publishSequenceNumber = 0;
+            PublisherConfirmationInfo? publisherConfirmationInfo = null;
             try
             {
-                if (_publisherConfirmationsEnabled)
-                {
-                    await _confirmSemaphore.WaitAsync(cancellationToken)
+                publisherConfirmationInfo =
+                    await MaybeStartPublisherConfirmationTracking(cancellationToken)
                         .ConfigureAwait(false);
 
-                    publishSequenceNumber = _nextPublishSeqNo;
-
-                    if (_publisherConfirmationTrackingEnabled)
-                    {
-                        publisherConfirmationTcs = new(TaskCreationOptions.RunContinuationsAsynchronously);
-                        _confirmsTaskCompletionSources[publishSequenceNumber] = publisherConfirmationTcs;
-                    }
-
-                    _nextPublishSeqNo++;
-                }
-
                 await EnforceFlowControlAsync(cancellationToken)
                     .ConfigureAwait(false);
 
@@ -998,6 +965,12 @@ await EnforceFlowControlAsync(cancellationToken)
                     ? RabbitMQActivitySource.Send(routingKey, exchange, body.Length)
                     : default;
 
+                ulong publishSequenceNumber = 0;
+                if (publisherConfirmationInfo is not null)
+                {
+                    publishSequenceNumber = publisherConfirmationInfo.PublishSequenceNumber;
+                }
+
                 BasicProperties? props = PopulateBasicPropertiesHeaders(basicProperties, sendActivity, publishSequenceNumber);
                 if (props is null)
                 {
@@ -1012,35 +985,16 @@ await ModelSendAsync(in cmd, in props, body, cancellationToken)
             }
             catch (Exception ex)
             {
-                if (_publisherConfirmationsEnabled)
-                {
-                    _nextPublishSeqNo--;
-                    if (_publisherConfirmationTrackingEnabled)
-                    {
-                        _confirmsTaskCompletionSources.TryRemove(publishSequenceNumber, out _);
-                    }
-                }
-
-                if (publisherConfirmationTcs is not null)
-                {
-                    publisherConfirmationTcs.SetException(ex);
-                }
-                else
+                bool exceptionWasHandled =
+                    MaybeHandleExceptionWithEnabledPublisherConfirmations(publisherConfirmationInfo, ex);
+                if (!exceptionWasHandled)
                 {
                     throw;
                 }
             }
             finally
             {
-                if (_publisherConfirmationsEnabled)
-                {
-                    _confirmSemaphore.Release();
-                }
-            }
-
-            if (publisherConfirmationTcs is not null)
-            {
-                await publisherConfirmationTcs.Task.WaitAsync(cancellationToken)
+                await MaybeEndPublisherConfirmationTracking(publisherConfirmationInfo, cancellationToken)
                     .ConfigureAwait(false);
             }
         }