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"; +}