perf: consolidate test dedup + hook register guards (#5528 A)

TUnit.Engine/Framework/TUnitServiceProvider.cs

@@ -179,7 +179,6 @@ public TUnitServiceProvider(IExtension extension,
 
         TestExecutor = Register(new TestExecutor(hookExecutor, lifecycleCoordinator, beforeHookTaskCache, afterHookPairTracker, ContextProvider, EventReceiverOrchestrator));
 
-        var testExecutionGuard = Register(new TestExecutionGuard());
         var testStateManager = Register(new TestStateManager());
         var testContextRestorer = Register(new TestContextRestorer());
         var testMethodInvoker = Register(new TestMethodInvoker());
@@ -223,7 +222,6 @@ public TUnitServiceProvider(IExtension extension,
         // Create the new TestCoordinator that orchestrates the granular services
         var testCoordinator = Register<ITestCoordinator>(
             new TestCoordinator(
-                testExecutionGuard,
                 testStateManager,
                 MessageBus,
                 testContextRestorer,

TUnit.Engine/Scheduling/TestRunner.cs

@@ -39,6 +39,9 @@ internal TestRunner(
         _parallelLimitLockProvider = parallelLimitLockProvider;
     }
 
+    // Dedup ledger for re-entrant ExecuteTestAsync calls (dependency recursion, scheduler races).
+    // Entries are intentionally retained for the session: a late dependency lookup must still
+    // observe the in-flight or completed TCS. Session-scoped lifetime bounds growth to O(test count).
     private readonly ConcurrentDictionary<string, TaskCompletionSource<bool>> _executingTests = new();
     private Exception? _firstFailFastException;
 

TUnit.Engine/Services/AfterHookPairTracker.cs

@@ -23,7 +23,21 @@ internal sealed class AfterHookPairTracker
     // Ensure only the first call to RegisterAfterTestSessionHook registers a callback.
     // Subsequent calls (e.g. from per-test timeout tokens) are ignored so that
     // a test timeout cannot prematurely trigger session-level After hooks.
-    private volatile bool _sessionHookRegistered;
+    // Use Interlocked.CompareExchange to avoid TOCTOU race where two threads both
+    // observe 0 and both proceed to register.
+    private int _sessionHookRegistered;
+
+    // Per-test callers would otherwise register one CancellationTokenRegistration per test
+    // for every assembly/class — 10k tests across 5 assemblies = 50k redundant registrations.
+    // First registration wins; subsequent calls short-circuit.
+    //
+    // Safety of single registration: The CancellationToken passed in is always the session-scoped
+    // token (or a derivative from Parallel.ForEachAsync which is itself linked to the session CT).
+    // Per-test timeout tokens are applied inside TestExecutor via TimeoutHelper.CreateLinkedTokenSource
+    // scoped to the test body only — they never reach this method. Therefore when the session cancels,
+    // the first test's registered CT still fires regardless of whether that test has completed.
+    private readonly ConcurrentHashSet<Assembly> _assemblyHookRegistered = new();
+    private readonly ConcurrentHashSet<Type> _classHookRegistered = new();
 
     // Track cancellation registrations for cleanup
     private readonly ConcurrentBag<CancellationTokenRegistration> _registrations = [];
@@ -37,13 +51,11 @@ public void RegisterAfterTestSessionHook(
         CancellationToken cancellationToken,
         Func<ValueTask<List<Exception>>> afterHookExecutor)
     {
-        if (_sessionHookRegistered)
+        if (Interlocked.CompareExchange(ref _sessionHookRegistered, 1, 0) != 0)
         {
             return;
         }
 
-        _sessionHookRegistered = true;
-
         // Register callback to run After hook on cancellation
         var registration = cancellationToken.Register(static state =>
         {
@@ -65,6 +77,11 @@ public void RegisterAfterAssemblyHook(
         CancellationToken cancellationToken,
         Func<Assembly, ValueTask<List<Exception>>> afterHookExecutor)
     {
+        if (!_assemblyHookRegistered.Add(assembly))
+        {
+            return;
+        }
+
         var registration = cancellationToken.Register(static state =>
         {
             var (pairTracker, assembly, afterHookExecutor) = ((AfterHookPairTracker, Assembly, Func<Assembly, ValueTask<List<Exception>>>))state!;
@@ -86,6 +103,11 @@ public void RegisterAfterClassHook(
         HookExecutor hookExecutor,
         CancellationToken cancellationToken)
     {
+        if (!_classHookRegistered.Add(testClass))
+        {
+            return;
+        }
+
         var registration = cancellationToken.Register(static state =>
         {
             var (pairTracker, testClass, hookExecutor) = ((AfterHookPairTracker, Type, HookExecutor))state!;

TUnit.Engine/Services/TestExecution/TestCoordinator.cs

@@ -15,7 +15,6 @@ namespace TUnit.Engine.Services.TestExecution;
 /// </summary>
 internal sealed class TestCoordinator : ITestCoordinator
 {
-    private readonly TestExecutionGuard _executionGuard;
     private readonly TestStateManager _stateManager;
     private readonly ITUnitMessageBus _messageBus;
     private readonly TestContextRestorer _contextRestorer;
@@ -26,7 +25,6 @@ internal sealed class TestCoordinator : ITestCoordinator
     private readonly EventReceiverOrchestrator _eventReceiverOrchestrator;
 
     public TestCoordinator(
-        TestExecutionGuard executionGuard,
         TestStateManager stateManager,
         ITUnitMessageBus messageBus,
         TestContextRestorer contextRestorer,
@@ -36,7 +34,6 @@ public TestCoordinator(
         TUnitFrameworkLogger logger,
         EventReceiverOrchestrator eventReceiverOrchestrator)
     {
-        _executionGuard = executionGuard;
         _stateManager = stateManager;
         _messageBus = messageBus;
         _contextRestorer = contextRestorer;
@@ -47,9 +44,11 @@ public TestCoordinator(
         _eventReceiverOrchestrator = eventReceiverOrchestrator;
     }
 
+    // Dedup happens in TestRunner via its own ConcurrentDictionary<string, TCS<bool>> —
+    // it's the single entry point for both scheduler and dependency recursion, so a second
+    // guard here would just double the TCS/dict allocations per test.
     public ValueTask ExecuteTestAsync(AbstractExecutableTest test, CancellationToken cancellationToken)
-        => _executionGuard.TryStartExecutionAsync(test.TestId,
-            () => ExecuteTestInternalAsync(test, cancellationToken));
+        => ExecuteTestInternalAsync(test, cancellationToken);
 
     private async ValueTask ExecuteTestInternalAsync(AbstractExecutableTest test, CancellationToken cancellationToken)
     {