[Fusion] Pool the request cancellation source across operation executions (#10080)

Author: michaelstaib

src/HotChocolate/Fusion/src/Fusion.Execution/Execution/OperationPlanContext.Pooling.cs

@@ -1,5 +1,6 @@
 using System.Collections.Immutable;
 using System.Diagnostics;
+using System.Runtime.CompilerServices;
 using HotChocolate.Buffers;
 using HotChocolate.Execution;
 using HotChocolate.Fusion.Diagnostics;
@@ -10,6 +11,8 @@ namespace HotChocolate.Fusion.Execution;
 
 public sealed partial class OperationPlanContext
 {
+    private CancellationTokenSource _engineCancellationSource = new();
+
     internal OperationPlanContext(
         INodeIdParser nodeIdParser,
         IFusionExecutionDiagnosticEvents diagnosticEvents,
@@ -66,6 +69,36 @@ internal void Initialize(
         EnsureNodeArrayCapacity(operationPlan.MaxNodeId);
     }
 
+    /// <summary>
+    /// Arms the pooled engine cancellation source against the request token and returns it for
+    /// this operation. The engine source halts the execution engine without cancelling the request
+    /// pipeline. The returned registration links the request token into the engine source so that
+    /// client-abort and server-shutdown still propagate, and it must be disposed before the source
+    /// is returned for reuse.
+    /// </summary>
+    internal (CancellationTokenSource Source, CancellationTokenRegistration Registration) RentEngineCancellation(
+        CancellationToken cancellationToken)
+    {
+        var registration = cancellationToken.UnsafeRegister(
+            static state => Unsafe.As<CancellationTokenSource>(state)!.Cancel(),
+            _engineCancellationSource);
+        return (_engineCancellationSource, registration);
+    }
+
+    /// <summary>
+    /// Returns the engine cancellation source for reuse. If it was cancelled and cannot be reset,
+    /// it is disposed and replaced with a fresh source. The caller must dispose the registration
+    /// from <see cref="RentEngineCancellation"/> before calling this.
+    /// </summary>
+    internal void ReturnEngineCancellation()
+    {
+        if (!_engineCancellationSource.TryReset())
+        {
+            _engineCancellationSource.Dispose();
+            _engineCancellationSource = new CancellationTokenSource();
+        }
+    }
+
     /// <summary>
     /// Cleans the context for return to the pool.
     /// Releases per-request state while retaining reusable buffers.
@@ -117,6 +150,7 @@ internal void Destroy()
     {
         _resultStore.Dispose();
         _executionState.Destroy();
+        _engineCancellationSource.Dispose();
     }
 
     public async ValueTask DisposeAsync()

src/HotChocolate/Fusion/src/Fusion.Execution/Execution/OperationPlanExecutor.cs

@@ -20,35 +20,48 @@ public static async Task<IExecutionResult> ExecuteAsync(
         OperationPlan operationPlan,
         CancellationToken cancellationToken)
     {
-        // We create a new CancellationTokenSource that can be used to halt the execution engine,
-        // without also cancelling the entire request pipeline.
-        using var executionCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
-
         await using var context = requestContext.Schema.Services.GetRequiredService<OperationPlanContextPool>().Rent();
-        context.Initialize(requestContext, variables, operationPlan, executionCts);
 
-        context.Begin();
+        // We reuse a pooled CancellationTokenSource that can halt the execution engine without
+        // cancelling the entire request pipeline. The request token is linked in so that
+        // client-abort / server-shutdown still propagates.
+        var (executionCts, cancellationRegistration) = context.RentEngineCancellation(cancellationToken);
 
-        switch (operationPlan.Operation.Definition.Operation)
+        try
         {
-            case OperationType.Query:
-                await ExecuteQueryAsync(context, operationPlan, executionCts.Token);
-                break;
+            context.Initialize(requestContext, variables, operationPlan, executionCts);
 
-            case OperationType.Mutation:
-                await ExecuteMutationAsync(context, operationPlan, executionCts.Token);
-                break;
+            context.Begin();
 
-            default:
-                throw new InvalidOperationException("Only queries and mutations can be executed.");
-        }
+            switch (operationPlan.Operation.Definition.Operation)
+            {
+                case OperationType.Query:
+                    await ExecuteQueryAsync(context, operationPlan, executionCts.Token);
+                    break;
 
-        // If the original CancellationToken of the request was cancelled,
-        // the Execution nodes and the PlanExecutor should have been gracefully cancelled,
-        // so we throw here to properly cancel the request execution.
-        cancellationToken.ThrowIfCancellationRequested();
+                case OperationType.Mutation:
+                    await ExecuteMutationAsync(context, operationPlan, executionCts.Token);
+                    break;
+
+                default:
+                    throw new InvalidOperationException("Only queries and mutations can be executed.");
+            }
 
-        return context.Complete();
+            // If the original CancellationToken of the request was cancelled,
+            // the Execution nodes and the PlanExecutor should have been gracefully cancelled,
+            // so we throw here to properly cancel the request execution.
+            cancellationToken.ThrowIfCancellationRequested();
+
+            return context.Complete();
+        }
+        finally
+        {
+            // Dispose the parent-token registration BEFORE returning the source for reuse so a
+            // stale registration can never fire into a reset source. DisposeAsync waits for an
+            // in-flight cancel callback to finish.
+            await cancellationRegistration.DisposeAsync();
+            context.ReturnEngineCancellation();
+        }
     }
 
     public static async Task<IExecutionResult> ExecuteWithDeferAsync(