From 4feb8cebea4374a73218879fea7472772187a7f8 Mon Sep 17 00:00:00 2001 From: Misaki Kasumi Date: Wed, 23 Oct 2024 16:02:38 +0800 Subject: [PATCH] Add `sched_getscheduler` & `sched_setscheduler` --- src/backend/libc/process/syscalls.rs | 27 +++++++++++++ src/backend/libc/process/types.rs | 46 +++++++++++++++++++++++ src/backend/linux_raw/process/syscalls.rs | 29 +++++++++++++- src/backend/linux_raw/process/types.rs | 35 +++++++++++++++++ src/process/sched.rs | 42 +++++++++++++++++++++ tests/process/sched.rs | 22 +++++++++++ 6 files changed, 200 insertions(+), 1 deletion(-) diff --git a/src/backend/libc/process/syscalls.rs b/src/backend/libc/process/syscalls.rs index efb5a77f0..d0f8da3ce 100644 --- a/src/backend/libc/process/syscalls.rs +++ b/src/backend/libc/process/syscalls.rs @@ -2,6 +2,8 @@ #[cfg(any(freebsdlike, linux_kernel, target_os = "fuchsia"))] use super::types::RawCpuSet; +#[cfg(any(freebsdlike, linux_kernel))] +use super::types::{SchedParam, SchedPolicy}; use crate::backend::c; #[cfg(not(any(target_os = "wasi", target_os = "fuchsia")))] use crate::backend::conv::borrowed_fd; @@ -66,6 +68,8 @@ use crate::process::{Resource, Rlimit}; target_os = "wasi" )))] use crate::process::{WaitId, WaitidOptions, WaitidStatus}; +#[cfg(any(freebsdlike, linux_kernel))] +use bitflags::Flags; use core::mem::MaybeUninit; #[cfg(target_os = "linux")] use { @@ -213,6 +217,29 @@ pub(crate) fn sched_setaffinity(pid: Option, cpuset: &RawCpuSet) -> io::Res } } +#[cfg(any(freebsdlike, linux_kernel))] +#[inline] +pub(crate) fn sched_getscheduler(pid: Option) -> io::Result { + unsafe { ret_c_int(c::sched_getscheduler(Pid::as_raw(pid) as _)) } + .map(|r| SchedPolicy::from_bits_retain(r as _)) +} + +#[cfg(any(freebsdlike, linux_kernel))] +#[inline] +pub(crate) fn sched_setscheduler( + pid: Option, + policy: SchedPolicy, + param: &SchedParam, +) -> io::Result<()> { + unsafe { + ret(c::sched_setscheduler( + Pid::as_raw(pid) as _, + policy.bits() as _, + core::ptr::from_ref(param) as _, + )) + } +} + #[inline] pub(crate) fn sched_yield() { unsafe { diff --git a/src/backend/libc/process/types.rs b/src/backend/libc/process/types.rs index f914128df..b7b5e157f 100644 --- a/src/backend/libc/process/types.rs +++ b/src/backend/libc/process/types.rs @@ -1,5 +1,7 @@ #[cfg(not(any(target_os = "espidf", target_os = "vita")))] use crate::backend::c; +#[cfg(any(freebsdlike, linux_kernel, target_os = "fuchsia"))] +use bitflags::bitflags; /// A command for use with [`membarrier`] and [`membarrier_cpu`]. /// @@ -170,3 +172,47 @@ pub(crate) fn raw_cpu_set_new() -> RawCpuSet { #[cfg(any(freebsdlike, linux_kernel, target_os = "fuchsia"))] pub(crate) const CPU_SETSIZE: usize = c::CPU_SETSIZE as usize; + +///`SCHED_*` constants for use with [`sched_getscheduler`] and [`sched_setscheduler`]. +#[cfg(any(freebsdlike, linux_kernel))] +bitflags! { + #[repr(transparent)] + #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)] + pub struct SchedPolicy: u32 { + /// `SCHED_OTHER` + const Other = c::SCHED_OTHER as u32; + /// `SCHED_FIFO` + const FIFO = c::SCHED_FIFO as u32; + /// `SCHED_RR` + const RoundRobin = c::SCHED_RR as u32; + /// `SCHED_BATCH` + #[cfg(linux_kernel)] + const Batch = c::SCHED_BATCH as u32; + /// `SCHED_IDLE` + #[cfg(linux_kernel)] + const Idle = c::SCHED_IDLE as u32; + /// `SCHED_DEADLINE` + #[cfg(linux_kernel)] + const Deadline = c::SCHED_DEADLINE as u32; + /// `SCHED_RESET_ON_FORK` + #[cfg(linux_kernel)] + const ResetOnFork = c::SCHED_RESET_ON_FORK as u32; + /// + const _ = !0; + } +} + +impl Default for SchedPolicy { + fn default() -> Self { + Self::Other + } +} + +/// `sched_param` for use with [`sched_getscheduler`] and [`sched_setscheduler`]. +#[cfg(any(freebsdlike, linux_kernel))] +#[derive(Clone, Debug, Default)] +#[repr(C)] +pub struct SchedParam { + /// Scheduling priority. + pub sched_priority: i32, +} diff --git a/src/backend/linux_raw/process/syscalls.rs b/src/backend/linux_raw/process/syscalls.rs index 85c6fbbad..8f01c7ebb 100644 --- a/src/backend/linux_raw/process/syscalls.rs +++ b/src/backend/linux_raw/process/syscalls.rs @@ -5,7 +5,7 @@ //! See the `rustix::backend` module documentation for details. #![allow(unsafe_code, clippy::undocumented_unsafe_blocks)] -use super::types::RawCpuSet; +use super::types::{RawCpuSet, SchedParam, SchedPolicy}; use crate::backend::c; #[cfg(all(feature = "alloc", feature = "fs"))] use crate::backend::conv::slice_mut; @@ -205,6 +205,33 @@ pub(crate) fn sched_yield() { } } +#[inline] +pub(crate) fn sched_getscheduler(pid: Option) -> io::Result { + unsafe { + ret_c_int(syscall_readonly!( + __NR_sched_getscheduler, + c_int(Pid::as_raw(pid)) + )) + } + .map(|policy| SchedPolicy::from_bits_retain(bitcast!(policy))) +} + +#[inline] +pub(crate) fn sched_setscheduler( + pid: Option, + policy: SchedPolicy, + param: &SchedParam, +) -> io::Result<()> { + unsafe { + ret(syscall_readonly!( + __NR_sched_setscheduler, + c_int(Pid::as_raw(pid)), + c_uint(policy.bits()), + core::ptr::from_ref(param) + )) + } +} + #[cfg(feature = "fs")] #[inline] pub(crate) fn umask(mode: Mode) -> Mode { diff --git a/src/backend/linux_raw/process/types.rs b/src/backend/linux_raw/process/types.rs index 841668a60..493e6452e 100644 --- a/src/backend/linux_raw/process/types.rs +++ b/src/backend/linux_raw/process/types.rs @@ -1,3 +1,4 @@ +use bitflags::bitflags; use linux_raw_sys::general::membarrier_cmd; /// A command for use with [`membarrier`] and [`membarrier_cpu`]. @@ -102,3 +103,37 @@ pub(crate) fn raw_cpu_set_new() -> RawCpuSet { } pub(crate) const CPU_SETSIZE: usize = 8 * core::mem::size_of::(); + +///`SCHED_*` constants for use with [`sched_getscheduler`] and [`sched_setscheduler`]. +bitflags! { + #[repr(transparent)] + #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, Default)] + pub struct SchedPolicy: u32 { + /// `SCHED_OTHER` + const Other = 0; + /// `SCHED_FIFO` + const FIFO = 1; + /// `SCHED_RR` + const RoundRobin = 2; + /// `SCHED_BATCH` + const Batch = 3; + /// `SCHED_ISO`; only presented in some patched kernels. + const Isochronous = 4; + /// `SCHED_IDLE` + const Idle = 5; + /// `SCHED_DEADLINE` + const Deadline = 6; + /// `SCHED_RESET_ON_FORK` + const ResetOnFork = 0x40000000; + /// + const _ = !0; + } +} + +/// `sched_param` for use with [`sched_getscheduler`] and [`sched_setscheduler`]. +#[derive(Clone, Debug, Default)] +#[repr(C)] +pub struct SchedParam { + /// Scheduling priority. + pub sched_priority: i32, +} diff --git a/src/process/sched.rs b/src/process/sched.rs index f9614e8a0..6277f92e1 100644 --- a/src/process/sched.rs +++ b/src/process/sched.rs @@ -159,3 +159,45 @@ pub fn sched_getaffinity(pid: Option) -> io::Result { pub fn sched_getcpu() -> usize { backend::process::syscalls::sched_getcpu() } + +pub use backend::process::types::{SchedParam, SchedPolicy}; + +/// `sched_getscheduler`—Get a thread's current scheduling policy. +/// +/// `pid` is the thread ID to check. If pid is `None`, then the current thread +/// is checked. +/// +/// # References +/// - [Linux] +/// - [FreeBSD] +/// +/// [Linux]: https://www.man7.org/linux/man-pages/man2/sched_setscheduler.2.html +/// [FreeBSD]: https://man.freebsd.org/cgi/man.cgi?query=sched_setscheduler +/// Other *BSD also support `sched_getscheduler`. +#[cfg(any(freebsdlike, linux_kernel))] +#[inline] +pub fn sched_getscheduler(pid: Option) -> io::Result { + backend::process::syscalls::sched_getscheduler(pid) +} + +/// `sched_setscheduler`—Set a thread's scheduling policy and parameters. +/// +/// `pid` is the thread ID to set scheduling policy and parameters. +/// If pid is `None`, then the target is the current thread. +/// +/// # References +/// - [Linux] +/// - [FreeBSD] +/// +/// [Linux]: https://www.man7.org/linux/man-pages/man2/sched_setscheduler.2.html +/// [FreeBSD]: https://man.freebsd.org/cgi/man.cgi?query=sched_setscheduler +/// Other *BSD also support `sched_setscheduler`. +#[cfg(any(freebsdlike, linux_kernel))] +#[inline] +pub fn sched_setscheduler( + pid: Option, + policy: SchedPolicy, + param: &SchedParam, +) -> io::Result<()> { + backend::process::syscalls::sched_setscheduler(pid, policy, param) +} diff --git a/tests/process/sched.rs b/tests/process/sched.rs index 1aacc8e42..5f6bf27f3 100644 --- a/tests/process/sched.rs +++ b/tests/process/sched.rs @@ -12,3 +12,25 @@ fn test_sched_getcpu() { let n = rustix::process::sched_getcpu(); assert!(n < rustix::process::CpuSet::MAX_CPU); } + +#[cfg(any(freebsdlike, linux_kernel))] +#[test] +fn test_sched_scheduler() { + use rustix::process::{SchedPolicy, SchedParam}; + + let policy = if cfg!(linux_kernel) { + SchedPolicy::Batch + } else { + // we cannot change policy in *BSD because we do not have priviledge + SchedPolicy::default() + }; + + // backup + let original_policy = rustix::process::sched_getscheduler(None).unwrap(); + + rustix::process::sched_setscheduler(None, policy, &SchedParam::default()).unwrap(); + assert_eq!(rustix::process::sched_getscheduler(None).unwrap(), policy); + + // restore + rustix::process::sched_setscheduler(None, original_policy, &SchedParam::default()).unwrap(); +}