[release/2.1] Memory Leak in Microsoft.Extensions.Caching.Memory when handling exceptions (dotnet/extensions#3536)

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

* Update patchConfig.props

* Try standardizing use of SetValue

* Apply feedback

* Fix syntax

* Undo breaking change

* Feedback

* Fixup

* Fixup patchConfig

* Add direct ref

* Fix another memory leak when one cache depends on another cache

* Fixup\n\nCommit migrated from dotnet/extensions@dc1dab7

Author: wtgodbe

src/Caching/Abstractions/src/MemoryCacheExtensions.cs

@@ -35,39 +35,43 @@ public static bool TryGetValue<TItem>(this IMemoryCache cache, object key, out T
 
         public static TItem Set<TItem>(this IMemoryCache cache, object key, TItem value)
         {
-            var entry = cache.CreateEntry(key);
-            entry.Value = value;
-            entry.Dispose();
+            using (var entry = cache.CreateEntry(key))
+            {
+                entry.Value = value;
+            }
 
             return value;
         }
 
         public static TItem Set<TItem>(this IMemoryCache cache, object key, TItem value, DateTimeOffset absoluteExpiration)
         {
-            var entry = cache.CreateEntry(key);
-            entry.AbsoluteExpiration = absoluteExpiration;
-            entry.Value = value;
-            entry.Dispose();
+            using (var entry = cache.CreateEntry(key))
+            {
+                entry.AbsoluteExpiration = absoluteExpiration;
+                entry.Value = value;
+            }
 
             return value;
         }
 
         public static TItem Set<TItem>(this IMemoryCache cache, object key, TItem value, TimeSpan absoluteExpirationRelativeToNow)
         {
-            var entry = cache.CreateEntry(key);
-            entry.AbsoluteExpirationRelativeToNow = absoluteExpirationRelativeToNow;
-            entry.Value = value;
-            entry.Dispose();
+            using (var entry = cache.CreateEntry(key))
+            {
+                entry.AbsoluteExpirationRelativeToNow = absoluteExpirationRelativeToNow;
+                entry.Value = value;
+            }
 
             return value;
         }
 
         public static TItem Set<TItem>(this IMemoryCache cache, object key, TItem value, IChangeToken expirationToken)
         {
-            var entry = cache.CreateEntry(key);
-            entry.AddExpirationToken(expirationToken);
-            entry.Value = value;
-            entry.Dispose();
+            using (var entry = cache.CreateEntry(key))
+            {
+                entry.AddExpirationToken(expirationToken);
+                entry.Value = value;
+            }
 
             return value;
         }
@@ -91,13 +95,11 @@ public static TItem GetOrCreate<TItem>(this IMemoryCache cache, object key, Func
         {
             if (!cache.TryGetValue(key, out object result))
             {
-                var 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();
+                using (var entry = cache.CreateEntry(key))
+                {
+                    result = factory(entry);
+                    entry.Value = result;
+                }
             }
 
             return (TItem)result;
@@ -107,13 +109,11 @@ public static async Task<TItem> GetOrCreateAsync<TItem>(this IMemoryCache cache,
         {
             if (!cache.TryGetValue(key, out object result))
             {
-                var entry = cache.CreateEntry(key);
-                result = await 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();
+                using (ICacheEntry entry = cache.CreateEntry(key))
+                {
+                    result = await factory(entry).ConfigureAwait(false);
+                    entry.Value = result;
+                }
             }
 
             return (TItem)result;

src/Caching/Memory/src/CacheEntry.cs

@@ -11,10 +11,10 @@ namespace Microsoft.Extensions.Caching.Memory
 {
     internal class CacheEntry : ICacheEntry
     {
-        private bool _added = false;
+        private bool _disposed = false;
         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;
@@ -25,22 +25,24 @@ 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)
         {
             if (key == null)
             {
                 throw new ArgumentNullException(nameof(key));
             }
 
-            if (notifyCacheEntryDisposed == null)
+            if (notifyCacheEntryCommit == null)
             {
-                throw new ArgumentNullException(nameof(notifyCacheEntryDisposed));
+                throw new ArgumentNullException(nameof(notifyCacheEntryCommit));
             }
 
             if (notifyCacheOfExpiration == null)
@@ -49,7 +51,7 @@ internal CacheEntry(
             }
 
             Key = key;
-            _notifyCacheEntryDisposed = notifyCacheEntryDisposed;
+            _notifyCacheEntryCommit = notifyCacheEntryCommit;
             _notifyCacheOfExpiration = notifyCacheOfExpiration;
 
             _scope = CacheEntryHelper.EnterScope(this);
@@ -173,20 +175,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/Caching/Memory/test/MemoryCacheSetAndRemoveTests.cs

@@ -2,6 +2,7 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 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,21 @@ 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()
+        {
+            var cache = CreateCache();
+
+            ICacheEntry entry = cache.CreateEntry("myKey");
+            Assert.NotNull(GetScope(entry));
+
+            entry.Dispose();
+            Assert.Null(GetScope(entry));
         }
 
         [Fact]
@@ -625,6 +644,13 @@ public async Task GetOrCreateAsyncFromCacheWithNullKeyThrows()
             await Assert.ThrowsAsync<ArgumentNullException>(async () => await cache.GetOrCreateAsync<object>(null, null));
         }
 
+        private object GetScope(ICacheEntry entry)
+        {
+            return entry.GetType()
+                .GetField("_scope", BindingFlags.NonPublic | BindingFlags.Instance)
+                .GetValue(entry);
+        }
+
         private class TestKey
         {
             public override bool Equals(object obj) => true;

src/Caching/Memory/test/Microsoft.Extensions.Caching.Memory.Tests.csproj

@@ -6,6 +6,7 @@
   </PropertyGroup>
 
   <ItemGroup>
+  	<Reference Include="Microsoft.Extensions.Caching.Abstractions" />
     <Reference Include="Microsoft.Extensions.Caching.Memory" />
     <Reference Include="Microsoft.Extensions.DependencyInjection" />
   </ItemGroup>