diff --git a/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/DeferDisconnectTests.cs b/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/DeferDisconnectTests.cs new file mode 100644 index 00000000000..ef3ef6cab63 --- /dev/null +++ b/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/DeferDisconnectTests.cs @@ -0,0 +1,121 @@ +using System.Text; +#if !NET11_0_OR_GREATER +using System.Net.Http.Json; +#endif +using HotChocolate.AspNetCore.Tests.Utilities; +using Microsoft.AspNetCore.Builder; +using Microsoft.AspNetCore.TestHost; +using Microsoft.Extensions.DependencyInjection; + +namespace HotChocolate.AspNetCore; + +public class DeferDisconnectTests(TestServerFactory serverFactory) : ServerTestBase(serverFactory) +{ + [Theory] + [InlineData("multipart/mixed")] + [InlineData("text/event-stream")] + public async Task Post_Should_ReturnNoErrors_When_ClientDisconnectedDuringDeferredResolver( + string acceptHeader) + { + // arrange + var gates = new DisconnectGates(); + using var server = CreateServer(gates); + var client = server.CreateClient(); + using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(30)); + using var disconnect = new CancellationTokenSource(); + + using var request = new HttpRequestMessage(HttpMethod.Post, "/graphql"); + request.Content = JsonContent.Create(new { query = "{ fast ... @defer { slow } }" }); + request.Headers.Add("Accept", acceptHeader); + + // act + var response = await client.SendAsync( + request, + HttpCompletionOption.ResponseHeadersRead, + disconnect.Token); + + var body = await response.Content.ReadAsStreamAsync(timeout.Token); + await ReadUntilAsync(body, "\"hasNext\":true", timeout.Token); + await gates.SlowEntered.Task.WaitAsync(timeout.Token); + + await disconnect.CancelAsync(); + await body.DisposeAsync(); + response.Dispose(); + await Task.Delay(TimeSpan.FromMilliseconds(300), timeout.Token); + + gates.SlowGate.SetResult(); + await gates.SlowReturned.Task.WaitAsync(timeout.Token); + await Task.Delay(TimeSpan.FromMilliseconds(300), timeout.Token); + + using var followUp = await client.PostAsync( + "/graphql", + JsonContent.Create(new { query = "{ fast }" }), + timeout.Token); + var followUpBody = await followUp.Content.ReadAsStringAsync(timeout.Token); + + // assert + followUpBody.MatchInlineSnapshot( + """ + {"data":{"fast":"fast-value"}} + """); + } + + private static async Task ReadUntilAsync( + Stream stream, + string expected, + CancellationToken cancellationToken) + { + var buffer = new byte[4096]; + var text = new StringBuilder(); + + while (!text.ToString().Contains(expected, StringComparison.Ordinal)) + { + var read = await stream.ReadAsync(buffer, cancellationToken); + + if (read == 0) + { + throw new InvalidOperationException("The response ended before the initial payload arrived."); + } + + text.Append(Encoding.UTF8.GetString(buffer, 0, read)); + } + } + + private TestServer CreateServer(DisconnectGates gates) + => ServerFactory.Create( + services => services + .AddSingleton(gates) + .AddRouting() + .AddGraphQLServer() + .AddQueryType() + .AddDefaultBatchDispatcher() + .ModifyOptions(o => o.EnableDefer = true), + app => app + .UseRouting() + .UseEndpoints(endpoints => endpoints.MapGraphQL())); + + public sealed class DisconnectGates + { + public TaskCompletionSource SlowEntered { get; } = Create(); + + public TaskCompletionSource SlowGate { get; } = Create(); + + public TaskCompletionSource SlowReturned { get; } = Create(); + + private static TaskCompletionSource Create() + => new(TaskCreationOptions.RunContinuationsAsynchronously); + } + + public sealed class DisconnectQuery + { + public string GetFast() => "fast-value"; + + public async Task GetSlowAsync([Service] DisconnectGates gates) + { + gates.SlowEntered.TrySetResult(); + await gates.SlowGate.Task; + gates.SlowReturned.TrySetResult(); + return null!; + } + } +} diff --git a/src/HotChocolate/Core/src/Types/Execution/DependencyInjection/InternalServiceCollectionExtensions.cs b/src/HotChocolate/Core/src/Types/Execution/DependencyInjection/InternalServiceCollectionExtensions.cs index 599340b1c94..5e5cd979abb 100644 --- a/src/HotChocolate/Core/src/Types/Execution/DependencyInjection/InternalServiceCollectionExtensions.cs +++ b/src/HotChocolate/Core/src/Types/Execution/DependencyInjection/InternalServiceCollectionExtensions.cs @@ -178,6 +178,13 @@ public override bool Return(OperationContext obj) return true; } + // a context with outstanding tasks must stay untouched, because those tasks still + // use it. It is discarded instead of reused. + if (obj.HasOutstandingTasks) + { + return false; + } + if (obj.IsSharedScheduler) { obj.ResetScheduler(); diff --git a/src/HotChocolate/Core/src/Types/Execution/Processing/OperationContext.Execution.cs b/src/HotChocolate/Core/src/Types/Execution/Processing/OperationContext.Execution.cs index d808befd2d5..223be9271c8 100644 --- a/src/HotChocolate/Core/src/Types/Execution/Processing/OperationContext.Execution.cs +++ b/src/HotChocolate/Core/src/Types/Execution/Processing/OperationContext.Execution.cs @@ -37,7 +37,14 @@ public int ExecutionBranchId } } - public OperationResultBuilder Result { get; } = new(); + public OperationResultBuilder Result + { + get + { + AssertInitialized(); + return _result; + } + } public RequestContext RequestContext { @@ -58,6 +65,7 @@ public ResolverTask CreateResolverTask( { AssertInitialized(); + TaskCreated(); var resolverTask = _resolverTaskFactory.Create(); resolverTask.Initialize( @@ -80,6 +88,7 @@ public BatchResolverTask CreateBatchResolverTask( { AssertInitialized(); + TaskCreated(); var batchTask = _batchResolverTaskFactory.Create(); batchTask.Initialize(this, field, selectionPath, branchId, deferUsage); return batchTask; diff --git a/src/HotChocolate/Core/src/Types/Execution/Processing/OperationContext.Pooling.cs b/src/HotChocolate/Core/src/Types/Execution/Processing/OperationContext.Pooling.cs index 5408e188673..78a6f4772de 100644 --- a/src/HotChocolate/Core/src/Types/Execution/Processing/OperationContext.Pooling.cs +++ b/src/HotChocolate/Core/src/Types/Execution/Processing/OperationContext.Pooling.cs @@ -21,6 +21,7 @@ internal sealed partial class OperationContext private readonly BranchTracker _branchTracker = new(); private readonly WorkScheduler _workScheduler; private readonly DeferExecutionCoordinator _deferExecutionCoordinator = new(); + private readonly OperationResultBuilder _result = new(); private WorkScheduler _currentWorkScheduler; private BranchTracker _currentBranchTracker; private DeferExecutionCoordinator _currentDeferExecutionCoordinator; @@ -41,6 +42,9 @@ internal sealed partial class OperationContext private MemoryArena? _memory; private int _branchId; private int _variableIndex; + private int _outstandingTasks; + private OperationContext? _taskRoot; + private TaskCompletionSource? _idleSignal; private object? _rootValue; private bool _propagateNullValues; private bool _isInitialized; @@ -65,6 +69,65 @@ public OperationContext( public bool IsSharedScheduler => !ReferenceEquals(_workScheduler, _currentWorkScheduler); + /// + /// Gets whether tasks created from this context, or from a defer context derived from it, + /// have not yet been returned to their pool. + /// + public bool HasOutstandingTasks => Volatile.Read(ref _outstandingTasks) > 0; + + /// + /// Registers a task that was created from this context. + /// + public void TaskCreated() + { + Interlocked.Increment(ref _outstandingTasks); + + if (_taskRoot is { } root) + { + Interlocked.Increment(ref root._outstandingTasks); + } + } + + /// + /// Releases a task that was registered with . + /// + public void TaskReturned() + { + var root = _taskRoot; + + if (Interlocked.Decrement(ref _outstandingTasks) == 0) + { + Volatile.Read(ref _idleSignal)?.TrySetResult(); + } + + if (root is not null) + { + Interlocked.Decrement(ref root._outstandingTasks); + } + } + + /// + /// Waits until all tasks created from this context have been returned to their pool. + /// Only one caller may wait at a time. + /// + public ValueTask WaitForOutstandingTasksAsync() + { + if (!HasOutstandingTasks) + { + return ValueTask.CompletedTask; + } + + var signal = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + Interlocked.Exchange(ref _idleSignal, signal); + + if (!HasOutstandingTasks) + { + signal.TrySetResult(); + } + + return new ValueTask(signal.Task); + } + public void Initialize( RequestContext requestContext, IServiceProvider scopedServices, @@ -77,6 +140,7 @@ public void Initialize( CancellationToken requestAbortedOverride = default, MemoryArena? memoryArena = null) { + _result.Reset(); _requestContext = requestContext; _schema = Unsafe.As(requestContext.Schema); _memory = memoryArena @@ -109,9 +173,9 @@ public void Initialize( IncludeConditionFlags = operation.CreateIncludeConditionFlags(variables); DeferConditionFlags = operation.CreateDeferConditionFlags(variables); - Result.Data = new ResultDocument(_memory, operation, IncludeConditionFlags); - Result.RequestIndex = _requestContext.RequestIndex; - Result.VariableIndex = variableIndex; + _result.Data = new ResultDocument(_memory, operation, IncludeConditionFlags); + _result.RequestIndex = _requestContext.RequestIndex; + _result.VariableIndex = variableIndex; _currentBranchTracker = _branchTracker; _currentWorkScheduler = _workScheduler; @@ -131,6 +195,7 @@ public void InitializeDeferContext( int executionBranchId, DeferUsage deferUsage) { + _result.Reset(); _requestContext = context._requestContext; _schema = context._schema; _errorHandler = context._errorHandler; @@ -149,6 +214,7 @@ public void InitializeDeferContext( _currentBranchTracker = context._currentBranchTracker; _currentWorkScheduler = context._currentWorkScheduler; _currentDeferExecutionCoordinator = context._currentDeferExecutionCoordinator; + _taskRoot = context._taskRoot ?? context; _propagateNullValues = context._propagateNullValues; _branchId = executionBranchId; _isInitialized = true; @@ -156,7 +222,7 @@ public void InitializeDeferContext( IncludeConditionFlags = context.IncludeConditionFlags; DeferConditionFlags = context.DeferConditionFlags; - Result.Data = new ResultDocument( + _result.Data = new ResultDocument( context._memory!, context.Operation, selectionSet, @@ -164,8 +230,8 @@ public void InitializeDeferContext( context.IncludeConditionFlags, context.DeferConditionFlags, deferUsage); - Result.RequestIndex = _requestContext.RequestIndex; - Result.VariableIndex = context._variableIndex; + _result.RequestIndex = _requestContext.RequestIndex; + _result.VariableIndex = context._variableIndex; } public void InitializeWorkSchedulerFrom(OperationContext context) @@ -203,11 +269,13 @@ public void Clean() _batchDispatcher = null!; _memory = null; _branchId = int.MinValue; + _taskRoot = null; + _idleSignal = null; _propagateNullValues = false; _isInitialized = false; IncludeConditionFlags = default; DeferConditionFlags = default; - Result.Reset(); + _result.Reset(); } } diff --git a/src/HotChocolate/Core/src/Types/Execution/Processing/Tasks/BatchResolverTask.cs b/src/HotChocolate/Core/src/Types/Execution/Processing/Tasks/BatchResolverTask.cs index c1a7cc19638..a2469e53203 100644 --- a/src/HotChocolate/Core/src/Types/Execution/Processing/Tasks/BatchResolverTask.cs +++ b/src/HotChocolate/Core/src/Types/Execution/Processing/Tasks/BatchResolverTask.cs @@ -30,6 +30,7 @@ internal sealed class BatchResolverTask : IResolverTask private ObjectField _field = null!; private SelectionPath _selectionPath = null!; private int _branchId; + private bool _isContextReleased; public BatchResolverTask( ObjectPool objectPool, @@ -164,7 +165,10 @@ private async ValueTask ExecuteAsync(CancellationToken cancellationToken) } finally { - _operationContext.Scheduler.Complete(this); + // the context can be reused as soon as it is released, so the scheduler is read first. + var scheduler = _operationContext.Scheduler; + ReleaseOperationContext(); + scheduler.Complete(this); for (var i = 0; i < contexts.Length; i++) { @@ -440,6 +444,7 @@ public void Initialize( DeferUsage? deferUsage) { _operationContext = operationContext; + _isContextReleased = false; _field = field; _selectionPath = selectionPath; _branchId = branchId; @@ -463,6 +468,7 @@ internal bool Reset() _rentedArgs.Clear(); _branchIds.Clear(); + ReleaseOperationContext(); _operationContext = null!; _field = null!; _selectionPath = null!; @@ -477,6 +483,26 @@ internal bool Reset() return true; } + /// + /// Tells the operation context that this task and its resolver tasks no longer use it. + /// Calling it again has no effect. + /// + private void ReleaseOperationContext() + { + foreach (var task in _resolverTasks) + { + task.ReleaseOperationContext(); + } + + if (_isContextReleased || _operationContext is null) + { + return; + } + + _isContextReleased = true; + _operationContext.TaskReturned(); + } + /// /// Represents a single entry in the batch — one parent object and its result location. /// diff --git a/src/HotChocolate/Core/src/Types/Execution/Processing/Tasks/DeferTask.cs b/src/HotChocolate/Core/src/Types/Execution/Processing/Tasks/DeferTask.cs index 2a3223c6e54..60e3978e44f 100644 --- a/src/HotChocolate/Core/src/Types/Execution/Processing/Tasks/DeferTask.cs +++ b/src/HotChocolate/Core/src/Types/Execution/Processing/Tasks/DeferTask.cs @@ -62,11 +62,12 @@ protected override async ValueTask ExecuteAsync(CancellationToken cancellationTo deferContext.DeferExecutionCoordinator.EnqueueResult(deferContext.BuildResult(), _executionBranchId); } - protected override ValueTask OnAfterCompletedAsync(CancellationToken cancellationToken) + protected override async ValueTask OnAfterCompletedAsync(CancellationToken cancellationToken) { - // TODO : we need to give the context back here and not rest once we have a pool. + // the branch can still be running when the request was aborted, so the defer context + // is only given back once all of its tasks have been returned. + await _deferContextOwner.OperationContext.WaitForOutstandingTasksAsync().ConfigureAwait(false); Reset(); - return ValueTask.CompletedTask; } public void Initialize( diff --git a/src/HotChocolate/Core/src/Types/Execution/Processing/Tasks/ResolverTask.Execute.cs b/src/HotChocolate/Core/src/Types/Execution/Processing/Tasks/ResolverTask.Execute.cs index 4a92d31e6f9..98aaa62c261 100644 --- a/src/HotChocolate/Core/src/Types/Execution/Processing/Tasks/ResolverTask.Execute.cs +++ b/src/HotChocolate/Core/src/Types/Execution/Processing/Tasks/ResolverTask.Execute.cs @@ -58,7 +58,7 @@ private async ValueTask ExecuteAsync(CancellationToken cancellationToken) } finally { - _operationContext.Scheduler.Complete(this); + CompleteOnScheduler(); objectPool.Return(this); } } @@ -253,7 +253,7 @@ public async ValueTask CompleteUnsafeAsync() } Status = _completionStatus; - _operationContext.Scheduler.Complete(this); + CompleteOnScheduler(); objectPool.Return(this); } } diff --git a/src/HotChocolate/Core/src/Types/Execution/Processing/Tasks/ResolverTask.Pooling.cs b/src/HotChocolate/Core/src/Types/Execution/Processing/Tasks/ResolverTask.Pooling.cs index a3bd0dffebd..4312c62b1f9 100644 --- a/src/HotChocolate/Core/src/Types/Execution/Processing/Tasks/ResolverTask.Pooling.cs +++ b/src/HotChocolate/Core/src/Types/Execution/Processing/Tasks/ResolverTask.Pooling.cs @@ -18,6 +18,7 @@ public void Initialize( DeferUsage? deferUsage) { _operationContext = operationContext; + _isContextReleased = false; _selection = selection; _context.Initialize(parent, selection, resultValue, operationContext, deferUsage, scopedContextData); _context.BranchId = executionBranchId; @@ -32,6 +33,7 @@ public void Initialize( /// Always true. internal bool Reset() { + ReleaseOperationContext(); _completionStatus = ExecutionTaskStatus.Completed; _operationContext = null!; _selection = null!; @@ -48,4 +50,26 @@ internal bool Reset() _args.Clear(); return true; } + + /// + /// Tells the operation context that this task no longer uses it. Calling it again has no effect. + /// + internal void ReleaseOperationContext() + { + if (_isContextReleased || _operationContext is null) + { + return; + } + + _isContextReleased = true; + _operationContext.TaskReturned(); + } + + private void CompleteOnScheduler() + { + // the context can be reused as soon as it is released, so the scheduler is read first. + var scheduler = _operationContext.Scheduler; + ReleaseOperationContext(); + scheduler.Complete(this); + } } diff --git a/src/HotChocolate/Core/src/Types/Execution/Processing/Tasks/ResolverTask.cs b/src/HotChocolate/Core/src/Types/Execution/Processing/Tasks/ResolverTask.cs index c68e8eb1139..7506130ba20 100644 --- a/src/HotChocolate/Core/src/Types/Execution/Processing/Tasks/ResolverTask.cs +++ b/src/HotChocolate/Core/src/Types/Execution/Processing/Tasks/ResolverTask.cs @@ -11,6 +11,7 @@ internal sealed partial class ResolverTask(ObjectPool objectPool) private readonly Dictionary _args = new Dictionary(StringComparer.Ordinal); private OperationContext _operationContext = null!; + private bool _isContextReleased; private Selection _selection = null!; private ExecutionTaskStatus _completionStatus = ExecutionTaskStatus.Completed; diff --git a/src/HotChocolate/Core/test/Execution.Tests/DeferContextPoolingTests.cs b/src/HotChocolate/Core/test/Execution.Tests/DeferContextPoolingTests.cs new file mode 100644 index 00000000000..b1cf7f1e4c2 --- /dev/null +++ b/src/HotChocolate/Core/test/Execution.Tests/DeferContextPoolingTests.cs @@ -0,0 +1,181 @@ +using Microsoft.Extensions.DependencyInjection; + +namespace HotChocolate.Execution; + +public class DeferContextPoolingTests +{ + [Fact] + public async Task Execute_Should_ReturnFullData_When_AbortedDeferredResolverFinishesDuringNextRequest() + { + // arrange + var gates = new PoolingGates(); + var executor = await CreateExecutorAsync(gates); + using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(30)); + using var abort = new CancellationTokenSource(); + var unobserved = new UnobservedExceptionCollector(); + + await using var deferred = await executor.ExecuteAsync( + OperationRequestBuilder.New().SetDocument("{ a ... @defer { slow } }").Build(), + abort.Token); + var deferredStream = deferred.ExpectResponseStream(); + var deferredResults = deferredStream.ReadResultsAsync().GetAsyncEnumerator(timeout.Token); + await deferredResults.MoveNextAsync(); + await deferredResults.Current.DisposeAsync(); + await gates.SlowEntered.Task.WaitAsync(timeout.Token); + + // act + await abort.CancelAsync(); + await Task.Delay(TimeSpan.FromMilliseconds(300), timeout.Token); + + var next = executor.ExecuteAsync("{ a b }", timeout.Token); + await gates.BEntered.Task.WaitAsync(timeout.Token); + gates.SlowGate.SetResult(); + await gates.SlowReturned.Task.WaitAsync(timeout.Token); + await Task.Delay(TimeSpan.FromMilliseconds(300), timeout.Token); + var nextCompletedEarly = next.IsCompleted; + gates.BGate.SetResult(); + await using var nextResult = await next.WaitAsync(timeout.Token); + + var unobservedExceptions = unobserved.Collect(); + + // assert + Assert.False(nextCompletedEarly); + nextResult.ExpectOperationResult().ToJson().MatchInlineSnapshot( + """ + { + "data": { + "a": "a", + "b": "b" + } + } + """); + Assert.Empty(unobservedExceptions); + } + + [Fact] + public async Task Execute_Should_ReturnNoErrors_When_AbortedDeferredResolverFinishesBeforeNextRequest() + { + // arrange + var gates = new PoolingGates { SlowReturnsNull = true }; + var executor = await CreateExecutorAsync(gates); + using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(30)); + using var abort = new CancellationTokenSource(); + var unobserved = new UnobservedExceptionCollector(); + + await using var deferred = await executor.ExecuteAsync( + OperationRequestBuilder.New().SetDocument("{ a ... @defer { slow } }").Build(), + abort.Token); + var deferredStream = deferred.ExpectResponseStream(); + var deferredResults = deferredStream.ReadResultsAsync().GetAsyncEnumerator(timeout.Token); + await deferredResults.MoveNextAsync(); + await deferredResults.Current.DisposeAsync(); + await gates.SlowEntered.Task.WaitAsync(timeout.Token); + + // act + await abort.CancelAsync(); + await Task.Delay(TimeSpan.FromMilliseconds(300), timeout.Token); + + gates.SlowGate.SetResult(); + await gates.SlowReturned.Task.WaitAsync(timeout.Token); + await Task.Delay(TimeSpan.FromMilliseconds(300), timeout.Token); + + await using var nextResult = await executor.ExecuteAsync("{ a }", timeout.Token); + + var unobservedExceptions = unobserved.Collect(); + + // assert + nextResult.ExpectOperationResult().ToJson().MatchInlineSnapshot( + """ + { + "data": { + "a": "a" + } + } + """); + Assert.Empty(unobservedExceptions); + } + + private static ValueTask CreateExecutorAsync(PoolingGates gates) + => new ServiceCollection() + .AddSingleton(gates) + .AddGraphQL() + .AddQueryType() + .ModifyOptions(o => o.EnableDefer = true) + .BuildRequestExecutorAsync(); + + private sealed class UnobservedExceptionCollector + { + private readonly List _exceptions = []; + + public UnobservedExceptionCollector() + { + TaskScheduler.UnobservedTaskException += OnUnobservedTaskException; + } + + public IReadOnlyList Collect() + { + // unobserved task exceptions surface when the faulted task is finalized. + GC.Collect(); + GC.WaitForPendingFinalizers(); + GC.Collect(); + + TaskScheduler.UnobservedTaskException -= OnUnobservedTaskException; + + lock (_exceptions) + { + return [.. _exceptions]; + } + } + + private void OnUnobservedTaskException(object? sender, UnobservedTaskExceptionEventArgs args) + { + var exception = args.Exception.ToString(); + + if (exception.Contains("HotChocolate.Execution.Processing.Tasks.", StringComparison.Ordinal)) + { + lock (_exceptions) + { + _exceptions.Add(exception); + } + } + } + } + + public sealed class PoolingGates + { + public bool SlowReturnsNull { get; init; } + + public TaskCompletionSource SlowEntered { get; } = Create(); + + public TaskCompletionSource SlowGate { get; } = Create(); + + public TaskCompletionSource SlowReturned { get; } = Create(); + + public TaskCompletionSource BEntered { get; } = Create(); + + public TaskCompletionSource BGate { get; } = Create(); + + private static TaskCompletionSource Create() + => new(TaskCreationOptions.RunContinuationsAsynchronously); + } + + public sealed class PoolingQuery + { + public string GetA() => "a"; + + public async Task GetBAsync([Service] PoolingGates gates) + { + gates.BEntered.TrySetResult(); + await gates.BGate.Task; + return "b"; + } + + public async Task GetSlowAsync([Service] PoolingGates gates) + { + gates.SlowEntered.TrySetResult(); + await gates.SlowGate.Task; + gates.SlowReturned.TrySetResult(); + return gates.SlowReturnsNull ? null! : "slow-value"; + } + } +}