CmdPal: Separate in-flight icon loads from cached results

src/modules/cmdpal/Microsoft.CmdPal.UI/Helpers/Icons/CachedIconSourceProvider.cs

@@ -2,6 +2,7 @@
 // The Microsoft Corporation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+using System.Collections.Concurrent;
 using System.Runtime.CompilerServices;
 using Microsoft.CmdPal.UI.ViewModels;
 using Microsoft.UI.Xaml.Controls;
@@ -12,18 +13,18 @@ namespace Microsoft.CmdPal.UI.Helpers;
 internal sealed class CachedIconSourceProvider : IIconSourceProvider
 {
     private readonly AdaptiveCache<IconCacheKey, Task<IconSource?>> _cache;
+    private readonly ConcurrentDictionary<IconCacheKey, Task<IconSource?>> _inFlight = new();
     private readonly Size _iconSize;
-    private readonly IconLoaderService _loader;
-    private readonly Lock _lock = new();
+    private readonly IIconLoaderService _loader;
 
-    public CachedIconSourceProvider(IconLoaderService loader, Size iconSize, int cacheSize)
+    public CachedIconSourceProvider(IIconLoaderService loader, Size iconSize, int cacheSize)
     {
         _loader = loader;
         _iconSize = iconSize;
         _cache = new AdaptiveCache<IconCacheKey, Task<IconSource?>>(cacheSize, TimeSpan.FromMinutes(60));
     }
 
-    public CachedIconSourceProvider(IconLoaderService loader, int iconSize, int cacheSize)
+    public CachedIconSourceProvider(IIconLoaderService loader, int iconSize, int cacheSize)
         : this(loader, new Size(iconSize, iconSize), cacheSize)
     {
     }
@@ -39,38 +40,54 @@ public CachedIconSourceProvider(IconLoaderService loader, int iconSize, int cach
 
     private Task<IconSource?> GetOrCreateSlowPath(IconCacheKey key, IconDataViewModel icon, double scale)
     {
-        lock (_lock)
-        {
-            if (_cache.TryGet(key, out var existingTask))
-            {
-                return existingTask;
-            }
-
-            var tcs = new TaskCompletionSource<IconSource?>(TaskCreationOptions.RunContinuationsAsynchronously);
+        var tcs = new TaskCompletionSource<IconSource?>(TaskCreationOptions.RunContinuationsAsynchronously);
+        var task = tcs.Task;
 
-            _loader.EnqueueLoad(
-                icon.Icon,
-                icon.FontFamily,
-                icon.Data?.Unsafe,
-                _iconSize,
-                scale,
-                tcs);
-
-            var task = tcs.Task;
+        var pending = _inFlight.GetOrAdd(key, task);
+        if (!ReferenceEquals(pending, task))
+        {
+            return pending;
+        }
 
-            _ = task.ContinueWith(
-                _ =>
+        _ = task.ContinueWith(
+            completed =>
+            {
+                try
                 {
-                    lock (_lock)
+                    if (completed.IsCompletedSuccessfully)
                     {
-                        _cache.TryRemove(key);
+                        _cache.Add(key, completed);
                     }
-                },
-                TaskContinuationOptions.OnlyOnFaulted);
-
-            _cache.Add(key, task);
-            return task;
+                }
+                finally
+                {
+                    _inFlight.TryRemove(new KeyValuePair<IconCacheKey, Task<IconSource?>>(key, completed));
+                }
+            },
+            CancellationToken.None,
+            TaskContinuationOptions.None,
+            TaskScheduler.Default);
+
+        try
+        {
+            if (!_loader.TryEnqueueLoad(
+                    icon.Icon,
+                    icon.FontFamily,
+                    icon.Data?.Unsafe,
+                    _iconSize,
+                    scale,
+                    tcs,
+                    IconLoadPriority.Low))
+            {
+                tcs.TrySetException(new ObjectDisposedException(nameof(IIconLoaderService)));
+            }
+        }
+        catch (Exception ex)
+        {
+            tcs.TrySetException(ex);
         }
+
+        return task;
     }
 
     private readonly struct IconCacheKey : IEquatable<IconCacheKey>

src/modules/cmdpal/Microsoft.CmdPal.UI/Helpers/Icons/IIconLoaderService.cs

@@ -10,7 +10,7 @@ namespace Microsoft.CmdPal.UI.Helpers;
 
 internal interface IIconLoaderService : IAsyncDisposable
 {
-    void EnqueueLoad(
+    bool TryEnqueueLoad(
         string? iconString,
         string? fontFamily,
         IRandomAccessStreamReference? streamRef,

src/modules/cmdpal/Microsoft.CmdPal.UI/Helpers/Icons/IconLoaderService.cs

@@ -40,7 +40,7 @@ public IconLoaderService(DispatcherQueue dispatcherQueue)
         }
     }
 
-    public void EnqueueLoad(
+    public bool TryEnqueueLoad(
         string? iconString,
         string? fontFamily,
         IRandomAccessStreamReference? streamRef,
@@ -55,15 +55,15 @@ public void EnqueueLoad(
         {
             if (_highPriorityQueue.Writer.TryWrite(workItem))
             {
-                return;
+                return true;
             }
 
 #if DEBUG
             Logger.LogDebug("High priority icon queue full, falling back to low priority");
 #endif
         }
 
-        _lowPriorityQueue.Writer.TryWrite(workItem);
+        return _lowPriorityQueue.Writer.TryWrite(workItem);
     }
 
     public async ValueTask DisposeAsync()

src/modules/cmdpal/Microsoft.CmdPal.UI/Helpers/Icons/IconSourceProvider.cs

@@ -10,18 +10,18 @@ namespace Microsoft.CmdPal.UI.Helpers;
 
 internal sealed class IconSourceProvider : IIconSourceProvider
 {
-    private readonly IconLoaderService _loader;
+    private readonly IIconLoaderService _loader;
     private readonly Size _iconSize;
     private readonly bool _isPriority;
 
-    public IconSourceProvider(IconLoaderService loader, Size iconSize, bool isPriority = false)
+    public IconSourceProvider(IIconLoaderService loader, Size iconSize, bool isPriority = false)
     {
         _loader = loader;
         _iconSize = iconSize;
         _isPriority = isPriority;
     }
 
-    public IconSourceProvider(IconLoaderService loader, int iconSize, bool isPriority = false)
+    public IconSourceProvider(IIconLoaderService loader, int iconSize, bool isPriority = false)
         : this(loader, new Size(iconSize, iconSize), isPriority)
     {
     }
@@ -30,14 +30,24 @@ public IconSourceProvider(IconLoaderService loader, int iconSize, bool isPriorit
     {
         var tcs = new TaskCompletionSource<IconSource?>(TaskCreationOptions.RunContinuationsAsynchronously);
 
-        _loader.EnqueueLoad(
-            icon.Icon,
-            icon.FontFamily,
-            icon.Data?.Unsafe,
-            _iconSize,
-            scale,
-            tcs,
-            _isPriority ? IconLoadPriority.High : IconLoadPriority.Low);
+        try
+        {
+            if (!_loader.TryEnqueueLoad(
+                    icon.Icon,
+                    icon.FontFamily,
+                    icon.Data?.Unsafe,
+                    _iconSize,
+                    scale,
+                    tcs,
+                    _isPriority ? IconLoadPriority.High : IconLoadPriority.Low))
+            {
+                tcs.TrySetException(new ObjectDisposedException(nameof(IIconLoaderService)));
+            }
+        }
+        catch (Exception ex)
+        {
+            tcs.TrySetException(ex);
+        }
 
         return tcs.Task;
     }

src/modules/cmdpal/Tests/Microsoft.CmdPal.UI.UnitTests/CachedIconSourceProviderTests.cs

@@ -0,0 +1,176 @@
+// Copyright (c) Microsoft Corporation
+// The Microsoft Corporation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System.Collections.Concurrent;
+using System.Reflection;
+using Microsoft.CmdPal.UI.Helpers;
+using Microsoft.CmdPal.UI.ViewModels;
+using Microsoft.UI.Xaml.Controls;
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+using Windows.Foundation;
+using Windows.Storage.Streams;
+
+namespace Microsoft.CmdPal.UI.UnitTests;
+
+[TestClass]
+public class CachedIconSourceProviderTests
+{
+    [TestMethod]
+    [Timeout(5_000)]
+    public async Task ConcurrentRequestsShareOneInFlightLoad()
+    {
+        var loader = new ControllableIconLoader();
+        var provider = new CachedIconSourceProvider(loader, new Size(20, 20), cacheSize: 16);
+        var icon = new IconDataViewModel { Icon = "test" };
+        var requests = new ConcurrentBag<Task<IconSource?>>();
+
+        Parallel.For(0, 32, _ => requests.Add(provider.GetIconSource(icon, 1.0)));
+
+        var requestArray = requests.ToArray();
+        Assert.HasCount(32, requestArray);
+        Assert.AreEqual(1, loader.EnqueueCount);
+        foreach (var request in requestArray)
+        {
+            Assert.AreSame(requestArray[0], request);
+        }
+
+        loader.CompleteNext(null);
+        await Task.WhenAll(requestArray);
+    }
+
+    [TestMethod]
+    [Timeout(5_000)]
+    public async Task SuccessfulLoadIsCachedBeforeInFlightEntryIsRemoved()
+    {
+        var loader = new ControllableIconLoader();
+        var provider = new CachedIconSourceProvider(loader, new Size(20, 20), cacheSize: 16);
+        var icon = new IconDataViewModel { Icon = "test" };
+
+        var first = provider.GetIconSource(icon, 1.0);
+        loader.CompleteNext(null);
+        await first;
+
+        Assert.IsTrue(
+            SpinWait.SpinUntil(() => GetInFlightCount(provider) == 0, TimeSpan.FromSeconds(2)),
+            "The completed load was not retired from the in-flight dictionary.");
+
+        var cached = provider.GetIconSource(icon, 1.0);
+
+        Assert.AreSame(first, cached);
+        Assert.AreEqual(1, loader.EnqueueCount);
+    }
+
+    [TestMethod]
+    [Timeout(5_000)]
+    public async Task FailedLoadIsRemovedAndCanBeRetried()
+    {
+        var loader = new ControllableIconLoader();
+        var provider = new CachedIconSourceProvider(loader, new Size(20, 20), cacheSize: 16);
+        var icon = new IconDataViewModel { Icon = "test" };
+
+        var failed = provider.GetIconSource(icon, 1.0);
+        loader.FailNext(new InvalidOperationException("Icon load failed."));
+
+        await Assert.ThrowsExactlyAsync<InvalidOperationException>(async () => await failed);
+        Assert.IsTrue(
+            SpinWait.SpinUntil(() => GetInFlightCount(provider) == 0, TimeSpan.FromSeconds(2)),
+            "The failed load was not retired from the in-flight dictionary.");
+
+        var retry = provider.GetIconSource(icon, 1.0);
+
+        Assert.AreNotSame(failed, retry);
+        Assert.AreEqual(2, loader.EnqueueCount);
+
+        loader.CompleteNext(null);
+        await retry;
+    }
+
+    [TestMethod]
+    [Timeout(5_000)]
+    public async Task RejectedLoadFaultsAndCanBeRetried()
+    {
+        var loader = new ControllableIconLoader { AcceptLoads = false };
+        var provider = new CachedIconSourceProvider(loader, new Size(20, 20), cacheSize: 16);
+        var icon = new IconDataViewModel { Icon = "test" };
+
+        var rejected = provider.GetIconSource(icon, 1.0);
+
+        await Assert.ThrowsExactlyAsync<ObjectDisposedException>(async () => await rejected);
+        Assert.IsTrue(
+            SpinWait.SpinUntil(() => GetInFlightCount(provider) == 0, TimeSpan.FromSeconds(2)),
+            "The rejected load was not retired from the in-flight dictionary.");
+
+        loader.AcceptLoads = true;
+        var retry = provider.GetIconSource(icon, 1.0);
+
+        Assert.AreNotSame(rejected, retry);
+        Assert.AreEqual(2, loader.EnqueueCount);
+
+        loader.CompleteNext(null);
+        await retry;
+    }
+
+    [TestMethod]
+    public async Task UncachedProviderFaultsRejectedLoad()
+    {
+        var loader = new ControllableIconLoader { AcceptLoads = false };
+        var provider = new IconSourceProvider(loader, new Size(16, 16));
+
+        var rejected = provider.GetIconSource(new IconDataViewModel { Icon = "test" }, 1.0);
+
+        await Assert.ThrowsExactlyAsync<ObjectDisposedException>(async () => await rejected);
+        Assert.AreEqual(1, loader.EnqueueCount);
+    }
+
+    private static int GetInFlightCount(CachedIconSourceProvider provider)
+    {
+        var field = typeof(CachedIconSourceProvider).GetField("_inFlight", BindingFlags.Instance | BindingFlags.NonPublic);
+        var inFlight = field!.GetValue(provider)!;
+        var countProperty = inFlight.GetType().GetProperty("Count");
+        return (int)countProperty!.GetValue(inFlight)!;
+    }
+
+    private sealed class ControllableIconLoader : IIconLoaderService
+    {
+        private readonly ConcurrentQueue<TaskCompletionSource<IconSource?>> _pending = new();
+        private int _enqueueCount;
+
+        public bool AcceptLoads { get; set; } = true;
+
+        public int EnqueueCount => Volatile.Read(ref _enqueueCount);
+
+        public bool TryEnqueueLoad(
+            string? iconString,
+            string? fontFamily,
+            IRandomAccessStreamReference? streamRef,
+            Size iconSize,
+            double scale,
+            TaskCompletionSource<IconSource?> tcs,
+            IconLoadPriority priority)
+        {
+            Interlocked.Increment(ref _enqueueCount);
+            if (!AcceptLoads)
+            {
+                return false;
+            }
+
+            _pending.Enqueue(tcs);
+            return true;
+        }
+
+        public void CompleteNext(IconSource? result)
+        {
+            Assert.IsTrue(_pending.TryDequeue(out var tcs), "No pending icon load was available.");
+            tcs.SetResult(result);
+        }
+
+        public void FailNext(Exception exception)
+        {
+            Assert.IsTrue(_pending.TryDequeue(out var tcs), "No pending icon load was available.");
+            tcs.SetException(exception);
+        }
+
+        public ValueTask DisposeAsync() => ValueTask.CompletedTask;
+    }
+}

src/modules/cmdpal/Tests/Microsoft.CmdPal.UI.UnitTests/IconDataStreamReference.cs

@@ -0,0 +1,12 @@
+// Copyright (c) Microsoft Corporation
+// The Microsoft Corporation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using Windows.Storage.Streams;
+
+namespace Microsoft.CmdPal.UI.ViewModels;
+
+internal sealed class IconDataStreamReference
+{
+    public IRandomAccessStreamReference? Unsafe { get; init; }
+}