From 65cd014431da27a51addac093b8b7f7f3b80a825 Mon Sep 17 00:00:00 2001
From: David Carlier <devnexen@gmail.com>
Date: Fri, 11 Apr 2025 12:46:53 +0100
Subject: [PATCH] sys::prctl: adding calls about core scheduling for isolation
 based on cookie.

---
 changelog/2629.added.md |  1 +
 src/sys/prctl.rs        | 52 +++++++++++++++++++++++++++++++++++++++++
 test/sys/test_prctl.rs  | 29 +++++++++++++++++++++++
 3 files changed, 82 insertions(+)
 create mode 100644 changelog/2629.added.md

diff --git a/changelog/2629.added.md b/changelog/2629.added.md
new file mode 100644
index 0000000000..1e71928002
--- /dev/null
+++ b/changelog/2629.added.md
@@ -0,0 +1 @@
+Add `get_sched_core`, `set_sched_core`, `share_sched_core_from` and `share_sched_core_to` for Linux
diff --git a/src/sys/prctl.rs b/src/sys/prctl.rs
index d183e105ef..f1145aed76 100644
--- a/src/sys/prctl.rs
+++ b/src/sys/prctl.rs
@@ -7,6 +7,7 @@
 
 use crate::errno::Errno;
 use crate::sys::signal::Signal;
+use crate::unistd::Pid;
 use crate::Result;
 
 use libc::{c_int, c_ulong, c_void};
@@ -226,3 +227,54 @@ pub fn set_vma_anon_name(addr: NonNull<c_void>, length: NonZeroUsize, name: Opti
 
     Errno::result(res).map(drop)
 }
+
+/// Get the (scheduling) cookie of the calling thread
+pub fn get_sched_core() -> Result<u32> {
+    let mut cookie: u32 = 0;
+    let res = unsafe {
+        libc::prctl(
+            libc::PR_SCHED_CORE,
+            libc::PR_SCHED_CORE_GET,
+            &mut cookie, 0, 0)
+    };
+
+    match Errno::result(res) {
+        Ok(_) => Ok(cookie),
+        Err(e) => Err(e),
+    }
+}
+
+/// Set a cookie for the calling thread
+pub fn set_sched_core(cookie: u32) -> Result<()> {
+    let res = unsafe {
+        libc::prctl(
+            libc::PR_SCHED_CORE,
+            libc::PR_SCHED_CORE_CREATE,
+            &cookie, 0, 0)
+    };
+
+    Errno::result(res).map(drop)
+}
+
+/// Share the cookie with another thread
+pub fn share_sched_core_from(pid: Pid) -> Result<()> {
+    let res = unsafe {
+        libc::prctl(
+            libc::PR_SCHED_CORE,
+            libc::PR_SCHED_CORE_SHARE_FROM,
+            &pid, 0, 0)
+    };
+
+    Errno::result(res).map(drop)
+}
+/// Calling thread shares the same cookie with another thread
+pub fn share_sched_core_to(pid: Pid) -> Result<()> {
+    let res = unsafe {
+        libc::prctl(
+            libc::PR_SCHED_CORE,
+            libc::PR_SCHED_CORE_SHARE_TO,
+            &pid, 0, 0)
+    };
+
+    Errno::result(res).map(drop)
+}
diff --git a/test/sys/test_prctl.rs b/test/sys/test_prctl.rs
index da072bba46..d2569661a0 100644
--- a/test/sys/test_prctl.rs
+++ b/test/sys/test_prctl.rs
@@ -164,4 +164,33 @@ mod test_prctl {
         .unwrap_or_default();
         prctl::set_vma_anon_name(ptr, sz, None).unwrap_or_default();
     }
+
+    #[test]
+    fn test_sched_core() {
+        use nix::errno::Errno;
+        use nix::unistd::Pid;
+
+        let c: u32 = 1234;
+        let p = Pid::from_raw(0);
+        let mut err = prctl::set_sched_core(c).unwrap_err();
+        match err {
+            Errno::EINVAL | Errno::ESRCH | Errno::ENODEV => (),
+            e => panic!("unexpected error {e}"),
+        }
+        err = prctl::get_sched_core().unwrap_err();
+        match err {
+            Errno::EINVAL | Errno::ESRCH | Errno::ENODEV => (),
+            e => panic!("unexpected error {e}"),
+        }
+        err = prctl::share_sched_core_from(p).unwrap_err();
+        match err {
+            Errno::EINVAL | Errno::ESRCH | Errno::ENODEV => (),
+            e => panic!("unexpected error {e}"),
+        }
+        err = prctl::share_sched_core_to(p).unwrap_err();
+        match err {
+            Errno::EINVAL | Errno::ESRCH | Errno::ENODEV => (),
+            e => panic!("unexpected error {e}"),
+        }
+    }
 }