diff --git a/BitFaster.Caching.UnitTests/Lfu/ConcurrentLfuSoakTests.cs b/BitFaster.Caching.UnitTests/Lfu/ConcurrentLfuSoakTests.cs index fee6902d..67a0cec0 100644 --- a/BitFaster.Caching.UnitTests/Lfu/ConcurrentLfuSoakTests.cs +++ b/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(1, 20, new BackgroundThreadScheduler(), EqualityComparer.Default); + TaskCompletionSource tcs = new TaskCompletionSource (); + + var removal = Task.Run(() => + { + while (!tcs.Task.IsCompleted) + { + cache.TryRemove(new KeyValuePair(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 CreateWithBackgroundScheduler() { var scheduler = new BackgroundThreadScheduler(); diff --git a/BitFaster.Caching/Lfu/ConcurrentLfuCore.cs b/BitFaster.Caching/Lfu/ConcurrentLfuCore.cs index 373a5ebb..67f6dbc1 100644 --- a/BitFaster.Caching/Lfu/ConcurrentLfuCore.cs +++ b/BitFaster.Caching/Lfu/ConcurrentLfuCore.cs @@ -317,20 +317,23 @@ public bool TryRemove(KeyValuePair item) { if (this.dictionary.TryGetValue(item.Key, out var node)) { - if (EqualityComparer.Default.Equals(node.Value, item.Value)) - { - var kvp = new KeyValuePair(item.Key, node); + lock (node) + { + if (EqualityComparer.Default.Equals(node.Value, item.Value)) + { + var kvp = new KeyValuePair(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>)this.dictionary).Remove(kvp)) + // https://devblogs.microsoft.com/pfxteam/little-known-gems-atomic-conditional-removals-from-concurrentdictionary/ + if (((ICollection>)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;