Make GetOrCreate(key) atomic on ConcurrentDictionary

Extensions.Test/DictionaryExtensionsTests.cs

@@ -90,6 +90,31 @@ public void GetOrCreateConcurrentDictionaryShouldReturnSameInstanceToParallelCal
 		Assert.HasCount(1000, dictionary["key1"]);
 	}
 
+	[TestMethod]
+	public void GetOrCreateConcurrentDictionaryWithoutDefaultShouldReturnStoredValueWhenAnotherCallerAddsFirst()
+	{
+		// The comparer adds a rival value for the key the second time it hashes it, which is the
+		// moment between a lookup that missed and the add that follows it
+		RacingComparer comparer = new();
+		ConcurrentDictionary<string, List<int>> dictionary = new(comparer);
+		List<int> rival = [];
+		comparer.OnSecondHash = () => dictionary.TryAdd("key1", rival);
+
+		List<int> result = dictionary.GetOrCreate("key1");
+
+		Assert.AreSame(dictionary["key1"], result);
+	}
+
+	[TestMethod]
+	public void GetOrCreateConcurrentDictionaryWithoutDefaultShouldReturnSameInstanceToParallelCallers()
+	{
+		ConcurrentDictionary<string, ConcurrentBag<int>> dictionary = new();
+
+		Parallel.For(0, 1000, i => dictionary.GetOrCreate("key1").Add(i));
+
+		Assert.HasCount(1000, dictionary["key1"]);
+	}
+
 	private sealed class RacingComparer : IEqualityComparer<string>
 	{
 		private int hashCount;
@@ -125,6 +150,22 @@ public void GetOrCreateShouldThrowArgumentNullExceptionWhenKeyIsNull()
 		Assert.ThrowsExactly<ArgumentNullException>(() => dictionary.GetOrCreate(null!));
 	}
 
+	[TestMethod]
+	public void GetOrCreateConcurrentDictionaryWithoutDefaultShouldThrowArgumentNullExceptionWhenDictionaryIsNull()
+	{
+		ConcurrentDictionary<string, int>? dictionary = null!;
+
+		Assert.ThrowsExactly<ArgumentNullException>(() => dictionary.GetOrCreate("key1"));
+	}
+
+	[TestMethod]
+	public void GetOrCreateConcurrentDictionaryWithoutDefaultShouldThrowArgumentNullExceptionWhenKeyIsNull()
+	{
+		ConcurrentDictionary<string, int> dictionary = new();
+
+		Assert.ThrowsExactly<ArgumentNullException>(() => dictionary.GetOrCreate(null!));
+	}
+
 	[TestMethod]
 	public void GetOrCreateShouldThrowArgumentNullExceptionWhenDefaultValueIsNull()
 	{

Extensions/DictionaryExtensions.cs

@@ -33,14 +33,14 @@
 #pragma warning disable KTSU0004 // Use Ensure.NotNull instead of manual null check
 		if (dictionary is null)
 		{
 			throw new ArgumentNullException(nameof(dictionary), "The dictionary cannot be null.");
 		}
 #pragma warning restore KTSU0004 // Use Ensure.NotNull instead of manual null check
 
 #pragma warning disable KTSU0004 // Use Ensure.NotNull instead of manual null check
 		if (key is null)
 		{
 			throw new ArgumentNullException(nameof(key), "The key cannot be null.");
 		}
 #pragma warning restore KTSU0004 // Use Ensure.NotNull instead of manual null check
 
@@ -60,6 +60,35 @@
 		return defaultValue;
 	}
 
+	/// <summary>
+	/// Method that gets a value from a dictionary if it exists, otherwise creates a new value and adds it to the dictionary.
+	/// </summary>
+	/// <typeparam name="TKey">The type of the keys in the dictionary.</typeparam>
+	/// <typeparam name="TVal">The type of the values in the dictionary.</typeparam>
+	/// <param name="dictionary">The dictionary to get the value from.</param>
+	/// <param name="key">The key to get the value for.</param>
+	/// <returns>The value for the key if it exists, otherwise a new value.</returns>
+	public static TVal GetOrCreate<TKey, TVal>(this ConcurrentDictionary<TKey, TVal> dictionary, TKey key) where TKey : notnull where TVal : new()
+	{
+#pragma warning disable KTSU0004 // Use Ensure.NotNull instead of manual null check
+		if (dictionary is null)
+		{
+			throw new ArgumentNullException(nameof(dictionary), "The dictionary cannot be null.");
+		}
+#pragma warning restore KTSU0004 // Use Ensure.NotNull instead of manual null check
+
+#pragma warning disable KTSU0004 // Use Ensure.NotNull instead of manual null check
+		if (key is null)
+		{
+			throw new ArgumentNullException(nameof(key), "The key cannot be null.");
+		}
+#pragma warning restore KTSU0004 // Use Ensure.NotNull instead of manual null check
+
+		// Without this overload the call binds to the IDictionary one, whose lookup-then-Add throws
+		// when another caller adds the key in between. GetOrAdd is atomic.
+		return dictionary.GetOrAdd(key, _ => new TVal());
+	}
+
 	/// <summary>
 	/// Method that gets a value from a dictionary if it exists, otherwise creates a new value and adds it to the dictionary.
 	/// </summary>