perf(gc): batch untracked package directory traversals

benches/benchsuite/benches/global_cache_tracker.rs

@@ -1,12 +1,15 @@
 //! Benchmarks for the global cache tracker.
 
+use cargo::ops::CleanContext;
 use cargo::util::GlobalContext;
 use cargo::util::cache_lock::CacheLockMode;
 use cargo::util::interning::InternedString;
+use cargo::workspace::gc::GcOpts;
 use cargo::workspace::global_cache_tracker::{self, DeferredGlobalLastUse, GlobalCacheTracker};
-use criterion::{Criterion, criterion_group, criterion_main};
+use criterion::{BatchSize, Criterion, criterion_group, criterion_main};
 use std::fs;
 use std::path::{Path, PathBuf};
+use std::time::Duration;
 
 // Samples of real-world data.
 const GLOBAL_CACHE_SAMPLE: &str = "global-cache-tracker/global-cache-sample";
@@ -149,10 +152,74 @@ fn global_tracker_update(c: &mut Criterion) {
     }
 }
 
+/// Creates a registry source directory holding `packages` package
+/// directories, as an older cargo that didn't track them would leave behind.
+fn untracked_src_dir(packages: usize) -> PathBuf {
+    let src = cargo_home().join("registry/src/bench-0000000000000000");
+    if src.exists() {
+        return src;
+    }
+    // The source directories are only tracked once their registry is, which
+    // happens by finding its index on disk.
+    fs::create_dir_all(cargo_home().join("registry/index/bench-0000000000000000/.cache")).unwrap();
+    for package in 0..packages {
+        let dir = src.join(format!("some-crate-{package}-0.1.0"));
+        fs::create_dir_all(dir.join("src")).unwrap();
+        fs::write(dir.join("Cargo.toml"), "[package]\n").unwrap();
+        for module in 0..5 {
+            fs::write(dir.join("src").join(format!("m{module}.rs")), "// x\n").unwrap();
+        }
+    }
+    src
+}
+
+/// Benchmarks synchronizing the database with a cache left by a cargo that
+/// didn't track it, which has to discover every package directory on disk.
+///
+/// Sizing those directories is a separate, much more expensive step, only
+/// taken when a size limit is given, so both are measured.
+fn global_tracker_sync(c: &mut Criterion) {
+    let gctx = initialize_context();
+    let _lock = gctx
+        .acquire_package_cache_lock(CacheLockMode::MutateExclusive)
+        .unwrap();
+    let db_path = GlobalCacheTracker::db_path(&gctx).into_path_unlocked();
+    untracked_src_dir(500);
+
+    let mut group = c.benchmark_group("global_tracker_sync");
+    for (name, max_download_size) in [("age_only", None), ("with_size", Some(u64::MAX))] {
+        let gc_opts = GcOpts {
+            max_src_age: Some(Duration::from_secs(0)),
+            max_download_size,
+            ..GcOpts::default()
+        };
+        group.bench_function(name, |b| {
+            b.iter_batched(
+                || {
+                    // Start each iteration with nothing tracked, so that the
+                    // whole directory is discovered again.
+                    if db_path.exists() {
+                        fs::remove_file(&db_path).unwrap();
+                    }
+                    GlobalCacheTracker::new(&gctx).unwrap()
+                },
+                |mut tracker| {
+                    let mut clean_ctx = CleanContext::new(&gctx);
+                    clean_ctx.dry_run = true;
+                    tracker.clean(&mut clean_ctx, &gc_opts).unwrap();
+                },
+                BatchSize::PerIteration,
+            )
+        });
+    }
+    group.finish();
+}
+
 criterion_group!(
     benches,
     global_tracker_init,
     global_tracker_empty_save,
-    global_tracker_update
+    global_tracker_update,
+    global_tracker_sync
 );
 criterion_main!(benches);

crates/cargo-util/src/du.rs

@@ -1,8 +1,9 @@
 //! A simple disk usage estimator.
 
+use std::collections::BTreeMap;
 use std::path::Path;
+use std::sync::Mutex;
 use std::sync::atomic::Ordering;
-use std::sync::{Arc, Mutex};
 
 use anyhow::{Context, Result};
 use ignore::overrides::OverrideBuilder;
@@ -22,17 +23,39 @@ use portable_atomic::AtomicU64;
 /// thus vastly undercounts directories with lots of small files). It would be
 /// nice to improve this or replace it with something better.
 pub fn du(path: &Path, patterns: &[&str]) -> Result<u64> {
-    du_inner(path, patterns).with_context(|| format!("failed to walk `{}`", path.display()))
+    Ok(du_each(path, &[path], patterns)?[0])
 }
 
-fn du_inner(path: &Path, patterns: &[&str]) -> Result<u64> {
-    let mut builder = OverrideBuilder::new(path);
+/// Determines the disk usage of each of the given directories, in order.
+///
+/// The patterns are as for [`du`], relative to `root`, which every path must
+/// be within. The paths must not be nested within one another.
+///
+/// A walk builds a thread pool of its own, which for a small directory costs
+/// far more than the walk itself. Prefer this over calling [`du`] in a loop:
+/// it walks every path with a single pool.
+pub fn du_each(root: &Path, paths: &[&Path], patterns: &[&str]) -> Result<Vec<u64>> {
+    du_each_inner(root, paths, patterns)
+        .with_context(|| format!("failed to walk `{}`", root.display()))
+}
+
+fn du_each_inner(root: &Path, paths: &[&Path], patterns: &[&str]) -> Result<Vec<u64>> {
+    let Some((first, rest)) = paths.split_first() else {
+        return Ok(Vec::new());
+    };
+
+    let mut builder = OverrideBuilder::new(root);
     for pattern in patterns {
         builder.add(pattern)?;
     }
     let overrides = builder.build()?;
 
-    let mut builder = WalkBuilder::new(path);
+    let mut builder = WalkBuilder::new(first);
+
+    for path in rest {
+        builder.add(path);
+    }
+
     builder
         .overrides(overrides)
         .hidden(false)
@@ -43,20 +66,30 @@ fn du_inner(path: &Path, patterns: &[&str]) -> Result<u64> {
         .git_exclude(false);
     let walker = builder.build_parallel();
 
-    let total = Arc::new(AtomicU64::new(0));
+    // BTreeMap, not HashMap, to stay off a faster-hasher dependency; this
+    // per-file lookup is dwarfed by the `stat` on each entry regardless.
+    let path_to_index: BTreeMap<&Path, usize> = paths.iter().copied().zip(0..).collect();
+    let total: Vec<_> = paths.iter().map(|_| AtomicU64::new(0)).collect();
 
     // A slot used to indicate there was an error while walking.
     //
     // It is possible that more than one error happens (such as in different
     // threads). The error returned is arbitrary in that case.
-    let err = Arc::new(Mutex::new(None));
+    let err = Mutex::new(None);
     walker.run(|| {
         Box::new(|entry| {
             match entry {
                 Ok(entry) => match entry.metadata() {
                     Ok(meta) => {
                         if meta.is_file() {
-                            total.fetch_add(meta.len(), Ordering::Relaxed);
+                            // Attribute each file to the directory it belongs under.
+                            // An entry's depth is counted from the path it was
+                            // reached through, which is that many levels up.
+                            let path = entry.path().ancestors().nth(entry.depth());
+
+                            if let Some(index) = path.and_then(|path| path_to_index.get(path)) {
+                                total[*index].fetch_add(meta.len(), Ordering::Relaxed);
+                            }
                         }
                     }
                     Err(e) => {
@@ -77,5 +110,69 @@ fn du_inner(path: &Path, patterns: &[&str]) -> Result<u64> {
         return Err(e);
     }
 
-    Ok(total.load(Ordering::Relaxed))
+    Ok(total
+        .iter()
+        .map(|total| total.load(Ordering::Relaxed))
+        .collect())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::fs;
+
+    /// Creates `<root>/<name>/{Cargo.toml, src/lib.rs, .git/objects}`,
+    /// each file holding a distinct number of bytes.
+    fn package(root: &Path, name: &str, size: u64) {
+        let dir = root.join(name);
+        fs::create_dir_all(dir.join("src")).unwrap();
+        fs::create_dir_all(dir.join(".git")).unwrap();
+        fs::write(dir.join("Cargo.toml"), vec![b'x'; size as usize]).unwrap();
+        fs::write(dir.join("src/lib.rs"), vec![b'x'; size as usize * 2]).unwrap();
+        fs::write(dir.join(".git/objects"), vec![b'x'; size as usize * 100]).unwrap();
+    }
+
+    #[test]
+    fn du_each_totals_each_path_separately() {
+        let tmp = tempfile::TempDir::new().unwrap();
+        let root = tmp.path();
+
+        package(root, "a", 1);
+        package(root, "b", 10);
+        package(root, "c", 100);
+
+        let paths = [&root.join("a"), &root.join("b"), &root.join("c")];
+        let paths: Vec<_> = paths.iter().map(|path| path.as_path()).collect();
+        assert_eq!(du_each(root, &paths, &[]).unwrap(), [103, 1030, 10300]);
+
+        // Walking each path on its own must agree.
+        let separately: Vec<_> = paths.iter().map(|path| du(path, &[]).unwrap()).collect();
+        assert_eq!(du_each(root, &paths, &[]).unwrap(), separately);
+    }
+
+    /// Patterns are anchored at the root shared by every path, rather than at
+    /// each path, so make sure they still apply within each of them.
+    #[test]
+    fn du_each_applies_patterns_within_each_path() {
+        let tmp = tempfile::TempDir::new().unwrap();
+        let root = tmp.path();
+        package(root, "a", 1);
+        package(root, "b", 10);
+
+        let paths = [&root.join("a"), &root.join("b")];
+        let paths: Vec<_> = paths.iter().map(|path| path.as_path()).collect();
+        assert_eq!(du_each(root, &paths, &["!.git"]).unwrap(), [3, 30]);
+
+        let separately: Vec<_> = paths
+            .iter()
+            .map(|path| du(path, &["!.git"]).unwrap())
+            .collect();
+        assert_eq!(du_each(root, &paths, &["!.git"]).unwrap(), separately);
+    }
+
+    #[test]
+    fn du_each_of_nothing() {
+        let tmp = tempfile::TempDir::new().unwrap();
+        assert_eq!(du_each(tmp.path(), &[], &[]).unwrap(), []);
+    }
 }

crates/cargo-util/src/lib.rs

@@ -7,7 +7,7 @@
 #![allow(clippy::disallowed_methods)]
 
 pub use self::read2::read2;
-pub use du::du;
+pub use du::{du, du_each};
 pub use process_builder::ProcessBuilder;
 pub use process_error::{ProcessError, exit_status_to_string, is_simple_exit_code};
 pub use sha256::Sha256;

src/workspace/global_cache_tracker.rs

@@ -975,6 +975,11 @@ impl GlobalCacheTracker {
             let index_path = base_path.join(id_name);
             let names = Self::list_dir_names(&index_path)?;
             let max = names.len();
+
+            // Gather the untracked directories before sizing them,
+            // so that they can all be walked together.
+            let mut untracked_directories = Vec::new();
+
             for (i, name) in names.iter().enumerate() {
                 if select_stmt.exists(params![id, name])? {
                     continue;
@@ -984,11 +989,26 @@ impl GlobalCacheTracker {
                     continue;
                 }
                 progress.tick(i, max, "")?;
-                let size = if populate_size {
-                    Some(du(&dir_path, table_name)?)
-                } else {
-                    None
-                };
+                untracked_directories.push((name, dir_path));
+            }
+
+            let sizes = if populate_size {
+                let paths: Vec<_> = untracked_directories
+                    .iter()
+                    .map(|(_, path)| path.as_path())
+                    .collect();
+
+                // Size each untracked directory all at once.
+                du_each(&index_path, &paths, table_name)?
+                    .into_iter()
+                    .map(Some)
+                    .collect()
+            } else {
+                // Skip walking entirely.
+                vec![None; untracked_directories.len()]
+            };
+
+            for ((name, _), size) in untracked_directories.iter().zip(sizes) {
                 insert_stmt.execute(params![id, name, size, now])?;
             }
         }
@@ -1027,13 +1047,20 @@ impl GlobalCacheTracker {
             })?
             .collect();
         let max = rows.len();
+        let mut row_ids = Vec::with_capacity(max);
+        let mut paths = Vec::with_capacity(max);
+
         for (i, row) in rows.into_iter().enumerate() {
             let (rowid, name, id_name): (i64, String, String) = row?;
-            let path = base_path.join(id_name).join(name);
             progress.tick(i, max, "")?;
-            // Missing files should have already been taken care of by
-            // update_db_for_removed.
-            let size = du(&path, table_name)?;
+            row_ids.push(rowid);
+            paths.push(base_path.join(id_name).join(name));
+        }
+
+        // Missing files should have already been taken care of by update_db_for_removed.
+        let paths: Vec<_> = paths.iter().map(|path| path.as_path()).collect();
+
+        for (rowid, size) in row_ids.iter().zip(du_each(base_path, &paths, table_name)?) {
             update_stmt.execute(params![size, rowid])?;
         }
         Ok(())
@@ -1822,19 +1849,18 @@ pub fn is_silent_error(e: &anyhow::Error) -> bool {
     false
 }
 
-/// Returns the disk usage for a git checkout directory.
-#[tracing::instrument]
-fn du_git_checkout(path: &Path) -> CargoResult<u64> {
-    // !.git is used because clones typically use hardlinks for the git
-    // contents. TODO: Verify behavior on Windows.
-    // TODO: Or even better, switch to worktrees, and remove this.
-    cargo_util::du(&path, &["!.git"])
-}
-
-fn du(path: &Path, table_name: &str) -> CargoResult<u64> {
-    if table_name == GIT_CO_TABLE {
-        du_git_checkout(path)
+/// Returns the disk usage of each of `paths`, which must all be within `root`.
+// skip(paths) to keep a potentially large slice out of tracing.
+#[tracing::instrument(skip(paths))]
+fn du_each(root: &Path, paths: &[&Path], table_name: &str) -> CargoResult<Vec<u64>> {
+    let patterns: &[&str] = if table_name == GIT_CO_TABLE {
+        // !.git is used because clones typically use hardlinks for the git
+        // contents. TODO: Verify behavior on Windows.
+        // TODO: Or even better, switch to worktrees, and remove this.
+        &["!.git"]
     } else {
-        cargo_util::du(&path, &[])
-    }
+        &[]
+    };
+
+    cargo_util::du_each(root, paths, patterns)
 }