Ensure broker-originated channel closure completes

lukebakken · lukebakken · commit 9bbcb5e2b889 · 2025-01-09T09:42:35.000-08:00
Fixes #1749 * Ensure `Dispose` and `DisposeAsync` are idempotent and thread-safe. * Use TaskCompletionSource when `HandleChannelCloseAsync` runs to allow dispose methods to wait. * Use `Interlocked` for thread safety. * I like `_isDisposing` better. So sue me! * Move the `Interlocked.Exchange` code to a getter, for readability. * Minor nullable change.
diff --git a/projects/RabbitMQ.Client/Impl/AutorecoveringChannel.cs b/projects/RabbitMQ.Client/Impl/AutorecoveringChannel.cs
@@ -31,6 +31,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
 using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
@@ -48,6 +49,7 @@ internal sealed class AutorecoveringChannel : IChannel, IRecoverable
         private AutorecoveringConnection _connection;
         private RecoveryAwareChannel _innerChannel;
         private bool _disposed;
+        private int _isDisposing;
 
         private ushort _prefetchCountConsumer;
         private ushort _prefetchCountGlobal;
@@ -256,19 +258,25 @@ public override string ToString()
 
         public async ValueTask DisposeAsync()
         {
-            if (_disposed)
+            if (IsDisposing)
             {
                 return;
             }
 
-            if (IsOpen)
+            try
             {
-                await this.AbortAsync()
-                    .ConfigureAwait(false);
-            }
+                if (IsOpen)
+                {
+                    await this.AbortAsync()
+                        .ConfigureAwait(false);
+                }
 
-            _recordedConsumerTags.Clear();
-            _disposed = true;
+                _recordedConsumerTags.Clear();
+            }
+            finally
+            {
+                _disposed = true;
+            }
         }
 
         public ValueTask<ulong> GetNextPublishSequenceNumberAsync(CancellationToken cancellationToken = default) => InnerChannel.GetNextPublishSequenceNumberAsync(cancellationToken);
@@ -482,7 +490,23 @@ private void ThrowIfDisposed()
                 ThrowDisposed();
             }
 
+            return;
+
+            [DoesNotReturn]
             static void ThrowDisposed() => throw new ObjectDisposedException(typeof(AutorecoveringChannel).FullName);
         }
+
+        private bool IsDisposing
+        {
+            get
+            {
+                if (Interlocked.Exchange(ref _isDisposing, 1) != 0)
+                {
+                    return true;
+                }
+
+                return false;
+            }
+        }
     }
 }
diff --git a/projects/RabbitMQ.Client/Impl/AutorecoveringConnection.cs b/projects/RabbitMQ.Client/Impl/AutorecoveringConnection.cs
@@ -31,6 +31,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
 using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
@@ -50,6 +51,7 @@ internal sealed partial class AutorecoveringConnection : IConnection
 
         private Connection _innerConnection;
         private bool _disposed;
+        private int _isDisposing;
 
         private Connection InnerConnection
         {
@@ -272,7 +274,7 @@ await RecordChannelAsync(autorecoveringChannel, channelsSemaphoreHeld: false, ca
 
         public async ValueTask DisposeAsync()
         {
-            if (_disposed)
+            if (IsDisposing)
             {
                 return;
             }
@@ -281,17 +283,18 @@ public async ValueTask DisposeAsync()
             {
                 await _innerConnection.DisposeAsync()
                     .ConfigureAwait(false);
+
+                _channels.Clear();
+                _recordedEntitiesSemaphore.Dispose();
+                _channelsSemaphore.Dispose();
+                _recoveryCancellationTokenSource.Dispose();
             }
             catch (OperationInterruptedException)
             {
                 // ignored, see rabbitmq/rabbitmq-dotnet-client#133
             }
             finally
             {
-                _channels.Clear();
-                _recordedEntitiesSemaphore.Dispose();
-                _channelsSemaphore.Dispose();
-                _recoveryCancellationTokenSource.Dispose();
                 _disposed = true;
             }
         }
@@ -307,7 +310,23 @@ private void ThrowIfDisposed()
                 ThrowDisposed();
             }
 
+            return;
+
+            [DoesNotReturn]
             static void ThrowDisposed() => throw new ObjectDisposedException(typeof(AutorecoveringConnection).FullName);
         }
+
+        private bool IsDisposing
+        {
+            get
+            {
+                if (Interlocked.Exchange(ref _isDisposing, 1) != 0)
+                {
+                    return true;
+                }
+
+                return false;
+            }
+        }
     }
 }
diff --git a/projects/RabbitMQ.Client/Impl/Channel.cs b/projects/RabbitMQ.Client/Impl/Channel.cs
@@ -59,8 +59,13 @@ internal partial class Channel : IChannel, IRecoverable
         private ShutdownEventArgs? _closeReason;
         public ShutdownEventArgs? CloseReason => Volatile.Read(ref _closeReason);
 
+        private TaskCompletionSource<bool>? _serverOriginatedChannelCloseTcs;
+
         internal readonly IConsumerDispatcher ConsumerDispatcher;
 
+        private bool _disposed;
+        private int _isDisposing;
+
         public Channel(ISession session, CreateChannelOptions createChannelOptions)
         {
             ContinuationTimeout = createChannelOptions.ContinuationTimeout;
@@ -514,22 +519,41 @@ public override string ToString()
 
         void IDisposable.Dispose()
         {
+            if (_disposed)
+            {
+                return;
+            }
+
             Dispose(true);
         }
 
         protected virtual void Dispose(bool disposing)
         {
+            if (IsDisposing)
+            {
+                return;
+            }
+
             if (disposing)
             {
-                if (IsOpen)
+                try
                 {
-                    this.AbortAsync().GetAwaiter().GetResult();
-                }
+                    if (IsOpen)
+                    {
+                        this.AbortAsync().GetAwaiter().GetResult();
+                    }
 
-                ConsumerDispatcher.Dispose();
-                _rpcSemaphore.Dispose();
-                _confirmSemaphore.Dispose();
-                _outstandingPublisherConfirmationsRateLimiter?.Dispose();
+                    _serverOriginatedChannelCloseTcs?.Task.Wait(TimeSpan.FromSeconds(5));
+
+                    ConsumerDispatcher.Dispose();
+                    _rpcSemaphore.Dispose();
+                    _confirmSemaphore.Dispose();
+                    _outstandingPublisherConfirmationsRateLimiter?.Dispose();
+                }
+                finally
+                {
+                    _disposed = true;
+                }
             }
         }
 
@@ -543,18 +567,37 @@ await DisposeAsyncCore()
 
         protected virtual async ValueTask DisposeAsyncCore()
         {
-            if (IsOpen)
+            if (IsDisposing)
             {
-                await this.AbortAsync().ConfigureAwait(false);
+                return;
             }
 
-            ConsumerDispatcher.Dispose();
-            _rpcSemaphore.Dispose();
-            _confirmSemaphore.Dispose();
-            if (_outstandingPublisherConfirmationsRateLimiter is not null)
+            try
             {
-                await _outstandingPublisherConfirmationsRateLimiter.DisposeAsync()
-                    .ConfigureAwait(false);
+                if (IsOpen)
+                {
+                    await this.AbortAsync().ConfigureAwait(false);
+                }
+
+                if (_serverOriginatedChannelCloseTcs is not null)
+                {
+                    await _serverOriginatedChannelCloseTcs.Task.WaitAsync(TimeSpan.FromSeconds(5))
+                        .ConfigureAwait(false);
+                }
+
+                ConsumerDispatcher.Dispose();
+                _rpcSemaphore.Dispose();
+                _confirmSemaphore.Dispose();
+
+                if (_outstandingPublisherConfirmationsRateLimiter is not null)
+                {
+                    await _outstandingPublisherConfirmationsRateLimiter.DisposeAsync()
+                        .ConfigureAwait(false);
+                }
+            }
+            finally
+            {
+                _disposed = true;
             }
         }
 
@@ -651,23 +694,41 @@ protected virtual ulong AdjustDeliveryTag(ulong deliveryTag)
 
         protected async Task<bool> HandleChannelCloseAsync(IncomingCommand cmd, CancellationToken cancellationToken)
         {
-            var channelClose = new ChannelClose(cmd.MethodSpan);
-            SetCloseReason(new ShutdownEventArgs(ShutdownInitiator.Peer,
-                channelClose._replyCode,
-                channelClose._replyText,
-                channelClose._classId,
-                channelClose._methodId));
+            TaskCompletionSource<bool>? serverOriginatedChannelCloseTcs = _serverOriginatedChannelCloseTcs;
+            if (serverOriginatedChannelCloseTcs is null)
+            {
+                // Attempt to assign the new TCS only if _tcs is still null
+                _ = Interlocked.CompareExchange(ref _serverOriginatedChannelCloseTcs,
+                    new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously), null);
+            }
 
-            await Session.CloseAsync(_closeReason, notify: false)
-                .ConfigureAwait(false);
+            try
+            {
+                var channelClose = new ChannelClose(cmd.MethodSpan);
+                SetCloseReason(new ShutdownEventArgs(ShutdownInitiator.Peer,
+                    channelClose._replyCode,
+                    channelClose._replyText,
+                    channelClose._classId,
+                    channelClose._methodId));
 
-            var method = new ChannelCloseOk();
-            await ModelSendAsync(in method, cancellationToken)
-                .ConfigureAwait(false);
+                await Session.CloseAsync(_closeReason, notify: false)
+                    .ConfigureAwait(false);
 
-            await Session.NotifyAsync(cancellationToken)
-                .ConfigureAwait(false);
-            return true;
+                var method = new ChannelCloseOk();
+                await ModelSendAsync(in method, cancellationToken)
+                    .ConfigureAwait(false);
+
+                await Session.NotifyAsync(cancellationToken)
+                    .ConfigureAwait(false);
+
+                _serverOriginatedChannelCloseTcs?.TrySetResult(true);
+                return true;
+            }
+            catch (Exception ex)
+            {
+                _serverOriginatedChannelCloseTcs?.TrySetException(ex);
+                throw;
+            }
         }
 
         protected async Task<bool> HandleChannelCloseOkAsync(IncomingCommand cmd, CancellationToken cancellationToken)
@@ -1587,5 +1648,18 @@ private Task<bool> DispatchCommandAsync(IncomingCommand cmd, CancellationToken c
                     }
             }
         }
+
+        private bool IsDisposing
+        {
+            get
+            {
+                if (Interlocked.Exchange(ref _isDisposing, 1) != 0)
+                {
+                    return true;
+                }
+
+                return false;
+            }
+        }
     }
 }
diff --git a/projects/RabbitMQ.Client/Impl/Connection.cs b/projects/RabbitMQ.Client/Impl/Connection.cs
@@ -32,6 +32,7 @@
 using System;
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.Diagnostics.CodeAnalysis;
 using System.IO;
 using System.Runtime.CompilerServices;
 using System.Threading;
@@ -46,6 +47,7 @@ namespace RabbitMQ.Client.Framing
     internal sealed partial class Connection : IConnection
     {
         private bool _disposed;
+        private int _isDisposing;
         private volatile bool _closed;
 
         private readonly ConnectionConfig _config;
@@ -489,7 +491,7 @@ internal ValueTask WriteAsync(RentedMemory frames, CancellationToken cancellatio
 
         public async ValueTask DisposeAsync()
         {
-            if (_disposed)
+            if (IsDisposing)
             {
                 return;
             }
@@ -523,23 +525,37 @@ private void ThrowIfDisposed()
         {
             if (_disposed)
             {
-                ThrowObjectDisposedException();
+                ThrowDisposed();
             }
 
-            static void ThrowObjectDisposedException()
-            {
-                throw new ObjectDisposedException(typeof(Connection).FullName);
-            }
+            return;
+
+            [DoesNotReturn]
+            static void ThrowDisposed() => throw new ObjectDisposedException(typeof(Connection).FullName);
         }
 
         public override string ToString()
         {
             return $"Connection({_id},{Endpoint})";
         }
 
+        [DoesNotReturn]
         private static void ThrowAlreadyClosedException(ShutdownEventArgs closeReason)
         {
             throw new AlreadyClosedException(closeReason);
         }
+
+        private bool IsDisposing
+        {
+            get
+            {
+                if (Interlocked.Exchange(ref _isDisposing, 1) != 0)
+                {
+                    return true;
+                }
+
+                return false;
+            }
+        }
     }
 }
diff --git a/projects/Test/Integration/TestChannelShutdown.cs b/projects/Test/Integration/TestChannelShutdown.cs
diff --git a/projects/Test/Integration/TestQueueDeclare.cs b/projects/Test/Integration/TestQueueDeclare.cs

Original file line number	Diff line number	Diff line change
`@@ -31,6 +31,7 @@`
`31`	`31`
`32`	`32`	`using System;`
`33`	`33`	`using System.Collections.Generic;`
	`34`	`+using System.Diagnostics.CodeAnalysis;`
`34`	`35`	`using System.Runtime.CompilerServices;`
`35`	`36`	`using System.Threading;`
`36`	`37`	`using System.Threading.Tasks;`
`@@ -50,6 +51,7 @@ internal sealed partial class AutorecoveringConnection : IConnection`
`50`	`51`
`51`	`52`	`private Connection _innerConnection;`
`52`	`53`	`private bool _disposed;`
	`54`	`+ private int _isDisposing;`
`53`	`55`
`54`	`56`	`private Connection InnerConnection`
`55`	`57`	`{`
`@@ -272,7 +274,7 @@ await RecordChannelAsync(autorecoveringChannel, channelsSemaphoreHeld: false, ca`
`272`	`274`
`273`	`275`	`public async ValueTask DisposeAsync()`
`274`	`276`	`{`
`275`		`- if (_disposed)`
	`277`	`+ if (IsDisposing)`
`276`	`278`	`{`
`277`	`279`	`return;`
`278`	`280`	`}`
`@@ -281,17 +283,18 @@ public async ValueTask DisposeAsync()`
`281`	`283`	`{`
`282`	`284`	`await _innerConnection.DisposeAsync()`
`283`	`285`	`.ConfigureAwait(false);`
	`286`	`+`
	`287`	`+ _channels.Clear();`
	`288`	`+ _recordedEntitiesSemaphore.Dispose();`
	`289`	`+ _channelsSemaphore.Dispose();`
	`290`	`+ _recoveryCancellationTokenSource.Dispose();`
`284`	`291`	`}`
`285`	`292`	`catch (OperationInterruptedException)`
`286`	`293`	`{`
`287`	`294`	`// ignored, see rabbitmq/rabbitmq-dotnet-client#133`
`288`	`295`	`}`
`289`	`296`	`finally`
`290`	`297`	`{`
`291`		`- _channels.Clear();`
`292`		`- _recordedEntitiesSemaphore.Dispose();`
`293`		`- _channelsSemaphore.Dispose();`
`294`		`- _recoveryCancellationTokenSource.Dispose();`
`295`	`298`	`_disposed = true;`
`296`	`299`	`}`
`297`	`300`	`}`
`@@ -307,7 +310,23 @@ private void ThrowIfDisposed()`
`307`	`310`	`ThrowDisposed();`
`308`	`311`	`}`
`309`	`312`
	`313`	`+ return;`
	`314`	`+`
	`315`	`+ [DoesNotReturn]`
`310`	`316`	`static void ThrowDisposed() => throw new ObjectDisposedException(typeof(AutorecoveringConnection).FullName);`
`311`	`317`	`}`
	`318`	`+`
	`319`	`+ private bool IsDisposing`
	`320`	`+ {`
	`321`	`+ get`
	`322`	`+ {`
	`323`	`+ if (Interlocked.Exchange(ref _isDisposing, 1) != 0)`
	`324`	`+ {`
	`325`	`+ return true;`
	`326`	`+ }`
	`327`	`+`
	`328`	`+ return false;`
	`329`	`+ }`
	`330`	`+ }`
`312`	`331`	`}`
`313`	`332`	`}`
Original file line number	Diff line number	Diff line change
`@@ -32,6 +32,7 @@`
`32`	`32`	`using System;`
`33`	`33`	`using System.Collections.Generic;`
`34`	`34`	`using System.Diagnostics;`
	`35`	`+using System.Diagnostics.CodeAnalysis;`
`35`	`36`	`using System.IO;`
`36`	`37`	`using System.Runtime.CompilerServices;`
`37`	`38`	`using System.Threading;`
`@@ -46,6 +47,7 @@ namespace RabbitMQ.Client.Framing`
`46`	`47`	`internal sealed partial class Connection : IConnection`
`47`	`48`	`{`
`48`	`49`	`private bool _disposed;`
	`50`	`+ private int _isDisposing;`
`49`	`51`	`private volatile bool _closed;`
`50`	`52`
`51`	`53`	`private readonly ConnectionConfig _config;`
`@@ -489,7 +491,7 @@ internal ValueTask WriteAsync(RentedMemory frames, CancellationToken cancellatio`
`489`	`491`
`490`	`492`	`public async ValueTask DisposeAsync()`
`491`	`493`	`{`
`492`		`- if (_disposed)`
	`494`	`+ if (IsDisposing)`
`493`	`495`	`{`
`494`	`496`	`return;`
`495`	`497`	`}`
`@@ -523,23 +525,37 @@ private void ThrowIfDisposed()`
`523`	`525`	`{`
`524`	`526`	`if (_disposed)`
`525`	`527`	`{`
`526`		`- ThrowObjectDisposedException();`
	`528`	`+ ThrowDisposed();`
`527`	`529`	`}`
`528`	`530`
`529`		`- static void ThrowObjectDisposedException()`
`530`		`- {`
`531`		`- throw new ObjectDisposedException(typeof(Connection).FullName);`
`532`		`- }`
	`531`	`+ return;`
	`532`	`+`
	`533`	`+ [DoesNotReturn]`
	`534`	`+ static void ThrowDisposed() => throw new ObjectDisposedException(typeof(Connection).FullName);`
`533`	`535`	`}`
`534`	`536`
`535`	`537`	`public override string ToString()`
`536`	`538`	`{`
`537`	`539`	`return $"Connection({_id},{Endpoint})";`
`538`	`540`	`}`
`539`	`541`
	`542`	`+ [DoesNotReturn]`
`540`	`543`	`private static void ThrowAlreadyClosedException(ShutdownEventArgs closeReason)`
`541`	`544`	`{`
`542`	`545`	`throw new AlreadyClosedException(closeReason);`
`543`	`546`	`}`
	`547`	`+`
	`548`	`+ private bool IsDisposing`
	`549`	`+ {`
	`550`	`+ get`
	`551`	`+ {`
	`552`	`+ if (Interlocked.Exchange(ref _isDisposing, 1) != 0)`
	`553`	`+ {`
	`554`	`+ return true;`
	`555`	`+ }`
	`556`	`+`
	`557`	`+ return false;`
	`558`	`+ }`
	`559`	`+ }`
`544`	`560`	`}`
`545`	`561`	`}`