LRU: atomic ConcurrentDictionary.Remove() (#67)

* atomic remove

Author: bitfaster

BitFaster.Caching.UnitTests/Lru/ConcurrentLruTests.cs

@@ -385,6 +385,19 @@ public void WhenKeyExistsTryUpdateUpdatesValueAndReturnsTrue()
             value.Should().Be("2");
         }
 
+        [Fact]
+        public void WhenKeyExistsTryUpdateDisposesOldValue()
+        {
+            var lruOfDisposable = new ConcurrentLru<int, DisposableItem>(1, 6, EqualityComparer<int>.Default);
+            var disposableValueFactory = new DisposableValueFactory();
+            var newValue = new DisposableItem();
+
+            lruOfDisposable.GetOrAdd(1, disposableValueFactory.Create);
+            lruOfDisposable.TryUpdate(1, newValue);
+
+            disposableValueFactory.Items[1].IsDisposed.Should().BeTrue();
+        }
+
         [Fact]
         public void WhenKeyDoesNotExistTryUpdateReturnsFalse()
         {
@@ -403,7 +416,7 @@ public void WhenKeyDoesNotExistAddOrUpdateAddsNewItem()
         }
 
         [Fact]
-        public void WhenKeyExistsAddOrUpdatUpdatesExistingItem()
+        public void WhenKeyExistsAddOrUpdateUpdatesExistingItem()
         {
             lru.AddOrUpdate(1, "1");
             lru.AddOrUpdate(1, "2");
@@ -412,6 +425,19 @@ public void WhenKeyExistsAddOrUpdatUpdatesExistingItem()
             value.Should().Be("2");
         }
 
+        [Fact]
+        public void WhenKeyExistsAddOrUpdateDisposesOldValue()
+        {
+            var lruOfDisposable = new ConcurrentLru<int, DisposableItem>(1, 6, EqualityComparer<int>.Default);
+            var disposableValueFactory = new DisposableValueFactory();
+            var newValue = new DisposableItem();
+
+            lruOfDisposable.GetOrAdd(1, disposableValueFactory.Create);
+            lruOfDisposable.AddOrUpdate(1, newValue);
+
+            disposableValueFactory.Items[1].IsDisposed.Should().BeTrue();
+        }
+
         [Fact]
         public void WhenKeyDoesNotExistAddOrUpdateMaintainsLruOrder()
         {

BitFaster.Caching/Lru/TemplateConcurrentLru.cs

@@ -176,48 +176,34 @@ public async Task<V> GetOrAddAsync(K key, Func<K, Task<V>> valueFactory)
         ///<inheritdoc/>
         public bool TryRemove(K key)
         {
-            // Possible race condition:
-            // Thread A TryRemove(1), removes LruItem1, has reference to removed item but not yet marked as removed
-            // Thread B GetOrAdd(1) => Adds LruItem1*
-            // Thread C GetOrAdd(2), Cycle, Move(LruItem1, Removed)
-            // 
-            // Thread C can run and remove LruItem1* from this.dictionary before Thread A has marked LruItem1 as removed.
-            // 
-            // In this situation, a subsequent attempt to fetch 1 will be a miss. The queues will still contain LruItem1*, 
-            // and it will not be marked as removed. If key 1 is fetched while LruItem1* is still in the queue, there will 
-            // be two queue entries for key 1, and neither is marked as removed. Thus when LruItem1 * ages out, it will  
-            // incorrectly remove 1 from the dictionary, and this cycle can repeat.
             if (this.dictionary.TryGetValue(key, out var existing))
             {
-                if (existing.WasRemoved)
-                {
-                    return false;
-                }
+                var kvp = new KeyValuePair<K, I>(key, existing);
 
-                lock (existing)
-                {
-                    if (existing.WasRemoved)
-                    {
-                        return false;
-                    }
-
-                    existing.WasRemoved = true;
-                }
-
-                if (this.dictionary.TryRemove(key, out var removedItem))
+                // hidden atomic remove
+                // https://devblogs.microsoft.com/pfxteam/little-known-gems-atomic-conditional-removals-from-concurrentdictionary/
+                if (((ICollection<KeyValuePair<K, I>>)this.dictionary).Remove(kvp))
                 {
                     // Mark as not accessed, it will later be cycled out of the queues because it can never be fetched 
                     // from the dictionary. Note: Hot/Warm/Cold count will reflect the removed item until it is cycled 
                     // from the queue.
-                    removedItem.WasAccessed = false;
+                    existing.WasAccessed = false;
+                    existing.WasRemoved = true;
 
-                    if (removedItem.Value is IDisposable d)
+                    // serialize dispose (common case dispose not thread safe)
+                    lock (existing)
                     {
-                        d.Dispose();
+                        if (existing.Value is IDisposable d)
+                        {
+                            d.Dispose();
+                        }
                     }
 
                     return true;
                 }
+
+                // it existed, but we couldn't remove - this means value was replaced afer the TryGetValue (a race), try again
+                return TryRemove(key);
             }
 
             return false;
@@ -233,7 +219,14 @@ public bool TryUpdate(K key, V value)
                 {
                     if (!existing.WasRemoved)
                     {
+                        V oldValue = existing.Value;
                         existing.Value = value;
+
+                        if (oldValue is IDisposable d)
+                        {
+                            d.Dispose();
+                        }
+
                         return true;
                     }
                 }
@@ -247,16 +240,9 @@ public bool TryUpdate(K key, V value)
         public void AddOrUpdate(K key, V value)
         { 
             // first, try to update
-            if (this.dictionary.TryGetValue(key, out var existing))
-            {
-                lock (existing)
-                {
-                    if (!existing.WasRemoved)
-                    {
-                        existing.Value = value;
-                        return;
-                    }
-                }
+            if (this.TryUpdate(key, value))
+            { 
+                return;
             }
 
             // then try add
@@ -380,26 +366,24 @@ private void Move(I item, ItemDestination where)
                     Interlocked.Increment(ref this.coldCount);
                     break;
                 case ItemDestination.Remove:
-                    if (!item.WasRemoved)
-                    {   
-                        // avoid race where 2 threads could remove the same key - see TryRemove for details.
+
+                    var kvp = new KeyValuePair<K, I>(item.Key, item);
+
+                    // hidden atomic remove
+                    // https://devblogs.microsoft.com/pfxteam/little-known-gems-atomic-conditional-removals-from-concurrentdictionary/
+                    if (((ICollection<KeyValuePair<K, I>>)this.dictionary).Remove(kvp))
+                    {
+                        item.WasRemoved = true;
+
                         lock (item)
                         { 
-                            if (item.WasRemoved)
-                            {
-                                break;
-                            }
-
-                            if (this.dictionary.TryRemove(item.Key, out var removedItem))
+                            if (item.Value is IDisposable d)
                             {
-                                item.WasRemoved = true;
-                                if (removedItem.Value is IDisposable d)
-                                {
-                                    d.Dispose();
-                                }
-                            }
+                                d.Dispose();
+                            } 
                         }
                     }
+
                     break;
             }
         }