Reworked Fusion Query Plan (#5342)

src/CookieCrumble/src/CookieCrumble/CookieCrumble.csproj

@@ -20,4 +20,8 @@
     <ProjectReference Include="..\..\..\HotChocolate\Core\src\Core\HotChocolate.Core.csproj" />
   </ItemGroup>
 
+  <ItemGroup Condition="'$(TargetFramework)' == 'net7.0' OR '$(TargetFramework)' == 'net6.0'">
+    <ProjectReference Include="..\..\..\HotChocolate\Fusion\src\Core\HotChocolate.Fusion.csproj" />
+  </ItemGroup>
+
 </Project>

src/CookieCrumble/src/CookieCrumble/Formatters/QueryPlanSnapshotValueFormatter.cs

@@ -0,0 +1,14 @@
+#if NET6_0_OR_GREATER
+using System.Buffers;
+using HotChocolate.Fusion.Planning;
+
+namespace CookieCrumble.Formatters;
+
+internal sealed class QueryPlanSnapshotValueFormatter : SnapshotValueFormatter<QueryPlan>
+{
+    protected override void Format(IBufferWriter<byte> snapshot, QueryPlan value)
+    {
+        value.Format(snapshot);
+    }
+}
+#endif

src/CookieCrumble/src/CookieCrumble/Snapshot.cs

@@ -27,6 +27,9 @@ public sealed class Snapshot
             new SchemaSnapshotValueFormatter(),
             new ExceptionSnapshotValueFormatter(),
             new SchemaErrorSnapshotValueFormatter(),
+#if NET6_0_OR_GREATER
+            new QueryPlanSnapshotValueFormatter(),
+#endif
         });
     private static readonly JsonSnapshotValueFormatter _defaultFormatter = new();
 

src/HotChocolate/Core/src/Execution/Processing/Operation.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Collections;
 using System.Collections.Generic;
 using HotChocolate.Language;
 using HotChocolate.Types;
@@ -115,5 +116,19 @@ public long CreateIncludeFlags(IVariableValueCollection variables)
         return context;
     }
 
+    public IEnumerator<ISelectionSet> GetEnumerator()
+    {
+        foreach (var selectionVariant in _selectionVariants)
+        {
+            foreach (var objectType in selectionVariant.GetPossibleTypes())
+            {
+                yield return selectionVariant.GetSelectionSet(objectType);
+            }
+        }
+    }
+
+    IEnumerator IEnumerable.GetEnumerator()
+        => GetEnumerator();
+
     public override string ToString() => OperationPrinter.Print(this);
 }

src/HotChocolate/Core/src/Execution/Processing/OperationCompiler.cs

@@ -29,6 +29,7 @@ public sealed partial class OperationCompiler
     private IncludeCondition[] _includeConditions = Array.Empty<IncludeCondition>();
     private CompilerContext? _deferContext;
     private int _nextSelectionId;
+    private int _nextSelectionSetRefId;
     private int _nextSelectionSetId;
     private int _nextFragmentId;
 
@@ -86,7 +87,7 @@ public OperationCompiler(InputParser parser)
 
             // collect root fields
             var rootPath = SelectionPath.Root;
-            var id = GetOrCreateSelectionSetId(operationDefinition.SelectionSet, rootPath);
+            var id = GetOrCreateSelectionSetRefId(operationDefinition.SelectionSet, rootPath);
             var variants = GetOrCreateSelectionVariants(id);
             SelectionSetInfo[] infos = { new(operationDefinition.SelectionSet, 0) };
 
@@ -120,7 +121,8 @@ public OperationCompiler(InputParser parser)
         finally
         {
             _nextSelectionId = 0;
-            _nextSelectionSetId = 0;
+            _nextSelectionSetRefId = 0;
+            _nextSelectionId = 0;
             _nextFragmentId = 0;
 
             _backlog.Clear();
@@ -263,7 +265,7 @@ private void CompleteSelectionSet(CompilerContext context)
                 }
 
                 var selectionPath = context.Path.Append(selection.ResponseName);
-                selectionSetId = GetOrCreateSelectionSetId(selection.SelectionSet, selectionPath);
+                selectionSetId = GetOrCreateSelectionSetRefId(selection.SelectionSet, selectionPath);
                 var selectionVariants = GetOrCreateSelectionVariants(selectionSetId);
                 var possibleTypes = context.Schema.GetPossibleTypes(fieldType);
 
@@ -318,6 +320,7 @@ private void CompleteSelectionSet(CompilerContext context)
         }
 
         context.SelectionVariants.AddSelectionSet(
+            _nextSelectionSetId++,
             context.Type,
             selections,
             fragments,
@@ -501,7 +504,7 @@ private void CollectFields(CompilerContext context)
                     ifConditionFlags = GetSelectionIncludeCondition(ifCondition, includeCondition);
                 }
 
-                var id = GetOrCreateSelectionSetId(selectionSet, context.Path);
+                var id = GetOrCreateSelectionSetRefId(selectionSet, context.Path);
                 var variants = GetOrCreateSelectionVariants(id);
                 var infos = new SelectionSetInfo[] { new(selectionSet, includeCondition) };
 
@@ -586,13 +589,13 @@ private static bool DoesTypeApply(IType typeCondition, IObjectType current)
 
     private int GetNextFragmentId() => _nextFragmentId++;
 
-    private int GetOrCreateSelectionSetId(SelectionSetNode selectionSet, SelectionPath path)
+    private int GetOrCreateSelectionSetRefId(SelectionSetNode selectionSet, SelectionPath path)
     {
         var selectionSetRef = new SelectionSetRef(selectionSet, path);
 
         if (!_selectionSetIdLookup.TryGetValue(selectionSetRef, out var selectionSetId))
         {
-            selectionSetId = _nextSelectionSetId++;
+            selectionSetId = _nextSelectionSetRefId++;
             _selectionSetIdLookup.Add(selectionSetRef, selectionSetId);
         }
 

src/HotChocolate/Core/src/Execution/Processing/Result/ResultBuilder.cs

@@ -59,6 +59,14 @@ public void RegisterForCleanup(Func<ValueTask> action)
         }
     }
 
+    public void RegisterForCleanup<T>(T state, Func<T, ValueTask> action)
+    {
+        lock (_syncExtensions)
+        {
+            _cleanupTasks.Add(() => action(state));
+        }
+    }
+
     public void SetPath(Path? path)
         => _path = path;
 

src/HotChocolate/Core/src/Execution/Processing/SelectionSet.cs

@@ -19,24 +19,9 @@ internal sealed class SelectionSet : ISelectionSet
     /// <summary>
     /// Initializes a new instance of <see cref="SelectionSet"/>.
     /// </summary>
-    /// <param name="selections">
-    /// A list of executable field selections.
-    /// </param>
-    /// <param name="isConditional">
-    /// Defines if this list needs post processing for skip and include.
+    /// <param name="id">
+    /// The selection set unique id.
     /// </param>
-    public SelectionSet(
-        Selection[] selections,
-        bool isConditional)
-    {
-        _selections = selections;
-        _fragments = _empty;
-        _flags = isConditional ? Flags.Conditional : Flags.None;
-    }
-
-    /// <summary>
-    /// Initializes a new instance of <see cref="SelectionSet"/>.
-    /// </summary>
     /// <param name="selections">
     /// A list of executable field selections.
     /// </param>
@@ -48,15 +33,20 @@ internal sealed class SelectionSet : ISelectionSet
     /// Defines if this list needs post processing for skip and include.
     /// </param>
     public SelectionSet(
+        int id,
         Selection[] selections,
         Fragment[]? fragments,
         bool isConditional)
     {
+        Id = id;
         _selections = selections;
         _fragments = fragments ?? _empty;
         _flags = isConditional ? Flags.Conditional : Flags.None;
     }
 
+    /// <inheritdoc />
+    public int Id { get; }
+
     /// <inheritdoc />
     public bool IsConditional => (_flags & Flags.Conditional) == Flags.Conditional;
 

src/HotChocolate/Core/src/Execution/Processing/SelectionVariants.cs

@@ -78,6 +78,7 @@ internal bool ContainsSelectionSet(IObjectType typeContext)
     }
 
     internal void AddSelectionSet(
+        int id,
         ObjectType typeContext,
         Selection[] selections,
         Fragment[]? fragments,
@@ -88,7 +89,7 @@ internal bool ContainsSelectionSet(IObjectType typeContext)
             throw new NotSupportedException(SelectionVariants_ReadOnly);
         }
 
-        var selectionSet = new SelectionSet(selections, fragments, isConditional);
+        var selectionSet = new SelectionSet(id, selections, fragments, isConditional);
 
         if (_map is not null)
         {

src/HotChocolate/Core/src/Execution/Serialization/JsonQueryResultFormatter.cs

@@ -362,6 +362,7 @@ private static void WritePathValue(Utf8JsonWriter writer, Path path)
         for (var i = 0; i < objectResult.Capacity; i++)
         {
             var field = Unsafe.Add(ref searchSpace, i);
+
             if (field.IsInitialized)
             {
                 writer.WritePropertyName(field.Name);
@@ -497,6 +498,10 @@ private static void WritePathValue(Utf8JsonWriter writer, Path path)
             case JsonElement element:
                 WriteJsonElement(writer, element);
                 break;
+
+            case RawJsonValue rawJsonValue:
+                writer.WriteRawValue(rawJsonValue.Value.Span, true);
+                break;
 #endif
             case Dictionary<string, object?> dict:
                 WriteDictionary(writer, dict);

src/HotChocolate/Core/src/Execution/Serialization/RawJsonValue.cs

@@ -0,0 +1,13 @@
+using System;
+
+namespace HotChocolate.Execution.Serialization;
+
+internal readonly struct RawJsonValue
+{
+    public RawJsonValue(ReadOnlyMemory<byte> value)
+    {
+        Value = value;
+    }
+
+    public ReadOnlyMemory<byte> Value { get; }
+}

src/HotChocolate/Core/src/Types/Execution/Processing/IOperation.cs

@@ -10,7 +10,7 @@ namespace HotChocolate.Execution.Processing;
 /// <summary>
 /// Represents a compiled GraphQL operation.
 /// </summary>
-public interface IOperation : IHasReadOnlyContextData
+public interface IOperation : IHasReadOnlyContextData, IEnumerable<ISelectionSet>
 {
     /// <summary>
     /// Gets the internal unique identifier for this operation.

src/HotChocolate/Core/src/Types/Execution/Processing/ISelectionSet.cs

@@ -11,6 +11,11 @@ namespace HotChocolate.Execution.Processing;
 /// </summary>
 public interface ISelectionSet
 {
+    /// <summary>
+    /// Gets an operation unique selection-set identifier of this selection.
+    /// </summary>
+    int Id { get; }
+
     /// <summary>
     /// Defines if this list needs post processing for skip and include.
     /// </summary>

src/HotChocolate/Fusion/src/Core/Execution/ExecutionState.cs

@@ -0,0 +1,114 @@
+using System.Diagnostics.CodeAnalysis;
+using HotChocolate.Execution.Processing;
+
+namespace HotChocolate.Fusion.Execution;
+
+internal sealed class ExecutionState : IExecutionState
+{
+    private readonly Dictionary<ISelectionSet, List<WorkItem>> _map = new();
+    private readonly HashSet<ISelectionSet> _immutable = new();
+
+    public bool ContainsState(ISelectionSet selectionSet)
+    {
+        var taken = false;
+        Monitor.Enter(_map, ref taken);
+
+        try
+        {
+            return _map.ContainsKey(selectionSet);
+        }
+        finally
+        {
+            if (taken)
+            {
+                Monitor.Exit(_map);
+            }
+        }
+    }
+
+    public bool TryGetState(
+        ISelectionSet selectionSet,
+        [NotNullWhen(true)] out IReadOnlyList<WorkItem>? values)
+    {
+        var taken = false;
+        Monitor.Enter(_map, ref taken);
+
+        try
+        {
+            // We mark a value immutable on first read.
+            //
+            // After we accessed the first time the state of a selection set its no longer allowed
+            // to mutate it.
+            //
+            // The query plan should actually be ordered in a way that there are no mutations after
+            // the state is being read from nodes.
+            _immutable.Add(selectionSet);
+
+            if (_map.TryGetValue(selectionSet, out var local))
+            {
+                values = local;
+                return true;
+            }
+            else
+            {
+                values = null;
+                return false;
+            }
+        }
+        finally
+        {
+            if (taken)
+            {
+                Monitor.Exit(_map);
+            }
+        }
+    }
+
+    public void RegisterState(WorkItem value)
+    {
+        var taken = false;
+        List<WorkItem>? values;
+        Monitor.Enter(_map, ref taken);
+
+        try
+        {
+            if (_immutable.Contains(value.SelectionSet))
+            {
+                throw new InvalidOperationException(
+                    $"The state for the selection set `{value.SelectionSet.Id}` is immutable.");
+            }
+
+            if (!_map.TryGetValue(value.SelectionSet, out values))
+            {
+                var temp = new List<WorkItem> { value };
+                _map.Add(value.SelectionSet, temp);
+            }
+        }
+        finally
+        {
+            if (taken)
+            {
+                Monitor.Exit(_map);
+            }
+        }
+
+        if (values is not null)
+        {
+            taken = false;
+            Monitor.Enter(values, ref taken);
+
+            try
+            {
+                values.Add(value);
+            }
+            finally
+            {
+                if (taken)
+                {
+                    Monitor.Exit(values);
+                }
+            }
+        }
+    }
+}
+