[release/5.0-rc2] Memory Leak in Microsoft.Extensions.Caching.Memory when handling exceptions

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]