Batch receive fixes (Azure#21566)

* Batch receive fixes

* Use AwaitWithCancellation

* volatile

* Remove unnecessary catch

* Fix tests

* Fix flaky tests

* Fix test

Author: JoshLove-msft

sdk/servicebus/Microsoft.Azure.WebJobs.Extensions.ServiceBus/src/Listeners/ServiceBusListener.cs

@@ -14,7 +14,8 @@
     <PackageReference Include="Microsoft.Azure.WebJobs.Sources" />
     <PackageReference Include="Microsoft.Azure.WebJobs" />
     <PackageReference Include="Microsoft.Extensions.Azure" />
-    <PackageReference Include="Azure.Messaging.ServiceBus" />
+      <!--Revert after Service Bus release-->
+      <!--<PackageReference Include="Azure.Messaging.ServiceBus" />-->
   </ItemGroup>
 
   <ItemGroup>
@@ -28,4 +29,9 @@
     <Compile Include="..\..\..\extensions\Microsoft.Azure.WebJobs.Extensions.Clients\src\Shared\WebJobsConfigurationExtensions.cs" LinkBase="Shared" />
 
   </ItemGroup>
+
+  <ItemGroup>
+    <!--Revert after Service Bus release-->
+    <ProjectReference Include="..\..\Azure.Messaging.ServiceBus\src\Azure.Messaging.ServiceBus.csproj" />
+  </ItemGroup>
 </Project>

sdk/servicebus/Microsoft.Azure.WebJobs.Extensions.ServiceBus/src/Microsoft.Azure.WebJobs.Extensions.ServiceBus.csproj

@@ -17,6 +17,8 @@ public class ServiceBusMessageActions
         private readonly ProcessMessageEventArgs _eventArgs;
         private readonly ProcessSessionMessageEventArgs _sessionEventArgs;
 
+        internal HashSet<ServiceBusReceivedMessage> SettledMessages { get; } = new();
+
         internal ServiceBusMessageActions(ProcessSessionMessageEventArgs sessionEventArgs)
         {
             _sessionEventArgs = sessionEventArgs;
@@ -50,6 +52,8 @@ public virtual async Task AbandonMessageAsync(
             {
                 await _sessionEventArgs.AbandonMessageAsync(message, propertiesToModify, cancellationToken).ConfigureAwait(false);
             }
+
+            SettledMessages.Add(message);
         }
 
         ///<inheritdoc cref="ServiceBusReceiver.CompleteMessageAsync(ServiceBusReceivedMessage, CancellationToken)"/>
@@ -69,6 +73,8 @@ public virtual async Task CompleteMessageAsync(
             {
                 await _sessionEventArgs.CompleteMessageAsync(message, cancellationToken).ConfigureAwait(false);
             }
+
+            SettledMessages.Add(message);
         }
 
         ///<inheritdoc cref="ServiceBusReceiver.DeadLetterMessageAsync(ServiceBusReceivedMessage, string, string, CancellationToken)"/>
@@ -105,6 +111,8 @@ await _sessionEventArgs.DeadLetterMessageAsync(
                     cancellationToken)
                 .ConfigureAwait(false);
             }
+
+            SettledMessages.Add(message);
         }
 
         ///<inheritdoc cref="ServiceBusReceiver.DeadLetterMessageAsync(ServiceBusReceivedMessage, IDictionary{string, object}, CancellationToken)"/>
@@ -137,6 +145,8 @@ await _sessionEventArgs.DeadLetterMessageAsync(
                     cancellationToken)
                 .ConfigureAwait(false);
             }
+
+            SettledMessages.Add(message);
         }
 
         ///<inheritdoc cref="ServiceBusReceiver.DeferMessageAsync(ServiceBusReceivedMessage, IDictionary{string, object}, CancellationToken)"/>
@@ -169,6 +179,8 @@ await _sessionEventArgs.DeferMessageAsync(
                     cancellationToken)
                 .ConfigureAwait(false);
             }
+
+            SettledMessages.Add(message);
         }
     }
 }

sdk/servicebus/Microsoft.Azure.WebJobs.Extensions.ServiceBus/src/Primitives/ServiceBusMessageActions.cs

@@ -101,7 +101,7 @@ private bool GetAutoCompleteMessagesOptionToUse(ServiceBusTriggerAttribute attri
         {
             if (attribute.IsAutoCompleteMessagesOptionSet)
             {
-                _logger.LogInformation($"The 'AutoCompleteMessages' option has been overrriden to '{attribute.AutoCompleteMessages}' value for '{functionName}' function.");
+                _logger.LogInformation($"The 'AutoCompleteMessages' option has been overriden to '{attribute.AutoCompleteMessages}' value for '{functionName}' function.");
 
                 return attribute.AutoCompleteMessages;
             }