diff --git a/Cargo.lock b/Cargo.lock index 9c5a084a58..bcdb93c025 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4162,6 +4162,17 @@ dependencies = [ "num_cpus", ] +[[package]] +name = "tikv-jemalloc-ctl" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "661f1f6a57b3a36dc9174a2c10f19513b4866816e13425d3e418b11cc37bc24c" +dependencies = [ + "libc", + "paste", + "tikv-jemalloc-sys", +] + [[package]] name = "tikv-jemalloc-sys" version = "0.6.1+5.3.0-1-ge13ca993e8ccb9ba9847cc330696e02839f328f7" @@ -4383,6 +4394,7 @@ dependencies = [ "insta", "insta-cmd", "jiff", + "mimalloc", "rayon", "regex", "ruff_db", @@ -4390,6 +4402,8 @@ dependencies = [ "ruff_python_trivia", "salsa", "tempfile", + "tikv-jemalloc-ctl", + "tikv-jemallocator", "toml", "tracing", "tracing-flame", diff --git a/Cargo.toml b/Cargo.toml index dbd9808fdd..5bb0e368f3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -180,6 +180,7 @@ tempfile = { version = "3.9.0" } test-case = { version = "3.3.1" } thiserror = { version = "2.0.0" } tikv-jemallocator = { version = "0.6.0" } +tikv-jemalloc-ctl = { version = "0.6.0", features = ["stats"] } toml = { version = "0.9.0" } tracing = { version = "0.1.40" } tracing-flame = { version = "0.2.0" } diff --git a/crates/ty/Cargo.toml b/crates/ty/Cargo.toml index 4189758bb1..70bbaffb91 100644 --- a/crates/ty/Cargo.toml +++ b/crates/ty/Cargo.toml @@ -51,5 +51,22 @@ regex = { workspace = true } tempfile = { workspace = true } toml = { workspace = true } +[features] +default = [] +# Prefer mimalloc over jemalloc on platforms that support both +mimalloc = ["dep:mimalloc"] + +# Platform-specific allocator dependencies +# Windows always uses mimalloc +[target.'cfg(target_os = "windows")'.dependencies] +mimalloc = { workspace = true } + +# Non-Windows platforms (except OpenBSD, AIX, Android) on supported architectures +# Use jemalloc by default, mimalloc if the feature is enabled +[target.'cfg(all(not(target_os = "windows"), not(target_os = "openbsd"), not(target_os = "aix"), not(target_os = "android"), any(target_arch = "x86_64", target_arch = "aarch64", target_arch = "powerpc64", target_arch = "riscv64")))'.dependencies] +tikv-jemallocator = { workspace = true } +tikv-jemalloc-ctl = { workspace = true } +mimalloc = { workspace = true, optional = true } + [lints] workspace = true diff --git a/crates/ty/src/allocator.rs b/crates/ty/src/allocator.rs new file mode 100644 index 0000000000..ce8dd04e1f --- /dev/null +++ b/crates/ty/src/allocator.rs @@ -0,0 +1,299 @@ +//! Global allocator configuration for ty. +//! +//! By default: +//! - Windows uses mimalloc +//! - Unix-like platforms (on supported architectures) use jemalloc +//! - Other platforms use the system allocator +//! +//! The `mimalloc` feature can be enabled to prefer mimalloc over jemalloc +//! on platforms that support both. + +use std::fmt::Write; + +// Condition for platforms where we can use either jemalloc or mimalloc +#[cfg(all( + not(target_os = "windows"), + not(target_os = "openbsd"), + not(target_os = "aix"), + not(target_os = "android"), + any( + target_arch = "x86_64", + target_arch = "aarch64", + target_arch = "powerpc64", + target_arch = "riscv64" + ) +))] +mod unix_allocator { + #[cfg(feature = "mimalloc")] + #[global_allocator] + pub(super) static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; + + #[cfg(not(feature = "mimalloc"))] + #[global_allocator] + pub(super) static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; +} + +// Windows always uses mimalloc +#[cfg(target_os = "windows")] +#[global_allocator] +static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; + +/// Returns the name of the allocator currently in use. +#[must_use] +pub(crate) fn allocator_name() -> &'static str { + #[cfg(target_os = "windows")] + { + "mimalloc" + } + + #[cfg(all( + not(target_os = "windows"), + not(target_os = "openbsd"), + not(target_os = "aix"), + not(target_os = "android"), + any( + target_arch = "x86_64", + target_arch = "aarch64", + target_arch = "powerpc64", + target_arch = "riscv64" + ) + ))] + { + #[cfg(feature = "mimalloc")] + { + "mimalloc" + } + + #[cfg(not(feature = "mimalloc"))] + { + "jemalloc" + } + } + + #[cfg(not(any( + target_os = "windows", + all( + not(target_os = "openbsd"), + not(target_os = "aix"), + not(target_os = "android"), + any( + target_arch = "x86_64", + target_arch = "aarch64", + target_arch = "powerpc64", + target_arch = "riscv64" + ) + ) + )))] + { + "system" + } +} + +/// Collects and formats memory usage statistics from the allocator. +/// +/// Returns a formatted string with memory statistics specific to the +/// allocator in use, or `None` if memory statistics are not available +/// for the current allocator. +#[must_use] +pub(crate) fn memory_usage_stats() -> Option { + #[cfg(target_os = "windows")] + { + mimalloc_stats() + } + + #[cfg(all( + not(target_os = "windows"), + not(target_os = "openbsd"), + not(target_os = "aix"), + not(target_os = "android"), + any( + target_arch = "x86_64", + target_arch = "aarch64", + target_arch = "powerpc64", + target_arch = "riscv64" + ) + ))] + { + #[cfg(feature = "mimalloc")] + { + mimalloc_stats() + } + + #[cfg(not(feature = "mimalloc"))] + { + jemalloc_stats() + } + } + + #[cfg(not(any( + target_os = "windows", + all( + not(target_os = "openbsd"), + not(target_os = "aix"), + not(target_os = "android"), + any( + target_arch = "x86_64", + target_arch = "aarch64", + target_arch = "powerpc64", + target_arch = "riscv64" + ) + ) + )))] + { + None + } +} + +/// Collect jemalloc memory statistics +#[cfg(all( + not(target_os = "windows"), + not(target_os = "openbsd"), + not(target_os = "aix"), + not(target_os = "android"), + any( + target_arch = "x86_64", + target_arch = "aarch64", + target_arch = "powerpc64", + target_arch = "riscv64" + ), + not(feature = "mimalloc") +))] +fn jemalloc_stats() -> Option { + use tikv_jemalloc_ctl::{epoch, stats}; + + // Advance the epoch to get fresh statistics + if epoch::advance().is_err() { + return None; + } + + let allocated = stats::allocated::read().ok()?; + let active = stats::active::read().ok()?; + let resident = stats::resident::read().ok()?; + let mapped = stats::mapped::read().ok()?; + let retained = stats::retained::read().ok()?; + let metadata = stats::metadata::read().ok()?; + + let mut output = String::new(); + writeln!(output, "Allocator: jemalloc").ok()?; + writeln!(output, " Allocated: {} ({} bytes)", format_bytes(allocated), allocated).ok()?; + writeln!(output, " Active: {} ({} bytes)", format_bytes(active), active).ok()?; + writeln!(output, " Resident: {} ({} bytes)", format_bytes(resident), resident).ok()?; + writeln!(output, " Mapped: {} ({} bytes)", format_bytes(mapped), mapped).ok()?; + writeln!(output, " Retained: {} ({} bytes)", format_bytes(retained), retained).ok()?; + writeln!(output, " Metadata: {} ({} bytes)", format_bytes(metadata), metadata).ok()?; + writeln!(output).ok()?; + writeln!(output, " Fragmentation: {:.2}%", fragmentation_percent(allocated, resident)).ok()?; + + Some(output) +} + +/// Collect mimalloc memory statistics +#[cfg(any( + target_os = "windows", + all( + not(target_os = "openbsd"), + not(target_os = "aix"), + not(target_os = "android"), + any( + target_arch = "x86_64", + target_arch = "aarch64", + target_arch = "powerpc64", + target_arch = "riscv64" + ), + feature = "mimalloc" + ) +))] +fn mimalloc_stats() -> Option { + // mimalloc doesn't have a simple stats API like jemalloc-ctl + // We can use the heap stats from the default heap + let mut output = String::new(); + writeln!(output, "Allocator: mimalloc").ok()?; + writeln!(output, " (Detailed stats available via MIMALLOC_SHOW_STATS=1 environment variable)").ok()?; + + // Try to get basic heap stats if available + // mimalloc::heap::stats() is not always available, so we provide basic info + writeln!(output).ok()?; + writeln!(output, " Tip: Set MIMALLOC_SHOW_STATS=1 to see detailed allocation statistics on exit").ok()?; + writeln!(output, " Tip: Set MIMALLOC_VERBOSE=1 for even more detailed output").ok()?; + + Some(output) +} + +/// Format bytes in a human-readable format +#[cfg(any( + test, + all( + not(target_os = "windows"), + not(target_os = "openbsd"), + not(target_os = "aix"), + not(target_os = "android"), + any( + target_arch = "x86_64", + target_arch = "aarch64", + target_arch = "powerpc64", + target_arch = "riscv64" + ), + not(feature = "mimalloc") + ) +))] +fn format_bytes(bytes: usize) -> String { + const KB: usize = 1024; + const MB: usize = KB * 1024; + const GB: usize = MB * 1024; + + if bytes >= GB { + format!("{:.2} GB", bytes as f64 / GB as f64) + } else if bytes >= MB { + format!("{:.2} MB", bytes as f64 / MB as f64) + } else if bytes >= KB { + format!("{:.2} KB", bytes as f64 / KB as f64) + } else { + format!("{bytes} B") + } +} + +/// Calculate fragmentation percentage +#[cfg(all( + not(target_os = "windows"), + not(target_os = "openbsd"), + not(target_os = "aix"), + not(target_os = "android"), + any( + target_arch = "x86_64", + target_arch = "aarch64", + target_arch = "powerpc64", + target_arch = "riscv64" + ), + not(feature = "mimalloc") +))] +fn fragmentation_percent(allocated: usize, resident: usize) -> f64 { + if resident == 0 { + 0.0 + } else { + ((resident - allocated) as f64 / resident as f64) * 100.0 + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_allocator_name() { + let name = allocator_name(); + assert!( + name == "jemalloc" || name == "mimalloc" || name == "system", + "Unexpected allocator name: {name}" + ); + } + + #[test] + fn test_format_bytes() { + assert_eq!(format_bytes(0), "0 B"); + assert_eq!(format_bytes(512), "512 B"); + assert_eq!(format_bytes(1024), "1.00 KB"); + assert_eq!(format_bytes(1536), "1.50 KB"); + assert_eq!(format_bytes(1024 * 1024), "1.00 MB"); + assert_eq!(format_bytes(1024 * 1024 * 1024), "1.00 GB"); + } +} diff --git a/crates/ty/src/main.rs b/crates/ty/src/main.rs index 6dbae583fe..2a1c88afe3 100644 --- a/crates/ty/src/main.rs +++ b/crates/ty/src/main.rs @@ -1,9 +1,11 @@ +mod allocator; + use colored::Colorize; use std::io; use ty::{ExitStatus, run}; pub fn main() -> ExitStatus { - run().unwrap_or_else(|error| { + let result = run().unwrap_or_else(|error| { use io::Write; // Use `writeln` instead of `eprintln` to avoid panicking when the stderr pipe is broken. @@ -29,5 +31,27 @@ pub fn main() -> ExitStatus { } ExitStatus::Error - }) + }); + + // Print allocator memory usage if TY_ALLOCATOR_STATS is set + if std::env::var("TY_ALLOCATOR_STATS").is_ok() { + use io::Write; + let mut stderr = io::stderr().lock(); + + if let Some(stats) = allocator::memory_usage_stats() { + writeln!(stderr).ok(); + writeln!(stderr, "{}", "Memory Usage Statistics:".bold()).ok(); + write!(stderr, "{stats}").ok(); + } else { + writeln!(stderr).ok(); + writeln!( + stderr, + "Allocator: {} (no detailed stats available)", + allocator::allocator_name() + ) + .ok(); + } + } + + result }