feat(starry-kernel): support cgroup2 process migration

os/StarryOS/kernel/src/cgroup/mod.rs

@@ -8,6 +8,7 @@ use core::fmt::Write;
 use ax_errno::{AxError, AxResult, LinuxError};
 use ax_kspin::SpinNoIrq;
 use spin::LazyLock;
+use starry_process::Pid;
 
 pub type CgroupId = u64;
 
@@ -150,6 +151,66 @@ pub fn remove_child(parent: CgroupId, name: &str) -> AxResult<()> {
     Ok(())
 }
 
+pub fn register_process(id: CgroupId) {
+    let mut tree = CGROUP_TREE.lock();
+    if let Some(node) = tree.nodes.get_mut(&id) {
+        node.live_processes = node.live_processes.saturating_add(1);
+    }
+}
+
+pub fn unregister_process(id: CgroupId) {
+    let mut tree = CGROUP_TREE.lock();
+    if let Some(node) = tree.nodes.get_mut(&id) {
+        node.live_processes = node.live_processes.saturating_sub(1);
+    }
+}
+
+pub fn attach_process(target: CgroupId, pid: Pid) -> AxResult<()> {
+    if pid == 0 {
+        return Err(AxError::from(LinuxError::EINVAL));
+    }
+
+    let proc_data =
+        crate::task::get_process_data(pid).map_err(|_| AxError::from(LinuxError::ESRCH))?;
+
+    let mut tree = CGROUP_TREE.lock();
+    if !tree.nodes.contains_key(&target) {
+        return Err(AxError::NotFound);
+    }
+
+    let old = proc_data.cgroup_id();
+    if !tree.nodes.contains_key(&old) {
+        return Err(AxError::NotFound);
+    }
+    if old == target {
+        return Ok(());
+    }
+
+    let target_live_processes = tree
+        .nodes
+        .get(&target)
+        .expect("target was checked above")
+        .live_processes
+        .checked_add(1)
+        .ok_or(AxError::NoMemory)?;
+    let old_live_processes = tree
+        .nodes
+        .get(&old)
+        .expect("old cgroup was checked above")
+        .live_processes
+        .saturating_sub(1);
+    tree.nodes
+        .get_mut(&old)
+        .expect("old cgroup was checked above")
+        .live_processes = old_live_processes;
+    tree.nodes
+        .get_mut(&target)
+        .expect("target was checked above")
+        .live_processes = target_live_processes;
+    proc_data.set_cgroup_id(target);
+    Ok(())
+}
+
 pub fn path(id: CgroupId) -> AxResult<String> {
     let tree = CGROUP_TREE.lock();
     let mut current = id;
@@ -180,12 +241,10 @@ pub fn path(id: CgroupId) -> AxResult<String> {
 
 pub fn procs_text(id: CgroupId) -> AxResult<String> {
     ensure_node_exists(id)?;
-    if id != ROOT_ID {
-        return Ok(String::new());
-    }
 
     let mut pids: Vec<_> = crate::task::processes()
         .into_iter()
+        .filter(|proc_data| proc_data.cgroup_id() == id)
         .map(|proc_data| proc_data.proc.pid())
         .collect();
     pids.sort_unstable();
@@ -197,6 +256,13 @@ pub fn procs_text(id: CgroupId) -> AxResult<String> {
     Ok(text)
 }
 
+pub fn proc_cgroup_text(proc_data: &crate::task::ProcessData) -> AxResult<String> {
+    let path = path(proc_data.cgroup_id())?;
+    let mut text = String::new();
+    let _ = writeln!(text, "0::{path}");
+    Ok(text)
+}
+
 pub fn controllers_text(id: CgroupId) -> AxResult<&'static str> {
     ensure_node_exists(id)?;
     Ok("")
@@ -207,9 +273,10 @@ pub fn subtree_control_text(id: CgroupId) -> AxResult<&'static str> {
     Ok("")
 }
 
-pub fn write_procs(id: CgroupId, _data: &[u8]) -> AxResult<()> {
+pub fn write_procs(id: CgroupId, data: &[u8]) -> AxResult<()> {
     ensure_node_exists(id)?;
-    Err(AxError::from(LinuxError::EOPNOTSUPP))
+    let pid = parse_procs_pid(data)?;
+    attach_process(id, pid)
 }
 
 pub fn write_subtree_control(id: CgroupId, _data: &[u8]) -> AxResult<()> {
@@ -223,3 +290,41 @@ pub fn ensure_node_exists(id: CgroupId) -> AxResult<()> {
     debug_assert_eq!(node.id, id);
     Ok(())
 }
+
+pub fn parse_procs_pid(data: &[u8]) -> AxResult<Pid> {
+    let data = trim_ascii_whitespace(data);
+    if data.is_empty() {
+        return Err(AxError::from(LinuxError::EINVAL));
+    }
+
+    let mut value = 0u64;
+    for byte in data {
+        if !byte.is_ascii_digit() {
+            return Err(AxError::from(LinuxError::EINVAL));
+        }
+        value = value
+            .checked_mul(10)
+            .and_then(|value| value.checked_add(u64::from(byte - b'0')))
+            .ok_or_else(|| AxError::from(LinuxError::EINVAL))?;
+        if value > u64::from(Pid::MAX) {
+            return Err(AxError::from(LinuxError::EINVAL));
+        }
+    }
+
+    if value == 0 {
+        return Err(AxError::from(LinuxError::EINVAL));
+    }
+    Ok(value as Pid)
+}
+
+fn trim_ascii_whitespace(data: &[u8]) -> &[u8] {
+    let start = data
+        .iter()
+        .position(|byte| !byte.is_ascii_whitespace())
+        .unwrap_or(data.len());
+    let end = data
+        .iter()
+        .rposition(|byte| !byte.is_ascii_whitespace())
+        .map_or(start, |index| index + 1);
+    &data[start..end]
+}

os/StarryOS/kernel/src/pseudofs/proc.rs

@@ -684,6 +684,7 @@ impl SimpleDirOps for ThreadDir {
                 "mounts",
                 "cmdline",
                 "comm",
+                "cgroup",
                 "exe",
                 "fd",
             ]
@@ -812,6 +813,13 @@ impl SimpleDirOps for ThreadDir {
                 }),
             )
             .into(),
+            "cgroup" => {
+                let proc_data = task.as_thread().proc_data.clone();
+                SimpleFile::new_regular(fs, move || {
+                    Ok(crate::cgroup::proc_cgroup_text(&proc_data)?.into_bytes())
+                })
+                .into()
+            }
             "exe" => SimpleFile::new(fs, NodeType::Symlink, move || {
                 Ok(task.as_thread().proc_data.exe_path.read().clone())
             })

os/StarryOS/kernel/src/task/mod.rs

@@ -14,7 +14,7 @@ use alloc::{boxed::Box, string::String, sync::Arc, vec::Vec};
 use core::{
     cell::RefCell,
     ops::Deref,
-    sync::atomic::{AtomicBool, AtomicI32, AtomicU8, AtomicU32, AtomicUsize, Ordering},
+    sync::atomic::{AtomicBool, AtomicI32, AtomicU8, AtomicU32, AtomicU64, AtomicUsize, Ordering},
 };
 
 use ax_runtime::hal::{cpu::uspace::UserContext, time::TimeValue};
@@ -502,6 +502,8 @@ impl ProcessImage {
 pub struct ProcessData {
     /// The process.
     pub proc: Arc<Process>,
+    /// Current cgroup v2 membership in the global hierarchy.
+    cgroup_id: AtomicU64,
     /// The executable path
     pub exe_path: RwLock<String>,
     /// The command line arguments
@@ -673,6 +675,7 @@ impl ProcessData {
     ) -> Arc<Self> {
         let this = Arc::new(Self {
             proc,
+            cgroup_id: AtomicU64::new(crate::cgroup::root_id()),
             exe_path: RwLock::new(image.exe_path),
             cmdline: RwLock::new(image.cmdline),
             auxv: RwLock::new(image.auxv),
@@ -741,9 +744,20 @@ impl ProcessData {
         // expression would nest a sleepable lock inside atomic context.
         let aspace_arc = this.aspace.lock().clone();
         crate::mm::attach_process_slot(&aspace_arc);
+        crate::cgroup::register_process(crate::cgroup::root_id());
         this
     }
 
+    /// Return the current cgroup v2 membership id.
+    pub fn cgroup_id(&self) -> crate::cgroup::CgroupId {
+        self.cgroup_id.load(Ordering::Acquire)
+    }
+
+    /// Update cgroup membership. This is serialized by the cgroup core lock.
+    pub(crate) fn set_cgroup_id(&self, id: crate::cgroup::CgroupId) {
+        self.cgroup_id.store(id, Ordering::Release);
+    }
+
     /// Whether this process shares its VM address space (`CLONE_VM`).
     #[inline]
     pub fn vm_aspace_shared(&self) -> bool {
@@ -1369,6 +1383,7 @@ impl ProcessData {
 
 impl Drop for ProcessData {
     fn drop(&mut self) {
+        crate::cgroup::unregister_process(self.cgroup_id());
         self.release_aspace_slot_if_needed();
     }
 }

test-suit/starryos/normal/qemu-smp1/cgroup-basic/c/src/main.c

@@ -151,6 +151,24 @@ static void expect_empty_file(const char *path, const char *msg)
     CHECK(nread == 0, msg);
 }
 
+static void expect_file_equals(const char *path, const char *expected, const char *msg)
+{
+    char buf[4096];
+    errno = 0;
+    ssize_t nread = read_text_file(path, buf, sizeof(buf));
+    int saved_errno = errno;
+    errno = saved_errno;
+    if (nread < 0) {
+        CHECK(0, msg);
+        return;
+    }
+
+    CHECK(strcmp(buf, expected) == 0, msg);
+    if (strcmp(buf, expected) != 0) {
+        printf("  OBSERVE | expected='%s' got='%s'\n", expected, buf);
+    }
+}
+
 static int buffer_contains_pid(const char *buf, pid_t pid)
 {
     const char *cursor = buf;
@@ -191,6 +209,50 @@ static void expect_write_errno(const char *path, const char *data,
     CHECK(written == -1 && saved_errno == expected_errno, msg);
 }
 
+static void expect_write_ok(const char *path, const char *data, const char *msg)
+{
+    int fd = open(path, O_WRONLY);
+    if (fd < 0) {
+        CHECK(0, msg);
+        return;
+    }
+
+    size_t len = strlen(data);
+    errno = 0;
+    ssize_t written = write(fd, data, len);
+    int saved_errno = errno;
+    close(fd);
+    errno = saved_errno;
+    CHECK(written == (ssize_t)len, msg);
+}
+
+static void expect_write_pid_ok(const char *path, pid_t pid, const char *msg)
+{
+    char data[32];
+    snprintf(data, sizeof(data), "%ld", (long)pid);
+    expect_write_ok(path, data, msg);
+}
+
+static void expect_file_contains_pid(const char *path, pid_t pid, const char *msg)
+{
+    char buf[4096];
+    errno = 0;
+    ssize_t nread = read_text_file(path, buf, sizeof(buf));
+    int saved_errno = errno;
+    errno = saved_errno;
+    CHECK(nread >= 0 && buffer_contains_pid(buf, pid), msg);
+}
+
+static void expect_file_not_contains_pid(const char *path, pid_t pid, const char *msg)
+{
+    char buf[4096];
+    errno = 0;
+    ssize_t nread = read_text_file(path, buf, sizeof(buf));
+    int saved_errno = errno;
+    errno = saved_errno;
+    CHECK(nread >= 0 && !buffer_contains_pid(buf, pid), msg);
+}
+
 static void expect_link_errno(const char *old_path, const char *new_path,
                               int expected_errno, const char *msg)
 {
@@ -315,6 +377,48 @@ int main(void)
     expect_path_missing(CGROUP2_PATH "/renamed", "rename failure leaves destination missing");
     expect_rmdir_ok(CGROUP2_PATH "/ren", "cleanup rename negative test cgroup");
 
+    pid_t self_pid = getpid();
+    expect_file_equals("/proc/self/cgroup", "0::/\n",
+                       "proc self cgroup initially points to root");
+
+    expect_mkdir_ok(CGROUP2_PATH "/migrate", "mkdir migrate cgroup succeeds");
+    expect_write_errno(CGROUP2_PATH "/migrate/cgroup.procs", " \n", EINVAL,
+                       "writing whitespace-only cgroup.procs fails with EINVAL");
+    expect_write_errno(CGROUP2_PATH "/migrate/cgroup.procs", "not-a-pid", EINVAL,
+                       "writing non-number cgroup.procs fails with EINVAL");
+    expect_write_errno(CGROUP2_PATH "/migrate/cgroup.procs", "0", EINVAL,
+                       "writing pid 0 cgroup.procs fails with EINVAL");
+    expect_write_errno(CGROUP2_PATH "/migrate/cgroup.procs", "99999999", ESRCH,
+                       "writing missing pid cgroup.procs fails with ESRCH");
+    expect_file_equals("/proc/self/cgroup", "0::/\n",
+                       "failed cgroup.procs writes keep process in root");
+
+    expect_write_pid_ok(CGROUP2_PATH "/migrate/cgroup.procs", self_pid,
+                        "writing current pid to child cgroup.procs succeeds");
+    expect_file_equals("/proc/self/cgroup", "0::/migrate\n",
+                       "proc self cgroup points to migrated cgroup");
+    expect_file_contains_pid(CGROUP2_PATH "/migrate/cgroup.procs", self_pid,
+                             "child cgroup.procs contains migrated process");
+    expect_file_not_contains_pid(CGROUP2_PATH "/cgroup.procs", self_pid,
+                                 "root cgroup.procs no longer contains migrated process");
+    expect_rmdir_errno(CGROUP2_PATH "/migrate", EBUSY,
+                       "rmdir populated cgroup fails with EBUSY");
+    expect_path_exists(CGROUP2_PATH "/migrate",
+                       "populated cgroup remains after failed rmdir");
+
+    expect_write_pid_ok(CGROUP2_PATH "/cgroup.procs", self_pid,
+                        "writing current pid to root cgroup.procs succeeds");
+    expect_file_equals("/proc/self/cgroup", "0::/\n",
+                       "proc self cgroup points back to root");
+    expect_file_contains_pid(CGROUP2_PATH "/cgroup.procs", self_pid,
+                             "root cgroup.procs contains migrated-back process");
+    expect_empty_file(CGROUP2_PATH "/migrate/cgroup.procs",
+                      "child cgroup.procs is empty after migrating back");
+    expect_rmdir_ok(CGROUP2_PATH "/migrate", "rmdir empty migrate cgroup succeeds");
+
+    expect_write_errno(CGROUP2_PATH "/cgroup.controllers", "x", EACCES,
+                       "writing cgroup.controllers fails with EACCES");
+
     check_mkdir(CGROUP_V1_PATH, "mkdir cgroup v1 mountpoint");
 
     errno = 0;