feat(axtask): add work-stealing load balancing for SMP

components/axsched/src/cfs.rs

@@ -168,13 +168,34 @@ impl<T> BaseScheduler for CFScheduler<T> {
         }
     }
 
+    fn pick_next_task_matching(
+        &mut self,
+        predicate: impl Fn(&Self::SchedItem) -> bool,
+    ) -> Option<Self::SchedItem> {
+        let key = {
+            let mut found = None;
+            for ((v, id), task) in &self.ready_queue {
+                if predicate(task) {
+                    found = Some((*v, *id));
+                    break;
+                }
+            }
+            found
+        };
+        key.and_then(|k| self.ready_queue.remove(&k))
+    }
+
     fn put_prev_task(&mut self, prev: Self::SchedItem, _preempt: bool) {
         let taskid = self.id_pool.fetch_add(1, Ordering::Release);
         prev.set_id(taskid);
         self.ready_queue
             .insert((prev.clone().get_vruntime(), taskid), prev);
     }
 
+    fn is_empty(&self) -> bool {
+        self.ready_queue.is_empty()
+    }
+
     fn task_tick(&mut self, current: &Self::SchedItem) -> bool {
         current.task_tick();
         if self.ready_queue.is_empty() {

components/axsched/src/fifo.rs

@@ -54,10 +54,38 @@ impl<T> BaseScheduler for FifoScheduler<T> {
         self.ready_queue.pop_front()
     }
 
+    fn pick_next_task_matching(
+        &mut self,
+        predicate: impl Fn(&Self::SchedItem) -> bool,
+    ) -> Option<Self::SchedItem> {
+        let mut skipped = alloc::vec::Vec::new();
+        loop {
+            match self.ready_queue.pop_front() {
+                Some(task) if predicate(&task) => {
+                    for t in skipped.into_iter().rev() {
+                        self.ready_queue.push_front(t);
+                    }
+                    return Some(task);
+                }
+                Some(task) => skipped.push(task),
+                None => {
+                    for t in skipped.into_iter().rev() {
+                        self.ready_queue.push_front(t);
+                    }
+                    return None;
+                }
+            }
+        }
+    }
+
     fn put_prev_task(&mut self, prev: Self::SchedItem, _preempt: bool) {
         self.ready_queue.push_back(prev);
     }
 
+    fn is_empty(&self) -> bool {
+        self.ready_queue.is_empty()
+    }
+
     fn task_tick(&mut self, _current: &Self::SchedItem) -> bool {
         false // no reschedule
     }

components/axsched/src/lib.rs

@@ -41,6 +41,17 @@ pub trait BaseScheduler {
     /// Returns [`None`] if there is not runnable task.
     fn pick_next_task(&mut self) -> Option<Self::SchedItem>;
 
+    /// Picks the next task matching a predicate, removing it from the
+    /// scheduler. Returns [`None`] if no matching task exists.
+    ///
+    /// The predicate check and removal happen atomically within a single
+    /// lock acquisition — the task is never left in an "out of queue"
+    /// window observable by other CPUs.
+    fn pick_next_task_matching(
+        &mut self,
+        predicate: impl Fn(&Self::SchedItem) -> bool,
+    ) -> Option<Self::SchedItem>;
+
     /// Puts the previous task back to the scheduler. The previous task is
     /// usually placed at the end of the ready queue, making it less likely
     /// to be re-scheduled.
@@ -58,4 +69,11 @@ pub trait BaseScheduler {
 
     /// set priority for a task
     fn set_priority(&mut self, task: &Self::SchedItem, prio: isize) -> bool;
+
+    /// Returns `true` if the ready queue has no runnable tasks.
+    ///
+    /// Used as a fast, racy heuristic (without holding the lock) to skip
+    /// empty queues before attempting work-stealing — mirrors Linux's
+    /// `src_rq->nr_running > 1` check in `idle_balance()`.
+    fn is_empty(&self) -> bool;
 }

components/axsched/src/round_robin.rs

@@ -102,6 +102,30 @@ impl<T, const S: usize> BaseScheduler for RRScheduler<T, S> {
         self.ready_queue.pop_front()
     }
 
+    fn pick_next_task_matching(
+        &mut self,
+        predicate: impl Fn(&Self::SchedItem) -> bool,
+    ) -> Option<Self::SchedItem> {
+        let mut skipped = alloc::vec::Vec::new();
+        loop {
+            match self.ready_queue.pop_front() {
+                Some(task) if predicate(&task) => {
+                    for t in skipped.into_iter().rev() {
+                        self.ready_queue.push_front(t);
+                    }
+                    return Some(task);
+                }
+                Some(task) => skipped.push(task),
+                None => {
+                    for t in skipped.into_iter().rev() {
+                        self.ready_queue.push_front(t);
+                    }
+                    return None;
+                }
+            }
+        }
+    }
+
     fn put_prev_task(&mut self, prev: Self::SchedItem, preempt: bool) {
         if prev.time_slice() > 0 && preempt {
             self.ready_queue.push_front(prev)
@@ -111,6 +135,10 @@ impl<T, const S: usize> BaseScheduler for RRScheduler<T, S> {
         }
     }
 
+    fn is_empty(&self) -> bool {
+        self.ready_queue.is_empty()
+    }
+
     fn task_tick(&mut self, current: &Self::SchedItem) -> bool {
         let old_slice = current.time_slice.fetch_sub(1, Ordering::Release);
         old_slice <= 1

components/kspin/src/base.rs

@@ -300,6 +300,25 @@ impl<G: BaseGuard, T: ?Sized> BaseSpinLock<G, T> {
         // there's no need to lock the inner mutex.
         unsafe { &mut *self.data.get() }
     }
+
+    /// Returns a shared reference to the inner data without acquiring the
+    /// lock.  The data may be concurrently modified — the result is stale
+    /// the instant this returns.  Only use for racy heuristics analoguous
+    /// to Linux's `src_rq->nr_running > 1` check.
+    ///
+    /// This is a shared (`&`) reference rather than `&mut`, so it does not
+    /// carry `noalias` — the compiler will not assume exclusivity and will
+    /// not miscompile concurrent accesses to the same data from other CPUs.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that no other thread is concurrently writing
+    /// to the inner data in a way that would cause a data race for the
+    /// specific fields being read.
+    #[inline(always)]
+    pub unsafe fn data_unchecked(&self) -> &T {
+        unsafe { &*self.data.get() }
+    }
 }
 
 impl<G: BaseGuard, T: Default> Default for BaseSpinLock<G, T> {

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

@@ -17,7 +17,7 @@ pub mod timerfd;
 mod wext;
 
 use alloc::{borrow::Cow, sync::Arc};
-use core::{ffi::c_int, time::Duration};
+use core::{ffi::c_int, sync::atomic::AtomicI32, time::Duration};
 
 use ax_errno::{AxError, AxResult};
 use ax_fs_ng::vfs::{FS_CONTEXT, FileBackend, FileFlags, OpenOptions};
@@ -262,9 +262,65 @@ pub struct FileDescriptor {
     pub cloexec: bool,
 }
 
+/// Shared file descriptor table with task-count tracking.
+///
+/// Mirrors Linux's `files_struct.count`: `task_count` tracks how many live
+/// tasks share this table via CLONE_FILES. Unlike `Arc::strong_count`, it is
+/// decremented synchronously in `close_all_fds()` (the exiting task's context),
+/// so the last exiting task always closes FDs in user-task context — never in
+/// the GC task's context.
+pub struct FdTable {
+    inner: RwLock<FlattenObjects<FileDescriptor, AX_FILE_LIMIT>>,
+    task_count: AtomicI32,
+}
+
+impl FdTable {
+    pub fn read(
+        &self,
+    ) -> impl core::ops::Deref<Target = FlattenObjects<FileDescriptor, AX_FILE_LIMIT>> + '_ {
+        self.inner.read()
+    }
+
+    pub fn write(
+        &self,
+    ) -> impl core::ops::DerefMut<Target = FlattenObjects<FileDescriptor, AX_FILE_LIMIT>> + '_ {
+        self.inner.write()
+    }
+
+    /// Record a new task sharing this table (clone with CLONE_FILES).
+    pub fn inc_task_count(&self) {
+        self.task_count
+            .fetch_add(1, core::sync::atomic::Ordering::Relaxed);
+    }
+
+    /// Record a task exiting. Returns `true` if this was the last live task.
+    pub fn dec_task_count(&self) -> bool {
+        self.task_count
+            .fetch_sub(1, core::sync::atomic::Ordering::Relaxed)
+            == 1
+    }
+
+    /// Wrap an existing inner table (used in unshare)
+    pub fn from_inner(inner: RwLock<FlattenObjects<FileDescriptor, AX_FILE_LIMIT>>) -> Self {
+        Self {
+            inner,
+            task_count: AtomicI32::new(1),
+        }
+    }
+}
+
+impl Default for FdTable {
+    fn default() -> Self {
+        Self {
+            inner: RwLock::new(FlattenObjects::default()),
+            task_count: AtomicI32::new(1),
+        }
+    }
+}
+
 scope_local::scope_local! {
     /// The current file descriptor table.
-    pub static FD_TABLE: Arc<RwLock<FlattenObjects<FileDescriptor, AX_FILE_LIMIT>>> = Arc::default();
+    pub static FD_TABLE: Arc<FdTable> = Arc::default();
 }
 
 /// Get a file-like object by `fd`.
@@ -336,8 +392,10 @@ fn notify_close_write(fd: &FileDescriptor) {
 pub fn release_locks_on_close(fd: FileDescriptor) {
     let key = fd.inner.inode_key();
     notify_close_write(&fd);
-    if let Some(k) = key {
-        let pid = current().as_thread().proc_data.proc.pid();
+    if let Some(k) = key
+        && let Some(thr) = current().try_as_thread()
+    {
+        let pid = thr.proc_data.proc.pid();
         crate::syscall::release_inode_posix_locks(pid, k);
     }
     drop(fd);
@@ -352,20 +410,23 @@ pub fn release_locks_on_close(fd: FileDescriptor) {
 /// This must be called when a process exits, so that pipe write ends and other
 /// resources are properly released. Without this, parent processes blocking on
 /// pipe reads will never receive EOF.
+///
+/// Uses per-FdTable `task_count` (mirroring Linux's `files_struct.count`) to
+/// determine the last live task sharing the table.  The count is decremented
+/// inside the FD_TABLE write lock, synchronized with the CLONE_FILES path
+/// that holds the read lock before incrementing it.
 pub fn close_all_fds() {
-    // Acquire the write lock before checking strong_count. The clone(CLONE_FILES)
-    // path in syscall/task/clone.rs also acquires FD_TABLE.read() before cloning
-    // the Arc, creating a shared synchronization boundary. This ensures:
-    // - If close_all_fds acquires the write lock first, clone blocks on read lock
-    //   until we release, so strong_count cannot change during our check.
-    // - If clone holds the read lock first, we block on write lock, and by the
-    //   time we proceed strong_count already reflects the clone.
     let mut table = FD_TABLE.write();
 
     // CLONE_FILES may share the same fd table across multiple tasks/processes.
     // In that case, an exiting sharer must not clear the whole table, or other
     // live sharers (including the parent) will lose stdout/stderr unexpectedly.
-    if Arc::strong_count(&FD_TABLE) > 1 {
+    //
+    // dec_task_count() returns true when this is the last live task — the
+    // check is done under FD_TABLE.write() so it serialises with the
+    // inc_task_count() in the clone(CLONE_FILES) path (which holds
+    // FD_TABLE.read()).
+    if !FD_TABLE.dec_task_count() {
         return;
     }
 

os/StarryOS/kernel/src/syscall/fs/fd_ops.rs

@@ -14,7 +14,7 @@ use starry_vm::{VmMutPtr, VmPtr};
 
 use crate::{
     file::{
-        Directory, FD_TABLE, File, FileDescriptor, FileLike, NsFd, Pipe, add_file_like,
+        Directory, FD_TABLE, FdTable, File, FileDescriptor, FileLike, NsFd, Pipe, add_file_like,
         close_file_like, get_file_like, memfd::Memfd, with_fs,
     },
     mm::vm_load_string,
@@ -377,7 +377,11 @@ pub fn sys_close_range(first: i32, last: i32, flags: u32) -> AxResult<isize> {
     if flags.contains(CloseRangeFlags::UNSHARE) {
         let curr = current();
         let proc_data = &curr.as_thread().proc_data;
-        let new_files = Arc::new(spin::RwLock::new(FD_TABLE.read().clone()));
+        let new_inner = spin::RwLock::new(FD_TABLE.read().clone());
+        let new_files = Arc::new(FdTable::from_inner(new_inner));
+        // Unshare stops sharing the old fd table; mirror the dec_task_count
+        // that the CLONE_FILES path pairs with inc_task_count.
+        FD_TABLE.dec_task_count();
         proc_data.with_current_scope_mut(|scope| {
             *FD_TABLE.scope_mut(scope).deref_mut() = new_files;
         });

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

@@ -49,27 +49,25 @@ pub fn handle_syscall(uctx: &mut UserContext) {
     };
 
     trace!("Syscall {sysno:?}");
-    match ax_task::current()
-        .as_thread()
-        .seccomp_state()
-        .evaluate(uctx)
-    {
-        SeccompDecision::Allow => {}
-        SeccompDecision::Errno(errno) => {
-            uctx.set_retval(seccomp_errno(errno));
-            return;
-        }
-        SeccompDecision::KillProcess => {
-            do_exit(Signo::SIGSYS as i32, true);
-            return;
-        }
-        SeccompDecision::KillThread => {
-            do_exit(Signo::SIGSYS as i32, false);
-            return;
-        }
-        SeccompDecision::UnsupportedAction => {
-            uctx.set_retval(-LinuxError::ENOSYS.code() as usize);
-            return;
+    if let Some(thr) = ax_task::current().try_as_thread() {
+        match thr.seccomp_state().evaluate(uctx) {
+            SeccompDecision::Allow => {}
+            SeccompDecision::Errno(errno) => {
+                uctx.set_retval(seccomp_errno(errno));
+                return;
+            }
+            SeccompDecision::KillProcess => {
+                do_exit(Signo::SIGSYS as i32, true);
+                return;
+            }
+            SeccompDecision::KillThread => {
+                do_exit(Signo::SIGSYS as i32, false);
+                return;
+            }
+            SeccompDecision::UnsupportedAction => {
+                uctx.set_retval(-LinuxError::ENOSYS.code() as usize);
+                return;
+            }
         }
     }
 
@@ -953,8 +951,8 @@ pub fn handle_syscall(uctx: &mut UserContext) {
         Sysno::timer_delete => sys_timer_delete(uctx.arg0() as _),
 
         _ => {
-            let tid = ax_task::current().as_thread().tid();
-            warn!("Unimplemented syscall: {sysno} (tid={tid})");
+            let tid = ax_task::current().try_as_thread().map(|t| t.tid());
+            warn!("Unimplemented syscall: {sysno} (tid={tid:?})");
             Err(AxError::Unsupported)
         }
     };

os/StarryOS/kernel/src/syscall/task/clone.rs

@@ -325,10 +325,11 @@ impl CloneArgs {
                 let mut scope = proc_data.scope.write();
                 if flags.contains(CloneFlags::FILES) {
                     // Synchronize with close_all_fds: holding a read lock
-                    // ensures close_all_fds either observes our strong_count
+                    // ensures close_all_fds either observes our task_count
                     // increment or blocks on write lock until we release.
                     let _guard = FD_TABLE.read();
                     FD_TABLE.scope_mut(&mut scope).clone_from(&FD_TABLE);
+                    FD_TABLE.inc_task_count();
                 } else {
                     FD_TABLE
                         .scope_mut(&mut scope)