[Fusion] Reduce forwarded variable allocations

src/HotChocolate/Fusion-vnext/src/Fusion.Execution/Execution/Clients/SourceSchemaRequestDispatcher.cs

@@ -1,6 +1,5 @@
 using System.Collections.Immutable;
 using System.Runtime.InteropServices;
-using HotChocolate.Fusion.Properties;
 using HotChocolate.Language;
 using static HotChocolate.Fusion.Execution.Clients.SourceSchemaClientCapabilities;
 using static HotChocolate.Fusion.Properties.FusionExecutionResources;
@@ -21,6 +20,11 @@ internal sealed class SourceSchemaRequestDispatcher
     : ISourceSchemaScheduler
     , ISourceSchemaDispatcher
 {
+    private const int NodeStateUnregistered = -1;
+    private const int NodeStatePending = 0;
+    private const int NodeStateSubmitted = 1;
+    private const int NodeStateSkipped = 2;
+
 #if NET9_0_OR_GREATER
     private readonly Lock _sync = new();
 #else
@@ -179,7 +183,7 @@ public void RegisterGroup(int groupId, IReadOnlyList<int> nodeIds)
                 }
 
                 _groupByNodeIdSlots[nodeId] = groupId;
-                _nodeStateSlots[nodeId] = 0;
+                _nodeStateSlots[nodeId] = NodeStatePending;
             }
         }
     }
@@ -415,7 +419,7 @@ private void RemoveGroup(GroupState group)
             if ((uint)nodeId < (uint)_groupByNodeIdSlots.Length)
             {
                 _groupByNodeIdSlots[nodeId] = -1;
-                _nodeStateSlots[nodeId] = -1;
+                _nodeStateSlots[nodeId] = NodeStateUnregistered;
             }
         }
     }
@@ -432,7 +436,7 @@ private void ClearNodeIdSlots()
             if ((uint)nodeId < (uint)_groupByNodeIdSlots.Length)
             {
                 _groupByNodeIdSlots[nodeId] = -1;
-                _nodeStateSlots[nodeId] = -1;
+                _nodeStateSlots[nodeId] = NodeStateUnregistered;
             }
         }
 
@@ -456,7 +460,7 @@ private void EnsureNodeIdSlotCapacity(int minCapacity)
         var groupByNodeIdSlots = new int[newCapacity];
         var nodeStateSlots = new int[newCapacity];
         Array.Fill(groupByNodeIdSlots, -1);
-        Array.Fill(nodeStateSlots, -1);
+        Array.Fill(nodeStateSlots, NodeStateUnregistered);
 
         if (_groupByNodeIdSlots.Length > 0)
         {
@@ -471,15 +475,15 @@ private void EnsureNodeIdSlotCapacity(int minCapacity)
     private sealed class GroupState(int id, int initialCapacity)
     {
         private readonly List<int> _nodeIds = new(initialCapacity);
-        private readonly Dictionary<int, PendingRequest> _pendingRequests = new(initialCapacity);
+        private readonly List<PendingRequest> _pendingRequests = new(initialCapacity);
         private int _remainingNodes;
         private bool _dispatchCreated;
 
         public int Id { get; } = id;
 
         public IEnumerable<int> NodeIds => _nodeIds;
 
-        public IEnumerable<PendingRequest> PendingRequests => _pendingRequests.Values;
+        public IEnumerable<PendingRequest> PendingRequests => _pendingRequests;
 
         public void RegisterNode(int nodeId)
         {
@@ -493,35 +497,40 @@ public bool TrySubmit(
             out PendingRequest? pendingRequest)
         {
             var nodeId = request.Node.Id;
+            var nodeState =
+                (uint)nodeId < (uint)nodeStateSlots.Length
+                    ? nodeStateSlots[nodeId]
+                    : NodeStateUnregistered;
 
-            if (_pendingRequests.ContainsKey(nodeId))
+            if (nodeState == NodeStateSubmitted)
             {
                 throw new InvalidOperationException(
                     string.Format(
                         SourceSchemaRequestDispatcher_DuplicateNodeSubmission,
                         nodeId));
             }
 
-            if ((uint)nodeId >= (uint)nodeStateSlots.Length || nodeStateSlots[nodeId] != 0)
+            if (nodeState != NodeStatePending)
             {
                 pendingRequest = null;
                 return false;
             }
 
-            nodeStateSlots[nodeId] = 1;
+            nodeStateSlots[nodeId] = NodeStateSubmitted;
             _remainingNodes--;
 
             pendingRequest = new PendingRequest(request);
-            _pendingRequests.Add(nodeId, pendingRequest);
+            _pendingRequests.Add(pendingRequest);
 
             return true;
         }
 
         public void Skip(int nodeId, int[] nodeStateSlots)
         {
-            if ((uint)nodeId < (uint)nodeStateSlots.Length && nodeStateSlots[nodeId] == 0)
+            if ((uint)nodeId < (uint)nodeStateSlots.Length
+                && nodeStateSlots[nodeId] == NodeStatePending)
             {
-                nodeStateSlots[nodeId] = 1;
+                nodeStateSlots[nodeId] = NodeStateSkipped;
                 _remainingNodes--;
             }
         }
@@ -542,7 +551,7 @@ public bool TryCreateDispatch(out ImmutableArray<PendingRequest> pendingRequests
                 return true;
             }
 
-            pendingRequests = [.. _pendingRequests.Values];
+            pendingRequests = [.. _pendingRequests];
             return true;
         }
     }

src/HotChocolate/Fusion-vnext/src/Fusion.Execution/Execution/Nodes/OperationBatchExecutionNode.cs

@@ -1,5 +1,6 @@
 using System.Buffers;
 using System.Collections.Immutable;
+using System.Runtime.InteropServices;
 using HotChocolate.Fusion.Execution.Clients;
 
 namespace HotChocolate.Fusion.Execution.Nodes;
@@ -127,6 +128,7 @@ protected override async ValueTask<ExecutionStatus> OnExecuteAsync(
         var index = 0;
         var bufferLength = 0;
         SourceSchemaResult[]? buffer = null;
+        SourceSchemaResult? singleResult = null;
         var hasSomeErrors = false;
 
         try
@@ -145,14 +147,34 @@ protected override async ValueTask<ExecutionStatus> OnExecuteAsync(
                 totalPathCount += variables[i].AdditionalPaths.Length;
             }
 
-            bufferLength = Math.Max(totalPathCount, 1);
-            buffer = ArrayPool<SourceSchemaResult>.Shared.Rent(bufferLength);
+            var initialBufferLength = Math.Max(totalPathCount, 2);
 
             await foreach (var result in response.ReadAsResultStreamAsync(cancellationToken))
             {
-                buffer[index++] = result;
+                // Store the first result without renting a buffer,
+                // since it might be the only one (e.g. a request-level error).
+                if (index == 0)
+                {
+                    singleResult = result;
+                    index = 1;
+                }
+                else
+                {
+                    // Once we see a second result, we know there are multiple,
+                    // so we rent a buffer and move the first result into it.
+                    if (buffer is null)
+                    {
+                        bufferLength = initialBufferLength;
+                        buffer = ArrayPool<SourceSchemaResult>.Shared.Rent(bufferLength);
+                        buffer[0] = singleResult!;
+                    }
 
-                if (result.HasErrors)
+                    buffer[index++] = result;
+                }
+
+                // Parsing errors here allows the result store to reuse the cached value
+                // and avoids a second document lookup per result.
+                if (result.Errors is not null)
                 {
                     hasSomeErrors = true;
                 }
@@ -182,14 +204,42 @@ protected override async ValueTask<ExecutionStatus> OnExecuteAsync(
                 buffer.AsSpan(0, index).Clear();
                 ArrayPool<SourceSchemaResult>.Shared.Return(buffer);
             }
+            else if (singleResult is not null)
+            {
+                singleResult.Dispose();
+            }
 
             AddErrors(context, exception, variables, _responseNames);
             return ExecutionStatus.Failed;
         }
 
         try
         {
-            context.AddPartialResults(_source, buffer.AsSpan(0, index), _responseNames);
+            if (buffer is not null)
+            {
+                context.AddPartialResults(
+                    _source,
+                    buffer.AsSpan(0, index),
+                    _responseNames,
+                    hasSomeErrors);
+            }
+            else if (singleResult is not null)
+            {
+                var firstResult = singleResult;
+                context.AddPartialResults(
+                    _source,
+                    MemoryMarshal.CreateReadOnlySpan(ref firstResult, 1),
+                    _responseNames,
+                    hasSomeErrors);
+            }
+            else
+            {
+                context.AddPartialResults(
+                    _source,
+                    [],
+                    _responseNames,
+                    hasSomeErrors);
+            }
         }
         catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
         {
@@ -206,8 +256,11 @@ protected override async ValueTask<ExecutionStatus> OnExecuteAsync(
         }
         finally
         {
-            buffer.AsSpan(0, index).Clear();
-            ArrayPool<SourceSchemaResult>.Shared.Return(buffer);
+            if (buffer is not null)
+            {
+                buffer.AsSpan(0, index).Clear();
+                ArrayPool<SourceSchemaResult>.Shared.Return(buffer);
+            }
         }
 
         return hasSomeErrors ? ExecutionStatus.PartialSuccess : ExecutionStatus.Success;

src/HotChocolate/Fusion-vnext/src/Fusion.Execution/Execution/Nodes/OperationExecutionNode.cs

@@ -134,6 +134,7 @@ protected override async ValueTask<ExecutionStatus> OnExecuteAsync(
         var index = 0;
         var bufferLength = 0;
         SourceSchemaResult[]? buffer = null;
+        SourceSchemaResult? singleResult = null;
         var hasSomeErrors = false;
 
         try
@@ -152,14 +153,32 @@ protected override async ValueTask<ExecutionStatus> OnExecuteAsync(
                 totalPathCount += variables[i].AdditionalPaths.Length;
             }
 
-            bufferLength = Math.Max(totalPathCount, 1);
-            buffer = ArrayPool<SourceSchemaResult>.Shared.Rent(bufferLength);
+            var initialBufferLength = Math.Max(totalPathCount, 2);
 
             await foreach (var result in response.ReadAsResultStreamAsync(cancellationToken))
             {
-                buffer[index++] = result;
+                // If there is only one response, we skip the buffer rental.
+                if (index == 0)
+                {
+                    singleResult = result;
+                    index = 1;
+                }
+                else
+                {
+                    // If we have more than one response, we rent a buffer and move the first result into it.
+                    if (buffer is null)
+                    {
+                        bufferLength = initialBufferLength;
+                        buffer = ArrayPool<SourceSchemaResult>.Shared.Rent(bufferLength);
+                        buffer[0] = singleResult!;
+                    }
 
-                if (result.HasErrors)
+                    buffer[index++] = result;
+                }
+
+                // Parsing errors here allows the result store to reuse the cached value
+                // and avoids a second document lookup per result.
+                if (result.Errors is not null)
                 {
                     hasSomeErrors = true;
                 }
@@ -189,14 +208,42 @@ protected override async ValueTask<ExecutionStatus> OnExecuteAsync(
                 buffer.AsSpan(0, index).Clear();
                 ArrayPool<SourceSchemaResult>.Shared.Return(buffer);
             }
+            else if (singleResult is not null)
+            {
+                singleResult.Dispose();
+            }
 
             AddErrors(context, exception, variables, _responseNames);
             return ExecutionStatus.Failed;
         }
 
         try
         {
-            context.AddPartialResults(_source, buffer.AsSpan(0, index), _responseNames);
+            if (buffer is not null)
+            {
+                context.AddPartialResults(
+                    _source,
+                    buffer.AsSpan(0, index),
+                    _responseNames,
+                    hasSomeErrors);
+            }
+            else if (singleResult is not null)
+            {
+                var firstResult = singleResult;
+                context.AddPartialResults(
+                    _source,
+                    MemoryMarshal.CreateReadOnlySpan(ref firstResult, 1),
+                    _responseNames,
+                    hasSomeErrors);
+            }
+            else
+            {
+                context.AddPartialResults(
+                    _source,
+                    [],
+                    _responseNames,
+                    hasSomeErrors);
+            }
         }
         catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
         {
@@ -213,8 +260,11 @@ protected override async ValueTask<ExecutionStatus> OnExecuteAsync(
         }
         finally
         {
-            buffer.AsSpan(0, index).Clear();
-            ArrayPool<SourceSchemaResult>.Shared.Return(buffer);
+            if (buffer is not null)
+            {
+                buffer.AsSpan(0, index).Clear();
+                ArrayPool<SourceSchemaResult>.Shared.Return(buffer);
+            }
         }
 
         return hasSomeErrors ? ExecutionStatus.PartialSuccess : ExecutionStatus.Success;

src/HotChocolate/Fusion-vnext/src/Fusion.Execution/Execution/OperationPlanContext.cs

@@ -260,9 +260,11 @@ internal ImmutableArray<VariableValues> CreateVariableValueSets(
     internal void AddPartialResults(
         SelectionPath sourcePath,
         ReadOnlySpan<SourceSchemaResult> results,
-        ReadOnlySpan<string> responseNames)
+        ReadOnlySpan<string> responseNames,
+        bool containsErrors = true)
     {
-        var canExecutionContinue = _resultStore.AddPartialResults(sourcePath, results, responseNames);
+        var canExecutionContinue =
+            _resultStore.AddPartialResults(sourcePath, results, responseNames, containsErrors);
 
         if (!canExecutionContinue)
         {
@@ -369,15 +371,15 @@ internal OperationResult Complete(bool reusable = false)
         return operationResult;
     }
 
-    private List<ObjectFieldNode> GetPathThroughVariables(
+    private IReadOnlyList<ObjectFieldNode> GetPathThroughVariables(
         ReadOnlySpan<string> forwardedVariables)
     {
         if (Variables.IsEmpty || forwardedVariables.Length == 0)
         {
-            return [];
+            return Array.Empty<ObjectFieldNode>();
         }
 
-        var variables = new List<ObjectFieldNode>();
+        var variables = new List<ObjectFieldNode>(forwardedVariables.Length);
 
         foreach (var variableName in forwardedVariables)
         {
@@ -400,7 +402,9 @@ private List<ObjectFieldNode> GetPathThroughVariables(
             }
         }
 
-        return variables;
+        return variables.Count == 0
+            ? Array.Empty<ObjectFieldNode>()
+            : variables;
     }
 
     public ISourceSchemaClient GetClient(string schemaName, OperationType operationType)