Ensure broker-originated channel closure completes

lukebakken · lukebakken · commit c4200ab969cb · 2025-01-06T18:31:41.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.
diff --git a/projects/RabbitMQ.Client/Impl/AutorecoveringChannel.cs b/projects/RabbitMQ.Client/Impl/AutorecoveringChannel.cs
@@ -48,6 +48,8 @@ internal sealed class AutorecoveringChannel : IChannel, IRecoverable
         private AutorecoveringConnection _connection;
         private RecoveryAwareChannel _innerChannel;
         private bool _disposed;
+        private bool _isDisposing;
+        private readonly object _locker = new();
 
         private ushort _prefetchCountConsumer;
         private ushort _prefetchCountGlobal;
@@ -252,7 +254,15 @@ await _connection.DeleteRecordedChannelAsync(this,
         public override string ToString()
             => InnerChannel.ToString();
 
-        public void Dispose() => DisposeAsync().AsTask().GetAwaiter().GetResult();
+        public void Dispose()
+        {
+            if (_disposed)
+            {
+                return;
+            }
+
+            DisposeAsync().AsTask().GetAwaiter().GetResult();
+        }
 
         public async ValueTask DisposeAsync()
         {
@@ -261,14 +271,30 @@ public async ValueTask DisposeAsync()
                 return;
             }
 
-            if (IsOpen)
+            lock (_locker)
             {
-                await this.AbortAsync()
-                    .ConfigureAwait(false);
+                if (_isDisposing)
+                {
+                    return;
+                }
+                _isDisposing = true;
             }
 
-            _recordedConsumerTags.Clear();
-            _disposed = true;
+            try
+            {
+                if (IsOpen)
+                {
+                    await this.AbortAsync()
+                        .ConfigureAwait(false);
+                }
+
+                _recordedConsumerTags.Clear();
+            }
+            finally
+            {
+                _disposed = true;
+                _isDisposing = false;
+            }
         }
 
         public ValueTask<ulong> GetNextPublishSequenceNumberAsync(CancellationToken cancellationToken = default) => InnerChannel.GetNextPublishSequenceNumberAsync(cancellationToken);
diff --git a/projects/RabbitMQ.Client/Impl/AutorecoveringConnection.cs b/projects/RabbitMQ.Client/Impl/AutorecoveringConnection.cs
@@ -50,6 +50,8 @@ internal sealed partial class AutorecoveringConnection : IConnection
 
         private Connection _innerConnection;
         private bool _disposed;
+        private bool _isDisposing;
+        private readonly object _locker = new();
 
         private Connection InnerConnection
         {
@@ -268,7 +270,15 @@ await RecordChannelAsync(autorecoveringChannel, channelsSemaphoreHeld: false, ca
             return autorecoveringChannel;
         }
 
-        public void Dispose() => DisposeAsync().AsTask().GetAwaiter().GetResult();
+        public void Dispose()
+        {
+            if (_disposed)
+            {
+                return;
+            }
+
+            DisposeAsync().AsTask().GetAwaiter().GetResult();
+        }
 
         public async ValueTask DisposeAsync()
         {
@@ -277,22 +287,33 @@ public async ValueTask DisposeAsync()
                 return;
             }
 
+            lock (_locker)
+            {
+                if (_isDisposing)
+                {
+                    return;
+                }
+                _isDisposing = true;
+            }
+
             try
             {
                 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;
+                _isDisposing = false;
             }
         }
 
diff --git a/projects/RabbitMQ.Client/Impl/Channel.cs b/projects/RabbitMQ.Client/Impl/Channel.cs
@@ -59,8 +59,15 @@ 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 bool _isDisposing;
+
+        private readonly object _locker = new();
+
         public Channel(ISession session, CreateChannelOptions createChannelOptions)
         {
             ContinuationTimeout = createChannelOptions.ContinuationTimeout;
@@ -514,22 +521,54 @@ public override string ToString()
 
         void IDisposable.Dispose()
         {
+            if (_disposed)
+            {
+                return;
+            }
+
             Dispose(true);
         }
 
         protected virtual void Dispose(bool disposing)
         {
-            if (disposing)
+            if (_disposed)
             {
-                if (IsOpen)
+                return;
+            }
+
+            lock (_locker)
+            {
+                if (_isDisposing)
                 {
-                    this.AbortAsync().GetAwaiter().GetResult();
+                    return;
                 }
+                _isDisposing = true;
+            }
 
-                ConsumerDispatcher.Dispose();
-                _rpcSemaphore.Dispose();
-                _confirmSemaphore.Dispose();
-                _outstandingPublisherConfirmationsRateLimiter?.Dispose();
+            if (disposing)
+            {
+                try
+                {
+                    if (IsOpen)
+                    {
+                        this.AbortAsync().GetAwaiter().GetResult();
+                    }
+
+                    if (_serverOriginatedChannelCloseTcs is not null)
+                    {
+                        _serverOriginatedChannelCloseTcs.Task.Wait(TimeSpan.FromSeconds(5));
+                    }
+
+                    ConsumerDispatcher.Dispose();
+                    _rpcSemaphore.Dispose();
+                    _confirmSemaphore.Dispose();
+                    _outstandingPublisherConfirmationsRateLimiter?.Dispose();
+                }
+                finally
+                {
+                    _disposed = true;
+                    _isDisposing = false;
+                }
             }
         }
 
@@ -543,18 +582,46 @@ await DisposeAsyncCore()
 
         protected virtual async ValueTask DisposeAsyncCore()
         {
-            if (IsOpen)
+            if (_disposed)
             {
-                await this.AbortAsync().ConfigureAwait(false);
+                return;
             }
 
-            ConsumerDispatcher.Dispose();
-            _rpcSemaphore.Dispose();
-            _confirmSemaphore.Dispose();
-            if (_outstandingPublisherConfirmationsRateLimiter is not null)
+            lock (_locker)
             {
-                await _outstandingPublisherConfirmationsRateLimiter.DisposeAsync()
-                    .ConfigureAwait(false);
+                if (_isDisposing)
+                {
+                    return;
+                }
+                _isDisposing = true;
+            }
+
+            try
+            {
+                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;
+                _isDisposing = false;
             }
         }
 
@@ -651,23 +718,38 @@ 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));
+            lock (_locker)
+            {
+                _serverOriginatedChannelCloseTcs ??= new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+            }
 
-            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)
diff --git a/projects/RabbitMQ.Client/Impl/Connection.cs b/projects/RabbitMQ.Client/Impl/Connection.cs
@@ -46,6 +46,8 @@ namespace RabbitMQ.Client.Framing
     internal sealed partial class Connection : IConnection
     {
         private bool _disposed;
+        private bool _isDisposing;
+        private readonly object _locker = new();
         private volatile bool _closed;
 
         private readonly ConnectionConfig _config;
@@ -485,7 +487,15 @@ internal ValueTask WriteAsync(RentedMemory frames, CancellationToken cancellatio
             return _frameHandler.WriteAsync(frames, cancellationToken);
         }
 
-        public void Dispose() => DisposeAsync().AsTask().GetAwaiter().GetResult();
+        public void Dispose()
+        {
+            if (_disposed)
+            {
+                return;
+            }
+
+            DisposeAsync().AsTask().GetAwaiter().GetResult();
+        }
 
         public async ValueTask DisposeAsync()
         {
@@ -494,6 +504,15 @@ public async ValueTask DisposeAsync()
                 return;
             }
 
+            lock (_locker)
+            {
+                if (_isDisposing)
+                {
+                    return;
+                }
+                _isDisposing = true;
+            }
+
             try
             {
                 if (IsOpen)
@@ -515,6 +534,7 @@ await _channel0.DisposeAsync()
             finally
             {
                 _disposed = true;
+                _isDisposing = 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
`@@ -50,6 +50,8 @@ internal sealed partial class AutorecoveringConnection : IConnection`
`50`	`50`
`51`	`51`	`private Connection _innerConnection;`
`52`	`52`	`private bool _disposed;`
	`53`	`+ private bool _isDisposing;`
	`54`	`+ private readonly object _locker = new();`
`53`	`55`
`54`	`56`	`private Connection InnerConnection`
`55`	`57`	`{`
`@@ -268,7 +270,15 @@ await RecordChannelAsync(autorecoveringChannel, channelsSemaphoreHeld: false, ca`
`268`	`270`	`return autorecoveringChannel;`
`269`	`271`	`}`
`270`	`272`
`271`		`- public void Dispose() => DisposeAsync().AsTask().GetAwaiter().GetResult();`
	`273`	`+ public void Dispose()`
	`274`	`+ {`
	`275`	`+ if (_disposed)`
	`276`	`+ {`
	`277`	`+ return;`
	`278`	`+ }`
	`279`	`+`
	`280`	`+ DisposeAsync().AsTask().GetAwaiter().GetResult();`
	`281`	`+ }`
`272`	`282`
`273`	`283`	`public async ValueTask DisposeAsync()`
`274`	`284`	`{`
`@@ -277,22 +287,33 @@ public async ValueTask DisposeAsync()`
`277`	`287`	`return;`
`278`	`288`	`}`
`279`	`289`
	`290`	`+ lock (_locker)`
	`291`	`+ {`
	`292`	`+ if (_isDisposing)`
	`293`	`+ {`
	`294`	`+ return;`
	`295`	`+ }`
	`296`	`+ _isDisposing = true;`
	`297`	`+ }`
	`298`	`+`
`280`	`299`	`try`
`281`	`300`	`{`
`282`	`301`	`await _innerConnection.DisposeAsync()`
`283`	`302`	`.ConfigureAwait(false);`
	`303`	`+`
	`304`	`+ _channels.Clear();`
	`305`	`+ _recordedEntitiesSemaphore.Dispose();`
	`306`	`+ _channelsSemaphore.Dispose();`
	`307`	`+ _recoveryCancellationTokenSource.Dispose();`
`284`	`308`	`}`
`285`	`309`	`catch (OperationInterruptedException)`
`286`	`310`	`{`
`287`	`311`	`// ignored, see rabbitmq/rabbitmq-dotnet-client#133`
`288`	`312`	`}`
`289`	`313`	`finally`
`290`	`314`	`{`
`291`		`- _channels.Clear();`
`292`		`- _recordedEntitiesSemaphore.Dispose();`
`293`		`- _channelsSemaphore.Dispose();`
`294`		`- _recoveryCancellationTokenSource.Dispose();`
`295`	`315`	`_disposed = true;`
	`316`	`+ _isDisposing = false;`
`296`	`317`	`}`
`297`	`318`	`}`
`298`	`319`
Original file line number	Diff line number	Diff line change
`@@ -46,6 +46,8 @@ namespace RabbitMQ.Client.Framing`
`46`	`46`	`internal sealed partial class Connection : IConnection`
`47`	`47`	`{`
`48`	`48`	`private bool _disposed;`
	`49`	`+ private bool _isDisposing;`
	`50`	`+ private readonly object _locker = new();`
`49`	`51`	`private volatile bool _closed;`
`50`	`52`
`51`	`53`	`private readonly ConnectionConfig _config;`
`@@ -485,7 +487,15 @@ internal ValueTask WriteAsync(RentedMemory frames, CancellationToken cancellatio`
`485`	`487`	`return _frameHandler.WriteAsync(frames, cancellationToken);`
`486`	`488`	`}`
`487`	`489`
`488`		`- public void Dispose() => DisposeAsync().AsTask().GetAwaiter().GetResult();`
	`490`	`+ public void Dispose()`
	`491`	`+ {`
	`492`	`+ if (_disposed)`
	`493`	`+ {`
	`494`	`+ return;`
	`495`	`+ }`
	`496`	`+`
	`497`	`+ DisposeAsync().AsTask().GetAwaiter().GetResult();`
	`498`	`+ }`
`489`	`499`
`490`	`500`	`public async ValueTask DisposeAsync()`
`491`	`501`	`{`
`@@ -494,6 +504,15 @@ public async ValueTask DisposeAsync()`
`494`	`504`	`return;`
`495`	`505`	`}`
`496`	`506`
	`507`	`+ lock (_locker)`
	`508`	`+ {`
	`509`	`+ if (_isDisposing)`
	`510`	`+ {`
	`511`	`+ return;`
	`512`	`+ }`
	`513`	`+ _isDisposing = true;`
	`514`	`+ }`
	`515`	`+`
`497`	`516`	`try`
`498`	`517`	`{`
`499`	`518`	`if (IsOpen)`
`@@ -515,6 +534,7 @@ await _channel0.DisposeAsync()`
`515`	`534`	`finally`
`516`	`535`	`{`
`517`	`536`	`_disposed = true;`
	`537`	`+ _isDisposing = false;`
`518`	`538`	`}`
`519`	`539`	`}`
`520`	`540`