Skip to content

Commit 2a1e88e

Browse files
youknowoneyouknowone
committed
rework SchedParam
1 parent 811801b commit 2a1e88e

File tree

1 file changed

+166
-204
lines changed

1 file changed

+166
-204
lines changed

‎crates/vm/src/stdlib/posix.rs‎

Lines changed: 166 additions & 204 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,20 @@ pub fn set_inheritable(fd: BorrowedFd<'_>, inheritable: bool) -> nix::Result<()>
1515
Ok(())
1616
}
1717

18-
#[pymodule(name = "posix", with(super::os::_os))]
18+
#[pymodule(name = "posix", with(
19+
super::os::_os,
20+
#[cfg(any(
21+
target_os = "linux",
22+
target_os = "netbsd",
23+
target_os = "freebsd",
24+
target_os = "android"
25+
))]
26+
posix_sched
27+
))]
1928
pub mod module {
2029
use crate::{
21-
AsObject, Py, PyObjectRef, PyPayload, PyResult, VirtualMachine,
22-
builtins::{PyDictRef, PyInt, PyListRef, PyStr, PyTupleRef, PyType},
30+
AsObject, Py, PyObjectRef, PyResult, VirtualMachine,
31+
builtins::{PyDictRef, PyInt, PyListRef, PyStr, PyTupleRef},
2332
convert::{IntoPyException, ToPyObject, TryFromObject},
2433
exceptions::OSErrorBuilder,
2534
function::{Either, KwArgs, OptionalArg},
@@ -28,7 +37,6 @@ pub mod module {
2837
_os, DirFd, FollowSymlinks, SupportFunc, TargetIsDirectory, fs_metadata,
2938
warn_if_bool_fd,
3039
},
31-
types::{Constructor, Representable},
3240
};
3341
#[cfg(any(
3442
target_os = "android",
@@ -889,206 +897,6 @@ pub mod module {
889897
nix::sched::sched_yield().map_err(|e| e.into_pyexception(vm))
890898
}
891899

892-
#[pyattr]
893-
#[pyclass(name = "sched_param")]
894-
#[derive(Debug, PyPayload)]
895-
struct SchedParam {
896-
sched_priority: PyObjectRef,
897-
}
898-
899-
impl TryFromObject for SchedParam {
900-
fn try_from_object(_vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
901-
Ok(Self {
902-
sched_priority: obj,
903-
})
904-
}
905-
}
906-
907-
#[pyclass(with(Constructor, Representable))]
908-
impl SchedParam {
909-
#[pygetset]
910-
fn sched_priority(&self, vm: &VirtualMachine) -> PyObjectRef {
911-
self.sched_priority.clone().to_pyobject(vm)
912-
}
913-
914-
#[pymethod]
915-
fn __reduce__(zelf: crate::PyRef<Self>, vm: &VirtualMachine) -> PyTupleRef {
916-
vm.new_tuple((zelf.class().to_owned(), (zelf.sched_priority.clone(),)))
917-
}
918-
919-
#[pymethod]
920-
fn __replace__(
921-
zelf: crate::PyRef<Self>,
922-
args: crate::function::FuncArgs,
923-
vm: &VirtualMachine,
924-
) -> PyResult<Self> {
925-
if !args.args.is_empty() {
926-
return Err(
927-
vm.new_type_error("__replace__() takes no positional arguments".to_owned())
928-
);
929-
}
930-
let sched_priority = match args.kwargs.get("sched_priority") {
931-
Some(v) => v.clone(),
932-
None => zelf.sched_priority.clone(),
933-
};
934-
// Check for unexpected keyword arguments
935-
for key in args.kwargs.keys() {
936-
if key.as_str() != "sched_priority" {
937-
return Err(vm.new_type_error(format!(
938-
"__replace__() got an unexpected keyword argument '{key}'"
939-
)));
940-
}
941-
}
942-
Ok(Self { sched_priority })
943-
}
944-
945-
#[cfg(any(
946-
target_os = "linux",
947-
target_os = "netbsd",
948-
target_os = "freebsd",
949-
target_os = "android"
950-
))]
951-
#[cfg(not(target_env = "musl"))]
952-
fn try_to_libc(&self, vm: &VirtualMachine) -> PyResult<libc::sched_param> {
953-
use crate::AsObject;
954-
let priority_class = self.sched_priority.class();
955-
let priority_type = priority_class.name();
956-
let priority = self.sched_priority.clone();
957-
let value = priority.downcast::<PyInt>().map_err(|_| {
958-
vm.new_type_error(format!("an integer is required (got type {priority_type})"))
959-
})?;
960-
let sched_priority = value.try_to_primitive(vm)?;
961-
Ok(libc::sched_param { sched_priority })
962-
}
963-
}
964-
965-
#[derive(FromArgs)]
966-
pub struct SchedParamArg {
967-
sched_priority: PyObjectRef,
968-
}
969-
970-
impl Constructor for SchedParam {
971-
type Args = SchedParamArg;
972-
973-
fn py_new(_cls: &Py<PyType>, arg: Self::Args, _vm: &VirtualMachine) -> PyResult<Self> {
974-
Ok(Self {
975-
sched_priority: arg.sched_priority,
976-
})
977-
}
978-
}
979-
980-
impl Representable for SchedParam {
981-
#[inline]
982-
fn repr_str(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<String> {
983-
let sched_priority_repr = zelf.sched_priority.repr(vm)?;
984-
Ok(format!(
985-
"posix.sched_param(sched_priority = {})",
986-
sched_priority_repr.as_str()
987-
))
988-
}
989-
}
990-
991-
#[cfg(any(
992-
target_os = "linux",
993-
target_os = "netbsd",
994-
target_os = "freebsd",
995-
target_os = "android"
996-
))]
997-
#[pyfunction]
998-
fn sched_getscheduler(pid: libc::pid_t, vm: &VirtualMachine) -> PyResult<i32> {
999-
let policy = unsafe { libc::sched_getscheduler(pid) };
1000-
if policy == -1 {
1001-
Err(vm.new_last_errno_error())
1002-
} else {
1003-
Ok(policy)
1004-
}
1005-
}
1006-
1007-
#[cfg(any(
1008-
target_os = "linux",
1009-
target_os = "netbsd",
1010-
target_os = "freebsd",
1011-
target_os = "android"
1012-
))]
1013-
#[derive(FromArgs)]
1014-
struct SchedSetschedulerArgs {
1015-
#[pyarg(positional)]
1016-
pid: i32,
1017-
#[pyarg(positional)]
1018-
policy: i32,
1019-
#[pyarg(positional)]
1020-
sched_param_obj: crate::PyRef<SchedParam>,
1021-
}
1022-
1023-
#[cfg(any(
1024-
target_os = "linux",
1025-
target_os = "netbsd",
1026-
target_os = "freebsd",
1027-
target_os = "android"
1028-
))]
1029-
#[cfg(not(target_env = "musl"))]
1030-
#[pyfunction]
1031-
fn sched_setscheduler(args: SchedSetschedulerArgs, vm: &VirtualMachine) -> PyResult<i32> {
1032-
let libc_sched_param = args.sched_param_obj.try_to_libc(vm)?;
1033-
let policy = unsafe { libc::sched_setscheduler(args.pid, args.policy, &libc_sched_param) };
1034-
if policy == -1 {
1035-
Err(vm.new_last_errno_error())
1036-
} else {
1037-
Ok(policy)
1038-
}
1039-
}
1040-
#[cfg(any(
1041-
target_os = "linux",
1042-
target_os = "netbsd",
1043-
target_os = "freebsd",
1044-
target_os = "android"
1045-
))]
1046-
#[pyfunction]
1047-
fn sched_getparam(pid: libc::pid_t, vm: &VirtualMachine) -> PyResult<SchedParam> {
1048-
let param = unsafe {
1049-
let mut param = core::mem::MaybeUninit::uninit();
1050-
if -1 == libc::sched_getparam(pid, param.as_mut_ptr()) {
1051-
return Err(vm.new_last_errno_error());
1052-
}
1053-
param.assume_init()
1054-
};
1055-
Ok(SchedParam {
1056-
sched_priority: param.sched_priority.to_pyobject(vm),
1057-
})
1058-
}
1059-
1060-
#[cfg(any(
1061-
target_os = "linux",
1062-
target_os = "netbsd",
1063-
target_os = "freebsd",
1064-
target_os = "android"
1065-
))]
1066-
#[derive(FromArgs)]
1067-
struct SchedSetParamArgs {
1068-
#[pyarg(positional)]
1069-
pid: i32,
1070-
#[pyarg(positional)]
1071-
sched_param_obj: crate::PyRef<SchedParam>,
1072-
}
1073-
1074-
#[cfg(any(
1075-
target_os = "linux",
1076-
target_os = "netbsd",
1077-
target_os = "freebsd",
1078-
target_os = "android"
1079-
))]
1080-
#[cfg(not(target_env = "musl"))]
1081-
#[pyfunction]
1082-
fn sched_setparam(args: SchedSetParamArgs, vm: &VirtualMachine) -> PyResult<i32> {
1083-
let libc_sched_param = args.sched_param_obj.try_to_libc(vm)?;
1084-
let ret = unsafe { libc::sched_setparam(args.pid, &libc_sched_param) };
1085-
if ret == -1 {
1086-
Err(vm.new_last_errno_error())
1087-
} else {
1088-
Ok(ret)
1089-
}
1090-
}
1091-
1092900
#[pyfunction]
1093901
fn get_inheritable(fd: BorrowedFd<'_>, vm: &VirtualMachine) -> PyResult<bool> {
1094902
let flags = fcntl::fcntl(fd, fcntl::FcntlArg::F_GETFD);
@@ -2699,3 +2507,157 @@ pub mod module {
26992507
Ok(())
27002508
}
27012509
}
2510+
2511+
#[cfg(any(
2512+
target_os = "linux",
2513+
target_os = "netbsd",
2514+
target_os = "freebsd",
2515+
target_os = "android"
2516+
))]
2517+
#[pymodule(sub)]
2518+
mod posix_sched {
2519+
use crate::{
2520+
AsObject, Py, PyObjectRef, PyResult, VirtualMachine,
2521+
builtins::{PyInt, PyTupleRef},
2522+
convert::{IntoPyException, ToPyObject},
2523+
function::FuncArgs,
2524+
types::PyStructSequence,
2525+
};
2526+
2527+
#[pystruct_sequence_data]
2528+
struct SchedParamData {
2529+
pub sched_priority: PyObjectRef,
2530+
}
2531+
2532+
#[pyattr]
2533+
#[pystruct_sequence(name = "sched_param", module = "posix", data = "SchedParamData")]
2534+
struct PySchedParam;
2535+
2536+
#[pyclass(with(PyStructSequence))]
2537+
impl PySchedParam {
2538+
#[pyslot]
2539+
fn slot_new(
2540+
cls: crate::builtins::PyTypeRef,
2541+
args: FuncArgs,
2542+
vm: &VirtualMachine,
2543+
) -> PyResult {
2544+
use crate::PyPayload;
2545+
let sched_priority: PyObjectRef = args.bind(vm)?;
2546+
let items = vec![sched_priority];
2547+
crate::builtins::PyTuple::new_unchecked(items.into_boxed_slice())
2548+
.into_ref_with_type(vm, cls)
2549+
.map(Into::into)
2550+
}
2551+
2552+
#[extend_class]
2553+
fn extend_pyclass(ctx: &crate::vm::Context, class: &'static Py<crate::builtins::PyType>) {
2554+
// Override __reduce__ to return (type, (sched_priority,))
2555+
// instead of the generic structseq (type, ((sched_priority,),)).
2556+
// The trait's extend_class checks contains_key before setting default.
2557+
const SCHED_PARAM_REDUCE: crate::function::PyMethodDef =
2558+
crate::function::PyMethodDef::new_const(
2559+
"__reduce__",
2560+
|zelf: crate::PyRef<crate::builtins::PyTuple>,
2561+
vm: &VirtualMachine|
2562+
-> PyTupleRef {
2563+
vm.new_tuple((zelf.class().to_owned(), (zelf[0].clone(),)))
2564+
},
2565+
crate::function::PyMethodFlags::METHOD,
2566+
None,
2567+
);
2568+
class.set_attr(
2569+
ctx.intern_str("__reduce__"),
2570+
SCHED_PARAM_REDUCE.to_proper_method(class, ctx),
2571+
);
2572+
}
2573+
}
2574+
2575+
#[cfg(not(target_env = "musl"))]
2576+
fn convert_sched_param(
2577+
obj: &PyObjectRef,
2578+
vm: &VirtualMachine,
2579+
) -> PyResult<libc::sched_param> {
2580+
use crate::{builtins::PyTuple, class::StaticType};
2581+
if !obj.fast_isinstance(PySchedParam::static_type()) {
2582+
return Err(vm.new_type_error("must have a sched_param object".to_owned()));
2583+
}
2584+
let tuple = obj.downcast_ref::<PyTuple>().unwrap();
2585+
let priority = tuple[0].clone();
2586+
let priority_type = priority.class().name().to_string();
2587+
let value = priority.downcast::<PyInt>().map_err(|_| {
2588+
vm.new_type_error(format!("an integer is required (got type {priority_type})"))
2589+
})?;
2590+
let sched_priority = value.try_to_primitive(vm)?;
2591+
Ok(libc::sched_param { sched_priority })
2592+
}
2593+
2594+
#[pyfunction]
2595+
fn sched_getscheduler(pid: libc::pid_t, vm: &VirtualMachine) -> PyResult<i32> {
2596+
let policy = unsafe { libc::sched_getscheduler(pid) };
2597+
if policy == -1 {
2598+
Err(vm.new_last_errno_error())
2599+
} else {
2600+
Ok(policy)
2601+
}
2602+
}
2603+
2604+
#[derive(FromArgs)]
2605+
struct SchedSetschedulerArgs {
2606+
#[pyarg(positional)]
2607+
pid: i32,
2608+
#[pyarg(positional)]
2609+
policy: i32,
2610+
#[pyarg(positional)]
2611+
sched_param: PyObjectRef,
2612+
}
2613+
2614+
#[cfg(not(target_env = "musl"))]
2615+
#[pyfunction]
2616+
fn sched_setscheduler(args: SchedSetschedulerArgs, vm: &VirtualMachine) -> PyResult<i32> {
2617+
let libc_sched_param = convert_sched_param(&args.sched_param, vm)?;
2618+
let policy = unsafe { libc::sched_setscheduler(args.pid, args.policy, &libc_sched_param) };
2619+
if policy == -1 {
2620+
Err(vm.new_last_errno_error())
2621+
} else {
2622+
Ok(policy)
2623+
}
2624+
}
2625+
2626+
#[pyfunction]
2627+
fn sched_getparam(pid: libc::pid_t, vm: &VirtualMachine) -> PyResult<PyTupleRef> {
2628+
let param = unsafe {
2629+
let mut param = core::mem::MaybeUninit::uninit();
2630+
if -1 == libc::sched_getparam(pid, param.as_mut_ptr()) {
2631+
return Err(vm.new_last_errno_error());
2632+
}
2633+
param.assume_init()
2634+
};
2635+
Ok(PySchedParam::from_data(
2636+
SchedParamData {
2637+
sched_priority: param.sched_priority.to_pyobject(vm),
2638+
},
2639+
vm,
2640+
))
2641+
}
2642+
2643+
#[derive(FromArgs)]
2644+
struct SchedSetParamArgs {
2645+
#[pyarg(positional)]
2646+
pid: i32,
2647+
#[pyarg(positional)]
2648+
sched_param: PyObjectRef,
2649+
}
2650+
2651+
#[cfg(not(target_env = "musl"))]
2652+
#[pyfunction]
2653+
fn sched_setparam(args: SchedSetParamArgs, vm: &VirtualMachine) -> PyResult<i32> {
2654+
let libc_sched_param = convert_sched_param(&args.sched_param, vm)?;
2655+
let ret = unsafe { libc::sched_setparam(args.pid, &libc_sched_param) };
2656+
if ret == -1 {
2657+
Err(vm.new_last_errno_error())
2658+
} else {
2659+
Ok(ret)
2660+
}
2661+
}
2662+
}
2663+

0 commit comments

Comments
 (0)