Replace syncronous countdown event with an async countdown latch based on TCS

xwipeoutx · mauroservienti · commit 844553a948d0 · 2025-11-28T13:51:06.000+01:00
[wip] initial cut of receiver, using tasks

Change from tasks to countdown event
diff --git a/src/NServiceBus.Transport.Sql.Shared/Receiving/AsyncCountdownLatch.cs b/src/NServiceBus.Transport.Sql.Shared/Receiving/AsyncCountdownLatch.cs
@@ -0,0 +1,31 @@
+﻿namespace NServiceBus.Transport.Sql.Shared;
+
+using System.Threading;
+using System.Threading.Tasks;
+
+class AsyncCountdownLatch
+{
+    int count;
+    readonly TaskCompletionSource completionSource;
+
+    public AsyncCountdownLatch(int count)
+    {
+        this.count = count;
+        completionSource = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        if (count <= 0)
+        {
+            completionSource.SetResult();
+        }
+    }
+
+    public Task WaitAsync() => completionSource.Task;
+
+    public void Signal()
+    {
+        if (Interlocked.Decrement(ref count) == 0)
+        {
+            completionSource.SetResult();
+        }
+    }
+}
diff --git a/src/NServiceBus.Transport.Sql.Shared/Receiving/MessageReceiver.cs b/src/NServiceBus.Transport.Sql.Shared/Receiving/MessageReceiver.cs
@@ -204,10 +204,13 @@ 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;
                 }
 
@@ -216,13 +219,19 @@ async Task ReceiveMessages(CancellationToken messageReceivingCancellationToken)
                 await localConcurrencyLimiter.WaitAsync(messageReceivingCancellationToken).ConfigureAwait(false);
 
                 _ = ProcessMessagesSwallowExceptionsAndReleaseConcurrencyLimiter(stopBatchCancellationSource,
-                    localConcurrencyLimiter, messageProcessingCancellationTokenSource.Token);
+                    localConcurrencyLimiter, receiveLatch, messageProcessingCancellationTokenSource.Token);
+            }
+
+            if (shouldWaitForReceiveTasks)
+            {
+                // Wait for all receive operations to complete before returning (and thus peeking again)
+                await receiveLatch.WaitAsync().ConfigureAwait(false);
             }
         }
 
         async Task ProcessMessagesSwallowExceptionsAndReleaseConcurrencyLimiter(
             CancellationTokenSource stopBatchCancellationTokenSource, SemaphoreSlim localConcurrencyLimiter,
-            CancellationToken messageProcessingCancellationToken)
+            AsyncCountdownLatch receiveLatch, CancellationToken messageProcessingCancellationToken)
         {
             try
             {
@@ -232,8 +241,8 @@ async Task ProcessMessagesSwallowExceptionsAndReleaseConcurrencyLimiter(
                     // in combination with TransactionScope will apply connection pooling and enlistment synchronous in ctor.
                     await Task.Yield();
 
-                    await processStrategy.ProcessMessage(stopBatchCancellationTokenSource,
-                            messageProcessingCancellationToken)
+                    await processStrategy.ProcessMessage(stopBatchCancellationTokenSource, receiveLatch,
+                        messageProcessingCancellationToken)
                         .ConfigureAwait(false);
 
                     messageProcessingCircuitBreaker.Success();
diff --git a/src/NServiceBus.Transport.Sql.Shared/Receiving/ProcessStrategy.cs b/src/NServiceBus.Transport.Sql.Shared/Receiving/ProcessStrategy.cs
@@ -35,7 +35,8 @@ public void Init(TableBasedQueue inputQueue, TableBasedQueue errorQueue, OnMessa
             this.criticalError = criticalError;
         }
 
-        public abstract Task ProcessMessage(CancellationTokenSource stopBatchCancellationTokenSource, CancellationToken cancellationToken = default);
+        public abstract Task ProcessMessage(CancellationTokenSource stopBatchCancellationTokenSource,
+            AsyncCountdownLatch receiveLatch, CancellationToken cancellationToken = default);
 
         protected async Task<bool> TryHandleMessage(Message message, TransportTransaction transportTransaction, ContextBag context, CancellationToken cancellationToken = default)
         {
diff --git a/src/NServiceBus.Transport.Sql.Shared/Receiving/ProcessWithNativeTransaction.cs b/src/NServiceBus.Transport.Sql.Shared/Receiving/ProcessWithNativeTransaction.cs
@@ -21,7 +21,8 @@ public ProcessWithNativeTransaction(TransactionOptions transactionOptions, DbCon
             isolationLevel = IsolationLevelMapper.Map(transactionOptions.IsolationLevel);
         }
 
-        public override async Task ProcessMessage(CancellationTokenSource stopBatchCancellationTokenSource, CancellationToken cancellationToken = default)
+        public override async Task ProcessMessage(CancellationTokenSource stopBatchCancellationTokenSource,
+            AsyncCountdownLatch receiveLatch, CancellationToken cancellationToken = default)
         {
             Message message = null;
             var context = new ContextBag();
@@ -31,7 +32,16 @@ public override async Task ProcessMessage(CancellationTokenSource stopBatchCance
                 using (var connection = await connectionFactory.OpenNewConnection(cancellationToken).ConfigureAwait(false))
                 using (var transaction = connection.BeginTransaction(isolationLevel))
                 {
-                    var receiveResult = await InputQueue.TryReceive(connection, transaction, cancellationToken).ConfigureAwait(false);
+                    MessageReadResult receiveResult;
+                    try
+                    {
+                        receiveResult = await InputQueue.TryReceive(connection, transaction, cancellationToken)
+                            .ConfigureAwait(false);
+                    }
+                    finally
+                    {
+                        receiveLatch.Signal();
+                    }
 
                     if (receiveResult == MessageReadResult.NoMessage)
                     {
diff --git a/src/NServiceBus.Transport.Sql.Shared/Receiving/ProcessWithNoTransaction.cs b/src/NServiceBus.Transport.Sql.Shared/Receiving/ProcessWithNoTransaction.cs
@@ -16,7 +16,8 @@ public ProcessWithNoTransaction(DbConnectionFactory connectionFactory, FailureIn
             this.exceptionClassifier = exceptionClassifier;
         }
 
-        public override async Task ProcessMessage(CancellationTokenSource stopBatchCancellationTokenSource, CancellationToken cancellationToken = default)
+        public override async Task ProcessMessage(CancellationTokenSource stopBatchCancellationTokenSource,
+            AsyncCountdownLatch receiveLatch, CancellationToken cancellationToken = default)
         {
             Message message = null;
             var context = new ContextBag();
@@ -27,8 +28,16 @@ public override async Task ProcessMessage(CancellationTokenSource stopBatchCance
                 {
                     using (var transaction = connection.BeginTransaction(IsolationLevel.ReadCommitted))
                     {
-                        var receiveResult = await InputQueue.TryReceive(connection, transaction, cancellationToken)
-                            .ConfigureAwait(false);
+                        MessageReadResult receiveResult;
+                        try
+                        {
+                            receiveResult = await InputQueue.TryReceive(connection, transaction, cancellationToken)
+                                .ConfigureAwait(false);
+                        }
+                        finally
+                        {
+                            receiveLatch.Signal();
+                        }
 
                         if (receiveResult == MessageReadResult.NoMessage)
                         {
diff --git a/src/NServiceBus.Transport.Sql.Shared/Receiving/ProcessWithTransactionScope.cs b/src/NServiceBus.Transport.Sql.Shared/Receiving/ProcessWithTransactionScope.cs
@@ -17,7 +17,8 @@ public ProcessWithTransactionScope(TransactionOptions transactionOptions, DbConn
             this.exceptionClassifier = exceptionClassifier;
         }
 
-        public override async Task ProcessMessage(CancellationTokenSource stopBatchCancellationTokenSource, CancellationToken cancellationToken = default)
+        public override async Task ProcessMessage(CancellationTokenSource stopBatchCancellationTokenSource,
+            AsyncCountdownLatch receiveLatch, CancellationToken cancellationToken = default)
         {
             Message message = null;
             var context = new ContextBag();
@@ -27,7 +28,15 @@ public override async Task ProcessMessage(CancellationTokenSource stopBatchCance
                 using (var scope = new TransactionScope(TransactionScopeOption.RequiresNew, transactionOptions, TransactionScopeAsyncFlowOption.Enabled))
                 using (var connection = await connectionFactory.OpenNewConnection(cancellationToken).ConfigureAwait(false))
                 {
-                    var receiveResult = await InputQueue.TryReceive(connection, null, cancellationToken).ConfigureAwait(false);
+                    MessageReadResult receiveResult;
+                    try
+                    {
+                        receiveResult = await InputQueue.TryReceive(connection, null, cancellationToken).ConfigureAwait(false);
+                    }
+                    finally
+                    {
+                        receiveLatch.Signal();
+                    }
 
                     if (receiveResult == MessageReadResult.NoMessage)
                     {

Original file line number	Diff line number	Diff line change
`@@ -204,10 +204,13 @@ async Task ReceiveMessages(CancellationToken messageReceivingCancellationToken)`
`204`	`204`	`? 1`
`205`	`205`	`: messageCount;`
`206`	`206`
	`207`	`+ bool shouldWaitForReceiveTasks = true;`
	`208`	`+ var receiveLatch = new AsyncCountdownLatch(maximumConcurrentProcessing);`
`207`	`209`	`for (var i = 0; i < maximumConcurrentProcessing; i++)`
`208`	`210`	`{`
`209`	`211`	`if (stopBatchCancellationSource.IsCancellationRequested)`
`210`	`212`	`{`
	`213`	`+ shouldWaitForReceiveTasks = false;`
`211`	`214`	`break;`
`212`	`215`	`}`
`213`	`216`
`@@ -216,13 +219,19 @@ async Task ReceiveMessages(CancellationToken messageReceivingCancellationToken)`
`216`	`219`	`await localConcurrencyLimiter.WaitAsync(messageReceivingCancellationToken).ConfigureAwait(false);`
`217`	`220`
`218`	`221`	`_ = ProcessMessagesSwallowExceptionsAndReleaseConcurrencyLimiter(stopBatchCancellationSource,`
`219`		`- localConcurrencyLimiter, messageProcessingCancellationTokenSource.Token);`
	`222`	`+ localConcurrencyLimiter, receiveLatch, messageProcessingCancellationTokenSource.Token);`
	`223`	`+ }`
	`224`	`+`
	`225`	`+ if (shouldWaitForReceiveTasks)`
	`226`	`+ {`
	`227`	`+ // Wait for all receive operations to complete before returning (and thus peeking again)`
	`228`	`+ await receiveLatch.WaitAsync().ConfigureAwait(false);`
`220`	`229`	`}`
`221`	`230`	`}`
`222`	`231`
`223`	`232`	`async Task ProcessMessagesSwallowExceptionsAndReleaseConcurrencyLimiter(`
`224`	`233`	`CancellationTokenSource stopBatchCancellationTokenSource, SemaphoreSlim localConcurrencyLimiter,`
`225`		`- CancellationToken messageProcessingCancellationToken)`
	`234`	`+ AsyncCountdownLatch receiveLatch, CancellationToken messageProcessingCancellationToken)`
`226`	`235`	`{`
`227`	`236`	`try`
`228`	`237`	`{`
`@@ -232,8 +241,8 @@ async Task ProcessMessagesSwallowExceptionsAndReleaseConcurrencyLimiter(`
`232`	`241`	`// in combination with TransactionScope will apply connection pooling and enlistment synchronous in ctor.`
`233`	`242`	`await Task.Yield();`
`234`	`243`
`235`		`- await processStrategy.ProcessMessage(stopBatchCancellationTokenSource,`
`236`		`- messageProcessingCancellationToken)`
	`244`	`+ await processStrategy.ProcessMessage(stopBatchCancellationTokenSource, receiveLatch,`
	`245`	`+ messageProcessingCancellationToken)`
`237`	`246`	`.ConfigureAwait(false);`
`238`	`247`
`239`	`248`	`messageProcessingCircuitBreaker.Success();`
Original file line number	Diff line number	Diff line change
`@@ -35,7 +35,8 @@ public void Init(TableBasedQueue inputQueue, TableBasedQueue errorQueue, OnMessa`
`35`	`35`	`this.criticalError = criticalError;`
`36`	`36`	`}`
`37`	`37`
`38`		`- public abstract Task ProcessMessage(CancellationTokenSource stopBatchCancellationTokenSource, CancellationToken cancellationToken = default);`
	`38`	`+ public abstract Task ProcessMessage(CancellationTokenSource stopBatchCancellationTokenSource,`
	`39`	`+ AsyncCountdownLatch receiveLatch, CancellationToken cancellationToken = default);`
`39`	`40`
`40`	`41`	`protected async Task<bool> TryHandleMessage(Message message, TransportTransaction transportTransaction, ContextBag context, CancellationToken cancellationToken = default)`
`41`	`42`	`{`
Original file line number	Diff line number	Diff line change
`@@ -16,7 +16,8 @@ public ProcessWithNoTransaction(DbConnectionFactory connectionFactory, FailureIn`
`16`	`16`	`this.exceptionClassifier = exceptionClassifier;`
`17`	`17`	`}`
`18`	`18`
`19`		`- public override async Task ProcessMessage(CancellationTokenSource stopBatchCancellationTokenSource, CancellationToken cancellationToken = default)`
	`19`	`+ public override async Task ProcessMessage(CancellationTokenSource stopBatchCancellationTokenSource,`
	`20`	`+ AsyncCountdownLatch receiveLatch, CancellationToken cancellationToken = default)`
`20`	`21`	`{`
`21`	`22`	`Message message = null;`
`22`	`23`	`var context = new ContextBag();`
`@@ -27,8 +28,16 @@ public override async Task ProcessMessage(CancellationTokenSource stopBatchCance`
`27`	`28`	`{`
`28`	`29`	`using (var transaction = connection.BeginTransaction(IsolationLevel.ReadCommitted))`
`29`	`30`	`{`
`30`		`- var receiveResult = await InputQueue.TryReceive(connection, transaction, cancellationToken)`
`31`		`- .ConfigureAwait(false);`
	`31`	`+ MessageReadResult receiveResult;`
	`32`	`+ try`
	`33`	`+ {`
	`34`	`+ receiveResult = await InputQueue.TryReceive(connection, transaction, cancellationToken)`
	`35`	`+ .ConfigureAwait(false);`
	`36`	`+ }`
	`37`	`+ finally`
	`38`	`+ {`
	`39`	`+ receiveLatch.Signal();`
	`40`	`+ }`
`32`	`41`
`33`	`42`	`if (receiveResult == MessageReadResult.NoMessage)`
`34`	`43`	`{`
Original file line number	Diff line number	Diff line change
`@@ -17,7 +17,8 @@ public ProcessWithTransactionScope(TransactionOptions transactionOptions, DbConn`
`17`	`17`	`this.exceptionClassifier = exceptionClassifier;`
`18`	`18`	`}`
`19`	`19`
`20`		`- public override async Task ProcessMessage(CancellationTokenSource stopBatchCancellationTokenSource, CancellationToken cancellationToken = default)`
	`20`	`+ public override async Task ProcessMessage(CancellationTokenSource stopBatchCancellationTokenSource,`
	`21`	`+ AsyncCountdownLatch receiveLatch, CancellationToken cancellationToken = default)`
`21`	`22`	`{`
`22`	`23`	`Message message = null;`
`23`	`24`	`var context = new ContextBag();`
`@@ -27,7 +28,15 @@ public override async Task ProcessMessage(CancellationTokenSource stopBatchCance`
`27`	`28`	`using (var scope = new TransactionScope(TransactionScopeOption.RequiresNew, transactionOptions, TransactionScopeAsyncFlowOption.Enabled))`
`28`	`29`	`using (var connection = await connectionFactory.OpenNewConnection(cancellationToken).ConfigureAwait(false))`
`29`	`30`	`{`
`30`		`- var receiveResult = await InputQueue.TryReceive(connection, null, cancellationToken).ConfigureAwait(false);`
	`31`	`+ MessageReadResult receiveResult;`
	`32`	`+ try`
	`33`	`+ {`
	`34`	`+ receiveResult = await InputQueue.TryReceive(connection, null, cancellationToken).ConfigureAwait(false);`
	`35`	`+ }`
	`36`	`+ finally`
	`37`	`+ {`
	`38`	`+ receiveLatch.Signal();`
	`39`	`+ }`
`31`	`40`
`32`	`41`	`if (receiveResult == MessageReadResult.NoMessage)`
`33`	`42`	`{`