Python
/
ruff
mirror of https://github.com/astral-sh/ruff
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add custom allocator support to ty with memory stats reporting 

- Copy Ruff's platform-specific allocator configuration to ty:
  - Windows: uses mimalloc
  - Unix-like (x86_64, aarch64, powerpc64, riscv64): uses jemalloc by default
  - Other platforms: uses system allocator

- Add `mimalloc` feature flag to prefer mimalloc over jemalloc on
  platforms that support both allocators

- Add allocator memory usage statistics:
  - Set TY_ALLOCATOR_STATS=1 to print memory stats on exit
  - jemalloc: shows allocated, active, resident, mapped, retained,
    metadata bytes and fragmentation percentage
  - mimalloc: provides guidance for using MIMALLOC_SHOW_STATS=1

- Add tikv-jemalloc-ctl workspace dependency with stats feature
This commit is contained in:
Claude 2025-12-13 08:30:47 +00:00
parent e2ec2bc306
commit ebab078df6
No known key found for this signature in database
5 changed files with 357 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
Cargo.lock generated
View File

@ -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",

1
Cargo.toml
View File

@ -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" }

17
crates/ty/Cargo.toml
View File

@ -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

299
crates/ty/src/allocator.rs Normal file
View File

@ -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<String> {
#[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<String> {
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<String> {
// 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");
}
}

28
crates/ty/src/main.rs
View File

@ -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
}