Merge pull request #3126 from zetwhite/posix_sched

youknowone · web-flow · commit e9c098b54e79 · 2021-09-25T14:45:44.000+09:00
posix : implement sched_set/getscheduler and sched_set/getparam
diff --git a/vm/src/stdlib/posix.rs b/vm/src/stdlib/posix.rs
@@ -472,12 +472,139 @@ pub mod module {
                 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<libc::sched_param> {
+            use crate::TypeProtocol;
+            let priority = self.sched_priority.clone();
+            let priority_type = priority.class().name();
+            let value = priority.downcast::<int::PyInt>().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"
+    ))]
+    #[pyfunction]
+    fn sched_getscheduler(pid: libc::pid_t, vm: &VirtualMachine) -> PyResult<i32> {
+        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"
+    ))]
+    #[derive(FromArgs)]
+    struct SchedSetschedulerArgs {
+        #[pyarg(positional)]
+        pid: i32,
+        #[pyarg(positional)]
+        policy: i32,
+        #[pyarg(positional)]
+        sched_param_obj: crate::PyRef<SchedParam>,
+    }
+
+    #[cfg(any(
+        target_os = "linux",
+        target_os = "netbsd",
+        target_os = "freebsd",
+        target_os = "android"
+    ))]
+    #[pyfunction]
+    fn sched_setscheduler(args: SchedSetschedulerArgs, vm: &VirtualMachine) -> PyResult<i32> {
+        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<SchedParam> {
+        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: crate::PyRef<SchedParam>,
+    }
+
+    #[cfg(any(
+        target_os = "linux",
+        target_os = "netbsd",
+        target_os = "freebsd",
+        target_os = "android"
+    ))]
+    #[pyfunction]
+    fn sched_setparam(args: SchedSetParamArgs, vm: &VirtualMachine) -> PyResult<i32> {
+        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)
         }
     }
 
