HybridCache: ensure that Size is always specified in L1

src/Libraries/Microsoft.Extensions.Caching.Hybrid/Internal/DefaultHybridCache.CacheItem.cs

@@ -87,6 +87,8 @@ protected virtual void OnFinalRelease() // any required release semantics
 
     internal abstract class CacheItem<T> : CacheItem
     {
+        public abstract bool TryGetSize(out long size);
+
         // attempt to get a value that was *not* previously reserved
         public abstract bool TryGetValue(out T value);
 

src/Libraries/Microsoft.Extensions.Caching.Hybrid/Internal/DefaultHybridCache.ImmutableCacheItem.cs

@@ -13,6 +13,8 @@ internal partial class DefaultHybridCache
 
         private T _value = default!; // deferred until SetValue
 
+        public long Size { get; private set; } = -1;
+
         public override bool DebugIsImmutable => true;
 
         // get a shared instance that passes as "reserved"; doesn't need to be 100% singleton,
@@ -30,14 +32,24 @@ public static ImmutableCacheItem<T> GetReservedShared()
             return obj;
         }
 
-        public void SetValue(T value) => _value = value;
+        public void SetValue(T value, long size)
+        {
+            _value = value;
+            Size = size;
+        }
 
         public override bool TryGetValue(out T value)
         {
             value = _value;
             return true; // always available
         }
 
+        public override bool TryGetSize(out long size)
+        {
+            size = Size;
+            return size >= 0;
+        }
+
         public override bool TryReserveBuffer(out BufferChunk buffer)
         {
             buffer = default;

src/Libraries/Microsoft.Extensions.Caching.Hybrid/Internal/DefaultHybridCache.L2.cs

@@ -95,13 +95,26 @@ internal void SetL1<T>(string key, CacheItem<T> value, HybridCacheEntryOptions?
         if (value.TryReserve())
         {
             // based on CacheExtensions.Set<TItem>, but with post-eviction recycling
-            using var cacheEntry = _localCache.CreateEntry(key);
+
+            // intentionally use manual Dispose rather than "using"; confusingly, it is Dispose()
+            // that actually commits the add - so: if we fault, we don't want to try
+            // committing a partially configured cache entry
+            var cacheEntry = _localCache.CreateEntry(key);
             cacheEntry.AbsoluteExpirationRelativeToNow = options?.LocalCacheExpiration ?? _defaultLocalCacheExpiration;
             cacheEntry.Value = value;
+
+            if (value.TryGetSize(out var size))
+            {
+                cacheEntry = cacheEntry.SetSize(size);
+            }
+
             if (value.NeedsEvictionCallback)
             {
-                _ = cacheEntry.RegisterPostEvictionCallback(CacheItem.SharedOnEviction);
+                cacheEntry = cacheEntry.RegisterPostEvictionCallback(CacheItem.SharedOnEviction);
             }
+
+            // commit
+            cacheEntry.Dispose();
         }
     }
 

src/Libraries/Microsoft.Extensions.Caching.Hybrid/Internal/DefaultHybridCache.MutableCacheItem.cs

@@ -21,34 +21,38 @@ public void SetValue(ref BufferChunk buffer, IHybridCacheSerializer<T> serialize
             buffer = default; // we're taking over the lifetime; the caller no longer has it!
         }
 
-        public void SetValue(T value, IHybridCacheSerializer<T> serializer, int maxLength)
-        {
-            _serializer = serializer;
-            var writer = RecyclableArrayBufferWriter<byte>.Create(maxLength);
-            serializer.Serialize(value, writer);
-
-            _buffer = new(writer.DetachCommitted(out var length), length, returnToPool: true);
-            writer.Dispose(); // no buffers left (we just detached them), but just in case of other logic
-        }
-
         public override bool TryGetValue(out T value)
         {
             // only if we haven't already burned
-            if (!TryReserve())
+            if (TryReserve())
             {
-                value = default!;
-                return false;
+                try
+                {
+                    value = _serializer.Deserialize(_buffer.AsSequence());
+                    return true;
+                }
+                finally
+                {
+                    _ = Release();
+                }
             }
 
-            try
-            {
-                value = _serializer.Deserialize(_buffer.AsSequence());
-                return true;
-            }
-            finally
+            value = default!;
+            return false;
+        }
+
+        public override bool TryGetSize(out long size)
+        {
+            // only if we haven't already burned
+            if (TryReserve())
             {
+                size = _buffer.Length;
                 _ = Release();
+                return true;
             }
+
+            size = 0;
+            return false;
         }
 
         public override bool TryReserveBuffer(out BufferChunk buffer)

src/Libraries/Microsoft.Extensions.Caching.Hybrid/Internal/DefaultHybridCache.StampedeStateT.cs

@@ -6,6 +6,7 @@
 using System.Diagnostics.CodeAnalysis;
 using System.Threading;
 using System.Threading.Tasks;
+using static Microsoft.Extensions.Caching.Hybrid.Internal.DefaultHybridCache;
 
 namespace Microsoft.Extensions.Caching.Hybrid.Internal;
 
@@ -175,29 +176,64 @@ private async Task BackgroundFetchAsync()
                 // nothing from L2; invoke the underlying data store
                 if ((Key.Flags & HybridCacheEntryFlags.DisableUnderlyingData) == 0)
                 {
-                    var cacheItem = SetResult(await _underlying!(_state!, SharedToken).ConfigureAwait(false));
-
-                    // note that at this point we've already released most or all of the waiting callers; everything
-                    // else here is background
-
-                    // write to L2 if appropriate
-                    if ((Key.Flags & HybridCacheEntryFlags.DisableDistributedCacheWrite) == 0)
+                    // invoke the callback supplied by the caller
+                    var newValue = await _underlying!(_state!, SharedToken).ConfigureAwait(false);
+
+                    // If we're writing this value *anywhere*, we're going to need to serialize; this is obvious
+                    // in the case of L2, but we also need it for L1, because MemoryCache might be enforcing
+                    // SizeLimit (we can't know - it is an abstraction), and for *that* we need to know the item size.
+                    // Likewise, if we're writing to a MutableCacheItem, we'll be serializing *anyway* for the payload.
+                    //
+                    // Rephrasing that: the only scenario in which we *do not* need to serialize is if:
+                    // - it is an ImmutableCacheItem
+                    // - we're writing neither to L1 nor L2
+
+                    const HybridCacheEntryFlags DisableL1AndL2 = HybridCacheEntryFlags.DisableLocalCacheWrite | HybridCacheEntryFlags.DisableDistributedCacheWrite;
+                    var cacheItem = CacheItem;
+                    bool skipSerialize = cacheItem is ImmutableCacheItem<T> && (Key.Flags & DisableL1AndL2) == DisableL1AndL2;
+
+                    if (skipSerialize)
                     {
-                        if (cacheItem.TryReserveBuffer(out var buffer))
-                        {
-                            // mutable: we've already serialized it for the shared cache item
-                            await Cache.SetL2Async(Key.Key, in buffer, _options, SharedToken).ConfigureAwait(false);
-                            _ = cacheItem.Release(); // because we reserved
-                        }
-                        else if (cacheItem.TryGetValue(out var value))
+                        SetImmutableResultWithoutSerialize(newValue);
+                    }
+                    else if (cacheItem.TryReserve())
+                    {
+                        // ^^^ The first thing we need to do is make sure we're not getting into a thread race over buffer disposal.
+                        // In particular, if this cache item is somehow so short-lived that the buffers would be released *before* we're
+                        // done writing them to L2, which happens *after* we've provided the value to consumers.
+                        var writer = RecyclableArrayBufferWriter<byte>.Create(MaximumPayloadBytes); // note this lifetime spans the SetL2Async
+                        var serializer = Cache.GetSerializer<T>();
+                        serializer.Serialize(newValue, writer);
+                        BufferChunk buffer = new(writer.DetachCommitted(out var length), length, returnToPool: true); // remove buffer ownership from the writer
+                        writer.Dispose(); // we're done with the writer
+
+                        // protect "buffer" (this is why we "reserved"); we don't want SetResult to nuke our local
+                        var snapshot = buffer;
+                        SetResultPreSerialized(newValue, ref snapshot, serializer);
+
+                        // Note that at this point we've already released most or all of the waiting callers. Everything
+                        // from this point onwards happens in the background, from the perspective of the calling code.
+
+                        // Write to L2 if appropriate.
+                        if ((Key.Flags & HybridCacheEntryFlags.DisableDistributedCacheWrite) == 0)
                         {
-                            // immutable: we'll need to do the serialize ourselves
-                            var writer = RecyclableArrayBufferWriter<byte>.Create(MaximumPayloadBytes); // note this lifetime spans the SetL2Async
-                            Cache.GetSerializer<T>().Serialize(value, writer);
-                            buffer = new(writer.GetBuffer(out var length), length, returnToPool: false); // writer still owns the buffer
+                            // We already have the payload serialized, so this is trivial to do.
                             await Cache.SetL2Async(Key.Key, in buffer, _options, SharedToken).ConfigureAwait(false);
-                            writer.Dispose(); // recycle on success
                         }
+
+                        // Release our hook on the CacheItem (only really important for "mutable").
+                        _ = cacheItem.Release();
+
+                        // Finally, recycle whatever was left over from SetResultPreSerialized; using "snapshot"
+                        // here is NOT a typo; if SetResultPreSerialized left this value alone (immutable), then
+                        // this is our recycle step; if SetResultPreSerialized transferred ownership to the (mutable)
+                        // CacheItem, then this becomes a no-op, and the buffer only gets fully recycled when the
+                        // CacheItem itself is fully clear.
+                        snapshot.RecycleIfAppropriate();
+                    }
+                    else
+                    {
+                        throw new InvalidOperationException("Internal HybridCache failure: unable to reserve cache item to assign result");
                     }
                 }
                 else
@@ -243,7 +279,7 @@ private void SetResultAndRecycleIfAppropriate(ref BufferChunk value)
             {
                 case ImmutableCacheItem<T> immutable:
                     // deserialize; and store object; buffer can be recycled now
-                    immutable.SetValue(serializer.Deserialize(new(value.Array!, 0, value.Length)));
+                    immutable.SetValue(serializer.Deserialize(new(value.Array!, 0, value.Length)), value.Length);
                     value.RecycleIfAppropriate();
                     cacheItem = immutable;
                     break;
@@ -261,20 +297,41 @@ private void SetResultAndRecycleIfAppropriate(ref BufferChunk value)
             SetResult(cacheItem);
         }
 
-        private CacheItem<T> SetResult(T value)
+        private void SetImmutableResultWithoutSerialize(T value)
         {
             // set a result from a value we calculated directly
             CacheItem<T> cacheItem;
             switch (CacheItem)
             {
                 case ImmutableCacheItem<T> immutable:
                     // no serialize needed
-                    immutable.SetValue(value);
+                    immutable.SetValue(value, size: -1);
+                    cacheItem = immutable;
+                    break;
+                default:
+                    cacheItem = ThrowUnexpectedCacheItem();
+                    break;
+            }
+
+            SetResult(cacheItem);
+        }
+
+        private void SetResultPreSerialized(T value, ref BufferChunk buffer, IHybridCacheSerializer<T> serializer)
+        {
+            // set a result from a value we calculated directly that
+            // has ALREADY BEEN SERIALIZED (we can optionally consume this buffer)
+            CacheItem<T> cacheItem;
+            switch (CacheItem)
+            {
+                case ImmutableCacheItem<T> immutable:
+                    // no serialize needed
+                    immutable.SetValue(value, size: buffer.Length);
                     cacheItem = immutable;
+
+                    // (but leave the buffer alone)
                     break;
                 case MutableCacheItem<T> mutable:
-                    // serialization happens here
-                    mutable.SetValue(value, Cache.GetSerializer<T>(), MaximumPayloadBytes);
+                    mutable.SetValue(ref buffer, serializer);
                     mutable.DebugOnlyTrackBuffer(Cache);
                     cacheItem = mutable;
                     break;
@@ -284,7 +341,6 @@ private CacheItem<T> SetResult(T value)
             }
 
             SetResult(cacheItem);
-            return cacheItem;
         }
 
         private void SetResult(CacheItem<T> value)

test/Libraries/Microsoft.Extensions.Caching.Hybrid.Tests/SizeTests.cs

@@ -0,0 +1,95 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using Microsoft.Extensions.Caching.Hybrid.Internal;
+using Microsoft.Extensions.DependencyInjection;
+
+namespace Microsoft.Extensions.Caching.Hybrid.Tests;
+
+public class SizeTests
+{
+    [Theory]
+    [InlineData(null, true)] // does not enforce size limits
+    [InlineData(8L, false)] // unreasonably small limit; chosen because our test string has length 12 - hence no expectation to find the second time
+    [InlineData(1024L, true)] // reasonable size limit
+    public async Task ValidateSizeLimit_Immutable(long? sizeLimit, bool expectFromL1)
+    {
+        var services = new ServiceCollection();
+        services.AddMemoryCache(options => options.SizeLimit = sizeLimit);
+        services.AddHybridCache();
+        using var provider = services.BuildServiceProvider();
+        var cache = Assert.IsType<DefaultHybridCache>(provider.GetRequiredService<HybridCache>());
+
+        const string Key = "abc";
+
+        // this looks weird; it is intentionally not a const - we want to check
+        // same instance without worrying about interning from raw literals
+        string expected = new("simple value".ToArray());
+        var actual = await cache.GetOrCreateAsync<string>(Key, ct => new(expected));
+
+        // expect same contents
+        Assert.Equal(expected, actual);
+
+        // expect same instance, because string is special-cased as a type
+        // that doesn't need defensive copies
+        Assert.Same(expected, actual);
+
+        // rinse and repeat, to check we get the value from L1
+        actual = await cache.GetOrCreateAsync<string>(Key, ct => new(Guid.NewGuid().ToString()));
+
+        if (expectFromL1)
+        {
+            // expect same contents from L1
+            Assert.Equal(expected, actual);
+
+            // expect same instance, because string is special-cased as a type
+            // that doesn't need defensive copies
+            Assert.Same(expected, actual);
+        }
+        else
+        {
+            // L1 cache not used
+            Assert.NotEqual(expected, actual);
+        }
+    }
+
+    [Theory]
+    [InlineData(null, true)] // does not enforce size limits
+    [InlineData(8L, false)] // unreasonably small limit; chosen because our test string has length 12 - hence no expectation to find the second time
+    [InlineData(1024L, true)] // reasonable size limit
+    public async Task ValidateSizeLimit_Mutable(long? sizeLimit, bool expectFromL1)
+    {
+        var services = new ServiceCollection();
+        services.AddMemoryCache(options => options.SizeLimit = sizeLimit);
+        services.AddHybridCache();
+        using var provider = services.BuildServiceProvider();
+        var cache = Assert.IsType<DefaultHybridCache>(provider.GetRequiredService<HybridCache>());
+
+        const string Key = "abc";
+
+        string expected = "simple value";
+        var actual = await cache.GetOrCreateAsync<MutablePoco>(Key, ct => new(new MutablePoco { Value = expected }));
+
+        // expect same contents
+        Assert.Equal(expected, actual.Value);
+
+        // rinse and repeat, to check we get the value from L1
+        actual = await cache.GetOrCreateAsync<MutablePoco>(Key, ct => new(new MutablePoco { Value = Guid.NewGuid().ToString() }));
+
+        if (expectFromL1)
+        {
+            // expect same contents from L1
+            Assert.Equal(expected, actual.Value);
+        }
+        else
+        {
+            // L1 cache not used
+            Assert.NotEqual(expected, actual.Value);
+        }
+    }
+
+    public class MutablePoco
+    {
+        public string Value { get; set; } = "";
+    }
+}