mirror of
https://github.com/bspeice/qadapt
synced 2024-11-15 02:28:11 -05:00
342 lines
9.8 KiB
Rust
342 lines
9.8 KiB
Rust
//! ## `debug_assert!` for your memory usage
|
|
//!
|
|
//! This allocator is a helper for writing high-performance code that is memory-sensitive;
|
|
//! a thread panic will be triggered if a function annotated with `#[no_alloc]`,
|
|
//! or code inside an `assert_no_alloc!` macro interacts with the allocator in any way.
|
|
//! Wanton allocations and unforeseen drops no more - this library lets you focus on
|
|
//! writing code without worrying if Rust properly managed to inline the variable into the stack.
|
|
//!
|
|
//! Now, an allocator blowing up in production is a scary thought; that's why QADAPT
|
|
//! is designed to strip its own code out whenever you're running with a release build.
|
|
//! Just like the [`debug_assert!` macro](https://doc.rust-lang.org/std/macro.debug_assert.html)
|
|
//! in Rust's standard library, it's safe to use without worrying about a unforeseen
|
|
//! circumstance causing your application to crash.
|
|
//!
|
|
//! # Usage
|
|
//!
|
|
//! Actually making use of QADAPT is straight-forward. To set up the allocator,
|
|
//! place the following snippet in either your program binaries (main.rs) or tests:
|
|
//!
|
|
//! ```rust,ignore
|
|
//! use qadapt::QADAPT;
|
|
//!
|
|
//! #[global_allocator]
|
|
//! static Q: QADAPT = QADAPT;
|
|
//! ```
|
|
//!
|
|
//! After that, there are two ways of telling QADAPT that it should trigger a panic:
|
|
//!
|
|
//! 1. Annotate functions with the `#[no_alloc]` proc macro:
|
|
//! ```rust,no_run
|
|
//! use qadapt::no_alloc;
|
|
//!
|
|
//! // This function is fine, there are no allocations here
|
|
//! #[no_alloc]
|
|
//! fn do_math() -> u8 {
|
|
//! 2 + 2
|
|
//! }
|
|
//!
|
|
//! // This function will trigger a panic when called
|
|
//! #[no_alloc]
|
|
//! fn does_panic() -> Box<u32> {
|
|
//! Box::new(5)
|
|
//! }
|
|
//!
|
|
//! fn main() {
|
|
//! do_math();
|
|
//! does_panic();
|
|
//! }
|
|
//! ```
|
|
//!
|
|
//! 2. Evaluate expressions with the `assert_no_alloc!` macro
|
|
//! ```rust,no_run
|
|
//! use qadapt::assert_no_alloc;
|
|
//!
|
|
//! fn do_work() {
|
|
//! // This code is allowed to trigger an allocation
|
|
//! let b = Box::new(8);
|
|
//!
|
|
//! // This code would panic if an allocation occurred inside it
|
|
//! let x = assert_no_alloc!(*b + 2);
|
|
//! assert_eq!(x, 10);
|
|
//! }
|
|
#![deny(missing_docs)]
|
|
|
|
// thread_id is necessary because `std::thread::current()` panics if we have not yet
|
|
// allocated a `thread_local!{}` it depends on.
|
|
use thread_id;
|
|
|
|
// Re-export the proc macros to use by other code
|
|
pub use qadapt_macro::*;
|
|
|
|
use spin::RwLock;
|
|
use std::alloc::GlobalAlloc;
|
|
use std::alloc::Layout;
|
|
use std::alloc::System;
|
|
use std::thread;
|
|
|
|
thread_local! {
|
|
static PROTECTION_LEVEL: RwLock<usize> = RwLock::new(0);
|
|
}
|
|
|
|
/// The QADAPT allocator itself
|
|
///
|
|
/// To make use of the allocator, include this code block in your program
|
|
/// binaries/tests:
|
|
///
|
|
/// ```rust,ignore
|
|
/// use qadapt::QADAPT;
|
|
///
|
|
/// #[global_allocator]
|
|
/// static Q: QADAPT = QADAPT;
|
|
/// ```
|
|
pub struct QADAPT;
|
|
|
|
static SYSTEM_ALLOC: System = System;
|
|
|
|
/// Let QADAPT know that we are now entering a protected region and that
|
|
/// panics should be triggered if allocations/drops happen while we are running.
|
|
///
|
|
/// **Example**:
|
|
///
|
|
/// ```rust,no_run
|
|
/// use qadapt::enter_protected;
|
|
/// use qadapt::exit_protected;
|
|
///
|
|
/// fn main() {
|
|
/// // Force an allocation by using a Box
|
|
/// let x = Box::new(2);
|
|
///
|
|
/// enter_protected();
|
|
/// // We're now in a memory-protected region - allocations and drops
|
|
/// // here will trigger thread panic
|
|
/// let y = *x * 4;
|
|
/// exit_protected();
|
|
///
|
|
/// // It's now safe to allocate/drop again
|
|
/// let z = Box::new(y);
|
|
/// }
|
|
pub fn enter_protected() {
|
|
#[cfg(debug_assertions)]
|
|
{
|
|
if thread::panicking() {
|
|
return;
|
|
}
|
|
|
|
if !*IS_ACTIVE.read() {
|
|
panic!("QADAPT not initialized when using allocation guards; please verify `#[global_allocator]` is set!");
|
|
}
|
|
|
|
PROTECTION_LEVEL
|
|
.try_with(|v| {
|
|
*v.write() += 1;
|
|
})
|
|
.unwrap_or_else(|_e| ());
|
|
}
|
|
}
|
|
|
|
/// Let QADAPT know that we are exiting a protected region. Will panic
|
|
/// if we attempt to [`exit_protected`] more times than we [`enter_protected`].
|
|
///
|
|
/// **Example**:
|
|
///
|
|
/// ```rust,no_run
|
|
/// use qadapt::enter_protected;
|
|
/// use qadapt::exit_protected;
|
|
///
|
|
/// fn main() {
|
|
/// // Force an allocation by using a Box
|
|
/// let x = Box::new(2);
|
|
///
|
|
/// enter_protected();
|
|
/// // We're now in a memory-protected region - allocations and drops
|
|
/// // here will trigger thread panic
|
|
/// let y = *x * 4;
|
|
/// exit_protected();
|
|
///
|
|
/// // It's now safe to allocate/drop again
|
|
/// let z = Box::new(y);
|
|
/// }
|
|
pub fn exit_protected() {
|
|
#[cfg(debug_assertions)]
|
|
{
|
|
if thread::panicking() {
|
|
return;
|
|
}
|
|
|
|
PROTECTION_LEVEL
|
|
.try_with(|v| {
|
|
let val = { *v.read() };
|
|
match val {
|
|
v if v == 0 => panic!("Attempt to exit protected too many times"),
|
|
_ => {
|
|
*v.write() -= 1;
|
|
}
|
|
}
|
|
})
|
|
.unwrap_or_else(|_e| ());
|
|
}
|
|
}
|
|
|
|
/// Get the result of an expression, guaranteeing that no memory accesses occur
|
|
/// during its evaluation.
|
|
///
|
|
/// **Example**:
|
|
///
|
|
/// ```rust,no_run
|
|
/// use qadapt::assert_no_alloc;
|
|
///
|
|
/// fn main() {
|
|
/// assert_no_alloc!(2 + 2);
|
|
/// }
|
|
/// ```
|
|
///
|
|
/// **Warning**: Unexpected behavior will occur when using the `return` keyword.
|
|
/// Because QADAPT doesn't have an opportunity to clean up, there may be a panic
|
|
/// in code that was not intended to be allocation-free. The compiler will warn you
|
|
/// that there is an unreachable statement if this happens.
|
|
///
|
|
/// ```rust,no_run
|
|
/// use qadapt::assert_no_alloc;
|
|
///
|
|
/// fn early_return() -> usize {
|
|
/// assert_no_alloc!(return 8);
|
|
/// }
|
|
///
|
|
/// fn main() {
|
|
/// let x = early_return();
|
|
///
|
|
/// // This triggers a panic - `Box::new` forces an allocation,
|
|
/// // and QADAPT still thinks we're in a protected region because
|
|
/// // of a return in the `early_return()` function
|
|
/// let b = Box::new(x);
|
|
/// }
|
|
#[macro_export]
|
|
macro_rules! assert_no_alloc {
|
|
($e:expr) => {{
|
|
::qadapt::enter_protected();
|
|
let e = { $e };
|
|
::qadapt::exit_protected();
|
|
e
|
|
}};
|
|
}
|
|
|
|
static IS_ACTIVE: RwLock<bool> = RwLock::new(false);
|
|
static INTERNAL_ALLOCATION: RwLock<usize> = RwLock::new(usize::max_value());
|
|
|
|
/// Get the current "protection level" in QADAPT: calls to enter_protected() - exit_protected()
|
|
///
|
|
/// **Example**:
|
|
///
|
|
/// ```rust,no_run
|
|
/// use qadapt::enter_protected;
|
|
/// use qadapt::exit_protected;
|
|
/// use qadapt::protection_level;
|
|
///
|
|
/// fn main() {
|
|
/// enter_protected();
|
|
/// // We're now in an allocation-protected code region
|
|
/// assert_eq!(1, protection_level());
|
|
///
|
|
/// enter_protected();
|
|
/// // We're still memory protected, but we'll now need to exit twice to be safe
|
|
/// assert_eq!(2, protection_level());
|
|
/// exit_protected();
|
|
/// exit_protected();
|
|
///
|
|
/// // It's now safe to allocate/drop
|
|
/// }
|
|
pub fn protection_level() -> usize {
|
|
if cfg!(debug_assertions) {
|
|
PROTECTION_LEVEL.try_with(|v| *v.read()).unwrap_or(0)
|
|
} else {
|
|
0
|
|
}
|
|
}
|
|
|
|
fn claim_internal_alloc() {
|
|
loop {
|
|
match INTERNAL_ALLOCATION.write() {
|
|
ref mut lock if **lock == usize::max_value() => {
|
|
**lock = thread_id::get();
|
|
break;
|
|
}
|
|
_ => (),
|
|
}
|
|
}
|
|
}
|
|
|
|
fn release_internal_alloc() {
|
|
match INTERNAL_ALLOCATION.write() {
|
|
ref mut lock if **lock == thread_id::get() => **lock = usize::max_value(),
|
|
_ => panic!("Internal allocation tracking error"),
|
|
}
|
|
}
|
|
|
|
fn alloc_immediate() -> bool {
|
|
thread::panicking() || *INTERNAL_ALLOCATION.read() == thread_id::get()
|
|
}
|
|
|
|
unsafe impl GlobalAlloc for QADAPT {
|
|
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
|
|
if !*IS_ACTIVE.read() {
|
|
*IS_ACTIVE.write() = true;
|
|
}
|
|
|
|
// If we're attempting to allocate our PROTECTION_LEVEL thread local,
|
|
// just allow it through
|
|
if alloc_immediate() {
|
|
return SYSTEM_ALLOC.alloc(layout);
|
|
}
|
|
|
|
// Because accessing PROTECTION_LEVEL has the potential to trigger an allocation,
|
|
// we need to acquire the INTERNAL_ALLOCATION lock for our thread.
|
|
claim_internal_alloc();
|
|
let protection_level: Result<usize, ()> =
|
|
PROTECTION_LEVEL.try_with(|v| *v.read()).or(Ok(0));
|
|
release_internal_alloc();
|
|
|
|
match protection_level {
|
|
Ok(v) if v == 0 => SYSTEM_ALLOC.alloc(layout),
|
|
Ok(v) => {
|
|
// Tripped a bad allocation, but make sure further memory access during unwind
|
|
// doesn't have issues
|
|
PROTECTION_LEVEL.with(|v| *v.write() = 0);
|
|
panic!(
|
|
"Unexpected allocation for size {}, protection level: {}",
|
|
layout.size(),
|
|
v
|
|
)
|
|
}
|
|
Err(_) => unreachable!(),
|
|
}
|
|
}
|
|
|
|
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
|
|
if alloc_immediate() {
|
|
return SYSTEM_ALLOC.dealloc(ptr, layout);
|
|
}
|
|
|
|
claim_internal_alloc();
|
|
let protection_level: Result<usize, ()> =
|
|
PROTECTION_LEVEL.try_with(|v| *v.read()).or(Ok(0));
|
|
release_internal_alloc();
|
|
|
|
// Free before checking panic to make sure we avoid leaks
|
|
SYSTEM_ALLOC.dealloc(ptr, layout);
|
|
match protection_level {
|
|
Ok(v) if v > 0 => {
|
|
// Tripped a bad drop, but make sure further memory access during unwind
|
|
// doesn't have issues
|
|
PROTECTION_LEVEL.with(|v| *v.write() = 0);
|
|
panic!(
|
|
"Unexpected deallocation for size {}, protection level: {}",
|
|
layout.size(),
|
|
v
|
|
)
|
|
}
|
|
_ => (),
|
|
}
|
|
}
|
|
}
|