#![allow(unsafe_code)]
#![deny(missing_docs)]
pub const NO_NODE: u32 = u32::MAX;
#[cfg(feature = "mock")]
pub mod mock {
use core::cell::RefCell;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum MockCall {
CurrentNode(u32),
BindRange {
base: usize,
len: usize,
node: u32,
},
ReserveOnNode {
size: usize,
align: usize,
node: u32,
},
}
std::thread_local! {
pub static CALLS: RefCell<Vec<MockCall>> = const { RefCell::new(Vec::new()) };
pub static CURRENT_NODE_SLOT: RefCell<u32> = const { RefCell::new(0) };
}
pub fn drain() -> Vec<MockCall> {
CALLS.with(|c| c.borrow_mut().drain(..).collect())
}
pub fn set_current_node(node: u32) {
CURRENT_NODE_SLOT.with(|c| *c.borrow_mut() = node);
}
pub(crate) fn current_node_slot() -> u32 {
CURRENT_NODE_SLOT.with(|c| *c.borrow())
}
pub(crate) fn record(call: MockCall) {
CALLS.with(|c| c.borrow_mut().push(call));
}
}
#[must_use]
pub fn current_node() -> Option<u32> {
#[cfg(feature = "mock")]
{
let n = mock::current_node_slot();
mock::record(mock::MockCall::CurrentNode(n));
return Some(n);
}
#[cfg(not(feature = "mock"))]
{
let raw = platform::current_node_impl();
if raw == NO_NODE {
None
} else {
Some(raw)
}
}
}
pub unsafe fn bind_range(base: *mut u8, len: usize, node: u32) {
if node == NO_NODE || len == 0 {
return;
}
#[cfg(feature = "mock")]
{
mock::record(mock::MockCall::BindRange {
base: base as usize,
len,
node,
});
return;
}
#[cfg(not(feature = "mock"))]
{
platform::bind_range_impl(base, len, node);
}
}
#[cfg(feature = "vmem-integration")]
#[must_use]
pub fn reserve_on_node(size: usize, align: usize, node: u32) -> Option<aligned_vmem::Reservation> {
#[cfg(feature = "mock")]
{
mock::record(mock::MockCall::ReserveOnNode { size, align, node });
let r = aligned_vmem::reserve_aligned(size, align)?;
if node != NO_NODE {
let base = r.as_ptr();
let len = r.len();
mock::record(mock::MockCall::BindRange {
base: base as usize,
len,
node,
});
}
return Some(r);
}
#[cfg(not(feature = "mock"))]
{
platform::reserve_on_node_impl(size, align, node)
}
}
#[cfg(all(target_os = "linux", not(miri)))]
mod platform {
use super::{bind_range_impl_linux, NO_NODE};
pub(super) fn current_node_impl() -> u32 {
let cpu = unsafe { libc_sched_getcpu() };
if cpu < 0 {
return NO_NODE;
}
cpu_to_numa_node(cpu as u32)
}
pub(super) fn bind_range_impl(base: *mut u8, len: usize, node: u32) {
unsafe { bind_range_impl_linux(base, len, node) };
}
#[cfg(feature = "vmem-integration")]
pub(super) fn reserve_on_node_impl(
size: usize,
align: usize,
node: u32,
) -> Option<aligned_vmem::Reservation> {
let r = aligned_vmem::reserve_aligned(size, align)?;
if node != NO_NODE {
let base = r.as_ptr();
let len = r.len();
unsafe { bind_range_impl_linux(base, len, node) };
}
Some(r)
}
fn cpu_to_numa_node(cpu_idx: u32) -> u32 {
for node in 0u32..64 {
if node_contains_cpu(node, cpu_idx) {
return node;
}
}
0
}
fn node_contains_cpu(node: u32, cpu_idx: u32) -> bool {
let mut path = [0u8; 64];
let path_str = format_sysfs_path(&mut path, node);
read_cpumap_contains_cpu(path_str, cpu_idx)
}
fn format_sysfs_path(buf: &mut [u8; 64], node: u32) -> &[u8] {
const PREFIX: &[u8] = b"/sys/devices/system/node/node";
const SUFFIX: &[u8] = b"/cpumap\0";
let mut pos = 0usize;
for &b in PREFIX {
buf[pos] = b;
pos += 1;
}
let mut tmp = [0u8; 4];
let mut n = node;
let mut digits = 0usize;
if n == 0 {
tmp[0] = b'0';
digits = 1;
} else {
while n > 0 {
tmp[digits] = b'0' + (n % 10) as u8;
n /= 10;
digits += 1;
}
tmp[..digits].reverse();
}
for i in 0..digits {
buf[pos] = tmp[i];
pos += 1;
}
for &b in SUFFIX {
buf[pos] = b;
pos += 1;
}
&buf[..pos]
}
fn read_cpumap_contains_cpu(path: &[u8], cpu_idx: u32) -> bool {
let fd = unsafe { libc_open(path.as_ptr() as *const core::ffi::c_char, 0) };
if fd < 0 {
return false;
}
let mut buf = [0u8; 256];
let n = unsafe { libc_read(fd, buf.as_mut_ptr() as *mut core::ffi::c_void, 256) };
unsafe { libc_close(fd) };
if n <= 0 {
return false;
}
parse_cpumap_contains_cpu(&buf[..n as usize], cpu_idx)
}
fn parse_cpumap_contains_cpu(data: &[u8], cpu_idx: u32) -> bool {
let data = trim_end(data);
let word_count = data.iter().filter(|&&b| b == b',').count() + 1;
let target_word = (cpu_idx / 32) as usize;
let bit_in_word = cpu_idx % 32;
if target_word >= word_count {
return false;
}
let left_index = word_count - 1 - target_word;
let word_str = match nth_token(data, left_index, b',') {
Some(s) => s,
None => return false,
};
let val = match parse_hex_u32(word_str) {
Some(v) => v,
None => return false,
};
(val >> bit_in_word) & 1 == 1
}
fn trim_end(data: &[u8]) -> &[u8] {
let mut end = data.len();
while end > 0 && (data[end - 1] == b'\n' || data[end - 1] == b'\r' || data[end - 1] == b' ')
{
end -= 1;
}
&data[..end]
}
fn nth_token(data: &[u8], n: usize, sep: u8) -> Option<&[u8]> {
let mut idx = 0usize;
let mut start = 0usize;
for (i, &b) in data.iter().enumerate() {
if b == sep {
if idx == n {
return Some(&data[start..i]);
}
idx += 1;
start = i + 1;
}
}
if idx == n {
Some(&data[start..])
} else {
None
}
}
fn parse_hex_u32(s: &[u8]) -> Option<u32> {
if s.is_empty() {
return None;
}
let mut val: u32 = 0;
for &b in s {
let digit = match b {
b'0'..=b'9' => b - b'0',
b'a'..=b'f' => b - b'a' + 10,
b'A'..=b'F' => b - b'A' + 10,
_ => return None,
};
val = val.wrapping_shl(4) | digit as u32;
}
Some(val)
}
extern "C" {
fn sched_getcpu() -> core::ffi::c_int;
fn open(path: *const core::ffi::c_char, flags: core::ffi::c_int, ...) -> core::ffi::c_int;
fn read(
fd: core::ffi::c_int,
buf: *mut core::ffi::c_void,
count: usize,
) -> core::ffi::c_long;
fn close(fd: core::ffi::c_int) -> core::ffi::c_int;
}
unsafe fn libc_sched_getcpu() -> core::ffi::c_int {
sched_getcpu()
}
unsafe fn libc_open(
path: *const core::ffi::c_char,
flags: core::ffi::c_int,
) -> core::ffi::c_int {
open(path, flags)
}
unsafe fn libc_read(
fd: core::ffi::c_int,
buf: *mut core::ffi::c_void,
count: usize,
) -> core::ffi::c_long {
read(fd, buf, count)
}
unsafe fn libc_close(fd: core::ffi::c_int) {
let _ = close(fd);
}
}
#[cfg(all(
target_os = "linux",
not(miri),
any(target_arch = "x86_64", target_arch = "aarch64")
))]
unsafe fn bind_range_impl_linux(base: *mut u8, len: usize, node: u32) {
if node == NO_NODE || node >= 64 {
return;
}
let nodemask: u64 = 1u64 << node;
let maxnode: u64 = 64;
libc_mbind(
base as *mut core::ffi::c_void,
len as u64,
MPOL_PREFERRED,
&nodemask as *const u64,
maxnode,
0,
);
}
#[cfg(all(
target_os = "linux",
not(miri),
not(any(target_arch = "x86_64", target_arch = "aarch64"))
))]
unsafe fn bind_range_impl_linux(_base: *mut u8, _len: usize, _node: u32) {
}
#[cfg(all(target_os = "linux", not(miri)))]
const MPOL_PREFERRED: i32 = 1;
#[cfg(all(target_os = "linux", not(miri), target_arch = "x86_64"))]
const SYS_MBIND: i64 = 237;
#[cfg(all(target_os = "linux", not(miri), target_arch = "aarch64"))]
const SYS_MBIND: i64 = 235;
#[cfg(all(
target_os = "linux",
not(miri),
any(target_arch = "x86_64", target_arch = "aarch64")
))]
extern "C" {
fn syscall(number: i64, ...) -> i64;
}
#[cfg(all(
target_os = "linux",
not(miri),
any(target_arch = "x86_64", target_arch = "aarch64")
))]
unsafe fn libc_mbind(
addr: *mut core::ffi::c_void,
len: u64,
mode: i32,
nodemask: *const u64,
maxnode: u64,
flags: u32,
) -> i64 {
syscall(
SYS_MBIND,
addr,
len as usize,
mode as i64,
nodemask,
maxnode as usize,
flags as i64,
)
}
#[cfg(all(windows, not(miri)))]
mod platform {
use super::NO_NODE;
pub(super) fn current_node_impl() -> u32 {
let mut proc_num = ProcessorNumber {
group: 0,
number: 0,
reserved: 0,
};
unsafe { GetCurrentProcessorNumberEx(&mut proc_num) };
let mut node: u16 = 0;
let ok = unsafe { GetNumaProcessorNodeEx(&proc_num, &mut node) };
if ok == 0 {
return NO_NODE;
}
node as u32
}
pub(super) fn bind_range_impl(_base: *mut u8, _len: usize, _node: u32) {
}
#[cfg(feature = "vmem-integration")]
pub(super) fn reserve_on_node_impl(
size: usize,
align: usize,
node: u32,
) -> Option<aligned_vmem::Reservation> {
if node == NO_NODE {
return aligned_vmem::reserve_aligned(size, align);
}
reserve_aligned_numa(size, align, node)
}
#[cfg(feature = "vmem-integration")]
fn reserve_aligned_numa(
size: usize,
align: usize,
node: u32,
) -> Option<aligned_vmem::Reservation> {
use aligned_vmem::PAGE;
if size == 0 || !align.is_power_of_two() || align < PAGE || size % PAGE != 0 {
return None;
}
let over = size.checked_add(align)?;
let raw = unsafe {
VirtualAllocExNuma(
GetCurrentProcess(),
core::ptr::null_mut(),
over,
MEM_RESERVE | MEM_COMMIT,
PAGE_READWRITE,
node,
)
};
if raw.is_null() {
return None;
}
let raw_u = raw as usize;
let base_u = (raw_u + align - 1) & !(align - 1);
let base = base_u as *mut u8;
let r = unsafe {
aligned_vmem::Reservation::from_raw_parts(base, size, raw as *mut u8, over, align)
};
Some(r)
}
#[repr(C)]
struct ProcessorNumber {
group: u16,
number: u8,
reserved: u8,
}
extern "system" {
fn GetCurrentProcessorNumberEx(proc_number: *mut ProcessorNumber);
fn GetNumaProcessorNodeEx(processor: *const ProcessorNumber, node_number: *mut u16) -> i32;
fn GetCurrentProcess() -> *mut core::ffi::c_void;
}
#[cfg(feature = "vmem-integration")]
extern "system" {
fn VirtualAllocExNuma(
h_process: *mut core::ffi::c_void,
lp_address: *mut core::ffi::c_void,
dw_size: usize,
fl_allocation_type: u32,
fl_protect: u32,
nnd_preferred: u32,
) -> *mut core::ffi::c_void;
}
#[cfg(feature = "vmem-integration")]
const MEM_RESERVE: u32 = 0x0000_2000;
#[cfg(feature = "vmem-integration")]
const MEM_COMMIT: u32 = 0x0000_1000;
#[cfg(feature = "vmem-integration")]
const PAGE_READWRITE: u32 = 0x04;
}
#[cfg(target_os = "macos")]
mod platform {
use super::NO_NODE;
pub(super) fn current_node_impl() -> u32 {
NO_NODE
}
pub(super) fn bind_range_impl(_base: *mut u8, _len: usize, _node: u32) {}
#[cfg(feature = "vmem-integration")]
pub(super) fn reserve_on_node_impl(
size: usize,
align: usize,
_node: u32,
) -> Option<aligned_vmem::Reservation> {
aligned_vmem::reserve_aligned(size, align)
}
}
#[cfg(miri)]
mod platform {
use super::NO_NODE;
pub(super) fn current_node_impl() -> u32 {
NO_NODE
}
pub(super) fn bind_range_impl(_base: *mut u8, _len: usize, _node: u32) {}
#[cfg(feature = "vmem-integration")]
pub(super) fn reserve_on_node_impl(
size: usize,
align: usize,
_node: u32,
) -> Option<aligned_vmem::Reservation> {
aligned_vmem::reserve_aligned(size, align)
}
}
#[cfg(not(any(target_os = "linux", windows, target_os = "macos", miri,)))]
mod platform {
use super::NO_NODE;
pub(super) fn current_node_impl() -> u32 {
NO_NODE
}
pub(super) fn bind_range_impl(_base: *mut u8, _len: usize, _node: u32) {}
#[cfg(feature = "vmem-integration")]
pub(super) fn reserve_on_node_impl(
size: usize,
align: usize,
_node: u32,
) -> Option<aligned_vmem::Reservation> {
aligned_vmem::reserve_aligned(size, align)
}
}