Simplify the latch and signaler usage

Co-Authored-By: Ramon Smits <[email protected]>
Co-Authored-By: Tomasz Masternak <[email protected]>
Co-Authored-By: Szymon Pobiega <[email protected]>

Author: 4 people

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

@@ -20,9 +20,15 @@ public AsyncCountdownLatch(int count)
         }
     }
 
-    public Task WaitAsync(CancellationToken cancellationToken = default) => completionSource.Task;
+#pragma warning disable PS0003
+    public Task WaitAsync(CancellationToken cancellationToken)
+#pragma warning restore PS0003
+    {
+        _ = cancellationToken.Register(completionSource.SetResult);
+        return completionSource.Task;
+    }
 
-    public Signaler GetSignaler() => new Signaler(this);
+    public Signaler GetSignaler() => new(this);
 
     void Signal()
     {

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

@@ -204,13 +204,11 @@ async Task ReceiveMessages(CancellationToken messageReceivingCancellationToken)
                     ? 1
                     : messageCount;
 
-            bool shouldWaitForReceiveTasks = true;
             var receiveLatch = new AsyncCountdownLatch(maximumConcurrentProcessing);
             for (var i = 0; i < maximumConcurrentProcessing; i++)
             {
                 if (stopBatchCancellationSource.IsCancellationRequested)
                 {
-                    shouldWaitForReceiveTasks = false;
                     break;
                 }
 
@@ -222,11 +220,8 @@ async Task ReceiveMessages(CancellationToken messageReceivingCancellationToken)
                     localConcurrencyLimiter, receiveLatch, messageProcessingCancellationTokenSource.Token);
             }
 
-            if (shouldWaitForReceiveTasks)
-            {
-                // Wait for all receive operations to complete before returning (and thus peeking again)
-                await receiveLatch.WaitAsync(CancellationToken.None).ConfigureAwait(false);
-            }
+            // Wait for all receive operations to complete before returning (and thus peeking again)
+            await receiveLatch.WaitAsync(stopBatchCancellationSource.Token).ConfigureAwait(false);
         }
 
         async Task ProcessMessagesSwallowExceptionsAndReleaseConcurrencyLimiter(