Remove heap allocations from Scoped<T> reference counting

BitFaster.Caching.UnitTests/ScopedTests.cs

@@ -8,6 +8,10 @@ namespace BitFaster.Caching.UnitTests
 {
     public class ScopedTests
     {
+#if NETCOREAPP3_1_OR_GREATER
+        private const long MaxExpectedBytesPerLifetime = 80L;
+#endif
+
         [Fact]
         public void WhenScopeIsCreatedThenScopeDisposedValueIsDisposed()
         {
@@ -49,6 +53,16 @@ public void WhenScopeIsCreatedThenLifetimeDisposedScopeDisposesValue()
             disposable.IsDisposed.Should().BeTrue();
         }
 
+        [Fact]
+        public void CreateLifetime_AfterIncrement_ReturnsIncrementedReferenceCount()
+        {
+            var scope = new Scoped<Disposable>(new Disposable());
+
+            using var lifetime = scope.CreateLifetime();
+
+            lifetime.ReferenceCount.Should().Be(2);
+        }
+
         [Fact]
         public void WhenScopeIsDisposedCreateScopeThrows()
         {
@@ -86,5 +100,28 @@ public void WhenScopedIsCreatedFromCacheItemHasExpectedLifetime()
 
             valueFactory.Disposable.IsDisposed.Should().BeTrue();
         }
+
+#if NETCOREAPP3_1_OR_GREATER
+        [Fact]
+        public void CreateLifetime_WhenCalledRepeatedly_DoesNotAllocateForReferenceCounting()
+        {
+            var scope = new Scoped<Disposable>(new Disposable());
+
+            using (scope.CreateLifetime())
+            {
+            }
+
+            long before = GC.GetAllocatedBytesForCurrentThread();
+
+            for (int i = 0; i < 256; i++)
+            {
+                using var lifetime = scope.CreateLifetime();
+            }
+
+            long allocated = GC.GetAllocatedBytesForCurrentThread() - before;
+
+            allocated.Should().BeLessThan(256 * MaxExpectedBytesPerLifetime);
+        }
+#endif
     }
 }

BitFaster.Caching/Lifetime.cs

@@ -4,15 +4,22 @@
 
 namespace BitFaster.Caching
 {
+    internal interface ILifetimeReleaser
+    {
+        void ReleaseLifetime();
+    }
+
     /// <summary>
     /// Represents the lifetime of a value. The value is alive and valid for use until the 
     /// lifetime is disposed.
     /// </summary>
     /// <typeparam name="T">The type of value</typeparam>
     public sealed class Lifetime<T> : IDisposable
     {
-        private readonly Action onDisposeAction;
-        private readonly ReferenceCount<T> refCount;
+        private readonly Action? releaseAction;
+        private readonly ILifetimeReleaser? releaser;
+        private readonly T value = default!;
+        private readonly int referenceCount;
         private bool isDisposed;
 
         /// <summary>
@@ -22,19 +29,27 @@ public sealed class Lifetime<T> : IDisposable
         /// <param name="onDisposeAction">The action to perform when the lifetime is terminated.</param>
         public Lifetime(ReferenceCount<T> value, Action onDisposeAction)
         {
-            this.refCount = value;
-            this.onDisposeAction = onDisposeAction;
+            this.value = value.Value;
+            this.referenceCount = value.Count;
+            this.releaseAction = onDisposeAction;
+        }
+
+        internal Lifetime(T value, int referenceCount, ILifetimeReleaser releaser)
+        {
+            this.value = value;
+            this.referenceCount = referenceCount;
+            this.releaser = releaser;
         }
 
         /// <summary>
         /// Gets the value.
         /// </summary>
-        public T Value => this.refCount.Value;
+        public T Value => this.value;
 
         /// <summary>
         /// Gets the count of Lifetime instances referencing the same value.
         /// </summary>
-        public int ReferenceCount => this.refCount.Count;
+        public int ReferenceCount => this.referenceCount;
 
         /// <summary>
         /// Terminates the lifetime and performs any cleanup required to release the value.
@@ -43,7 +58,15 @@ public void Dispose()
         {
             if (!this.isDisposed)
             {
-                this.onDisposeAction();
+                if (this.releaser is ILifetimeReleaser lifetimeReleaser)
+                {
+                    lifetimeReleaser.ReleaseLifetime();
+                }
+                else
+                {
+                    this.releaseAction!();
+                }
+
                 this.isDisposed = true;
             }
         }

BitFaster.Caching/Scoped.cs

@@ -13,24 +13,27 @@ namespace BitFaster.Caching
     /// <typeparam name="T">The type of scoped value.</typeparam>
     [DebuggerTypeProxy(typeof(Scoped<>.ScopedDebugView))]
     [DebuggerDisplay("{FormatDebug(),nq}")]
-    public sealed class Scoped<T> : IScoped<T>, IDisposable where T : IDisposable
+    public sealed class Scoped<T> : IScoped<T>, IDisposable, ILifetimeReleaser where T : IDisposable
     {
-        private ReferenceCount<T> refCount;
-        private int disposed = 0;
+        private const int DisposedFlag = unchecked((int)0x80000000);
+        private const int ReferenceCountMask = int.MaxValue;
+
+        private readonly T value;
+        private int state = 1;
 
         /// <summary>
         /// Initializes a new Scoped value.
         /// </summary>
         /// <param name="value">The value to scope.</param>
         public Scoped(T value)
         {
-            this.refCount = new ReferenceCount<T>(value);
+            this.value = value;
         }
 
         /// <summary>
         /// Gets a value indicating whether the scope is disposed.
         /// </summary>
-        public bool IsDisposed => Volatile.Read(ref this.disposed) == 1;
+        public bool IsDisposed => Volatile.Read(ref this.state) < 0;
 
         /// <summary>
         /// Attempts to create a lifetime for the scoped value. The lifetime guarantees the value is alive until 
@@ -42,19 +45,17 @@ public bool TryCreateLifetime([MaybeNullWhen(false)] out Lifetime<T> lifetime)
         {
             while (true)
             {
-                var oldRefCount = this.refCount;
+                int oldState = Volatile.Read(ref this.state);
 
-                // If old ref count is 0, the scoped object has been disposed and there was a race.
-                if (IsDisposed || oldRefCount.Count == 0)
+                if (oldState < 0)
                 {
                     lifetime = default;
                     return false;
                 }
 
-                if (oldRefCount == Interlocked.CompareExchange(ref this.refCount, oldRefCount.IncrementCopy(), oldRefCount))
+                if (Interlocked.CompareExchange(ref this.state, oldState + 1, oldState) == oldState)
                 {
-                    // When Lifetime is disposed, it calls DecrementReferenceCount
-                    lifetime = new Lifetime<T>(oldRefCount, this.DecrementReferenceCount);
+                    lifetime = new Lifetime<T>(this.value, (oldState & ReferenceCountMask) + 1, this);
                     return true;
                 }
             }
@@ -74,22 +75,23 @@ public Lifetime<T> CreateLifetime()
             return lifetime;
         }
 
-        private void DecrementReferenceCount()
+        void ILifetimeReleaser.ReleaseLifetime()
         {
             while (true)
             {
-                var oldRefCount = this.refCount;
+                int oldState = Volatile.Read(ref this.state);
+                int oldReferenceCount = oldState & ReferenceCountMask;
+                int newReferenceCount = oldReferenceCount - 1;
+                int newState = (oldState & DisposedFlag) | newReferenceCount;
 
-                if (oldRefCount == Interlocked.CompareExchange(ref this.refCount, oldRefCount.DecrementCopy(), oldRefCount))
+                if (Interlocked.CompareExchange(ref this.state, newState, oldState) == oldState)
                 {
-                    // Note this.refCount may be stale.
-                    // Old count == 1, thus new ref count is 0, dispose the value.
-                    if (oldRefCount.Count == 1)
+                    if (newReferenceCount == 0 && (oldState & DisposedFlag) != 0)
                     {
-                        oldRefCount.Value?.Dispose();
+                        this.value?.Dispose();
                     }
 
-                    break;
+                    return;
                 }
             }
         }
@@ -100,10 +102,28 @@ private void DecrementReferenceCount()
         /// </summary>
         public void Dispose()
         {
-            // Dispose exactly once, decrement via dispose exactly once
-            if (Interlocked.CompareExchange(ref this.disposed, 1, 0) == 0)
+            while (true)
             {
-                DecrementReferenceCount();
+                int oldState = Volatile.Read(ref this.state);
+
+                if ((oldState & DisposedFlag) != 0)
+                {
+                    return;
+                }
+
+                int oldReferenceCount = oldState & ReferenceCountMask;
+                int newReferenceCount = oldReferenceCount - 1;
+                int newState = DisposedFlag | newReferenceCount;
+
+                if (Interlocked.CompareExchange(ref this.state, newState, oldState) == oldState)
+                {
+                    if (newReferenceCount == 0)
+                    {
+                        this.value?.Dispose();
+                    }
+
+                    return;
+                }
             }
         }
 
@@ -115,7 +135,7 @@ internal string FormatDebug()
                 return "[Disposed Scope]";
             }
 
-            return this.refCount.Value?.ToString() ?? "[null]";
+            return this.value?.ToString() ?? "[null]";
         }
 
         [ExcludeFromCodeCoverage]
@@ -133,7 +153,7 @@ public ScopedDebugView(Scoped<T> scoped)
 
             public bool IsDisposed => this.scoped.IsDisposed;
 
-            public T Value => this.scoped.refCount.Value;
+            public T Value => this.scoped.value;
         }
     }
 }