numalloc 0.1.2

A blazing-fast, NUMA-aware memory allocator written in pure Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145

use std::ptr::NonNull;

/// NUMA topology information.
pub struct NumaTopology {
    pub num_nodes: usize,
}

/// Detect the NUMA topology of the current system.
pub fn detect_topology() -> NumaTopology {
    #[cfg(all(target_os = "linux", feature = "numa-binding"))]
    {
        let num_nodes = detect_numa_nodes_linux().unwrap_or(1);
        NumaTopology { num_nodes }
    }
    #[cfg(not(all(target_os = "linux", feature = "numa-binding")))]
    {
        NumaTopology { num_nodes: 1 }
    }
}

#[cfg(all(target_os = "linux", feature = "numa-binding"))]
fn detect_numa_nodes_linux() -> std::io::Result<usize> {
    let mut count = 0usize;
    for entry in std::fs::read_dir("/sys/devices/system/node/")? {
        let name = entry?.file_name();
        if name.to_string_lossy().starts_with("node") {
            count += 1;
        }
    }
    Ok(count.max(1))
}

/// Allocate anonymous memory via `mmap`.
///
/// # Safety
/// Caller must ensure `size > 0`.
pub unsafe fn mmap_anonymous(size: usize) -> Option<NonNull<u8>> {
    #[cfg(target_os = "macos")]
    let flags = libc::MAP_PRIVATE | libc::MAP_ANON;
    #[cfg(not(target_os = "macos"))]
    let flags = libc::MAP_PRIVATE | libc::MAP_ANONYMOUS;

    let ptr = unsafe {
        libc::mmap(
            std::ptr::null_mut(),
            size,
            libc::PROT_READ | libc::PROT_WRITE,
            flags,
            -1,
            0,
        )
    };
    if ptr == libc::MAP_FAILED {
        return None;
    }
    NonNull::new(ptr as *mut u8)
}

/// Release memory previously obtained from [`mmap_anonymous`].
///
/// # Safety
/// `ptr` must originate from `mmap_anonymous` and `size` must match.
pub unsafe fn munmap(ptr: NonNull<u8>, size: usize) {
    unsafe { libc::munmap(ptr.as_ptr() as *mut libc::c_void, size) };
}

/// Bind a memory region to a specific NUMA node via `mbind`.
///
/// # Safety
/// `ptr` and `size` must describe a valid, mmap'd region.
#[cfg(all(target_os = "linux", feature = "numa-binding"))]
pub unsafe fn bind_to_node(ptr: NonNull<u8>, size: usize, node: usize) {
    let nodemask: u64 = 1u64 << node;
    unsafe {
        libc::syscall(
            libc::SYS_mbind,
            ptr.as_ptr() as *mut libc::c_void,
            size,
            2i32, // MPOL_BIND
            &nodemask as *const u64,
            64u64,
            0u32,
        );
    }
}

#[cfg(not(all(target_os = "linux", feature = "numa-binding")))]
pub unsafe fn bind_to_node(_ptr: NonNull<u8>, _size: usize, _node: usize) {}

/// Bind the calling thread to all CPUs belonging to `node`.
#[cfg(all(target_os = "linux", feature = "numa-binding"))]
pub fn bind_thread_to_node(node: usize) {
    unsafe {
        let mut cpuset: libc::cpu_set_t = std::mem::zeroed();
        if let Ok(cpulist) =
            std::fs::read_to_string(format!("/sys/devices/system/node/node{node}/cpulist"))
        {
            for range in cpulist.trim().split(',') {
                if let Some((start, end)) = range.split_once('-') {
                    let s: usize = start.parse().unwrap_or(0);
                    let e: usize = end.parse().unwrap_or(s);
                    for cpu in s..=e {
                        libc::CPU_SET(cpu, &mut cpuset);
                    }
                } else if let Ok(cpu) = range.parse::<usize>() {
                    libc::CPU_SET(cpu, &mut cpuset);
                }
            }
            libc::sched_setaffinity(0, std::mem::size_of::<libc::cpu_set_t>(), &cpuset);
        }
    }
}

#[cfg(not(all(target_os = "linux", feature = "numa-binding")))]
pub fn bind_thread_to_node(_node: usize) {}

/// Return the system page size (cached after first call).
pub fn page_size() -> usize {
    use std::sync::atomic::{AtomicUsize, Ordering};
    static CACHED: AtomicUsize = AtomicUsize::new(0);
    let val = CACHED.load(Ordering::Relaxed);
    if val != 0 {
        return val;
    }
    let ps = unsafe { libc::sysconf(libc::_SC_PAGESIZE) as usize };
    CACHED.store(ps, Ordering::Relaxed);
    ps
}

/// Advise the kernel that the given memory range is no longer needed.
/// The kernel may reclaim the physical pages; future accesses will
/// zero-fault them back in.
///
/// # Safety
/// `ptr` and `size` must describe a valid mmap'd region.
pub unsafe fn madvise_dontneed(ptr: NonNull<u8>, size: usize) {
    #[cfg(target_os = "macos")]
    unsafe {
        libc::madvise(ptr.as_ptr() as *mut libc::c_void, size, libc::MADV_FREE);
    }
    #[cfg(not(target_os = "macos"))]
    unsafe {
        libc::madvise(ptr.as_ptr() as *mut libc::c_void, size, libc::MADV_DONTNEED);
    }
}