Replace lock with SemaphoreSlim

lukebakken · lukebakken · commit 17188c4a5acf · 2024-04-22T08:36:14.000-07:00
diff --git a/projects/RabbitMQ.Client/client/impl/ChannelBase.cs b/projects/RabbitMQ.Client/client/impl/ChannelBase.cs
@@ -59,8 +59,7 @@ internal abstract class ChannelBase : IChannel, IRecoverable
         private readonly RpcContinuationQueue _continuationQueue = new RpcContinuationQueue();
         private readonly ManualResetEventSlim _flowControlBlock = new ManualResetEventSlim(true);
 
-        // TODO replace with SemaphoreSlim
-        private object _confirmLock;
+        private SemaphoreSlim _confirmSemaphore;
         private readonly LinkedList<ulong> _pendingDeliveryTags = new LinkedList<ulong>();
 
         private bool _onlyAcksReceived = true;
@@ -420,7 +419,7 @@ internal void FinishClose()
             m_connectionStartCell?.TrySetResult(null);
         }
 
-        private bool ConfirmsAreEnabled => _confirmLock != null;
+        private bool ConfirmsAreEnabled => _confirmSemaphore != null;
 
         private async Task HandleCommandAsync(IncomingCommand cmd, CancellationToken cancellationToken)
         {
@@ -484,7 +483,8 @@ private void OnChannelShutdown(ShutdownEventArgs reason)
 
             if (ConfirmsAreEnabled)
             {
-                lock (_confirmLock)
+                _confirmSemaphore.Wait();
+                try
                 {
                     if (_confirmsTaskCompletionSources?.Count > 0)
                     {
@@ -497,6 +497,10 @@ private void OnChannelShutdown(ShutdownEventArgs reason)
                         _confirmsTaskCompletionSources.Clear();
                     }
                 }
+                finally
+                {
+                    _confirmSemaphore.Release();
+                }
             }
 
             _flowControlBlock.Set();
@@ -542,6 +546,7 @@ protected virtual void Dispose(bool disposing)
 
                 ConsumerDispatcher.Dispose();
                 _rpcSemaphore.Dispose();
+                _confirmSemaphore?.Dispose();
             }
         }
 
@@ -596,7 +601,8 @@ protected void HandleAckNack(ulong deliveryTag, bool multiple, bool isNack)
             if (ConfirmsAreEnabled)
             {
                 // let's take a lock so we can assume that deliveryTags are unique, never duplicated and always sorted
-                lock (_confirmLock)
+                _confirmSemaphore.Wait();
+                try
                 {
                     // No need to do anything if there are no delivery tags in the list
                     if (_pendingDeliveryTags.Count > 0)
@@ -633,6 +639,10 @@ protected void HandleAckNack(ulong deliveryTag, bool multiple, bool isNack)
                         _onlyAcksReceived = true;
                     }
                 }
+                finally
+                {
+                    _confirmSemaphore.Release();
+                }
             }
         }
 
@@ -1054,10 +1064,16 @@ public async ValueTask BasicPublishAsync<TProperties>(string exchange, string ro
         {
             if (ConfirmsAreEnabled)
             {
-                lock (_confirmLock)
+                await _confirmSemaphore.WaitAsync(cancellationToken)
+                    .ConfigureAwait(false);
+                try
                 {
                     _pendingDeliveryTags.AddLast(NextPublishSeqNo++);
                 }
+                finally
+                {
+                    _confirmSemaphore.Release();
+                }
             }
 
             try
@@ -1084,11 +1100,17 @@ await ModelSendAsync(in cmd, in basicProperties, body, cancellationToken)
             {
                 if (ConfirmsAreEnabled)
                 {
-                    lock (_confirmLock)
+                    await _confirmSemaphore.WaitAsync(cancellationToken)
+                        .ConfigureAwait(false);
+                    try
                     {
                         NextPublishSeqNo--;
                         _pendingDeliveryTags.RemoveLast();
                     }
+                    finally
+                    {
+                        _confirmSemaphore.Release();
+                    }
                 }
 
                 throw;
@@ -1102,10 +1124,16 @@ public async ValueTask BasicPublishAsync<TProperties>(CachedString exchange, Cac
         {
             if (ConfirmsAreEnabled)
             {
-                lock (_confirmLock)
+                await _confirmSemaphore.WaitAsync(cancellationToken)
+                    .ConfigureAwait(false);
+                try
                 {
                     _pendingDeliveryTags.AddLast(NextPublishSeqNo++);
                 }
+                finally
+                {
+                    _confirmSemaphore.Release();
+                }
             }
 
             try
@@ -1133,11 +1161,17 @@ await ModelSendAsync(in cmd, in basicProperties, body, cancellationToken)
             {
                 if (ConfirmsAreEnabled)
                 {
-                    lock (_confirmLock)
+                    await _confirmSemaphore.WaitAsync(cancellationToken)
+                        .ConfigureAwait(false);
+                    try
                     {
                         NextPublishSeqNo--;
                         _pendingDeliveryTags.RemoveLast();
                     }
+                    finally
+                    {
+                        _confirmSemaphore.Release();
+                    }
                 }
 
                 throw;
@@ -1242,7 +1276,7 @@ await ModelSendAsync(method, k.CancellationToken)
 
                 // Note:
                 // Non-null means confirms are enabled
-                _confirmLock = new object();
+                _confirmSemaphore = new SemaphoreSlim(1, 1);
 
                 return;
             }
@@ -1742,63 +1776,50 @@ await ModelSendAsync(method, k.CancellationToken)
 
         private List<TaskCompletionSource<bool>> _confirmsTaskCompletionSources;
 
-        public Task<bool> WaitForConfirmsAsync(CancellationToken token = default)
+        public async Task<bool> WaitForConfirmsAsync(CancellationToken cancellationToken = default)
         {
             if (false == ConfirmsAreEnabled)
             {
                 throw new InvalidOperationException("Confirms not selected");
             }
 
             TaskCompletionSource<bool> tcs;
-            lock (_confirmLock)
+            await _confirmSemaphore.WaitAsync(cancellationToken)
+                .ConfigureAwait(false);
+            try
             {
                 if (_pendingDeliveryTags.Count == 0)
                 {
                     if (_onlyAcksReceived == false)
                     {
                         _onlyAcksReceived = true;
-                        return Task.FromResult(false);
+                        return false;
                     }
 
-                    return Task.FromResult(true);
+                    return true;
                 }
 
                 tcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
                 _confirmsTaskCompletionSources.Add(tcs);
             }
-
-            if (!token.CanBeCanceled)
+            finally
             {
-                return tcs.Task;
+                _confirmSemaphore.Release();
             }
 
-            return WaitForConfirmsWithTokenAsync(tcs, token);
-        }
+            bool rv;
 
-        private async Task<bool> WaitForConfirmsWithTokenAsync(TaskCompletionSource<bool> tcs, CancellationToken token)
-        {
-            CancellationTokenRegistration tokenRegistration =
-#if NET6_0_OR_GREATER
-                token.UnsafeRegister(
-                    state => ((TaskCompletionSource<bool>)state).TrySetCanceled(), tcs);
-#else
-                token.Register(
-                    state => ((TaskCompletionSource<bool>)state).TrySetCanceled(),
-                    state: tcs, useSynchronizationContext: false);
-#endif
-            try
+            if (false == cancellationToken.CanBeCanceled)
             {
-                return await tcs.Task.ConfigureAwait(false);
+                rv = await tcs.Task.ConfigureAwait(false);
             }
-            finally
+            else
             {
-#if NET6_0_OR_GREATER
-                await tokenRegistration.DisposeAsync()
+                rv = await WaitForConfirmsWithTokenAsync(tcs, cancellationToken)
                     .ConfigureAwait(false);
-#else
-                tokenRegistration.Dispose();
-#endif
             }
+
+            return rv;
         }
 
         public async Task WaitForConfirmsOrDieAsync(CancellationToken token = default)
@@ -1830,6 +1851,33 @@ await CloseAsync(ea, false, token)
             }
         }
 
+        private async Task<bool> WaitForConfirmsWithTokenAsync(TaskCompletionSource<bool> tcs,
+            CancellationToken cancellationToken)
+        {
+            CancellationTokenRegistration tokenRegistration =
+#if NET6_0_OR_GREATER
+                cancellationToken.UnsafeRegister(
+                    state => ((TaskCompletionSource<bool>)state).TrySetCanceled(), tcs);
+#else
+                cancellationToken.Register(
+                    state => ((TaskCompletionSource<bool>)state).TrySetCanceled(),
+                    state: tcs, useSynchronizationContext: false);
+#endif
+            try
+            {
+                return await tcs.Task.ConfigureAwait(false);
+            }
+            finally
+            {
+#if NET6_0_OR_GREATER
+                await tokenRegistration.DisposeAsync()
+                    .ConfigureAwait(false);
+#else
+                tokenRegistration.Dispose();
+#endif
+            }
+        }
+
         private static BasicProperties PopulateActivityAndPropagateTraceId<TProperties>(TProperties basicProperties,
             Activity sendActivity) where TProperties : IReadOnlyBasicProperties, IAmqpHeader
         {
diff --git a/projects/RabbitMQ.Client/client/impl/Connection.Commands.cs b/projects/RabbitMQ.Client/client/impl/Connection.Commands.cs
@@ -204,7 +204,7 @@ private Task NotifyCredentialRefreshed(bool succesfully)
             if (succesfully)
             {
                 return UpdateSecretAsync(_config.CredentialsProvider.Password, "Token refresh",
-                    CancellationToken.None); // TODO
+                    CancellationToken.None); // TODO cancellation token
             }
             else
             {
diff --git a/projects/Test/Common/IntegrationFixture.cs b/projects/Test/Common/IntegrationFixture.cs
@@ -144,7 +144,10 @@ public virtual async Task InitializeAsync()
                     _channel = await _conn.CreateChannelAsync();
                 }
 
-                AddCallbackHandlers();
+                if (IsVerbose)
+                {
+                    AddCallbackHandlers();
+                }
             }
 
             if (_connFactory.AutomaticRecoveryEnabled)
@@ -182,43 +185,40 @@ public virtual async Task DisposeAsync()
 
         protected virtual void AddCallbackHandlers()
         {
-            if (IsVerbose)
+            if (_conn != null)
             {
-                if (_conn != null)
+                _conn.CallbackException += (o, ea) =>
                 {
-                    _conn.CallbackException += (o, ea) =>
-                    {
-                        _output.WriteLine("{0} connection callback exception: {1}",
-                            _testDisplayName, ea.Exception);
-                    };
+                    _output.WriteLine("{0} connection callback exception: {1}",
+                        _testDisplayName, ea.Exception);
+                };
 
-                    _conn.ConnectionShutdown += (o, ea) =>
+                _conn.ConnectionShutdown += (o, ea) =>
+                {
+                    HandleConnectionShutdown(_conn, ea, (args) =>
                     {
-                        HandleConnectionShutdown(_conn, ea, (args) =>
-                        {
-                            _output.WriteLine("{0} connection shutdown, args: {1}",
-                                _testDisplayName, args);
-                        });
-                    };
-                }
+                        _output.WriteLine("{0} connection shutdown, args: {1}",
+                            _testDisplayName, args);
+                    });
+                };
+            }
 
-                if (_channel != null)
+            if (_channel != null)
+            {
+                _channel.CallbackException += (o, ea) =>
                 {
-                    _channel.CallbackException += (o, ea) =>
-                    {
-                        _output.WriteLine("{0} channel callback exception: {1}",
-                            _testDisplayName, ea.Exception);
-                    };
+                    _output.WriteLine("{0} channel callback exception: {1}",
+                        _testDisplayName, ea.Exception);
+                };
 
-                    _channel.ChannelShutdown += (o, ea) =>
+                _channel.ChannelShutdown += (o, ea) =>
+                {
+                    HandleChannelShutdown(_channel, ea, (args) =>
                     {
-                        HandleChannelShutdown(_channel, ea, (args) =>
-                        {
-                            _output.WriteLine("{0} channel shutdown, args: {1}",
-                                _testDisplayName, args);
-                        });
-                    };
-                }
+                        _output.WriteLine("{0} channel shutdown, args: {1}",
+                            _testDisplayName, args);
+                    });
+                };
             }
         }
 
diff --git a/projects/Test/Integration/TestAsyncConsumer.cs b/projects/Test/Integration/TestAsyncConsumer.cs
@@ -392,7 +392,7 @@ public async Task TestBasicAckAsync()
                 var c = sender as AsyncEventingBasicConsumer;
                 Assert.NotNull(c);
                 await _channel.BasicAckAsync(args.DeliveryTag, false);
-                messagesReceived++;
+                Interlocked.Increment(ref messagesReceived);
                 if (messagesReceived == messageCount)
                 {
                     publishSyncSource.SetResult(true);
@@ -413,14 +413,12 @@ await _channel.BasicConsumeAsync(queue: queueName, autoAck: false,
                 {
                     byte[] _body = _encoding.GetBytes(Guid.NewGuid().ToString());
                     await _channel.BasicPublishAsync(string.Empty, queueName, _body);
+                    await _channel.WaitForConfirmsOrDieAsync();
                 }
             });
 
-            await _channel.WaitForConfirmsOrDieAsync();
             Assert.True(await publishSyncSource.Task);
-
             Assert.Equal(messageCount, messagesReceived);
-
             await _channel.CloseAsync(_closeArgs, false, CancellationToken.None);
         }
 

Original file line number	Diff line number	Diff line change
`@@ -59,8 +59,7 @@ internal abstract class ChannelBase : IChannel, IRecoverable`
`59`	`59`	`private readonly RpcContinuationQueue _continuationQueue = new RpcContinuationQueue();`
`60`	`60`	`private readonly ManualResetEventSlim _flowControlBlock = new ManualResetEventSlim(true);`
`61`	`61`
`62`		`- // TODO replace with SemaphoreSlim`
`63`		`- private object _confirmLock;`
	`62`	`+ private SemaphoreSlim _confirmSemaphore;`
`64`	`63`	`private readonly LinkedList<ulong> _pendingDeliveryTags = new LinkedList<ulong>();`
`65`	`64`
`66`	`65`	`private bool _onlyAcksReceived = true;`
`@@ -420,7 +419,7 @@ internal void FinishClose()`
`420`	`419`	`m_connectionStartCell?.TrySetResult(null);`
`421`	`420`	`}`
`422`	`421`
`423`		`- private bool ConfirmsAreEnabled => _confirmLock != null;`
	`422`	`+ private bool ConfirmsAreEnabled => _confirmSemaphore != null;`
`424`	`423`
`425`	`424`	`private async Task HandleCommandAsync(IncomingCommand cmd, CancellationToken cancellationToken)`
`426`	`425`	`{`
`@@ -484,7 +483,8 @@ private void OnChannelShutdown(ShutdownEventArgs reason)`
`484`	`483`
`485`	`484`	`if (ConfirmsAreEnabled)`
`486`	`485`	`{`
`487`		`- lock (_confirmLock)`
	`486`	`+ _confirmSemaphore.Wait();`
	`487`	`+ try`
`488`	`488`	`{`
`489`	`489`	`if (_confirmsTaskCompletionSources?.Count > 0)`
`490`	`490`	`{`
`@@ -497,6 +497,10 @@ private void OnChannelShutdown(ShutdownEventArgs reason)`
`497`	`497`	`_confirmsTaskCompletionSources.Clear();`
`498`	`498`	`}`
`499`	`499`	`}`
	`500`	`+ finally`
	`501`	`+ {`
	`502`	`+ _confirmSemaphore.Release();`
	`503`	`+ }`
`500`	`504`	`}`
`501`	`505`
`502`	`506`	`_flowControlBlock.Set();`
`@@ -542,6 +546,7 @@ protected virtual void Dispose(bool disposing)`
`542`	`546`
`543`	`547`	`ConsumerDispatcher.Dispose();`
`544`	`548`	`_rpcSemaphore.Dispose();`
	`549`	`+ _confirmSemaphore?.Dispose();`
`545`	`550`	`}`
`546`	`551`	`}`
`547`	`552`
`@@ -596,7 +601,8 @@ protected void HandleAckNack(ulong deliveryTag, bool multiple, bool isNack)`
`596`	`601`	`if (ConfirmsAreEnabled)`
`597`	`602`	`{`
`598`	`603`	`// let's take a lock so we can assume that deliveryTags are unique, never duplicated and always sorted`
`599`		`- lock (_confirmLock)`
	`604`	`+ _confirmSemaphore.Wait();`
	`605`	`+ try`
`600`	`606`	`{`
`601`	`607`	`// No need to do anything if there are no delivery tags in the list`
`602`	`608`	`if (_pendingDeliveryTags.Count > 0)`
`@@ -633,6 +639,10 @@ protected void HandleAckNack(ulong deliveryTag, bool multiple, bool isNack)`
`633`	`639`	`_onlyAcksReceived = true;`
`634`	`640`	`}`
`635`	`641`	`}`
	`642`	`+ finally`
	`643`	`+ {`
	`644`	`+ _confirmSemaphore.Release();`
	`645`	`+ }`
`636`	`646`	`}`
`637`	`647`	`}`
`638`	`648`
`@@ -1054,10 +1064,16 @@ public async ValueTask BasicPublishAsync<TProperties>(string exchange, string ro`
`1054`	`1064`	`{`
`1055`	`1065`	`if (ConfirmsAreEnabled)`
`1056`	`1066`	`{`
`1057`		`- lock (_confirmLock)`
	`1067`	`+ await _confirmSemaphore.WaitAsync(cancellationToken)`
	`1068`	`+ .ConfigureAwait(false);`
	`1069`	`+ try`
`1058`	`1070`	`{`
`1059`	`1071`	`_pendingDeliveryTags.AddLast(NextPublishSeqNo++);`
`1060`	`1072`	`}`
	`1073`	`+ finally`
	`1074`	`+ {`
	`1075`	`+ _confirmSemaphore.Release();`
	`1076`	`+ }`
`1061`	`1077`	`}`
`1062`	`1078`
`1063`	`1079`	`try`
`@@ -1084,11 +1100,17 @@ await ModelSendAsync(in cmd, in basicProperties, body, cancellationToken)`
`1084`	`1100`	`{`
`1085`	`1101`	`if (ConfirmsAreEnabled)`
`1086`	`1102`	`{`
`1087`		`- lock (_confirmLock)`
	`1103`	`+ await _confirmSemaphore.WaitAsync(cancellationToken)`
	`1104`	`+ .ConfigureAwait(false);`
	`1105`	`+ try`
`1088`	`1106`	`{`
`1089`	`1107`	`NextPublishSeqNo--;`
`1090`	`1108`	`_pendingDeliveryTags.RemoveLast();`
`1091`	`1109`	`}`
	`1110`	`+ finally`
	`1111`	`+ {`
	`1112`	`+ _confirmSemaphore.Release();`
	`1113`	`+ }`
`1092`	`1114`	`}`
`1093`	`1115`
`1094`	`1116`	`throw;`
`@@ -1102,10 +1124,16 @@ public async ValueTask BasicPublishAsync<TProperties>(CachedString exchange, Cac`
`1102`	`1124`	`{`
`1103`	`1125`	`if (ConfirmsAreEnabled)`
`1104`	`1126`	`{`
`1105`		`- lock (_confirmLock)`
	`1127`	`+ await _confirmSemaphore.WaitAsync(cancellationToken)`
	`1128`	`+ .ConfigureAwait(false);`
	`1129`	`+ try`
`1106`	`1130`	`{`
`1107`	`1131`	`_pendingDeliveryTags.AddLast(NextPublishSeqNo++);`
`1108`	`1132`	`}`
	`1133`	`+ finally`
	`1134`	`+ {`
	`1135`	`+ _confirmSemaphore.Release();`
	`1136`	`+ }`
`1109`	`1137`	`}`
`1110`	`1138`
`1111`	`1139`	`try`
`@@ -1133,11 +1161,17 @@ await ModelSendAsync(in cmd, in basicProperties, body, cancellationToken)`
`1133`	`1161`	`{`
`1134`	`1162`	`if (ConfirmsAreEnabled)`
`1135`	`1163`	`{`
`1136`		`- lock (_confirmLock)`
	`1164`	`+ await _confirmSemaphore.WaitAsync(cancellationToken)`
	`1165`	`+ .ConfigureAwait(false);`
	`1166`	`+ try`
`1137`	`1167`	`{`
`1138`	`1168`	`NextPublishSeqNo--;`
`1139`	`1169`	`_pendingDeliveryTags.RemoveLast();`
`1140`	`1170`	`}`
	`1171`	`+ finally`
	`1172`	`+ {`
	`1173`	`+ _confirmSemaphore.Release();`
	`1174`	`+ }`
`1141`	`1175`	`}`
`1142`	`1176`
`1143`	`1177`	`throw;`
`@@ -1242,7 +1276,7 @@ await ModelSendAsync(method, k.CancellationToken)`
`1242`	`1276`
`1243`	`1277`	`// Note:`
`1244`	`1278`	`// Non-null means confirms are enabled`
`1245`		`- _confirmLock = new object();`
	`1279`	`+ _confirmSemaphore = new SemaphoreSlim(1, 1);`
`1246`	`1280`
`1247`	`1281`	`return;`
`1248`	`1282`	`}`
`@@ -1742,63 +1776,50 @@ await ModelSendAsync(method, k.CancellationToken)`
`1742`	`1776`
`1743`	`1777`	`private List<TaskCompletionSource<bool>> _confirmsTaskCompletionSources;`
`1744`	`1778`
`1745`		`- public Task<bool> WaitForConfirmsAsync(CancellationToken token = default)`
	`1779`	`+ public async Task<bool> WaitForConfirmsAsync(CancellationToken cancellationToken = default)`
`1746`	`1780`	`{`
`1747`	`1781`	`if (false == ConfirmsAreEnabled)`
`1748`	`1782`	`{`
`1749`	`1783`	`throw new InvalidOperationException("Confirms not selected");`
`1750`	`1784`	`}`
`1751`	`1785`
`1752`	`1786`	`TaskCompletionSource<bool> tcs;`
`1753`		`- lock (_confirmLock)`
	`1787`	`+ await _confirmSemaphore.WaitAsync(cancellationToken)`
	`1788`	`+ .ConfigureAwait(false);`
	`1789`	`+ try`
`1754`	`1790`	`{`
`1755`	`1791`	`if (_pendingDeliveryTags.Count == 0)`
`1756`	`1792`	`{`
`1757`	`1793`	`if (_onlyAcksReceived == false)`
`1758`	`1794`	`{`
`1759`	`1795`	`_onlyAcksReceived = true;`
`1760`		`- return Task.FromResult(false);`
	`1796`	`+ return false;`
`1761`	`1797`	`}`
`1762`	`1798`
`1763`		`- return Task.FromResult(true);`
	`1799`	`+ return true;`
`1764`	`1800`	`}`
`1765`	`1801`
`1766`	`1802`	`tcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);`
`1767`	`1803`	`_confirmsTaskCompletionSources.Add(tcs);`
`1768`	`1804`	`}`
`1769`		`-`
`1770`		`- if (!token.CanBeCanceled)`
	`1805`	`+ finally`
`1771`	`1806`	`{`
`1772`		`- return tcs.Task;`
	`1807`	`+ _confirmSemaphore.Release();`
`1773`	`1808`	`}`
`1774`	`1809`
`1775`		`- return WaitForConfirmsWithTokenAsync(tcs, token);`
`1776`		`- }`
	`1810`	`+ bool rv;`
`1777`	`1811`
`1778`		`- private async Task<bool> WaitForConfirmsWithTokenAsync(TaskCompletionSource<bool> tcs, CancellationToken token)`
`1779`		`- {`
`1780`		`- CancellationTokenRegistration tokenRegistration =`
`1781`		`-#if NET6_0_OR_GREATER`
`1782`		`- token.UnsafeRegister(`
`1783`		`- state => ((TaskCompletionSource<bool>)state).TrySetCanceled(), tcs);`
`1784`		`-#else`
`1785`		`- token.Register(`
`1786`		`- state => ((TaskCompletionSource<bool>)state).TrySetCanceled(),`
`1787`		`- state: tcs, useSynchronizationContext: false);`
`1788`		`-#endif`
`1789`		`- try`
	`1812`	`+ if (false == cancellationToken.CanBeCanceled)`
`1790`	`1813`	`{`
`1791`		`- return await tcs.Task.ConfigureAwait(false);`
	`1814`	`+ rv = await tcs.Task.ConfigureAwait(false);`
`1792`	`1815`	`}`
`1793`		`- finally`
	`1816`	`+ else`
`1794`	`1817`	`{`
`1795`		`-#if NET6_0_OR_GREATER`
`1796`		`- await tokenRegistration.DisposeAsync()`
	`1818`	`+ rv = await WaitForConfirmsWithTokenAsync(tcs, cancellationToken)`
`1797`	`1819`	`.ConfigureAwait(false);`
`1798`		`-#else`
`1799`		`- tokenRegistration.Dispose();`
`1800`		`-#endif`
`1801`	`1820`	`}`
	`1821`	`+`
	`1822`	`+ return rv;`
`1802`	`1823`	`}`
`1803`	`1824`
`1804`	`1825`	`public async Task WaitForConfirmsOrDieAsync(CancellationToken token = default)`
`@@ -1830,6 +1851,33 @@ await CloseAsync(ea, false, token)`
`1830`	`1851`	`}`
`1831`	`1852`	`}`
`1832`	`1853`
	`1854`	`+ private async Task<bool> WaitForConfirmsWithTokenAsync(TaskCompletionSource<bool> tcs,`
	`1855`	`+ CancellationToken cancellationToken)`
	`1856`	`+ {`
	`1857`	`+ CancellationTokenRegistration tokenRegistration =`
	`1858`	`+#if NET6_0_OR_GREATER`
	`1859`	`+ cancellationToken.UnsafeRegister(`
	`1860`	`+ state => ((TaskCompletionSource<bool>)state).TrySetCanceled(), tcs);`
	`1861`	`+#else`
	`1862`	`+ cancellationToken.Register(`
	`1863`	`+ state => ((TaskCompletionSource<bool>)state).TrySetCanceled(),`
	`1864`	`+ state: tcs, useSynchronizationContext: false);`
	`1865`	`+#endif`
	`1866`	`+ try`
	`1867`	`+ {`
	`1868`	`+ return await tcs.Task.ConfigureAwait(false);`
	`1869`	`+ }`
	`1870`	`+ finally`
	`1871`	`+ {`
	`1872`	`+#if NET6_0_OR_GREATER`
	`1873`	`+ await tokenRegistration.DisposeAsync()`
	`1874`	`+ .ConfigureAwait(false);`
	`1875`	`+#else`
	`1876`	`+ tokenRegistration.Dispose();`
	`1877`	`+#endif`
	`1878`	`+ }`
	`1879`	`+ }`
	`1880`	`+`
`1833`	`1881`	`private static BasicProperties PopulateActivityAndPropagateTraceId<TProperties>(TProperties basicProperties,`
`1834`	`1882`	`Activity sendActivity) where TProperties : IReadOnlyBasicProperties, IAmqpHeader`
`1835`	`1883`	`{`
Original file line number	Diff line number	Diff line change
`@@ -204,7 +204,7 @@ private Task NotifyCredentialRefreshed(bool succesfully)`
`204`	`204`	`if (succesfully)`
`205`	`205`	`{`
`206`	`206`	`return UpdateSecretAsync(_config.CredentialsProvider.Password, "Token refresh",`
`207`		`- CancellationToken.None); // TODO`
	`207`	`+ CancellationToken.None); // TODO cancellation token`
`208`	`208`	`}`
`209`	`209`	`else`
`210`	`210`	`{`
Original file line number	Diff line number	Diff line change
`@@ -392,7 +392,7 @@ public async Task TestBasicAckAsync()`
`392`	`392`	`var c = sender as AsyncEventingBasicConsumer;`
`393`	`393`	`Assert.NotNull(c);`
`394`	`394`	`await _channel.BasicAckAsync(args.DeliveryTag, false);`
`395`		`- messagesReceived++;`
	`395`	`+ Interlocked.Increment(ref messagesReceived);`
`396`	`396`	`if (messagesReceived == messageCount)`
`397`	`397`	`{`
`398`	`398`	`publishSyncSource.SetResult(true);`
`@@ -413,14 +413,12 @@ await _channel.BasicConsumeAsync(queue: queueName, autoAck: false,`
`413`	`413`	`{`
`414`	`414`	`byte[] _body = _encoding.GetBytes(Guid.NewGuid().ToString());`
`415`	`415`	`await _channel.BasicPublishAsync(string.Empty, queueName, _body);`
	`416`	`+ await _channel.WaitForConfirmsOrDieAsync();`
`416`	`417`	`}`
`417`	`418`	`});`
`418`	`419`
`419`		`- await _channel.WaitForConfirmsOrDieAsync();`
`420`	`420`	`Assert.True(await publishSyncSource.Task);`
`421`		`-`
`422`	`421`	`Assert.Equal(messageCount, messagesReceived);`
`423`		`-`
`424`	`422`	`await _channel.CloseAsync(_closeArgs, false, CancellationToken.None);`
`425`	`423`	`}`
`426`	`424`