diff --git a/os/StarryOS/kernel/src/cgroup/mod.rs b/os/StarryOS/kernel/src/cgroup/mod.rs index 219e8b9805..6fc71a890d 100644 --- a/os/StarryOS/kernel/src/cgroup/mod.rs +++ b/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 { - 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 { + 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 { let mut current = id; let mut names = Vec::new(); loop { @@ -178,14 +273,24 @@ pub fn path(id: CgroupId) -> AxResult { Ok(path) } +pub fn path(id: CgroupId) -> AxResult { + let tree = CGROUP_TREE.lock(); + path_locked(&tree, id) +} + pub fn procs_text(id: CgroupId) -> AxResult { - 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 { Ok(text) } +pub fn proc_cgroup_text(proc_data: &crate::task::ProcessData) -> AxResult { + 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 { + 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] +} diff --git a/os/StarryOS/kernel/src/entry.rs b/os/StarryOS/kernel/src/entry.rs index ea53fe4027..2a3a058149 100644 --- a/os/StarryOS/kernel/src/entry.rs +++ b/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(); diff --git a/os/StarryOS/kernel/src/pseudofs/proc.rs b/os/StarryOS/kernel/src/pseudofs/proc.rs index 155da91801..1cf099be92 100644 --- a/os/StarryOS/kernel/src/pseudofs/proc.rs +++ b/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()) }) diff --git a/os/StarryOS/kernel/src/syscall/task/clone.rs b/os/StarryOS/kernel/src/syscall/task/clone.rs index 1be7b49157..f243d3a74c 100644 --- a/os/StarryOS/kernel/src/syscall/task/clone.rs +++ b/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()); diff --git a/os/StarryOS/kernel/src/task/mod.rs b/os/StarryOS/kernel/src/task/mod.rs index 7b354cdea5..bac110a65a 100644 --- a/os/StarryOS/kernel/src/task/mod.rs +++ b/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, + /// 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, /// The command line arguments @@ -663,6 +668,11 @@ pub struct ProcessData { cont_event: Arc, } +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>, exit_signal: Option, vm_aspace_shared: bool, - ) -> Arc { + cgroup_init: ProcessCgroupInit<'_>, + ) -> AxResult> { 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(); } } diff --git a/os/StarryOS/kernel/src/task/ops.rs b/os/StarryOS/kernel/src/task/ops.rs index 9d57885e22..41256ef161 100644 --- a/os/StarryOS/kernel/src/task/ops.rs +++ b/os/StarryOS/kernel/src/task/ops.rs @@ -147,7 +147,10 @@ pub fn get_process_data(pid: Pid) -> AxResult> { /// `NoSuchProcess` immediately, regardless of whether any other strong /// [`Arc`] 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). diff --git a/test-suit/starryos/normal/qemu-smp1/cgroup-basic/c/src/main.c b/test-suit/starryos/normal/qemu-smp1/cgroup-basic/c/src/main.c index f96bc1af75..45fd682eb1 100644 --- a/test-suit/starryos/normal/qemu-smp1/cgroup-basic/c/src/main.c +++ b/test-suit/starryos/normal/qemu-smp1/cgroup-basic/c/src/main.c @@ -10,6 +10,7 @@ #include #include #include +#include #include static int __pass = 0; @@ -151,6 +152,72 @@ 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; + } + + size_t end = (size_t)nread < sizeof(buf) ? (size_t)nread : sizeof(buf) - 1; + buf[end] = '\0'; + int matches = strcmp(buf, expected) == 0; + CHECK(matches, msg); + if (!matches) { + printf(" OBSERVE | expected='%s' got='%s'\n", expected, buf); + } +} + +static ssize_t write_all_fd(int fd, const void *buf, size_t len) +{ + const char *cursor = buf; + size_t written = 0; + + while (written < len) { + ssize_t ret = write(fd, cursor + written, len - written); + if (ret < 0) { + if (errno == EINTR) { + continue; + } + return -1; + } + if (ret == 0) { + errno = EIO; + return -1; + } + written += (size_t)ret; + } + + return (ssize_t)written; +} + +static ssize_t read_exact_fd(int fd, void *buf, size_t len) +{ + char *cursor = buf; + size_t nread = 0; + + while (nread < len) { + ssize_t ret = read(fd, cursor + nread, len - nread); + if (ret < 0) { + if (errno == EINTR) { + continue; + } + return -1; + } + if (ret == 0) { + break; + } + nread += (size_t)ret; + } + + return (ssize_t)nread; +} + static int buffer_contains_pid(const char *buf, pid_t pid) { const char *cursor = buf; @@ -191,6 +258,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_all_fd(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 +426,126 @@ 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"); + int ready_pipe[2]; + int release_pipe[2]; + errno = 0; + int ready_ok = pipe(ready_pipe); + int ready_errno = errno; + errno = ready_errno; + CHECK(ready_ok == 0, "create fork inheritance ready pipe"); + errno = 0; + int release_ok = ready_ok == 0 ? pipe(release_pipe) : -1; + int release_errno = errno; + errno = release_errno; + CHECK(release_ok == 0, "create fork inheritance release pipe"); + if (ready_ok == 0 && release_ok == 0) { + errno = 0; + pid_t child_pid = fork(); + int fork_errno = errno; + errno = fork_errno; + CHECK(child_pid >= 0, "fork after cgroup migration succeeds"); + if (child_pid == 0) { + close(ready_pipe[0]); + close(release_pipe[1]); + + char child_buf[128]; + ssize_t child_nread = + read_text_file("/proc/self/cgroup", child_buf, sizeof(child_buf)); + size_t end = child_nread >= 0 && (size_t)child_nread < sizeof(child_buf) + ? (size_t)child_nread + : sizeof(child_buf) - 1; + child_buf[end] = '\0'; + char inherited = child_nread >= 0 && strcmp(child_buf, "0::/migrate\n") == 0 + ? '1' + : '0'; + (void)write_all_fd(ready_pipe[1], &inherited, 1); + + char release_token; + ssize_t release_nread = read_exact_fd(release_pipe[0], &release_token, 1); + _exit(inherited == '1' && release_nread == 1 ? 0 : 1); + } + + close(ready_pipe[1]); + close(release_pipe[0]); + if (child_pid > 0) { + char inherited = '0'; + errno = 0; + ssize_t inherited_nread = read_exact_fd(ready_pipe[0], &inherited, 1); + int inherited_errno = errno; + errno = inherited_errno; + CHECK(inherited_nread == 1 && inherited == '1', + "fork child inherits migrated cgroup in procfs"); + expect_file_contains_pid(CGROUP2_PATH "/migrate/cgroup.procs", child_pid, + "child cgroup.procs contains inherited child process"); + expect_file_not_contains_pid(CGROUP2_PATH "/cgroup.procs", child_pid, + "root cgroup.procs excludes inherited child process"); + + char release_token = 'x'; + errno = 0; + ssize_t release_written = write_all_fd(release_pipe[1], &release_token, 1); + int release_write_errno = errno; + errno = release_write_errno; + CHECK(release_written == 1, "release fork inheritance child"); + + int child_status = 0; + errno = 0; + pid_t waited = waitpid(child_pid, &child_status, 0); + int wait_errno = errno; + errno = wait_errno; + CHECK(waited == child_pid && WIFEXITED(child_status) && + WEXITSTATUS(child_status) == 0, + "fork inheritance child exits cleanly"); + } + close(ready_pipe[0]); + close(release_pipe[1]); + } else { + if (ready_ok == 0) { + close(ready_pipe[0]); + close(ready_pipe[1]); + } + } + 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;