|
// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
using System;
using System.Collections.Generic;
using System.Diagnostics.CodeAnalysis;
namespace Microsoft.AspNetCore.Razor.Utilities;
/// <summary>
/// A thread-safe cache implementation that uses weak references to store values, allowing them to be garbage collected
/// when no longer referenced elsewhere. The cache periodically cleans up dead weak references to prevent unbounded growth.
/// </summary>
/// <typeparam name="TKey">The type of keys used to identify cached values. Must be non-null.</typeparam>
/// <typeparam name="TValue">The type of values stored in the cache. Must be a reference type.</typeparam>
/// <remarks>
/// This cache is designed for scenarios where you want to cache expensive-to-create objects but allow them to be
/// garbage collected when memory pressure occurs. The cache will automatically clean up dead references when the
/// number of add operations reaches the specified cleanup threshold.
/// </remarks>
internal class CleanableWeakCache<TKey, TValue>
where TKey : notnull
where TValue : class?
{
/// <summary>
/// The underlying dictionary that maps keys to weak references containing the cached values.
/// </summary>
private readonly Dictionary<TKey, WeakReference<TValue>> _cacheMap = [];
/// <summary>
/// Synchronization object to ensure thread-safe access to the cache.
/// </summary>
#if NET9_0_OR_GREATER
private readonly System.Threading.Lock _lock = new();
#else
private readonly object _lock = new();
#endif
/// <summary>
/// The number of add operations that must occur before triggering a cleanup of dead weak references.
/// </summary>
private readonly int _cleanUpThreshold;
/// <summary>
/// Counter tracking the number of add operations since the last cleanup was performed.
/// </summary>
private int _addsSinceLastCleanUp;
/// <summary>
/// Initializes a new instance of the <see cref="CleanableWeakCache{TKey, TValue}"/> class.
/// </summary>
/// <param name="cleanUpThreshold">
/// The number of add operations that must occur before triggering automatic cleanup of dead weak references.
/// Must be positive.
/// </param>
/// <exception cref="ArgumentOutOfRangeException">
/// Thrown when <paramref name="cleanUpThreshold"/> is zero or negative.
/// </exception>
public CleanableWeakCache(int cleanUpThreshold)
{
ArgHelper.ThrowIfNegativeOrZero(cleanUpThreshold);
_cleanUpThreshold = cleanUpThreshold;
}
/// <summary>
/// Gets the value associated with the specified key, or adds the provided value if the key is not found
/// or the previously cached value has been garbage collected.
/// </summary>
/// <param name="key">The key of the value to get or add.</param>
/// <param name="value">The value to add if the key is not found or the cached value is no longer available.</param>
/// <returns>
/// The existing value if found and still alive, otherwise the provided <paramref name="value"/>.
/// </returns>
public TValue GetOrAdd(TKey key, TValue value)
{
lock (_lock)
{
// Try to get the existing value or add the new one.
return GetOrAdd_NoLock(key, value);
}
}
/// <summary>
/// Gets the value associated with the specified key, or adds a new value created by the factory function
/// if the key is not found or the previously cached value has been garbage collected.
/// </summary>
/// <param name="key">The key of the value to get or add.</param>
/// <param name="valueFactory">A factory function to create the value if it needs to be added to the cache.</param>
/// <returns>
/// The existing value if found and still alive, otherwise a new value created by <paramref name="valueFactory"/>.
/// </returns>
public TValue GetOrAdd(TKey key, Func<TValue> valueFactory)
{
// First check without creating the value.
if (TryGet(key, out var value))
{
return value;
}
// Create the value outside the lock to avoid holding the lock
// while creating a potentially expensive object.
var newValue = valueFactory();
lock (_lock)
{
// Try to add the newly-created value or get the existing one.
return GetOrAdd_NoLock(key, newValue);
}
}
/// <summary>
/// Gets the value associated with the specified key, or adds a new value created by the factory function
/// using the provided argument if the key is not found or the previously cached value has been garbage collected.
/// </summary>
/// <typeparam name="TArg">The type of the argument passed to the value factory function.</typeparam>
/// <param name="key">The key of the value to get or add.</param>
/// <param name="arg">The argument to pass to the value factory function.</param>
/// <param name="valueFactory">A factory function to create the value using the provided argument.</param>
/// <returns>
/// The existing value if found and still alive, otherwise a new value created by <paramref name="valueFactory"/>.
/// </returns>
public TValue GetOrAdd<TArg>(TKey key, TArg arg, Func<TArg, TValue> valueFactory)
{
// First check without creating the value.
if (TryGet(key, out var value))
{
return value;
}
// Create the value outside the lock to avoid holding the lock
// while creating a potentially expensive object.
var newValue = valueFactory(arg);
lock (_lock)
{
// Try to add the newly-created value or get the existing one.
return GetOrAdd_NoLock(key, newValue);
}
}
/// <summary>
/// Attempts to add the specified key-value pair to the cache.
/// </summary>
/// <param name="key">The key of the value to add.</param>
/// <param name="value">The value to add to the cache.</param>
/// <returns>
/// <see langword="true"/> if the key-value pair was successfully added;
/// <see langword="false"/> if a live value already exists for the specified key.
/// </returns>
public bool TryAdd(TKey key, TValue value)
{
lock (_lock)
{
// Check if the key exists and the weak reference still has a live target
if (!_cacheMap.TryGetValue(key, out var weakRef))
{
// The key is not in the map. Add a new weak reference.
_cacheMap.Add(key, new(value));
}
else
{
// We have a weak reference, check if it is still alive
if (weakRef.TryGetTarget(out var existingValue))
{
// Yup, we have a live value for this key.
// However, we aren't returning it to the caller, so we should keep it alive
// until after we return to ensure the existence check remains valid.
GC.KeepAlive(existingValue);
return false;
}
// Set the weak ref's target to the new value.
weakRef.SetTarget(value);
}
// We added an item to the cache.
// Increment the add counter and trigger cleanup if needed.
CleanUpIfNeeded_NoLock();
return true;
}
}
/// <summary>
/// Attempts to get the value associated with the specified key.
/// </summary>
/// <param name="key">The key of the value to retrieve.</param>
/// <param name="value">
/// When this method returns, contains the value associated with the specified key if found and still alive;
/// otherwise, <see langword="null"/>.
/// </param>
/// <returns>
/// <see langword="true"/> if a live value was found for the specified key; otherwise, <see langword="false"/>.
/// </returns>
public bool TryGet(TKey key, [NotNullWhen(true)] out TValue? value)
{
lock (_lock)
{
// Check if the key exists and the weak reference still has a live target
if (_cacheMap.TryGetValue(key, out var weakRef) &&
weakRef.TryGetTarget(out value))
{
return true;
}
// Key not found or target was garbage collected
value = null;
return false;
}
}
/// <summary>
/// Internal method that attempts to retrieve an existing value from the cache or add a new value if none exists.
/// This method assumes the caller already holds the lock.
/// </summary>
/// <param name="key">The key of the value to get or add.</param>
/// <param name="value">The value to add if no existing live value is found.</param>
/// <returns>
/// The existing live value if one was found; otherwise, the provided <paramref name="value"/> after adding it to the cache.
/// </returns>
/// <remarks>
/// This method increments the add counter and triggers cleanup if the threshold is reached when adding a value.
/// </remarks>
private TValue GetOrAdd_NoLock(TKey key, TValue value)
{
if (_cacheMap.TryGetValue(key, out var weakRef))
{
if (weakRef.TryGetTarget(out var existingValue))
{
// There was already a value in the map. Return it!
return existingValue;
}
// The key was in the map, but the weak reference was collected.
// Set its target to the new value.
weakRef.SetTarget(value);
}
else
{
// The key is not in the map. Add a new weak reference.
_cacheMap.Add(key, new(value));
}
// We added an item to the cache.
// Increment the add counter and trigger cleanup if needed.
CleanUpIfNeeded_NoLock();
return value;
}
/// <summary>
/// Increments the add counter and removes all cache entries whose weak references no longer have live targets
/// if the cleanup threshold has been reached. This method assumes the caller already holds the lock.
/// </summary>
/// <remarks>
/// This method resets the add counter to zero after cleanup is performed.
/// </remarks>
private void CleanUpIfNeeded_NoLock()
{
if (++_addsSinceLastCleanUp < _cleanUpThreshold)
{
return;
}
// Use a memory builder to collect keys of dead weak references
using var deadKeys = new MemoryBuilder<TKey>(initialCapacity: _cacheMap.Count, clearArray: true);
// Identify all keys with dead weak references
foreach (var (key, weakRef) in _cacheMap)
{
if (!weakRef.TryGetTarget(out _))
{
deadKeys.Append(key);
}
}
// Remove all dead entries from the cache
foreach (var key in deadKeys.AsMemory().Span)
{
_cacheMap.Remove(key);
}
// Reset the add counter since we just performed cleanup
_addsSinceLastCleanUp = 0;
}
}
|