Implement rustix::process::sched_getcpu. (#921)

On some architectures, there is a `getcpu` entry in the vDSO, so use
that if available, and use the syscall otherwise.

Author: sunfishcode

src/backend/libc/process/mod.rs

@@ -1,4 +1,4 @@
-#[cfg(any(linux_kernel, target_os = "dragonfly", target_os = "fuchsia"))]
+#[cfg(any(freebsdlike, linux_kernel, target_os = "fuchsia"))]
 pub(crate) mod cpu_set;
 #[cfg(not(windows))]
 pub(crate) mod syscalls;

src/backend/libc/process/syscalls.rs

@@ -1,6 +1,6 @@
 //! libc syscalls supporting `rustix::process`.
 
-#[cfg(any(linux_kernel, target_os = "dragonfly", target_os = "fuchsia"))]
+#[cfg(any(freebsdlike, linux_kernel, target_os = "fuchsia"))]
 use super::types::RawCpuSet;
 use crate::backend::c;
 #[cfg(not(any(target_os = "wasi", target_os = "fuchsia")))]
@@ -72,6 +72,14 @@ use {
     super::super::conv::ret_owned_fd, crate::process::PidfdFlags, crate::process::PidfdGetfdFlags,
 };
 
+#[cfg(any(linux_kernel, target_os = "dragonfly"))]
+#[inline]
+pub(crate) fn sched_getcpu() -> usize {
+    let r = unsafe { libc::sched_getcpu() };
+    debug_assert!(r >= 0);
+    r as usize
+}
+
 #[cfg(feature = "fs")]
 #[cfg(not(target_os = "wasi"))]
 pub(crate) fn chdir(path: &CStr) -> io::Result<()> {
@@ -181,7 +189,7 @@ pub(crate) fn getpgrp() -> Pid {
     }
 }
 
-#[cfg(any(linux_kernel, target_os = "dragonfly", target_os = "fuchsia"))]
+#[cfg(any(freebsdlike, linux_kernel, target_os = "fuchsia"))]
 #[inline]
 pub(crate) fn sched_getaffinity(pid: Option<Pid>, cpuset: &mut RawCpuSet) -> io::Result<()> {
     unsafe {
@@ -193,7 +201,7 @@ pub(crate) fn sched_getaffinity(pid: Option<Pid>, cpuset: &mut RawCpuSet) -> io:
     }
 }
 
-#[cfg(any(linux_kernel, target_os = "dragonfly", target_os = "fuchsia"))]
+#[cfg(any(freebsdlike, linux_kernel, target_os = "fuchsia"))]
 #[inline]
 pub(crate) fn sched_setaffinity(pid: Option<Pid>, cpuset: &RawCpuSet) -> io::Result<()> {
     unsafe {

src/backend/libc/process/types.rs

@@ -155,18 +155,18 @@ pub type RawCpuid = u32;
 #[cfg(freebsdlike)]
 pub type RawId = c::id_t;
 
-#[cfg(any(linux_kernel, target_os = "dragonfly", target_os = "fuchsia"))]
+#[cfg(any(linux_kernel, target_os = "fuchsia"))]
 pub(crate) type RawCpuSet = c::cpu_set_t;
+#[cfg(freebsdlike)]
+pub(crate) type RawCpuSet = c::cpuset_t;
 
-#[cfg(any(linux_kernel, target_os = "dragonfly", target_os = "fuchsia"))]
+#[cfg(any(freebsdlike, linux_kernel, target_os = "fuchsia"))]
 #[inline]
 pub(crate) fn raw_cpu_set_new() -> RawCpuSet {
     let mut set = unsafe { core::mem::zeroed() };
     super::cpu_set::CPU_ZERO(&mut set);
     set
 }
 
-#[cfg(any(linux_kernel, target_os = "fuchsia"))]
+#[cfg(any(freebsdlike, linux_kernel, target_os = "fuchsia"))]
 pub(crate) const CPU_SETSIZE: usize = c::CPU_SETSIZE as usize;
-#[cfg(target_os = "dragonfly")]
-pub(crate) const CPU_SETSIZE: usize = 256;

src/backend/linux_raw/process/syscalls.rs

@@ -36,6 +36,32 @@ use {crate::backend::conv::ret_c_uint_infallible, crate::fs::Mode};
 #[cfg(feature = "alloc")]
 use {crate::backend::conv::slice_just_addr_mut, crate::process::Gid};
 
+// `sched_getcpu` has special optimizations via the vDSO on some architectures.
+#[cfg(any(
+    target_arch = "x86_64",
+    target_arch = "x86",
+    target_arch = "riscv64",
+    target_arch = "powerpc64"
+))]
+pub(crate) use crate::backend::vdso_wrappers::sched_getcpu;
+
+// `sched_getcpu` on platforms without a vDSO entry for it.
+#[cfg(not(any(
+    target_arch = "x86_64",
+    target_arch = "x86",
+    target_arch = "riscv64",
+    target_arch = "powerpc64"
+)))]
+#[inline]
+pub(crate) fn sched_getcpu() -> usize {
+    let mut cpu = MaybeUninit::<u32>::uninit();
+    unsafe {
+        let r = ret(syscall!(__NR_getcpu, &mut cpu, zero(), zero()));
+        debug_assert!(r.is_ok());
+        cpu.assume_init() as usize
+    }
+}
+
 #[cfg(feature = "fs")]
 #[inline]
 pub(crate) fn chdir(filename: &CStr) -> io::Result<()> {

src/backend/linux_raw/vdso_wrappers.rs

@@ -14,6 +14,14 @@ use super::reg::{ArgReg, RetReg, SyscallNumber, A0, A1, A2, A3, A4, A5, R0};
 use super::vdso;
 #[cfg(target_arch = "x86")]
 use core::arch::global_asm;
+#[cfg(feature = "process")]
+#[cfg(any(
+    target_arch = "x86_64",
+    target_arch = "x86",
+    target_arch = "riscv64",
+    target_arch = "powerpc64"
+))]
+use core::ffi::c_void;
 use core::mem::transmute;
 use core::ptr::null_mut;
 use core::sync::atomic::AtomicPtr;
@@ -96,6 +104,30 @@ pub(crate) fn clock_gettime_dynamic(which_clock: DynamicClockId<'_>) -> io::Resu
     }
 }
 
+#[cfg(feature = "process")]
+#[cfg(any(
+    target_arch = "x86_64",
+    target_arch = "x86",
+    target_arch = "riscv64",
+    target_arch = "powerpc64"
+))]
+#[inline]
+pub(crate) fn sched_getcpu() -> 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 callee = match transmute(GETCPU.load(Relaxed)) {
+            Some(callee) => callee,
+            None => init_getcpu(),
+        };
+        let r0 = callee(cpu.as_mut_ptr(), null_mut(), null_mut());
+        debug_assert_eq!(r0, 0);
+        cpu.assume_init() as usize
+    }
+}
+
 #[cfg(target_arch = "x86")]
 pub(super) mod x86_via_vdso {
     use super::{transmute, ArgReg, Relaxed, RetReg, SyscallNumber, A0, A1, A2, A3, A4, A5, R0};
@@ -223,6 +255,15 @@ pub(super) mod x86_via_vdso {
 #[cfg(feature = "time")]
 type ClockGettimeType = unsafe extern "C" fn(c::c_int, *mut Timespec) -> c::c_int;
 
+#[cfg(feature = "process")]
+#[cfg(any(
+    target_arch = "x86_64",
+    target_arch = "x86",
+    target_arch = "riscv64",
+    target_arch = "powerpc64"
+))]
+type GetcpuType = unsafe extern "C" fn(*mut u32, *mut u32, *mut c_void) -> c::c_int;
+
 /// The underlying syscall functions are only called from asm, using the
 /// special syscall calling convention to pass arguments and return values,
 /// which the signature here doesn't reflect.
@@ -239,6 +280,22 @@ fn init_clock_gettime() -> ClockGettimeType {
     unsafe { transmute(CLOCK_GETTIME.load(Relaxed)) }
 }
 
+/// Initialize `GETCPU` and return its value.
+#[cfg(feature = "process")]
+#[cfg(any(
+    target_arch = "x86_64",
+    target_arch = "x86",
+    target_arch = "riscv64",
+    target_arch = "powerpc64"
+))]
+#[cold]
+fn init_getcpu() -> GetcpuType {
+    init();
+    // SAFETY: Load the function address from static storage that we just
+    // initialized.
+    unsafe { transmute(GETCPU.load(Relaxed)) }
+}
+
 /// Initialize `SYSCALL` and return its value.
 #[cfg(target_arch = "x86")]
 #[cold]
@@ -254,6 +311,14 @@ fn init_syscall() -> SyscallType {
 struct Function;
 #[cfg(feature = "time")]
 static mut CLOCK_GETTIME: AtomicPtr<Function> = AtomicPtr::new(null_mut());
+#[cfg(feature = "process")]
+#[cfg(any(
+    target_arch = "x86_64",
+    target_arch = "x86",
+    target_arch = "riscv64",
+    target_arch = "powerpc64"
+))]
+static mut GETCPU: AtomicPtr<Function> = AtomicPtr::new(null_mut());
 #[cfg(target_arch = "x86")]
 static mut SYSCALL: AtomicPtr<Function> = AtomicPtr::new(null_mut());
 
@@ -315,6 +380,24 @@ unsafe fn _rustix_clock_gettime_via_syscall(
     ret(syscall!(__NR_clock_gettime, c_int(clockid), res))
 }
 
+#[cfg(feature = "process")]
+#[cfg(any(
+    target_arch = "x86_64",
+    target_arch = "x86",
+    target_arch = "riscv64",
+    target_arch = "powerpc64"
+))]
+unsafe extern "C" fn rustix_getcpu_via_syscall(
+    cpu: *mut u32,
+    node: *mut u32,
+    unused: *mut c_void,
+) -> c::c_int {
+    match ret(syscall!(__NR_getcpu, cpu, node, unused)) {
+        Ok(()) => 0,
+        Err(err) => err.raw_os_error().wrapping_neg(),
+    }
+}
+
 #[cfg(target_arch = "x86")]
 extern "C" {
     /// A symbol pointing to an `int 0x80` instruction. This “function” is only
@@ -361,6 +444,24 @@ fn minimal_init() {
                 .ok();
         }
 
+        #[cfg(feature = "process")]
+        #[cfg(any(
+            target_arch = "x86_64",
+            target_arch = "x86",
+            target_arch = "riscv64",
+            target_arch = "powerpc64"
+        ))]
+        {
+            GETCPU
+                .compare_exchange(
+                    null_mut(),
+                    rustix_getcpu_via_syscall as *mut Function,
+                    Relaxed,
+                    Relaxed,
+                )
+                .ok();
+        }
+
         #[cfg(target_arch = "x86")]
         {
             SYSCALL
@@ -430,6 +531,60 @@ fn init() {
             }
         }
 
+        #[cfg(feature = "process")]
+        #[cfg(any(
+            target_arch = "x86_64",
+            target_arch = "x86",
+            target_arch = "riscv64",
+            target_arch = "powerpc64"
+        ))]
+        {
+            // Look up the platform-specific `getcpu` symbol as documented
+            // [here].
+            //
+            // [here]: https://man7.org/linux/man-pages/man7/vdso.7.html
+            #[cfg(target_arch = "x86_64")]
+            let ptr = vdso.sym(cstr!("LINUX_2.6"), cstr!("__vdso_getcpu"));
+            #[cfg(target_arch = "x86")]
+            let ptr = vdso.sym(cstr!("LINUX_2.6"), cstr!("__vdso_getcpu"));
+            #[cfg(target_arch = "riscv64")]
+            let ptr = vdso.sym(cstr!("LINUX_4.15"), cstr!("__kernel_getcpu"));
+            #[cfg(target_arch = "powerpc64")]
+            let ptr = vdso.sym(cstr!("LINUX_2.6.15"), cstr!("__kernel_getcpu"));
+
+            #[cfg(any(
+                target_arch = "x86_64",
+                target_arch = "riscv64",
+                target_arch = "powerpc64"
+            ))]
+            let ok = true;
+
+            // On 32-bit x86, the symbol doesn't appear present sometimes.
+            #[cfg(target_arch = "x86")]
+            let ok = !ptr.is_null();
+
+            #[cfg(any(
+                target_arch = "aarch64",
+                target_arch = "arm",
+                target_arch = "mips",
+                target_arch = "mips32r6",
+                target_arch = "mips64",
+                target_arch = "mips64r6"
+            ))]
+            let ok = false;
+
+            if ok {
+                assert!(!ptr.is_null());
+
+                // SAFETY: Store the computed function addresses in static
+                // storage so that we don't need to compute it again (but if
+                // we do, it doesn't hurt anything).
+                unsafe {
+                    GETCPU.store(ptr.cast(), Relaxed);
+                }
+            }
+        }
+
         // On x86, also look up the vsyscall entry point.
         #[cfg(target_arch = "x86")]
         {

src/process/mod.rs

@@ -32,7 +32,7 @@ mod procctl;
     target_os = "wasi"
 )))]
 mod rlimit;
-#[cfg(any(linux_kernel, target_os = "dragonfly", target_os = "fuchsia"))]
+#[cfg(any(freebsdlike, linux_kernel, target_os = "fuchsia"))]
 mod sched;
 mod sched_yield;
 #[cfg(not(target_os = "wasi"))] // WASI doesn't have umask.
@@ -71,7 +71,7 @@ pub use procctl::*;
     target_os = "wasi"
 )))]
 pub use rlimit::*;
-#[cfg(any(linux_kernel, target_os = "dragonfly", target_os = "fuchsia"))]
+#[cfg(any(freebsdlike, linux_kernel, target_os = "fuchsia"))]
 pub use sched::*;
 pub use sched_yield::sched_yield;
 #[cfg(not(target_os = "wasi"))]

src/process/sched.rs

@@ -108,3 +108,18 @@ pub fn sched_getaffinity(pid: Option<Pid>) -> io::Result<CpuSet> {
     let mut cpuset = CpuSet::new();
     backend::process::syscalls::sched_getaffinity(pid, &mut cpuset.cpu_set).and(Ok(cpuset))
 }
+
+/// `sched_getcpu()`—Get the CPU that the current thread is currently on.
+///
+/// # References
+///  - [Linux]
+///  - [DragonFly BSD]
+///
+/// [Linux]: https://man7.org/linux/man-pages/man2/sched_getcpu.2.html
+/// [DragonFly BSD]: https://man.dragonflybsd.org/?command=sched_getcpu&section=2
+// FreeBSD added `sched_getcpu` in 13.0.
+#[cfg(any(linux_kernel, target_os = "dragonfly"))]
+#[inline]
+pub fn sched_getcpu() -> usize {
+    backend::process::syscalls::sched_getcpu()
+}

tests/process/main.rs

@@ -19,7 +19,7 @@ mod priority;
 mod procctl;
 #[cfg(not(any(target_os = "fuchsia", target_os = "redox", target_os = "wasi")))]
 mod rlimit;
-mod sched_yield;
+mod sched;
 #[cfg(not(target_os = "wasi"))] // WASI doesn't have umask.
 mod umask;
 #[cfg(not(any(target_os = "espidf", target_os = "wasi")))] // WASI doesn't have waitpid.

tests/process/sched.rs

@@ -0,0 +1,14 @@
+use rustix::process::sched_yield;
+
+#[test]
+fn test_sched_yield() {
+    // Just make sure we can call it.
+    sched_yield();
+}
+
+#[cfg(any(linux_kernel, target_os = "dragonfly"))]
+#[test]
+fn test_sched_getcpu() {
+    let n = rustix::process::sched_getcpu();
+    assert!(n < rustix::process::CpuSet::MAX_CPU);
+}