Make ServiceBusMessageActions thread-safe (Azure#25556)

* Make ServiceBusMessageActions thread-safe

* PR fb

Author: JoshLove-msft

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

@@ -343,7 +343,7 @@ await receiver.ReceiveMessagesAsync(
                                 foreach (ServiceBusReceivedMessage message in messagesArray)
                                 {
                                     // skip messages that were settled in the user's function
-                                    if (input.MessageActions.SettledMessages.Contains(message))
+                                    if (input.MessageActions.SettledMessages.ContainsKey(message))
                                     {
                                         continue;
                                     }
@@ -360,7 +360,7 @@ await receiver.ReceiveMessagesAsync(
                                 foreach (ServiceBusReceivedMessage message in messagesArray)
                                 {
                                     // skip messages that were settled in the user's function
-                                    if (input.MessageActions.SettledMessages.Contains(message))
+                                    if (input.MessageActions.SettledMessages.ContainsKey(message))
                                     {
                                         continue;
                                     }

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

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+using System.Collections.Concurrent;
 using Azure.Messaging.ServiceBus;
 using System.Collections.Generic;
 using System.Threading;
@@ -17,7 +18,7 @@ public class ServiceBusMessageActions
         private readonly ProcessMessageEventArgs _eventArgs;
         private readonly ProcessSessionMessageEventArgs _sessionEventArgs;
 
-        internal HashSet<ServiceBusReceivedMessage> SettledMessages { get; } = new();
+        internal ConcurrentDictionary<ServiceBusReceivedMessage, byte> SettledMessages { get; } = new();
 
         internal ServiceBusMessageActions(ProcessSessionMessageEventArgs sessionEventArgs)
         {
@@ -64,7 +65,7 @@ public virtual async Task AbandonMessageAsync(
                 await _sessionEventArgs.AbandonMessageAsync(message, propertiesToModify, cancellationToken).ConfigureAwait(false);
             }
 
-            SettledMessages.Add(message);
+            TrackMessageAsSettled(message);
         }
 
         ///<inheritdoc cref="ServiceBusReceiver.CompleteMessageAsync(ServiceBusReceivedMessage, CancellationToken)"/>
@@ -85,7 +86,7 @@ public virtual async Task CompleteMessageAsync(
                 await _sessionEventArgs.CompleteMessageAsync(message, cancellationToken).ConfigureAwait(false);
             }
 
-            SettledMessages.Add(message);
+            TrackMessageAsSettled(message);
         }
 
         ///<inheritdoc cref="ServiceBusReceiver.DeadLetterMessageAsync(ServiceBusReceivedMessage, string, string, CancellationToken)"/>
@@ -123,7 +124,7 @@ await _sessionEventArgs.DeadLetterMessageAsync(
                 .ConfigureAwait(false);
             }
 
-            SettledMessages.Add(message);
+            TrackMessageAsSettled(message);
         }
 
         ///<inheritdoc cref="ServiceBusReceiver.DeadLetterMessageAsync(ServiceBusReceivedMessage, IDictionary{string, object}, CancellationToken)"/>
@@ -157,7 +158,7 @@ await _sessionEventArgs.DeadLetterMessageAsync(
                 .ConfigureAwait(false);
             }
 
-            SettledMessages.Add(message);
+            TrackMessageAsSettled(message);
         }
 
         ///<inheritdoc cref="ServiceBusReceiver.DeferMessageAsync(ServiceBusReceivedMessage, IDictionary{string, object}, CancellationToken)"/>
@@ -191,7 +192,10 @@ await _sessionEventArgs.DeferMessageAsync(
                 .ConfigureAwait(false);
             }
 
-            SettledMessages.Add(message);
+            TrackMessageAsSettled(message);
         }
+
+        private void TrackMessageAsSettled(ServiceBusReceivedMessage message)
+            => SettledMessages[message] = 0;
     }
 }

sdk/servicebus/Microsoft.Azure.WebJobs.Extensions.ServiceBus/tests/MessageActionsTests.cs

@@ -0,0 +1,33 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
+using Azure.Messaging.ServiceBus;
+using Microsoft.Azure.WebJobs.ServiceBus;
+using Moq;
+using NUnit.Framework;
+
+namespace Microsoft.Azure.WebJobs.Extensions.ServiceBus.Tests
+{
+    public class MessageActionsTests
+    {
+        [Test]
+        public async Task CanCompleteMessagesConcurrently()
+        {
+            var mockReceiver = new Mock<ServiceBusReceiver>();
+            mockReceiver.Setup(r => r.CompleteMessageAsync(It.IsAny<ServiceBusReceivedMessage>(), It.IsAny<CancellationToken>()))
+                // simulate completing the message
+                .Returns(async() => await Task.Delay(20));
+            var actions = new ServiceBusMessageActions(mockReceiver.Object);
+            var tasks = new List<Task>();
+            for (int i = 0; i < 1000; i++)
+            {
+                tasks.Add(actions.CompleteMessageAsync(ServiceBusModelFactory.ServiceBusReceivedMessage()));
+            }
+
+            await Task.WhenAll(tasks);
+        }
+    }
+}