Skip to content

Commit

Permalink
Add sched_getscheduler & sched_setscheduler
Browse files Browse the repository at this point in the history
  • Loading branch information
ruihe774 committed Oct 23, 2024
1 parent 7d4a719 commit dba5d0d
Show file tree
Hide file tree
Showing 6 changed files with 203 additions and 1 deletion.
25 changes: 25 additions & 0 deletions src/backend/libc/process/syscalls.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -213,6 +215,29 @@ pub(crate) fn sched_setaffinity(pid: Option<Pid>, cpuset: &RawCpuSet) -> io::Res
}
}

#[cfg(any(freebsdlike, linux_kernel))]
#[inline]
pub(crate) fn sched_getscheduler(pid: Option<Pid>) -> io::Result<SchedPolicy> {
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<Pid>,
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 {
Expand Down
51 changes: 51 additions & 0 deletions src/backend/libc/process/types.rs
Original file line number Diff line number Diff line change
@@ -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`].
///
Expand Down Expand Up @@ -170,3 +172,52 @@ 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;

#[cfg(any(freebsdlike, linux_kernel))]
bitflags! {
/// `SCHED_*` constants for use with [`sched_getscheduler`] and [`sched_setscheduler`].
#[repr(transparent)]
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
pub struct SchedPolicy: u32 {
/// `SCHED_OTHER`
#[cfg(not(target_os = "android"))] // not defined in android
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(target_os = "linux")] // not defined in android
const Deadline = c::SCHED_DEADLINE as u32;
/// `SCHED_RESET_ON_FORK`
#[cfg(target_os = "linux")] // not defined in android
const ResetOnFork = c::SCHED_RESET_ON_FORK as u32;
/// <https://docs.rs/bitflags/*/bitflags/#externally-defined-flags>
const _ = !0;
}
}

#[cfg(any(freebsdlike, linux_kernel))]
impl Default for SchedPolicy {
fn default() -> Self {
#[cfg(not(target_os = "android"))]
return Self::Other;
#[cfg(target_os = "android")]
return Self::from_bits_retain(0);
}
}

/// `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,
}
29 changes: 28 additions & 1 deletion src/backend/linux_raw/process/syscalls.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -205,6 +205,33 @@ pub(crate) fn sched_yield() {
}
}

#[inline]
pub(crate) fn sched_getscheduler(pid: Option<Pid>) -> io::Result<SchedPolicy> {
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<Pid>,
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 {
Expand Down
35 changes: 35 additions & 0 deletions src/backend/linux_raw/process/types.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use bitflags::bitflags;
use linux_raw_sys::general::membarrier_cmd;

/// A command for use with [`membarrier`] and [`membarrier_cpu`].
Expand Down Expand Up @@ -102,3 +103,37 @@ pub(crate) fn raw_cpu_set_new() -> RawCpuSet {
}

pub(crate) const CPU_SETSIZE: usize = 8 * core::mem::size_of::<RawCpuSet>();

bitflags! {
/// `SCHED_*` constants for use with [`sched_getscheduler`] and [`sched_setscheduler`].
#[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;
/// <https://docs.rs/bitflags/*/bitflags/#externally-defined-flags>
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,
}
42 changes: 42 additions & 0 deletions src/process/sched.rs
Original file line number Diff line number Diff line change
Expand Up @@ -159,3 +159,45 @@ pub fn sched_getaffinity(pid: Option<Pid>) -> io::Result<CpuSet> {
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<Pid>) -> io::Result<SchedPolicy> {
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<Pid>,
policy: SchedPolicy,
param: &SchedParam,
) -> io::Result<()> {
backend::process::syscalls::sched_setscheduler(pid, policy, param)
}
22 changes: 22 additions & 0 deletions tests/process/sched.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::{SchedParam, SchedPolicy};

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();
}

0 comments on commit dba5d0d

Please sign in to comment.