@@ -44,6 +44,7 @@ use crate::query_plan::query_planning_traversal::BestQueryPlanInfo;
4444use crate :: query_plan:: query_planning_traversal:: QueryPlanningParameters ;
4545use crate :: query_plan:: query_planning_traversal:: QueryPlanningTraversal ;
4646use crate :: query_plan:: query_planning_traversal:: convert_type_from_subgraph;
47+ use crate :: query_plan:: query_planning_traversal:: non_local_selections_estimation;
4748use crate :: schema:: ValidFederationSchema ;
4849use crate :: schema:: position:: AbstractTypeDefinitionPosition ;
4950use crate :: schema:: position:: CompositeTypeDefinitionPosition ;
@@ -172,7 +173,7 @@ pub struct QueryPlanningStatistics {
172173 pub evaluated_plan_paths : Cell < usize > ,
173174}
174175
175- #[ derive( Default , Clone ) ]
176+ #[ derive( Clone ) ]
176177pub struct QueryPlanOptions < ' a > {
177178 /// A set of labels which will be used _during query planning_ to
178179 /// enable/disable edges with a matching label in their override condition.
@@ -183,6 +184,19 @@ pub struct QueryPlanOptions<'a> {
183184 pub override_conditions : Vec < String > ,
184185
185186 pub check_for_cooperative_cancellation : Option < & ' a dyn Fn ( ) -> ControlFlow < ( ) > > ,
187+ /// Impose a limit on the number of non-local selections, which can be a
188+ /// performance hazard. On by default.
189+ pub non_local_selections_limit_enabled : bool ,
190+ }
191+
192+ impl Default for QueryPlanOptions < ' _ > {
193+ fn default ( ) -> Self {
194+ Self {
195+ override_conditions : Vec :: new ( ) ,
196+ check_for_cooperative_cancellation : None ,
197+ non_local_selections_limit_enabled : true ,
198+ }
199+ }
186200}
187201
188202impl std:: fmt:: Debug for QueryPlanOptions < ' _ > {
@@ -197,6 +211,10 @@ impl std::fmt::Debug for QueryPlanOptions<'_> {
197211 & "None"
198212 } ,
199213 )
214+ . field (
215+ "non_local_selections_limit_enabled" ,
216+ & self . non_local_selections_limit_enabled ,
217+ )
200218 . finish ( )
201219 }
202220}
@@ -224,7 +242,7 @@ pub struct QueryPlanner {
224242 /// subgraphs.
225243 // PORT_NOTE: Named `inconsistentAbstractTypesRuntimes` in the JS codebase, which was slightly
226244 // confusing.
227- abstract_types_with_inconsistent_runtime_types : IndexSet < AbstractTypeDefinitionPosition > ,
245+ abstract_types_with_inconsistent_runtime_types : IndexSet < Name > ,
228246}
229247
230248impl QueryPlanner {
@@ -317,6 +335,7 @@ impl QueryPlanner {
317335 . get_types ( )
318336 . filter_map ( |position| AbstractTypeDefinitionPosition :: try_from ( position) . ok ( ) )
319337 . filter ( |position| is_inconsistent ( position. clone ( ) ) )
338+ . map ( |position| position. type_name ( ) . clone ( ) )
320339 . collect :: < IndexSet < _ > > ( ) ;
321340
322341 Ok ( Self {
@@ -443,10 +462,23 @@ impl QueryPlanner {
443462 fetch_id_generator : Arc :: new ( FetchIdGenerator :: new ( ) ) ,
444463 } ;
445464
465+ let mut non_local_selection_state = options
466+ . non_local_selections_limit_enabled
467+ . then ( non_local_selections_estimation:: State :: default) ;
446468 let root_node = if !defer_conditions. is_empty ( ) {
447- compute_plan_for_defer_conditionals ( & mut parameters, & mut processor, defer_conditions)
469+ compute_plan_for_defer_conditionals (
470+ & mut parameters,
471+ & mut processor,
472+ defer_conditions,
473+ & mut non_local_selection_state,
474+ )
448475 } else {
449- compute_plan_internal ( & mut parameters, & mut processor, has_defers)
476+ compute_plan_internal (
477+ & mut parameters,
478+ & mut processor,
479+ has_defers,
480+ & mut non_local_selection_state,
481+ )
450482 } ?;
451483
452484 let root_node = match root_node {
@@ -521,6 +553,7 @@ impl QueryPlanner {
521553fn compute_root_serial_dependency_graph (
522554 parameters : & QueryPlanningParameters ,
523555 has_defers : bool ,
556+ non_local_selection_state : & mut Option < non_local_selections_estimation:: State > ,
524557) -> Result < Vec < FetchDependencyGraph > , FederationError > {
525558 let QueryPlanningParameters {
526559 supergraph_schema,
@@ -547,14 +580,24 @@ fn compute_root_serial_dependency_graph(
547580 mut fetch_dependency_graph,
548581 path_tree : mut prev_path,
549582 ..
550- } = compute_root_parallel_best_plan ( parameters, selection_set, has_defers) ?;
583+ } = compute_root_parallel_best_plan (
584+ parameters,
585+ selection_set,
586+ has_defers,
587+ non_local_selection_state,
588+ ) ?;
551589 let mut prev_subgraph = only_root_subgraph ( & fetch_dependency_graph) ?;
552590 for selection_set in split_roots {
553591 let BestQueryPlanInfo {
554592 fetch_dependency_graph : new_dep_graph,
555593 path_tree : new_path,
556594 ..
557- } = compute_root_parallel_best_plan ( parameters, selection_set, has_defers) ?;
595+ } = compute_root_parallel_best_plan (
596+ parameters,
597+ selection_set,
598+ has_defers,
599+ non_local_selection_state,
600+ ) ?;
558601 let new_subgraph = only_root_subgraph ( & new_dep_graph) ?;
559602 if new_subgraph == prev_subgraph {
560603 // The new operation (think 'mutation' operation) is on the same subgraph than the previous one, so we can concat them in a single fetch
@@ -677,10 +720,16 @@ pub(crate) fn compute_root_fetch_groups(
677720fn compute_root_parallel_dependency_graph (
678721 parameters : & QueryPlanningParameters ,
679722 has_defers : bool ,
723+ non_local_selection_state : & mut Option < non_local_selections_estimation:: State > ,
680724) -> Result < FetchDependencyGraph , FederationError > {
681725 trace ! ( "Starting process to construct a parallel fetch dependency graph" ) ;
682726 let selection_set = parameters. operation . selection_set . clone ( ) ;
683- let best_plan = compute_root_parallel_best_plan ( parameters, selection_set, has_defers) ?;
727+ let best_plan = compute_root_parallel_best_plan (
728+ parameters,
729+ selection_set,
730+ has_defers,
731+ non_local_selection_state,
732+ ) ?;
684733 snapshot ! (
685734 "FetchDependencyGraph" ,
686735 best_plan. fetch_dependency_graph. to_dot( ) ,
@@ -693,13 +742,15 @@ fn compute_root_parallel_best_plan(
693742 parameters : & QueryPlanningParameters ,
694743 selection : SelectionSet ,
695744 has_defers : bool ,
745+ non_local_selection_state : & mut Option < non_local_selections_estimation:: State > ,
696746) -> Result < BestQueryPlanInfo , FederationError > {
697747 let planning_traversal = QueryPlanningTraversal :: new (
698748 parameters,
699749 selection,
700750 has_defers,
701751 parameters. operation . root_kind ,
702752 FetchDependencyGraphToCostProcessor ,
753+ non_local_selection_state. as_mut ( ) ,
703754 ) ?;
704755
705756 // Getting no plan means the query is essentially unsatisfiable (it's a valid query, but we can prove it will never return a result),
@@ -713,11 +764,16 @@ fn compute_plan_internal(
713764 parameters : & mut QueryPlanningParameters ,
714765 processor : & mut FetchDependencyGraphToQueryPlanProcessor ,
715766 has_defers : bool ,
767+ non_local_selection_state : & mut Option < non_local_selections_estimation:: State > ,
716768) -> Result < Option < PlanNode > , FederationError > {
717769 let root_kind = parameters. operation . root_kind ;
718770
719771 let ( main, deferred, primary_selection) = if root_kind == SchemaRootDefinitionKind :: Mutation {
720- let dependency_graphs = compute_root_serial_dependency_graph ( parameters, has_defers) ?;
772+ let dependency_graphs = compute_root_serial_dependency_graph (
773+ parameters,
774+ has_defers,
775+ non_local_selection_state,
776+ ) ?;
721777 let mut main = None ;
722778 let mut deferred = vec ! [ ] ;
723779 let mut primary_selection = None :: < SelectionSet > ;
@@ -741,7 +797,11 @@ fn compute_plan_internal(
741797 }
742798 ( main, deferred, primary_selection)
743799 } else {
744- let mut dependency_graph = compute_root_parallel_dependency_graph ( parameters, has_defers) ?;
800+ let mut dependency_graph = compute_root_parallel_dependency_graph (
801+ parameters,
802+ has_defers,
803+ non_local_selection_state,
804+ ) ?;
745805
746806 let ( main, deferred) = dependency_graph. process ( & mut * processor, root_kind) ?;
747807 snapshot ! (
@@ -769,13 +829,14 @@ fn compute_plan_for_defer_conditionals(
769829 parameters : & mut QueryPlanningParameters ,
770830 processor : & mut FetchDependencyGraphToQueryPlanProcessor ,
771831 defer_conditions : IndexMap < Name , IndexSet < String > > ,
832+ non_local_selection_state : & mut Option < non_local_selections_estimation:: State > ,
772833) -> Result < Option < PlanNode > , FederationError > {
773834 generate_condition_nodes (
774835 parameters. operation . clone ( ) ,
775836 defer_conditions. iter ( ) ,
776837 & mut |op| {
777838 parameters. operation = op;
778- compute_plan_internal ( parameters, processor, true )
839+ compute_plan_internal ( parameters, processor, true , non_local_selection_state )
779840 } ,
780841 )
781842}
0 commit comments