From 1f85deb371b2294545f3234e40d4e014c8a9fb49 Mon Sep 17 00:00:00 2001 From: Michael Staib Date: Fri, 10 Jul 2026 09:52:23 +0200 Subject: [PATCH] [Fusion] Pool the request cancellation source across operation executions --- .../Execution/OperationPlanContext.Pooling.cs | 34 ++++++++++++ .../Execution/OperationPlanExecutor.cs | 55 ++++++++++++------- 2 files changed, 68 insertions(+), 21 deletions(-) diff --git a/src/HotChocolate/Fusion/src/Fusion.Execution/Execution/OperationPlanContext.Pooling.cs b/src/HotChocolate/Fusion/src/Fusion.Execution/Execution/OperationPlanContext.Pooling.cs index 6f1530036ed..5694882d87b 100644 --- a/src/HotChocolate/Fusion/src/Fusion.Execution/Execution/OperationPlanContext.Pooling.cs +++ b/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); } + /// + /// 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. + /// + internal (CancellationTokenSource Source, CancellationTokenRegistration Registration) RentEngineCancellation( + CancellationToken cancellationToken) + { + var registration = cancellationToken.UnsafeRegister( + static state => Unsafe.As(state)!.Cancel(), + _engineCancellationSource); + return (_engineCancellationSource, registration); + } + + /// + /// 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 before calling this. + /// + internal void ReturnEngineCancellation() + { + if (!_engineCancellationSource.TryReset()) + { + _engineCancellationSource.Dispose(); + _engineCancellationSource = new CancellationTokenSource(); + } + } + /// /// 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() diff --git a/src/HotChocolate/Fusion/src/Fusion.Execution/Execution/OperationPlanExecutor.cs b/src/HotChocolate/Fusion/src/Fusion.Execution/Execution/OperationPlanExecutor.cs index 7d6f4809219..08c8be313be 100644 --- a/src/HotChocolate/Fusion/src/Fusion.Execution/Execution/OperationPlanExecutor.cs +++ b/src/HotChocolate/Fusion/src/Fusion.Execution/Execution/OperationPlanExecutor.cs @@ -20,35 +20,48 @@ public static async Task 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().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 ExecuteWithDeferAsync(