Add Operation Planner Guardrails

src/HotChocolate/Fusion-vnext/src/Fusion.Execution/Configuration/FusionGatewaySetup.cs

@@ -2,6 +2,7 @@
 using HotChocolate.Features;
 using HotChocolate.Fusion.Execution;
 using HotChocolate.Fusion.Execution.Clients;
+using HotChocolate.Fusion.Planning;
 using HotChocolate.Validation;
 using Microsoft.Extensions.DependencyInjection;
 
@@ -15,6 +16,8 @@ internal sealed class FusionGatewaySetup
 
     public List<Action<FusionRequestOptions>> RequestOptionsModifiers { get; } = [];
 
+    public List<Action<OperationPlannerOptions>> PlannerOptionsModifiers { get; } = [];
+
     public List<Action<FusionParserOptions>> ParserOptionsModifiers { get; } = [];
 
     public List<Action<IServiceProvider, IServiceCollection>> SchemaServiceModifiers { get; } = [];

src/HotChocolate/Fusion-vnext/src/Fusion.Execution/DependencyInjection/CoreFusionGatewayBuilderExtensions.Options.cs

@@ -1,10 +1,21 @@
 using HotChocolate.Fusion.Configuration;
 using HotChocolate.Fusion.Execution;
+using HotChocolate.Fusion.Planning;
 
 namespace Microsoft.Extensions.DependencyInjection;
 
+/// <summary>
+/// Extension methods for configuring Fusion gateway options on <see cref="IFusionGatewayBuilder"/>.
+/// </summary>
 public static partial class CoreFusionGatewayBuilderExtensions
 {
+    /// <summary>
+    /// Registers a callback to modify the core <see cref="FusionOptions"/>
+    /// (cache sizes, eviction timeout, error handling mode, etc.).
+    /// </summary>
+    /// <param name="builder">The gateway builder.</param>
+    /// <param name="configure">A delegate that configures the options.</param>
+    /// <returns>The <see cref="IFusionGatewayBuilder"/> for chaining.</returns>
     public static IFusionGatewayBuilder ModifyOptions(
         this IFusionGatewayBuilder builder,
         Action<FusionOptions> configure)
@@ -17,6 +28,13 @@ public static IFusionGatewayBuilder ModifyOptions(
             options => options.OptionsModifiers.Add(configure));
     }
 
+    /// <summary>
+    /// Registers a callback to modify the <see cref="FusionRequestOptions"/>
+    /// (persisted operations, exception details, etc.).
+    /// </summary>
+    /// <param name="builder">The gateway builder.</param>
+    /// <param name="configure">A delegate that configures the request options.</param>
+    /// <returns>The <see cref="IFusionGatewayBuilder"/> for chaining.</returns>
     public static IFusionGatewayBuilder ModifyRequestOptions(
         this IFusionGatewayBuilder builder,
         Action<FusionRequestOptions> configure)
@@ -28,4 +46,23 @@ public static IFusionGatewayBuilder ModifyRequestOptions(
             builder,
             options => options.RequestOptionsModifiers.Add(configure));
     }
+
+    /// <summary>
+    /// Registers a callback to modify the <see cref="OperationPlannerOptions"/>
+    /// (planning guardrails such as max planning time, max expanded nodes, etc.).
+    /// </summary>
+    /// <param name="builder">The gateway builder.</param>
+    /// <param name="configure">A delegate that configures the planner options.</param>
+    /// <returns>The <see cref="IFusionGatewayBuilder"/> for chaining.</returns>
+    public static IFusionGatewayBuilder ModifyPlannerOptions(
+        this IFusionGatewayBuilder builder,
+        Action<OperationPlannerOptions> configure)
+    {
+        ArgumentNullException.ThrowIfNull(builder);
+        ArgumentNullException.ThrowIfNull(configure);
+
+        return FusionSetupUtilities.Configure(
+            builder,
+            options => options.PlannerOptionsModifiers.Add(configure));
+    }
 }

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

@@ -176,6 +176,7 @@ private FusionRequestExecutor CreateRequestExecutor(
 
         var options = CreateOptions(setup);
         var requestOptions = CreateRequestOptions(setup);
+        var plannerOptions = CreatePlannerOptions(setup);
         var parserOptions = CreateParserOptions(setup);
         var clientConfigurations = CreateClientConfigurations(setup, configuration.Settings.Document);
         var features = CreateSchemaFeatures(
@@ -184,7 +185,7 @@ private FusionRequestExecutor CreateRequestExecutor(
             requestOptions,
             parserOptions,
             clientConfigurations);
-        var schemaServices = CreateSchemaServices(setup, options, requestOptions);
+        var schemaServices = CreateSchemaServices(setup, options, requestOptions, plannerOptions);
 
         var schema = CreateSchema(schemaName, configuration.Schema, schemaServices, features);
         var pipeline = CreatePipeline(setup, schema, schemaServices, requestOptions);
@@ -259,6 +260,20 @@ private static FusionRequestOptions CreateRequestOptions(FusionGatewaySetup setu
         return options;
     }
 
+    private static OperationPlannerOptions CreatePlannerOptions(FusionGatewaySetup setup)
+    {
+        var options = new OperationPlannerOptions();
+
+        foreach (var configure in setup.PlannerOptionsModifiers)
+        {
+            configure.Invoke(options);
+        }
+
+        options.MakeReadOnly();
+
+        return options;
+    }
+
     private static ParserOptions CreateParserOptions(FusionGatewaySetup setup)
     {
         var options = new FusionParserOptions();
@@ -373,12 +388,13 @@ private static Dictionary<string, ITypeResolverInterceptor> CreateTypeResolverIn
     private ServiceProvider CreateSchemaServices(
         FusionGatewaySetup setup,
         FusionOptions options,
-        FusionRequestOptions requestOptions)
+        FusionRequestOptions requestOptions,
+        OperationPlannerOptions plannerOptions)
     {
         var schemaServices = new ServiceCollection();
 
         AddCoreServices(schemaServices, options, requestOptions);
-        AddOperationPlanner(schemaServices);
+        AddOperationPlanner(schemaServices, plannerOptions);
         AddParserServices(schemaServices);
         AddDocumentValidator(setup, schemaServices);
         AddDiagnosticEvents(schemaServices);
@@ -428,7 +444,9 @@ private void AddCoreServices(
         services.AddTransient<CompositeTypeInterceptor>(static _ => new IntrospectionFieldInterceptor());
     }
 
-    private static void AddOperationPlanner(IServiceCollection services)
+    private static void AddOperationPlanner(
+        IServiceCollection services,
+        OperationPlannerOptions plannerOptions)
     {
         services.TryAddSingleton<ObjectPoolProvider>(
             static _ => new DefaultObjectPoolProvider());
@@ -450,10 +468,13 @@ private static void AddOperationPlanner(IServiceCollection services)
                 sp.GetRequiredService<FusionSchemaDefinition>(),
                 sp.GetRequiredService<ObjectPool<OrderedDictionary<string, List<FieldSelectionNode>>>>()));
 
+        services.AddSingleton(plannerOptions);
+
         services.AddSingleton(
             static sp => new OperationPlanner(
                 sp.GetRequiredService<FusionSchemaDefinition>(),
-                sp.GetRequiredService<OperationCompiler>()));
+                sp.GetRequiredService<OperationCompiler>(),
+                sp.GetRequiredService<OperationPlannerOptions>()));
     }
 
     private static void AddParserServices(IServiceCollection services)

src/HotChocolate/Fusion-vnext/src/Fusion.Execution/Execution/Pipeline/OperationPlanMiddleware.cs

@@ -65,7 +65,13 @@ private void PlanOperation(
             var operation = rewritten.GetOperation(context.Request.OperationName);
 
             // After optimizing the query structure we can begin the planning process.
-            var operationPlan = _planner.CreatePlan(operationId, operationHash, operationShortHash, operation);
+            var operationPlan =
+                _planner.CreatePlan(
+                    operationId,
+                    operationHash,
+                    operationShortHash,
+                    operation,
+                    context.RequestAborted);
             OnAfterPlanCompleted(operationDocumentInfo, operationPlan);
             context.SetOperationPlan(operationPlan);
         }

src/HotChocolate/Fusion-vnext/src/Fusion.Execution/Planning/OperationPlanner.cs

@@ -56,18 +56,24 @@ public OperationPlanner(
     /// <param name="hash">The hash of the operation document.</param>
     /// <param name="shortHash">The short hash of the operation document.</param>
     /// <param name="operationDefinition">The operation definition to create a plan for.</param>
+    /// <param name="cancellationToken">A token that can be used to cancel planning.</param>
     /// <returns>The operation plan.</returns>
     public OperationPlan CreatePlan(
         string id,
         string hash,
         string shortHash,
-        OperationDefinitionNode operationDefinition)
+        OperationDefinitionNode operationDefinition,
+        CancellationToken cancellationToken = default)
     {
         ArgumentException.ThrowIfNullOrEmpty(id);
         ArgumentException.ThrowIfNullOrEmpty(hash);
         ArgumentException.ThrowIfNullOrEmpty(shortHash);
         ArgumentNullException.ThrowIfNull(operationDefinition);
 
+        // We make sire that the cancellation token is observed right at the beginning of the method,
+        // so that if the caller passed in an already canceled token we don't do any unnecessary work.
+        cancellationToken.ThrowIfCancellationRequested();
+
         var eventSource = PlannerEventSource.Log;
         var eventSourceEnabled = eventSource.IsEnabled();
         var operationType = operationDefinition.Operation.ToString();
@@ -127,7 +133,7 @@ node with
                 }
 
                 // Now that we have seeded the possible plans we can start planning.
-                var plan = Plan(id, possiblePlans, eventSourceEnabled);
+                var plan = Plan(id, possiblePlans, eventSourceEnabled, cancellationToken);
 
                 if (!plan.HasValue)
                 {
@@ -183,27 +189,46 @@ node with
         }
     }
 
-    private PlanResult? Plan(string operationId, PlanQueue possiblePlans, bool emitPlannerEvents)
+    private PlanResult? Plan(
+        string operationId,
+        PlanQueue possiblePlans,
+        bool emitPlannerEvents,
+        CancellationToken cancellationToken)
     {
         var eventSource = PlannerEventSource.Log;
         var searchSpace = possiblePlans.Count;
         var expandedNodes = 0;
+        var maxPlanningTime = _options.MaxPlanningTime;
+        var maxExpandedNodes = _options.MaxExpandedNodes;
+        var maxQueueSize = _options.MaxQueueSize;
+        var maxGeneratedOptionsPerWorkItem = _options.MaxGeneratedOptionsPerWorkItem;
+        var planningStartedAt = maxPlanningTime.HasValue ? Stopwatch.GetTimestamp() : 0L;
 
         // TryBuildGreedyCompletePlan quickly builds one full plan by always choosing the currently
         // cheapest next option at each step.
         //
         // It gives the planner an initial best known complete cost, so the main search can skip branches
         // that are already worse. If it cannot finish a full plan, it returns null and the planner
         // continues without that early shortcut.
-        var bestCompletePlan = TryBuildGreedyCompletePlan(possiblePlans);
+        var bestCompletePlan = TryBuildGreedyCompletePlan(possiblePlans, cancellationToken);
         var bestCompletePlanCost = bestCompletePlan is null ? double.PositiveInfinity : bestCompletePlan.PathCost;
 
         while (possiblePlans.TryDequeue(out var current, out _))
         {
+            // we evaluate the cancellationToken at the beginning of each plan evaluation loop,
+            // so that we throw ones a request was canceled so that no unnecessary work is done.
+            cancellationToken.ThrowIfCancellationRequested();
+
             expandedNodes++;
             var possiblePlansCount = possiblePlans.Count;
             searchSpace = Math.Max(possiblePlansCount, searchSpace);
 
+            // before we get into another planning iteration, we check if we have
+            // exceeded any of the configured guardrails and throw if so.
+            EnsurePlanningTimeGuardrail();
+            EnsureExpandedNodesGuardrail(expandedNodes);
+            EnsureQueueSizeGuardrail(possiblePlansCount);
+
             var backlog = current.Backlog;
 
             if (emitPlannerEvents)
@@ -246,6 +271,7 @@ node with
             // the current possible plan. It's not guaranteed that this plan will work
             // out or that it is efficient.
             backlog = current.Backlog.Pop(out var workItem);
+            var queueCountBeforeExpansion = possiblePlans.Count;
 
             switch (workItem)
             {
@@ -286,6 +312,13 @@ node with
                     throw new NotSupportedException(
                         "The work item type is not supported.");
             }
+
+            // after we have expanded the current plan node into possible next steps,
+            // we check how many new plans we have created and if we have exceeded
+            // the guardrail for generated options per work item.
+            var queueCountAfterExpansion = possiblePlans.Count;
+            searchSpace = Math.Max(queueCountAfterExpansion, searchSpace);
+            EnsureGeneratedOptionsGuardrail(queueCountBeforeExpansion, queueCountAfterExpansion);
         }
 
         if (bestCompletePlan is null)
@@ -312,9 +345,99 @@ static string FormatWorkItemName(WorkItem workItem)
                 NodeLookupWorkItem => "NodeLookupBound",
                 _ => "Unknown"
             };
+
+        void EnsurePlanningTimeGuardrail()
+        {
+            if (maxPlanningTime is not { } planningTimeLimit)
+            {
+                return;
+            }
+
+            var elapsed = Stopwatch.GetElapsedTime(planningStartedAt);
+            if (elapsed < planningTimeLimit)
+            {
+                return;
+            }
+
+            ThrowGuardrailExceeded(
+                OperationPlannerGuardrailReason.MaxPlanningTimeExceeded,
+                ToGuardrailMilliseconds(planningTimeLimit),
+                ToGuardrailMilliseconds(elapsed));
+        }
+
+        void EnsureExpandedNodesGuardrail(int currentExpandedNodes)
+        {
+            if (maxExpandedNodes is not { } expandedNodesLimit
+                || currentExpandedNodes <= expandedNodesLimit)
+            {
+                return;
+            }
+
+            ThrowGuardrailExceeded(
+                OperationPlannerGuardrailReason.MaxExpandedNodesExceeded,
+                expandedNodesLimit,
+                currentExpandedNodes);
+        }
+
+        void EnsureQueueSizeGuardrail(int queueSize)
+        {
+            if (maxQueueSize is not { } queueSizeLimit
+                || queueSize <= queueSizeLimit)
+            {
+                return;
+            }
+
+            ThrowGuardrailExceeded(
+                OperationPlannerGuardrailReason.MaxQueueSizeExceeded,
+                queueSizeLimit,
+                queueSize);
+        }
+
+        void EnsureGeneratedOptionsGuardrail(int queueCountBeforeExpansion, int queueCountAfterExpansion)
+        {
+            if (maxGeneratedOptionsPerWorkItem is not { } generatedOptionsLimit)
+            {
+                return;
+            }
+
+            var generatedOptions = queueCountAfterExpansion - queueCountBeforeExpansion;
+            if (generatedOptions <= generatedOptionsLimit)
+            {
+                return;
+            }
+
+            ThrowGuardrailExceeded(
+                OperationPlannerGuardrailReason.MaxGeneratedOptionsPerWorkItemExceeded,
+                generatedOptionsLimit,
+                generatedOptions);
+        }
+
+        void ThrowGuardrailExceeded(
+            OperationPlannerGuardrailReason reason,
+            long limit,
+            long observed)
+        {
+            if (emitPlannerEvents)
+            {
+                eventSource.PlanGuardrailExceeded(
+                    operationId,
+                    reason.ToString(),
+                    limit,
+                    observed);
+            }
+
+            throw new OperationPlannerGuardrailException(
+                operationId,
+                reason,
+                limit,
+                observed);
+        }
+
+        static long ToGuardrailMilliseconds(TimeSpan value)
+            => checked((long)Math.Ceiling(value.TotalMilliseconds));
     }
 
-    private PlanNode? TryBuildGreedyCompletePlan(PlanQueue possiblePlans)
+    private PlanNode? TryBuildGreedyCompletePlan(PlanQueue possiblePlans, CancellationToken cancellationToken)
     {
         if (!possiblePlans.TryPeek(out var current, out _))
         {
@@ -325,6 +448,7 @@ static string FormatWorkItemName(WorkItem workItem)
 
         while (true)
         {
+            cancellationToken.ThrowIfCancellationRequested();
             var backlog = current.Backlog;
 
             if (backlog.IsEmpty)