diff --git a/cli/src/cli_util.rs b/cli/src/cli_util.rs index 103ea17322..92b07ce278 100644 --- a/cli/src/cli_util.rs +++ b/cli/src/cli_util.rs @@ -23,7 +23,7 @@ use std::path::{Path, PathBuf}; use std::process::ExitCode; use std::rc::Rc; use std::str::FromStr; -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; use std::time::SystemTime; use std::{fs, str}; @@ -31,6 +31,7 @@ use clap::builder::{NonEmptyStringValueParser, TypedValueParser, ValueParserFact use clap::{Arg, ArgAction, ArgMatches, Command, FromArgMatches}; use indexmap::{IndexMap, IndexSet}; use itertools::Itertools; +use jj_cbits::mimalloc; use jj_lib::backend::{ChangeId, CommitId, MergedTreeId}; use jj_lib::commit::Commit; use jj_lib::git_backend::GitBackend; @@ -2151,6 +2152,57 @@ pub struct EarlyArgs { pub config_toml: Vec, } +#[derive(clap::Args, Clone, Debug)] +pub struct ShowAllocStats { + /// Show memory allocation statistics from the internal heap allocator + /// on `stdout`, when the program exits. + #[arg(long, global = true)] + show_heap_stats: bool, +} + +/// Lazy global static. Used only to defer printing mimalloc stats until the +/// program exits, if set to `true`. +static PRINT_HEAP_STATS: OnceLock = OnceLock::new(); + +/// Enable heap statistics for the user interface; should be used with +/// [`CliRunner::add_global_args`]. Does nothing if the memory allocator is +/// unused, i.e. `#[global_allocator]` is not set to mimalloc in your program. +pub fn heap_stats_enable(_ui: &mut Ui, opts: ShowAllocStats) -> Result<(), CommandError> { + if opts.show_heap_stats { + PRINT_HEAP_STATS.set(true).unwrap(); + } + Ok(()) +} + +/// Reset heap allocation statistics for the memory allocator. Often used at the +/// very beginning of the program (to clear allocations that may happen before +/// `_main` execution.) +pub fn heap_stats_reset() { + mimalloc::stats_reset(); +} + +/// Print heap allocation statistics to `stderr`, if enabled. +pub fn heap_stats_print() { + // NOTE (aseipp): can we do our own custom printing here? it's kind of ugly + if PRINT_HEAP_STATS.get() == Some(&true) { + eprintln!("========================================"); + eprintln!("mimalloc memory allocation statistics:\n"); + mimalloc::stats_print(&|l| { + eprint!("{}", l.to_string_lossy()); + }); + } +} + +/// Wrap a given closure with calls to [`heap_stats_reset`] and +/// [`heap_stats_print`], and return the result of the closure. Useful for +/// conveiently printing heap allocation statistics for a given function body. +pub fn heap_stats_with_closure(f: impl FnOnce() -> R) -> R { + heap_stats_reset(); + let result = f(); + heap_stats_print(); + result +} + /// Create a description from a list of paragraphs. /// /// Based on the Git CLI behavior. See `opt_parse_m()` and `cleanup_mode` in diff --git a/cli/src/main.rs b/cli/src/main.rs index 38ea1a6870..07a9cc3c5e 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -12,12 +12,19 @@ // See the License for the specific language governing permissions and // limitations under the License. +//! Primary entrypoint for the stock `jj` command-line tool. mimalloc enabled. + use jj_cbits::mimalloc::MiMalloc; -use jj_cli::cli_util::CliRunner; +use jj_cli::cli_util::{heap_stats_enable, heap_stats_with_closure, CliRunner}; #[global_allocator] static ALLOC: MiMalloc = MiMalloc; fn main() -> std::process::ExitCode { - CliRunner::init().version(env!("JJ_VERSION")).run() + heap_stats_with_closure(|| { + CliRunner::init() + .add_global_args(heap_stats_enable) + .version(env!("JJ_VERSION")) + .run() + }) } diff --git a/cli/tests/cli-reference@.md.snap b/cli/tests/cli-reference@.md.snap index 82a27bfc24..9039562ba9 100644 --- a/cli/tests/cli-reference@.md.snap +++ b/cli/tests/cli-reference@.md.snap @@ -163,6 +163,10 @@ repository. Possible values: `true`, `false` * `--config-toml ` — Additional configuration options (can be repeated) +* `--show-heap-stats` — Show memory allocation statistics from the internal heap allocator on `stdout`, when the program exits + + Possible values: `true`, `false` + diff --git a/cli/tests/test_global_opts.rs b/cli/tests/test_global_opts.rs index 3596460ed9..533cc92243 100644 --- a/cli/tests/test_global_opts.rs +++ b/cli/tests/test_global_opts.rs @@ -451,6 +451,8 @@ fn test_help() { --color When to colorize output (always, never, auto) --no-pager Disable the pager --config-toml Additional configuration options (can be repeated) + --show-heap-stats Show memory allocation statistics from the internal heap + allocator on `stdout`, when the program exits "###); }