From bebe4b70d3fe0382f3ab3ea02da7b0efbe28f176 Mon Sep 17 00:00:00 2001 From: Michael Staib Date: Tue, 25 Aug 2026 14:44:30 +0200 Subject: [PATCH 1/2] Fix deferred variable batch response streams --- .../Pipeline/OperationExecutionMiddleware.cs | 28 ++++- .../Execution/Processing/QueryExecutor.cs | 38 +++--- .../Core/test/Execution.Tests/DeferTests.cs | 118 ++++++++++++++++++ 3 files changed, 162 insertions(+), 22 deletions(-) diff --git a/src/HotChocolate/Core/src/Types/Execution/Pipeline/OperationExecutionMiddleware.cs b/src/HotChocolate/Core/src/Types/Execution/Pipeline/OperationExecutionMiddleware.cs index 9d073c48e16..f0b4254f906 100644 --- a/src/HotChocolate/Core/src/Types/Execution/Pipeline/OperationExecutionMiddleware.cs +++ b/src/HotChocolate/Core/src/Types/Execution/Pipeline/OperationExecutionMiddleware.cs @@ -129,6 +129,12 @@ await _queryExecutor.ExecuteBatchAsync( resultBuffer, variableSets.Length); + // Incremental batch items return response streams that are consumed after the request + // pipeline has completed. We transfer ownership of each streamed item's operation + // context into its stream so the context is only cleaned and returned to the pool once + // the stream has been fully consumed, mirroring the single operation path. + TransferStreamedContextOwnership(operationContextBuffer, resultBuffer, variableSets.Length); + context.Result = new OperationResultBatch([.. resultBuffer.AsSpan(0, variableSets.Length)]); } catch (OperationCanceledException) @@ -172,6 +178,24 @@ static void Initialize( operationContexts[variableIndex] = operationContextOwner; } + static void TransferStreamedContextOwnership( + OperationContextOwner[] operationContextBuffer, + IExecutionResult[] resultBuffer, + int length) + { + for (var i = 0; i < length; i++) + { + if (resultBuffer[i].IsStreamResult() && operationContextBuffer[i] is { } contextOwner) + { + resultBuffer[i].RegisterForCleanup(contextOwner); + + // Ownership now belongs to the stream, so we drop it from the buffer to keep + // ReleaseResources from disposing (and pooling) the context a second time. + operationContextBuffer[i] = null!; + } + } + } + static void AbandonContexts(ref OperationContextOwner[]? operationContextBuffer, int length) { if (operationContextBuffer is not null) @@ -201,7 +225,9 @@ static void ReleaseResources( foreach (var contextOwner in contextOwners) { - contextOwner.Dispose(); + // Owners transferred to a response stream are cleared from the buffer and are + // disposed by that stream once it has been consumed, so we skip them here. + contextOwner?.Dispose(); } contextOwners.Clear(); diff --git a/src/HotChocolate/Core/src/Types/Execution/Processing/QueryExecutor.cs b/src/HotChocolate/Core/src/Types/Execution/Processing/QueryExecutor.cs index 4bcee2ca818..d3090e19740 100644 --- a/src/HotChocolate/Core/src/Types/Execution/Processing/QueryExecutor.cs +++ b/src/HotChocolate/Core/src/Types/Execution/Processing/QueryExecutor.cs @@ -159,23 +159,25 @@ private async Task ExecuteBatchIncrementalAsync( for (var i = 0; i < length; ++i) { - if (i == 0) - { - var branchId = parentContext.ExecutionBranchId; - await scheduler.WaitForCompletionAsync(branchId).ConfigureAwait(false); - parentContext.DeferExecutionCoordinator.EnqueueResult(parentContext.BuildResult()); - results[i] = new ResponseStream(CreateStreamAndComplete, ExecutionResultKind.DeferredResult); - } - else - { - var context = operationContexts[i].OperationContext; - var branchId = context.ExecutionBranchId; - await scheduler.WaitForCompletionAsync(branchId).ConfigureAwait(false); - context.DeferExecutionCoordinator.EnqueueResult(context.BuildResult()); - results[i] = new ResponseStream(CreateStream, ExecutionResultKind.DeferredResult); - } + var context = operationContexts[i].OperationContext; + var branchId = context.ExecutionBranchId; + await scheduler.WaitForCompletionAsync(branchId).ConfigureAwait(false); + context.DeferExecutionCoordinator.EnqueueResult(context.BuildResult()); + + // The first item's stream drives completion of the shared execution: once it has + // delivered every payload it awaits the scheduler and seals the shared arena. Every + // item reads from its own coordinator and honors its own cancellation token, so a + // payload is never enqueued into one coordinator and read from another. + results[i] = i == 0 + ? new ResponseStream(CreateStreamAndComplete, ExecutionResultKind.DeferredResult) + : CreateItemStream(context); } + static ResponseStream CreateItemStream(OperationContext context) + => new( + () => context.DeferExecutionCoordinator.ReadResultsAsync(context.RequestAborted), + ExecutionResultKind.DeferredResult); + async IAsyncEnumerable CreateStreamAndComplete() { var requestAborted = parentContext.RequestAborted; @@ -192,12 +194,6 @@ async IAsyncEnumerable CreateStreamAndComplete() await execution.ConfigureAwait(false); memory.Seal(); } - - IAsyncEnumerable CreateStream() - { - var requestAborted = parentContext.RequestAborted; - return parentContext.DeferExecutionCoordinator.ReadResultsAsync(requestAborted); - } } private static void FillSchedulerWithWork( diff --git a/src/HotChocolate/Core/test/Execution.Tests/DeferTests.cs b/src/HotChocolate/Core/test/Execution.Tests/DeferTests.cs index 29513fc1bc5..1b67699e53b 100644 --- a/src/HotChocolate/Core/test/Execution.Tests/DeferTests.cs +++ b/src/HotChocolate/Core/test/Execution.Tests/DeferTests.cs @@ -1,3 +1,4 @@ +using System.Text.Json; using HotChocolate.AspNetCore; using Microsoft.AspNetCore.Http; using Microsoft.Extensions.DependencyInjection; @@ -6,6 +7,123 @@ namespace HotChocolate.Execution; public class DeferTests { + [Fact] + public async Task VariableBatch_Defer_Should_Deliver_Each_Items_Payloads_When_Executed() + { + // arrange + // the timeout turns a lost-coordinator hang into a test failure instead of a block. + var executor = await DeferAndStreamTestSchema.CreateAsync(); + using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(15)); + + var request = OperationRequestBuilder + .New() + .SetDocument( + """ + query ($id: ID!) { + person(id: $id) { + id + ... @defer { + name + } + } + } + """) + .SetVariableValues( + new List> + { + new Dictionary { ["id"] = "UGVyc29uOjE=" }, + new Dictionary { ["id"] = "UGVyc29uOjI=" } + }) + .Build(); + + // act + await using var result = await executor.ExecuteAsync(request, cts.Token); + + var summaries = new List(); + foreach (var item in result.ExpectOperationResultBatch().Results) + { + summaries.Add(await SummarizeStreamAsync(item.ExpectResponseStream(), cts.Token)); + } + + // assert + summaries.MatchInlineSnapshots( + [ + "person.id=UGVyc29uOjE=; name=Pascal; hasNext=True,False", + "person.id=UGVyc29uOjI=; name=Rafi; hasNext=True,False" + ]); + } + + // Reads a single item's response stream to completion and projects it to a delivery-shape + // summary that pins the initial data, the deferred data, and the per-payload hasNext sequence + // without depending on the non-deterministic branch identifiers in the incremental envelope. + private static async Task SummarizeStreamAsync( + ResponseStream stream, + CancellationToken cancellationToken) + { + string? personId = null; + string? name = null; + var hasNext = new List(); + + await foreach (var payload in stream.ReadResultsAsync().WithCancellation(cancellationToken)) + { + await using var payloadCleanup = payload; + using var document = JsonDocument.Parse(payload.ToJson()); + var root = document.RootElement; + + hasNext.Add(root.TryGetProperty("hasNext", out var hasNextValue) && hasNextValue.GetBoolean()); + + if (root.TryGetProperty("data", out var data) + && data.ValueKind is JsonValueKind.Object + && data.TryGetProperty("person", out var person) + && person.TryGetProperty("id", out var id)) + { + personId = id.GetString(); + } + + // The deferred name must arrive inside an incremental entry's data (the deferred + // payload), never in the initial data. Reading only incremental[].data keeps the + // non-deterministic branch identifiers carried elsewhere in the envelope out of + // the summary. + if (root.TryGetProperty("incremental", out var incremental) + && incremental.ValueKind is JsonValueKind.Array) + { + foreach (var entry in incremental.EnumerateArray()) + { + if (entry.TryGetProperty("data", out var incrementalData) + && TryFindName(incrementalData, out var deferredName)) + { + name = deferredName; + } + } + } + } + + return $"person.id={personId}; name={name}; hasNext={string.Join(",", hasNext)}"; + } + + private static bool TryFindName(JsonElement element, out string? name) + { + if (element.ValueKind is JsonValueKind.Object) + { + foreach (var property in element.EnumerateObject()) + { + if (property.Name is "name" && property.Value.ValueKind is JsonValueKind.String) + { + name = property.Value.GetString(); + return true; + } + + if (TryFindName(property.Value, out name)) + { + return true; + } + } + } + + name = null; + return false; + } + [Fact] public async Task InlineFragment_Defer() { From b4160f8dbca1c09d2bc7297c024e2227e4f639f2 Mon Sep 17 00:00:00 2001 From: Michael Staib Date: Tue, 25 Aug 2026 15:54:42 +0200 Subject: [PATCH 2/2] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- src/HotChocolate/Core/test/Execution.Tests/DeferTests.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/HotChocolate/Core/test/Execution.Tests/DeferTests.cs b/src/HotChocolate/Core/test/Execution.Tests/DeferTests.cs index 1b67699e53b..83efa01e0d8 100644 --- a/src/HotChocolate/Core/test/Execution.Tests/DeferTests.cs +++ b/src/HotChocolate/Core/test/Execution.Tests/DeferTests.cs @@ -8,7 +8,7 @@ namespace HotChocolate.Execution; public class DeferTests { [Fact] - public async Task VariableBatch_Defer_Should_Deliver_Each_Items_Payloads_When_Executed() + public async Task VariableBatch_Defer_Should_Deliver_Payloads_Per_Item_When_Executed() { // arrange // the timeout turns a lost-coordinator hang into a test failure instead of a block.