Skip to content

Commit

Permalink
Add cancellation checks to potentially heavy loops in query planning
Browse files Browse the repository at this point in the history
  • Loading branch information
@sachindshinde
sachindshinde committed Feb 28, 2025
1 parent d740eea commit d051bb8
Show file tree
Hide file tree
Showing 7 changed files with 94 additions and 12 deletions.
23 changes: 17 additions & 6 deletions apollo-federation/src/operation/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1541,9 +1541,12 @@ impl SelectionSet {
Ok(())
}

pub(crate) fn expand_all_fragments(&self) -> Result<SelectionSet, FederationError> {
pub(crate) fn expand_all_fragments(
&self,
check_cancellation: &impl Fn() -> Result<(), SingleFederationError>,
) -> Result<SelectionSet, FederationError> {
let mut expanded_selections = vec![];
SelectionSet::expand_selection_set(&mut expanded_selections, self)?;
SelectionSet::expand_selection_set(&mut expanded_selections, self, check_cancellation)?;

let mut expanded = SelectionSet {
schema: self.schema.clone(),
Expand All @@ -1557,12 +1560,14 @@ impl SelectionSet {
fn expand_selection_set(
destination: &mut Vec<Selection>,
selection_set: &SelectionSet,
check_cancellation: &impl Fn() -> Result<(), SingleFederationError>,
) -> Result<(), FederationError> {
for value in selection_set.selections.values() {
check_cancellation()?;
match value {
Selection::Field(field_selection) => {
let selections = match &field_selection.selection_set {
Some(s) => Some(s.expand_all_fragments()?),
Some(s) => Some(s.expand_all_fragments(check_cancellation)?),
None => None,
};
destination.push(Selection::from_field(
Expand All @@ -1580,12 +1585,14 @@ impl SelectionSet {
SelectionSet::expand_selection_set(
destination,
&spread_selection.selection_set,
check_cancellation,
)?;
} else {
// convert to inline fragment
let expanded = InlineFragmentSelection::from_fragment_spread_selection(
selection_set.type_position.clone(), // the parent type of this inline selection
spread_selection,
check_cancellation,
)?;
destination.push(Selection::InlineFragment(Arc::new(expanded)));
}
Expand All @@ -1594,7 +1601,9 @@ impl SelectionSet {
destination.push(
InlineFragmentSelection::new(
inline_selection.inline_fragment.clone(),
inline_selection.selection_set.expand_all_fragments()?,
inline_selection
.selection_set
.expand_all_fragments(check_cancellation)?,
)
.into(),
);
Expand Down Expand Up @@ -2718,6 +2727,7 @@ impl InlineFragmentSelection {
pub(crate) fn from_fragment_spread_selection(
parent_type_position: CompositeTypeDefinitionPosition,
fragment_spread_selection: &Arc<FragmentSpreadSelection>,
check_cancellation: &impl Fn() -> Result<(), SingleFederationError>,
) -> Result<InlineFragmentSelection, FederationError> {
let schema = fragment_spread_selection.spread.schema.schema();
for directive in fragment_spread_selection.spread.directives.iter() {
Expand Down Expand Up @@ -2755,7 +2765,7 @@ impl InlineFragmentSelection {
},
fragment_spread_selection
.selection_set
.expand_all_fragments()?,
.expand_all_fragments(check_cancellation)?,
))
}

Expand Down Expand Up @@ -3751,10 +3761,11 @@ pub(crate) fn normalize_operation(
named_fragments: NamedFragments,
schema: &ValidFederationSchema,
interface_types_with_interface_objects: &IndexSet<InterfaceTypeDefinitionPosition>,
check_cancellation: &impl Fn() -> Result<(), SingleFederationError>,
) -> Result<Operation, FederationError> {
let mut normalized_selection_set =
SelectionSet::from_selection_set(&operation.selection_set, &named_fragments, schema)?;
normalized_selection_set = normalized_selection_set.expand_all_fragments()?;
normalized_selection_set = normalized_selection_set.expand_all_fragments(check_cancellation)?;
// We clear up the fragments since we've expanded all.
// Also note that expanding fragment usually generate unnecessary fragments/inefficient
// selections, so it basically always make sense to flatten afterwards. Besides, fragment
Expand Down
26 changes: 25 additions & 1 deletion apollo-federation/src/operation/tests/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ pub(super) fn parse_and_expand(
.expect("must have anonymous operation");
let fragments = NamedFragments::new(&doc.fragments, schema);

normalize_operation(operation, fragments, schema, &Default::default())
normalize_operation(operation, fragments, schema, &Default::default(), &None)
}

#[test]
Expand Down Expand Up @@ -100,6 +100,7 @@ type Foo {
NamedFragments::new(&executable_document.fragments, &schema),
&schema,
&IndexSet::default(),
&None,
)
.unwrap();
normalized_operation.named_fragments = Default::default();
Expand Down Expand Up @@ -154,6 +155,7 @@ type Foo {
NamedFragments::new(&executable_document.fragments, &schema),
&schema,
&IndexSet::default(),
&None,
)
.unwrap();
normalized_operation.named_fragments = Default::default();
Expand Down Expand Up @@ -196,6 +198,7 @@ type Query {
NamedFragments::new(&executable_document.fragments, &schema),
&schema,
&IndexSet::default(),
&None,
)
.unwrap();

Expand Down Expand Up @@ -231,6 +234,7 @@ type T {
NamedFragments::new(&executable_document.fragments, &schema),
&schema,
&IndexSet::default(),
&None,
)
.unwrap();
let expected = r#"query Test {
Expand Down Expand Up @@ -274,6 +278,7 @@ type T {
NamedFragments::new(&executable_document.fragments, &schema),
&schema,
&IndexSet::default(),
&None,
)
.unwrap();
let expected = r#"query Test($skipIf: Boolean!) {
Expand Down Expand Up @@ -320,6 +325,7 @@ type T {
NamedFragments::new(&executable_document.fragments, &schema),
&schema,
&IndexSet::default(),
&None,
)
.unwrap();
let expected = r#"query Test($skipIf: Boolean!) {
Expand Down Expand Up @@ -364,6 +370,7 @@ type T {
NamedFragments::new(&executable_document.fragments, &schema),
&schema,
&IndexSet::default(),
&None,
)
.unwrap();
let expected = r#"query Test($skipIf: Boolean!) {
Expand Down Expand Up @@ -410,6 +417,7 @@ type T {
NamedFragments::new(&executable_document.fragments, &schema),
&schema,
&IndexSet::default(),
&None,
)
.unwrap();
let expected = r#"query Test($skip1: Boolean!, $skip2: Boolean!) {
Expand Down Expand Up @@ -461,6 +469,7 @@ type T {
NamedFragments::new(&executable_document.fragments, &schema),
&schema,
&IndexSet::default(),
&None,
)
.unwrap();
let expected = r#"query Test {
Expand Down Expand Up @@ -527,6 +536,7 @@ type V {
NamedFragments::new(&executable_document.fragments, &schema),
&schema,
&IndexSet::default(),
&None,
)
.unwrap();
let expected = r#"query Test {
Expand Down Expand Up @@ -586,6 +596,7 @@ type T {
NamedFragments::new(&executable_document.fragments, &schema),
&schema,
&IndexSet::default(),
&None,
)
.unwrap();
let expected = r#"query Test {
Expand Down Expand Up @@ -632,6 +643,7 @@ type T {
NamedFragments::new(&executable_document.fragments, &schema),
&schema,
&IndexSet::default(),
&None,
)
.unwrap();
let expected = r#"query Test($skipIf: Boolean!) {
Expand Down Expand Up @@ -682,6 +694,7 @@ type T {
NamedFragments::new(&executable_document.fragments, &schema),
&schema,
&IndexSet::default(),
&None,
)
.unwrap();
let expected = r#"query Test($skipIf: Boolean!) {
Expand Down Expand Up @@ -730,6 +743,7 @@ type T {
NamedFragments::new(&executable_document.fragments, &schema),
&schema,
&IndexSet::default(),
&None,
)
.unwrap();
let expected = r#"query Test($skipIf: Boolean!) {
Expand Down Expand Up @@ -778,6 +792,7 @@ type T {
NamedFragments::new(&executable_document.fragments, &schema),
&schema,
&IndexSet::default(),
&None,
)
.unwrap();
let expected = r#"query Test($skip1: Boolean!, $skip2: Boolean!) {
Expand Down Expand Up @@ -830,6 +845,7 @@ type T {
NamedFragments::new(&executable_document.fragments, &schema),
&schema,
&IndexSet::default(),
&None,
)
.unwrap();
let expected = r#"query Test {
Expand Down Expand Up @@ -898,6 +914,7 @@ type V {
NamedFragments::new(&executable_document.fragments, &schema),
&schema,
&IndexSet::default(),
&None,
)
.unwrap();
let expected = r#"query Test {
Expand Down Expand Up @@ -944,6 +961,7 @@ type Foo {
NamedFragments::new(&executable_document.fragments, &schema),
&schema,
&IndexSet::default(),
&None,
)
.unwrap();
let expected = r#"query TestQuery {
Expand Down Expand Up @@ -983,6 +1001,7 @@ type Foo {
NamedFragments::new(&executable_document.fragments, &schema),
&schema,
&IndexSet::default(),
&None,
)
.unwrap();
let expected = r#"query TestQuery {
Expand Down Expand Up @@ -1033,6 +1052,7 @@ scalar FieldSet
NamedFragments::new(&executable_document.fragments, &schema),
&schema,
&interface_objects,
&None,
)
.unwrap();
let expected = r#"query TestQuery {
Expand Down Expand Up @@ -1172,6 +1192,7 @@ mod make_selection_tests {
Default::default(),
&schema,
&Default::default(),
&None,
)
.unwrap();

Expand Down Expand Up @@ -1271,6 +1292,7 @@ mod lazy_map_tests {
Default::default(),
&schema,
&Default::default(),
&None,
)
.unwrap();

Expand Down Expand Up @@ -1329,6 +1351,7 @@ mod lazy_map_tests {
Default::default(),
&schema,
&Default::default(),
&None,
)
.unwrap();

Expand Down Expand Up @@ -1534,6 +1557,7 @@ fn test_expand_all_fragments1() {
NamedFragments::new(&executable_document.fragments, &schema),
&schema,
&IndexSet::default(),
&None,
)
.unwrap();
normalized_operation.named_fragments = Default::default();
Expand Down
Loading

0 comments on commit d051bb8

Please sign in to comment.