Add memory benchmark through jemalloc allocation trace

rust/rubydex/Cargo.toml

@@ -44,5 +44,13 @@ assert_cmd = "2.0"
 predicates = "3.1"
 regex = "1.10"
 
+[target.'cfg(not(target_os = "windows"))'.dev-dependencies]
+tikv-jemallocator = { version = "0.7.0", features = ["stats"] }
+tikv-jemalloc-ctl = { version = "0.7.0", features = ["stats"] }
+
+[[bench]]
+name = "graph_memory"
+harness = false
+
 [lints]
 workspace = true

rust/rubydex/benches/graph_memory.rs

@@ -0,0 +1,52 @@
+//! Jemalloc does not compile on Windows, so this memory benchmark executable is a no-op
+
+#[cfg(not(target_os = "windows"))]
+mod imp {
+    use std::collections::HashSet;
+
+    use rubydex::{
+        indexing::{self, IndexerBackend},
+        listing,
+        model::graph::Graph,
+        resolution::Resolver,
+    };
+    use tikv_jemalloc_ctl::{epoch, stats};
+
+    // Substitute the global allocator by jemalloc, so that we can track all allocations and measure the graph usage
+    #[global_allocator]
+    static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
+    /// Advance the jemalloc epoch (stats are cached between epochs) and read the
+    /// number of bytes currently allocated by the application.
+    fn allocated_bytes() -> usize {
+        epoch::advance().expect("failed to advance jemalloc epoch");
+        stats::allocated::read().expect("failed to read stats.allocated (is the `stats` feature on?)")
+    }
+
+    pub fn run() {
+        let paths: Vec<String> = std::env::args().skip(1).collect();
+        let paths = if paths.is_empty() { vec![".".to_string()] } else { paths };
+        let (file_paths, _) = listing::collect_file_paths(paths, &HashSet::new());
+
+        let mut graph = Graph::new();
+        let _ = indexing::index_files(&mut graph, file_paths, IndexerBackend::RubyIndexer);
+        Resolver::new(&mut graph).resolve();
+
+        // Compare the total memory used in the allocator before and after dropping the graph
+        let before_drop = allocated_bytes();
+        drop(graph);
+        let after_drop = allocated_bytes();
+
+        let total_graph_memory = before_drop.saturating_sub(after_drop);
+
+        #[allow(clippy::cast_precision_loss)]
+        let mega_bytes = total_graph_memory as f64 / 1024.0 / 1024.0;
+
+        println!("Total graph memory: {mega_bytes:.2} MB");
+    }
+}
+
+fn main() {
+    #[cfg(not(target_os = "windows"))]
+    imp::run();
+}

utils/bench-graph-memory

@@ -0,0 +1,66 @@
+#!/bin/bash
+
+# Measure the total memory used by the graph. Allows comparing a branch against main
+#
+# Usage:
+#   utils/bench-graph-memory <path>             # current branch only
+#   utils/bench-graph-memory --compare <path>   # main, then current branch
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
+RUST_DIR="$PROJECT_ROOT/rust"
+COMPARE=0
+TARGET=""
+BRANCH="$(git branch --show-current)"
+
+for arg in "$@"; do
+    case "$arg" in
+        --compare) COMPARE=1 ;;
+        *) TARGET="$arg" ;;
+    esac
+done
+
+if [ -z "$TARGET" ]; then
+    echo "Usage: $0 [--compare] <path>" >&2
+    exit 1
+fi
+
+# Expand the target based on Rubydex's root directory
+case "$TARGET" in
+    /*) ;;
+    *) TARGET="$PROJECT_ROOT/$TARGET" ;;
+esac
+
+run_bench() {
+    cd "$RUST_DIR" && cargo bench --bench graph_memory -- "$TARGET"
+}
+
+if [ "$COMPARE" -eq 1 ]; then
+    if ! git merge-base --is-ancestor main HEAD; then
+        echo "Error: $BRANCH is not based on the current local main" >&2
+        echo "Please rebase first and run again." >&2
+        exit 1
+    fi
+
+    STASHED=0
+
+    if [ -n "$(git status --porcelain)" ]; then
+        git stash push --include-untracked --quiet
+        STASHED=1
+    fi
+
+    git checkout --quiet main
+    echo "########### Measuring main ############"
+    run_bench
+
+    git checkout --quiet "$BRANCH"
+
+    if [ "$STASHED" -eq 1 ]; then
+        git stash pop --quiet
+    fi
+fi
+
+echo "########### Measuring $BRANCH ############"
+run_bench