perf: use GetOrAdd args overload to eliminate closure allocations in event receivers (thomhurst#5222)

Replace TryGetValue + GetOrAdd(closure) pattern with GetOrAdd<TArg>(static lambda, args)
in First* event receiver methods. This removes redundant TryGetValue pre-checks
(GetOrAdd already does this internally) and avoids closure allocation by passing
context, sessionContext/assemblyContext/classContext, and cancellationToken as
explicit tuple args to a static lambda.

Author: thomhurst

TUnit.Core/Data/ThreadSafeDictionary.cs

@@ -1,4 +1,4 @@
-using System.Collections.Concurrent;
+using System.Collections.Concurrent;
 using System.Diagnostics.CodeAnalysis;
 
 namespace TUnit.Core.Data;
@@ -7,111 +7,69 @@ namespace TUnit.Core.Data;
 /// Provides a thread-safe dictionary with lazy value initialization.
 /// </summary>
 /// <typeparam name="TKey">The type of keys in the dictionary. Must be non-null.</typeparam>
-/// <typeparam name="TValue">The type of values in the dictionary. Must have a public parameterless constructor.</typeparam>
+/// <typeparam name="TValue">The type of values in the dictionary.</typeparam>
 /// <remarks>
-/// <para>
 /// This class provides a concurrent dictionary where values are lazily initialized on first access.
 /// Each value is created exactly once per key, even when accessed concurrently from multiple threads.
-/// </para>
-/// <para>
-/// The lazy initialization is thread-safe and uses <see cref="LazyThreadSafetyMode.ExecutionAndPublication"/>
-/// to ensure that the factory function is executed only once per key, even under concurrent access.
-/// </para>
-/// <para>
-/// This class is marked as hidden in the editor browsable state in non-debug builds as it's primarily
-/// intended for internal use within the TUnit framework.
-/// </para>
 /// </remarks>
-/// <example>
-/// <code>
-/// var cache = new ThreadSafeDictionary&lt;string, ExpensiveResource&gt;();
-///
-/// // Values are created lazily on first access
-/// var resource1 = cache.GetOrAdd("key1", k => new ExpensiveResource(k));
-///
-/// // Same key returns the same instance (thread-safe)
-/// var resource2 = cache.GetOrAdd("key1", k => new ExpensiveResource(k));
-/// // resource1 == resource2
-/// </code>
-/// </example>
 #if !DEBUG
 [System.ComponentModel.EditorBrowsable(System.ComponentModel.EditorBrowsableState.Never)]
 #endif
-public class ThreadSafeDictionary<TKey,
-    [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicParameterlessConstructor)] TValue>
+public class ThreadSafeDictionary<TKey, [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicParameterlessConstructor)] TValue>
     where TKey : notnull
 {
     private readonly ConcurrentDictionary<TKey, Lazy<TValue>> _innerDictionary = new();
 
     /// <summary>
     /// Gets a collection containing the keys in the dictionary.
     /// </summary>
-    /// <value>A collection of keys present in the dictionary.</value>
     public ICollection<TKey> Keys => _innerDictionary.Keys;
 
     /// <summary>
     /// Gets an enumerable collection of values in the dictionary.
     /// </summary>
-    /// <value>An enumerable of initialized values.</value>
-    /// <remarks>
-    /// Accessing this property will force initialization of all lazy values in the dictionary.
-    /// </remarks>
     public IEnumerable<TValue> Values => _innerDictionary.Values.Select(static lazy => lazy.Value);
 
     /// <summary>
     /// Gets the value associated with the specified key, or creates it if it doesn't exist.
+    /// The factory is guaranteed to run at most once per key even under concurrent access.
     /// </summary>
-    /// <param name="key">The key of the value to get or create.</param>
-    /// <param name="func">The factory function to create the value if the key doesn't exist.</param>
-    /// <returns>
-    /// The value for the key. This will be either the existing value or a newly created value from the factory function.
-    /// </returns>
     /// <remarks>
-    /// This method is thread-safe. If multiple threads call this method simultaneously with the same key,
-    /// the factory function will be executed only once, and all threads will receive the same instance.
-    /// This implementation uses a two-phase approach: TryGetValue for the fast path (when key exists),
-    /// and GetOrAdd with a pre-created Lazy for the slow path (new key). This prevents the factory
-    /// from being invoked multiple times during concurrent access.
+    /// Uses a fast-path TryGetValue to avoid allocating a Lazy on repeated calls,
+    /// then falls back to GetOrAdd with a pre-created Lazy to prevent the factory
+    /// running multiple times during a race.
     /// </remarks>
     public TValue GetOrAdd(TKey key, Func<TKey, TValue> func)
     {
-        // Fast path: Check if key already exists (lock-free read)
         if (_innerDictionary.TryGetValue(key, out var existingLazy))
         {
             return existingLazy.Value;
         }
 
-        return GetOrAddSlow(key, func);
+        var newLazy = new Lazy<TValue>(() => func(key), LazyThreadSafetyMode.ExecutionAndPublication);
+        var winning = _innerDictionary.GetOrAdd(key, newLazy);
+        return winning.Value;
     }
 
-    private TValue GetOrAddSlow(TKey key, Func<TKey, TValue> func)
+    /// <summary>
+    /// Gets the value associated with the specified key, or creates it using the factory and arg if it doesn't exist.
+    /// Avoids closure allocation at the call site by accepting the factory argument explicitly.
+    /// </summary>
+    public TValue GetOrAdd<TArg>(TKey key, Func<TKey, TArg, TValue> func, TArg arg)
     {
-        // Slow path: Key not found, need to create
-        // Create Lazy instance OUTSIDE of GetOrAdd to prevent factory from running during race
-        var newLazy = new Lazy<TValue>(() => func(key), LazyThreadSafetyMode.ExecutionAndPublication);
-
-        // Use GetOrAdd with VALUE (not factory) - atomic operation that either:
-        // 1. Adds our newLazy if key still doesn't exist
-        // 2. Returns existing Lazy if another thread just added one
-        var winningLazy = _innerDictionary.GetOrAdd(key, newLazy);
+        if (_innerDictionary.TryGetValue(key, out var existingLazy))
+        {
+            return existingLazy.Value;
+        }
 
-        // CRITICAL: Always return value from the Lazy that's actually in the dictionary
-        // This ensures only ONE factory execution even if multiple Lazy instances were created
-        return winningLazy.Value;
+        var newLazy = new Lazy<TValue>(() => func(key, arg), LazyThreadSafetyMode.ExecutionAndPublication);
+        var winning = _innerDictionary.GetOrAdd(key, newLazy);
+        return winning.Value;
     }
 
     /// <summary>
     /// Tries to get the value associated with the specified key.
     /// </summary>
-    /// <param name="key">The key of the value to get.</param>
-    /// <param name="value">
-    /// When this method returns, contains the value associated with the specified key if found;
-    /// otherwise, the default value for the type.
-    /// </param>
-    /// <returns><see langword="true"/> if the dictionary contains an element with the specified key; otherwise, <see langword="false"/>.</returns>
-    /// <remarks>
-    /// If the key exists, accessing the value will trigger lazy initialization if it hasn't occurred yet.
-    /// </remarks>
     public bool TryGetValue(TKey key, [NotNullWhen(true)] out TValue? value)
     {
         if (_innerDictionary.TryGetValue(key, out var lazy))
@@ -120,40 +78,28 @@ public bool TryGetValue(TKey key, [NotNullWhen(true)] out TValue? value)
             return true;
         }
 
-        value = default!;
+        value = default;
         return false;
     }
 
     /// <summary>
     /// Removes the value with the specified key from the dictionary.
+    /// Returns the value only if it had already been initialized; otherwise returns default.
     /// </summary>
-    /// <param name="key">The key of the value to remove.</param>
-    /// <returns>
-    /// The value that was removed, or the default value for the type if the key was not found.
-    /// </returns>
-    /// <remarks>
-    /// If the key exists and the value has been lazily initialized, that instance is returned.
-    /// Otherwise, the default value for <typeparamref name="TValue"/> is returned.
-    /// </remarks>
     public TValue? Remove(TKey key)
     {
-        if (_innerDictionary.TryRemove(key, out var lazy))
+        if (_innerDictionary.TryRemove(key, out var lazy) && lazy.IsValueCreated)
         {
             return lazy.Value;
         }
 
-        return default(TValue?);
+        return default;
     }
 
     /// <summary>
     /// Gets the value associated with the specified key.
     /// </summary>
-    /// <param name="key">The key of the value to get.</param>
-    /// <returns>The value associated with the specified key.</returns>
-    /// <exception cref="KeyNotFoundException">Thrown when the specified key is not found in the dictionary.</exception>
-    /// <remarks>
-    /// Accessing this indexer will trigger lazy initialization of the value if it hasn't occurred yet.
-    /// </remarks>
+    /// <exception cref="KeyNotFoundException">Thrown when the specified key is not found.</exception>
     public TValue this[TKey key] => _innerDictionary.TryGetValue(key, out var lazy)
         ? lazy.Value
         : throw new KeyNotFoundException($"Key '{key}' not found in dictionary");

TUnit.Engine/Services/EventReceiverOrchestrator.cs

@@ -321,13 +321,9 @@ public ValueTask InvokeFirstTestInSessionEventReceiversAsync(
             return default;
         }
 
-        if (_firstTestInSessionTasks.TryGetValue("session", out var existingTask))
-        {
-            return new ValueTask(existingTask);
-        }
-
         var task = _firstTestInSessionTasks.GetOrAdd("session",
-            _ => InvokeFirstTestInSessionEventReceiversCoreAsync(context, sessionContext, cancellationToken));
+            static (_, args) => args.self.InvokeFirstTestInSessionEventReceiversCoreAsync(args.context, args.sessionContext, args.cancellationToken),
+            (self: this, context, sessionContext, cancellationToken));
         return new ValueTask(task);
     }
 
@@ -357,13 +353,9 @@ public ValueTask InvokeFirstTestInAssemblyEventReceiversAsync(
 
         var assemblyName = assemblyContext.Assembly.GetName().FullName ?? "";
 
-        if (_firstTestInAssemblyTasks.TryGetValue(assemblyName, out var existingTask))
-        {
-            return new ValueTask(existingTask);
-        }
-
         var task = _firstTestInAssemblyTasks.GetOrAdd(assemblyName,
-            _ => InvokeFirstTestInAssemblyEventReceiversCoreAsync(context, assemblyContext, cancellationToken));
+            static (_, args) => args.self.InvokeFirstTestInAssemblyEventReceiversCoreAsync(args.context, args.assemblyContext, args.cancellationToken),
+            (self: this, context, assemblyContext, cancellationToken));
         return new ValueTask(task);
     }
 
@@ -393,13 +385,9 @@ public ValueTask InvokeFirstTestInClassEventReceiversAsync(
 
         var classType = classContext.ClassType;
 
-        if (_firstTestInClassTasks.TryGetValue(classType, out var existingTask))
-        {
-            return new ValueTask(existingTask);
-        }
-
         var task = _firstTestInClassTasks.GetOrAdd(classType,
-            _ => InvokeFirstTestInClassEventReceiversCoreAsync(context, classContext, cancellationToken));
+            static (_, args) => args.self.InvokeFirstTestInClassEventReceiversCoreAsync(args.context, args.classContext, args.cancellationToken),
+            (self: this, context, classContext, cancellationToken));
         return new ValueTask(task);
     }
 

TUnit.PublicAPI/Tests.Core_Library_Has_No_API_Changes.DotNet10_0.verified.txt

@@ -1787,6 +1787,7 @@ namespace .Data
         public .<TKey> Keys { get; }
         public .<TValue> Values { get; }
         public TValue GetOrAdd(TKey key, <TKey, TValue> func) { }
+        public TValue GetOrAdd<TArg>(TKey key, <TKey, TArg, TValue> func, TArg arg) { }
         public TValue? Remove(TKey key) { }
         public bool TryGetValue(TKey key, [.(true)] out TValue? value) { }
     }

TUnit.PublicAPI/Tests.Core_Library_Has_No_API_Changes.DotNet8_0.verified.txt

@@ -1787,6 +1787,7 @@ namespace .Data
         public .<TKey> Keys { get; }
         public .<TValue> Values { get; }
         public TValue GetOrAdd(TKey key, <TKey, TValue> func) { }
+        public TValue GetOrAdd<TArg>(TKey key, <TKey, TArg, TValue> func, TArg arg) { }
         public TValue? Remove(TKey key) { }
         public bool TryGetValue(TKey key, [.(true)] out TValue? value) { }
     }

TUnit.PublicAPI/Tests.Core_Library_Has_No_API_Changes.DotNet9_0.verified.txt

@@ -1787,6 +1787,7 @@ namespace .Data
         public .<TKey> Keys { get; }
         public .<TValue> Values { get; }
         public TValue GetOrAdd(TKey key, <TKey, TValue> func) { }
+        public TValue GetOrAdd<TArg>(TKey key, <TKey, TArg, TValue> func, TArg arg) { }
         public TValue? Remove(TKey key) { }
         public bool TryGetValue(TKey key, [.(true)] out TValue? value) { }
     }

TUnit.PublicAPI/Tests.Core_Library_Has_No_API_Changes.Net4_7.verified.txt

@@ -1733,6 +1733,7 @@ namespace .Data
         public .<TKey> Keys { get; }
         public .<TValue> Values { get; }
         public TValue GetOrAdd(TKey key, <TKey, TValue> func) { }
+        public TValue GetOrAdd<TArg>(TKey key, <TKey, TArg, TValue> func, TArg arg) { }
         public TValue? Remove(TKey key) { }
         public bool TryGetValue(TKey key, [.(true)] out TValue? value) { }
     }