change ChangeConcurrency implementation to prevent closing connection midflight. (#1577)

ramonsmits · timbussmann · web-flow · commit 57bed5481fff · 2025-03-10T14:01:01.000+01:00
Co-authored-by: Tim Bussmann &lt;timbussmann@users.noreply.github.com&gt;
diff --git a/src/NServiceBus.Transport.RabbitMQ.TransportTests/Changing_concurrency.cs b/src/NServiceBus.Transport.RabbitMQ.TransportTests/Changing_concurrency.cs
@@ -0,0 +1,96 @@
+﻿namespace NServiceBus.Transport.RabbitMQ.TransportTests
+{
+    using System;
+    using System.Collections.Generic;
+    using System.Threading;
+    using System.Threading.Tasks;
+    using NServiceBus.TransportTests;
+    using NUnit.Framework;
+
+    public class Changing_concurrency : NServiceBusTransportTest
+    {
+        [TestCase(TransportTransactionMode.ReceiveOnly)]
+        [TestCase(TransportTransactionMode.SendsAtomicWithReceive)]
+        [TestCase(TransportTransactionMode.TransactionScope)]
+        public async Task Should_complete_current_message(TransportTransactionMode transactionMode)
+        {
+            var triggeredChangeConcurrency = CreateTaskCompletionSource();
+            var concurrencyChanged = CreateTaskCompletionSource();
+            int invocationCounter = 0;
+
+            await StartPump(async (context, ct) =>
+                {
+                    Interlocked.Increment(ref invocationCounter);
+
+                    _ = Task.Run(async () =>
+                    {
+                        var task = receiver.ChangeConcurrency(new PushRuntimeSettings(1), ct);
+                        triggeredChangeConcurrency.SetResult();
+                        await task;
+                        concurrencyChanged.SetResult();
+                    }, ct);
+
+                    await triggeredChangeConcurrency.Task;
+
+                }, (_, ct) =>
+                {
+                    Assert.Fail("Message processing should not fail");
+                    return Task.FromResult(ErrorHandleResult.RetryRequired);
+                },
+                transactionMode);
+
+            await SendMessage(InputQueueName);
+            await concurrencyChanged.Task;
+            await StopPump();
+            Assert.AreEqual(1, invocationCounter, "message should successfully complete on first processing attempt");
+        }
+
+        [TestCase(TransportTransactionMode.ReceiveOnly)]
+        [TestCase(TransportTransactionMode.SendsAtomicWithReceive)]
+        [TestCase(TransportTransactionMode.TransactionScope)]
+        public async Task Should_dispatch_delayed_retries(TransportTransactionMode transactionMode)
+        {
+            int invocationCounter = 0;
+
+            var triggeredChangeConcurrency = CreateTaskCompletionSource();
+            var sentMessageReceived = CreateTaskCompletionSource();
+
+            await StartPump((context, _) =>
+                {
+                    Interlocked.Increment(ref invocationCounter);
+                    if (context.Headers.TryGetValue("FromOnError", out var value) && value == bool.TrueString)
+                    {
+                        sentMessageReceived.SetResult();
+                    }
+                    else
+                    {
+                        throw new Exception("triggering recoverability pipeline");
+                    }
+
+                    return Task.CompletedTask;
+                }, async (context, ct) =>
+                {
+                    // same behavior as delayed retries
+                    _ = Task.Run(async () =>
+                    {
+                        var task = receiver.ChangeConcurrency(new PushRuntimeSettings(1), ct);
+                        triggeredChangeConcurrency.SetResult();
+                        await task;
+                    }, ct);
+
+                    await triggeredChangeConcurrency.Task;
+                    await SendMessage(InputQueueName,
+                        new Dictionary<string, string>() { { "FromOnError", bool.TrueString } },
+                        context.TransportTransaction, cancellationToken: ct);
+                    return ErrorHandleResult.Handled;
+                },
+                transactionMode);
+
+            await SendMessage(InputQueueName);
+
+            await sentMessageReceived.Task;
+            await StopPump();
+            Assert.AreEqual(2, invocationCounter, "there should be exactly 2 messages (initial message and new message from recoverability pipeline)");
+        }
+    }
+}
diff --git a/src/NServiceBus.Transport.RabbitMQ.TransportTests/When_restarting_endpoint.cs b/src/NServiceBus.Transport.RabbitMQ.TransportTests/When_restarting_endpoint.cs
@@ -0,0 +1,112 @@
+﻿namespace NServiceBus.TransportTests
+{
+    using System;
+    using System.Collections.Concurrent;
+    using System.Collections.Generic;
+    using System.Threading.Tasks;
+    using NUnit.Framework;
+    using Transport;
+
+    public class When_restarting_endpoint : NServiceBusTransportTest
+    {
+        [TestCase(TransportTransactionMode.None)]
+        [TestCase(TransportTransactionMode.ReceiveOnly)]
+        [TestCase(TransportTransactionMode.SendsAtomicWithReceive)]
+        [TestCase(TransportTransactionMode.TransactionScope)]
+        public async Task Should_handle_multiple_stop_calls_to_started_receiver(TransportTransactionMode transactionMode)
+        {
+            await StartPump(
+                (context, token) => Task.CompletedTask,
+                (context, token) => Task.FromResult(ErrorHandleResult.Handled),
+                transactionMode);
+
+            await receiver.StopReceive();
+            Assert.DoesNotThrowAsync(() => receiver.StopReceive());
+        }
+
+        [TestCase(TransportTransactionMode.None)]
+        [TestCase(TransportTransactionMode.ReceiveOnly)]
+        [TestCase(TransportTransactionMode.SendsAtomicWithReceive)]
+        [TestCase(TransportTransactionMode.TransactionScope)]
+        public async Task Should_allow_restarting_receivers(TransportTransactionMode transactionMode)
+        {
+            var messageReceived = CreateTaskCompletionSource();
+
+            await StartPump((context, token) =>
+                {
+                    messageReceived.SetResult();
+                    return Task.CompletedTask;
+                },
+                (context, token) => Task.FromResult(ErrorHandleResult.Handled), transactionMode);
+
+            await receiver.StopReceive();
+            await receiver.StartReceive();
+
+            await SendMessage(InputQueueName);
+            await messageReceived.Task;
+        }
+
+        [TestCase(TransportTransactionMode.None)]
+        [TestCase(TransportTransactionMode.ReceiveOnly)]
+        [TestCase(TransportTransactionMode.SendsAtomicWithReceive)]
+        [TestCase(TransportTransactionMode.TransactionScope)]
+        public async Task Should_gracefully_restart_processing(TransportTransactionMode transactionMode)
+        {
+            var receivedMessages = new ConcurrentQueue<string>();
+            var pumpStopTriggered = CreateTaskCompletionSource();
+            var followupMessageReceived = CreateTaskCompletionSource();
+            await StartPump(async (context, token) =>
+            {
+                var messageType = context.Headers["Type"];
+                receivedMessages.Enqueue(messageType);
+                TestContext.WriteLine("Received message " + messageType);
+
+                switch (messageType)
+                {
+                    case "Start":
+
+                        // run async because the pump might block the return until all inflight messages are processed.
+                        var stopTask = Task.Run(async () =>
+                        {
+                            var t = receiver.StopReceive(token);
+                            pumpStopTriggered.SetResult();
+                            await t;
+                            TestContext.WriteLine("Stopped receiver");
+                        }, token);
+
+                        await pumpStopTriggered.Task;
+                        await SendMessage(InputQueueName, new Dictionary<string, string>() { { "Type", "Followup" } },
+                            context.TransportTransaction, cancellationToken: token);
+
+                        _ = stopTask.ContinueWith(async _ =>
+                        {
+                            await receiver.StartReceive(token);
+                            TestContext.WriteLine("Started receiver");
+                        }, token);
+
+                        break;
+                    case "Followup":
+                        followupMessageReceived.SetResult();
+                        break;
+                    default:
+                        throw new ArgumentException();
+                }
+            },
+                (context, token) =>
+                {
+                    Assert.Fail($"Message failed processing: {context.Exception}");
+                    return Task.FromResult(ErrorHandleResult.Handled);
+                }, transactionMode);
+
+            await SendMessage(InputQueueName, new Dictionary<string, string>() { { "Type", "Start" } });
+            await followupMessageReceived.Task;
+            await StopPump();
+
+            Assert.AreEqual(2, receivedMessages.Count);
+            Assert.IsTrue(receivedMessages.TryDequeue(out var firstMessageType));
+            Assert.AreEqual("Start", firstMessageType);
+            Assert.IsTrue(receivedMessages.TryDequeue(out var secondMessageType));
+            Assert.AreEqual("Followup", secondMessageType);
+        }
+    }
+}
diff --git a/src/NServiceBus.Transport.RabbitMQ/RabbitMQTransportInfrastructure.cs b/src/NServiceBus.Transport.RabbitMQ/RabbitMQTransportInfrastructure.cs
@@ -54,15 +54,12 @@ internal void SetupInfrastructure(string[] sendingQueues)
             }
         }
 
-        public override Task Shutdown(CancellationToken cancellationToken = default)
+        public override async Task Shutdown(CancellationToken cancellationToken = default)
         {
-            foreach (IMessageReceiver receiver in Receivers.Values)
-            {
-                ((MessagePump)receiver).Dispose();
-            }
+            await Task.WhenAll(Receivers.Values.Select(r => r.StopReceive(cancellationToken)))
+                .ConfigureAwait(false);
 
             channelProvider.Dispose();
-            return Task.CompletedTask;
         }
 
         public override string ToTransportAddress(QueueAddress address) => TranslateAddress(address);
diff --git a/src/NServiceBus.Transport.RabbitMQ/Receiving/MessagePump.cs b/src/NServiceBus.Transport.RabbitMQ/Receiving/MessagePump.cs
@@ -12,7 +12,7 @@
     using global::RabbitMQ.Client.Exceptions;
     using Logging;
 
-    sealed class MessagePump : IMessageReceiver, IDisposable
+    sealed class MessagePump : IMessageReceiver
     {
         static readonly ILog Logger = LogManager.GetLogger(typeof(MessagePump));
         static readonly TransportTransaction transportTransaction = new();
@@ -31,7 +31,6 @@ sealed class MessagePump : IMessageReceiver, IDisposable
         readonly FastConcurrentLru<string, int> deliveryAttempts = new(1_000);
         readonly FastConcurrentLru<string, bool> failedBasicAckMessages = new(1_000);
 
-        bool disposed;
         OnMessage onMessage;
         OnError onError;
         int maxConcurrency;
@@ -113,12 +112,13 @@ public Task StartReceive(CancellationToken cancellationToken = default)
             return Task.CompletedTask;
         }
 
-        public Task ChangeConcurrency(PushRuntimeSettings limitations, CancellationToken cancellationToken = default)
+        public async Task ChangeConcurrency(PushRuntimeSettings limitations, CancellationToken cancellationToken = default)
         {
-            maxConcurrency = limitations.MaxConcurrency;
             Logger.InfoFormat("Calling a change concurrency and reconnecting with new value {0}.", limitations.MaxConcurrency);
-            _ = Task.Run(() => Reconnect(messageProcessingCancellationTokenSource.Token), cancellationToken);
-            return Task.CompletedTask;
+
+            await StopReceive(cancellationToken).ConfigureAwait(false);
+            maxConcurrency = limitations.MaxConcurrency;
+            await StartReceive(CancellationToken.None).ConfigureAwait(false);
         }
 
         void ConnectToBroker()
@@ -148,6 +148,12 @@ void ConnectToBroker()
 
         public async Task StopReceive(CancellationToken cancellationToken = default)
         {
+            if (messagePumpCancellationTokenSource is null)
+            {
+                // Receiver hasn't been started or is already stopped
+                return;
+            }
+
             messagePumpCancellationTokenSource?.Cancel();
 
             using (cancellationToken.Register(() => messageProcessingCancellationTokenSource?.Cancel()))
@@ -183,6 +189,12 @@ public async Task StopReceive(CancellationToken cancellationToken = default)
 
                 await connectionShutdownCompleted.Task.ConfigureAwait(false);
             }
+
+            messagePumpCancellationTokenSource?.Dispose();
+            messagePumpCancellationTokenSource = null;
+            messageProcessingCancellationTokenSource?.Dispose();
+            connection.Dispose();
+            circuitBreaker?.Dispose();
         }
 
 #pragma warning disable PS0018 // A task-returning method should have a CancellationToken parameter unless it has a parameter implementing ICancellableContext
@@ -511,24 +523,5 @@ async Task MovePoisonMessage(AsyncEventingBasicConsumer consumer, BasicDeliverEv
                 Logger.Warn($"Failed to acknowledge poison message because the channel was closed. The message was sent to queue '{queue}' but also returned to the original queue.", ex);
             }
         }
-
-        public void Dispose()
-        {
-            if (disposed)
-            {
-                return;
-            }
-
-            circuitBreaker?.Dispose();
-            messagePumpCancellationTokenSource?.Cancel();
-            messagePumpCancellationTokenSource?.Dispose();
-            // This makes sure that if the stop token hasn't been canceled the processing source is canceled
-            // so that any possible reconnect attempt has the possibility to gracefully stop too.
-            messageProcessingCancellationTokenSource?.Cancel();
-            messageProcessingCancellationTokenSource?.Dispose();
-
-            connection?.Dispose();
-            disposed = true;
-        }
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -54,15 +54,12 @@ internal void SetupInfrastructure(string[] sendingQueues)`
`54`	`54`	`}`
`55`	`55`	`}`
`56`	`56`
`57`		`- public override Task Shutdown(CancellationToken cancellationToken = default)`
	`57`	`+ public override async Task Shutdown(CancellationToken cancellationToken = default)`
`58`	`58`	`{`
`59`		`- foreach (IMessageReceiver receiver in Receivers.Values)`
`60`		`- {`
`61`		`- ((MessagePump)receiver).Dispose();`
`62`		`- }`
	`59`	`+ await Task.WhenAll(Receivers.Values.Select(r => r.StopReceive(cancellationToken)))`
	`60`	`+ .ConfigureAwait(false);`
`63`	`61`
`64`	`62`	`channelProvider.Dispose();`
`65`		`- return Task.CompletedTask;`
`66`	`63`	`}`
`67`	`64`
`68`	`65`	`public override string ToTransportAddress(QueueAddress address) => TranslateAddress(address);`