Refactor to a disposable signaler class

SzymonPobiega · SzymonPobiega · commit 1b6cf4e9a560 · 2025-12-10T09:27:45.000+01:00
diff --git a/src/NServiceBus.Transport.Sql.Shared/Receiving/AsyncCountdownLatch.cs b/src/NServiceBus.Transport.Sql.Shared/Receiving/AsyncCountdownLatch.cs
@@ -1,5 +1,6 @@
 ﻿namespace NServiceBus.Transport.Sql.Shared;
 
+using System;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -21,11 +22,33 @@ public AsyncCountdownLatch(int count)
 
     public Task WaitAsync(CancellationToken cancellationToken = default) => completionSource.Task;
 
-    public void Signal()
+    public Signaler GetSignaler() => new Signaler(this);
+
+    void Signal()
     {
         if (Interlocked.Decrement(ref count) == 0)
         {
             completionSource.SetResult();
         }
     }
+
+    public class Signaler(AsyncCountdownLatch parent) : IDisposable
+    {
+        bool signalled;
+
+        public void Signal()
+        {
+            parent.Signal();
+            signalled = true;
+        }
+
+        public void Dispose()
+        {
+            if (!signalled)
+            {
+                parent.Signal();
+                signalled = true;
+            }
+        }
+    }
 }
diff --git a/src/NServiceBus.Transport.Sql.Shared/Receiving/MessageReceiver.cs b/src/NServiceBus.Transport.Sql.Shared/Receiving/MessageReceiver.cs
@@ -233,6 +233,7 @@ async Task ProcessMessagesSwallowExceptionsAndReleaseConcurrencyLimiter(
             CancellationTokenSource stopBatchCancellationTokenSource, SemaphoreSlim localConcurrencyLimiter,
             AsyncCountdownLatch receiveLatch, CancellationToken messageProcessingCancellationToken)
         {
+            using var latchSignaler = receiveLatch.GetSignaler();
             try
             {
                 try
@@ -241,7 +242,7 @@ async Task ProcessMessagesSwallowExceptionsAndReleaseConcurrencyLimiter(
                     // in combination with TransactionScope will apply connection pooling and enlistment synchronous in ctor.
                     await Task.Yield();
 
-                    await processStrategy.ProcessMessage(stopBatchCancellationTokenSource, receiveLatch,
+                    await processStrategy.ProcessMessage(stopBatchCancellationTokenSource, latchSignaler,
                         messageProcessingCancellationToken)
                         .ConfigureAwait(false);
 
diff --git a/src/NServiceBus.Transport.Sql.Shared/Receiving/ProcessStrategy.cs b/src/NServiceBus.Transport.Sql.Shared/Receiving/ProcessStrategy.cs
@@ -36,7 +36,7 @@ public void Init(TableBasedQueue inputQueue, TableBasedQueue errorQueue, OnMessa
         }
 
         public abstract Task ProcessMessage(CancellationTokenSource stopBatchCancellationTokenSource,
-            AsyncCountdownLatch receiveLatch, CancellationToken cancellationToken = default);
+            AsyncCountdownLatch.Signaler 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
@@ -22,28 +22,18 @@ public ProcessWithNativeTransaction(TransactionOptions transactionOptions, DbCon
         }
 
         public override async Task ProcessMessage(CancellationTokenSource stopBatchCancellationTokenSource,
-            AsyncCountdownLatch receiveLatch, CancellationToken cancellationToken = default)
+            AsyncCountdownLatch.Signaler receiveLatch, CancellationToken cancellationToken = default)
         {
             Message message = null;
             var context = new ContextBag();
-            var hasLatchBeenSignalled = false;
 
             try
             {
                 using (var connection = await connectionFactory.OpenNewConnection(cancellationToken).ConfigureAwait(false))
                 using (var transaction = connection.BeginTransaction(isolationLevel))
                 {
-                    MessageReadResult receiveResult;
-                    try
-                    {
-                        receiveResult = await InputQueue.TryReceive(connection, transaction, cancellationToken)
-                            .ConfigureAwait(false);
-                    }
-                    finally
-                    {
-                        receiveLatch.Signal();
-                        hasLatchBeenSignalled = true;
-                    }
+                    var receiveResult = await InputQueue.TryReceive(connection, transaction, cancellationToken).ConfigureAwait(false);
+                    receiveLatch.Signal();
 
                     if (receiveResult == MessageReadResult.NoMessage)
                     {
@@ -89,13 +79,6 @@ public override async Task ProcessMessage(CancellationTokenSource stopBatchCance
                 }
                 failureInfoStorage.RecordFailureInfoForMessage(message.TransportId, ex, context);
             }
-            finally
-            {
-                if (!hasLatchBeenSignalled)
-                {
-                    receiveLatch.Signal();
-                }
-            }
         }
 
         async Task<bool> TryProcess(Message message, TransportTransaction transportTransaction, ContextBag context, CancellationToken cancellationToken)
diff --git a/src/NServiceBus.Transport.Sql.Shared/Receiving/ProcessWithNoTransaction.cs b/src/NServiceBus.Transport.Sql.Shared/Receiving/ProcessWithNoTransaction.cs
@@ -17,29 +17,19 @@ public ProcessWithNoTransaction(DbConnectionFactory connectionFactory, FailureIn
         }
 
         public override async Task ProcessMessage(CancellationTokenSource stopBatchCancellationTokenSource,
-            AsyncCountdownLatch receiveLatch, CancellationToken cancellationToken = default)
+            AsyncCountdownLatch.Signaler receiveLatch, CancellationToken cancellationToken = default)
         {
             Message message = null;
             var context = new ContextBag();
-            var hasLatchBeenSignalled = false;
 
             using (var connection = await connectionFactory.OpenNewConnection(cancellationToken).ConfigureAwait(false))
             {
                 try
                 {
                     using (var transaction = connection.BeginTransaction(IsolationLevel.ReadCommitted))
                     {
-                        MessageReadResult receiveResult;
-                        try
-                        {
-                            receiveResult = await InputQueue.TryReceive(connection, transaction, cancellationToken)
-                                .ConfigureAwait(false);
-                        }
-                        finally
-                        {
-                            receiveLatch.Signal();
-                            hasLatchBeenSignalled = true;
-                        }
+                        var receiveResult = await InputQueue.TryReceive(connection, transaction, cancellationToken).ConfigureAwait(false);
+                        receiveLatch.Signal();
 
                         if (receiveResult == MessageReadResult.NoMessage)
                         {
@@ -89,13 +79,6 @@ await ErrorQueue
                     // Since this is TransactionMode.None, we don't care whether error handling says handled or retry. Message is gone either way.
                     _ = await HandleError(ex, message, transportTransaction, 1, context, cancellationToken).ConfigureAwait(false);
                 }
-                finally
-                {
-                    if (!hasLatchBeenSignalled)
-                    {
-                        receiveLatch.Signal();
-                    }
-                }
             }
         }
 
diff --git a/src/NServiceBus.Transport.Sql.Shared/Receiving/ProcessWithTransactionScope.cs b/src/NServiceBus.Transport.Sql.Shared/Receiving/ProcessWithTransactionScope.cs
@@ -18,27 +18,18 @@ public ProcessWithTransactionScope(TransactionOptions transactionOptions, DbConn
         }
 
         public override async Task ProcessMessage(CancellationTokenSource stopBatchCancellationTokenSource,
-            AsyncCountdownLatch receiveLatch, CancellationToken cancellationToken = default)
+            AsyncCountdownLatch.Signaler receiveLatch, CancellationToken cancellationToken = default)
         {
             Message message = null;
             var context = new ContextBag();
-            var hasLatchBeenSignalled = false;
 
             try
             {
                 using (var scope = new TransactionScope(TransactionScopeOption.RequiresNew, transactionOptions, TransactionScopeAsyncFlowOption.Enabled))
                 using (var connection = await connectionFactory.OpenNewConnection(cancellationToken).ConfigureAwait(false))
                 {
-                    MessageReadResult receiveResult;
-                    try
-                    {
-                        receiveResult = await InputQueue.TryReceive(connection, null, cancellationToken).ConfigureAwait(false);
-                    }
-                    finally
-                    {
-                        receiveLatch.Signal();
-                        hasLatchBeenSignalled = true;
-                    }
+                    var receiveResult = await InputQueue.TryReceive(connection, null, cancellationToken).ConfigureAwait(false);
+                    receiveLatch.Signal();
 
                     if (receiveResult == MessageReadResult.NoMessage)
                     {
@@ -81,13 +72,6 @@ public override async Task ProcessMessage(CancellationTokenSource stopBatchCance
                 }
                 failureInfoStorage.RecordFailureInfoForMessage(message.TransportId, ex, context);
             }
-            finally
-            {
-                if (!hasLatchBeenSignalled)
-                {
-                    receiveLatch.Signal();
-                }
-            }
         }
 
         async Task<bool> TryProcess(Message message, TransportTransaction transportTransaction, ContextBag context, CancellationToken cancellationToken)

Original file line number	Diff line number	Diff line change
`@@ -233,6 +233,7 @@ async Task ProcessMessagesSwallowExceptionsAndReleaseConcurrencyLimiter(`
`233`	`233`	`CancellationTokenSource stopBatchCancellationTokenSource, SemaphoreSlim localConcurrencyLimiter,`
`234`	`234`	`AsyncCountdownLatch receiveLatch, CancellationToken messageProcessingCancellationToken)`
`235`	`235`	`{`
	`236`	`+ using var latchSignaler = receiveLatch.GetSignaler();`
`236`	`237`	`try`
`237`	`238`	`{`
`238`	`239`	`try`
`@@ -241,7 +242,7 @@ async Task ProcessMessagesSwallowExceptionsAndReleaseConcurrencyLimiter(`
`241`	`242`	`// in combination with TransactionScope will apply connection pooling and enlistment synchronous in ctor.`
`242`	`243`	`await Task.Yield();`
`243`	`244`
`244`		`- await processStrategy.ProcessMessage(stopBatchCancellationTokenSource, receiveLatch,`
	`245`	`+ await processStrategy.ProcessMessage(stopBatchCancellationTokenSource, latchSignaler,`
`245`	`246`	`messageProcessingCancellationToken)`
`246`	`247`	`.ConfigureAwait(false);`
`247`	`248`
Original file line number	Diff line number	Diff line change
`@@ -36,7 +36,7 @@ public void Init(TableBasedQueue inputQueue, TableBasedQueue errorQueue, OnMessa`
`36`	`36`	`}`
`37`	`37`
`38`	`38`	`public abstract Task ProcessMessage(CancellationTokenSource stopBatchCancellationTokenSource,`
`39`		`- AsyncCountdownLatch receiveLatch, CancellationToken cancellationToken = default);`
	`39`	`+ AsyncCountdownLatch.Signaler receiveLatch, CancellationToken cancellationToken = default);`
`40`	`40`
`41`	`41`	`protected async Task<bool> TryHandleMessage(Message message, TransportTransaction transportTransaction, ContextBag context, CancellationToken cancellationToken = default)`
`42`	`42`	`{`
Original file line number	Diff line number	Diff line change
`@@ -22,28 +22,18 @@ public ProcessWithNativeTransaction(TransactionOptions transactionOptions, DbCon`
`22`	`22`	`}`
`23`	`23`
`24`	`24`	`public override async Task ProcessMessage(CancellationTokenSource stopBatchCancellationTokenSource,`
`25`		`- AsyncCountdownLatch receiveLatch, CancellationToken cancellationToken = default)`
	`25`	`+ AsyncCountdownLatch.Signaler receiveLatch, CancellationToken cancellationToken = default)`
`26`	`26`	`{`
`27`	`27`	`Message message = null;`
`28`	`28`	`var context = new ContextBag();`
`29`		`- var hasLatchBeenSignalled = false;`
`30`	`29`
`31`	`30`	`try`
`32`	`31`	`{`
`33`	`32`	`using (var connection = await connectionFactory.OpenNewConnection(cancellationToken).ConfigureAwait(false))`
`34`	`33`	`using (var transaction = connection.BeginTransaction(isolationLevel))`
`35`	`34`	`{`
`36`		`- MessageReadResult receiveResult;`
`37`		`- try`
`38`		`- {`
`39`		`- receiveResult = await InputQueue.TryReceive(connection, transaction, cancellationToken)`
`40`		`- .ConfigureAwait(false);`
`41`		`- }`
`42`		`- finally`
`43`		`- {`
`44`		`- receiveLatch.Signal();`
`45`		`- hasLatchBeenSignalled = true;`
`46`		`- }`
	`35`	`+ var receiveResult = await InputQueue.TryReceive(connection, transaction, cancellationToken).ConfigureAwait(false);`
	`36`	`+ receiveLatch.Signal();`
`47`	`37`
`48`	`38`	`if (receiveResult == MessageReadResult.NoMessage)`
`49`	`39`	`{`
`@@ -89,13 +79,6 @@ public override async Task ProcessMessage(CancellationTokenSource stopBatchCance`
`89`	`79`	`}`
`90`	`80`	`failureInfoStorage.RecordFailureInfoForMessage(message.TransportId, ex, context);`
`91`	`81`	`}`
`92`		`- finally`
`93`		`- {`
`94`		`- if (!hasLatchBeenSignalled)`
`95`		`- {`
`96`		`- receiveLatch.Signal();`
`97`		`- }`
`98`		`- }`
`99`	`82`	`}`
`100`	`83`
`101`	`84`	`async Task<bool> TryProcess(Message message, TransportTransaction transportTransaction, ContextBag context, CancellationToken cancellationToken)`
Original file line number	Diff line number	Diff line change
`@@ -17,29 +17,19 @@ public ProcessWithNoTransaction(DbConnectionFactory connectionFactory, FailureIn`
`17`	`17`	`}`
`18`	`18`
`19`	`19`	`public override async Task ProcessMessage(CancellationTokenSource stopBatchCancellationTokenSource,`
`20`		`- AsyncCountdownLatch receiveLatch, CancellationToken cancellationToken = default)`
	`20`	`+ AsyncCountdownLatch.Signaler receiveLatch, CancellationToken cancellationToken = default)`
`21`	`21`	`{`
`22`	`22`	`Message message = null;`
`23`	`23`	`var context = new ContextBag();`
`24`		`- var hasLatchBeenSignalled = false;`
`25`	`24`
`26`	`25`	`using (var connection = await connectionFactory.OpenNewConnection(cancellationToken).ConfigureAwait(false))`
`27`	`26`	`{`
`28`	`27`	`try`
`29`	`28`	`{`
`30`	`29`	`using (var transaction = connection.BeginTransaction(IsolationLevel.ReadCommitted))`
`31`	`30`	`{`
`32`		`- MessageReadResult receiveResult;`
`33`		`- try`
`34`		`- {`
`35`		`- receiveResult = await InputQueue.TryReceive(connection, transaction, cancellationToken)`
`36`		`- .ConfigureAwait(false);`
`37`		`- }`
`38`		`- finally`
`39`		`- {`
`40`		`- receiveLatch.Signal();`
`41`		`- hasLatchBeenSignalled = true;`
`42`		`- }`
	`31`	`+ var receiveResult = await InputQueue.TryReceive(connection, transaction, cancellationToken).ConfigureAwait(false);`
	`32`	`+ receiveLatch.Signal();`
`43`	`33`
`44`	`34`	`if (receiveResult == MessageReadResult.NoMessage)`
`45`	`35`	`{`
`@@ -89,13 +79,6 @@ await ErrorQueue`
`89`	`79`	`// Since this is TransactionMode.None, we don't care whether error handling says handled or retry. Message is gone either way.`
`90`	`80`	`_ = await HandleError(ex, message, transportTransaction, 1, context, cancellationToken).ConfigureAwait(false);`
`91`	`81`	`}`
`92`		`- finally`
`93`		`- {`
`94`		`- if (!hasLatchBeenSignalled)`
`95`		`- {`
`96`		`- receiveLatch.Signal();`
`97`		`- }`
`98`		`- }`
`99`	`82`	`}`
`100`	`83`	`}`
`101`	`84`
Original file line number	Diff line number	Diff line change
`@@ -18,27 +18,18 @@ public ProcessWithTransactionScope(TransactionOptions transactionOptions, DbConn`
`18`	`18`	`}`
`19`	`19`
`20`	`20`	`public override async Task ProcessMessage(CancellationTokenSource stopBatchCancellationTokenSource,`
`21`		`- AsyncCountdownLatch receiveLatch, CancellationToken cancellationToken = default)`
	`21`	`+ AsyncCountdownLatch.Signaler receiveLatch, CancellationToken cancellationToken = default)`
`22`	`22`	`{`
`23`	`23`	`Message message = null;`
`24`	`24`	`var context = new ContextBag();`
`25`		`- var hasLatchBeenSignalled = false;`
`26`	`25`
`27`	`26`	`try`
`28`	`27`	`{`
`29`	`28`	`using (var scope = new TransactionScope(TransactionScopeOption.RequiresNew, transactionOptions, TransactionScopeAsyncFlowOption.Enabled))`
`30`	`29`	`using (var connection = await connectionFactory.OpenNewConnection(cancellationToken).ConfigureAwait(false))`
`31`	`30`	`{`
`32`		`- MessageReadResult receiveResult;`
`33`		`- try`
`34`		`- {`
`35`		`- receiveResult = await InputQueue.TryReceive(connection, null, cancellationToken).ConfigureAwait(false);`
`36`		`- }`
`37`		`- finally`
`38`		`- {`
`39`		`- receiveLatch.Signal();`
`40`		`- hasLatchBeenSignalled = true;`
`41`		`- }`
	`31`	`+ var receiveResult = await InputQueue.TryReceive(connection, null, cancellationToken).ConfigureAwait(false);`
	`32`	`+ receiveLatch.Signal();`
`42`	`33`
`43`	`34`	`if (receiveResult == MessageReadResult.NoMessage)`
`44`	`35`	`{`
`@@ -81,13 +72,6 @@ public override async Task ProcessMessage(CancellationTokenSource stopBatchCance`
`81`	`72`	`}`
`82`	`73`	`failureInfoStorage.RecordFailureInfoForMessage(message.TransportId, ex, context);`
`83`	`74`	`}`
`84`		`- finally`
`85`		`- {`
`86`		`- if (!hasLatchBeenSignalled)`
`87`		`- {`
`88`		`- receiveLatch.Signal();`
`89`		`- }`
`90`		`- }`
`91`	`75`	`}`
`92`	`76`
`93`	`77`	`async Task<bool> TryProcess(Message message, TransportTransaction transportTransaction, ContextBag context, CancellationToken cancellationToken)`