cli: global --show-heap-stats argument for EOL heap stats

Summary: One of the more useful features of mimalloc. In the future this
interface could be expanded to support printing to a callback, so that it
could be triggered and dumped e.g. on debugging endpoints for longer running
processes, when we get there.

However, it's not enough to just add it to the `main` module; not only would it
be ugly and overspecific, it should be reusable so other custom `jj` commands
can opt-in to statistics trivially as well. Users should be strongly encouraged
to do this -- enable mimalloc and stat tracking -- since that's what upstream
will now work and test with; but, of course, all things are optional.

Thus, we expose an API from the `cli_utils` module for use with
`cli_util::CliRunner::add_global_args`, but ultimately it's up to the client to
call it.

Signed-off-by: Austin Seipp <[email protected]>
Change-Id: I6abe4b962bbe7c62ebca97be630b750c

cli/src/cli_util.rs

@@ -22,14 +22,15 @@ 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::{iter, str};
 
 use clap::builder::{NonEmptyStringValueParser, TypedValueParser, ValueParserFactory};
 use clap::{Arg, ArgAction, ArgMatches, Command, FromArgMatches};
 use indexmap::{IndexMap, IndexSet};
 use itertools::Itertools;
+use jj_cbits::mimalloc;
 use jj_lib::backend::{BackendError, ChangeId, CommitId, MergedTreeId, ObjectId};
 use jj_lib::commit::Commit;
 use jj_lib::git::{
@@ -2586,6 +2587,57 @@ pub struct EarlyArgs {
     pub config_toml: Vec<String>,
 }
 
+#[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<bool> = 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<R>(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

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

cli/tests/test_global_opts.rs

@@ -425,6 +425,8 @@ fn test_help() {
                                  specified
           --from <FROM>          Show changes from this revision. Defaults to @ if --to is specified
           --to <TO>              Edit changes in this revision. Defaults to @ if --from is specified
+          --show-heap-stats      Show memory allocation statistics from the internal heap allocator on
+                                 `stdout`, when the program exits
       -h, --help                 Print help (see more with '--help')
 
     Global Options: