Stop pumps on shutdown (#1468)

* Stop pumps on shutdown

* Allow stopping multiple times

* Stop the pumps on shutdown

* Cleanup

---------

Co-authored-by: Daniel Marbach <[email protected]>

Author: danielmarbach

src/NServiceBus.Transport.PostgreSql/PostgreSqlTransportInfrastructure.cs

@@ -51,9 +51,15 @@ public PostgreSqlTransportInfrastructure(PostgreSqlTransport transport, HostSett
         exceptionClassifier = new PostgreSqlExceptionClassifier();
     }
 
-    public override Task Shutdown(CancellationToken cancellationToken = default)
+    public override async Task Shutdown(CancellationToken cancellationToken = default)
     {
-        return dueDelayedMessageProcessor?.Stop(cancellationToken) ?? Task.FromResult(0);
+        await Task.WhenAll(Receivers.Values.Select(pump => pump.StopReceive(cancellationToken)))
+            .ConfigureAwait(false);
+
+        if (dueDelayedMessageProcessor != null)
+        {
+            await dueDelayedMessageProcessor.Stop(cancellationToken).ConfigureAwait(false);
+        }
     }
 
     public override string ToTransportAddress(Transport.QueueAddress address)

src/NServiceBus.Transport.PostgreSql/Receiving/QueueCreator.cs

@@ -23,41 +23,36 @@ public QueueCreator(ISqlConstants sqlConstants, DbConnectionFactory connectionFa
 
         public async Task CreateQueueIfNecessary(string[] queueAddresses, CanonicalQueueAddress delayedQueueAddress, CancellationToken cancellationToken = default)
         {
-            using (var connection = await connectionFactory.OpenNewConnection(cancellationToken).ConfigureAwait(false))
+            using var connection = await connectionFactory.OpenNewConnection(cancellationToken).ConfigureAwait(false);
+            foreach (var address in queueAddresses)
             {
-                foreach (var address in queueAddresses)
-                {
-                    await CreateQueue(sqlConstants.CreateQueueText, addressTranslator(address), connection, createMessageBodyColumn, cancellationToken).ConfigureAwait(false);
-                }
+                await CreateQueue(sqlConstants.CreateQueueText, addressTranslator(address), connection, createMessageBodyColumn, cancellationToken).ConfigureAwait(false);
+            }
 
-                if (delayedQueueAddress != null)
-                {
-                    await CreateQueue(sqlConstants.CreateDelayedMessageStoreText, delayedQueueAddress, connection, createMessageBodyColumn, cancellationToken).ConfigureAwait(false);
-                }
+            if (delayedQueueAddress != null)
+            {
+                await CreateQueue(sqlConstants.CreateDelayedMessageStoreText, delayedQueueAddress, connection, createMessageBodyColumn, cancellationToken).ConfigureAwait(false);
             }
         }
 
         async Task CreateQueue(string creationScript, CanonicalQueueAddress canonicalQueueAddress, DbConnection connection, bool createMessageBodyColumn, CancellationToken cancellationToken)
         {
             try
             {
-                using (var transaction = connection.BeginTransaction())
-                {
-                    var tableName = canonicalQueueAddress.QualifiedTableName;
+                using var transaction = await connection.BeginTransactionAsync(cancellationToken).ConfigureAwait(false);
+                var tableName = canonicalQueueAddress.QualifiedTableName;
 
-                    var sql = string.Format(creationScript, tableName, canonicalQueueAddress.Table);
-                    using (var command = connection.CreateCommand())
-                    {
-                        command.Transaction = transaction;
-                        command.CommandText = sql;
-                        command.CommandType = CommandType.Text;
-
-                        await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
-
-                    }
+                var sql = string.Format(creationScript, tableName, canonicalQueueAddress.Table);
+                using (var command = connection.CreateCommand())
+                {
+                    command.Transaction = transaction;
+                    command.CommandText = sql;
+                    command.CommandType = CommandType.Text;
 
-                    transaction.Commit();
+                    await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
                 }
+
+                await transaction.CommitAsync(cancellationToken).ConfigureAwait(false);
             }
             catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
             {
@@ -75,19 +70,17 @@ async Task CreateQueue(string creationScript, CanonicalQueueAddress canonicalQue
                 var advisoryLockId = CalculateLockId(canonicalQueueAddress.QualifiedTableName);
                 var bodyStringSql = string.Format(sqlConstants.AddMessageBodyStringColumn, canonicalQueueAddress.Schema, canonicalQueueAddress.Table, advisoryLockId);
 
-                using (var transaction = connection.BeginTransaction())
+                using var transaction = await connection.BeginTransactionAsync(cancellationToken).ConfigureAwait(false);
+                using (var command = connection.CreateCommand())
                 {
-                    using (var command = connection.CreateCommand())
-                    {
-                        command.Transaction = transaction;
-                        command.CommandText = bodyStringSql;
-                        command.CommandType = CommandType.Text;
-
-                        await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
-                    }
+                    command.Transaction = transaction;
+                    command.CommandText = bodyStringSql;
+                    command.CommandType = CommandType.Text;
 
-                    transaction.Commit();
+                    await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
                 }
+
+                await transaction.CommitAsync(cancellationToken).ConfigureAwait(false);
             }
         }
 

src/NServiceBus.Transport.Sql.Shared/Receiving/MessageReceiver.cs

@@ -120,9 +120,15 @@ public async Task ChangeConcurrency(PushRuntimeSettings newLimitations,
 
         public async Task StopReceive(CancellationToken cancellationToken = default)
         {
-            messageReceivingCancellationTokenSource?.Cancel();
+            if (messageReceivingCancellationTokenSource == null)
+            {
+                // already stopped or never started
+                return;
+            }
+
+            await messageReceivingCancellationTokenSource.CancelAsync().ConfigureAwait(false);
 
-            using (cancellationToken.Register(() => messageProcessingCancellationTokenSource?.Cancel()))
+            await using (cancellationToken.Register(() => messageProcessingCancellationTokenSource?.Cancel()))
             {
                 await messageReceivingTask.ConfigureAwait(false);
 
@@ -140,8 +146,9 @@ public async Task StopReceive(CancellationToken cancellationToken = default)
             messageReceivingCircuitBreaker.Dispose();
             messageProcessingCircuitBreaker.Dispose();
             concurrencyLimiter.Dispose();
-            messageReceivingCancellationTokenSource?.Dispose();
-            messageProcessingCancellationTokenSource?.Dispose();
+            messageReceivingCancellationTokenSource.Dispose();
+            messageReceivingCancellationTokenSource = null;
+            messageProcessingCancellationTokenSource.Dispose();
         }
 
         async Task ReceiveMessagesAndSwallowExceptions(CancellationToken messageReceivingCancellationToken)

src/NServiceBus.Transport.SqlServer/SqlServerTransportInfrastructure.cs

@@ -342,11 +342,16 @@ public void ConfigureSendInfrastructure()
                 connectionFactory);
         }
 
-        public override Task Shutdown(CancellationToken cancellationToken = default)
+        public override async Task Shutdown(CancellationToken cancellationToken = default)
         {
-            return dueDelayedMessageProcessor?.Stop(cancellationToken) ?? Task.FromResult(0);
-        }
+            await Task.WhenAll(Receivers.Values.Select(pump => pump.StopReceive(cancellationToken)))
+                .ConfigureAwait(false);
 
+            if (dueDelayedMessageProcessor != null)
+            {
+                await dueDelayedMessageProcessor.Stop(cancellationToken).ConfigureAwait(false);
+            }
+        }
 
         readonly SqlServerTransport transport;
         readonly HostSettings hostSettings;