use core::cell::Cell;
use core::ffi::c_void;
use core::sync::atomic::{AtomicBool, AtomicPtr, Ordering};
#[cfg(target_has_atomic = "64")]
use core::sync::atomic::{AtomicI64, AtomicU64};
#[cfg(all(not(target_has_atomic = "64"), feature = "std"))]
use portable_atomic::{AtomicI64, AtomicU64};
#[cfg(all(not(target_has_atomic = "64"), not(feature = "std")))]
use split64::{AtomicI64, AtomicU64};
#[cfg(any(all(not(target_has_atomic = "64"), not(feature = "std")), test))]
mod split64 {
use core::sync::atomic::{AtomicU32, Ordering};
const fn split(v: u64) -> (u32, u32) {
(v as u32, (v >> 32) as u32)
}
const fn join(lo: u32, hi: u32) -> u64 {
((hi as u64) << 32) | lo as u64
}
#[derive(Debug)]
pub struct AtomicU64 {
lo: AtomicU32,
hi: AtomicU32,
}
impl AtomicU64 {
pub const fn new(v: u64) -> Self {
let (lo, hi) = split(v);
Self {
lo: AtomicU32::new(lo),
hi: AtomicU32::new(hi),
}
}
pub fn load(&self, _ord: Ordering) -> u64 {
join(
self.lo.load(Ordering::Acquire),
self.hi.load(Ordering::Acquire),
)
}
pub fn store(&self, v: u64, _ord: Ordering) {
let (lo, hi) = split(v);
self.lo.store(lo, Ordering::Release);
self.hi.store(hi, Ordering::Release);
}
pub fn fetch_add(&self, v: u64, ord: Ordering) -> u64 {
let prev = self.load(ord);
self.store(prev.wrapping_add(v), ord);
prev
}
}
#[derive(Debug)]
pub struct AtomicI64(AtomicU64);
impl AtomicI64 {
pub const fn new(v: i64) -> Self {
Self(AtomicU64::new(v as u64))
}
pub fn load(&self, ord: Ordering) -> i64 {
self.0.load(ord) as i64
}
pub fn store(&self, v: i64, ord: Ordering) {
self.0.store(v as u64, ord);
}
pub fn compare_exchange(
&self,
current: i64,
new: i64,
success: Ordering,
_failure: Ordering,
) -> Result<i64, i64> {
let seen = self.load(success);
if seen == current {
self.store(new, success);
Ok(seen)
} else {
Err(seen)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn every_ordering_options_uses_is_accepted() {
let v = AtomicI64::new(i64::MIN);
v.store(-1, Ordering::Release);
assert_eq!(v.load(Ordering::Acquire), -1);
assert_eq!(
v.compare_exchange(-1, 7, Ordering::AcqRel, Ordering::Acquire),
Ok(-1)
);
assert_eq!(v.load(Ordering::Acquire), 7);
assert_eq!(
v.compare_exchange(-1, 9, Ordering::AcqRel, Ordering::Acquire),
Err(7)
);
let u = AtomicU64::new(u64::from(u32::MAX));
assert_eq!(u.fetch_add(1, Ordering::Relaxed), u64::from(u32::MAX));
assert_eq!(u.load(Ordering::Relaxed), 1u64 << 32);
}
}
}
pub const OPTION_COUNT: usize = 38;
pub const GENERIC_COLLECT: usize = 36;
const _: () = assert!(GENERIC_COLLECT < OPTION_COUNT);
#[cfg(not(ra_small_profile))]
pub const GENERIC_COLLECT_DEFAULT: i64 = 10_000;
#[cfg(ra_small_profile)]
pub const GENERIC_COLLECT_DEFAULT: i64 = if cfg!(ra_generic_collect = "64") {
64
} else if cfg!(ra_generic_collect = "4096") {
4096
} else if cfg!(ra_generic_collect = "65536") {
65_536
} else {
512
};
pub const OPTION_NAMES: [&str; OPTION_COUNT] = [
"show_errors",
"show_stats",
"verbose",
"eager_commit",
"arena_eager_commit",
"purge_decommits",
"allow_large_os_pages",
"reserve_huge_os_pages",
"reserve_huge_os_pages_at",
"reserve_os_memory",
"deprecated_segment_cache",
"deprecated_page_reset",
"abandoned_page_purge",
"deprecated_segment_reset",
"eager_commit_delay",
"purge_delay",
"use_numa_nodes",
"disallow_os_alloc",
"os_tag",
"max_errors",
"max_warnings",
"max_segment_reclaim",
"destroy_on_exit",
"arena_reserve",
"arena_purge_mult",
"purge_extend_delay",
"abandoned_reclaim_on_free",
"disallow_arena_alloc",
"retry_on_oom",
"visit_abandoned",
"guarded_min",
"guarded_max",
"guarded_precise",
"guarded_sample_rate",
"guarded_sample_seed",
"target_segments_per_thread",
"generic_collect",
"allow_thp",
];
const DEFAULTS: [i64; OPTION_COUNT] = [
0, 0, 0, 1, 2, 1, 0, 0, -1, 0, 0,
0,
0,
1, 1, -1, 0, 0, 100, 32, 32, 10, 0, 1_048_576, 10, 1, 1, 0, 400, 0, 0,
0,
0, 1000, 0, 0, GENERIC_COLLECT_DEFAULT, 1, ];
static VALUES: [AtomicI64; OPTION_COUNT] = [const { AtomicI64::new(i64::MIN) }; OPTION_COUNT];
static ENV_PARSED: AtomicBool = AtomicBool::new(false);
#[inline(always)]
fn ensure_init() {
if !ENV_PARSED.load(Ordering::Acquire) {
ensure_init_slow();
}
}
#[cold]
#[inline(never)]
fn ensure_init_slow() {
if ENV_PARSED.swap(true, Ordering::AcqRel) {
return;
}
for i in 0..OPTION_COUNT {
VALUES[i].store(DEFAULTS[i], Ordering::Release);
}
#[cfg(all(
feature = "std",
not(all(target_arch = "wasm32", target_os = "unknown"))
))]
env::pass();
}
#[cfg(all(
feature = "std",
not(all(target_arch = "wasm32", target_os = "unknown"))
))]
mod env {
use core::sync::atomic::Ordering;
use super::{OPTION_COUNT, OPTION_NAMES, VALUES};
const MAX_NAME_LEN: usize = {
let mut m = 0;
let mut i = 0;
while i < OPTION_COUNT {
if OPTION_NAMES[i].len() > m {
m = OPTION_NAMES[i].len();
}
i += 1;
}
m
};
const KEY_CAP: usize = "RUSTY_ALLOC_".len() + MAX_NAME_LEN + 1;
const VALUE_CAP: usize = 64;
pub(super) fn pass() {
#[cfg(all(target_os = "linux", not(miri)))]
{
scan_pass();
}
#[cfg(not(all(target_os = "linux", not(miri))))]
{
keyed_pass();
}
}
#[cfg(all(target_os = "linux", not(miri)))]
fn scan_pass() {
let mut rusty = [core::ptr::null::<u8>(); OPTION_COUNT];
let mut mi = [core::ptr::null::<u8>(); OPTION_COUNT];
crate::prim::env_for_each(|e| {
unsafe {
let (table, rest) = if let Some(r) = strip(e, b"RUSTY_ALLOC_") {
(&mut rusty, r)
} else if let Some(r) = strip(e, b"MIMALLOC_") {
(&mut mi, r)
} else {
return;
};
for i in 0..OPTION_COUNT {
if let Some(v) = match_name(rest, OPTION_NAMES[i].as_bytes()) {
if table[i].is_null() {
table[i] = v;
}
break;
}
}
}
});
let mut val = [0u8; VALUE_CAP];
for i in 0..OPTION_COUNT {
let got =
unsafe { copy_value(rusty[i], &mut val).or_else(|| copy_value(mi[i], &mut val)) };
if let Some(v) = got.and_then(|len| parse_value(&val[..len])) {
VALUES[i].store(v, Ordering::Release);
}
}
}
#[cfg(all(target_os = "linux", not(miri)))]
unsafe fn strip(e: *const u8, prefix: &[u8]) -> Option<*const u8> {
for (k, &b) in prefix.iter().enumerate() {
if unsafe { *e.add(k) } != b {
return None;
}
}
Some(unsafe { e.add(prefix.len()) })
}
#[cfg(all(target_os = "linux", not(miri)))]
unsafe fn match_name(rest: *const u8, name: &[u8]) -> Option<*const u8> {
for (k, &b) in name.iter().enumerate() {
if unsafe { *rest.add(k) } != b.to_ascii_uppercase() {
return None;
}
}
unsafe { (*rest.add(name.len()) == b'=').then(|| rest.add(name.len() + 1)) }
}
#[cfg(all(target_os = "linux", not(miri)))]
unsafe fn copy_value(v: *const u8, out: &mut [u8]) -> Option<usize> {
if v.is_null() {
return None;
}
for (n, slot) in out.iter_mut().enumerate() {
let b = unsafe { *v.add(n) };
if b == 0 {
return Some(n);
}
*slot = b;
}
None }
#[cfg_attr(all(target_os = "linux", not(miri)), allow(dead_code))]
fn keyed_pass() {
let mut key = [0u8; KEY_CAP];
let mut val = [0u8; VALUE_CAP];
for i in 0..OPTION_COUNT {
let name = OPTION_NAMES[i].as_bytes();
let n = build_key(&mut key, b"RUSTY_ALLOC_", name);
let got = lookup(&key[..n], &mut val).or_else(|| {
let n = build_key(&mut key, b"MIMALLOC_", name);
lookup(&key[..n], &mut val)
});
if let Some(v) = got.and_then(|len| parse_value(&val[..len])) {
VALUES[i].store(v, Ordering::Release);
}
}
}
fn build_key(key: &mut [u8; KEY_CAP], prefix: &[u8], name: &[u8]) -> usize {
let mut n = 0;
for &b in prefix {
key[n] = b;
n += 1;
}
for &b in name {
key[n] = b.to_ascii_uppercase();
n += 1;
}
key[n] = 0;
n + 1
}
#[cfg(all(any(windows, unix), not(miri)))]
fn lookup(key: &[u8], out: &mut [u8]) -> Option<usize> {
crate::prim::getenv(key, out)
}
#[cfg(not(all(any(windows, unix), not(miri))))]
fn lookup(key: &[u8], out: &mut [u8]) -> Option<usize> {
let name = core::str::from_utf8(key.split_last()?.1).ok()?;
let v = std::env::var_os(name)?;
let v = v.to_str()?.as_bytes();
if v.len() >= out.len() {
return None;
}
out[..v.len()].copy_from_slice(v);
Some(v.len())
}
fn parse_value(v: &[u8]) -> Option<i64> {
let t = core::str::from_utf8(v).ok()?.trim();
if ["", "1", "true", "yes", "on"]
.iter()
.any(|w| t.eq_ignore_ascii_case(w))
{
return Some(1);
}
if ["0", "false", "no", "off"]
.iter()
.any(|w| t.eq_ignore_ascii_case(w))
{
return Some(0);
}
t.parse::<i64>().ok()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn keys_are_the_uppercased_names_with_a_nul_and_always_fit() {
let mut key = [0u8; KEY_CAP];
let n = build_key(&mut key, b"RUSTY_ALLOC_", b"purge_delay");
assert_eq!(&key[..n], b"RUSTY_ALLOC_PURGE_DELAY\0");
let n = build_key(&mut key, b"MIMALLOC_", b"show_stats");
assert_eq!(&key[..n], b"MIMALLOC_SHOW_STATS\0");
for name in OPTION_NAMES {
let n = build_key(&mut key, b"RUSTY_ALLOC_", name.as_bytes());
assert!(n <= KEY_CAP, "{name}");
assert_eq!(key[n - 1], 0);
}
}
#[test]
fn values_follow_mimalloc_grammar() {
assert_eq!(parse_value(b""), Some(1));
assert_eq!(parse_value(b" On "), Some(1));
assert_eq!(parse_value(b"TRUE"), Some(1));
assert_eq!(parse_value(b"off"), Some(0));
assert_eq!(parse_value(b"No"), Some(0));
assert_eq!(parse_value(b" -1 "), Some(-1));
assert_eq!(parse_value(b"1048576"), Some(1_048_576));
assert_eq!(parse_value(b"maybe"), None);
assert_eq!(parse_value(&[0xff, 0xfe]), None);
}
#[test]
fn lookup_reads_the_process_environment_in_place() {
unsafe {
std::env::set_var("RUSTY_ALLOC_TEST_PROBE", " 42 ");
std::env::set_var("RUSTY_ALLOC_TEST_LONG_PROBE", "x".repeat(VALUE_CAP));
}
let mut out = [0u8; VALUE_CAP];
let n = lookup(b"RUSTY_ALLOC_TEST_PROBE\0", &mut out).expect("set");
assert_eq!(&out[..n], b" 42 ");
assert_eq!(parse_value(&out[..n]), Some(42));
assert_eq!(lookup(b"RUSTY_ALLOC_TEST_PROBE_UNSET\0", &mut out), None);
assert_eq!(lookup(b"RUSTY_ALLOC_TEST_LONG_PROBE\0", &mut out), None);
}
}
}
pub fn get(option: usize) -> i64 {
if option >= OPTION_COUNT {
return 0;
}
if crate::ONE_REGION {
return DEFAULTS[option];
}
ensure_init();
let v = VALUES[option].load(Ordering::Acquire);
if v == i64::MIN { DEFAULTS[option] } else { v }
}
pub fn set(option: usize, value: i64) {
if crate::ONE_REGION {
return;
}
if option < OPTION_COUNT {
ensure_init();
VALUES[option].store(value, Ordering::Release);
}
}
pub fn set_default(option: usize, value: i64) {
if crate::ONE_REGION {
return;
}
if option < OPTION_COUNT {
ensure_init();
let _ = VALUES[option].compare_exchange(
DEFAULTS[option],
value,
Ordering::AcqRel,
Ordering::Acquire,
);
}
}
pub fn is_enabled(option: usize) -> bool {
get(option) != 0
}
pub fn get_clamp(option: usize, min: i64, max: i64) -> i64 {
get(option).clamp(min.min(max), min.max(max))
}
pub fn get_size(option: usize) -> usize {
let v = get(option).max(0) as usize;
match option {
9 | 23 => v.saturating_mul(1024),
_ => v,
}
}
#[cfg(feature = "std")]
pub fn print() {
ensure_init();
for (i, name) in OPTION_NAMES.iter().enumerate() {
out_fmt(&format!("option '{name}': {}\n", get(i)));
}
}
pub type OutputFun = unsafe extern "C" fn(msg: *const core::ffi::c_char, arg: *mut c_void);
pub type ErrorFun = unsafe extern "C" fn(err: i32, arg: *mut c_void);
pub type DeferredFreeFun = unsafe extern "C" fn(force: bool, heartbeat: u64, arg: *mut c_void);
static OUTPUT_FUN: AtomicUsize2 = AtomicUsize2::new();
static ERROR_FUN: AtomicUsize2 = AtomicUsize2::new();
static DEFERRED_FUN: AtomicUsize2 = AtomicUsize2::new();
static HEARTBEAT: AtomicU64 = AtomicU64::new(0);
ra_thread_local! {
static IN_DEFERRED: Cell<bool> = const { Cell::new(false) };
static IN_ERROR: Cell<bool> = const { Cell::new(false) };
static IN_OUTPUT: Cell<bool> = const { Cell::new(false) };
}
struct HookExit(fn());
impl Drop for HookExit {
fn drop(&mut self) {
(self.0)();
}
}
struct AtomicUsize2 {
f: AtomicPtr<c_void>,
a: AtomicPtr<c_void>,
}
impl AtomicUsize2 {
#[inline]
fn load_fun(&self) -> *mut c_void {
self.f.load(Ordering::Acquire)
}
const fn new() -> Self {
AtomicUsize2 {
f: AtomicPtr::new(core::ptr::null_mut()),
a: AtomicPtr::new(core::ptr::null_mut()),
}
}
fn set(&self, f: *mut c_void, a: *mut c_void) {
self.a.store(a, Ordering::Release);
self.f.store(f, Ordering::Release);
}
fn load(&self) -> (*mut c_void, *mut c_void) {
(
self.f.load(Ordering::Acquire),
self.a.load(Ordering::Acquire),
)
}
}
pub fn register_output(f: Option<OutputFun>, arg: *mut c_void) {
OUTPUT_FUN.set(f.map_or(core::ptr::null_mut(), |f| f as *mut c_void), arg);
}
pub fn register_error(f: Option<ErrorFun>, arg: *mut c_void) {
ERROR_FUN.set(f.map_or(core::ptr::null_mut(), |f| f as *mut c_void), arg);
}
pub fn register_deferred_free(f: Option<DeferredFreeFun>, arg: *mut c_void) {
DEFERRED_FUN.set(f.map_or(core::ptr::null_mut(), |f| f as *mut c_void), arg);
}
pub fn out_fmt(msg: &str) {
let (f, a) = OUTPUT_FUN.load();
if f.is_null() {
#[cfg(feature = "std")]
{
use std::io::Write;
let _ = std::io::stderr().write_all(msg.as_bytes());
}
return;
}
if IN_OUTPUT.with(|c| c.get()) {
return;
}
IN_OUTPUT.with(|c| c.set(true));
let _clear = HookExit(|| IN_OUTPUT.with(|c| c.set(false)));
let bytes = msg.as_bytes();
let mut buf = [0u8; 512];
let n = bytes.len().min(511);
buf[..n].copy_from_slice(&bytes[..n]);
unsafe {
let fun: OutputFun = core::mem::transmute::<*mut c_void, OutputFun>(f);
fun(buf.as_ptr().cast(), a);
}
}
fn render_error(buf: &mut [u8; 32], err: i32) -> &str {
const PREFIX: &[u8] = b"rusty_alloc: error ";
buf[..PREFIX.len()].copy_from_slice(PREFIX);
let mut n = PREFIX.len();
if err < 0 {
buf[n] = b'-';
n += 1;
}
let mut v = err.unsigned_abs();
let mut digits = [0u8; 10];
let mut d = 0;
loop {
digits[d] = b'0' + (v % 10) as u8;
d += 1;
v /= 10;
if v == 0 {
break;
}
}
while d > 0 {
d -= 1;
buf[n] = digits[d];
n += 1;
}
buf[n] = b'\n';
n += 1;
core::str::from_utf8(&buf[..n]).unwrap_or(
"rusty_alloc: error
",
)
}
pub fn deferred_free(force: bool) {
if DEFERRED_FUN.load_fun().is_null() {
return;
}
fire_deferred(force);
}
#[cold]
#[inline(never)]
fn fire_deferred(force: bool) {
if IN_DEFERRED.with(|c| c.get()) {
return;
}
IN_DEFERRED.with(|c| c.set(true));
let _clear = HookExit(|| IN_DEFERRED.with(|c| c.set(false)));
let (f, a) = DEFERRED_FUN.load();
if !f.is_null() {
let hb = HEARTBEAT.fetch_add(1, Ordering::Relaxed);
unsafe {
let fun: DeferredFreeFun = core::mem::transmute::<*mut c_void, DeferredFreeFun>(f);
fun(force, hb, a);
}
}
}
pub fn error(err: i32) {
let (f, a) = ERROR_FUN.load();
if !f.is_null() {
if IN_ERROR.with(|c| c.get()) {
return;
}
IN_ERROR.with(|c| c.set(true));
let _clear = HookExit(|| IN_ERROR.with(|c| c.set(false)));
unsafe {
let fun: ErrorFun = core::mem::transmute::<*mut c_void, ErrorFun>(f);
fun(err, a);
}
} else if is_enabled(0) {
let mut buf = [0u8; 32];
out_fmt(render_error(&mut buf, err));
}
}
#[cfg(test)]
mod render_error_tests {
use super::render_error;
#[test]
fn renders_every_shape_without_a_formatter() {
let mut b = [0u8; 32];
assert_eq!(render_error(&mut b, 0), "rusty_alloc: error 0\n");
assert_eq!(render_error(&mut b, 7), "rusty_alloc: error 7\n");
assert_eq!(render_error(&mut b, 12345), "rusty_alloc: error 12345\n");
assert_eq!(render_error(&mut b, -1), "rusty_alloc: error -1\n");
assert_eq!(
render_error(&mut b, i32::MAX),
"rusty_alloc: error 2147483647\n"
);
assert_eq!(
render_error(&mut b, i32::MIN),
"rusty_alloc: error -2147483648\n"
);
}
}