Pass cancellation token to IFrameHandler.CloseAsync

lukebakken · lukebakken · commit d9d916620e4f · 2024-01-31T15:11:54.000-08:00
diff --git a/projects/RabbitMQ.Client/client/impl/AutorecoveringConnection.Recovery.cs b/projects/RabbitMQ.Client/client/impl/AutorecoveringConnection.Recovery.cs
@@ -94,7 +94,7 @@ await Task.Delay(_config.NetworkRecoveryInterval, token)
                         .ConfigureAwait(false);
                     success = await TryPerformAutomaticRecoveryAsync(token)
                         .ConfigureAwait(false);
-                } while (!success && !token.IsCancellationRequested);
+                } while (false == success && false == token.IsCancellationRequested);
             }
             catch (OperationCanceledException)
             {
diff --git a/projects/RabbitMQ.Client/client/impl/AutorecoveringConnection.cs b/projects/RabbitMQ.Client/client/impl/AutorecoveringConnection.cs
@@ -196,7 +196,7 @@ public override string ToString()
 
         internal Task CloseFrameHandlerAsync()
         {
-            return InnerConnection.FrameHandler.CloseAsync();
+            return InnerConnection.FrameHandler.CloseAsync(CancellationToken.None);
         }
 
         ///<summary>API-side invocation of updating the secret.</summary>
@@ -208,17 +208,44 @@ public Task UpdateSecretAsync(string newSecret, string reason)
         }
 
         ///<summary>Asynchronous API-side invocation of connection.close with timeout.</summary>
-        public async Task CloseAsync(ushort reasonCode, string reasonText, TimeSpan timeout, bool abort, CancellationToken cancellationToken = default)
+        public async Task CloseAsync(ushort reasonCode, string reasonText, TimeSpan timeout, bool abort,
+            CancellationToken cancellationToken = default)
         {
             ThrowIfDisposed();
-            // TODO
-            // what to do if this method throws OperationCanceledException? We still need to close the connection.
-            await StopRecoveryLoopAsync(cancellationToken)
-                .ConfigureAwait(false);
-            if (_innerConnection.IsOpen)
+
+            Task CloseInnerConnectionAsync()
+            {
+                if (_innerConnection.IsOpen)
+                {
+                    return _innerConnection.CloseAsync(reasonCode, reasonText, timeout, abort, cancellationToken);
+                }
+                else
+                {
+                    return Task.CompletedTask;
+                }
+            }
+
+            try
             {
-                await _innerConnection.CloseAsync(reasonCode, reasonText, timeout, abort, cancellationToken)
+                await StopRecoveryLoopAsync(cancellationToken)
                     .ConfigureAwait(false);
+
+                await CloseInnerConnectionAsync()
+                    .ConfigureAwait(false);
+            }
+            catch (Exception ex)
+            {
+                try
+                {
+                    await CloseInnerConnectionAsync()
+                        .ConfigureAwait(false);
+                }
+                catch (Exception innerConnectionException)
+                {
+                    throw new AggregateException(ex, innerConnectionException);
+                }
+
+                throw;
             }
         }
 
diff --git a/projects/RabbitMQ.Client/client/impl/Connection.Commands.cs b/projects/RabbitMQ.Client/client/impl/Connection.Commands.cs
@@ -114,7 +114,7 @@ await _frameHandler.SendProtocolHeaderAsync(cancellationToken)
             {
                 TerminateMainloop();
                 // TODO hmmm
-                FinishCloseAsync().EnsureCompleted();
+                FinishCloseAsync(CancellationToken.None).EnsureCompleted();
                 throw new ProtocolVersionMismatchException(Protocol.MajorVersion, Protocol.MinorVersion, serverVersion.Major, serverVersion.Minor);
             }
 
diff --git a/projects/RabbitMQ.Client/client/impl/Connection.Heartbeat.cs b/projects/RabbitMQ.Client/client/impl/Connection.Heartbeat.cs
@@ -120,7 +120,7 @@ private void HeartbeatReadTimerCallback(object? state)
                 {
                     TerminateMainloop();
                     // TODO hmmm
-                    FinishCloseAsync().EnsureCompleted();
+                    FinishCloseAsync(CancellationToken.None).EnsureCompleted();
                 }
                 else
                 {
diff --git a/projects/RabbitMQ.Client/client/impl/Connection.Receive.cs b/projects/RabbitMQ.Client/client/impl/Connection.Receive.cs
@@ -45,18 +45,18 @@ internal sealed partial class Connection
         private readonly IFrameHandler _frameHandler;
         private Task _mainLoopTask;
 
-        // TODO cancellation token
         private async Task MainLoop()
         {
-            CancellationToken _mainLoopToken = _mainLoopCts.Token;
+            CancellationToken mainLoopToken = _mainLoopCts.Token;
             try
             {
-                await ReceiveLoopAsync(_mainLoopToken)
+                await ReceiveLoopAsync(mainLoopToken)
                     .ConfigureAwait(false);
             }
             catch (OperationCanceledException)
             {
                 // TODO what to do here?
+                // Debug log?
             }
             catch (EndOfStreamException eose)
             {
@@ -68,7 +68,7 @@ await ReceiveLoopAsync(_mainLoopToken)
             }
             catch (HardProtocolException hpe)
             {
-                await HardProtocolExceptionHandlerAsync(hpe, _mainLoopToken)
+                await HardProtocolExceptionHandlerAsync(hpe, mainLoopToken)
                     .ConfigureAwait(false);
             }
             catch (FileLoadException fileLoadException)
@@ -91,7 +91,9 @@ await HardProtocolExceptionHandlerAsync(hpe, _mainLoopToken)
                 HandleMainLoopException(ea);
             }
 
-            await FinishCloseAsync();
+            // TODO is this the best way?
+            using var cts = new CancellationTokenSource(InternalConstants.DefaultConnectionCloseTimeout);
+            await FinishCloseAsync(cts.Token);
         }
 
         private async Task ReceiveLoopAsync(CancellationToken mainLoopCancelllationToken)
diff --git a/projects/RabbitMQ.Client/client/impl/Connection.cs b/projects/RabbitMQ.Client/client/impl/Connection.cs
@@ -151,7 +151,7 @@ public event EventHandler<ShutdownEventArgs> ConnectionShutdown
             add
             {
                 ThrowIfDisposed();
-                var reason = CloseReason;
+                ShutdownEventArgs? reason = CloseReason;
                 if (reason is null)
                 {
                     _connectionShutdownWrapper.AddHandler(value);
@@ -379,8 +379,7 @@ await _mainLoopTask.WaitAsync(timeout, cancellationToken)
             {
                 try
                 {
-                    // TODO cancellation token?
-                    await _frameHandler.CloseAsync()
+                    await _frameHandler.CloseAsync(cancellationToken)
                         .ConfigureAwait(false);
                 }
                 catch
@@ -411,13 +410,13 @@ internal void ClosedViaPeer(ShutdownEventArgs reason)
         }
 
         // Only call at the end of the Mainloop or HeartbeatLoop
-        private async Task FinishCloseAsync()
+        private async Task FinishCloseAsync(CancellationToken cancellationToken)
         {
             _mainLoopCts.Cancel();
             _closed = true;
             MaybeStopHeartbeatTimers();
 
-            await _frameHandler.CloseAsync();
+            await _frameHandler.CloseAsync(cancellationToken);
             _channel0.SetCloseReason(CloseReason);
             _channel0.FinishClose();
             RabbitMqClientEventSource.Log.ConnectionClosed();
@@ -483,6 +482,17 @@ public void Dispose()
                 {
                     throw new InvalidOperationException("Connection must be closed before calling Dispose!");
                 }
+
+                _session0.Dispose();
+                /*
+                 * TODO dispose channel?
+                if (_channel0 != null && _channel0.IsOpen)
+                {
+                    IChannel ch0 = _channel0 as IChannel;
+                    ch0?.Dispose();
+                }
+                */
+                _mainLoopCts.Dispose();
             }
             catch (OperationInterruptedException)
             {
diff --git a/projects/RabbitMQ.Client/client/impl/IFrameHandler.cs b/projects/RabbitMQ.Client/client/impl/IFrameHandler.cs
@@ -56,7 +56,7 @@ internal interface IFrameHandler
         ///<summary>Socket write timeout. System.Threading.Timeout.InfiniteTimeSpan signals "infinity".</summary>
         TimeSpan WriteTimeout { set; }
 
-        Task CloseAsync();
+        Task CloseAsync(CancellationToken cancellationToken);
 
         ///<summary>Read a frame from the underlying
         ///transport. Returns null if the read operation timed out
diff --git a/projects/RabbitMQ.Client/client/impl/MainSession.cs b/projects/RabbitMQ.Client/client/impl/MainSession.cs
@@ -43,11 +43,10 @@
 namespace RabbitMQ.Client.Impl
 {
     ///<summary>Small ISession implementation used only for channel 0.</summary>
-    internal sealed class MainSession : Session
+    internal sealed class MainSession : Session, IDisposable
     {
         private volatile bool _closeIsServerInitiated;
         private volatile bool _closing;
-        // TODO dispose
         private readonly SemaphoreSlim _closingSemaphore = new SemaphoreSlim(1, 1);
 
         public MainSession(Connection connection) : base(connection, 0)
@@ -167,5 +166,7 @@ public override ValueTask TransmitAsync<T>(in T cmd, CancellationToken cancellat
 
             return base.TransmitAsync(in cmd, cancellationToken);
         }
+
+        public void Dispose() => ((IDisposable)_closingSemaphore).Dispose();
     }
 }
diff --git a/projects/RabbitMQ.Client/client/impl/SocketFrameHandler.cs b/projects/RabbitMQ.Client/client/impl/SocketFrameHandler.cs
@@ -218,7 +218,7 @@ public async Task ConnectAsync(CancellationToken cancellationToken)
                 }
                 catch
                 {
-                    await CloseAsync()
+                    await CloseAsync(cancellationToken)
                         .ConfigureAwait(false);
                     throw;
                 }
@@ -231,18 +231,17 @@ await CloseAsync()
             _connected = true;
         }
 
-        // TODO cancellationToken
-        public async Task CloseAsync()
+        public async Task CloseAsync(CancellationToken cancellationToken)
         {
             if (_closed || _socket == null)
             {
                 return;
             }
 
-            await _closingSemaphore.WaitAsync()
-                .ConfigureAwait(false);
             try
             {
+                await _closingSemaphore.WaitAsync(cancellationToken)
+                    .ConfigureAwait(false);
                 try
                 {
                     _channelWriter.Complete();
@@ -269,9 +268,13 @@ await _pipeReader.CompleteAsync()
                     // ignore, we are closing anyway
                 }
             }
+            catch
+            {
+            }
             finally
             {
                 _closingSemaphore.Release();
+                _closingSemaphore.Dispose();
                 _closed = true;
             }
         }
@@ -441,7 +444,7 @@ await tcpClient.ConnectAsync(endpoint.Address, endpoint.Port, linkedTokenSource.
             {
                 if (timeoutTokenSource.Token.IsCancellationRequested)
                 {
-                    // TODO do not use System.TimeoutException here
+                    // TODO maybe do not use System.TimeoutException here
                     var timeoutException = new TimeoutException(msg, e);
                     throw new ConnectFailureException(msg, timeoutException);
                 }
diff --git a/projects/Test/Common/IntegrationFixtureBase.cs b/projects/Test/Common/IntegrationFixtureBase.cs
@@ -390,11 +390,11 @@ protected static async Task WaitAsync(TaskCompletionSource<bool> tcs, TimeSpan t
                 await tcs.Task.WaitAsync(timeSpan);
                 bool result = await tcs.Task;
                 Assert.True((true == result) && (tcs.Task.IsCompletedSuccessfully()),
-                    $"waiting {timeSpan.TotalSeconds} seconds on a tcs for '{desc}' timed out");
+                    $"waiting {timeSpan.TotalSeconds} seconds on a tcs for '{desc}' failed");
             }
-            catch (TimeoutException)
+            catch (TimeoutException ex)
             {
-                Assert.Fail($"waiting {timeSpan.TotalSeconds} seconds on a tcs for '{desc}' timed out");
+                Assert.Fail($"waiting {timeSpan.TotalSeconds} seconds on a tcs for '{desc}' timed out, ex: {ex}");
             }
         }
 

Original file line number	Diff line number	Diff line change
`@@ -94,7 +94,7 @@ await Task.Delay(_config.NetworkRecoveryInterval, token)`
`94`	`94`	`.ConfigureAwait(false);`
`95`	`95`	`success = await TryPerformAutomaticRecoveryAsync(token)`
`96`	`96`	`.ConfigureAwait(false);`
`97`		`- } while (!success && !token.IsCancellationRequested);`
	`97`	`+ } while (false == success && false == token.IsCancellationRequested);`
`98`	`98`	`}`
`99`	`99`	`catch (OperationCanceledException)`
`100`	`100`	`{`
Original file line number	Diff line number	Diff line change
`@@ -114,7 +114,7 @@ await _frameHandler.SendProtocolHeaderAsync(cancellationToken)`
`114`	`114`	`{`
`115`	`115`	`TerminateMainloop();`
`116`	`116`	`// TODO hmmm`
`117`		`- FinishCloseAsync().EnsureCompleted();`
	`117`	`+ FinishCloseAsync(CancellationToken.None).EnsureCompleted();`
`118`	`118`	`throw new ProtocolVersionMismatchException(Protocol.MajorVersion, Protocol.MinorVersion, serverVersion.Major, serverVersion.Minor);`
`119`	`119`	`}`
`120`	`120`
Original file line number	Diff line number	Diff line change
`@@ -120,7 +120,7 @@ private void HeartbeatReadTimerCallback(object? state)`
`120`	`120`	`{`
`121`	`121`	`TerminateMainloop();`
`122`	`122`	`// TODO hmmm`
`123`		`- FinishCloseAsync().EnsureCompleted();`
	`123`	`+ FinishCloseAsync(CancellationToken.None).EnsureCompleted();`
`124`	`124`	`}`
`125`	`125`	`else`
`126`	`126`	`{`
Original file line number	Diff line number	Diff line change
`@@ -45,18 +45,18 @@ internal sealed partial class Connection`
`45`	`45`	`private readonly IFrameHandler _frameHandler;`
`46`	`46`	`private Task _mainLoopTask;`
`47`	`47`
`48`		`- // TODO cancellation token`
`49`	`48`	`private async Task MainLoop()`
`50`	`49`	`{`
`51`		`- CancellationToken _mainLoopToken = _mainLoopCts.Token;`
	`50`	`+ CancellationToken mainLoopToken = _mainLoopCts.Token;`
`52`	`51`	`try`
`53`	`52`	`{`
`54`		`- await ReceiveLoopAsync(_mainLoopToken)`
	`53`	`+ await ReceiveLoopAsync(mainLoopToken)`
`55`	`54`	`.ConfigureAwait(false);`
`56`	`55`	`}`
`57`	`56`	`catch (OperationCanceledException)`
`58`	`57`	`{`
`59`	`58`	`// TODO what to do here?`
	`59`	`+ // Debug log?`
`60`	`60`	`}`
`61`	`61`	`catch (EndOfStreamException eose)`
`62`	`62`	`{`
`@@ -68,7 +68,7 @@ await ReceiveLoopAsync(_mainLoopToken)`
`68`	`68`	`}`
`69`	`69`	`catch (HardProtocolException hpe)`
`70`	`70`	`{`
`71`		`- await HardProtocolExceptionHandlerAsync(hpe, _mainLoopToken)`
	`71`	`+ await HardProtocolExceptionHandlerAsync(hpe, mainLoopToken)`
`72`	`72`	`.ConfigureAwait(false);`
`73`	`73`	`}`
`74`	`74`	`catch (FileLoadException fileLoadException)`
`@@ -91,7 +91,9 @@ await HardProtocolExceptionHandlerAsync(hpe, _mainLoopToken)`
`91`	`91`	`HandleMainLoopException(ea);`
`92`	`92`	`}`
`93`	`93`
`94`		`- await FinishCloseAsync();`
	`94`	`+ // TODO is this the best way?`
	`95`	`+ using var cts = new CancellationTokenSource(InternalConstants.DefaultConnectionCloseTimeout);`
	`96`	`+ await FinishCloseAsync(cts.Token);`
`95`	`97`	`}`
`96`	`98`
`97`	`99`	`private async Task ReceiveLoopAsync(CancellationToken mainLoopCancelllationToken)`
Original file line number	Diff line number	Diff line change
`@@ -151,7 +151,7 @@ public event EventHandler<ShutdownEventArgs> ConnectionShutdown`
`151`	`151`	`add`
`152`	`152`	`{`
`153`	`153`	`ThrowIfDisposed();`
`154`		`- var reason = CloseReason;`
	`154`	`+ ShutdownEventArgs? reason = CloseReason;`
`155`	`155`	`if (reason is null)`
`156`	`156`	`{`
`157`	`157`	`_connectionShutdownWrapper.AddHandler(value);`
`@@ -379,8 +379,7 @@ await _mainLoopTask.WaitAsync(timeout, cancellationToken)`
`379`	`379`	`{`
`380`	`380`	`try`
`381`	`381`	`{`
`382`		`- // TODO cancellation token?`
`383`		`- await _frameHandler.CloseAsync()`
	`382`	`+ await _frameHandler.CloseAsync(cancellationToken)`
`384`	`383`	`.ConfigureAwait(false);`
`385`	`384`	`}`
`386`	`385`	`catch`
`@@ -411,13 +410,13 @@ internal void ClosedViaPeer(ShutdownEventArgs reason)`
`411`	`410`	`}`
`412`	`411`
`413`	`412`	`// Only call at the end of the Mainloop or HeartbeatLoop`
`414`		`- private async Task FinishCloseAsync()`
	`413`	`+ private async Task FinishCloseAsync(CancellationToken cancellationToken)`
`415`	`414`	`{`
`416`	`415`	`_mainLoopCts.Cancel();`
`417`	`416`	`_closed = true;`
`418`	`417`	`MaybeStopHeartbeatTimers();`
`419`	`418`
`420`		`- await _frameHandler.CloseAsync();`
	`419`	`+ await _frameHandler.CloseAsync(cancellationToken);`
`421`	`420`	`_channel0.SetCloseReason(CloseReason);`
`422`	`421`	`_channel0.FinishClose();`
`423`	`422`	`RabbitMqClientEventSource.Log.ConnectionClosed();`
`@@ -483,6 +482,17 @@ public void Dispose()`
`483`	`482`	`{`
`484`	`483`	`throw new InvalidOperationException("Connection must be closed before calling Dispose!");`
`485`	`484`	`}`
	`485`	`+`
	`486`	`+ _session0.Dispose();`
	`487`	`+ /*`
	`488`	`+ * TODO dispose channel?`
	`489`	`+ if (_channel0 != null && _channel0.IsOpen)`
	`490`	`+ {`
	`491`	`+ IChannel ch0 = _channel0 as IChannel;`
	`492`	`+ ch0?.Dispose();`
	`493`	`+ }`
	`494`	`+ */`
	`495`	`+ _mainLoopCts.Dispose();`
`486`	`496`	`}`
`487`	`497`	`catch (OperationInterruptedException)`
`488`	`498`	`{`
Original file line number	Diff line number	Diff line change
`@@ -43,11 +43,10 @@`
`43`	`43`	`namespace RabbitMQ.Client.Impl`
`44`	`44`	`{`
`45`	`45`	`///<summary>Small ISession implementation used only for channel 0.</summary>`
`46`		`- internal sealed class MainSession : Session`
	`46`	`+ internal sealed class MainSession : Session, IDisposable`
`47`	`47`	`{`
`48`	`48`	`private volatile bool _closeIsServerInitiated;`
`49`	`49`	`private volatile bool _closing;`
`50`		`- // TODO dispose`
`51`	`50`	`private readonly SemaphoreSlim _closingSemaphore = new SemaphoreSlim(1, 1);`
`52`	`51`
`53`	`52`	`public MainSession(Connection connection) : base(connection, 0)`
`@@ -167,5 +166,7 @@ public override ValueTask TransmitAsync<T>(in T cmd, CancellationToken cancellat`
`167`	`166`
`168`	`167`	`return base.TransmitAsync(in cmd, cancellationToken);`
`169`	`168`	`}`
	`169`	`+`
	`170`	`+ public void Dispose() => ((IDisposable)_closingSemaphore).Dispose();`
`170`	`171`	`}`
`171`	`172`	`}`
Original file line number	Diff line number	Diff line change
`@@ -218,7 +218,7 @@ public async Task ConnectAsync(CancellationToken cancellationToken)`
`218`	`218`	`}`
`219`	`219`	`catch`
`220`	`220`	`{`
`221`		`- await CloseAsync()`
	`221`	`+ await CloseAsync(cancellationToken)`
`222`	`222`	`.ConfigureAwait(false);`
`223`	`223`	`throw;`
`224`	`224`	`}`
`@@ -231,18 +231,17 @@ await CloseAsync()`
`231`	`231`	`_connected = true;`
`232`	`232`	`}`
`233`	`233`
`234`		`- // TODO cancellationToken`
`235`		`- public async Task CloseAsync()`
	`234`	`+ public async Task CloseAsync(CancellationToken cancellationToken)`
`236`	`235`	`{`
`237`	`236`	`if (_closed \|\| _socket == null)`
`238`	`237`	`{`
`239`	`238`	`return;`
`240`	`239`	`}`
`241`	`240`
`242`		`- await _closingSemaphore.WaitAsync()`
`243`		`- .ConfigureAwait(false);`
`244`	`241`	`try`
`245`	`242`	`{`
	`243`	`+ await _closingSemaphore.WaitAsync(cancellationToken)`
	`244`	`+ .ConfigureAwait(false);`
`246`	`245`	`try`
`247`	`246`	`{`
`248`	`247`	`_channelWriter.Complete();`
`@@ -269,9 +268,13 @@ await _pipeReader.CompleteAsync()`
`269`	`268`	`// ignore, we are closing anyway`
`270`	`269`	`}`
`271`	`270`	`}`
	`271`	`+ catch`
	`272`	`+ {`
	`273`	`+ }`
`272`	`274`	`finally`
`273`	`275`	`{`
`274`	`276`	`_closingSemaphore.Release();`
	`277`	`+ _closingSemaphore.Dispose();`
`275`	`278`	`_closed = true;`
`276`	`279`	`}`
`277`	`280`	`}`
`@@ -441,7 +444,7 @@ await tcpClient.ConnectAsync(endpoint.Address, endpoint.Port, linkedTokenSource.`
`441`	`444`	`{`
`442`	`445`	`if (timeoutTokenSource.Token.IsCancellationRequested)`
`443`	`446`	`{`
`444`		`- // TODO do not use System.TimeoutException here`
	`447`	`+ // TODO maybe do not use System.TimeoutException here`
`445`	`448`	`var timeoutException = new TimeoutException(msg, e);`
`446`	`449`	`throw new ConnectFailureException(msg, timeoutException);`
`447`	`450`	`}`
Original file line number	Diff line number	Diff line change
`@@ -390,11 +390,11 @@ protected static async Task WaitAsync(TaskCompletionSource<bool> tcs, TimeSpan t`
`390`	`390`	`await tcs.Task.WaitAsync(timeSpan);`
`391`	`391`	`bool result = await tcs.Task;`
`392`	`392`	`Assert.True((true == result) && (tcs.Task.IsCompletedSuccessfully()),`
`393`		`- $"waiting {timeSpan.TotalSeconds} seconds on a tcs for '{desc}' timed out");`
	`393`	`+ $"waiting {timeSpan.TotalSeconds} seconds on a tcs for '{desc}' failed");`
`394`	`394`	`}`
`395`		`- catch (TimeoutException)`
	`395`	`+ catch (TimeoutException ex)`
`396`	`396`	`{`
`397`		`- Assert.Fail($"waiting {timeSpan.TotalSeconds} seconds on a tcs for '{desc}' timed out");`
	`397`	`+ Assert.Fail($"waiting {timeSpan.TotalSeconds} seconds on a tcs for '{desc}' timed out, ex: {ex}");`
`398`	`398`	`}`
`399`	`399`	`}`
`400`	`400`