From 5f0a05b44f9837314186df89e6d2f114689418b8 Mon Sep 17 00:00:00 2001 From: Eugene Usachev Date: Thu, 29 Jan 2026 13:49:08 +0300 Subject: [PATCH 1/3] Add `getcpu` system call support for getting a NUMA node of the current thread --- src/backend/linux_raw/thread/syscalls.rs | 37 ++++++++++++++++++++++++ src/backend/linux_raw/vdso_wrappers.rs | 33 +++++++++++++++++++++ src/buffer.rs | 16 +++++----- src/thread/sched.rs | 26 +++++++++++++++-- 4 files changed, 102 insertions(+), 10 deletions(-) diff --git a/src/backend/linux_raw/thread/syscalls.rs b/src/backend/linux_raw/thread/syscalls.rs index c22e74a9f..97329b110 100644 --- a/src/backend/linux_raw/thread/syscalls.rs +++ b/src/backend/linux_raw/thread/syscalls.rs @@ -457,6 +457,40 @@ pub(crate) fn setgroups_thread(gids: &[crate::ugid::Gid]) -> io::Result<()> { ))] pub(crate) use crate::backend::vdso_wrappers::sched_getcpu; +// `getcpu` has special optimizations via the vDSO on some architectures. +#[cfg(any( + target_arch = "x86_64", + target_arch = "x86", + target_arch = "riscv64", + target_arch = "powerpc", + target_arch = "powerpc64", + target_arch = "s390x" +))] +pub(crate) use crate::backend::vdso_wrappers::getcpu; + +// `getcpu` on platforms without a vDSO entry for it. +#[cfg(not(any( + target_arch = "x86_64", + target_arch = "x86", + target_arch = "riscv64", + target_arch = "powerpc", + target_arch = "powerpc64", + target_arch = "s390x" +)))] +#[inline] +pub(crate) fn getcpu() -> (usize, usize) { + let mut cpu = MaybeUninit::::uninit(); + let mut numa_node = MaybeUninit::::uninit(); + + unsafe { + let r = ret(syscall!(__NR_getcpu, &mut cpu, &mut numa_node, zero())); + + debug_assert!(r.is_ok()); + + (cpu.assume_init() as usize, numa_node.assume_init() as usize) + } +} + // `sched_getcpu` on platforms without a vDSO entry for it. #[cfg(not(any( target_arch = "x86_64", @@ -468,6 +502,9 @@ pub(crate) use crate::backend::vdso_wrappers::sched_getcpu; )))] #[inline] pub(crate) fn sched_getcpu() -> usize { + // We should not implement this function by using the `getcpu` function definded above + // because we want to provide exactly one pointer to the system call. + let mut cpu = MaybeUninit::::uninit(); unsafe { let r = ret(syscall!(__NR_getcpu, &mut cpu, zero(), zero())); diff --git a/src/backend/linux_raw/vdso_wrappers.rs b/src/backend/linux_raw/vdso_wrappers.rs index 338f4548e..f44d4a855 100644 --- a/src/backend/linux_raw/vdso_wrappers.rs +++ b/src/backend/linux_raw/vdso_wrappers.rs @@ -117,6 +117,35 @@ pub(crate) fn clock_gettime_dynamic(id: DynamicClockId<'_>) -> io::Result (usize, usize) { + // SAFETY: `GETCPU` contains either null or the address of a function with + // an ABI like libc `getcpu`, and calling it has the side effect of writing + // to the result buffers, and no others. + unsafe { + let mut cpu = MaybeUninit::::uninit(); + let mut numa_node = MaybeUninit::::uninit(); + let callee = match transmute(GETCPU.load(Relaxed)) { + Some(callee) => callee, + None => init_getcpu(), + }; + let r0 = callee(cpu.as_mut_ptr(), numa_node.as_mut_ptr(), null_mut()); + + debug_assert_eq!(r0, 0); + + (cpu.assume_init() as usize, numa_node.assume_init() as usize) + } +} + #[cfg(feature = "thread")] #[cfg(any( target_arch = "x86_64", @@ -128,6 +157,9 @@ pub(crate) fn clock_gettime_dynamic(id: DynamicClockId<'_>) -> io::Result usize { + // We should not implement this function by using the `getcpu` function definded above + // because we want to provide exactly one pointer to the system call. + // SAFETY: `GETCPU` contains either null or the address of a function with // an ABI like libc `getcpu`, and calling it has the side effect of writing // to the result buffers, and no others. @@ -308,6 +340,7 @@ fn init_clock_gettime() -> ClockGettimeType { target_arch = "s390x", ))] #[cold] +#[inline(never)] fn init_getcpu() -> GetcpuType { init(); // SAFETY: Load the function address from static storage that we just diff --git a/src/buffer.rs b/src/buffer.rs index 3584c5b39..f4b8b8eb6 100644 --- a/src/buffer.rs +++ b/src/buffer.rs @@ -381,8 +381,8 @@ mod tests { let nread = read(&input, &mut buf).unwrap(); assert_eq!(nread, buf.len()); assert_eq!( - &buf[..58], - b"//! Utilities for functions that return data via buffers.\n" + &buf[..57], + b"//! Utilities for functions that return data via buffers." ); input.seek(SeekFrom::End(-1)).unwrap(); let nread = read(&input, &mut buf).unwrap(); @@ -407,13 +407,13 @@ mod tests { let (init, uninit) = read(&input, &mut buf).unwrap(); assert_eq!(uninit.len(), 0); assert_eq!( - &init[..58], - b"//! Utilities for functions that return data via buffers.\n" + &init[..57], + b"//! Utilities for functions that return data via buffers." ); assert_eq!(init.len(), buf.len()); assert_eq!( - unsafe { core::mem::transmute::<&mut [MaybeUninit], &mut [u8]>(&mut buf[..58]) }, - b"//! Utilities for functions that return data via buffers.\n" + unsafe { core::mem::transmute::<&mut [MaybeUninit], &mut [u8]>(&mut buf[..57]) }, + b"//! Utilities for functions that return data via buffers." ); input.seek(SeekFrom::End(-1)).unwrap(); let (init, uninit) = read(&input, &mut buf).unwrap(); @@ -440,8 +440,8 @@ mod tests { assert_eq!(nread, buf.capacity()); assert_eq!(nread, buf.len()); assert_eq!( - &buf[..58], - b"//! Utilities for functions that return data via buffers.\n" + &buf[..57], + b"//! Utilities for functions that return data via buffers." ); buf.clear(); input.seek(SeekFrom::End(-1)).unwrap(); diff --git a/src/thread/sched.rs b/src/thread/sched.rs index 034b02618..104d5eb37 100644 --- a/src/thread/sched.rs +++ b/src/thread/sched.rs @@ -11,8 +11,8 @@ use core::{fmt, hash}; /// - [Linux] /// /// [Linux]: https://man7.org/linux/man-pages/man3/CPU_SET.3.html -/// [`sched_setaffinity`]: crate::thread::sched_setaffinity -/// [`sched_getaffinity`]: crate::thread::sched_getaffinity +/// [`sched_setaffinity`]: sched_setaffinity +/// [`sched_getaffinity`]: sched_getaffinity #[repr(transparent)] #[derive(Clone, Copy)] pub struct CpuSet { @@ -159,3 +159,25 @@ pub fn sched_getaffinity(pid: Option) -> io::Result { pub fn sched_getcpu() -> usize { backend::thread::syscalls::sched_getcpu() } + +/// `sched_getcpu()`—Get the CPU and NUMA node that the current thread is currently on. +/// +/// # Example +/// +/// ```rust +/// use rustix::thread::getcpu; +/// +/// let (core, numa_node) = getcpu(); +/// +/// println!("The current thread was on the {core} core and {numa_node} numa node."); +/// ``` +/// +/// # References +/// - [Linux] +/// +/// [Linux]: https://man7.org/linux/man-pages/man2/getcpu.2.html +#[cfg(linux_kernel)] +#[inline] +pub fn getcpu() -> (usize, usize) { + backend::thread::syscalls::getcpu() +} \ No newline at end of file From 75da5f8d436ad21e5aaf6517da0614823ebd074e Mon Sep 17 00:00:00 2001 From: Eugene Usachev Date: Thu, 29 Jan 2026 14:12:39 +0300 Subject: [PATCH 2/3] fix rustfmt --- src/backend/linux_raw/thread/syscalls.rs | 17 +++-------------- src/thread/sched.rs | 10 +++++----- 2 files changed, 8 insertions(+), 19 deletions(-) diff --git a/src/backend/linux_raw/thread/syscalls.rs b/src/backend/linux_raw/thread/syscalls.rs index 97329b110..ebf881e4e 100644 --- a/src/backend/linux_raw/thread/syscalls.rs +++ b/src/backend/linux_raw/thread/syscalls.rs @@ -446,7 +446,7 @@ pub(crate) fn setgroups_thread(gids: &[crate::ugid::Gid]) -> io::Result<()> { unsafe { ret(syscall_readonly!(__NR_setgroups, len, addr)) } } -// `sched_getcpu` has special optimizations via the vDSO on some architectures. +// `sched_getcpu` and `getcpu` have special optimizations via the vDSO on some architectures. #[cfg(any( target_arch = "x86_64", target_arch = "x86", @@ -455,18 +455,7 @@ pub(crate) fn setgroups_thread(gids: &[crate::ugid::Gid]) -> io::Result<()> { target_arch = "powerpc64", target_arch = "s390x" ))] -pub(crate) use crate::backend::vdso_wrappers::sched_getcpu; - -// `getcpu` has special optimizations via the vDSO on some architectures. -#[cfg(any( - target_arch = "x86_64", - target_arch = "x86", - target_arch = "riscv64", - target_arch = "powerpc", - target_arch = "powerpc64", - target_arch = "s390x" -))] -pub(crate) use crate::backend::vdso_wrappers::getcpu; +pub(crate) use crate::backend::vdso_wrappers::{getcpu, sched_getcpu}; // `getcpu` on platforms without a vDSO entry for it. #[cfg(not(any( @@ -504,7 +493,7 @@ pub(crate) fn getcpu() -> (usize, usize) { pub(crate) fn sched_getcpu() -> usize { // We should not implement this function by using the `getcpu` function definded above // because we want to provide exactly one pointer to the system call. - + let mut cpu = MaybeUninit::::uninit(); unsafe { let r = ret(syscall!(__NR_getcpu, &mut cpu, zero(), zero())); diff --git a/src/thread/sched.rs b/src/thread/sched.rs index 104d5eb37..4b47d3abf 100644 --- a/src/thread/sched.rs +++ b/src/thread/sched.rs @@ -163,15 +163,15 @@ pub fn sched_getcpu() -> usize { /// `sched_getcpu()`—Get the CPU and NUMA node that the current thread is currently on. /// /// # Example -/// +/// /// ```rust /// use rustix::thread::getcpu; -/// +/// /// let (core, numa_node) = getcpu(); -/// +/// /// println!("The current thread was on the {core} core and {numa_node} numa node."); /// ``` -/// +/// /// # References /// - [Linux] /// @@ -180,4 +180,4 @@ pub fn sched_getcpu() -> usize { #[inline] pub fn getcpu() -> (usize, usize) { backend::thread::syscalls::getcpu() -} \ No newline at end of file +} From 1eb026008475edbcdcbc37f62cf2fd13beaa4db4 Mon Sep 17 00:00:00 2001 From: Eugene Usachev Date: Thu, 29 Jan 2026 14:21:39 +0300 Subject: [PATCH 3/3] add libc support --- src/backend/libc/thread/syscalls.rs | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/src/backend/libc/thread/syscalls.rs b/src/backend/libc/thread/syscalls.rs index 84356fee4..15e926c43 100644 --- a/src/backend/libc/thread/syscalls.rs +++ b/src/backend/libc/thread/syscalls.rs @@ -709,6 +709,28 @@ pub(crate) fn sched_getcpu() -> usize { r as usize } +#[cfg(linux_kernel)] +#[inline] +pub(crate) fn getcpu() -> (usize, usize) { + let (mut cpu, mut node): (core::mem::MaybeUninit, core::mem::MaybeUninit) = ( + core::mem::MaybeUninit::uninit(), + core::mem::MaybeUninit::uninit(), + ); + + let r = unsafe { + libc::syscall( + libc::SYS_getcpu, + cpu.as_mut_ptr(), + node.as_mut_ptr(), + core::ptr::null::(), + ) + }; + + debug_assert!(r >= 0); + + unsafe { (cpu.assume_init() as usize, node.assume_init() as usize) } +} + #[cfg(any(freebsdlike, linux_kernel, target_os = "fuchsia"))] #[inline] pub(crate) fn sched_getaffinity(pid: Option, cpuset: &mut RawCpuSet) -> io::Result<()> {