From 79bb7733f28304df84fe8872c9710ff6701873a5 Mon Sep 17 00:00:00 2001
From: Austin Seipp <aseipp@pobox.com>
Date: Wed, 1 Nov 2023 22:48:43 -0500
Subject: [PATCH] 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.

Consolidating this into `main.rs` is generally the right choice I think, since
users of `jj_cli` and related crates will have to write their own `main.rs`
anyway.

Signed-off-by: Austin Seipp <aseipp@pobox.com>
Change-Id: I6abe4b962bbe7c62ebca97be630b750c
---
 cli/src/main.rs                  | 58 +++++++++++++++++++++++++++++++-
 cli/tests/cli-reference@.md.snap |  4 +++
 cli/tests/test_global_opts.rs    |  2 ++
 3 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/cli/src/main.rs b/cli/src/main.rs
index 725aa9ee4e..ccbc0edada 100644
--- a/cli/src/main.rs
+++ b/cli/src/main.rs
@@ -12,14 +12,70 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::sync::OnceLock;
+
 #[cfg(feature = "mimalloc")]
 use jj_cbits::mimalloc::MiMalloc;
 use jj_cli::cli_util::CliRunner;
+use jj_cli::command_error::CommandError;
+
+/// 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();
+
+#[derive(clap::Args, Clone, Debug)]
+pub struct ShowAllocStats {
+    /// Show memory allocation statistics from the internal heap allocator
+    /// on `stderr`, when the program exits.
+    #[arg(long, global = true)]
+    show_heap_stats: bool,
+}
+
+/// 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 jj_cli::ui::Ui,
+    opts: ShowAllocStats,
+) -> Result<(), CommandError> {
+    if opts.show_heap_stats {
+        PRINT_HEAP_STATS.set(true).unwrap();
+    }
+    Ok(())
+}
 
 #[cfg(feature = "mimalloc")]
 #[global_allocator]
 static ALLOC: MiMalloc = MiMalloc;
 
 fn main() -> std::process::ExitCode {
-    CliRunner::init().version(env!("JJ_VERSION")).run()
+    let result = CliRunner::init()
+        // NOTE (aseipp): always attach heap_stats_enable here, even if compiled
+        // without mimalloc; we don't want the test suite or other users to have
+        // to worry about if the command exists or not based on the build
+        // configuration
+        .add_global_args(heap_stats_enable)
+        .version(env!("JJ_VERSION"))
+        .run();
+
+    if PRINT_HEAP_STATS.get() == Some(&true) {
+        #[cfg(feature = "mimalloc")]
+        {
+            // NOTE (aseipp): can we do our own custom printing here? it's kind of ugly
+            eprintln!("========================================");
+            eprintln!("mimalloc memory allocation statistics:\n");
+            jj_cbits::mimalloc::stats_print(&|l| {
+                eprint!("{}", l.to_string_lossy());
+            });
+        }
+
+        #[cfg(not(feature = "mimalloc"))]
+        {
+            eprintln!(
+                "Note: heap statistics requested, but custom memory allocator (mimalloc) is not \
+                 enabled."
+            );
+        }
+    }
+    result
 }
diff --git a/cli/tests/cli-reference@.md.snap b/cli/tests/cli-reference@.md.snap
index 1649395938..a4ef95e4cb 100644
--- a/cli/tests/cli-reference@.md.snap
+++ b/cli/tests/cli-reference@.md.snap
@@ -168,6 +168,10 @@ To get started, see the tutorial at https://github.com/martinvonz/jj/blob/main/d
   Possible values: `true`, `false`
 
 * `--config-toml <TOML>` — Additional configuration options (can be repeated)
+* `--show-heap-stats` — Show memory allocation statistics from the internal heap allocator on `stderr`, 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 cf358b9b90..76ee1b000e 100644
--- a/cli/tests/test_global_opts.rs
+++ b/cli/tests/test_global_opts.rs
@@ -551,6 +551,8 @@ fn test_help() {
           --quiet                        Silence non-primary command output
           --no-pager                     Disable the pager
           --config-toml <TOML>           Additional configuration options (can be repeated)
+          --show-heap-stats              Show memory allocation statistics from the internal heap
+                                         allocator on `stderr`, when the program exits
     "###);
 }