Merge pull request #1122 from Particular/transport-pump-stop-release-3.2

Transport pump stop for release 3.2
Particular · Dec 19, 2024 · 9e73e6e · 9e73e6e
2 parents bf2cb4a + 1fb094d
commit 9e73e6e
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 8 deletions.
diff --git a/src/Transport/AzureServiceBusTransportInfrastructure.cs b/src/Transport/AzureServiceBusTransportInfrastructure.cs
@@ -1,4 +1,6 @@
-namespace NServiceBus.Transport.AzureServiceBus
+#nullable enable
+
+namespace NServiceBus.Transport.AzureServiceBus
 {
     using System.Linq;
     using System.Text;
@@ -16,7 +18,12 @@ sealed class AzureServiceBusTransportInfrastructure : TransportInfrastructure
         readonly ServiceBusClient defaultClient;
         readonly (ReceiveSettings receiveSettings, ServiceBusClient client)[] receiveSettingsAndClientPairs;
 
-        public AzureServiceBusTransportInfrastructure(AzureServiceBusTransport transportSettings, HostSettings hostSettings, (ReceiveSettings receiveSettings, ServiceBusClient client)[] receiveSettingsAndClientPairs, ServiceBusClient defaultClient)
+        public AzureServiceBusTransportInfrastructure(
+            AzureServiceBusTransport transportSettings,
+            HostSettings hostSettings,
+            (ReceiveSettings receiveSettings, ServiceBusClient client)[] receiveSettingsAndClientPairs,
+            ServiceBusClient defaultClient
+            )
         {
             this.transportSettings = transportSettings;
 
@@ -26,7 +33,10 @@ public AzureServiceBusTransportInfrastructure(AzureServiceBusTransport transport
 
             messageSenderRegistry = new MessageSenderRegistry(defaultClient);
 
-            Dispatcher = new MessageDispatcher(messageSenderRegistry, transportSettings.Topology.TopicToPublishTo);
+            Dispatcher = new MessageDispatcher(
+                messageSenderRegistry,
+                transportSettings.Topology.TopicToPublishTo
+                );
             Receivers = receiveSettingsAndClientPairs.ToDictionary(static settingsAndClient =>
             {
                 var (receiveSettings, _) = settingsAndClient;
@@ -65,15 +75,16 @@ IMessageReceiver CreateMessagePump(ReceiveSettings receiveSettings, ServiceBusCl
                 hostSettings.CriticalErrorAction,
                 receiveSettings.UsePublishSubscribe
                     ? new SubscriptionManager(receiveAddress, transportSettings, defaultClient)
-                    : null);
+                    : null
+                );
         }
 
         public override async Task Shutdown(CancellationToken cancellationToken = default)
         {
-            if (messageSenderRegistry != null)
-            {
-                await messageSenderRegistry.Close(cancellationToken).ConfigureAwait(false);
-            }
+            await Task.WhenAll(Receivers.Values.Select(r => r.StopReceive(cancellationToken)))
+                .ConfigureAwait(false);
+
+            await messageSenderRegistry.Close(cancellationToken).ConfigureAwait(false);
 
             foreach (var (_, serviceBusClient) in receiveSettingsAndClientPairs)
             {

diff --git a/src/Transport/Receiving/MessagePump.cs b/src/Transport/Receiving/MessagePump.cs
@@ -199,6 +199,12 @@ void UpdateProcessingCapacity(int maxConcurrency)
 
         public async Task StopReceive(CancellationToken cancellationToken = default)
         {
+            if (messageProcessingCancellationTokenSource is null)
+            {
+                // Receiver hasn't been started or is already stopped
+                return;
+            }
+
             // Wiring up the stop token to trigger the cancellation token that is being
             // used inside the message handling pipeline
             using var _ = cancellationToken