David Fowler
2021-02-05T03:59:16Z

ConcurrentDictionaryExtensions.cs

| using System;
| using System.Threading.Tasks;
|
| namespace System.Collections.Concurrent
| {
| public static class ConcurrentDictionaryExtensions
| {
| ///
| /// Provides an alternative to <see cref="ConcurrentDictionary{TKey, TValue}.GetOrAdd(TKey, Func{TKey, TValue})"/> that disposes values that implement <see cref="IDisposable"/>.
| ///
| /// <typeparam name="TKey"></typeparam>
| /// <typeparam name="TValue"></typeparam>
| ///
| ///
| ///
| /// <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();
| }
| }
|
|
| ///
| /// 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.
| ///
| /// <typeparam name="TKey"></typeparam>
| /// <typeparam name="TValue"></typeparam>
| ///
| ///
| ///
| /// <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;
| }
| }
| }
| }
| }
| }