Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
thomhurst merged 1 commit into from
Feb 6, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 39 additions & 34 deletions TUnit.Engine/Scheduling/ConstraintKeyScheduler.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
using System.Collections.Concurrent;
using TUnit.Core;
using TUnit.Core.Logging;
using TUnit.Engine.Logging;
Expand Down Expand Up @@ -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>();

Expand All @@ -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)
Expand All @@ -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);
}
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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;
Expand All @@ -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)
{
Expand All @@ -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;
Expand All @@ -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
}
}

Expand All @@ -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);
}
}
Expand Down
99 changes: 99 additions & 0 deletions TUnit.Engine/Scheduling/WaitingTestIndex.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;
}
}
Loading