Improve concurrency handling in `ExportAssetDownloader` (#1002)

Co-authored-by: Oleksii Holub <1935960+Tyrrrz@users.noreply.github.com>

DiscordChatExporter.Cli.Tests/Infra/ExportWrapper.cs

@@ -9,6 +9,7 @@ using System.Threading;
 using System.Threading.Tasks;
 using AngleSharp.Dom;
 using AngleSharp.Html.Dom;
+using AsyncKeyedLock;
 using CliFx.Infrastructure;
 using DiscordChatExporter.Cli.Commands;
 using DiscordChatExporter.Cli.Tests.Utils;
@@ -20,7 +21,11 @@ namespace DiscordChatExporter.Cli.Tests.Infra;
 
 public static class ExportWrapper
 {
-    private static readonly ConcurrentDictionary<string, SemaphoreSlim> Locks = new(StringComparer.Ordinal);
+    private static readonly AsyncKeyedLocker<string> Locker = new(o =>
+    {
+        o.PoolSize = 20;
+        o.PoolInitialFill = 1;
+    });
 
     private static readonly string DirPath = Path.Combine(
         Path.GetDirectoryName(Assembly.GetExecutingAssembly().Location) ?? Directory.GetCurrentDirectory(),
@@ -46,10 +51,7 @@ public static class ExportWrapper
         var filePath = Path.Combine(DirPath, fileName);
 
         // Lock separately for each channel and format
-        var channelLock = Locks.GetOrAdd(filePath, _ => new SemaphoreSlim(1, 1));
-        await channelLock.WaitAsync();
-
-        try
+        using (await Locker.LockAsync(filePath))
         {
             // Perform export only if it hasn't been done before
             if (!File.Exists(filePath))
@@ -65,10 +67,6 @@ public static class ExportWrapper
 
             return await File.ReadAllTextAsync(filePath);
         }
-        finally
-        {
-            channelLock.Release();
-        }
     }
 
     public static async ValueTask<IHtmlDocument> ExportAsHtmlAsync(Snowflake channelId) => Html.Parse(

DiscordChatExporter.Core/DiscordChatExporter.Core.csproj

@@ -5,6 +5,7 @@
   </PropertyGroup>
 
   <ItemGroup>
+    <PackageReference Include="AsyncKeyedLock" Version="6.1.1" />
     <PackageReference Include="Gress" Version="2.0.1" />
     <PackageReference Include="JsonExtensions" Version="1.2.0" />
     <PackageReference Include="Polly" Version="7.2.3" />

DiscordChatExporter.Core/Exporting/ExportAssetDownloader.cs

@@ -7,6 +7,7 @@ using System.Text;
 using System.Text.RegularExpressions;
 using System.Threading;
 using System.Threading.Tasks;
+using AsyncKeyedLock;
 using DiscordChatExporter.Core.Utils;
 using DiscordChatExporter.Core.Utils.Extensions;
 
@@ -14,6 +15,11 @@ namespace DiscordChatExporter.Core.Exporting;
 
 internal partial class ExportAssetDownloader
 {
+    private static readonly AsyncKeyedLocker<string> _locker = new(o =>
+    {
+        o.PoolSize = 20;
+        o.PoolInitialFill = 1;
+    });
     private readonly string _workingDirPath;
     private readonly bool _reuse;
 
@@ -28,51 +34,54 @@ internal partial class ExportAssetDownloader
 
     public async ValueTask<string> DownloadAsync(string url, CancellationToken cancellationToken = default)
     {
-        if (_pathCache.TryGetValue(url, out var cachedFilePath))
-            return cachedFilePath;
-
         var fileName = GetFileNameFromUrl(url);
         var filePath = Path.Combine(_workingDirPath, fileName);
 
-        // Reuse existing files if we're allowed to
-        if (!_reuse || !File.Exists(filePath))
+        using (await _locker.LockAsync(filePath, cancellationToken).ConfigureAwait(false))
         {
-            Directory.CreateDirectory(_workingDirPath);
+            if (_pathCache.TryGetValue(url, out var cachedFilePath))
+                return cachedFilePath;
 
-            await Http.ResiliencePolicy.ExecuteAsync(async () =>
+            // Reuse existing files if we're allowed to
+            if (!_reuse || !File.Exists(filePath))
             {
-                // Download the file
-                using var response = await Http.Client.GetAsync(url, cancellationToken);
-                await using (var output = File.Create(filePath))
-                    await response.Content.CopyToAsync(output, cancellationToken);
+                Directory.CreateDirectory(_workingDirPath);
 
-                // Try to set the file date according to the last-modified header
-                try
+                await Http.ResiliencePolicy.ExecuteAsync(async () =>
                 {
-                    var lastModified = response.Content.Headers.TryGetValue("Last-Modified")?.Pipe(s =>
-                        DateTimeOffset.TryParse(s, CultureInfo.InvariantCulture, DateTimeStyles.None, out var instant)
-                            ? instant
-                            : (DateTimeOffset?) null
-                    );
+                    // Download the file
+                    using var response = await Http.Client.GetAsync(url, cancellationToken);
+                    await using (var output = File.Create(filePath))
+                        await response.Content.CopyToAsync(output, cancellationToken);
 
-                    if (lastModified is not null)
+                    // Try to set the file date according to the last-modified header
+                    try
                     {
-                        File.SetCreationTimeUtc(filePath, lastModified.Value.UtcDateTime);
-                        File.SetLastWriteTimeUtc(filePath, lastModified.Value.UtcDateTime);
-                        File.SetLastAccessTimeUtc(filePath, lastModified.Value.UtcDateTime);
+                        var lastModified = response.Content.Headers.TryGetValue("Last-Modified")?.Pipe(s =>
+                            DateTimeOffset.TryParse(s, CultureInfo.InvariantCulture, DateTimeStyles.None, out var instant)
+                                ? instant
+                                : (DateTimeOffset?)null
+                        );
+
+                        if (lastModified is not null)
+                        {
+                            File.SetCreationTimeUtc(filePath, lastModified.Value.UtcDateTime);
+                            File.SetLastWriteTimeUtc(filePath, lastModified.Value.UtcDateTime);
+                            File.SetLastAccessTimeUtc(filePath, lastModified.Value.UtcDateTime);
+                        }
                     }
-                }
-                catch
-                {
-                    // This can apparently fail for some reason.
-                    // https://github.com/Tyrrrz/DiscordChatExporter/issues/585
-                    // Updating file dates is not a critical task, so we'll just
-                    // ignore exceptions thrown here.
-                }
-            });
-        }
+                    catch
+                    {
+                        // This can apparently fail for some reason.
+                        // https://github.com/Tyrrrz/DiscordChatExporter/issues/585
+                        // Updating file dates is not a critical task, so we'll just
+                        // ignore exceptions thrown here.
+                    }
+                });
+            }
 
-        return _pathCache[url] = filePath;
+            return _pathCache[url] = filePath;
+        }
     }
 }