Create cancellation token from timeout

Fixes #1759

When passing a timeout of 0, `DisposeAsync` would block forever after closing a connection.

This change ensures that the timeout is used in a cancellation token.

Author: lukebakken

projects/RabbitMQ.Client/IConnectionExtensions.cs

@@ -59,10 +59,11 @@ public static Task CloseAsync(this IConnection connection, ushort reasonCode, st
         /// To wait infinitely for the close operations to complete use <see cref="System.Threading.Timeout.InfiniteTimeSpan"/>.
         /// </para>
         /// </remarks>
-        public static Task CloseAsync(this IConnection connection, TimeSpan timeout)
+        public static async Task CloseAsync(this IConnection connection, TimeSpan timeout)
         {
-            return connection.CloseAsync(Constants.ReplySuccess, "Goodbye", timeout, false,
-                CancellationToken.None);
+            using var cts = new CancellationTokenSource(timeout);
+            await connection.CloseAsync(Constants.ReplySuccess, "Goodbye", timeout, false, cts.Token)
+                .ConfigureAwait(false);
         }
 
         /// <summary>
@@ -82,10 +83,11 @@ public static Task CloseAsync(this IConnection connection, TimeSpan timeout)
         /// Operation timeout.
         /// </para>
         /// </remarks>
-        public static Task CloseAsync(this IConnection connection, ushort reasonCode, string reasonText, TimeSpan timeout)
+        public static async Task CloseAsync(this IConnection connection, ushort reasonCode, string reasonText, TimeSpan timeout)
         {
-            return connection.CloseAsync(reasonCode, reasonText, timeout, false,
-                CancellationToken.None);
+            using var cts = new CancellationTokenSource(timeout);
+            await connection.CloseAsync(reasonCode, reasonText, timeout, false, cts.Token)
+                .ConfigureAwait(false);
         }
 
         /// <summary>
@@ -97,10 +99,12 @@ public static Task CloseAsync(this IConnection connection, ushort reasonCode, st
         /// <see cref="IOException"/> during closing connection.
         ///This method waits infinitely for the in-progress close operation to complete.
         /// </remarks>
-        public static Task AbortAsync(this IConnection connection)
+        public static async Task AbortAsync(this IConnection connection)
         {
-            return connection.CloseAsync(Constants.ReplySuccess, "Connection close forced", InternalConstants.DefaultConnectionAbortTimeout, true,
-                CancellationToken.None);
+            using var cts = new CancellationTokenSource(InternalConstants.DefaultConnectionAbortTimeout);
+            await connection.CloseAsync(Constants.ReplySuccess,
+                "Connection close forced", InternalConstants.DefaultConnectionAbortTimeout, true, cts.Token)
+                .ConfigureAwait(false);
         }
 
         /// <summary>
@@ -116,10 +120,12 @@ public static Task AbortAsync(this IConnection connection)
         /// A message indicating the reason for closing the connection
         /// </para>
         /// </remarks>
-        public static Task AbortAsync(this IConnection connection, ushort reasonCode, string reasonText)
+        public static async Task AbortAsync(this IConnection connection, ushort reasonCode, string reasonText)
         {
-            return connection.CloseAsync(reasonCode, reasonText, InternalConstants.DefaultConnectionAbortTimeout, true,
-                CancellationToken.None);
+            using var cts = new CancellationTokenSource(InternalConstants.DefaultConnectionAbortTimeout);
+            await connection.CloseAsync(reasonCode,
+                reasonText, InternalConstants.DefaultConnectionAbortTimeout, true, cts.Token)
+                .ConfigureAwait(false);
         }
 
         /// <summary>
@@ -135,10 +141,12 @@ public static Task AbortAsync(this IConnection connection, ushort reasonCode, st
         /// To wait infinitely for the close operations to complete use <see cref="Timeout.Infinite"/>.
         /// </para>
         /// </remarks>
-        public static Task AbortAsync(this IConnection connection, TimeSpan timeout)
+        public static async Task AbortAsync(this IConnection connection, TimeSpan timeout)
         {
-            return connection.CloseAsync(Constants.ReplySuccess, "Connection close forced", timeout, true,
-                CancellationToken.None);
+            using var cts = new CancellationTokenSource(InternalConstants.DefaultConnectionAbortTimeout);
+            await connection.CloseAsync(Constants.ReplySuccess,
+                "Connection close forced", timeout, true, cts.Token)
+                .ConfigureAwait(false);
         }
 
         /// <summary>
@@ -155,10 +163,12 @@ public static Task AbortAsync(this IConnection connection, TimeSpan timeout)
         /// A message indicating the reason for closing the connection.
         /// </para>
         /// </remarks>
-        public static Task AbortAsync(this IConnection connection, ushort reasonCode, string reasonText, TimeSpan timeout)
+        public static async Task AbortAsync(this IConnection connection, ushort reasonCode, string reasonText, TimeSpan timeout)
         {
-            return connection.CloseAsync(reasonCode, reasonText, timeout, true,
-                CancellationToken.None);
+            using var cts = new CancellationTokenSource(InternalConstants.DefaultConnectionAbortTimeout);
+            await connection.CloseAsync(reasonCode,
+                reasonText, timeout, true, cts.Token)
+                .ConfigureAwait(false);
         }
     }
 }

projects/RabbitMQ.Client/Impl/Connection.cs

@@ -330,8 +330,6 @@ internal async Task CloseAsync(ShutdownEventArgs reason, bool abort, TimeSpan ti
             }
             else
             {
-                cancellationToken.ThrowIfCancellationRequested();
-
                 await OnShutdownAsync(reason)
                     .ConfigureAwait(false);
                 await _session0.SetSessionClosingAsync(false, cancellationToken)
@@ -518,7 +516,6 @@ await this.AbortAsync()
                 }
 
                 _session0.Dispose();
-                _mainLoopCts.Dispose();
 
                 await _channel0.DisposeAsync()
                     .ConfigureAwait(false);
@@ -529,6 +526,7 @@ await _channel0.DisposeAsync()
             }
             finally
             {
+                _mainLoopCts.Dispose();
                 _disposed = true;
             }
         }

projects/RabbitMQ.Client/Impl/MainSession.cs

@@ -47,6 +47,7 @@ internal sealed class MainSession : Session, IDisposable
         private volatile bool _closeIsServerInitiated;
         private volatile bool _closing;
         private readonly SemaphoreSlim _closingSemaphore = new SemaphoreSlim(1, 1);
+        private bool _disposed = false;
 
         public MainSession(Connection connection, uint maxBodyLength)
             : base(connection, 0, maxBodyLength)
@@ -83,6 +84,13 @@ public override Task HandleFrameAsync(InboundFrame frame, CancellationToken canc
 
         public async Task SetSessionClosingAsync(bool closeIsServerInitiated, CancellationToken cancellationToken)
         {
+            if (_disposed)
+            {
+                _closing = true;
+                _closeIsServerInitiated = closeIsServerInitiated;
+                return;
+            }
+
             if (await _closingSemaphore.WaitAsync(InternalConstants.DefaultConnectionAbortTimeout, cancellationToken)
                     .ConfigureAwait(false))
             {
@@ -122,6 +130,24 @@ public override ValueTask TransmitAsync<T>(in T cmd, CancellationToken cancellat
             return base.TransmitAsync(in cmd, cancellationToken);
         }
 
-        public void Dispose() => ((IDisposable)_closingSemaphore).Dispose();
+        public void Dispose()
+        {
+            if (_disposed)
+            {
+                return;
+            }
+
+            try
+            {
+                _closingSemaphore.Dispose();
+            }
+            catch
+            {
+            }
+            finally
+            {
+                _disposed = true;
+            }
+        }
     }
 }

projects/Test/Integration/GH/TestGitHubIssues.cs

@@ -131,5 +131,21 @@ public async Task TestHeartbeatTimeoutValue_GH1756()
 
             Assert.True(_conn.Heartbeat != default);
         }
+
+        [Fact]
+        public async Task DisposeWhileCatchingTimeoutDeadlocksRepro_GH1759()
+        {
+            _connFactory = new ConnectionFactory();
+            _conn = await _connFactory.CreateConnectionAsync();
+            try
+            {
+                await _conn.CloseAsync(TimeSpan.Zero);
+            }
+            catch (Exception)
+            {
+            }
+
+            await _conn.DisposeAsync();
+        }
     }
 }