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,8 +151,102 @@ pub fn remove_child(parent: CgroupId, name: &str) -> AxResult<()> {
     Ok(())
 }
 
-pub fn path(id: CgroupId) -> AxResult<String> {
-    let tree = CGROUP_TREE.lock();
+pub fn register_process(id: CgroupId) -> AxResult<()> {
+    let mut tree = CGROUP_TREE.lock();
+    let node = tree.nodes.get_mut(&id).ok_or(AxError::NotFound)?;
+    node.live_processes = node
+        .live_processes
+        .checked_add(1)
+        .ok_or_else(|| AxError::from(LinuxError::EINVAL))?;
+    Ok(())
+}
+
+pub fn register_fork_child(parent: &crate::task::ProcessData) -> AxResult<CgroupId> {
+    let mut tree = CGROUP_TREE.lock();
+    if !parent.is_cgroup_membership_active() {
+        return Err(AxError::from(LinuxError::ESRCH));
+    }
+
+    let id = parent.cgroup_id();
+    let node = tree
+        .nodes
+        .get_mut(&id)
+        .ok_or_else(|| AxError::from(LinuxError::ESRCH))?;
+    node.live_processes = node
+        .live_processes
+        .checked_add(1)
+        .ok_or_else(|| AxError::from(LinuxError::EINVAL))?;
+    Ok(id)
+}
+
+fn unregister_process_locked(tree: &mut CgroupTree, id: CgroupId) {
+    if let Some(node) = tree.nodes.get_mut(&id) {
+        if let Some(live_processes) = node.live_processes.checked_sub(1) {
+            node.live_processes = live_processes;
+        } else {
+            debug_assert!(false, "cgroup live_processes underflow during unregister");
+        }
+    }
+}
+
+pub fn release_process_membership(proc_data: &crate::task::ProcessData) {
+    let mut tree = CGROUP_TREE.lock();
+    if proc_data.deactivate_cgroup_membership() {
+        unregister_process_locked(&mut tree, proc_data.cgroup_id());
+    }
+}
+
+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);
+    }
+    if !proc_data.is_cgroup_membership_active() {
+        return Err(AxError::from(LinuxError::ESRCH));
+    }
+
+    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_else(|| AxError::from(LinuxError::EINVAL))?;
+    let old_live_processes = tree
+        .nodes
+        .get(&old)
+        .expect("old cgroup was checked above")
+        .live_processes
+        .checked_sub(1)
+        .ok_or_else(|| AxError::from(LinuxError::EINVAL))?;
+    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(())
+}
+
+fn path_locked(tree: &CgroupTree, id: CgroupId) -> AxResult<String> {
     let mut current = id;
     let mut names = Vec::new();
     loop {
@@ -178,14 +273,24 @@ pub fn path(id: CgroupId) -> AxResult<String> {
     Ok(path)
 }
 
+pub fn path(id: CgroupId) -> AxResult<String> {
+    let tree = CGROUP_TREE.lock();
+    path_locked(&tree, id)
+}
+
 pub fn procs_text(id: CgroupId) -> AxResult<String> {
-    ensure_node_exists(id)?;
-    if id != ROOT_ID {
-        return Ok(String::new());
-    }
+    // Lock order: callers that need both process state and the cgroup tree must
+    // first take a process-table snapshot and then lock CGROUP_TREE. Do not
+    // hold CGROUP_TREE while entering task lookup/table code: remove_process()
+    // removes from PROCESS_TABLE first and then releases cgroup membership.
+    let processes = crate::task::processes();
+    let tree = CGROUP_TREE.lock();
+    let node = tree.nodes.get(&id).ok_or(AxError::NotFound)?;
+    debug_assert_eq!(node.id, id);
 
-    let mut pids: Vec<_> = crate::task::processes()
+    let mut pids: Vec<_> = processes
         .into_iter()
+        .filter(|proc_data| proc_data.is_cgroup_membership_active() && proc_data.cgroup_id() == id)
         .map(|proc_data| proc_data.proc.pid())
         .collect();
     pids.sort_unstable();
@@ -197,6 +302,17 @@ pub fn procs_text(id: CgroupId) -> AxResult<String> {
     Ok(text)
 }
 
+pub fn proc_cgroup_text(proc_data: &crate::task::ProcessData) -> AxResult<String> {
+    let tree = CGROUP_TREE.lock();
+    if !proc_data.is_cgroup_membership_active() {
+        return Err(AxError::from(LinuxError::ESRCH));
+    }
+    let path = path_locked(&tree, 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 +323,9 @@ pub fn subtree_control_text(id: CgroupId) -> AxResult<&'static str> {
     Ok("")
 }
 
-pub fn write_procs(id: CgroupId, _data: &[u8]) -> AxResult<()> {
-    ensure_node_exists(id)?;
-    Err(AxError::from(LinuxError::EOPNOTSUPP))
+pub fn write_procs(id: CgroupId, data: &[u8]) -> AxResult<()> {
+    let pid = parse_procs_pid(data)?;
+    attach_process(id, pid)
 }
 
 pub fn write_subtree_control(id: CgroupId, _data: &[u8]) -> AxResult<()> {
@@ -223,3 +339,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/entry.rs

@@ -77,7 +77,9 @@ pub fn init(args: &[String], envs: &[String]) {
         Arc::default(),
         None,
         false,
-    );
+        crate::task::ProcessCgroupInit::Root,
+    )
+    .expect("Failed to create init process data");
 
     {
         let mut scope = proc.scope.write();

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

@@ -760,6 +760,7 @@ impl SimpleDirOps for ThreadDir {
                 "mounts",
                 "cmdline",
                 "comm",
+                "cgroup",
                 "exe",
                 "fd",
                 "ns",
@@ -889,6 +890,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/syscall/task/clone.rs

@@ -14,7 +14,10 @@ use starry_vm::VmMutPtr;
 use crate::{
     file::{FD_TABLE, FileLike, PidFd, close_file_like},
     mm::copy_from_kernel,
-    task::{AsThread, ProcessData, ProcessImage, Thread, add_task_to_table, new_user_task},
+    task::{
+        AsThread, ProcessCgroupInit, ProcessData, ProcessImage, Thread, add_task_to_table,
+        new_user_task,
+    },
 };
 
 bitflags! {
@@ -227,7 +230,8 @@ impl CloneArgs {
                 signal_actions,
                 exit_signal,
                 flags.contains(CloneFlags::VM),
-            );
+                ProcessCgroupInit::Inherit(old_proc_data),
+            )?;
             proc_data.set_umask(old_proc_data.umask());
             proc_data.set_nice(old_proc_data.nice());
             proc_data.set_heap_top(old_proc_data.get_heap_top());

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

@@ -14,9 +14,10 @@ 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_errno::AxResult;
 use ax_runtime::hal::{cpu::uspace::UserContext, time::TimeValue};
 use ax_sync::{Mutex, spin::SpinNoIrq};
 use ax_task::{TaskExt, TaskInner};
@@ -502,6 +503,10 @@ impl ProcessImage {
 pub struct ProcessData {
     /// The process.
     pub proc: Arc<Process>,
+    /// Current cgroup v2 membership in the global hierarchy.
+    cgroup_id: AtomicU64,
+    /// Whether this process still contributes to its cgroup's live count.
+    cgroup_membership_active: AtomicBool,
     /// The executable path
     pub exe_path: RwLock<String>,
     /// The command line arguments
@@ -663,6 +668,11 @@ pub struct ProcessData {
     cont_event: Arc<PollSet>,
 }
 
+pub enum ProcessCgroupInit<'a> {
+    Root,
+    Inherit(&'a ProcessData),
+}
+
 impl ProcessData {
     /// Create a new [`ProcessData`].
     pub fn new(
@@ -672,9 +682,12 @@ impl ProcessData {
         signal_actions: Arc<SpinNoIrq<SignalActions>>,
         exit_signal: Option<Signo>,
         vm_aspace_shared: bool,
-    ) -> Arc<Self> {
+        cgroup_init: ProcessCgroupInit<'_>,
+    ) -> AxResult<Arc<Self>> {
         let this = Arc::new(Self {
             proc,
+            cgroup_id: AtomicU64::new(crate::cgroup::root_id()),
+            cgroup_membership_active: AtomicBool::new(false),
             exe_path: RwLock::new(image.exe_path),
             cmdline: RwLock::new(image.cmdline),
             auxv: RwLock::new(image.auxv),
@@ -745,7 +758,42 @@ 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);
-        this
+        let cgroup_id = match cgroup_init {
+            ProcessCgroupInit::Root => {
+                let id = crate::cgroup::root_id();
+                crate::cgroup::register_process(id)?;
+                id
+            }
+            ProcessCgroupInit::Inherit(parent) => crate::cgroup::register_fork_child(parent)?,
+        };
+        this.cgroup_id.store(cgroup_id, Ordering::Release);
+        this.cgroup_membership_active.store(true, Ordering::Release);
+        Ok(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 still has an active cgroup membership.
+    pub(crate) fn is_cgroup_membership_active(&self) -> bool {
+        self.cgroup_membership_active.load(Ordering::Acquire)
+    }
+
+    /// Mark this process's cgroup membership inactive.
+    pub(crate) fn deactivate_cgroup_membership(&self) -> bool {
+        self.cgroup_membership_active.swap(false, Ordering::AcqRel)
+    }
+
+    /// Release this process's cgroup live-count membership once.
+    pub fn release_cgroup_membership_if_needed(&self) {
+        crate::cgroup::release_process_membership(self);
     }
 
     /// Whether this process shares its VM address space (`CLONE_VM`).
@@ -1373,6 +1421,7 @@ impl ProcessData {
 
 impl Drop for ProcessData {
     fn drop(&mut self) {
+        self.release_cgroup_membership_if_needed();
         self.release_aspace_slot_if_needed();
     }
 }

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

@@ -147,7 +147,10 @@ pub fn get_process_data(pid: Pid) -> AxResult<Arc<ProcessData>> {
 /// `NoSuchProcess` immediately, regardless of whether any other strong
 /// [`Arc<ProcessData>`] references (e.g. task objects) are still alive.
 pub fn remove_process(pid: Pid) {
-    PROCESS_TABLE.write().remove(&pid);
+    let proc_data = PROCESS_TABLE.write().remove(&pid);
+    if let Some(proc_data) = proc_data {
+        proc_data.release_cgroup_membership_if_needed();
+    }
 }
 
 /// Records a PID as zombie (exited but not yet reaped).