Fix race in LFU concurrent update and TryRemove(kvp) (#618)

Author: bitfaster

BitFaster.Caching.UnitTests/Lfu/ConcurrentLfuSoakTests.cs

@@ -290,6 +290,32 @@ await Threaded.Run(threads, i =>
             RunIntegrityCheck(cache, this.output);
         }
 
+        [Fact]
+        public async Task WhenConcurrentUpdateAndRemoveKvp()
+        {
+            var cache = new ConcurrentLfu<int, string>(1, 20, new BackgroundThreadScheduler(), EqualityComparer<int>.Default);
+            TaskCompletionSource<int> tcs = new TaskCompletionSource<int> ();
+
+            var removal = Task.Run(() =>
+            {
+                while (!tcs.Task.IsCompleted)
+                {
+                    cache.TryRemove(new KeyValuePair<int, string>(5, "x"));
+                }
+            });
+
+            for (var i = 0; i < 100_000; i++)
+            {
+                cache.AddOrUpdate(5, "a");
+                cache.TryGet(5, out _).Should().BeTrue("key 'a' should not be deleted");
+                cache.AddOrUpdate(5, "x");
+            }
+
+            tcs.SetResult(int.MaxValue);
+
+            await removal;
+        }
+
         private ConcurrentLfu<int, string> CreateWithBackgroundScheduler()
         {
             var scheduler = new BackgroundThreadScheduler();

BitFaster.Caching/Lfu/ConcurrentLfuCore.cs

@@ -317,20 +317,23 @@ public bool TryRemove(KeyValuePair<K, V> item)
         {
             if (this.dictionary.TryGetValue(item.Key, out var node))
             {
-                if (EqualityComparer<V>.Default.Equals(node.Value, item.Value))
-                {
-                    var kvp = new KeyValuePair<K, N>(item.Key, node);
+                lock (node)
+                { 
+                    if (EqualityComparer<V>.Default.Equals(node.Value, item.Value))
+                    {
+                        var kvp = new KeyValuePair<K, N>(item.Key, node);
 
 #if NET6_0_OR_GREATER
-                    if (this.dictionary.TryRemove(kvp))
+                        if (this.dictionary.TryRemove(kvp))
 #else
-                    // https://devblogs.microsoft.com/pfxteam/little-known-gems-atomic-conditional-removals-from-concurrentdictionary/
-                    if (((ICollection<KeyValuePair<K, N>>)this.dictionary).Remove(kvp))
+                        // https://devblogs.microsoft.com/pfxteam/little-known-gems-atomic-conditional-removals-from-concurrentdictionary/
+                        if (((ICollection<KeyValuePair<K, N>>)this.dictionary).Remove(kvp))
 #endif
-                    {
-                        node.WasRemoved = true;
-                        AfterWrite(node);
-                        return true;
+                        {
+                            node.WasRemoved = true;
+                            AfterWrite(node);
+                            return true;
+                        }
                     }
                 }
             }
@@ -361,14 +364,20 @@ public bool TryUpdate(K key, V value)
         {
             if (this.dictionary.TryGetValue(key, out var node))
             {
-                node.Value = value;
+                lock (node)
+                { 
+                    if (!node.WasRemoved)
+                    {
+                         node.Value = value;
 
-                // It's ok for this to be lossy, since the node is already tracked
-                // and we will just lose ordering/hit count, but not orphan the node.
-                this.writeBuffer.TryAdd(node);
-                TryScheduleDrain();
-                this.policy.OnWrite(node);
-                return true;
+                        // It's ok for this to be lossy, since the node is already tracked
+                        // and we will just lose ordering/hit count, but not orphan the node.
+                        this.writeBuffer.TryAdd(node);
+                        TryScheduleDrain();
+                        this.policy.OnWrite(node);
+                        return true;
+                    }
+                }
             }
 
             return false;