Thread-Safe Garbage Collectible Memoizer

Question

I am working on a little thread-safe, garbage collectible memoizer for Funcs in C#.

The goals:

• Make it easy to Memoize deterministic functions.
• Make sure invocation of a memoized Func with a given input is only actually computed once, even if invoked from multiple threads simultaneously.
• Make the memoizer's cache optionally garbage collectible so keys that are not incommon use can be disposed.

First, the IMemoizeThings interface:

/// <summary>
/// Defines the methods needed to implement a Memoizer.
/// </summary>
/// <typeparam name="TKey">The type of memoizer's keys.</typeparam>
/// <typeparam name="TValue">The type of the memoizer's return values.</typeparam>
public interface IMemoizeThings<in TKey, out TValue>
{
    /// <summary>
    /// Gets or adds a key/value pair from the memoizer if it already exists.  Adds a key/value pair to the memoizer if it does not already exist.
    /// </summary>
    /// <param name="key">The key of the element to get or add.</param>
    /// <returns>The value for the key.  This will be either the existing value for the key if the key is already in the memoizer, or the new value if the key was not in the memoizer.</returns>
    TValue GetOrAdd(TKey key);
}

And here is the InvokeWithWriteLock extension method. I'm thinking this might should just be moved into the base class, as it probably isn't useful as an extension method outside the library.

/// <summary>
/// Enters a write lock on the given <see cref="ReaderWriterLockSlim"/>, invokes the given <see cref="Func{TResult}"/>, exits the write lock, and returns the result of the <paramref name="func"/>.
/// </summary>
/// <typeparam name="TResult">The return type.</typeparam>
/// <param name="func">The function to invoke inside a write lock.</param>
/// <param name="locker">The <see cref="ReaderWriterLockSlim"/> to use for locking.</param>
/// <returns>The result of the <paramref name="func"/>.</returns>
public static TResult InvokeWithWriteLock<TResult>(this Func<TResult> func, ReaderWriterLockSlim locker)
{
    locker.EnterWriteLock();
    try
    {
        return func();
    }
    finally
    {
        locker.ExitWriteLock();
    }
}

Next, the abstract base class, which contains the common workflow for all derived memoizers:

/// <summary>
/// Represents a thread-safe memoizer for a given <see cref="Func{TKey, TValue}"/>.
/// </summary>
/// <typeparam name="TKey">The type of the memoizer's keys.</typeparam>
/// <typeparam name="TValue">The type of the memoizer's return values.</typeparam>
/// <typeparam name="TCachedValue">The type of the memoizer's cached values.</typeparam>
public abstract class MemoizerBase<TKey, TValue, TCachedValue> : IMemoizeThings<TKey, TValue>
{
    protected MemoizerBase(Func<TKey, TValue> func)
        : this()
    {
        Func = func;
    }

    private MemoizerBase()
    {
        Cache = new Dictionary<TKey, TCachedValue>();
        CacheLock = new ReaderWriterLockSlim(LockRecursionPolicy.SupportsRecursion);
    }

    protected Dictionary<TKey, TCachedValue> Cache { get; }

    protected ReaderWriterLockSlim CacheLock { get; }

    protected Func<TKey, TValue> Func { get; private set; }

    /// <summary>
    /// Gets or adds a key/value pair from the memoizer if it already exists.  Adds a key/value pair to the memoizer if it does not already exist.
    /// </summary>
    /// <param name="key">The key of the element to get or add.</param>
    /// <returns>The value for the key.  This will be either the existing value for the key if the key is already in the memoizer, or the new value if the key was not in the memoizer.</returns>
    public TValue GetOrAdd(TKey key)
    {
        TValue result;

        CacheLock.EnterUpgradeableReadLock();
        try
        {
            if (ContainsKey(key))
            {
                return GetValue(key);
            }
            else
            {
                Func<TValue> func = () => GetOrSetValue(key);
                result = func.InvokeWithWriteLock(CacheLock);
            }
        }
        finally
        {
            CacheLock.ExitUpgradeableReadLock();
        }

        return result;
    }

    protected abstract bool ContainsKey(TKey key);

    protected abstract TValue SetValue(TKey key);

    protected abstract TValue GetValue(TKey key);

    private TValue GetOrSetValue(TKey key)
    {
        return ContainsKey(key)
            ? GetValue(key)
            : SetValue(key);
    }
}

Next, the non-garbage collectable implementation:

/// <summary>
/// Represents a persistent thread-safe memoizer for a given <see cref="Func{TKey, TValue}"/>.
/// </summary>
/// <typeparam name="TKey">The type of the memoizer's keys.</typeparam>
/// <typeparam name="TValue">The type of the memoizer's return values.</typeparam>
public class Memoizer<TKey, TValue> : MemoizerBase<TKey, TValue, TValue>
{
    /// <summary>
    /// Initializes a new instance of the <see cref="Memoizer{TKey, TValue}"/> class.
    /// </summary>
    /// <param name="func">The encapsulated method that this memoizer will memoize.</param>
    public Memoizer(Func<TKey, TValue> func)
        : base(func)
    {
    }

    protected override TValue SetValue(TKey key)
    {
        var value = Func(key);
        Cache.Add(key, value);
        return value;
    }

    protected override TValue GetValue(TKey key)
    {
        return Cache[key];
    }

    protected override bool ContainsKey(TKey key)
    {
        return Cache.ContainsKey(key);
    }
}

Then the garbage collectable version, using weak references, which I hope is the right way to approach this:

/// <summary>
/// Represents a thread-safe memoizer for a given <see cref="Func{TKey, TValue}"/>.  The memoizer's cached values can be expire and be garbage-collected.
/// </summary>
/// <typeparam name="TKey">The type of the memoizer's keys.</typeparam>
/// <typeparam name="TValue">The type of the memoizer's return values.</typeparam>
public class ExpirableMemoizer<TKey, TValue> : MemoizerBase<TKey, TValue, WeakReference>
{
    /// <summary>
    /// Initializes a new instance of the <see cref="ExpirableMemoizer{TKey, TValue}"/> class.
    /// </summary>
    /// <param name="func">The encapsulated method that this memoizer will memoize.</param>
    public ExpirableMemoizer(Func<TKey, TValue> func)
        : base(func)
    {
    }

    protected override TValue SetValue(TKey key)
    {
        var value = Func(key);
        Cache.Add(key, new WeakReference(value));
        return value;
    }

    protected override TValue GetValue(TKey key)
    {
        TValue value;
        var weakReference = Cache[key];
        if (weakReference.Target == null)
        {
            value = Func(key);
            Cache[key].Target = value;
        }
        else
        {
            value = (TValue)weakReference.Target;
        }

        return value;
    }

    protected override bool ContainsKey(TKey key)
    {
        return Cache.ContainsKey(key);
    }
}

And finally the extension method that I wanted to build in the first place:

/// <summary>
/// Memoizes an encapsulated method that has 1 parameter and returns a value of the type specified by the <typeparamref name="TResult"/> parameter.
/// </summary>
/// <typeparam name="T">The type of the parameter of the encapsulated method that this delegate will memoize.</typeparam>
/// <typeparam name="TResult">The type of the retun value of the encapsulated method that this delegate will memoize.</typeparam>
/// <param name="func">The encapsulated method that this delegate will memoize.</param>
/// <param name="isExpirable">A value that specifies whether the garbage collector can collect the memoized values.</param>
/// <returns>A memoized version of the encapsulated method represented by the <paramref name="func"/> parameter.</returns>
public static Func<T, TResult> Memoize<T, TResult>(
    this Func<T, TResult> func,
    bool isExpirable = false)
{
    IMemoizeThings<T, TResult> memoizer;
    if (isExpirable)
    {
        memoizer = new ExpirableMemoizer<T, TResult>(func);
    }
    else
    {
        memoizer = new Memoizer<T, TResult>(func);
    }

    return argument => memoizer.GetOrAdd(argument);
}

Now we need a Func to memoize. I'm using a quick Fibonacci function. I know there are better approaches to fibonacci, but it serves as a simple example:

public long Fibonacci(long n)
{
    long a = 0;
    long b = 1;

    for (long i = 0; i < n; ++i)
    {
        var temp = a;
        a = b;
        b = temp + b;
    }

    return a;
}

And the actual usage:

Func<long, long> fibonacci = Fibonacci;
var memoized = fibonacci.Memoize(true);

for (var i = 0; i < 1000000000; i++)
{
    // The underlying "Fibonacci" function is only invoke once.
    var hugeResult = memoized(100);
    Console.WriteLine($"{hugeResult}");
}

Answer 1

I believe that might be too much stuff happening here, with a sole purpose of solving a rather simple problem: "I need an in-memory cache". The .NET standard library already has a cache class (System.Runtime.Caching.MemoryCache), so the only reason why you would want to replace it is if you tested it and found out that is doesn't fit your needs for certain reasons.

Also, while abstract classes allow you to reuse code, they also tend to create strong dependencies on their interface - you have a lot of plumbing here necessary just to get things going. So my suggestion would be to start with a more generic data structure, instead of adding a layer of methods in an abstract base class:

1. For the simpler case (without weak references), the entire code could be (almost) replaced with a single method:

   public static Func<T, TResult> Memoize<T, TResult>(this Func<T, TResult> func)
   {
       var dict = new ConcurrentDictionary<T, TResult>();
       return key => dict.GetOrAdd(key, func);
   }
   

   The ConcurrentDictionary class is lockless, thread-safe, and usually outperforms a plain Mutex or a ReaderWriterLockSlim (ReaderWriterLockSlim being the slowest of the three in most cases).

   So the only difference you would have with the method above is that ConcurrentDictionary doesn't lock the call to func (but only one result is atomically stored to the dictionary at the end). To resolve this, you could start by writing a ConcurrentDictionary extension method similar to:

   // This basically extends ConcurrentDictionary.GetOrAdd to 
   // acquire a mutex before calling func.
   
   /// <summary>
   /// Returns the value associated with the specified key if there already is
   /// one, or calls the specified delegate to create a new value which is
   /// stored and returned. This method will not lock if the value already
   /// exists, but it will lock the entire transaction if a new value needs to
   /// be instantiated.
   /// </summary>
   public static TValue GetOrAddSafe<TKey, TValue>(
       this ConcurrentDictionary<TKey, TValue> @this,
       object lockInstance,
       TKey key,
       Func<TKey, TValue> valueProvider)
   {
       TValue value;
   
       // note that we don't need to lock here at all,
       // ConcurrentDictionary does all the lockless magic
       if (@this.TryGetValue(key, out value))
       {
           return value;
       }
   
       // fallback to mutex
       lock (lockInstance)
       {
           return @this.GetOrAdd(key, valueProvider);
       }
   }
   

   and then

   static readonly object _dictlock = new object();    
   public static Func<T, TResult> Memoize<T, TResult>(this Func<T, TResult> func)
   {
       var dict = new ConcurrentDictionary<T, TResult>();
       return key => dict.GetOrAddSafe(_dictlock, key, func);
   }
   

   So you end up with a reusable generic extension method and get rid of everything else at the same time.

2. Regarding the weak-referenced version, you could also easily squeeze the functionality in a single method using a similar approach to the one above. But a more important question is whether this functionality will be useful, since GC will probably collect these items very soon, and you don't have any control over the process. Are you sure you don't want to check out MemoryCache instead?

Answer 2

1) I would use a simple lock to do the synchronization. ReaderWriterLockSlim outperforms regular lock only in very specific scenarios and it only makes sense if you want to write on one thread and read on another thread in parallel. Which is not the use case that you have. Alternatively, consider using ConcurrentDictionary, that has a built-in thread-safe GetOrAdd method.

2) Func<TValue> func = () => GetOrSetValue(key); - this lambda looks like an unnecessary overhead. Just move InvokeWithWriteLock into the base class, and remove the lambda.

3) By calling both ContainsKey and GetValue you are doing the dictionary lookup twice on every request. You should use Dictionary.TryGetValue method instead.

4) The weak reference implementation looks weird. It looks like in most scenarios the values would be removed from the cache straight away, unless the reference to those values was stored elsewhere. Is there really a usecase for such cache? A more common approach is to attach a timestamp to key or value, and update it every time the key/value is accessed. Then you should periodically check the dictionary and remove items that were not accessed long enough.

5) I think ICache or IMemoizer are better names for the interface.