Reduce locking complexity

Author: MarkCiliaVincenti

DBADash/DBADash.csproj

@@ -104,7 +104,7 @@
     </None>
   </ItemGroup>
   <ItemGroup>
-    <PackageReference Include="AsyncKeyedLock" Version="7.1.7" />
+    <PackageReference Include="AsyncKeyedLock" Version="7.1.8" />
     <PackageReference Include="AWSSDK.Core" Version="4.0.3.1" />
     <PackageReference Include="AWSSDK.S3" Version="4.0.11.3" />
     <PackageReference Include="AWSSDK.SQS" Version="4.0.2.3" />

DBADash/Messaging/SQSMessageProcessing.cs

@@ -1,17 +1,17 @@
 using Amazon.SQS;
 using Amazon.SQS.Model;
+using AsyncKeyedLock;
+using Microsoft.Data.SqlClient;
+using Polly;
+using Polly.Retry;
 using Serilog;
 using System;
 using System.Collections.Generic;
-using System.Threading.Tasks;
-using Microsoft.Data.SqlClient;
-using Polly.Retry;
-using Polly;
-using System.Collections.Concurrent;
 using System.Data;
 using System.Linq;
 using System.Text;
 using System.Threading;
+using System.Threading.Tasks;
 
 namespace DBADash.Messaging
 {
@@ -24,17 +24,17 @@ public class SQSMessageProcessing
         private const int delayAfterReceivingMessageForDifferentAgent = 1000; // ms
         private const int delayBetweenMessages = 100; // ms
         private const int errorDelay = 1000; // ms
-        private readonly ConcurrentDictionary<string, SemaphoreSlim> _semaphores = new();
         private const int MaxDegreeOfParallelism = 2;
+        private readonly AsyncKeyedLocker<string> _semaphores = new(o => o.MaxCount = MaxDegreeOfParallelism);
 
         private static readonly ResiliencePipeline pipeline = new ResiliencePipelineBuilder()
             .AddRetry(new RetryStrategyOptions
-        {
-            ShouldHandle = new PredicateBuilder().Handle<Exception>(),
-            BackoffType = DelayBackoffType.Constant,
-            MaxRetryAttempts = 2,
-            Delay = TimeSpan.FromSeconds(1)
-        })
+            {
+                ShouldHandle = new PredicateBuilder().Handle<Exception>(),
+                BackoffType = DelayBackoffType.Constant,
+                MaxRetryAttempts = 2,
+                Delay = TimeSpan.FromSeconds(1)
+            })
         .Build();
 
         public SQSMessageProcessing(CollectionConfig config)
@@ -227,9 +227,6 @@ private async Task ProcessMessageAsync(Message message, string DBADashAgentIdent
             if (msg == null) return;
 
             // Implementations of MessageBase will process the message and return a DataSet or null
-            var semaphore =
-                _semaphores.GetOrAdd(destinationConnectionHash,
-                    _ => new SemaphoreSlim(MaxDegreeOfParallelism));
             try
             {
                 if (msg is CancellationMessage cancellationMessage) // Process cancellation message immediately without waiting for semaphore
@@ -239,7 +236,8 @@ await SendReplyMessage(DBADashAgentIdentifier, handle, destinationConnectionHash
                         ResponseMessage.ResponseTypes.Failure, "Message Cancelled").ConfigureAwait(false);
                     return;
                 }
-                if (!await semaphore.WaitAsync(msg.SemaphoreTimeout)) // Semaphore used to limit concurrent processing per connection
+                using var semaphore = await _semaphores.LockOrNullAsync(destinationConnectionHash, msg.SemaphoreTimeout);
+                if (semaphore is null) // Semaphore used to limit concurrent processing per connection
                 {
                     Log.Warning("Semaphore timeout for message {id} of type {type} with handle {handle}.  Service is busy.", msg.Id, msg.GetType().ToString(), handle);
                     // Semaphore timed out, service is busy
@@ -248,17 +246,9 @@ await SendReplyMessage(DBADashAgentIdentifier, handle, destinationConnectionHash
 
                     return;
                 }
-
-                try
-                {
-                    await DoProcessMessageAsync(msg, DBADashAgentIdentifier, handle, destinationConnectionHash,
-                        replySQS,
-                        replyAgent);
-                }
-                finally
-                {
-                    semaphore.Release();
-                }
+                await DoProcessMessageAsync(msg, DBADashAgentIdentifier, handle, destinationConnectionHash,
+                    replySQS,
+                    replyAgent);
             }
             catch (Exception ex)
             {
@@ -384,11 +374,11 @@ private async Task ProcessReply(Message message, Guid handle, string destination
             if (responseMessage.Data != null && responseMessage.Data.Tables.Contains("DBADash"))
             {
                 Log.Debug("Writing data to SQL");
-             
+
                 await DestinationHandling.WriteDBAsync(responseMessage.Data, destination.ConnectionString,
                         Config);
-         
-                if (responseMessage.Exception==null && responseMessage.Data.Tables.Contains("Errors") && responseMessage.Data.Tables["Errors"]!.Rows.Count>0)
+
+                if (responseMessage.Exception == null && responseMessage.Data.Tables.Contains("Errors") && responseMessage.Data.Tables["Errors"]!.Rows.Count > 0)
                 {
                     var dtErrors = responseMessage.Data.Tables["Errors"];
                     var sbErrors = new StringBuilder();

DBADashService/DBADashJob.cs

@@ -80,18 +80,11 @@ public async Task Execute(IJobExecutionContext context)
                             Log.Debug("Wait for lock {0}", context.JobDetail.Key);
                             // Ensures that S3 folder can only be processed by 1 job instance at a time.
                             // Note: DisallowConcurrentExecution didn't prevent triggered at startup job from overlapping with the scheduled one
-                            var semaphore = ScheduleService.Locker.GetOrAdd(cfg.ConnectionString, _ => new SemaphoreSlim(1, 1));
-                            await semaphore.WaitAsync();
-                            try
+                            using (await ScheduleService.Locker.LockAsync(cfg.ConnectionString))
                             {
                                 Log.Debug("Lock acquired {0}", context.JobDetail.Key);
                                 await CollectS3(cfg);
                             }
-                            finally
-                            {
-                                semaphore.Release();
-                            }
-
                             break;
                         }
                     case ConnectionType.SQL:

DBADashService/SchedulerService.cs

@@ -1,4 +1,5 @@
-using DBADash;
+using AsyncKeyedLock;
+using DBADash;
 using DBADash.InstanceMetadata;
 using DBADash.Messaging;
 using Microsoft.Extensions.Hosting;
@@ -30,7 +31,7 @@ public class ScheduleService : BackgroundService
         private System.Timers.Timer folderCleanupTimer;
         private readonly CollectionSchedules schedules;
         private MessageProcessing messageProcessing;
-        public static readonly ConcurrentDictionary<string, SemaphoreSlim> Locker = new();
+        public static readonly AsyncKeyedLocker<string> Locker = new();
 
         private static readonly ResiliencePipeline pipeline = new ResiliencePipelineBuilder()
             .AddRetry(new RetryStrategyOptions