perf: index-based constraint key scheduler for O(k) lookups (#4678)

Replace O(n) full queue drain with WaitingTestIndex that maps constraint
keys to waiting tests. When keys are released, only tests waiting on those
specific keys are examined. This reduces per-completion work from O(n) to
O(k) where k is the number of released keys, eliminating O(n^2) total
overhead for n constrained tests.

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: thomhurst

TUnit.Engine/Scheduling/ConstraintKeyScheduler.cs

@@ -1,4 +1,3 @@
-using System.Collections.Concurrent;
 using TUnit.Core;
 using TUnit.Core.Logging;
 using TUnit.Engine.Logging;
@@ -41,14 +40,14 @@ public async ValueTask ExecuteTestsWithConstraintsAsync(
         var lockedKeys = new HashSet<string>();
         var lockObject = new object();
 
-        // Queue for tests waiting for their constraint keys to become available
-        var waitingTests = new ConcurrentQueue<(AbstractExecutableTest Test, IReadOnlyList<string> ConstraintKeys, TaskCompletionSource<bool> StartSignal)>();
+        // Indexed structure for tests waiting for their constraint keys to become available
+        var waitingTestIndex = new WaitingTestIndex();
 
         // Active test tasks
         var activeTasks = new List<Task>();
 
         // Process each test
-        foreach (var (test, constraintKeys, _) in sortedTests)
+        foreach (var (test, constraintKeys, priority) in sortedTests)
         {
             var startSignal = new TaskCompletionSource<bool>();
 
@@ -75,6 +74,17 @@ public async ValueTask ExecuteTestsWithConstraintsAsync(
                         lockedKeys.Add(constraintKeys[i]);
                     }
                 }
+                else
+                {
+                    // Add to the indexed waiting structure while still under lock
+                    waitingTestIndex.Add(new WaitingTest
+                    {
+                        TestId = test.TestId,
+                        ConstraintKeys = constraintKeys,
+                        StartSignal = startSignal,
+                        Priority = priority
+                    });
+                }
             }
 
             if (canStart)
@@ -84,18 +94,17 @@ public async ValueTask ExecuteTestsWithConstraintsAsync(
                     await _logger.LogDebugAsync($"Starting test {test.TestId} with constraint keys: {string.Join(", ", constraintKeys)}").ConfigureAwait(false);
                 startSignal.SetResult(true);
 
-                var testTask = ExecuteTestAndReleaseKeysAsync(test, constraintKeys, lockedKeys, lockObject, waitingTests, cancellationToken);
+                var testTask = ExecuteTestAndReleaseKeysAsync(test, constraintKeys, lockedKeys, lockObject, waitingTestIndex, cancellationToken);
                 test.ExecutionTask = testTask;
                 activeTasks.Add(testTask);
             }
             else
             {
-                // Queue the test to wait for its keys
+                // Test was already added to the waiting index inside the lock above
                 if (_logger.IsDebugEnabled)
                     await _logger.LogDebugAsync($"Queueing test {test.TestId} waiting for constraint keys: {string.Join(", ", constraintKeys)}").ConfigureAwait(false);
-                waitingTests.Enqueue((test, constraintKeys, startSignal));
 
-                var testTask = WaitAndExecuteTestAsync(test, constraintKeys, startSignal, lockedKeys, lockObject, waitingTests, cancellationToken);
+                var testTask = WaitAndExecuteTestAsync(test, constraintKeys, startSignal, lockedKeys, lockObject, waitingTestIndex, cancellationToken);
                 test.ExecutionTask = testTask;
                 activeTasks.Add(testTask);
             }
@@ -114,7 +123,7 @@ private async Task WaitAndExecuteTestAsync(
         TaskCompletionSource<bool> startSignal,
         HashSet<string> lockedKeys,
         object lockObject,
-        ConcurrentQueue<(AbstractExecutableTest Test, IReadOnlyList<string> ConstraintKeys, TaskCompletionSource<bool> StartSignal)> waitingTests,
+        WaitingTestIndex waitingTestIndex,
         CancellationToken cancellationToken)
     {
         // Wait for signal to start
@@ -123,7 +132,7 @@ private async Task WaitAndExecuteTestAsync(
         if (_logger.IsDebugEnabled)
             await _logger.LogDebugAsync($"Starting previously queued test {test.TestId} with constraint keys: {string.Join(", ", constraintKeys)}").ConfigureAwait(false);
 
-        await ExecuteTestAndReleaseKeysAsync(test, constraintKeys, lockedKeys, lockObject, waitingTests, cancellationToken).ConfigureAwait(false);
+        await ExecuteTestAndReleaseKeysAsync(test, constraintKeys, lockedKeys, lockObject, waitingTestIndex, cancellationToken).ConfigureAwait(false);
     }
 
     #if NET6_0_OR_GREATER
@@ -134,7 +143,7 @@ private async Task ExecuteTestAndReleaseKeysAsync(
         IReadOnlyList<string> constraintKeys,
         HashSet<string> lockedKeys,
         object lockObject,
-        ConcurrentQueue<(AbstractExecutableTest Test, IReadOnlyList<string> ConstraintKeys, TaskCompletionSource<bool> StartSignal)> waitingTests,
+        WaitingTestIndex waitingTestIndex,
         CancellationToken cancellationToken)
     {
         SemaphoreSlim? parallelLimiterSemaphore = null;
@@ -158,9 +167,7 @@ private async Task ExecuteTestAndReleaseKeysAsync(
             parallelLimiterSemaphore?.Release();
 
             // Release the constraint keys and check if any waiting tests can now run
-            // Pre-allocate lists outside the lock to minimize lock duration
-            var testsToStart = new List<(AbstractExecutableTest Test, IReadOnlyList<string> ConstraintKeys, TaskCompletionSource<bool> StartSignal)>();
-            var testsToRequeue = new List<(AbstractExecutableTest Test, IReadOnlyList<string> ConstraintKeys, TaskCompletionSource<bool> StartSignal)>();
+            var testsToStart = new List<WaitingTest>();
 
             lock (lockObject)
             {
@@ -170,16 +177,23 @@ private async Task ExecuteTestAndReleaseKeysAsync(
                     lockedKeys.Remove(key);
                 }
 
-                // Check waiting tests to see if any can now run
+                // Only examine tests that are waiting on the keys we just released (O(k) lookup)
+                var candidates = waitingTestIndex.GetCandidatesForReleasedKeys(constraintKeys);
 
-                while (waitingTests.TryDequeue(out var waitingTest))
+                // Sort candidates by priority to respect ordering
+                // Use a simple list + sort rather than a SortedSet to avoid per-element allocation
+                var sortedCandidates = new List<WaitingTest>(candidates.Count);
+                sortedCandidates.AddRange(candidates);
+                sortedCandidates.Sort(static (a, b) => a.Priority.CompareTo(b.Priority));
+
+                foreach (var candidate in sortedCandidates)
                 {
-                    // Check if all constraint keys are available for this waiting test - manual loop avoids LINQ allocation
+                    // Check if all constraint keys are available for this candidate
                     var canStart = true;
-                    var waitingKeyCount = waitingTest.ConstraintKeys.Count;
+                    var waitingKeyCount = candidate.ConstraintKeys.Count;
                     for (var i = 0; i < waitingKeyCount; i++)
                     {
-                        if (lockedKeys.Contains(waitingTest.ConstraintKeys[i]))
+                        if (lockedKeys.Contains(candidate.ConstraintKeys[i]))
                         {
                             canStart = false;
                             break;
@@ -191,23 +205,14 @@ private async Task ExecuteTestAndReleaseKeysAsync(
                         // Lock the keys for this test
                         for (var i = 0; i < waitingKeyCount; i++)
                         {
-                            lockedKeys.Add(waitingTest.ConstraintKeys[i]);
+                            lockedKeys.Add(candidate.ConstraintKeys[i]);
                         }
 
-                        // Mark test to start after we exit the lock
-                        testsToStart.Add(waitingTest);
+                        // Remove from the index and mark for starting
+                        waitingTestIndex.Remove(candidate);
+                        testsToStart.Add(candidate);
                     }
-                    else
-                    {
-                        // Still can't run, keep it in the queue
-                        testsToRequeue.Add(waitingTest);
-                    }
-                }
-
-                // Re-add tests that still can't run
-                foreach (var waitingTestItem in testsToRequeue)
-                {
-                    waitingTests.Enqueue(waitingTestItem);
+                    // If can't start, leave it in the index for future key releases
                 }
             }
 
@@ -218,7 +223,7 @@ private async Task ExecuteTestAndReleaseKeysAsync(
             foreach (var testToStart in testsToStart)
             {
                 if (_logger.IsDebugEnabled)
-                    await _logger.LogDebugAsync($"Unblocking waiting test {testToStart.Test.TestId} with constraint keys: {string.Join(", ", testToStart.ConstraintKeys)}").ConfigureAwait(false);
+                    await _logger.LogDebugAsync($"Unblocking waiting test {testToStart.TestId} with constraint keys: {string.Join(", ", testToStart.ConstraintKeys)}").ConfigureAwait(false);
                 testToStart.StartSignal.SetResult(true);
             }
         }

TUnit.Engine/Scheduling/WaitingTestIndex.cs

@@ -0,0 +1,99 @@
+namespace TUnit.Engine.Scheduling;
+
+/// <summary>
+/// Represents a test that is waiting for its constraint keys to become available.
+/// </summary>
+internal sealed class WaitingTest
+{
+    public required string TestId { get; init; }
+    public required IReadOnlyList<string> ConstraintKeys { get; init; }
+    public required TaskCompletionSource<bool> StartSignal { get; init; }
+    public required int Priority { get; init; }
+}
+
+/// <summary>
+/// Index structure that maps constraint keys to waiting tests, enabling O(k) lookup
+/// when keys are released instead of scanning the entire queue.
+/// All operations must be performed under the caller's lock.
+/// </summary>
+internal sealed class WaitingTestIndex
+{
+    // Maps each constraint key to the set of tests waiting on that key
+    private readonly Dictionary<string, HashSet<WaitingTest>> _keyToTests = new();
+
+    // Master set of all waiting tests (for fast membership checks and count)
+    private readonly HashSet<WaitingTest> _allTests = new();
+
+    /// <summary>
+    /// Gets the number of waiting tests currently in the index.
+    /// </summary>
+    public int Count => _allTests.Count;
+
+    /// <summary>
+    /// Adds a waiting test to all key indexes.
+    /// </summary>
+    public void Add(WaitingTest waitingTest)
+    {
+        _allTests.Add(waitingTest);
+
+        var keys = waitingTest.ConstraintKeys;
+        var keyCount = keys.Count;
+        for (var i = 0; i < keyCount; i++)
+        {
+            var key = keys[i];
+            if (!_keyToTests.TryGetValue(key, out var tests))
+            {
+                tests = new HashSet<WaitingTest>();
+                _keyToTests[key] = tests;
+            }
+            tests.Add(waitingTest);
+        }
+    }
+
+    /// <summary>
+    /// Removes a waiting test from all key indexes.
+    /// </summary>
+    public void Remove(WaitingTest waitingTest)
+    {
+        if (!_allTests.Remove(waitingTest))
+        {
+            return;
+        }
+
+        var keys = waitingTest.ConstraintKeys;
+        var keyCount = keys.Count;
+        for (var i = 0; i < keyCount; i++)
+        {
+            var key = keys[i];
+            if (_keyToTests.TryGetValue(key, out var tests))
+            {
+                tests.Remove(waitingTest);
+                if (tests.Count == 0)
+                {
+                    _keyToTests.Remove(key);
+                }
+            }
+        }
+    }
+
+    /// <summary>
+    /// Returns a deduplicated set of waiting tests that are associated with any of the released keys.
+    /// These are candidates that might be unblocked (but still need to be checked against locked keys).
+    /// </summary>
+    public HashSet<WaitingTest> GetCandidatesForReleasedKeys(IReadOnlyList<string> releasedKeys)
+    {
+        var candidates = new HashSet<WaitingTest>();
+
+        var keyCount = releasedKeys.Count;
+        for (var i = 0; i < keyCount; i++)
+        {
+            if (_keyToTests.TryGetValue(releasedKeys[i], out var tests))
+            {
+                // HashSet.UnionWith handles deduplication
+                candidates.UnionWith(tests);
+            }
+        }
+
+        return candidates;
+    }
+}