Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
using System;
using System.Threading.Tasks;
namespace System.Collections.Concurrent
{
public static class ConcurrentDictionaryExtensions
{
/// <summary>
/// Provides an alternative to <see cref="ConcurrentDictionary{TKey, TValue}.GetOrAdd(TKey, Func{TKey, TValue})"/> that disposes values that implement <see cref="IDisposable"/>.
/// </summary>
/// <typeparam name="TKey"></typeparam>
/// <typeparam name="TValue"></typeparam>
/// <param name="dictionary"></param>
/// <param name="key"></param>
/// <param name="valueFactory"></param>
/// <returns></returns>
public static TValue GetOrAddWithDispose<TKey, TValue>(
this ConcurrentDictionary<TKey, TValue> dictionary,
TKey key,
Func<TKey, TValue> valueFactory) where TValue : IDisposable
{
while (true)
{
if (dictionary.TryGetValue(key, out var value))
{
// Try to get the value
return value;
}
/// Try to add the value
value = valueFactory(key);
if (dictionary.TryAdd(key, value))
{
// Won the race, so return the instance
return value;
}
// Lost the race, dispose the created object
value.Dispose();
}
}
/// <summary>
/// Provides an alternative to <see cref="ConcurrentDictionary{TKey, TValue}.GetOrAdd(TKey, Func{TKey, TValue})"/> specifically for asynchronous values. The factory method will only run once.
/// </summary>
/// <typeparam name="TKey"></typeparam>
/// <typeparam name="TValue"></typeparam>
/// <param name="dictionary"></param>
/// <param name="key"></param>
/// <param name="valueFactory"></param>
/// <returns></returns>
public static async Task<TValue> GetOrAddAsync<TKey, TValue>(
this ConcurrentDictionary<TKey, Task<TValue>> dictionary,
TKey key,
Func<TKey, Task<TValue>> valueFactory)
{
while (true)
{
if (dictionary.TryGetValue(key, out var task))
{
return await task;
}
// This is the task that we'll return to all waiters. We'll complete it when the factory is complete
var tcs = new TaskCompletionSource<TValue>(TaskCreationOptions.RunContinuationsAsynchronously);
if (dictionary.TryAdd(key, tcs.Task))
{
try
{
var value = await valueFactory(key);
tcs.TrySetResult(value);
return await tcs.Task;
}
catch (Exception ex)
{
// Make sure all waiters see the exception
tcs.SetException(ex);
// We remove the entry if the factory failed so it's not a permanent failure
// and future gets can retry (this could be a pluggable policy)
dictionary.TryRemove(key, out _);
throw;
}
}
}
}
}
}
@gabrielsimas

This comment has been minimized.

Copy link

@gabrielsimas gabrielsimas commented Feb 5, 2021

Thanks for sharing

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment