[release/5.0-rc2] Memory Leak in Microsoft.Extensions.Caching.Memory …

…when handling exceptions (#42494)

* Memory Leak in Microsoft.Extensions.Caching.Memory when handling exceptions

When an exception is thrown inside MemoryCache.GetOrCreate, we are leaking CacheEntry objects. This is because they are not being Disposed properly, and the async local CacheEntryStack is growing indefinitely.

The fix is to ensure the CacheEntry objects are disposed correctly. In order to do this, I set a flag to indicate whether the CacheEntry.Value has been set. If it hasn't, Disposing the CacheEntry won't add it to the cache.

Fix #42321

* Fix another memory leak when one cache depends on another cache. The inner cache's entries will reference the outer cache entries through the ScopeLease object.

Null'ing out the CacheEntry._scope field when it is disposed fixes this issue.

Co-authored-by: Eric Erhardt <[email protected]>

src/libraries/Microsoft.Extensions.Caching.Abstractions/src/MemoryCacheExtensions.cs

@@ -43,69 +43,60 @@ public static bool TryGetValue<TItem>(this IMemoryCache cache, object key, out T
 
         public static TItem Set<TItem>(this IMemoryCache cache, object key, TItem value)
         {
-            ICacheEntry entry = cache.CreateEntry(key);
+            using ICacheEntry entry = cache.CreateEntry(key);
             entry.Value = value;
-            entry.Dispose();
 
             return value;
         }
 
         public static TItem Set<TItem>(this IMemoryCache cache, object key, TItem value, DateTimeOffset absoluteExpiration)
         {
-            ICacheEntry entry = cache.CreateEntry(key);
+            using ICacheEntry entry = cache.CreateEntry(key);
             entry.AbsoluteExpiration = absoluteExpiration;
             entry.Value = value;
-            entry.Dispose();
 
             return value;
         }
 
         public static TItem Set<TItem>(this IMemoryCache cache, object key, TItem value, TimeSpan absoluteExpirationRelativeToNow)
         {
-            ICacheEntry entry = cache.CreateEntry(key);
+            using ICacheEntry entry = cache.CreateEntry(key);
             entry.AbsoluteExpirationRelativeToNow = absoluteExpirationRelativeToNow;
             entry.Value = value;
-            entry.Dispose();
 
             return value;
         }
 
         public static TItem Set<TItem>(this IMemoryCache cache, object key, TItem value, IChangeToken expirationToken)
         {
-            ICacheEntry entry = cache.CreateEntry(key);
+            using ICacheEntry entry = cache.CreateEntry(key);
             entry.AddExpirationToken(expirationToken);
             entry.Value = value;
-            entry.Dispose();
 
             return value;
         }
 
         public static TItem Set<TItem>(this IMemoryCache cache, object key, TItem value, MemoryCacheEntryOptions options)
         {
-            using (ICacheEntry entry = cache.CreateEntry(key))
+            using ICacheEntry entry = cache.CreateEntry(key);
+            if (options != null)
             {
-                if (options != null)
-                {
-                    entry.SetOptions(options);
-                }
-
-                entry.Value = value;
+                entry.SetOptions(options);
             }
 
+            entry.Value = value;
+
             return value;
         }
 
         public static TItem GetOrCreate<TItem>(this IMemoryCache cache, object key, Func<ICacheEntry, TItem> factory)
         {
             if (!cache.TryGetValue(key, out object result))
             {
-                ICacheEntry entry = cache.CreateEntry(key);
+                using ICacheEntry entry = cache.CreateEntry(key);
+
                 result = factory(entry);
-                entry.SetValue(result);
-                // need to manually call dispose instead of having a using
-                // in case the factory passed in throws, in which case we
-                // do not want to add the entry to the cache
-                entry.Dispose();
+                entry.Value = result;
             }
 
             return (TItem)result;
@@ -115,13 +106,10 @@ public static async Task<TItem> GetOrCreateAsync<TItem>(this IMemoryCache cache,
         {
             if (!cache.TryGetValue(key, out object result))
             {
-                ICacheEntry entry = cache.CreateEntry(key);
+                using ICacheEntry entry = cache.CreateEntry(key);
+
                 result = await factory(entry).ConfigureAwait(false);
-                entry.SetValue(result);
-                // need to manually call dispose instead of having a using
-                // in case the factory passed in throws, in which case we
-                // do not want to add the entry to the cache
-                entry.Dispose();
+                entry.Value = result;
             }
 
             return (TItem)result;

src/libraries/Microsoft.Extensions.Caching.Memory/src/CacheEntry.cs

@@ -12,10 +12,10 @@ namespace Microsoft.Extensions.Caching.Memory
 {
     internal class CacheEntry : ICacheEntry
     {
-        private bool _added;
+        private bool _disposed;
         private static readonly Action<object> ExpirationCallback = ExpirationTokensExpired;
         private readonly Action<CacheEntry> _notifyCacheOfExpiration;
-        private readonly Action<CacheEntry> _notifyCacheEntryDisposed;
+        private readonly Action<CacheEntry> _notifyCacheEntryCommit;
         private IList<IDisposable> _expirationTokenRegistrations;
         private IList<PostEvictionCallbackRegistration> _postEvictionCallbacks;
         private bool _isExpired;
@@ -27,12 +27,14 @@ internal class CacheEntry : ICacheEntry
         private TimeSpan? _slidingExpiration;
         private long? _size;
         private IDisposable _scope;
+        private object _value;
+        private bool _valueHasBeenSet;
 
         internal readonly object _lock = new object();
 
         internal CacheEntry(
             object key,
-            Action<CacheEntry> notifyCacheEntryDisposed,
+            Action<CacheEntry> notifyCacheEntryCommit,
             Action<CacheEntry> notifyCacheOfExpiration,
             ILogger logger)
         {
@@ -41,9 +43,9 @@ internal CacheEntry(
                 throw new ArgumentNullException(nameof(key));
             }
 
-            if (notifyCacheEntryDisposed == null)
+            if (notifyCacheEntryCommit == null)
             {
-                throw new ArgumentNullException(nameof(notifyCacheEntryDisposed));
+                throw new ArgumentNullException(nameof(notifyCacheEntryCommit));
             }
 
             if (notifyCacheOfExpiration == null)
@@ -57,7 +59,7 @@ internal CacheEntry(
             }
 
             Key = key;
-            _notifyCacheEntryDisposed = notifyCacheEntryDisposed;
+            _notifyCacheEntryCommit = notifyCacheEntryCommit;
             _notifyCacheOfExpiration = notifyCacheOfExpiration;
 
             _scope = CacheEntryHelper.EnterScope(this);
@@ -182,20 +184,39 @@ public long? Size
 
         public object Key { get; private set; }
 
-        public object Value { get; set; }
+        public object Value
+        {
+            get => _value;
+            set
+            {
+                _value = value;
+                _valueHasBeenSet = true;
+            }
+        }
 
         internal DateTimeOffset LastAccessed { get; set; }
 
         internal EvictionReason EvictionReason { get; private set; }
 
         public void Dispose()
         {
-            if (!_added)
+            if (!_disposed)
             {
-                _added = true;
+                _disposed = true;
+
+                // Ensure the _scope reference is cleared because it can reference other CacheEntry instances.
+                // This CacheEntry is going to be put into a MemoryCache, and we don't want to root unnecessary objects.
                 _scope.Dispose();
-                _notifyCacheEntryDisposed(this);
-                PropagateOptions(CacheEntryHelper.Current);
+                _scope = null;
+
+                // Don't commit or propagate options if the CacheEntry Value was never set.
+                // We assume an exception occurred causing the caller to not set the Value successfully,
+                // so don't use this entry.
+                if (_valueHasBeenSet)
+                {
+                    _notifyCacheEntryCommit(this);
+                    PropagateOptions(CacheEntryHelper.Current);
+                }
             }
         }
 

src/libraries/Microsoft.Extensions.Caching.Memory/tests/MemoryCacheSetAndRemoveTests.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
+using System.Reflection;
 using System.Threading;
 using System.Threading.Tasks;
 using Xunit;
@@ -180,6 +181,9 @@ public void GetOrCreate_WillNotCreateEmptyValue_WhenFactoryThrows()
             }
 
             Assert.False(cache.TryGetValue(key, out int obj));
+
+            // verify that throwing an exception doesn't leak CacheEntry objects
+            Assert.Null(CacheEntryHelper.Current);
         }
 
         [Fact]
@@ -199,6 +203,28 @@ await cache.GetOrCreateAsync<int>(key, entry =>
             }
 
             Assert.False(cache.TryGetValue(key, out int obj));
+
+            // verify that throwing an exception doesn't leak CacheEntry objects
+            Assert.Null(CacheEntryHelper.Current);
+        }
+
+        [Fact]
+        public void DisposingCacheEntryReleasesScope()
+        {
+            object GetScope(ICacheEntry entry)
+            {
+                return entry.GetType()
+                    .GetField("_scope", BindingFlags.NonPublic | BindingFlags.Instance)
+                    .GetValue(entry);
+            }
+
+            var cache = CreateCache();
+
+            ICacheEntry entry = cache.CreateEntry("myKey");
+            Assert.NotNull(GetScope(entry));
+
+            entry.Dispose();
+            Assert.Null(GetScope(entry));
         }
 
         [Fact]