Add getcpu system call support for getting a NUMA node of the current thread

Author: Eugene-Usachev

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::<u32>::uninit();
+    let mut numa_node = MaybeUninit::<u32>::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::<u32>::uninit();
     unsafe {
         let r = ret(syscall!(__NR_getcpu, &mut cpu, zero(), zero()));

src/backend/linux_raw/vdso_wrappers.rs

@@ -117,6 +117,35 @@ pub(crate) fn clock_gettime_dynamic(id: DynamicClockId<'_>) -> io::Result<Timesp
     }
 }
 
+#[cfg(feature = "thread")]
+#[cfg(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) {
+    // 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::<u32>::uninit();
+        let mut numa_node = MaybeUninit::<u32>::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<Timesp
 ))]
 #[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.
+
     // 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

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<u8>], &mut [u8]>(&mut buf[..58]) },
-            b"//! Utilities for functions that return data via buffers.\n"
+            unsafe { core::mem::transmute::<&mut [MaybeUninit<u8>], &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();

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<Pid>) -> io::Result<CpuSet> {
 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()
+}