Fix object disposed exception in SingletonCache.Release (#66)

* fix obj disposed ex

Author: bitfaster

BitFaster.Caching.UnitTests/AssemblyInitialize.cs

@@ -0,0 +1,28 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading;
+using System.Threading.Tasks;
+using Xunit.Abstractions;
+using Xunit.Sdk;
+
+[assembly: Xunit.TestFramework("BitFaster.Caching.UnitTests.AssemblyInitialize", "BitFaster.Caching.UnitTests")]
+
+namespace BitFaster.Caching.UnitTests
+{
+    public class AssemblyInitialize : XunitTestFramework
+    {
+        public AssemblyInitialize(IMessageSink messageSink)
+            : base(messageSink)
+        {
+            ThreadPool.SetMinThreads(16, 16);
+        }
+
+        public new void Dispose()
+        {
+            // Place tear down code here
+            base.Dispose();
+        }
+    }
+}

BitFaster.Caching.UnitTests/ReferenceCountTests.cs

@@ -52,21 +52,5 @@ public void WhenObjectsAreDifferentHashcodesAreDifferent()
 
             a.GetHashCode().Should().NotBe(b.GetHashCode());
         }
-
-        [Fact]
-        public void WhenObjectDisposed()
-        {
-            var a = new ReferenceCount<Disposable>(new Disposable());
-            var b = a.DecrementCopy();
-
-            b.Invoking(rc => rc.IncrementCopy()).Should().Throw<ObjectDisposedException>();
-        }
-
-        private class Disposable : IDisposable
-        {
-            public void Dispose()
-            {
-            }
-        }
     }
 }

BitFaster.Caching.UnitTests/SingletonCacheTests.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using FluentAssertions;
@@ -145,21 +146,48 @@ public void WhenValueIsDisposableItIsDisposedWhenReleased()
         {
             var cache = new SingletonCache<string, DisposeTest>();
 
+            DisposeTest value = null;
+
             using (var lifetime = cache.Acquire("Foo"))
             {
-                DisposeTest.WasDisposed.Should().BeFalse();
+                value = lifetime.Value;
+                value.IsDisposed.Should().BeFalse();
             }
 
-            DisposeTest.WasDisposed.Should().BeTrue();
+            value.IsDisposed.Should().BeTrue();
+        }
+
+        [Fact]
+        public async Task DisposeOnManyDifferentThreadsAlwaysReturnsActiveValue()
+        {
+            var cache = new SingletonCache<string, DisposeTest>();
+
+            var tasks = Enumerable.Range(0, 64).Select(i => Task.Run(() => 
+            { 
+                using (var handle = cache.Acquire("Foo"))
+                { 
+                    handle.Value.ThrowIfDisposed();    
+                }
+            }));
+
+            await Task.WhenAll(tasks);
         }
 
         public class DisposeTest : IDisposable
         {
-            public static bool WasDisposed { get; set; }
+            public bool IsDisposed { get; private set; }
+
+            public void ThrowIfDisposed()
+            {
+                if (this.IsDisposed)
+                {
+                    throw new ObjectDisposedException("Error");
+                }
+            }
 
             public void Dispose()
             {
-                WasDisposed = true;
+                this.IsDisposed = true;
             }
         }
     }

BitFaster.Caching/BitFaster.Caching.csproj

@@ -8,7 +8,7 @@
     <Description>High performance, thread-safe in-memory caching primitives for .NET.</Description>
     <PackageLicenseFile>LICENSE</PackageLicenseFile>
     <GeneratePackageOnBuild>true</GeneratePackageOnBuild>
-    <Version>1.0.2</Version>
+    <Version>1.0.3</Version>
     <Copyright>Copyright © Alex Peck 2020</Copyright>
     <PackageProjectUrl></PackageProjectUrl>
     <RepositoryUrl>https://github.com/bitfaster/BitFaster.Caching</RepositoryUrl>
@@ -18,8 +18,8 @@
     <RootNamespace>BitFaster.Caching</RootNamespace>
     <IncludeSource>True</IncludeSource>
     <IncludeSymbols>True</IncludeSymbols>
-    <AssemblyVersion>1.0.2.0</AssemblyVersion>
-    <FileVersion>1.0.2.0</FileVersion>
+    <AssemblyVersion>1.0.3.0</AssemblyVersion>
+    <FileVersion>1.0.3.0</FileVersion>
   </PropertyGroup>
 
   <ItemGroup>

BitFaster.Caching/ReferenceCount.cs

@@ -53,11 +53,6 @@ public int Count
         /// <returns>A copy of the ReferenceCount with the count incremented by 1.</returns>
         public ReferenceCount<TValue> IncrementCopy()
         {
-            if (this.count <= 0 && this.value is IDisposable)
-            {
-                throw new ObjectDisposedException($"{typeof(TValue).Name} is disposed.");
-            }
-
             return new ReferenceCount<TValue>(this.value, this.count + 1);
         }
 

BitFaster.Caching/SingletonCache.cs

@@ -59,20 +59,32 @@ private void Release(TKey key)
         {
             while (true)
             {
-                var oldRefCount = this.cache[key];
-                var newRefCount = oldRefCount.DecrementCopy();
-                if (this.cache.TryUpdate(key, newRefCount, oldRefCount))
+                if (!this.cache.TryGetValue(key, out var oldRefCount))
                 {
-                    if (newRefCount.Count == 0)
+                    // already released, exit
+                    break;
+                }
+
+                // if count is 1, the caller's decrement makes refcount 0: it is unreferenced and eligible to remove
+                if (oldRefCount.Count == 1)
+                {
+                    var kvp = new KeyValuePair<TKey, ReferenceCount<TValue>>(key, oldRefCount);
+
+                    // hidden atomic remove
+                    // https://devblogs.microsoft.com/pfxteam/little-known-gems-atomic-conditional-removals-from-concurrentdictionary/
+                    if (((ICollection<KeyValuePair<TKey, ReferenceCount<TValue>>>)this.cache).Remove(kvp))
                     {
-                        if (((IDictionary<TKey, ReferenceCount<TValue>>)this.cache).Remove(new KeyValuePair<TKey, ReferenceCount<TValue>>(key, newRefCount)))
+                        // no longer in cache, safe to dispose and exit
+                        if (oldRefCount.Value is IDisposable d)
                         {
-                            if (newRefCount.Value is IDisposable d)
-                            {
-                                d.Dispose();
-                            }
+                            d.Dispose();
                         }
+                        break;
                     }
+                }
+                else if (this.cache.TryUpdate(key, oldRefCount.DecrementCopy(), oldRefCount))
+                {
+                    // replaced with decremented copy, exit
                     break;
                 }
             }