[release/2.1] Memory Leak in Microsoft.Extensions.Caching.Memory when handling exceptions (#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

Author: wtgodbe

eng/PatchConfig.props

@@ -13,6 +13,8 @@
  </PropertyGroup>
  <PropertyGroup Condition=" '$(VersionPrefix)' == '2.1.23' ">
  <PackagesInPatch>
+ Microsoft.Extensions.Caching.Abstractions;
+ Microsoft.Extensions.Caching.Memory;
  </PackagesInPatch>
  </PropertyGroup>
 

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>