Fix econnaborted seen in integration test runs

This error was introduced with commit 1fa0562 The main frame reading loop was cancelled too soon.
rabbitmq · Mar 27, 2024 · eaa7e33 · eaa7e33
1 parent f664259
commit eaa7e33
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 8 deletions.
diff --git a/projects/RabbitMQ.Client/client/impl/Connection.Commands.cs b/projects/RabbitMQ.Client/client/impl/Connection.Commands.cs
@@ -50,7 +50,7 @@ public Task UpdateSecretAsync(string newSecret, string reason)
 
         internal void NotifyReceivedCloseOk()
         {
-            TerminateMainloop();
+            MaybeTerminateMainloopAndStopHeartbeatTimers(cancelMainLoop: true);
             _closed = true;
         }
 
@@ -112,7 +112,11 @@ await _frameHandler.SendProtocolHeaderAsync(cancellationToken)
             var serverVersion = new AmqpVersion(connectionStart.m_versionMajor, connectionStart.m_versionMinor);
             if (!serverVersion.Equals(Protocol.Version))
             {
-                TerminateMainloop();
+                /*
+                 * Note:
+                 * FinishCloseAsync will cancel the main loop
+                 */
+                MaybeTerminateMainloopAndStopHeartbeatTimers();
                 await FinishCloseAsync(cancellationToken);
                 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
@@ -116,8 +116,13 @@ private async void HeartbeatReadTimerCallback(object? state)
 
                 if (shouldTerminate)
                 {
-                    TerminateMainloop();
-                    await FinishCloseAsync(_mainLoopCts.Token)
+                    MaybeTerminateMainloopAndStopHeartbeatTimers();
+                    /*
+                     * Note: do NOT use the main loop cancellation token,
+                     * because FininshCloseAsync immediately cancels it
+                     */
+                    using var cts = new CancellationTokenSource(InternalConstants.DefaultConnectionAbortTimeout);
+                    await FinishCloseAsync(cts.Token)
                         .ConfigureAwait(false);
                 }
                 else

diff --git a/projects/RabbitMQ.Client/client/impl/Connection.Receive.cs b/projects/RabbitMQ.Client/client/impl/Connection.Receive.cs
@@ -165,9 +165,12 @@ private async Task ProcessFrameAsync(InboundFrame frame, CancellationToken cance
         ///<remarks>
         /// May be called more than once. Should therefore be idempotent.
         ///</remarks>
-        private void TerminateMainloop()
+        private void MaybeTerminateMainloopAndStopHeartbeatTimers(bool cancelMainLoop = false)
         {
-            _mainLoopCts.Cancel();
+            if (cancelMainLoop)
+            {
+                _mainLoopCts.Cancel();
+            }
             MaybeStopHeartbeatTimers();
         }
 

diff --git a/projects/RabbitMQ.Client/client/impl/Connection.cs b/projects/RabbitMQ.Client/client/impl/Connection.cs
@@ -362,7 +362,11 @@ await _session0.TransmitAsync(method, cancellationToken)
                 }
                 finally
                 {
-                    TerminateMainloop();
+                    /*
+                     * Note:
+                     * NotifyReceivedCloseOk will cancel the main loop
+                     */
+                    MaybeTerminateMainloopAndStopHeartbeatTimers();
                 }
             }
 
@@ -402,7 +406,7 @@ internal void ClosedViaPeer(ShutdownEventArgs reason)
 
             OnShutdown(reason);
             _session0.SetSessionClosing(true);
-            TerminateMainloop();
+            MaybeTerminateMainloopAndStopHeartbeatTimers(cancelMainLoop: true);
         }
 
         // Only call at the end of the Mainloop or HeartbeatLoop