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/Interfaces/ITestCoordinator.cs

@@ -5,8 +5,14 @@
 namespace TUnit.Engine.Interfaces;
 
 /// <summary>
-/// Interface for executing a single test
+/// Interface for executing a single test.
 /// </summary>
+/// <remarks>
+/// Callers MUST invoke <see cref="ExecuteTestAsync"/> through <c>TestRunner</c>, not directly.
+/// <c>TestRunner</c> owns the dedup ledger that prevents double-execution when a test is reached
+/// via both the scheduler and dependency recursion. Bypassing <c>TestRunner</c> loses that
+/// guarantee and can cause the same test to run twice.
+/// </remarks>
 internal interface ITestCoordinator
 {
     ValueTask ExecuteTestAsync(AbstractExecutableTest test, CancellationToken cancellationToken);

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 = [];
@@ -34,18 +48,16 @@ internal sealed class AfterHookPairTracker
     /// This prevents per-test timeout tokens from prematurely firing session hooks.
     /// </summary>
     public void RegisterAfterTestSessionHook(
-        CancellationToken cancellationToken,
+        CancellationToken sessionCancellationToken,
         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 =>
+        var registration = sessionCancellationToken.Register(static state =>
         {
             var (pairTracker, afterHookExecutor) = ((AfterHookPairTracker, Func<ValueTask<List<Exception>>>))state!;
             // Use sync-over-async here because CancellationToken.Register requires Action (not Func<Task>)
@@ -60,12 +72,22 @@ public void RegisterAfterTestSessionHook(
     /// Registers Assembly After hooks to run on cancellation or normal completion.
     /// Ensures After hooks run exactly once even if called both ways.
     /// </summary>
+    /// <param name="sessionCancellationToken">
+    /// MUST be the session-scoped token (or a derivative linked to it). Only the first caller per
+    /// assembly registers a callback; subsequent calls short-circuit. Passing a narrower per-test
+    /// token would cause After hooks to miss cancellation when later tests' tokens fire.
+    /// </param>
     public void RegisterAfterAssemblyHook(
         Assembly assembly,
-        CancellationToken cancellationToken,
+        CancellationToken sessionCancellationToken,
         Func<Assembly, ValueTask<List<Exception>>> afterHookExecutor)
     {
-        var registration = cancellationToken.Register(static state =>
+        if (!_assemblyHookRegistered.Add(assembly))
+        {
+            return;
+        }
+
+        var registration = sessionCancellationToken.Register(static state =>
         {
             var (pairTracker, assembly, afterHookExecutor) = ((AfterHookPairTracker, Assembly, Func<Assembly, ValueTask<List<Exception>>>))state!;
             _ = pairTracker.GetOrCreateAfterAssemblyTask(assembly, afterHookExecutor);
@@ -78,15 +100,25 @@ public void RegisterAfterAssemblyHook(
     /// Registers Class After hooks to run on cancellation or normal completion.
     /// Ensures After hooks run exactly once even if called both ways.
     /// </summary>
+    /// <param name="sessionCancellationToken">
+    /// MUST be the session-scoped token (or a derivative linked to it). Only the first caller per
+    /// class registers a callback; subsequent calls short-circuit. Passing a narrower per-test
+    /// token would cause After hooks to miss cancellation when later tests' tokens fire.
+    /// </param>
      [UnconditionalSuppressMessage("Trimming", "IL2077",
             Justification = "Type parameter is annotated at the method boundary.")]
     public void RegisterAfterClassHook(
         [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors | DynamicallyAccessedMemberTypes.PublicProperties | DynamicallyAccessedMemberTypes.PublicMethods)]
         Type testClass,
         HookExecutor hookExecutor,
-        CancellationToken cancellationToken)
+        CancellationToken sessionCancellationToken)
     {
-        var registration = cancellationToken.Register(static state =>
+        if (!_classHookRegistered.Add(testClass))
+        {
+            return;
+        }
+
+        var registration = sessionCancellationToken.Register(static state =>
         {
             var (pairTracker, testClass, hookExecutor) = ((AfterHookPairTracker, Type, HookExecutor))state!;
             _ = pairTracker.GetOrCreateAfterClassTask(testClass, hookExecutor, CancellationToken.None);

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)
     {