How to achieve remove_if functionality in .NET ConcurrentDictionary

Question

I have a scenario where I have to keep reference counted object for given key in the ConcurrentDictionary, if reference count reaches 0, I want to delete the key. This has to be thread safe hence I am planning to use the ConcurrentDictionary.

Sample program as follows. In the concurrent dictionary, I have key and value , the value is KeyValuePair which holds my custom object and reference count.

ConcurrentDictionary<string, KeyValuePair<object, int>> ccd = 
    new ConcurrentDictionary<string, KeyValuePair<object, int>>();

// following code adds the key, if not exists with reference 
// count   for  my custom object to 1
// if the key already exists it increments the reference count

var addOrUpdateValue = ccd.AddOrUpdate("mykey",
    new KeyValuePair<object, int>(new object(), 1),
    (k, pair) => new KeyValuePair<object, int>(pair.Key, pair.Value + 1));

Now I want a way to remove the key when the reference count reaches to 0. I was thinking , remove method on ConcurrentDictionary which takes key and predicate , removes the key if the predicate return 'true'. Example.

ConcurrentDictionary.remove(TKey, Predicate<TValue> ). 

There is no such method on ConcurrentDictionary, question is how to do the same in thread safe way ?.

Answer 1

.NET doesn't expose a RemoveIf directly, but it does expose the building blocks necessary to make it work without doing your own locking.

ConcurrentDictionary implements ICollection<T>, which has a Remove that takes and tests for a full KeyValuePair instead of just a key. Despite being hidden, this Remove is still thread-safe and we'll use it to implement this. One caveat for this to work is that Remove uses EqualityComparer<T>.Default to test the value, so it must be equality comparable. Your current one is not, so we'll re-implement that as such:

struct ObjectCount : IEquatable<ObjectCount>
{
    public object Object { get; }
    public int Count { get; }

    public ObjectCount(object o, int c)
    {
        Object = o;
        Count = c;
    }

    public bool Equals(ObjectCount o) =>
       object.Equals(Object, o.Object) && Count == o.Count;

    public override bool Equals(object o) =>
       (o as ObjectCount?)?.Equals(this) == true;

    // this hash combining will work but you can do better.
    // it is not actually used by any of this code.
    public override int GetHashCode() =>
       (Object?.GetHashCode() ?? 0) ^ Count.GetHashCode();
}

And finally, we'll define a method to increment/decrement counts from your dictionary:

void UpdateCounts(ConcurrentDictionary<string, ObjectCount> dict, string key, int toAdd)
{
    var addOrUpdateValue = dict.AddOrUpdate(key,
        new ObjectCount(new object(), 1),
        (k, pair) => new ObjectCount(pair.Key, pair.Value + toAdd));

    if(addOrUpdateValue.Count == 0)
    {
        ((ICollection<KeyValuePair<string, ObjectCount>>)dict).Remove(
            new KeyValuePair<string, ObjectCount>(key, addOrUpdateValue));
    }
}

The value for that key might be changed between the calls of AddOrUpdate and Remove, but that doesn't matter to us: because Remove tests the full KeyValuePair, it will only remove it if the value hasn't changed since the update.

This is the common lock-free pattern of setting up a change and then using a final thread-safe op to safely "commit" the change only if our data structure hasn't been updated in the meantime.