diff --git a/vm/src/stdlib/os.rs b/vm/src/stdlib/os.rs index bcc50da1b4..c90bc1b981 100644 --- a/vm/src/stdlib/os.rs +++ b/vm/src/stdlib/os.rs @@ -2494,12 +2494,140 @@ mod posix { sched_priority_repr.as_str() )) } + + #[cfg(any( + target_os = "linux", + target_os = "netbsd", + target_os = "freebsd", + target_os = "android" + ))] + fn try_to_libc(&self, vm: &VirtualMachine) -> PyResult { + let priority = self.sched_priority.clone(); + let priority_type = priority.class().name(); + let value = priority.downcast::().map_err(|_| { + vm.new_type_error(format!( + "an integer is required (got type {})", + priority_type + )) + })?; + let sched_priority = int::try_to_primitive(value.as_bigint(), vm)?; + Ok(libc::sched_param { sched_priority }) + } + } + + #[derive(FromArgs)] + pub struct SchedParamArg { + #[pyarg(any)] + sched_priority: PyObjectRef, } impl SlotConstructor for SchedParam { - type Args = SchedParam; - fn py_new(cls: PyTypeRef, sched_param: Self::Args, vm: &VirtualMachine) -> PyResult { - sched_param.into_pyresult_with_type(vm, cls) + type Args = SchedParamArg; + fn py_new(cls: PyTypeRef, arg: Self::Args, vm: &VirtualMachine) -> PyResult { + SchedParam { + sched_priority: arg.sched_priority, + } + .into_pyresult_with_type(vm, cls) + } + } + + #[cfg(any( + target_os = "linux", + target_os = "netbsd", + target_os = "freebsd", + target_os = "android" + ))] + #[derive(FromArgs)] + struct SchedSetschedulerArgs { + #[pyarg(positional)] + pid: i32, + #[pyarg(positional)] + policy: i32, + #[pyarg(positional)] + sched_param_obj: PyRef, + } + + #[cfg(any( + target_os = "linux", + target_os = "netbsd", + target_os = "freebsd", + target_os = "android" + ))] + #[pyfunction] + fn sched_getscheduler(pid: libc::pid_t, vm: &VirtualMachine) -> PyResult { + let policy = unsafe { libc::sched_getscheduler(pid) }; + if policy == -1 { + Err(errno_err(vm)) + } else { + Ok(policy) + } + } + + #[cfg(any( + target_os = "linux", + target_os = "netbsd", + target_os = "freebsd", + target_os = "android" + ))] + #[pyfunction] + fn sched_setscheduler(args: SchedSetschedulerArgs, vm: &VirtualMachine) -> PyResult { + let libc_sched_param = args.sched_param_obj.try_to_libc(vm)?; + let policy = unsafe { libc::sched_setscheduler(args.pid, args.policy, &libc_sched_param) }; + if policy == -1 { + Err(errno_err(vm)) + } else { + Ok(policy) + } + } + + #[cfg(any( + target_os = "linux", + target_os = "netbsd", + target_os = "freebsd", + target_os = "android" + ))] + #[pyfunction] + fn sched_getparam(pid: libc::pid_t, vm: &VirtualMachine) -> PyResult { + let param = unsafe { + let mut param = std::mem::MaybeUninit::uninit(); + if -1 == libc::sched_getparam(pid, param.as_mut_ptr()) { + return Err(errno_err(vm)); + } + param.assume_init() + }; + Ok(SchedParam { + sched_priority: param.sched_priority.into_pyobject(vm), + }) + } + + #[cfg(any( + target_os = "linux", + target_os = "netbsd", + target_os = "freebsd", + target_os = "android" + ))] + #[derive(FromArgs)] + struct SchedSetParamArgs { + #[pyarg(positional)] + pid: i32, + #[pyarg(positional)] + sched_param_obj: PyRef, + } + + #[cfg(any( + target_os = "linux", + target_os = "netbsd", + target_os = "freebsd", + target_os = "android" + ))] + #[pyfunction] + fn sched_setparam(args: SchedSetParamArgs, vm: &VirtualMachine) -> PyResult { + let libc_sched_param = args.sched_param_obj.try_to_libc(vm)?; + let ret = unsafe { libc::sched_setparam(args.pid, &libc_sched_param) }; + if ret == -1 { + Err(errno_err(vm)) + } else { + Ok(ret) } }