HybridCache - minor post-PR fixes (#55251)

* additional HybridCache post-PR fixes

Author: mgravell

src/Caching/Hybrid/src/HybridCacheBuilderExtensions.cs

@@ -1,12 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using System;
-using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
-using System.Linq;
-using System.Text;
-using System.Threading.Tasks;
 using Microsoft.Extensions.DependencyInjection;
 
 namespace Microsoft.Extensions.Caching.Hybrid;

src/Caching/Hybrid/src/HybridCacheOptions.cs

@@ -1,13 +1,6 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using System;
-using System.Collections.Generic;
-using System.Linq;
-using System.Text;
-using System.Threading.Tasks;
-using Microsoft.Extensions.Options;
-
 namespace Microsoft.Extensions.Caching.Hybrid;
 
 /// <summary>

src/Caching/Hybrid/src/HybridCacheServiceExtensions.cs

@@ -2,15 +2,9 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
-using System.Collections.Generic;
-using System.Diagnostics.CodeAnalysis;
-using System.Linq;
-using System.Text;
-using System.Threading.Tasks;
 using Microsoft.Extensions.Caching.Hybrid.Internal;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.DependencyInjection.Extensions;
-using Microsoft.Extensions.Internal;
 
 namespace Microsoft.Extensions.Caching.Hybrid;
 

src/Caching/Hybrid/src/IHybridCacheBuilder.cs

@@ -1,11 +1,6 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using System;
-using System.Collections.Generic;
-using System.Linq;
-using System.Text;
-using System.Threading.Tasks;
 using Microsoft.Extensions.DependencyInjection;
 
 namespace Microsoft.Extensions.Caching.Hybrid;

src/Caching/Hybrid/src/Internal/BufferChunk.cs

@@ -9,7 +9,8 @@
 namespace Microsoft.Extensions.Caching.Hybrid.Internal;
 
 // used to convey buffer status; like ArraySegment<byte>, but Offset is always
-// zero, and we use the MSB of the length to track whether or not to recycle this value
+// zero, and we use the most significant bit (MSB, the sign flag) of the length
+// to track whether or not to recycle this value
 internal readonly struct BufferChunk
 {
     private const int MSB = (1 << 31);
@@ -40,6 +41,10 @@ public BufferChunk(byte[] array)
         Array = array;
         _lengthAndPoolFlag = array.Length;
         // assume not pooled, if exact-sized
+        // (we don't expect array.Length to be negative; we're really just saying
+        // "we expect the result of assigning array.Length to _lengthAndPoolFlag
+        // to give the expected Length *and* not have the MSB set; we're just
+        // checking that we haven't fat-fingered our MSB logic)
         Debug.Assert(!ReturnToPool, "do not return right-sized arrays");
         Debug.Assert(Length == array.Length, "array length not respected");
     }

src/Caching/Hybrid/src/Internal/DefaultHybridCache.CacheItem.cs

@@ -2,6 +2,8 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
+using System.Diagnostics;
+using System.Threading;
 using Microsoft.Extensions.Caching.Memory;
 
 namespace Microsoft.Extensions.Caching.Hybrid.Internal;
@@ -10,37 +12,85 @@ partial class DefaultHybridCache
 {
     internal abstract class CacheItem
     {
+        private int _refCount = 1; // the number of pending operations against this cache item
+        // note: the ref count is the number of callers anticipating this value at any given time; initially,
+        // it is one for a simple "get the value" flow, but if another call joins with us, it'll be incremented;
+        // if either cancels, it will get decremented, with the entire flow being cancelled if it ever becomes
+        // zero
+        // this counter also drives cache lifetime, with the cache itself incrementing the count by one; in the
+        // case of mutable data, cache eviction may reduce this to zero (in cooperation with any concurrent readers,
+        // who incr/decr around their fetch), allowing safe buffer recycling
+
+        internal int RefCount => Volatile.Read(ref _refCount);
+
         internal static readonly PostEvictionDelegate _sharedOnEviction = static (key, value, reason, state) =>
         {
             if (value is CacheItem item)
             {
-                // perform per-item clean-up; this could be buffer recycling (if defensive copies needed),
-                // or could be disposal
-                item.OnEviction();
+                item.Release();
             }
         };
 
-        public virtual void Release() { } // for recycling purposes
+        public virtual bool NeedsEvictionCallback => false; // do we need to call Release when evicted?
 
-        public abstract bool NeedsEvictionCallback { get; } // do we need to call Release when evicted?
-
-        public virtual void OnEviction() { } // only invoked if NeedsEvictionCallback reported true
+        protected virtual void OnFinalRelease() { } // any required release semantics
 
         public abstract bool TryReserveBuffer(out BufferChunk buffer);
 
         public abstract bool DebugIsImmutable { get; }
+
+        public bool Release() // returns true ONLY for the final release step
+        {
+            var newCount = Interlocked.Decrement(ref _refCount);
+            Debug.Assert(newCount >= 0, "over-release detected");
+            if (newCount == 0)
+            {
+                // perform per-item clean-up, i.e. buffer recycling (if defensive copies needed)
+                OnFinalRelease();
+                return true;
+            }
+            return false;
+        }
+
+        public bool TryReserve()
+        {
+            // this is basically interlocked increment, but with a check against:
+            // a) incrementing upwards from zero
+            // b) overflowing *back* to zero
+            var oldValue = Volatile.Read(ref _refCount);
+            do
+            {
+                if (oldValue is 0 or -1)
+                {
+                    return false; // already burned, or about to roll around back to zero
+                }
+
+                var updated = Interlocked.CompareExchange(ref _refCount, oldValue + 1, oldValue);
+                if (updated == oldValue)
+                {
+                    return true; // we exchanged
+                }
+                oldValue = updated; // we failed, but we have an updated state
+            } while (true);
+        }
+
     }
 
     internal abstract class CacheItem<T> : CacheItem
     {
+        internal static CacheItem<T> Create() => ImmutableTypeCache<T>.IsImmutable ? new ImmutableCacheItem<T>() : new MutableCacheItem<T>();
+
+        // attempt to get a value that was *not* previously reserved
         public abstract bool TryGetValue(out T value);
 
-        public T GetValue()
+        // get a value that *was* reserved, countermanding our reservation in the process
+        public T GetReservedValue()
         {
             if (!TryGetValue(out var value))
             {
                 Throw();
             }
+            Release();
             return value;
 
             static void Throw() => throw new ObjectDisposedException("The cache item has been recycled before the value was obtained");

src/Caching/Hybrid/src/Internal/DefaultHybridCache.Debug.cs

@@ -24,45 +24,45 @@ internal bool DebugTryGetCacheItem(string key, [NotNullWhen(true)] out CacheItem
 
     private int _outstandingBufferCount;
 
-    internal int DebugGetOutstandingBuffers(bool flush = false)
+    internal int DebugOnlyGetOutstandingBuffers(bool flush = false)
                 => flush ? Interlocked.Exchange(ref _outstandingBufferCount, 0) : Volatile.Read(ref _outstandingBufferCount);
 
     [Conditional("DEBUG")]
-    internal void DebugDecrementOutstandingBuffers()
+    internal void DebugOnlyDecrementOutstandingBuffers()
     {
         Interlocked.Decrement(ref _outstandingBufferCount);
     }
 
     [Conditional("DEBUG")]
-    internal void DebugIncrementOutstandingBuffers()
+    internal void DebugOnlyIncrementOutstandingBuffers()
     {
         Interlocked.Increment(ref _outstandingBufferCount);
     }
 #endif
 
     partial class MutableCacheItem<T>
     {
-        partial void DebugDecrementOutstandingBuffers();
-        partial void DebugTrackBufferCore(DefaultHybridCache cache);
+        partial void DebugOnlyDecrementOutstandingBuffers();
+        partial void DebugOnlyTrackBufferCore(DefaultHybridCache cache);
 
         [Conditional("DEBUG")]
-        internal void DebugTrackBuffer(DefaultHybridCache cache) => DebugTrackBufferCore(cache);
+        internal void DebugOnlyTrackBuffer(DefaultHybridCache cache) => DebugOnlyTrackBufferCore(cache);
 
 #if DEBUG
         private DefaultHybridCache? _cache; // for buffer-tracking - only enabled in DEBUG
-        partial void DebugDecrementOutstandingBuffers()
+        partial void DebugOnlyDecrementOutstandingBuffers()
         {
             if (_buffer.ReturnToPool)
             {
-                _cache?.DebugDecrementOutstandingBuffers();
+                _cache?.DebugOnlyDecrementOutstandingBuffers();
             }
         }
-        partial void DebugTrackBufferCore(DefaultHybridCache cache)
+        partial void DebugOnlyTrackBufferCore(DefaultHybridCache cache)
         {
             _cache = cache;
             if (_buffer.ReturnToPool)
             {
-                _cache?.DebugIncrementOutstandingBuffers();
+                _cache?.DebugOnlyIncrementOutstandingBuffers();
             }
         }
 #endif

src/Caching/Hybrid/src/Internal/DefaultHybridCache.ImmutableCacheItem.cs

@@ -1,33 +1,34 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using System;
-using System.Diagnostics;
-using System.Diagnostics.CodeAnalysis;
+using System.Threading;
 
 namespace Microsoft.Extensions.Caching.Hybrid.Internal;
 
 partial class DefaultHybridCache
 {
     private sealed class ImmutableCacheItem<T> : CacheItem<T> // used to hold types that do not require defensive copies
     {
-        private readonly T _value;
-        public ImmutableCacheItem(T value) => _value = value;
+        private T _value = default!; // deferred until SetValue
 
-        private static ImmutableCacheItem<T>? _sharedDefault;
+        public void SetValue(T value) => _value = value;
 
-        // this is only used when the underlying store is disabled; we don't need 100% singleton; "good enough is"
-        public static ImmutableCacheItem<T> Default => _sharedDefault ??= new(default!);
+        private static ImmutableCacheItem<T>? _sharedDefault;
 
-        public override void OnEviction()
+        // get a shared instance that passes as "reserved"; doesn't need to be 100% singleton,
+        // but we don't want to break the reservation rules either; if we can't reserve: create new
+        public static ImmutableCacheItem<T> GetReservedShared()
         {
-            var obj = _value as IDisposable;
-            Debug.Assert(obj is not null, "shouldn't be here for non-disposable types");
-            obj?.Dispose();
+            var obj = Volatile.Read(ref _sharedDefault);
+            if (obj is null || !obj.TryReserve())
+            {
+                obj = new();
+                obj.TryReserve(); // this is reliable on a new instance
+                Volatile.Write(ref _sharedDefault, obj);
+            }
+            return obj;
         }
 
-        public override bool NeedsEvictionCallback => ImmutableTypeCache<T>.IsDisposable;
-
         public override bool TryGetValue(out T value)
         {
             value = _value;

src/Caching/Hybrid/src/Internal/DefaultHybridCache.L2.cs

@@ -113,13 +113,16 @@ private DistributedCacheEntryOptions GetOptions(HybridCacheEntryOptions? options
 
     internal void SetL1<T>(string key, CacheItem<T> value, HybridCacheEntryOptions? options)
     {
-        // based on CacheExtensions.Set<TItem>, but with post-eviction recycling
-        using var cacheEntry = _localCache.CreateEntry(key);
-        cacheEntry.AbsoluteExpirationRelativeToNow = options?.LocalCacheExpiration ?? _defaultLocalCacheExpiration;
-        cacheEntry.Value = value;
-        if (value.NeedsEvictionCallback)
+        if (value.TryReserve()) // incr ref-count for the the cache itself; this *may* be released via the NeedsEvictionCallback path
         {
-            cacheEntry.RegisterPostEvictionCallback(CacheItem._sharedOnEviction);
+            // based on CacheExtensions.Set<TItem>, but with post-eviction recycling
+            using var cacheEntry = _localCache.CreateEntry(key);
+            cacheEntry.AbsoluteExpirationRelativeToNow = options?.LocalCacheExpiration ?? _defaultLocalCacheExpiration;
+            cacheEntry.Value = value;
+            if (value.NeedsEvictionCallback)
+            {
+                cacheEntry.RegisterPostEvictionCallback(CacheItem._sharedOnEviction);
+            }
         }
     }
 }

src/Caching/Hybrid/src/Internal/DefaultHybridCache.MutableCacheItem.cs

@@ -1,28 +1,23 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using System;
-using System.Diagnostics;
-using System.Threading;
-
 namespace Microsoft.Extensions.Caching.Hybrid.Internal;
 
 partial class DefaultHybridCache
 {
     private sealed partial class MutableCacheItem<T> : CacheItem<T> // used to hold types that require defensive copies
     {
-        private readonly IHybridCacheSerializer<T> _serializer;
-        private readonly BufferChunk _buffer;
-        private int _refCount = 1; // buffer released when this becomes zero
+        private IHybridCacheSerializer<T> _serializer = null!; // deferred until SetValue
+        private BufferChunk _buffer;
 
-        public MutableCacheItem(ref BufferChunk buffer, IHybridCacheSerializer<T> serializer)
+        public void SetValue(ref BufferChunk buffer, IHybridCacheSerializer<T> serializer)
         {
             _serializer = serializer;
             _buffer = buffer;
             buffer = default; // we're taking over the lifetime; the caller no longer has it!
         }
 
-        public MutableCacheItem(T value, IHybridCacheSerializer<T> serializer, int maxLength)
+        public void SetValue(T value, IHybridCacheSerializer<T> serializer, int maxLength)
         {
             _serializer = serializer;
             var writer = RecyclableArrayBufferWriter<byte>.Create(maxLength);
@@ -34,35 +29,10 @@ public MutableCacheItem(T value, IHybridCacheSerializer<T> serializer, int maxLe
 
         public override bool NeedsEvictionCallback => _buffer.ReturnToPool;
 
-        public override void OnEviction() => Release();
-
-        public override void Release()
+        protected override void OnFinalRelease()
         {
-            var newCount = Interlocked.Decrement(ref _refCount);
-            if (newCount == 0)
-            {
-                DebugDecrementOutstandingBuffers();
-                _buffer.RecycleIfAppropriate();
-            }
-        }
-
-        public bool TryReserve()
-        {
-            var oldValue = Volatile.Read(ref _refCount);
-            do
-            {
-                if (oldValue is 0 or -1)
-                {
-                    return false; // already burned, or about to roll around back to zero
-                }
-
-                var updated = Interlocked.CompareExchange(ref _refCount, oldValue + 1, oldValue);
-                if (updated == oldValue)
-                {
-                    return true; // we exchanged
-                }
-                oldValue = updated; // we failed, but we have an updated state
-            } while (true);
+            DebugOnlyDecrementOutstandingBuffers();
+            _buffer.RecycleIfAppropriate();
         }
 
         public override bool TryGetValue(out T value)