Merge pull request #1121 from Particular/transport-pump-stop-release-4.0

Transport pump stop for release 4.0
Particular · Dec 19, 2024 · 3d34e8c · 3d34e8c
2 parents 68235c5 + 312a633
commit 3d34e8c
Show file tree

Hide file tree

Showing 2 changed files with 24 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;
@@ -70,10 +80,10 @@ IMessageReceiver CreateMessagePump(ReceiveSettings receiveSettings, ServiceBusCl
 
         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)
             {
@@ -100,4 +110,4 @@ public override string ToTransportAddress(QueueAddress address)
             return queue.ToString();
         }
     }
-}
+}
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
             await using var _ = cancellationToken