Fix a race condition in WorkScheduler

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;
 
+    /// <summary>
+    /// Gets a value indicating whether taken tasks have not been completed yet.
+    /// The scheduler lock must be held.
+    /// </summary>
     public bool HasRunningTasks => _running > 0;
 
+    /// <summary>
+    /// Marks a taken task as completed and returns <c>true</c> if no taken tasks remain.
+    /// The scheduler lock must be held.
+    /// </summary>
     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;
     }
 
+    /// <summary>
+    /// Takes the next task, preferring immediate over deferred tasks, and counts it as running.
+    /// The scheduler lock must be held.
+    /// </summary>
     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)
         }
     }
 
+    /// <summary>
+    /// Removes all tasks and resets the running count.
+    /// Only valid on an idle or discarded scheduler, and does not take the scheduler lock.
+    /// </summary>
     public void Clear()
     {
         _immediateStack.Clear();

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)
     /// </summary>
     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)
     /// </summary>
     private void TryDispatchPendingBatchesUnsafe()
     {
+        AssertLockHeld();
+
         if (_pendingBatches.Count == 0)
         {
             return;
@@ -117,6 +123,8 @@ private void TryDispatchPendingBatchesUnsafe()
     /// </summary>
     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;

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;

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;

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<string>();
+        var failures = new ConcurrentQueue<string>();
+
+        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<StressQuery>()
+            .AddDataLoader<StressRelatedDataLoader>()
+            .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<string> 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<StressItem> 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<string> GetNameAsync()
+        {
+            await Task.Yield();
+            return "item-" + id;
+        }
+
+        public async Task<StressItem> GetChildAsync()
+        {
+            await Task.Yield();
+            return new StressItem(id * 1000);
+        }
+
+        public async Task<StressItem> GetRelatedAsync(
+            StressRelatedDataLoader dataLoader,
+            CancellationToken cancellationToken)
+            => (await dataLoader.LoadAsync(id, cancellationToken))!;
+    }
+
+    public sealed class StressRelatedDataLoader(IBatchScheduler batchScheduler, DataLoaderOptions options)
+        : BatchDataLoader<int, StressItem>(batchScheduler, options)
+    {
+        protected override async Task<IReadOnlyDictionary<int, StressItem>> LoadBatchAsync(
+            IReadOnlyList<int> 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";
+}