diff --git a/src/HotChocolate/Core/src/Types/Execution/Processing/WorkQueue.cs b/src/HotChocolate/Core/src/Types/Execution/Processing/WorkQueue.cs
index 194bbe97d38..1dd31ca3f89 100644
--- a/src/HotChocolate/Core/src/Types/Execution/Processing/WorkQueue.cs
+++ b/src/HotChocolate/Core/src/Types/Execution/Processing/WorkQueue.cs
@@ -10,11 +10,19 @@ internal sealed class WorkQueue
public bool IsEmpty => _immediateStack.Count == 0 && _deferredStack.Count == 0;
+ ///
+ /// Gets a value indicating whether taken tasks have not been completed yet.
+ /// The scheduler lock must be held.
+ ///
public bool HasRunningTasks => _running > 0;
+ ///
+ /// Marks a taken task as completed and returns true if no taken tasks remain.
+ /// The scheduler lock must be held.
+ ///
public bool Complete()
{
- var value = Interlocked.Decrement(ref _running);
+ var value = --_running;
if (value < 0)
{
@@ -24,12 +32,16 @@ public bool Complete()
return value is 0;
}
+ ///
+ /// Takes the next task, preferring immediate over deferred tasks, and counts it as running.
+ /// The scheduler lock must be held.
+ ///
public bool TryTake([MaybeNullWhen(false)] out IExecutionTask executionTask)
{
if (_immediateStack.TryPop(out executionTask)
|| _deferredStack.TryPop(out executionTask))
{
- Interlocked.Increment(ref _running);
+ _running++;
return true;
}
@@ -50,6 +62,10 @@ public void Push(IExecutionTask executionTask)
}
}
+ ///
+ /// Removes all tasks and resets the running count.
+ /// Only valid on an idle or discarded scheduler, and does not take the scheduler lock.
+ ///
public void Clear()
{
_immediateStack.Clear();
diff --git a/src/HotChocolate/Core/src/Types/Execution/Processing/WorkScheduler.Batching.cs b/src/HotChocolate/Core/src/Types/Execution/Processing/WorkScheduler.Batching.cs
index 64d3ca71017..c0e14f96467 100644
--- a/src/HotChocolate/Core/src/Types/Execution/Processing/WorkScheduler.Batching.cs
+++ b/src/HotChocolate/Core/src/Types/Execution/Processing/WorkScheduler.Batching.cs
@@ -50,6 +50,8 @@ public void RegisterBatchEntry(
[MethodImpl(MethodImplOptions.AggressiveInlining)]
private void IncrementPathCountUnsafe(SelectionPath path)
{
+ AssertLockHeld();
+
ref var count = ref CollectionsMarshal.GetValueRefOrAddDefault(_activePaths, path, out _);
count++;
}
@@ -60,6 +62,8 @@ private void IncrementPathCountUnsafe(SelectionPath path)
///
private void DecrementPathCountUnsafe(SelectionPath path)
{
+ AssertLockHeld();
+
ref var count = ref CollectionsMarshal.GetValueRefOrNullRef(_activePaths, path);
if (!Unsafe.IsNullRef(ref count) && --count <= 0)
@@ -75,6 +79,8 @@ private void DecrementPathCountUnsafe(SelectionPath path)
///
private void TryDispatchPendingBatchesUnsafe()
{
+ AssertLockHeld();
+
if (_pendingBatches.Count == 0)
{
return;
@@ -117,6 +123,8 @@ private void TryDispatchPendingBatchesUnsafe()
///
private bool CanDispatchBatchUnsafe(SelectionPath batchPath)
{
+ AssertLockHeld();
+
// Walk up ancestor paths. If any ancestor still has active tasks,
// more entries could still be added to this batch.
var ancestor = batchPath.Parent;
diff --git a/src/HotChocolate/Core/src/Types/Execution/Processing/WorkScheduler.Execute.cs b/src/HotChocolate/Core/src/Types/Execution/Processing/WorkScheduler.Execute.cs
index 3f471aea09f..0cbd7c59000 100644
--- a/src/HotChocolate/Core/src/Types/Execution/Processing/WorkScheduler.Execute.cs
+++ b/src/HotChocolate/Core/src/Types/Execution/Processing/WorkScheduler.Execute.cs
@@ -208,6 +208,8 @@ private void TryDispatchOrComplete()
private void TryDispatchOrCompleteUnsafe(bool isWaitingForTaskCompletion)
{
+ AssertLockHeld();
+
if (_isCompleted)
{
return;
diff --git a/src/HotChocolate/Core/src/Types/Execution/Processing/WorkScheduler.cs b/src/HotChocolate/Core/src/Types/Execution/Processing/WorkScheduler.cs
index 6506e8f4092..0de4506318c 100644
--- a/src/HotChocolate/Core/src/Types/Execution/Processing/WorkScheduler.cs
+++ b/src/HotChocolate/Core/src/Types/Execution/Processing/WorkScheduler.cs
@@ -1,3 +1,4 @@
+using System.Diagnostics;
using System.Runtime.CompilerServices;
namespace HotChocolate.Execution.Processing;
@@ -104,11 +105,11 @@ public void Complete(IExecutionTask task)
switch (task)
{
case Tasks.ResolverTask resolverTask:
- CompleteBranchTask(task.BranchId);
-
- if (work.Complete())
+ lock (_sync)
{
- lock (_sync)
+ CompleteBranchTaskUnsafe(task.BranchId);
+
+ if (CompleteWorkUnsafe(work))
{
_completed.Add(resolverTask.Id);
DecrementPathCountUnsafe(resolverTask.FieldSelectionPath);
@@ -124,7 +125,7 @@ public void Complete(IExecutionTask task)
CompleteBranchTaskUnsafe(additionalBranchId);
}
- if (work.Complete())
+ if (CompleteWorkUnsafe(work))
{
_completed.Add(task.Id);
DecrementPathCountUnsafe(batchResolverTask.FieldSelectionPath);
@@ -133,11 +134,11 @@ public void Complete(IExecutionTask task)
break;
default:
- CompleteBranchTask(task.BranchId);
-
- if (work.Complete())
+ lock (_sync)
{
- lock (_sync)
+ CompleteBranchTaskUnsafe(task.BranchId);
+
+ if (CompleteWorkUnsafe(work))
{
_completed.Add(task.Id);
}
@@ -148,8 +149,26 @@ public void Complete(IExecutionTask task)
_signal.Set();
}
+ private bool CompleteWorkUnsafe(WorkQueue work)
+ {
+ AssertLockHeld();
+ return work.Complete();
+ }
+
+ [Conditional("DEBUG")]
+ private void AssertLockHeld()
+ {
+#if NET9_0_OR_GREATER
+ Debug.Assert(_sync.IsHeldByCurrentThread);
+#else
+ Debug.Assert(Monitor.IsEntered(_sync));
+#endif
+ }
+
private void RegisterBranchTaskUnsafe(int branchId)
{
+ AssertLockHeld();
+
if (branchId == BranchTracker.SystemBranchId)
{
return;
@@ -164,27 +183,11 @@ private void RegisterBranchTaskUnsafe(int branchId)
branch.RegisterTask();
}
- private void CompleteBranchTask(int branchId)
- {
- if (branchId == BranchTracker.SystemBranchId)
- {
- return;
- }
-
- lock (_sync)
- {
- if (_activeBranches.TryGetValue(branchId, out var branch)
- && branch.CompleteTask())
- {
- _activeBranches.Remove(branchId);
- branch.Complete();
- }
- }
- }
-
[MethodImpl(MethodImplOptions.AggressiveInlining)]
private void CompleteBranchTaskUnsafe(int branchId)
{
+ AssertLockHeld();
+
if (branchId == BranchTracker.SystemBranchId)
{
return;
diff --git a/src/HotChocolate/Core/test/Execution.Tests/Processing/WorkSchedulerStressTests.cs b/src/HotChocolate/Core/test/Execution.Tests/Processing/WorkSchedulerStressTests.cs
new file mode 100644
index 00000000000..7b9365ab329
--- /dev/null
+++ b/src/HotChocolate/Core/test/Execution.Tests/Processing/WorkSchedulerStressTests.cs
@@ -0,0 +1,183 @@
+using System.Collections.Concurrent;
+using System.Diagnostics;
+using GreenDonut;
+using Microsoft.Extensions.DependencyInjection;
+
+namespace HotChocolate.Execution.Processing;
+
+// Opt-in stress test for the completion bookkeeping of the WorkScheduler. It only runs when
+// HC_STRESS_TESTS=1 is set; HC_STRESS_SECONDS overrides the default duration of 20 seconds.
+// It lives in its own non-parallel collection because it observes process-wide unobserved task
+// exceptions, which concurrently running tests could otherwise contribute to.
+[Collection(WorkSchedulerStressCollection.Name)]
+public class WorkSchedulerStressTests(ITestOutputHelper output)
+{
+ private const string EnableVariable = "HC_STRESS_TESTS";
+ private const string SecondsVariable = "HC_STRESS_SECONDS";
+ private const int DefaultSeconds = 20;
+ private const int Concurrency = 64;
+ private static readonly TimeSpan s_requestTimeout = TimeSpan.FromSeconds(30);
+
+ private const string Query =
+ """
+ {
+ items(count: 16) {
+ id
+ name
+ related { id name }
+ child { id name related { id } }
+ }
+ }
+ """;
+
+ [Fact]
+ public async Task Execute_Should_NotProduceUnobservedExceptions_When_RequestsRunConcurrently()
+ {
+ // arrange
+ if (Environment.GetEnvironmentVariable(EnableVariable) != "1")
+ {
+ Assert.Skip($"Set {EnableVariable}=1 to run the WorkScheduler stress test.");
+ }
+
+ var duration = TimeSpan.FromSeconds(
+ int.TryParse(Environment.GetEnvironmentVariable(SecondsVariable), out var seconds) && seconds > 0
+ ? seconds
+ : DefaultSeconds);
+ var cancellationToken = TestContext.Current.CancellationToken;
+
+ var unobserved = new ConcurrentQueue();
+ var failures = new ConcurrentQueue();
+
+ void OnUnobservedException(object? sender, UnobservedTaskExceptionEventArgs args)
+ {
+ foreach (var exception in args.Exception.InnerExceptions)
+ {
+ unobserved.Enqueue(exception.ToString());
+ }
+
+ args.SetObserved();
+ }
+
+ var executor = await new ServiceCollection()
+ .AddGraphQL()
+ .AddQueryType()
+ .AddDataLoader()
+ .BuildRequestExecutorAsync(cancellationToken: cancellationToken);
+
+ long requests = 0;
+ TaskScheduler.UnobservedTaskException += OnUnobservedException;
+
+ try
+ {
+ // act
+ var stopwatch = Stopwatch.StartNew();
+
+ while (stopwatch.Elapsed < duration)
+ {
+ var round = new Task[Concurrency];
+
+ for (var i = 0; i < Concurrency; i++)
+ {
+ round[i] = Task.Run(
+ () => ExecuteRequestAsync(executor, failures, cancellationToken),
+ cancellationToken);
+ }
+
+ await Task.WhenAll(round);
+ requests += Concurrency;
+
+ GC.Collect();
+ GC.WaitForPendingFinalizers();
+ GC.Collect();
+ }
+ }
+ finally
+ {
+ TaskScheduler.UnobservedTaskException -= OnUnobservedException;
+ }
+
+ // assert
+ output.WriteLine($"requests={requests} failures={failures.Count} unobserved={unobserved.Count}");
+
+ Assert.True(requests > 0);
+ Assert.Empty(failures);
+ Assert.Empty(unobserved);
+ }
+
+ private static async Task ExecuteRequestAsync(
+ IRequestExecutor executor,
+ ConcurrentQueue failures,
+ CancellationToken cancellationToken)
+ {
+ try
+ {
+ using var requestCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+ requestCts.CancelAfter(s_requestTimeout);
+
+ // The extra WaitAsync bounds a request that stops observing its cancellation token.
+ await using var result = await executor
+ .ExecuteAsync(Query, requestCts.Token)
+ .WaitAsync(s_requestTimeout + TimeSpan.FromSeconds(5), cancellationToken);
+
+ if (result.ExpectOperationResult().Errors is { Count: > 0 } errors)
+ {
+ failures.Enqueue(string.Join(" | ", errors.Select(e => e.Message)));
+ }
+ }
+ catch (Exception ex) when (!cancellationToken.IsCancellationRequested)
+ {
+ failures.Enqueue(ex.ToString());
+ }
+ }
+
+ public sealed class StressQuery
+ {
+ public IEnumerable GetItems(int count)
+ {
+ for (var i = 0; i < count; i++)
+ {
+ yield return new StressItem(i);
+ }
+ }
+ }
+
+ public sealed class StressItem(int id)
+ {
+ public int Id => id;
+
+ public async Task GetNameAsync()
+ {
+ await Task.Yield();
+ return "item-" + id;
+ }
+
+ public async Task GetChildAsync()
+ {
+ await Task.Yield();
+ return new StressItem(id * 1000);
+ }
+
+ public async Task GetRelatedAsync(
+ StressRelatedDataLoader dataLoader,
+ CancellationToken cancellationToken)
+ => (await dataLoader.LoadAsync(id, cancellationToken))!;
+ }
+
+ public sealed class StressRelatedDataLoader(IBatchScheduler batchScheduler, DataLoaderOptions options)
+ : BatchDataLoader(batchScheduler, options)
+ {
+ protected override async Task> LoadBatchAsync(
+ IReadOnlyList keys,
+ CancellationToken cancellationToken)
+ {
+ await Task.Yield();
+ return keys.ToDictionary(k => k, k => new StressItem(k + 100000));
+ }
+ }
+}
+
+[CollectionDefinition(Name, DisableParallelization = true)]
+public sealed class WorkSchedulerStressCollection
+{
+ public const string Name = "WorkSchedulerStress";
+}