CmdPal: Avoid all-stripe locking in the adaptive icon cache

PowerToys.slnx

@@ -381,6 +381,10 @@
       <Platform Solution="*|ARM64" Project="ARM64" />
       <Platform Solution="*|x64" Project="x64" />
     </Project>
+    <Project Path="src/modules/cmdpal/Tests/Microsoft.CmdPal.UI.UnitTests/Microsoft.CmdPal.UI.UnitTests.csproj">
+      <Platform Solution="*|ARM64" Project="ARM64" />
+      <Platform Solution="*|x64" Project="x64" />
+    </Project>
     <Project Path="src/modules/cmdpal/Tests/Microsoft.CmdPal.UI.ViewModels.UnitTests/Microsoft.CmdPal.UI.ViewModels.UnitTests.csproj">
       <Platform Solution="*|ARM64" Project="ARM64" />
       <Platform Solution="*|x64" Project="x64" />

src/modules/cmdpal/Microsoft.CmdPal.UI/Helpers/AdaptiveCache`2.cs

@@ -24,10 +24,15 @@ internal sealed class AdaptiveCache<TKey, TValue>
     private readonly ConcurrentStack<CacheEntry> _pool = [];
     private readonly WaitCallback _maintenanceCallback;
 
+    // ConcurrentDictionary.Count acquires every stripe lock. Keep an approximate count so
+    // cache maintenance never makes the XAML UI thread wait for all dictionary locks.
+    private int _entryCount;
     private long _currentTick;
     private long _lastDecayTicks = DateTime.UtcNow.Ticks;
     private InterlockedBoolean _maintenanceSwitch = new(false);
 
+    internal int ApproximateCount => Volatile.Read(ref _entryCount);
+
     public AdaptiveCache(int capacity = 384, TimeSpan? decayInterval = null, double decayFactor = 0.5)
     {
         _capacity = capacity;
@@ -66,7 +71,11 @@ public TValue GetOrAdd<TArg>(TKey key, Func<TKey, TArg, TValue> factory, TArg ar
         var tick = Interlocked.Increment(ref _currentTick);
         newEntry.Initialize(key, value, 1.0, tick);
 
-        if (!_map.TryAdd(key, newEntry))
+        if (_map.TryAdd(key, newEntry))
+        {
+            Interlocked.Increment(ref _entryCount);
+        }
+        else
         {
             newEntry.Clear();
             _pool.Push(newEntry);
@@ -117,7 +126,11 @@ public void Add(TKey key, TValue value)
 
         newEntry.Initialize(key, value, 1.0, tick);
 
-        if (!_map.TryAdd(key, newEntry))
+        if (_map.TryAdd(key, newEntry))
+        {
+            Interlocked.Increment(ref _entryCount);
+        }
+        else
         {
             newEntry.Clear();
             _pool.Push(newEntry);
@@ -133,6 +146,7 @@ public bool TryRemove(TKey key)
     {
         if (_map.TryRemove(key, out var evicted))
         {
+            Interlocked.Decrement(ref _entryCount);
             evicted.Clear();
             _pool.Push(evicted);
             return true;
@@ -143,7 +157,9 @@ public bool TryRemove(TKey key)
 
     public void Clear()
     {
-        foreach (var key in _map.Keys)
+        // Enumerate the dictionary rather than _map.Keys: the enumerator is lock-free,
+        // while Keys snapshots under every stripe lock.
+        foreach (var (key, _) in _map)
         {
             TryRemove(key);
         }
@@ -153,7 +169,7 @@ public void Clear()
 
     private bool ShouldMaintenanceRun()
     {
-        return _map.Count > _capacity || (DateTime.UtcNow.Ticks - Interlocked.Read(ref _lastDecayTicks)) > _decayInterval.Ticks;
+        return ApproximateCount > _capacity || (DateTime.UtcNow.Ticks - Interlocked.Read(ref _lastDecayTicks)) > _decayInterval.Ticks;
     }
 
     private void TryRunMaintenance()
@@ -184,13 +200,9 @@ private void PerformCleanup()
 
             var score = CalculateScore(entry, currentTick);
 
-            if (score < 0.1 || _map.Count > _capacity)
+            if (score < 0.1 || ApproximateCount > _capacity)
             {
-                if (_map.TryRemove(key, out var evicted))
-                {
-                    evicted.Clear();
-                    _pool.Push(evicted);
-                }
+                TryRemove(key);
             }
         }
     }

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

@@ -0,0 +1,102 @@
+// 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;
+using System.Collections.Concurrent;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CmdPal.UI.Helpers;
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+namespace Microsoft.CmdPal.UI.UnitTests;
+
+[TestClass]
+public class AdaptiveCacheTests
+{
+    [TestMethod]
+    public void ApproximateCountChangesOnlyForSuccessfulMutations()
+    {
+        var cache = new AdaptiveCache<int, int>(capacity: 8);
+
+        cache.Add(1, 1);
+        cache.Add(1, 2);
+        _ = cache.GetOrAdd(1, static (key, _) => key, 0);
+
+        Assert.AreEqual(1, cache.ApproximateCount);
+        Assert.IsFalse(cache.TryRemove(2));
+        Assert.AreEqual(1, cache.ApproximateCount);
+        Assert.IsTrue(cache.TryRemove(1));
+        Assert.AreEqual(0, cache.ApproximateCount);
+
+        _ = cache.GetOrAdd(2, static (key, _) => key, 0);
+        Assert.AreEqual(1, cache.ApproximateCount);
+
+        cache.Clear();
+        Assert.AreEqual(0, cache.ApproximateCount);
+    }
+
+    [TestMethod]
+    [Timeout(15_000)]
+    public async Task ConcurrentCleanupAndFailedLoadsKeepApproximateCountConsistent()
+    {
+        const int capacity = 1;
+        const int workerCount = 8;
+        const int itemsPerWorker = 256;
+        var cache = new AdaptiveCache<int, Task<int>>(capacity, TimeSpan.FromHours(1));
+        var faultRemovals = new ConcurrentBag<Task>();
+
+        cache.Add(-2, Task.FromResult(-2));
+        cache.Add(-1, Task.FromResult(-1));
+
+        Assert.IsTrue(
+            SpinWait.SpinUntil(() => cache.ApproximateCount <= capacity, TimeSpan.FromSeconds(5)),
+            "Capacity cleanup did not run.");
+
+        var workers = Enumerable.Range(0, workerCount)
+            .Select(worker => Task.Run(() =>
+            {
+                for (var item = 0; item < itemsPerWorker; item++)
+                {
+                    var key = (worker * itemsPerWorker) + item;
+
+                    if (item % 4 == 0)
+                    {
+                        var completionSource = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
+                        cache.Add(key, completionSource.Task);
+
+                        faultRemovals.Add(completionSource.Task.ContinueWith(
+                            completedTask =>
+                            {
+                                _ = completedTask.Exception;
+                                cache.TryRemove(key);
+                            },
+                            CancellationToken.None,
+                            TaskContinuationOptions.OnlyOnFaulted,
+                            TaskScheduler.Default));
+
+                        completionSource.SetException(new InvalidOperationException("Icon load failed."));
+                    }
+                    else
+                    {
+                        _ = cache.GetOrAdd(key, static (currentKey, _) => Task.FromResult(currentKey), 0);
+                    }
+
+                    _ = cache.TryGet(key, out _);
+                }
+            }))
+            .ToArray();
+
+        await Task.WhenAll(workers);
+        await Task.WhenAll(faultRemovals);
+
+        cache.Clear();
+
+        Assert.AreEqual(0, cache.ApproximateCount);
+        for (var key = 0; key < workerCount * itemsPerWorker; key++)
+        {
+            Assert.IsFalse(cache.TryGet(key, out _), $"Cache still contains key {key} after Clear.");
+        }
+    }
+}

src/modules/cmdpal/Tests/Microsoft.CmdPal.UI.UnitTests/Microsoft.CmdPal.UI.UnitTests.csproj

@@ -0,0 +1,28 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <!-- Look at Directory.Build.props in root for common stuff as well -->
+  <Import Project="$(RepoRoot)src\Common.Dotnet.CsWinRT.props" />
+
+  <PropertyGroup>
+    <IsPackable>false</IsPackable>
+    <IsTestProject>true</IsTestProject>
+    <RootNamespace>Microsoft.CmdPal.UI.UnitTests</RootNamespace>
+    <OutputPath>$(RepoRoot)$(Platform)\$(Configuration)\WinUI3Apps\CmdPal\tests\</OutputPath>
+    <AppendTargetFrameworkToOutputPath>false</AppendTargetFrameworkToOutputPath>
+    <AppendRuntimeIdentifierToOutputPath>false</AppendRuntimeIdentifierToOutputPath>
+    <ImplicitUsings>enable</ImplicitUsings>
+    <LangVersion>preview</LangVersion>
+    <Nullable>enable</Nullable>
+  </PropertyGroup>
+
+  <ItemGroup>
+    <PackageReference Include="MSTest" />
+  </ItemGroup>
+
+  <ItemGroup>
+    <ProjectReference Include="..\..\Microsoft.CmdPal.Common\Microsoft.CmdPal.Common.csproj" />
+  </ItemGroup>
+
+  <ItemGroup>
+    <Compile Include="..\..\Microsoft.CmdPal.UI\Helpers\AdaptiveCache`2.cs" Link="Helpers\AdaptiveCache`2.cs" />
+  </ItemGroup>
+</Project>