@@ -42,6 +42,7 @@ use crate::query_plan::query_planning_traversal::BestQueryPlanInfo;
4242use crate :: query_plan:: query_planning_traversal:: QueryPlanningParameters ;
4343use crate :: query_plan:: query_planning_traversal:: QueryPlanningTraversal ;
4444use crate :: query_plan:: query_planning_traversal:: convert_type_from_subgraph;
45+ use crate :: query_plan:: query_planning_traversal:: non_local_selections_estimation;
4546use crate :: schema:: ValidFederationSchema ;
4647use crate :: schema:: position:: AbstractTypeDefinitionPosition ;
4748use crate :: schema:: position:: CompositeTypeDefinitionPosition ;
@@ -170,7 +171,7 @@ pub struct QueryPlanningStatistics {
170171 pub evaluated_plan_paths : Cell < usize > ,
171172}
172173
173- #[ derive( Debug , Default , Clone ) ]
174+ #[ derive( Debug , Clone ) ]
174175pub struct QueryPlanOptions {
175176 /// A set of labels which will be used _during query planning_ to
176177 /// enable/disable edges with a matching label in their override condition.
@@ -179,6 +180,18 @@ pub struct QueryPlanOptions {
179180 /// progressive @override feature.
180181 // PORT_NOTE: In JS implementation this was a Map
181182 pub override_conditions : Vec < String > ,
183+ /// Impose a limit on the number of non-local selections, which can be a
184+ /// performance hazard. On by default.
185+ pub non_local_selections_limit_enabled : bool ,
186+ }
187+
188+ impl Default for QueryPlanOptions {
189+ fn default ( ) -> Self {
190+ Self {
191+ override_conditions : Vec :: new ( ) ,
192+ non_local_selections_limit_enabled : true ,
193+ }
194+ }
182195}
183196
184197#[ derive( Debug , Default , Clone ) ]
@@ -204,7 +217,7 @@ pub struct QueryPlanner {
204217 /// subgraphs.
205218 // PORT_NOTE: Named `inconsistentAbstractTypesRuntimes` in the JS codebase, which was slightly
206219 // confusing.
207- abstract_types_with_inconsistent_runtime_types : IndexSet < AbstractTypeDefinitionPosition > ,
220+ abstract_types_with_inconsistent_runtime_types : IndexSet < Name > ,
208221}
209222
210223impl QueryPlanner {
@@ -297,6 +310,7 @@ impl QueryPlanner {
297310 . get_types ( )
298311 . filter_map ( |position| AbstractTypeDefinitionPosition :: try_from ( position) . ok ( ) )
299312 . filter ( |position| is_inconsistent ( position. clone ( ) ) )
313+ . map ( |position| position. type_name ( ) . clone ( ) )
300314 . collect :: < IndexSet < _ > > ( ) ;
301315
302316 Ok ( Self {
@@ -417,10 +431,23 @@ impl QueryPlanner {
417431 fetch_id_generator : Arc :: new ( FetchIdGenerator :: new ( ) ) ,
418432 } ;
419433
434+ let mut non_local_selection_state = options
435+ . non_local_selections_limit_enabled
436+ . then ( non_local_selections_estimation:: State :: default) ;
420437 let root_node = if !defer_conditions. is_empty ( ) {
421- compute_plan_for_defer_conditionals ( & mut parameters, & mut processor, defer_conditions)
438+ compute_plan_for_defer_conditionals (
439+ & mut parameters,
440+ & mut processor,
441+ defer_conditions,
442+ & mut non_local_selection_state,
443+ )
422444 } else {
423- compute_plan_internal ( & mut parameters, & mut processor, has_defers)
445+ compute_plan_internal (
446+ & mut parameters,
447+ & mut processor,
448+ has_defers,
449+ & mut non_local_selection_state,
450+ )
424451 } ?;
425452
426453 let root_node = match root_node {
@@ -495,6 +522,7 @@ impl QueryPlanner {
495522fn compute_root_serial_dependency_graph (
496523 parameters : & QueryPlanningParameters ,
497524 has_defers : bool ,
525+ non_local_selection_state : & mut Option < non_local_selections_estimation:: State > ,
498526) -> Result < Vec < FetchDependencyGraph > , FederationError > {
499527 let QueryPlanningParameters {
500528 supergraph_schema,
@@ -521,14 +549,24 @@ fn compute_root_serial_dependency_graph(
521549 mut fetch_dependency_graph,
522550 path_tree : mut prev_path,
523551 ..
524- } = compute_root_parallel_best_plan ( parameters, selection_set, has_defers) ?;
552+ } = compute_root_parallel_best_plan (
553+ parameters,
554+ selection_set,
555+ has_defers,
556+ non_local_selection_state,
557+ ) ?;
525558 let mut prev_subgraph = only_root_subgraph ( & fetch_dependency_graph) ?;
526559 for selection_set in split_roots {
527560 let BestQueryPlanInfo {
528561 fetch_dependency_graph : new_dep_graph,
529562 path_tree : new_path,
530563 ..
531- } = compute_root_parallel_best_plan ( parameters, selection_set, has_defers) ?;
564+ } = compute_root_parallel_best_plan (
565+ parameters,
566+ selection_set,
567+ has_defers,
568+ non_local_selection_state,
569+ ) ?;
532570 let new_subgraph = only_root_subgraph ( & new_dep_graph) ?;
533571 if new_subgraph == prev_subgraph {
534572 // The new operation (think 'mutation' operation) is on the same subgraph than the previous one, so we can concat them in a single fetch
@@ -648,10 +686,16 @@ pub(crate) fn compute_root_fetch_groups(
648686fn compute_root_parallel_dependency_graph (
649687 parameters : & QueryPlanningParameters ,
650688 has_defers : bool ,
689+ non_local_selection_state : & mut Option < non_local_selections_estimation:: State > ,
651690) -> Result < FetchDependencyGraph , FederationError > {
652691 trace ! ( "Starting process to construct a parallel fetch dependency graph" ) ;
653692 let selection_set = parameters. operation . selection_set . clone ( ) ;
654- let best_plan = compute_root_parallel_best_plan ( parameters, selection_set, has_defers) ?;
693+ let best_plan = compute_root_parallel_best_plan (
694+ parameters,
695+ selection_set,
696+ has_defers,
697+ non_local_selection_state,
698+ ) ?;
655699 snapshot ! (
656700 "FetchDependencyGraph" ,
657701 best_plan. fetch_dependency_graph. to_dot( ) ,
@@ -664,13 +708,15 @@ fn compute_root_parallel_best_plan(
664708 parameters : & QueryPlanningParameters ,
665709 selection : SelectionSet ,
666710 has_defers : bool ,
711+ non_local_selection_state : & mut Option < non_local_selections_estimation:: State > ,
667712) -> Result < BestQueryPlanInfo , FederationError > {
668713 let planning_traversal = QueryPlanningTraversal :: new (
669714 parameters,
670715 selection,
671716 has_defers,
672717 parameters. operation . root_kind ,
673718 FetchDependencyGraphToCostProcessor ,
719+ non_local_selection_state. as_mut ( ) ,
674720 ) ?;
675721
676722 // Getting no plan means the query is essentially unsatisfiable (it's a valid query, but we can prove it will never return a result),
@@ -684,11 +730,16 @@ fn compute_plan_internal(
684730 parameters : & mut QueryPlanningParameters ,
685731 processor : & mut FetchDependencyGraphToQueryPlanProcessor ,
686732 has_defers : bool ,
733+ non_local_selection_state : & mut Option < non_local_selections_estimation:: State > ,
687734) -> Result < Option < PlanNode > , FederationError > {
688735 let root_kind = parameters. operation . root_kind ;
689736
690737 let ( main, deferred, primary_selection) = if root_kind == SchemaRootDefinitionKind :: Mutation {
691- let dependency_graphs = compute_root_serial_dependency_graph ( parameters, has_defers) ?;
738+ let dependency_graphs = compute_root_serial_dependency_graph (
739+ parameters,
740+ has_defers,
741+ non_local_selection_state,
742+ ) ?;
692743 let mut main = None ;
693744 let mut deferred = vec ! [ ] ;
694745 let mut primary_selection = None :: < SelectionSet > ;
@@ -712,7 +763,11 @@ fn compute_plan_internal(
712763 }
713764 ( main, deferred, primary_selection)
714765 } else {
715- let mut dependency_graph = compute_root_parallel_dependency_graph ( parameters, has_defers) ?;
766+ let mut dependency_graph = compute_root_parallel_dependency_graph (
767+ parameters,
768+ has_defers,
769+ non_local_selection_state,
770+ ) ?;
716771
717772 let ( main, deferred) = dependency_graph. process ( & mut * processor, root_kind) ?;
718773 snapshot ! (
@@ -740,13 +795,14 @@ fn compute_plan_for_defer_conditionals(
740795 parameters : & mut QueryPlanningParameters ,
741796 processor : & mut FetchDependencyGraphToQueryPlanProcessor ,
742797 defer_conditions : IndexMap < Name , IndexSet < String > > ,
798+ non_local_selection_state : & mut Option < non_local_selections_estimation:: State > ,
743799) -> Result < Option < PlanNode > , FederationError > {
744800 generate_condition_nodes (
745801 parameters. operation . clone ( ) ,
746802 defer_conditions. iter ( ) ,
747803 & mut |op| {
748804 parameters. operation = op;
749- compute_plan_internal ( parameters, processor, true )
805+ compute_plan_internal ( parameters, processor, true , non_local_selection_state )
750806 } ,
751807 )
752808}
0 commit comments