Thread safe Pooled Dictionary

Question

I have an object that is fairly expensive to create and has a tendency to get created multiple times because each distinct user may have one or many instances of this object. This object can expire as its internal authentication becomes stale and no longer valid.

Because of this, I want need to keep a dictionary so I can keep a list of objects for each user.... but....I need this dictionary to not retain objects indefinitely since they can expire. ( Managing the expiration is outside the scope of this )

I used this article as a starting point, but took the concept a little further to manage the number of objects stored internally to keep the pool from getting too large.

public sealed class PooledDictionary<TKey, TValue>
{
    private readonly byte m_poolSize;
    private readonly ConcurrentDictionary<TKey, ConcurrentQueue<TValue>> m_pool;
    private readonly ConcurrentQueue<TKey> m_keyQueue;
    public PooledDictionary(byte maxPoolSize)
    {
        if (maxPoolSize <= 0)
            maxPoolSize = (byte)Math.Min(Environment.ProcessorCount, byte.MaxValue);

        m_poolSize = maxPoolSize;
        m_keyQueue = new ConcurrentQueue<TKey>();
        m_pool = new ConcurrentDictionary<TKey, ConcurrentQueue<TValue>>(Environment.ProcessorCount * 2, m_poolSize);
    }

    public bool TryAdd(TKey key, TValue value)
    {
        if (key == null)
            return false;

        if (!m_pool.ContainsKey(key) && m_pool.TryAdd(key, new ConcurrentQueue<TValue>()))
        {
            m_keyQueue.Enqueue(key);

            while (m_pool.Count > m_poolSize)
            {
                TKey localKey;
                if (m_keyQueue.TryDequeue(out localKey))
                {
                    TryRemove(localKey);
                }
            }
        }

        ConcurrentQueue<TValue> q;
        if (m_pool.TryGetValue(key, out q))
        {
            q.Enqueue(value);
            while (q.Count > m_poolSize)
            {
                TValue localValue;
                q.TryDequeue(out localValue);
            }
        }
        else
        {
            return false;
        }
        return true;
    }

    public bool TryRemove(TKey key)
    {
        if (key == null)
            return false;

        ConcurrentQueue<TValue> value;
        return m_pool.TryRemove(key, out value);
    }

    public int Count(TKey key)
    {
        if (key == null)
            return 0;

        if (!m_pool.ContainsKey(key)) m_pool.TryAdd(key, new ConcurrentQueue<TValue>());
        ConcurrentQueue<TValue> q;
        if (m_pool.TryGetValue(key, out q))
        {
            return q.Count;
        }
        return 0;
    }

    public TValue Request(TKey key, Func<TValue> creator = null)
    {
        if (key == null)
            return default(TValue);

        if (!m_pool.ContainsKey(key)) m_pool.TryAdd(key, new ConcurrentQueue<TValue>());
        ConcurrentQueue<TValue> q;
        if (m_pool.TryGetValue(key, out q))
        {
            TValue v;
            if (q.TryDequeue(out v)) return v;
            if (null != creator) return creator();
        }
        return default(TValue);
    }

    public bool Release(TKey key, TValue value)
    {
        return TryAdd(key, value); //just adds it back to key's queue
    }
}

Answer 1

Anytime that you check a count then do something on that you have to assume another thread did something in between.

for example this code

while (q.Count > m_poolSize)
{
    TValue localValue;
    q.TryDequeue(out localValue);
}

Because no locking is happening two threads could read the count and see it's over and both could remove an item. Since you are using ConcurrentQueue both will remove different items. This would cause your pool to drop below max size. Maybe that's ok in your situation.

Same issue is applied when having the Keys being a ConcurrentQueue. In theory if you had your max pool size lower than the processor count times 2 and each thread hit that line at the same time you could get an exception, since you would try to dequeue when the queue has no items in it. Granted this is an extreme example and one not likely to hit but at a min you would want to catch that exception.

If you want to have the max stored in the pool you will need to do locking on an object and if we are going to do then both the ConcurrentQueue can go away. The Keys I would replace as a list since that will do the array manipulation for us and change the ConcurrentDictionary values to just be a normal Queue.

private readonly ConcurrentDictionary<TKey, Queue<TValue>> m_pool;
private IList<TKey> m_keys;
private readonly object m_lock = new object();

We can simplify the Count method to use the concurrentDictionary GetOrAdd method

public int Count(TKey key)
{
    if (key == null)
        return 0;

    var q = m_pool.GetOrAdd(key, _ => new Queue<TValue>());
    return q.Count;
}

The Queue class implements ICollection and ICollection has a SyncRoot object. To make it easy I would create a helper method so I don't have to cast to the Icollection each time I wanted the SyncRoot

private object SyncLock(ICollection collection)
{
    return collection.SyncRoot;
}

For the Request Method I would use the GetOrAdd again and lock on the return queue to know the count hasn't changed when Dequeued

public TValue Request(TKey key, Func<TValue> creator = null)
{
    if (key == null)
       return default(TValue);

    var q = m_pool.GetOrAdd(key, _ => new Queue<TValue>());
    lock (SyncLock(q))
    {
        if (q.Count > 0)
        {
            return q.Dequeue();
        }
    }
    if (creator != null)
    {
        return creator();
    }
    return default(TValue);

}

TryAdd method now needs locking as well. Plus we will use the ConcurrentDictionary GetOrAdd method and pass in the value to add so we can tell if it added to the dictionary or if it returned back an existing value.

public bool TryAdd(TKey key, TValue value)
{
    if (key == null)
       return false;

    var addedQueue = new Queue<TValue>();
    var q = m_pool.GetOrAdd(key, addedQueue);

    lock (SyncLock(q))
    {
        q.Enqueue(value);
        if (q.Count > m_poolSize)
        {
            q.Dequeue();
        }
    }

    // If these are the same then we added to the concurrent dictionary 
    if (q == addedQueue)
    {
        lock (m_lock)
        {
            m_keys.Add(key);
            if (m_keys.Count > m_poolSize)
            {
                var ejected = m_keys[0];
                m_keys.RemoveAt(0);
                TryRemove(ejected);
            }
        }
    }

    return true;
}

Since we add to both the keys and concurrentdictionary then check to see if over we should init them with plus one of the max count

public PooledDictionary(byte maxPoolSize)
{
    if (maxPoolSize <= 0)
        maxPoolSize = (byte)Math.Min(Environment.ProcessorCount, byte.MaxValue);

    m_poolSize = maxPoolSize;
    m_keys = new List<TKey>(maxPoolSize + 1);

    m_pool = new ConcurrentDictionary<TKey, Queue<TValue>>(Environment.ProcessorCount * 2, m_poolSize + 1);
}

In this implementation there is no need for while(q.Count > m_poolsize) as we controlling the adding an could only ever be one over the size. Which is our trigger to remove one.