refactor: optimize task handling by replacing Task with ValueTask for improved performance (#3528)

thomhurst · web-flow · commit b3f06cfac19d · 2025-10-26T17:29:04.000Z
diff --git a/TUnit.Engine/Events/EventReceiverRegistry.cs b/TUnit.Engine/Events/EventReceiverRegistry.cs
@@ -26,6 +26,7 @@ private enum EventTypes
     
     private volatile EventTypes _registeredEvents = EventTypes.None;
     private readonly Dictionary<Type, object[]> _receiversByType = new();
+    private readonly Dictionary<Type, Array> _cachedTypedReceivers = new();
     private readonly ReaderWriterLockSlim _lock = new();
     
     /// <summary>
@@ -66,7 +67,7 @@ public void RegisterReceiver(object receiver)
     private void RegisterReceiverInternal(object receiver)
     {
         UpdateEventFlags(receiver);
-        
+
         // Register for each interface type the object implements
         // We use a simpler approach that doesn't require reflection
         RegisterIfImplements<ITestStartEventReceiver>(receiver);
@@ -79,6 +80,8 @@ private void RegisterReceiverInternal(object receiver)
         RegisterIfImplements<ILastTestInAssemblyEventReceiver>(receiver);
         RegisterIfImplements<IFirstTestInClassEventReceiver>(receiver);
         RegisterIfImplements<ILastTestInClassEventReceiver>(receiver);
+
+        _cachedTypedReceivers.Clear();
     }
     
     private void RegisterIfImplements<T>(object receiver) where T : class
@@ -136,30 +139,67 @@ private void RegisterIfImplements<T>(object receiver) where T : class
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     public bool HasAnyReceivers() => _registeredEvents != EventTypes.None;
     
-    /// <summary>
-    /// Get receivers of specific type (for invocation)
-    /// </summary>
     public T[] GetReceiversOfType<T>() where T : class
     {
+        var typeKey = typeof(T);
+
         _lock.EnterReadLock();
         try
         {
-            if (_receiversByType.TryGetValue(typeof(T), out var receivers))
+            if (_cachedTypedReceivers.TryGetValue(typeKey, out var cached))
+            {
+                return (T[])cached;
+            }
+        }
+        finally
+        {
+            _lock.ExitReadLock();
+        }
+
+        _lock.EnterUpgradeableReadLock();
+        try
+        {
+            if (_cachedTypedReceivers.TryGetValue(typeKey, out var cached))
+            {
+                return (T[])cached;
+            }
+
+            if (_receiversByType.TryGetValue(typeKey, out var receivers))
             {
-                // Cast array to specific type
                 var typedArray = new T[receivers.Length];
                 for (var i = 0; i < receivers.Length; i++)
                 {
                     typedArray[i] = (T)receivers[i];
                 }
+
+                _lock.EnterWriteLock();
+                try
+                {
+                    _cachedTypedReceivers[typeKey] = typedArray;
+                }
+                finally
+                {
+                    _lock.ExitWriteLock();
+                }
+
                 return typedArray;
             }
-            return [
-            ];
+
+            T[] emptyArray = [];
+            _lock.EnterWriteLock();
+            try
+            {
+                _cachedTypedReceivers[typeKey] = emptyArray;
+            }
+            finally
+            {
+                _lock.ExitWriteLock();
+            }
+            return emptyArray;
         }
         finally
         {
-            _lock.ExitReadLock();
+            _lock.ExitUpgradeableReadLock();
         }
     }
     
diff --git a/TUnit.Engine/Services/BeforeHookTaskCache.cs b/TUnit.Engine/Services/BeforeHookTaskCache.cs
@@ -17,33 +17,34 @@ internal sealed class BeforeHookTaskCache
     private Task? _beforeTestSessionTask;
     private readonly object _testSessionLock = new();
 
-    public Task GetOrCreateBeforeTestSessionTask(Func<Task> taskFactory)
+    public ValueTask GetOrCreateBeforeTestSessionTask(Func<ValueTask> taskFactory)
     {
         if (_beforeTestSessionTask != null)
         {
-            return _beforeTestSessionTask;
+            return new ValueTask(_beforeTestSessionTask);
         }
 
         lock (_testSessionLock)
         {
-            // Double-check after acquiring lock
             if (_beforeTestSessionTask == null)
             {
-                _beforeTestSessionTask = taskFactory();
+                _beforeTestSessionTask = taskFactory().AsTask();
             }
-            return _beforeTestSessionTask;
+            return new ValueTask(_beforeTestSessionTask);
         }
     }
 
-    public Task GetOrCreateBeforeAssemblyTask(Assembly assembly, Func<Assembly, Task> taskFactory)
+    public ValueTask GetOrCreateBeforeAssemblyTask(Assembly assembly, Func<Assembly, ValueTask> taskFactory)
     {
-        return _beforeAssemblyTasks.GetOrAdd(assembly, taskFactory);
+        var task = _beforeAssemblyTasks.GetOrAdd(assembly, a => taskFactory(a).AsTask());
+        return new ValueTask(task);
     }
 
-    public Task GetOrCreateBeforeClassTask(
+    public ValueTask GetOrCreateBeforeClassTask(
         [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors | DynamicallyAccessedMemberTypes.PublicProperties | DynamicallyAccessedMemberTypes.PublicMethods)]
-        Type testClass, Func<Type, Task> taskFactory)
+        Type testClass, Func<Type, ValueTask> taskFactory)
     {
-        return _beforeClassTasks.GetOrAdd(testClass, taskFactory);
+        var task = _beforeClassTasks.GetOrAdd(testClass, t => taskFactory(t).AsTask());
+        return new ValueTask(task);
     }
 }
diff --git a/TUnit.Engine/Services/EventReceiverOrchestrator.cs b/TUnit.Engine/Services/EventReceiverOrchestrator.cs
@@ -214,20 +214,19 @@ public async ValueTask InvokeHookRegistrationEventReceiversAsync(HookRegisteredC
 
     // First/Last event methods with fast-path checks
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    public async ValueTask InvokeFirstTestInSessionEventReceiversAsync(
+    public ValueTask InvokeFirstTestInSessionEventReceiversAsync(
         TestContext context,
         TestSessionContext sessionContext,
         CancellationToken cancellationToken)
     {
         if (!_registry.HasFirstTestInSessionReceivers())
         {
-            return;
+            return default;
         }
 
-        // Use GetOrAdd to ensure exactly one task is created per session and all tests await it
         var task = _firstTestInSessionTasks.GetOrAdd("session",
             _ => InvokeFirstTestInSessionEventReceiversCoreAsync(context, sessionContext, cancellationToken));
-        await task;
+        return new ValueTask(task);
     }
 
     private async Task InvokeFirstTestInSessionEventReceiversCoreAsync(
@@ -244,21 +243,20 @@ private async Task InvokeFirstTestInSessionEventReceiversCoreAsync(
     }
 
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    public async ValueTask InvokeFirstTestInAssemblyEventReceiversAsync(
+    public ValueTask InvokeFirstTestInAssemblyEventReceiversAsync(
         TestContext context,
         AssemblyHookContext assemblyContext,
         CancellationToken cancellationToken)
     {
         if (!_registry.HasFirstTestInAssemblyReceivers())
         {
-            return;
+            return default;
         }
 
         var assemblyName = assemblyContext.Assembly.GetName().FullName ?? "";
-        // Use GetOrAdd to ensure exactly one task is created per assembly and all tests await it
         var task = _firstTestInAssemblyTasks.GetOrAdd(assemblyName,
             _ => InvokeFirstTestInAssemblyEventReceiversCoreAsync(context, assemblyContext, cancellationToken));
-        await task;
+        return new ValueTask(task);
     }
 
     private async Task InvokeFirstTestInAssemblyEventReceiversCoreAsync(
@@ -275,21 +273,20 @@ private async Task InvokeFirstTestInAssemblyEventReceiversCoreAsync(
     }
 
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    public async ValueTask InvokeFirstTestInClassEventReceiversAsync(
+    public ValueTask InvokeFirstTestInClassEventReceiversAsync(
         TestContext context,
         ClassHookContext classContext,
         CancellationToken cancellationToken)
     {
         if (!_registry.HasFirstTestInClassReceivers())
         {
-            return;
+            return default;
         }
 
         var classType = classContext.ClassType;
-        // Use GetOrAdd to ensure exactly one task is created per class and all tests await it
         var task = _firstTestInClassTasks.GetOrAdd(classType,
             _ => InvokeFirstTestInClassEventReceiversCoreAsync(context, classContext, cancellationToken));
-        await task;
+        return new ValueTask(task);
     }
 
     private async Task InvokeFirstTestInClassEventReceiversCoreAsync(
diff --git a/TUnit.Engine/Services/HookExecutor.cs b/TUnit.Engine/Services/HookExecutor.cs

Original file line number	Diff line number	Diff line change
`@@ -17,33 +17,34 @@ internal sealed class BeforeHookTaskCache`
`17`	`17`	`private Task? _beforeTestSessionTask;`
`18`	`18`	`private readonly object _testSessionLock = new();`
`19`	`19`
`20`		`- public Task GetOrCreateBeforeTestSessionTask(Func<Task> taskFactory)`
	`20`	`+ public ValueTask GetOrCreateBeforeTestSessionTask(Func<ValueTask> taskFactory)`
`21`	`21`	`{`
`22`	`22`	`if (_beforeTestSessionTask != null)`
`23`	`23`	`{`
`24`		`- return _beforeTestSessionTask;`
	`24`	`+ return new ValueTask(_beforeTestSessionTask);`
`25`	`25`	`}`
`26`	`26`
`27`	`27`	`lock (_testSessionLock)`
`28`	`28`	`{`
`29`		`- // Double-check after acquiring lock`
`30`	`29`	`if (_beforeTestSessionTask == null)`
`31`	`30`	`{`
`32`		`- _beforeTestSessionTask = taskFactory();`
	`31`	`+ _beforeTestSessionTask = taskFactory().AsTask();`
`33`	`32`	`}`
`34`		`- return _beforeTestSessionTask;`
	`33`	`+ return new ValueTask(_beforeTestSessionTask);`
`35`	`34`	`}`
`36`	`35`	`}`
`37`	`36`
`38`		`- public Task GetOrCreateBeforeAssemblyTask(Assembly assembly, Func<Assembly, Task> taskFactory)`
	`37`	`+ public ValueTask GetOrCreateBeforeAssemblyTask(Assembly assembly, Func<Assembly, ValueTask> taskFactory)`
`39`	`38`	`{`
`40`		`- return _beforeAssemblyTasks.GetOrAdd(assembly, taskFactory);`
	`39`	`+ var task = _beforeAssemblyTasks.GetOrAdd(assembly, a => taskFactory(a).AsTask());`
	`40`	`+ return new ValueTask(task);`
`41`	`41`	`}`
`42`	`42`
`43`		`- public Task GetOrCreateBeforeClassTask(`
	`43`	`+ public ValueTask GetOrCreateBeforeClassTask(`
`44`	`44`	`[DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors \| DynamicallyAccessedMemberTypes.PublicProperties \| DynamicallyAccessedMemberTypes.PublicMethods)]`
`45`		`- Type testClass, Func<Type, Task> taskFactory)`
	`45`	`+ Type testClass, Func<Type, ValueTask> taskFactory)`
`46`	`46`	`{`
`47`		`- return _beforeClassTasks.GetOrAdd(testClass, taskFactory);`
	`47`	`+ var task = _beforeClassTasks.GetOrAdd(testClass, t => taskFactory(t).AsTask());`
	`48`	`+ return new ValueTask(task);`
`48`	`49`	`}`
`49`	`50`	`}`
Original file line number	Diff line number	Diff line change
`@@ -214,20 +214,19 @@ public async ValueTask InvokeHookRegistrationEventReceiversAsync(HookRegisteredC`
`214`	`214`
`215`	`215`	`// First/Last event methods with fast-path checks`
`216`	`216`	`[MethodImpl(MethodImplOptions.AggressiveInlining)]`
`217`		`- public async ValueTask InvokeFirstTestInSessionEventReceiversAsync(`
	`217`	`+ public ValueTask InvokeFirstTestInSessionEventReceiversAsync(`
`218`	`218`	`TestContext context,`
`219`	`219`	`TestSessionContext sessionContext,`
`220`	`220`	`CancellationToken cancellationToken)`
`221`	`221`	`{`
`222`	`222`	`if (!_registry.HasFirstTestInSessionReceivers())`
`223`	`223`	`{`
`224`		`- return;`
	`224`	`+ return default;`
`225`	`225`	`}`
`226`	`226`
`227`		`- // Use GetOrAdd to ensure exactly one task is created per session and all tests await it`
`228`	`227`	`var task = _firstTestInSessionTasks.GetOrAdd("session",`
`229`	`228`	`_ => InvokeFirstTestInSessionEventReceiversCoreAsync(context, sessionContext, cancellationToken));`
`230`		`- await task;`
	`229`	`+ return new ValueTask(task);`
`231`	`230`	`}`
`232`	`231`
`233`	`232`	`private async Task InvokeFirstTestInSessionEventReceiversCoreAsync(`
`@@ -244,21 +243,20 @@ private async Task InvokeFirstTestInSessionEventReceiversCoreAsync(`
`244`	`243`	`}`
`245`	`244`
`246`	`245`	`[MethodImpl(MethodImplOptions.AggressiveInlining)]`
`247`		`- public async ValueTask InvokeFirstTestInAssemblyEventReceiversAsync(`
	`246`	`+ public ValueTask InvokeFirstTestInAssemblyEventReceiversAsync(`
`248`	`247`	`TestContext context,`
`249`	`248`	`AssemblyHookContext assemblyContext,`
`250`	`249`	`CancellationToken cancellationToken)`
`251`	`250`	`{`
`252`	`251`	`if (!_registry.HasFirstTestInAssemblyReceivers())`
`253`	`252`	`{`
`254`		`- return;`
	`253`	`+ return default;`
`255`	`254`	`}`
`256`	`255`
`257`	`256`	`var assemblyName = assemblyContext.Assembly.GetName().FullName ?? "";`
`258`		`- // Use GetOrAdd to ensure exactly one task is created per assembly and all tests await it`
`259`	`257`	`var task = _firstTestInAssemblyTasks.GetOrAdd(assemblyName,`
`260`	`258`	`_ => InvokeFirstTestInAssemblyEventReceiversCoreAsync(context, assemblyContext, cancellationToken));`
`261`		`- await task;`
	`259`	`+ return new ValueTask(task);`
`262`	`260`	`}`
`263`	`261`
`264`	`262`	`private async Task InvokeFirstTestInAssemblyEventReceiversCoreAsync(`
`@@ -275,21 +273,20 @@ private async Task InvokeFirstTestInAssemblyEventReceiversCoreAsync(`
`275`	`273`	`}`
`276`	`274`
`277`	`275`	`[MethodImpl(MethodImplOptions.AggressiveInlining)]`
`278`		`- public async ValueTask InvokeFirstTestInClassEventReceiversAsync(`
	`276`	`+ public ValueTask InvokeFirstTestInClassEventReceiversAsync(`
`279`	`277`	`TestContext context,`
`280`	`278`	`ClassHookContext classContext,`
`281`	`279`	`CancellationToken cancellationToken)`
`282`	`280`	`{`
`283`	`281`	`if (!_registry.HasFirstTestInClassReceivers())`
`284`	`282`	`{`
`285`		`- return;`
	`283`	`+ return default;`
`286`	`284`	`}`
`287`	`285`
`288`	`286`	`var classType = classContext.ClassType;`
`289`		`- // Use GetOrAdd to ensure exactly one task is created per class and all tests await it`
`290`	`287`	`var task = _firstTestInClassTasks.GetOrAdd(classType,`
`291`	`288`	`_ => InvokeFirstTestInClassEventReceiversCoreAsync(context, classContext, cancellationToken));`
`292`		`- await task;`
	`289`	`+ return new ValueTask(task);`
`293`	`290`	`}`
`294`	`291`
`295`	`292`	`private async Task InvokeFirstTestInClassEventReceiversCoreAsync(`