use core::cell::{Cell, UnsafeCell};
use core::ffi::c_void;
use core::ptr;
use core::sync::atomic::{AtomicBool, AtomicPtr, AtomicUsize, Ordering};
use crate::heap::Heap;
use crate::page::{DelayedList, XFLAG_NEVER, page_collect, page_set_flag};
use crate::segment::{self, Segment};
use crate::{os, prim};
#[repr(C)]
pub struct HeapBox {
pub delayed: DelayedList,
pub next_box: AtomicPtr<HeapBox>,
pub owner_tid: usize,
pub tag: i32,
pub allow_destroy: bool,
pub heap: UnsafeCell<Heap>,
}
const _: () = assert!(core::mem::offset_of!(HeapBox, delayed) == 0);
#[inline]
pub unsafe fn box_of_xheap(xheap: usize) -> *mut HeapBox {
debug_assert!(xheap != 0);
xheap as *mut HeapBox
}
static HEAPS_LOCK: AtomicBool = AtomicBool::new(false);
static HEAPS_HEAD: AtomicPtr<HeapBox> = AtomicPtr::new(ptr::null_mut());
fn heaps_lock() {
while HEAPS_LOCK
.compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
.is_err()
{
core::hint::spin_loop();
}
}
fn heaps_unlock() {
HEAPS_LOCK.store(false, Ordering::Release);
}
fn heaps_register(hb: *mut HeapBox) {
heaps_lock();
unsafe {
(*hb)
.next_box
.store(HEAPS_HEAD.load(Ordering::Relaxed), Ordering::Relaxed);
}
HEAPS_HEAD.store(hb, Ordering::Relaxed);
heaps_unlock();
}
fn heaps_unregister(hb: *mut HeapBox) {
heaps_lock();
let mut cur = HEAPS_HEAD.load(Ordering::Relaxed);
if cur == hb {
unsafe {
HEAPS_HEAD.store((*hb).next_box.load(Ordering::Relaxed), Ordering::Relaxed);
}
} else {
while !cur.is_null() {
unsafe {
let nxt = (*cur).next_box.load(Ordering::Relaxed);
if nxt == hb {
(*cur)
.next_box
.store((*nxt).next_box.load(Ordering::Relaxed), Ordering::Relaxed);
break;
}
cur = nxt;
}
}
}
heaps_unlock();
}
#[cfg(all(target_arch = "x86_64", target_os = "linux", not(miri)))]
mod heap_tls {
use super::HeapBox;
core::arch::global_asm!(
".section .tbss,\"awT\",@nobits",
".globl __ra_tls_heap",
".hidden __ra_tls_heap",
".type __ra_tls_heap,@object",
".p2align 3",
"__ra_tls_heap:",
".zero 8",
".size __ra_tls_heap,8",
".text",
);
#[inline(always)]
fn slot() -> *mut *mut HeapBox {
let off: usize;
unsafe {
core::arch::asm!(
"mov {o}, qword ptr [rip + __ra_tls_heap@GOTTPOFF]",
o = out(reg) off,
options(nostack, preserves_flags, readonly, pure),
);
}
core::ptr::with_exposed_provenance_mut(super::thread_id().wrapping_add(off))
}
#[inline(always)]
pub fn get() -> *mut HeapBox {
let v: *mut HeapBox;
unsafe {
core::arch::asm!(
"mov {t}, qword ptr [rip + __ra_tls_heap@GOTTPOFF]",
"mov {o}, qword ptr fs:[{t}]",
t = out(reg) _,
o = out(reg) v,
options(nostack, preserves_flags, readonly),
);
}
v
}
#[inline(always)]
pub fn set(hb: *mut HeapBox) {
unsafe { slot().write(hb) }
}
}
#[cfg(not(all(target_arch = "x86_64", target_os = "linux", not(miri))))]
mod heap_tls {
use super::HeapBox;
use core::cell::Cell;
use core::ptr;
std::thread_local! {
static HEAP_PTR: Cell<*mut HeapBox> = const { Cell::new(ptr::null_mut()) };
}
#[inline(always)]
pub fn get() -> *mut HeapBox {
HEAP_PTR.with(|c| c.get())
}
#[inline(always)]
pub fn set(hb: *mut HeapBox) {
HEAP_PTR.with(|c| c.set(hb));
}
}
std::thread_local! {
static TID: Cell<usize> = const { Cell::new(0) };
}
#[inline(always)]
pub fn thread_id() -> usize {
#[cfg(all(target_arch = "x86_64", target_os = "linux", not(miri)))]
{
let tp: usize;
unsafe {
core::arch::asm!("mov {}, fs:0", out(reg) tp,
options(nostack, preserves_flags, readonly));
}
tp
}
#[cfg(all(target_arch = "x86_64", target_os = "windows", not(miri)))]
{
let teb: usize;
unsafe {
core::arch::asm!("mov {}, gs:0x30", out(reg) teb,
options(nostack, preserves_flags, readonly));
}
teb
}
#[cfg(all(target_arch = "aarch64", not(miri)))]
{
let tp: usize;
unsafe {
core::arch::asm!("mrs {}, tpidr_el0", out(reg) tp,
options(nostack, preserves_flags));
}
tp
}
#[cfg(not(any(
all(target_arch = "x86_64", target_os = "linux", not(miri)),
all(target_arch = "x86_64", target_os = "windows", not(miri)),
all(target_arch = "aarch64", not(miri))
)))]
{
let t = TID.with(|c| c.get());
if t != 0 { t } else { init_tid() }
}
}
#[cold]
#[cfg_attr(
any(
all(target_arch = "x86_64", target_os = "linux", not(miri)),
all(target_arch = "x86_64", target_os = "windows", not(miri)),
all(target_arch = "aarch64", not(miri))
),
allow(dead_code)
)]
fn init_tid() -> usize {
let t = prim::thread_id();
TID.with(|c| c.set(t));
t
}
pub fn for_each_heap(f: &mut dyn FnMut(&crate::heap::Heap)) {
heaps_lock();
let mut hb = HEAPS_HEAD.load(Ordering::Acquire);
while !hb.is_null() {
unsafe {
let snapshot = core::ptr::read_volatile((*hb).heap.get());
f(&snapshot);
hb = (*hb).next_box.load(Ordering::Acquire);
}
}
heaps_unlock();
}
#[inline]
pub fn heap_box() -> *mut HeapBox {
let hb = heap_tls::get();
if !hb.is_null() {
hb
} else {
init_thread_heap()
}
}
pub fn create_heap(tag: i32, allow_destroy: bool, arena_id: i32) -> *mut HeapBox {
let size = core::mem::size_of::<HeapBox>();
let block = os::alloc_aligned(size, os::page_size(), true, false)
.expect("rusty_alloc: cannot allocate heap");
let hb: *mut HeapBox = block.ptr.cast();
unsafe {
ptr::write(
hb,
HeapBox {
delayed: DelayedList::new(),
next_box: AtomicPtr::new(ptr::null_mut()),
owner_tid: thread_id(),
tag,
allow_destroy,
heap: UnsafeCell::new(Heap::new()),
},
);
(*(*hb).heap.get()).delayed = &raw const (*hb).delayed;
(*(*hb).heap.get()).arena_id = arena_id;
(*(*hb).heap.get()).tag = tag;
(*(*hb).heap.get()).rng.reseed();
let rate = crate::options::get(33).max(0) as usize; let gmin = crate::options::get(30).max(0) as usize; let gmax = crate::options::get(31).max(0) as usize; if gmax > 0 {
(*(*hb).heap.get()).guarded_set_size_bound(gmin, gmax);
(*(*hb).heap.get())
.guarded_set_sample_rate(rate, crate::options::get(34).max(0) as usize);
}
}
heaps_register(hb);
hb
}
#[cold]
fn init_thread_heap() -> *mut HeapBox {
let hb = create_heap(0, false, -1);
heap_tls::set(hb);
BACKING_PTR.with(|c| c.set(hb));
done_slot().set(hb.cast::<c_void>());
hb
}
std::thread_local! {
static BACKING_PTR: Cell<*mut HeapBox> = const { Cell::new(ptr::null_mut()) };
}
pub fn backing_heap() -> *mut HeapBox {
let b = BACKING_PTR.with(|c| c.get());
if !b.is_null() {
b
} else {
let _ = heap_box(); BACKING_PTR.with(|c| c.get())
}
}
pub unsafe fn set_default_heap(hb: *mut HeapBox) -> *mut HeapBox {
let prev = heap_box();
heap_tls::set(hb);
prev
}
fn done_slot() -> prim::TlsSlot {
static RAW: AtomicUsize = AtomicUsize::new(0);
static INIT: AtomicBool = AtomicBool::new(false);
let raw = RAW.load(Ordering::Acquire);
if raw != 0 {
return unsafe { prim::TlsSlot::from_raw(raw - 1) };
}
if INIT
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
.is_ok()
{
let Some(slot) = prim::TlsSlot::new(Some(thread_done_cb)) else {
std::process::abort();
};
RAW.store(slot.into_raw() + 1, Ordering::Release); }
loop {
let raw = RAW.load(Ordering::Acquire);
if raw != 0 {
return unsafe { prim::TlsSlot::from_raw(raw - 1) };
}
core::hint::spin_loop();
}
}
#[cfg(all(windows, not(miri)))]
unsafe extern "system" fn thread_done_cb(v: *const c_void) {
if !v.is_null() {
unsafe { thread_done(v.cast_mut().cast()) };
}
}
#[cfg(any(not(windows), miri))]
unsafe extern "C" fn thread_done_cb(v: *mut c_void) {
if !v.is_null() {
unsafe { thread_done(v.cast()) };
}
}
pub unsafe fn thread_done(hb: *mut HeapBox) {
unsafe {
let h = &mut *(*hb).heap.get();
h.collect(true);
let mut seg = h.segments;
h.segments = ptr::null_mut();
while !seg.is_null() {
let next = (*seg).next;
let end = (*seg).next_free_slice as usize;
let mut idx = segment::HEADER_SLICES;
while idx < end {
let slot = &raw mut (*seg).pages[idx];
let len = ((*slot).slice_count as usize).max(1);
if (*slot).block_size > 0 {
page_collect(slot);
page_set_flag(slot, XFLAG_NEVER);
(*slot).xheap.store(0, Ordering::Release);
}
idx += len;
}
if (*seg).used_pages == 0 {
let _ = segment::segment_free(seg);
h.stats.segments_freed += 1;
} else {
(*seg).thread_id.store(0, Ordering::Release);
abandoned_push(seg);
}
seg = next;
}
let mut hseg = h.huge_segments;
h.huge_segments = ptr::null_mut();
while !hseg.is_null() {
let next = (*hseg).next;
let pg = &raw mut (*hseg).pages[1];
page_collect(pg);
if (*pg).used == 0 {
let _ = segment::huge_free(hseg);
h.stats.segments_freed += 1;
} else {
page_set_flag(pg, XFLAG_NEVER);
(*pg).xheap.store(0, Ordering::Release);
(*hseg).thread_id.store(0, Ordering::Release);
abandoned_push(hseg);
}
hseg = next;
}
let mut b = (*hb).delayed.head.swap(0, Ordering::AcqRel) as *mut crate::page::Block;
while !b.is_null() {
let next = (*b).next;
let seg = segment::segment_of(b.cast());
let pg = segment::page_of(seg, b.cast());
crate::page::remote_free(pg, b);
b = next;
}
heaps_unregister(hb);
let size = core::mem::size_of::<HeapBox>();
let blockdesc = os::OsBlock {
ptr: hb.cast(),
size: os::page_align_up(size),
is_large: false,
is_zero: false,
};
let _ = os::free(blockdesc);
}
heap_tls::set(ptr::null_mut());
}
pub unsafe fn heap_delete(hb: *mut HeapBox) {
let backing = backing_heap();
if hb == backing {
return; }
unsafe {
debug_assert_eq!((*hb).owner_tid, thread_id());
let h = &mut *(*hb).heap.get();
let bh = &mut *(*backing).heap.get();
h.collect(true);
let mut seg = h.segments;
h.segments = ptr::null_mut();
while !seg.is_null() {
let next = (*seg).next;
bh.adopt_segment(seg); seg = next;
}
let mut hseg = h.huge_segments;
h.huge_segments = ptr::null_mut();
while !hseg.is_null() {
let next = (*hseg).next;
let pg = &raw mut (*hseg).pages[1];
(*pg).xheap.store(bh.delayed as usize, Ordering::Release);
(*hseg).next = bh.huge_segments;
bh.huge_segments = hseg;
hseg = next;
}
merge_stats(&mut bh.stats, &h.stats);
release_heap_box(hb);
}
}
pub unsafe fn heap_destroy(hb: *mut HeapBox) {
unsafe {
if !(*hb).allow_destroy {
heap_delete(hb);
return;
}
debug_assert_eq!((*hb).owner_tid, thread_id());
let h = &mut *(*hb).heap.get();
let mut seg = h.segments;
while !seg.is_null() {
let end = (*seg).next_free_slice as usize;
let mut idx = segment::HEADER_SLICES;
while idx < end {
let slot = &raw mut (*seg).pages[idx];
let len = ((*slot).slice_count as usize).max(1);
if (*slot).block_size > 0 {
page_set_flag(slot, XFLAG_NEVER);
(*slot).xheap.store(0, Ordering::Release);
}
idx += len;
}
let next = (*seg).next;
let _ = segment::segment_free(seg);
seg = next;
}
let mut hseg = h.huge_segments;
while !hseg.is_null() {
let pg = &raw mut (*hseg).pages[1];
page_set_flag(pg, XFLAG_NEVER);
(*pg).xheap.store(0, Ordering::Release);
let next = (*hseg).next;
let _ = segment::huge_free(hseg);
hseg = next;
}
release_heap_box(hb);
}
}
fn merge_stats(into: &mut crate::heap::Stats, from: &crate::heap::Stats) {
into.allocs += from.allocs;
into.frees += from.frees;
into.generic += from.generic;
into.pages_fresh += from.pages_fresh;
into.segments += from.segments;
into.huge_allocs += from.huge_allocs;
into.extends += from.extends;
into.large_allocs += from.large_allocs;
into.pages_retired += from.pages_retired;
into.segments_freed += from.segments_freed;
into.realloc_in_place += from.realloc_in_place;
into.realloc_moved += from.realloc_moved;
into.delayed_frees += from.delayed_frees;
into.reclaims += from.reclaims;
}
unsafe fn release_heap_box(hb: *mut HeapBox) {
heaps_unregister(hb);
if heap_tls::get() == hb {
heap_tls::set(backing_heap());
}
let size = core::mem::size_of::<HeapBox>();
let blockdesc = os::OsBlock {
ptr: hb.cast(),
size: os::page_align_up(size),
is_large: false,
is_zero: false,
};
unsafe {
let _ = os::free(blockdesc);
}
}
const MAX_SUBPROCS: usize = 64;
static ABANDONED_LOCK: AtomicBool = AtomicBool::new(false);
static ABANDONED_HEADS: [AtomicPtr<Segment>; MAX_SUBPROCS] =
[const { AtomicPtr::new(ptr::null_mut()) }; MAX_SUBPROCS];
static SUBPROC_NEXT: AtomicUsize = AtomicUsize::new(1);
pub static ABANDONED_COUNT: AtomicUsize = AtomicUsize::new(0);
std::thread_local! {
static SUBPROC: Cell<usize> = const { Cell::new(0) };
}
pub fn subproc_main() -> usize {
0
}
pub fn subproc_new() -> usize {
let id = SUBPROC_NEXT.fetch_add(1, Ordering::AcqRel);
assert!(id < MAX_SUBPROCS, "out of subprocess ids");
id
}
pub fn subproc_delete(id: usize) {
if id == 0 || id >= MAX_SUBPROCS {
return;
}
abandoned_lock();
unsafe {
let mut seg = ABANDONED_HEADS[id].swap(ptr::null_mut(), Ordering::AcqRel);
while !seg.is_null() {
let next = (*seg).next;
(*seg).next = ABANDONED_HEADS[0].load(Ordering::Relaxed);
ABANDONED_HEADS[0].store(seg, Ordering::Relaxed);
seg = next;
}
}
abandoned_unlock();
}
pub fn subproc_add_current_thread(id: usize) {
assert!(id < MAX_SUBPROCS);
SUBPROC.with(|c| c.set(id));
}
fn my_subproc() -> usize {
SUBPROC.with(|c| c.get())
}
pub fn abandoned_visit_blocks(
subproc_id: usize,
heap_tag: i32,
visit_blocks: bool,
f: &mut dyn FnMut(&crate::heap::AreaInfo, *mut u8, usize) -> bool,
) -> bool {
if subproc_id >= MAX_SUBPROCS {
return false;
}
abandoned_lock();
let mut ok = true;
let mut seg = ABANDONED_HEADS[subproc_id].load(Ordering::Acquire);
while !seg.is_null() && ok {
unsafe {
ok = crate::heap::visit_segment_blocks(seg, heap_tag, visit_blocks, false, f);
seg = (*seg).next;
}
}
abandoned_unlock();
ok
}
fn abandoned_lock() {
while ABANDONED_LOCK
.compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
.is_err()
{
core::hint::spin_loop();
}
}
fn abandoned_unlock() {
ABANDONED_LOCK.store(false, Ordering::Release);
}
fn abandoned_push(seg: *mut Segment) {
let sp = my_subproc();
abandoned_lock();
unsafe {
(*seg).next = ABANDONED_HEADS[sp].load(Ordering::Relaxed);
}
ABANDONED_HEADS[sp].store(seg, Ordering::Relaxed);
ABANDONED_COUNT.fetch_add(1, Ordering::Relaxed);
abandoned_unlock();
}
pub fn abandoned_pop() -> *mut Segment {
let sp = my_subproc();
if ABANDONED_HEADS[sp].load(Ordering::Acquire).is_null() {
return ptr::null_mut(); }
abandoned_lock();
let seg = ABANDONED_HEADS[sp].load(Ordering::Relaxed);
if !seg.is_null() {
unsafe {
ABANDONED_HEADS[sp].store((*seg).next, Ordering::Relaxed);
(*seg).next = ptr::null_mut();
(*seg).thread_id.store(thread_id(), Ordering::Release);
}
ABANDONED_COUNT.fetch_sub(1, Ordering::Relaxed);
}
abandoned_unlock();
seg
}