Fix race in LFU maintenance for add-remove-add of the same key (#533)

Author: bitfaster

BitFaster.Caching.UnitTests/Lfu/ConcurrentLfuSoakTests.cs

@@ -376,14 +376,7 @@ private void VerifyLruInDictionary(LfuNodeList<K, V> lfuNodes, ITestOutputHelper
                 node.WasRemoved.Should().BeFalse();
                 node.WasDeleted.Should().BeFalse();
 
-                // This can occur if there is a race between:
-                // Thread 1: TryRemove, delete node from dictionary, set WasRemoved flag
-                // Thread 2: Check WasRemoved flag, if not add to lru
-                // It's not clear how WasRemoved can be false in this situation.
-                if (!cache.TryGet(node.Key, out _))
-                {
-                    output.WriteLine($"Orphaned node with key {node.Key} detected.");
-                }
+                cache.TryGet(node.Key, out _).Should().BeTrue($"Orphaned node with key {node.Key} detected.");
 
                 node = node.Next;
             }

BitFaster.Caching/Lfu/ConcurrentLfuCore.cs

@@ -133,7 +133,7 @@ public void AddOrUpdate(K key, V value)
 
         public void Clear()
         {
-            this.Trim(this.Count);
+            Trim(int.MaxValue);
 
             lock (maintenanceLock)
             {
@@ -145,15 +145,16 @@ public void Clear()
 
         public void Trim(int itemCount)
         {
-            itemCount = Math.Min(itemCount, this.Count);
-            var candidates = new List<LfuNode<K, V>>(itemCount);
-
-            // TODO: this is LRU order eviction, Caffeine is based on frequency
+            List<LfuNode<K, V>> candidates;
             lock (maintenanceLock)
             {
-                // flush all buffers
                 Maintenance();
 
+                int lruCount = this.windowLru.Count + this.probationLru.Count + this.protectedLru.Count;
+                itemCount = Math.Min(itemCount, lruCount);
+                candidates = new (itemCount);
+
+                // Note: this is LRU order eviction, Caffeine is based on frequency
                 // walk in lru order, get itemCount keys to evict
                 TakeCandidatesInLruOrder(this.probationLru, candidates, itemCount);
                 TakeCandidatesInLruOrder(this.protectedLru, candidates, itemCount);
@@ -608,12 +609,6 @@ private void OnWrite(LfuNode<K, V> node)
 
         private void PromoteProbation(LfuNode<K, V> node)
         {
-            if (node.list == null)
-            {
-                // Ignore stale accesses for an entry that is no longer present
-                return;
-            }
-
             this.probationLru.Remove(node);
             this.protectedLru.AddLast(node);
             node.Position = Position.Protected;
@@ -699,6 +694,22 @@ private void EvictFromMain(LfuNode<K, V> candidateNode)
                     break;
                 }
 
+                if (candidate.node.WasRemoved)
+                {
+                    var evictee = candidate.node;
+                    candidate.Next();
+                    Evict(evictee);
+                    continue;
+                }
+
+                if (victim.node.WasRemoved)
+                {
+                    var evictee = victim.node;
+                    victim.Next();
+                    Evict(evictee);
+                    continue;
+                }
+
                 // Evict the entry with the lowest frequency
                 if (candidate.freq > victim.freq)
                 {
@@ -749,7 +760,17 @@ private bool AdmitCandidate(K candidateKey, K victimKey)
 
         private void Evict(LfuNode<K, V> evictee)
         {
-            this.dictionary.TryRemove(evictee.Key, out var _);
+            evictee.WasRemoved = true;
+            evictee.WasDeleted = true;
+
+            // This handles the case where the same key exists in the write buffer both
+            // as added and removed. Remove via KVP ensures we don't remove added nodes.
+            var kvp = new KeyValuePair<K, N>(evictee.Key, (N)evictee);
+#if NET6_0_OR_GREATER
+            this.dictionary.TryRemove(kvp);
+#else
+            ((ICollection<KeyValuePair<K, N>>)this.dictionary).Remove(kvp);
+#endif
             evictee.list.Remove(evictee);
             Disposer<V>.Dispose(evictee.Value);
             this.metrics.evictedCount++;

BitFaster.Caching/Lfu/LfuNode.cs

@@ -21,12 +21,18 @@ public LfuNode(K k, V v)
 
         public Position Position { get; set; }
 
+        /// <summary>
+        /// Node was removed from the dictionary, but is still present in the LRU lists.
+        /// </summary>
         public bool WasRemoved
         {
             get => this.wasRemoved;
             set => this.wasRemoved = value;
         }
 
+        /// <summary>
+        /// Node has been removed both from the dictionary and the LRU lists.
+        /// </summary>
         public bool WasDeleted
         {
             get => this.wasDeleted;