From 4c7114d6af86934d46a2562430c5f9a77c170a30 Mon Sep 17 00:00:00 2001 From: Iryna Shestak Date: Mon, 24 Feb 2025 12:13:32 +0100 Subject: [PATCH] add dhat heap profiling tests for Supergaph and QueryPlanner functions (#6814) Building on the work in https://github.com/apollographql/router/pull/6743, this commit adds heap profiling tests for a few of the apollo_federation::Supergraph and apollo_federaion::query_plan::query_planner functions: * apollo_federation::Supergraph::new() * supergraph.to_api_schema() * supergraph.extract_subgraphs_to_supergraph() * apollo_federation::query_plan::query_planner::QueryPlanner::new() * planner.build_query_plan() This is not meant to be a comprehensive heap profiling testing framework, but merely an indication to us that something has dramatically changed in the most used query planning functionality. A smoke signal, if you will. Heap usage testing with the `dhat` crate is a bit finicky, and you can't really add them to an existing test file or combine multiple tests in a file ([more on this in the crate docs](https://docs.rs/dhat/latest/dhat/#heap-usage-testing)). Tests have to be added as individual tests, or we have to create a custom test harness. We already have a whole bunch of different test harnesses in this repo, which make it quite confusing to pick _which_ harness/integration test "framework" to write additional integration tests with. None of the existing ones really work for this kind of testing. _So_, in order to make these tests as easy as possible to reason about, they are added as their own individual integration tests in `apollo-federation`'s `Cargo.toml`. You will see them like as regular integration tests when you run `cargo nextest run`: Screenshot 2025-02-17 at 17 57 41 For the time being, the tests use a very tiny, not very complex supergraph and operation. It will be good to add additional tests with a more complex version of both, as the heap usage is going to be vastly different and will scale differently. --- apollo-federation/Cargo.toml | 12 ++- .../connectors_validation.rs} | 6 +- .../tests/dhat_profiling/query_plan.rs | 99 +++++++++++++++++++ .../tests/dhat_profiling/supergraph.rs | 67 +++++++++++++ 4 files changed, 180 insertions(+), 4 deletions(-) rename apollo-federation/tests/{connectors/validation.rs => dhat_profiling/connectors_validation.rs} (76%) create mode 100644 apollo-federation/tests/dhat_profiling/query_plan.rs create mode 100644 apollo-federation/tests/dhat_profiling/supergraph.rs diff --git a/apollo-federation/Cargo.toml b/apollo-federation/Cargo.toml index 7228a9a9fe8..4fb01d7ed2d 100644 --- a/apollo-federation/Cargo.toml +++ b/apollo-federation/Cargo.toml @@ -60,5 +60,13 @@ dhat = "0.3.3" name = "main" [[test]] -name = "isolated_connectors_validation" -path = "tests/connectors/validation.rs" +name = "connectors_validation_profiling" +path = "tests/dhat_profiling/connectors_validation.rs" + +[[test]] +name = "supergraph_creation_profiling" +path = "tests/dhat_profiling/supergraph.rs" + +[[test]] +name = "query_plan_creation_profiling" +path = "tests/dhat_profiling/query_plan.rs" diff --git a/apollo-federation/tests/connectors/validation.rs b/apollo-federation/tests/dhat_profiling/connectors_validation.rs similarity index 76% rename from apollo-federation/tests/connectors/validation.rs rename to apollo-federation/tests/dhat_profiling/connectors_validation.rs index df8640ce4ad..ec12a161591 100644 --- a/apollo-federation/tests/connectors/validation.rs +++ b/apollo-federation/tests/dhat_profiling/connectors_validation.rs @@ -1,12 +1,14 @@ #[global_allocator] pub(crate) static ALLOC: dhat::Alloc = dhat::Alloc; +// Failure of the test can be diagnosed using the dhat-heap.json file. + +// These values should be kept slightly larger (~10%) than the current heap usage to catch +// significant increases. #[test] fn valid_large_body() { const SCHEMA: &str = "src/sources/connect/validation/test_data/valid_large_body.graphql"; - // These values should be kept slightly larger (~10%) than the current heap usage to catch - // significant increases. Failure of the test can be diagnosed using the dhat-heap.json file. const MAX_BYTES: usize = 204_800; // 200 KiB const MAX_ALLOCATIONS: u64 = 22_300; diff --git a/apollo-federation/tests/dhat_profiling/query_plan.rs b/apollo-federation/tests/dhat_profiling/query_plan.rs new file mode 100644 index 00000000000..70e8c8ed7be --- /dev/null +++ b/apollo-federation/tests/dhat_profiling/query_plan.rs @@ -0,0 +1,99 @@ +#[global_allocator] +pub(crate) static ALLOC: dhat::Alloc = dhat::Alloc; + +// Failure of the test can be diagnosed using the dhat-heap.json file. + +// The figures have a 5% buffer from the actual profiling stats. This +// should help us keep an eye on allocation increases, (hopefully) without +// too much flakiness. +#[test] +fn valid_query_plan() { + const SCHEMA: &str = "../examples/graphql/supergraph.graphql"; + const OPERATION: &str = "query fetchUser { + me { + id + name + username + reviews { + ...reviews + } + } + recommendedProducts { + ...products + } + topProducts { + ...products + } + } + fragment products on Product { + upc + weight + price + shippingEstimate + reviews { + ...reviews + } + } + fragment reviews on Review { + id + author { + id + name + } + } + "; + + // Number of bytes when the heap size reached its global maximum with a 5% buffer. + // Actual number: 683_609. + const MAX_BYTES_QUERY_PLANNER: usize = 718_000; // ~720 KiB + + // Total number of allocations with a 5% buffer. + // Actual number: 13_891. + const MAX_ALLOCATIONS_QUERY_PLANNER: u64 = 14_600; + + // Number of bytes when the heap size reached its global maximum with a 5% buffer. + // Actual number: 816_807. + // + // Planning adds 133_198 bytes to heap max (816_807-683_609=133_198). + const MAX_BYTES_QUERY_PLAN: usize = 857_000; // ~860 KiB + + // Total number of allocations with a 5% buffer. + // Actual number: 21_277. + // + // Planning adds 7_386 allocations (21_277-13_891=7_386). + const MAX_ALLOCATIONS_QUERY_PLAN: u64 = 22_400; + + let schema = std::fs::read_to_string(SCHEMA).unwrap(); + + let _profiler = dhat::Profiler::builder().testing().build(); + + let supergraph = + apollo_federation::Supergraph::new(&schema).expect("supergraph should be valid"); + let api_options = apollo_federation::ApiSchemaOptions::default(); + let api_schema = supergraph + .to_api_schema(api_options) + .expect("api schema should be valid"); + let qp_config = apollo_federation::query_plan::query_planner::QueryPlannerConfig::default(); + let planner = apollo_federation::query_plan::query_planner::QueryPlanner::new( + &supergraph, + qp_config.clone(), + ) + .expect("query planner should be created"); + let stats = dhat::HeapStats::get(); + dhat::assert!(stats.max_bytes < MAX_BYTES_QUERY_PLANNER); + dhat::assert!(stats.total_blocks < MAX_ALLOCATIONS_QUERY_PLANNER); + + let document = apollo_compiler::ExecutableDocument::parse_and_validate( + api_schema.schema(), + OPERATION, + "operation.graphql", + ) + .expect("operation should be valid"); + let qp_options = apollo_federation::query_plan::query_planner::QueryPlanOptions::default(); + planner + .build_query_plan(&document, None, qp_options) + .expect("valid query plan"); + let stats = dhat::HeapStats::get(); + dhat::assert!(stats.max_bytes < MAX_BYTES_QUERY_PLAN); + dhat::assert!(stats.total_blocks < MAX_ALLOCATIONS_QUERY_PLAN); +} diff --git a/apollo-federation/tests/dhat_profiling/supergraph.rs b/apollo-federation/tests/dhat_profiling/supergraph.rs new file mode 100644 index 00000000000..92f04bd8c8b --- /dev/null +++ b/apollo-federation/tests/dhat_profiling/supergraph.rs @@ -0,0 +1,67 @@ +#[global_allocator] +pub(crate) static ALLOC: dhat::Alloc = dhat::Alloc; + +// Failure of the test can be diagnosed using the dhat-heap.json file. + +// The figures have a 5% buffer from the actual profiling stats. This +// should help us keep an eye on allocation increases, (hopefully) without +// too much flakiness. +#[test] +fn valid_supergraph_schema() { + const SCHEMA: &str = "../examples/graphql/supergraph.graphql"; + + // Number of bytes when the heap size reached its global maximum with a 5% buffer. + // Actual number: 128_605. + const MAX_BYTES_SUPERGRAPH: usize = 135_050; // ~135 KiB. actual number: 128605 + + // Total number of allocations with a 5% buffer. + // Actual number: 4889. + const MAX_ALLOCATIONS_SUPERGRAPH: u64 = 5_150; // number of allocations. actual number: 4889 + + // Number of bytes when the heap size reached its global maximum with a 5% buffer. + // Actual number: 188_420. + // + // API schema generation allocates additional 59_635 bytes (188_420-128_605=59_635). + const MAX_BYTES_API_SCHEMA: usize = 197_900; // ~200 KiB + + // Total number of allocations with a 5% buffer. + // Actual number: 5_535. + // + // API schema has an additional 646 allocations (5_535-4_889=646). + const MAX_ALLOCATIONS_API_SCHEMA: u64 = 5_800; + + // Number of bytes when the heap size reached its global maximum with a 5% buffer. + // Actual number: 552_781. + // + // Extract subgraphs allocates additional 364_361 bytes (552_781-188_420=364_361). + const MAX_BYTES_SUBGRAPHS: usize = 580_420; // ~600 KiB + + // Total number of allocations with a 5% buffer. + // Actual number: 12_185. + // + // Extract subgraphs from supergraph has an additional 6_650 allocations (12_185-5_535=6_650). + const MAX_ALLOCATIONS_SUBGRAPHS: u64 = 12_800; + + let schema = std::fs::read_to_string(SCHEMA).unwrap(); + + let _profiler = dhat::Profiler::builder().testing().build(); + + let supergraph = + apollo_federation::Supergraph::new(&schema).expect("supergraph should be valid"); + let stats = dhat::HeapStats::get(); + dhat::assert!(stats.max_bytes < MAX_BYTES_SUPERGRAPH); + dhat::assert!(stats.total_blocks < MAX_ALLOCATIONS_SUPERGRAPH); + + let api_options = apollo_federation::ApiSchemaOptions::default(); + let _api_schema = supergraph.to_api_schema(api_options); + let stats = dhat::HeapStats::get(); + dhat::assert!(stats.max_bytes < MAX_BYTES_API_SCHEMA); + dhat::assert!(stats.total_blocks < MAX_ALLOCATIONS_API_SCHEMA); + + let _subgraphs = supergraph + .extract_subgraphs() + .expect("subgraphs should be extracted"); + let stats = dhat::HeapStats::get(); + dhat::assert!(stats.max_bytes < MAX_BYTES_SUBGRAPHS); + dhat::assert!(stats.total_blocks < MAX_ALLOCATIONS_SUBGRAPHS); +}