Fix pooled operation contexts being reused while deferred resolvers still run

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<DisconnectQuery>()
+                .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<string> GetSlowAsync([Service] DisconnectGates gates)
+        {
+            gates.SlowEntered.TrySetResult();
+            await gates.SlowGate.Task;
+            gates.SlowReturned.TrySetResult();
+            return null!;
+        }
+    }
+}

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();

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;

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);
 
+    /// <summary>
+    /// Gets whether tasks created from this context, or from a defer context derived from it,
+    /// have not yet been returned to their pool.
+    /// </summary>
+    public bool HasOutstandingTasks => Volatile.Read(ref _outstandingTasks) > 0;
+
+    /// <summary>
+    /// Registers a task that was created from this context.
+    /// </summary>
+    public void TaskCreated()
+    {
+        Interlocked.Increment(ref _outstandingTasks);
+
+        if (_taskRoot is { } root)
+        {
+            Interlocked.Increment(ref root._outstandingTasks);
+        }
+    }
+
+    /// <summary>
+    /// Releases a task that was registered with <see cref="TaskCreated"/>.
+    /// </summary>
+    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);
+        }
+    }
+
+    /// <summary>
+    /// Waits until all tasks created from this context have been returned to their pool.
+    /// Only one caller may wait at a time.
+    /// </summary>
+    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<Schema>(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,23 +214,24 @@ public void InitializeDeferContext(
         _currentBranchTracker = context._currentBranchTracker;
         _currentWorkScheduler = context._currentWorkScheduler;
         _currentDeferExecutionCoordinator = context._currentDeferExecutionCoordinator;
+        _taskRoot = context._taskRoot ?? context;
         _propagateNullValues = context._propagateNullValues;
         _branchId = executionBranchId;
         _isInitialized = true;
 
         IncludeConditionFlags = context.IncludeConditionFlags;
         DeferConditionFlags = context.DeferConditionFlags;
 
-        Result.Data = new ResultDocument(
+        _result.Data = new ResultDocument(
             context._memory!,
             context.Operation,
             selectionSet,
             selectionPath,
             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();
         }
     }
 

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<BatchResolverTask> 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;
     }
 
+    /// <summary>
+    /// Tells the operation context that this task and its resolver tasks no longer use it.
+    /// Calling it again has no effect.
+    /// </summary>
+    private void ReleaseOperationContext()
+    {
+        foreach (var task in _resolverTasks)
+        {
+            task.ReleaseOperationContext();
+        }
+
+        if (_isContextReleased || _operationContext is null)
+        {
+            return;
+        }
+
+        _isContextReleased = true;
+        _operationContext.TaskReturned();
+    }
+
     /// <summary>
     /// Represents a single entry in the batch — one parent object and its result location.
     /// </summary>