perf: Convert JSON docs to the expected caching format in parallel

LukeMathWalker · LukeMathWalker · commit 7e4e89cebba3 · 2025-12-26T11:03:12.000+01:00
Insertion happens on the "main" thread, but the more expensive
format conversion is executed on a threadpool.
diff --git a/compiler/pavex_test_runner/Cargo.toml b/compiler/pavex_test_runner/Cargo.toml
@@ -16,6 +16,7 @@ anyhow = { workspace = true }
 ahash = { workspace = true }
 console = { workspace = true }
 cargo_metadata = { workspace = true }
+cargo-like-utils = { workspace = true }
 fs-err = { workspace = true }
 libtest-mimic = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
@@ -32,6 +33,7 @@ object-pool = { workspace = true }
 num_cpus = { workspace = true }
 globwalk = { workspace = true }
 tracing-subscriber = { workspace = true, features = ["env-filter", "fmt"] }
+pavex_cli_shell = { workspace = true }
 regex = { workspace = true }
 once_cell = { workspace = true }
 pavexc = { workspace = true }
diff --git a/compiler/pavex_test_runner/src/lib.rs b/compiler/pavex_test_runner/src/lib.rs
@@ -6,11 +6,13 @@ use std::process::{Command, Output};
 
 use ahash::HashSet;
 use anyhow::Context;
+use cargo_like_utils::shell::Shell;
 use cargo_metadata::diagnostic::DiagnosticLevel;
 use console::style;
 use guppy::graph::PackageGraph;
 use itertools::Itertools;
 use libtest_mimic::{Arguments, Conclusion, Failed, Trial};
+use pavex_cli_shell::try_init_shell;
 use pavexc::rustdoc::CrateCollection;
 use pavexc::{DEFAULT_DOCS_TOOLCHAIN, DiagnosticSink};
 use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
@@ -112,6 +114,7 @@ pub fn run_tests(
         // acquire the global target directory lock. If multiple tests end up
         // doing this at the same time, their execution will be serialized,
         // which would have a major impact on the overall test suite runtime.
+        try_init_shell(Shell::new());
         warm_up_rustdoc_cache(&package_graph, &test_name2test_data)?;
 
         // First battery of UI tests.
@@ -426,7 +429,6 @@ fn warm_up_rustdoc_cache(
     // We want to ensure that all invocations of `pavexc generate` hit the cache
     // thus avoiding the need to invoke `rustdoc` and acquire a contentious
     // lock over the target directory.
-    println!("Pre-computing JSON documentation for relevant crates");
     let crate_collection = CrateCollection::new(
         DEFAULT_DOCS_TOOLCHAIN.to_owned(),
         package_graph.clone(),
@@ -443,14 +445,20 @@ fn warm_up_rustdoc_cache(
         .workspace()
         .iter()
         .filter(|p| {
-            !(p.name().starts_with("application_")
+            !(
+                // Avoid computing JOSN docs for:
+                // - Generated code
+                p.name().starts_with("application_")
+                // - Integration test crates
                 || p.name().starts_with("integration_")
+                // - The workspace hack package
                 || p.name() == "workspace_hack"
-                || (p.name().starts_with("app_") && !app_names.contains(p.name())))
+                // - Tests that have been filtered out of this run
+                || (p.name().starts_with("app_") && !app_names.contains(p.name()))
+            )
         })
         .map(|p| p.name().to_owned())
         .collect();
-    crates.remove("workspace_hack");
     // Toolchain docs
     crates.insert("core".into());
     crates.insert("alloc".into());
@@ -471,16 +479,19 @@ fn warm_up_rustdoc_cache(
     crates.insert("yansi".into());
     crates.insert("serde".into());
     crates.insert("zerocopy".into());
-    let package_ids = crate_collection
+    let package_ids: Vec<_> = crate_collection
         .package_graph()
         .packages()
         .filter(|p| crates.contains(p.name()))
-        .map(|p| p.id().to_owned());
+        .map(|p| p.id().to_owned())
+        .collect();
+    let n_packages = package_ids.len();
+    println!("Pre-computing JSON documentation for {n_packages} crates",);
     crate_collection
-        .batch_compute_crates(package_ids)
+        .batch_compute_crates(package_ids.into_iter())
         .context("Failed to warm rustdoc JSON cache")?;
     println!(
-        "Pre-computed JSON documentation in {} seconds",
+        "Pre-computed JSON documentation for {n_packages} crates in {} seconds",
         timer.elapsed().as_secs()
     );
 
diff --git a/compiler/pavexc/src/rustdoc/compute/cache.rs b/compiler/pavexc/src/rustdoc/compute/cache.rs
@@ -137,8 +137,9 @@ impl RustdocGlobalFsCache {
         }
     }
 
-    /// Store the JSON documentation generated by `rustdoc` in the cache.
-    pub(crate) fn insert(
+    /// Convert the JSON documentation generated by `rustdoc` into the format used by our cache,
+    /// then store it.
+    pub(crate) fn convert_and_insert(
         &self,
         cache_key: &RustdocCacheKey,
         krate: &crate::rustdoc::Crate,
@@ -147,17 +148,53 @@ impl RustdocGlobalFsCache {
     ) -> Result<(), anyhow::Error> {
         let connection = self.connection_pool.get()?;
         match cache_key {
-            RustdocCacheKey::ThirdPartyCrate(metadata) => self.third_party_cache.insert(
-                metadata,
+            RustdocCacheKey::ThirdPartyCrate(metadata) => {
+                self.third_party_cache.convert_and_insert(
+                    metadata,
+                    krate,
+                    &self.cargo_fingerprint,
+                    &connection,
+                    cache_indexes,
+                    package_graph,
+                )
+            }
+            RustdocCacheKey::ToolchainCrate(name) => self.toolchain_cache.convert_and_insert(
+                name,
                 krate,
                 &self.cargo_fingerprint,
                 &connection,
-                cache_indexes,
-                package_graph,
             ),
+        }
+    }
+
+    /// Store the JSON documentation for a crate, which has already been converted to the expected
+    /// format for caching.
+    ///
+    /// This method should be preferred to [`Self::insert`] whenever multi-threading is involved,
+    /// since [`RustdocGlobalFsCache`] isn't thread-safe, but the conversion into the caching format
+    /// doesn't require its internals and can therefore be offloaded to another thread without issues.
+    pub(crate) fn insert(
+        &self,
+        cache_key: &RustdocCacheKey,
+        cache_entry: CacheEntry,
+        package_graph: &PackageGraph,
+    ) -> Result<(), anyhow::Error> {
+        let connection = self.connection_pool.get()?;
+        match cache_key {
+            RustdocCacheKey::ThirdPartyCrate(metadata) => {
+                let Some(cache_key) = self.third_party_cache.cache_key(
+                    metadata,
+                    &self.cargo_fingerprint,
+                    package_graph,
+                ) else {
+                    return Ok(());
+                };
+                self.third_party_cache
+                    .insert(cache_key, &connection, cache_entry)
+            }
             RustdocCacheKey::ToolchainCrate(name) => {
                 self.toolchain_cache
-                    .insert(name, krate, &self.cargo_fingerprint, &connection)
+                    .insert(name, cache_entry, &self.cargo_fingerprint, &connection)
             }
         }
     }
@@ -317,7 +354,7 @@ impl ToolchainCache {
         let import_path2id = row.get_ref_unwrap(7).as_bytes()?;
         let re_exports = row.get_ref_unwrap(8).as_bytes()?;
 
-        let krate = CachedData {
+        let krate = CacheEntry {
             root_item_id,
             external_crates: Cow::Borrowed(external_crates),
             paths: Cow::Borrowed(paths),
@@ -339,14 +376,26 @@ impl ToolchainCache {
 
     /// Store the JSON documentation for a toolchain crate in the cache.
     #[instrument(name = "Cache rustdoc output on disk", skip_all, level=tracing::Level::DEBUG, fields(crate.name = name))]
-    fn insert(
+    fn convert_and_insert(
         &self,
         name: &str,
         krate: &crate::rustdoc::Crate,
         cargo_fingerprint: &str,
         connection: &rusqlite::Connection,
     ) -> Result<(), anyhow::Error> {
-        let cached_data = CachedData::new(krate).context("Failed to serialize docs")?;
+        let cache_entry = CacheEntry::new(krate).context("Failed to serialize docs")?;
+        self.insert(name, cache_entry, cargo_fingerprint, connection)
+    }
+
+    /// Store the JSON documentation for a toolchain crate in the cache.
+    #[instrument(name = "Cache rustdoc output on disk", skip_all, level=tracing::Level::DEBUG, fields(crate.name = name))]
+    fn insert(
+        &self,
+        name: &str,
+        cache_entry: CacheEntry<'_>,
+        cargo_fingerprint: &str,
+        connection: &rusqlite::Connection,
+    ) -> Result<(), anyhow::Error> {
         let mut stmt = connection.prepare_cached(
             "INSERT INTO rustdoc_toolchain_crates_cache (
                 name,
@@ -365,23 +414,23 @@ impl ToolchainCache {
         stmt.execute(params![
             name,
             cargo_fingerprint,
-            cached_data.root_item_id,
-            cached_data.external_crates,
-            cached_data.paths,
-            cached_data.format_version,
-            cached_data.items,
-            cached_data.item_id2delimiters,
-            cached_data
+            cache_entry.root_item_id,
+            cache_entry.external_crates,
+            cache_entry.paths,
+            cache_entry.format_version,
+            cache_entry.items,
+            cache_entry.item_id2delimiters,
+            cache_entry
                 .secondary_indexes
                 .as_ref()
                 .expect("Indexing never fails for toolchain crates")
                 .import_index,
-            cached_data
+            cache_entry
                 .secondary_indexes
                 .as_ref()
                 .expect("Indexing never fails for toolchain crates")
                 .import_path2id,
-            cached_data
+            cache_entry
                 .secondary_indexes
                 .as_ref()
                 .expect("Indexing never fails for toolchain crates")
@@ -537,7 +586,7 @@ impl ThirdPartyCrateCache {
                     _ => None,
                 };
 
-            let krate = CachedData {
+            let krate = CacheEntry {
                 root_item_id,
                 external_crates: Cow::Borrowed(external_crates),
                 paths: Cow::Borrowed(paths),
@@ -570,14 +619,30 @@ impl ThirdPartyCrateCache {
         outcome
     }
 
-    /// Store the JSON documentation generated by `rustdoc` in the cache.
+    /// Compute the cache key for a given package.
+    fn cache_key<'a>(
+        &self,
+        package_metadata: &'a PackageMetadata,
+        cargo_fingerprint: &'a str,
+        package_graph: &PackageGraph,
+    ) -> Option<ThirdPartyCrateCacheKey<'a>> {
+        ThirdPartyCrateCacheKey::build(
+            package_graph,
+            package_metadata,
+            cargo_fingerprint,
+            self.cache_workspace_packages,
+        )
+    }
+
+    /// Convert the JSON documentation generated by `rustdoc` to the format used by our cache,
+    /// then store it.
     #[instrument(
-        name = "Cache third-party crate docs to disk",
+        name = "Convert and cache docs for a third-party crate to disk",
         skip_all,
         level=tracing::Level::DEBUG,
         fields(crate.id = %package_metadata.id(), cache_key = tracing::field::Empty))
     ]
-    fn insert(
+    fn convert_and_insert(
         &self,
         package_metadata: &PackageMetadata,
         krate: &crate::rustdoc::Crate,
@@ -596,11 +661,29 @@ impl ThirdPartyCrateCache {
         };
         tracing::Span::current().record("cache_key", tracing::field::debug(&cache_key));
         let cached_data = if cache_indexes {
-            CachedData::new(krate)
+            CacheEntry::new(krate)
         } else {
-            CachedData::raw(krate)
+            CacheEntry::raw(krate)
         }
         .context("Failed to serialize docs")?;
+        self.insert(cache_key, connection, cached_data)
+    }
+
+    /// Store the JSON documentation generated by `rustdoc` in the cache,
+    /// without having to perform the conversion towards the caching format.
+    #[instrument(
+        name = "Stored cache data for third-party crate docs to disk",
+        skip_all,
+        level=tracing::Level::DEBUG,
+        fields(cache_key = tracing::field::Empty))
+    ]
+    fn insert(
+        &self,
+        cache_key: ThirdPartyCrateCacheKey<'_>,
+        connection: &rusqlite::Connection,
+        cached_data: CacheEntry<'_>,
+    ) -> Result<(), anyhow::Error> {
+        tracing::Span::current().record("cache_key", tracing::field::debug(&cache_key));
         let mut stmt = connection.prepare_cached(
             "INSERT INTO rustdoc_3d_party_crates_cache (
                 crate_name,
@@ -692,7 +775,7 @@ impl ThirdPartyCrateCache {
 
 #[derive(Debug)]
 /// The serialized form of a crate's documentation, as stored in the cache.
-pub(super) struct CachedData<'a> {
+pub(in crate::rustdoc) struct CacheEntry<'a> {
     root_item_id: u32,
     external_crates: Cow<'a, [u8]>,
     paths: Cow<'a, [u8]>,
@@ -705,15 +788,15 @@ pub(super) struct CachedData<'a> {
 #[derive(Debug)]
 /// Data that can be computed starting from the raw JSON documentation for a crate,
 /// without having to re-invoke `rustdoc`.
-pub(super) struct SecondaryIndexes<'a> {
+pub(in crate::rustdoc) struct SecondaryIndexes<'a> {
     import_index: Cow<'a, [u8]>,
     annotated_items: Option<Cow<'a, [u8]>>,
     import_path2id: Cow<'a, [u8]>,
     re_exports: Cow<'a, [u8]>,
 }
 
-impl<'a> CachedData<'a> {
-    pub(super) fn new(krate: &'a crate::rustdoc::Crate) -> Result<CachedData<'a>, anyhow::Error> {
+impl<'a> CacheEntry<'a> {
+    pub fn new(krate: &'a crate::rustdoc::Crate) -> Result<CacheEntry<'a>, anyhow::Error> {
         let mut cached = Self::raw(krate)?;
 
         let import_index = bincode::serde::encode_to_vec(&krate.import_index, BINCODE_CONFIG)?;
@@ -732,7 +815,7 @@ impl<'a> CachedData<'a> {
     }
 
     /// Cache only `rustdoc`'s output, no secondary indexes.
-    pub(super) fn raw(krate: &'a crate::rustdoc::Crate) -> Result<CachedData<'a>, anyhow::Error> {
+    pub fn raw(krate: &'a crate::rustdoc::Crate) -> Result<CacheEntry<'a>, anyhow::Error> {
         let crate_data = &krate.core.krate;
         let CrateItemIndex::Eager(index) = &crate_data.index else {
             anyhow::bail!(
@@ -752,7 +835,7 @@ impl<'a> CachedData<'a> {
             bincode::serde::encode_to_vec(&crate_data.external_crates, BINCODE_CONFIG)?;
         let paths = bincode::serde::encode_to_vec(&crate_data.paths, BINCODE_CONFIG)?;
 
-        Ok(CachedData {
+        Ok(CacheEntry {
             root_item_id: crate_data.root_item_id.0,
             external_crates: Cow::Owned(external_crates),
             paths: Cow::Owned(paths),
diff --git a/compiler/pavexc/src/rustdoc/compute/mod.rs b/compiler/pavexc/src/rustdoc/compute/mod.rs
@@ -8,6 +8,7 @@ mod format;
 mod toolchain;
 
 use ahash::{HashMap, HashMapExt};
+pub(super) use cache::CacheEntry;
 pub(crate) use cache::{RustdocCacheKey, RustdocGlobalFsCache};
 
 use anyhow::Context;
diff --git a/compiler/pavexc/src/rustdoc/queries.rs b/compiler/pavexc/src/rustdoc/queries.rs
