diff --git a/TUnit.Engine/Framework/TUnitServiceProvider.cs b/TUnit.Engine/Framework/TUnitServiceProvider.cs index 3707790daf..b2a10470c6 100644 --- a/TUnit.Engine/Framework/TUnitServiceProvider.cs +++ b/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( new TestCoordinator( - testExecutionGuard, testStateManager, MessageBus, testContextRestorer, diff --git a/TUnit.Engine/Interfaces/ITestCoordinator.cs b/TUnit.Engine/Interfaces/ITestCoordinator.cs index 608fee019b..79dcb9aa9a 100644 --- a/TUnit.Engine/Interfaces/ITestCoordinator.cs +++ b/TUnit.Engine/Interfaces/ITestCoordinator.cs @@ -5,8 +5,14 @@ namespace TUnit.Engine.Interfaces; /// -/// Interface for executing a single test +/// Interface for executing a single test. /// +/// +/// Callers MUST invoke through TestRunner, not directly. +/// TestRunner owns the dedup ledger that prevents double-execution when a test is reached +/// via both the scheduler and dependency recursion. Bypassing TestRunner loses that +/// guarantee and can cause the same test to run twice. +/// internal interface ITestCoordinator { ValueTask ExecuteTestAsync(AbstractExecutableTest test, CancellationToken cancellationToken); diff --git a/TUnit.Engine/Scheduling/TestRunner.cs b/TUnit.Engine/Scheduling/TestRunner.cs index ab9999e8eb..5ff9208591 100644 --- a/TUnit.Engine/Scheduling/TestRunner.cs +++ b/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> _executingTests = new(); private Exception? _firstFailFastException; diff --git a/TUnit.Engine/Services/AfterHookPairTracker.cs b/TUnit.Engine/Services/AfterHookPairTracker.cs index cfb8be41a0..ca85515c3f 100644 --- a/TUnit.Engine/Services/AfterHookPairTracker.cs +++ b/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 _assemblyHookRegistered = new(); + private readonly ConcurrentHashSet _classHookRegistered = new(); // Track cancellation registrations for cleanup private readonly ConcurrentBag _registrations = []; @@ -34,18 +48,16 @@ internal sealed class AfterHookPairTracker /// This prevents per-test timeout tokens from prematurely firing session hooks. /// public void RegisterAfterTestSessionHook( - CancellationToken cancellationToken, + CancellationToken sessionCancellationToken, Func>> 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>>))state!; // Use sync-over-async here because CancellationToken.Register requires Action (not Func) @@ -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. /// + /// + /// 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. + /// public void RegisterAfterAssemblyHook( Assembly assembly, - CancellationToken cancellationToken, + CancellationToken sessionCancellationToken, Func>> 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>>))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. /// + /// + /// 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. + /// [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); diff --git a/TUnit.Engine/Services/TestExecution/TestCoordinator.cs b/TUnit.Engine/Services/TestExecution/TestCoordinator.cs index a54aedbf6a..9dcdf8c394 100644 --- a/TUnit.Engine/Services/TestExecution/TestCoordinator.cs +++ b/TUnit.Engine/Services/TestExecution/TestCoordinator.cs @@ -15,7 +15,6 @@ namespace TUnit.Engine.Services.TestExecution; /// 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> — + // 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) { diff --git a/TUnit.Engine/Services/TestExecution/TestExecutionGuard.cs b/TUnit.Engine/Services/TestExecution/TestExecutionGuard.cs deleted file mode 100644 index 32fdc49a77..0000000000 --- a/TUnit.Engine/Services/TestExecution/TestExecutionGuard.cs +++ /dev/null @@ -1,54 +0,0 @@ -using System.Collections.Concurrent; - -namespace TUnit.Engine.Services.TestExecution; - -/// -/// Prevents duplicate test execution using thread-safe mechanisms. -/// Single Responsibility: Execution deduplication. -/// -internal sealed class TestExecutionGuard -{ - private readonly ConcurrentDictionary> _executingTests = new(); - - public ValueTask TryStartExecutionAsync(string testId, Func executionFunc) - { - // Fast path: check if test is already executing without allocating a TCS - if (_executingTests.TryGetValue(testId, out var existingTcs)) - { - return new ValueTask(WaitForExistingExecutionAsync(existingTcs)); - } - - var tcs = new TaskCompletionSource(); - existingTcs = _executingTests.GetOrAdd(testId, tcs); - - if (existingTcs != tcs) - { - return new ValueTask(WaitForExistingExecutionAsync(existingTcs)); - } - - return ExecuteAndCompleteAsync(testId, tcs, executionFunc); - } - - private static async Task WaitForExistingExecutionAsync(TaskCompletionSource tcs) - { - await tcs.Task.ConfigureAwait(false); - } - - private async ValueTask ExecuteAndCompleteAsync(string testId, TaskCompletionSource tcs, Func executionFunc) - { - try - { - await executionFunc().ConfigureAwait(false); - tcs.SetResult(true); - } - catch (Exception ex) - { - tcs.SetException(ex); - throw; - } - finally - { - _executingTests.TryRemove(testId, out _); - } - } -} \ No newline at end of file