add dhat heap profiling tests for Supergaph and QueryPlanner functions (

#6814)

Building on the work in #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`:

<img width="887" alt="Screenshot 2025-02-17 at 17 57 41" src="https://github.com/user-attachments/assets/b91b372c-66b4-4db0-a6bd-7d27d15f9ac6" />

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

@@ -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"

apollo-federation/tests/connectors/validation.rs → 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;
 

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);
+}

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);
+}