Skip to content

Commit

Permalink
cli: global --show-heap-stats argument for EOL heap stats
Browse files Browse the repository at this point in the history
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 <[email protected]>
Change-Id: I6abe4b962bbe7c62ebca97be630b750c
  • Loading branch information
thoughtpolice committed Mar 23, 2024
1 parent e97f476 commit 5f30cf9
Show file tree
Hide file tree
Showing 3 changed files with 63 additions and 1 deletion.
58 changes: 57 additions & 1 deletion cli/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
4 changes: 4 additions & 0 deletions cli/tests/[email protected]
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,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`
Expand Down
2 changes: 2 additions & 0 deletions cli/tests/test_global_opts.rs
Original file line number Diff line number Diff line change
Expand Up @@ -481,6 +481,8 @@ fn test_help() {
--color <WHEN> When to colorize output (always, never, auto)
--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
"###);
}

Expand Down

0 comments on commit 5f30cf9

Please sign in to comment.