fix(axtask): improve might_sleep diagnostics and coverage

os/StarryOS/kernel/src/mm/access.rs

@@ -32,6 +32,7 @@ pub fn access_user_memory<R>(f: impl FnOnce() -> R) -> R {
         ax_runtime::hal::cpu::asm::irqs_enabled(),
         "faultable user memory access requires IRQs enabled"
     );
+    might_sleep();
 
     let curr = current();
     let Some(thr) = curr.try_as_thread() else {

os/StarryOS/kernel/src/syscall/io_mpx/select.rs

@@ -36,6 +36,15 @@ impl FdSet {
     }
 }
 
+fn write_fd_set(user: Option<&mut __kernel_fd_set>, selected: &FdSet, nfds: usize) {
+    if let Some(user) = user {
+        unsafe { FD_ZERO(user) };
+        for index in selected.0.into_iter().take(nfds) {
+            unsafe { FD_SET(index as _, user) };
+        }
+    }
+}
+
 impl fmt::Debug for FdSet {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_list().entries(&self.0).finish()
@@ -98,20 +107,14 @@ fn do_select(
     drop(fd_table);
     let fds = FdPollSet(fds);
 
-    if let Some(readfds) = readfds.as_deref_mut() {
-        unsafe { FD_ZERO(readfds) };
-    }
-    if let Some(writefds) = writefds.as_deref_mut() {
-        unsafe { FD_ZERO(writefds) };
-    }
-    if let Some(exceptfds) = exceptfds.as_deref_mut() {
-        unsafe { FD_ZERO(exceptfds) };
-    }
     with_blocked_signals(sigmask.copied(), || {
-        match block_on(future::timeout(
+        let result = block_on(future::timeout(
             timeout,
             poll_io(&fds, IoEvents::empty(), false, || {
                 let mut res = 0usize;
+                let mut selected_readfds = FdSet(Bitmap::new());
+                let mut selected_writefds = FdSet(Bitmap::new());
+                let mut selected_exceptfds = FdSet(Bitmap::new());
                 for ((fd, interested), index) in fds.0.iter().zip(fd_indices.iter().copied()) {
                     let events = fd.poll();
                     let always_report = events & IoEvents::ALWAYS_POLL;
@@ -123,28 +126,38 @@ fn do_select(
                     let selected_except =
                         selected.contains(IoEvents::ERR) && except_set.0.get(index);
 
-                    if selected_read && let Some(set) = readfds.as_deref_mut() {
+                    if selected_read {
                         res += 1;
-                        unsafe { FD_SET(index as _, set) };
+                        selected_readfds.0.set(index, true);
                     }
-                    if selected_write && let Some(set) = writefds.as_deref_mut() {
+                    if selected_write {
                         res += 1;
-                        unsafe { FD_SET(index as _, set) };
+                        selected_writefds.0.set(index, true);
                     }
-                    if selected_except && let Some(set) = exceptfds.as_deref_mut() {
+                    if selected_except {
                         res += 1;
-                        unsafe { FD_SET(index as _, set) };
+                        selected_exceptfds.0.set(index, true);
                     }
                 }
                 if res > 0 {
+                    write_fd_set(readfds.as_deref_mut(), &selected_readfds, nfds as _);
+                    write_fd_set(writefds.as_deref_mut(), &selected_writefds, nfds as _);
+                    write_fd_set(exceptfds.as_deref_mut(), &selected_exceptfds, nfds as _);
                     return Ok(res as _);
                 }
 
                 Err(AxError::WouldBlock)
             }),
-        )) {
+        ));
+        match result {
             Ok(r) => r,
-            Err(_) => Ok(0),
+            Err(_) => {
+                let empty = FdSet(Bitmap::new());
+                write_fd_set(readfds, &empty, nfds as _);
+                write_fd_set(writefds, &empty, nfds as _);
+                write_fd_set(exceptfds, &empty, nfds as _);
+                Ok(0)
+            }
         }
     })
 }

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

@@ -6,7 +6,7 @@ use starry_vm::{VmMutPtr, VmPtr};
 use syscalls::Sysno;
 
 use super::{
-    AsThread, SyscallRestartInfo, SyscallTraceState, TimerState, check_signals,
+    AsThread, SyscallRestartInfo, SyscallTraceState, TimerState, check_signals, poll_process_timer,
     ptrace_stop_current, ptrace_syscall_stop_current, raise_signal_fatal, set_timer_state,
     unblock_next_signal, wait_existing_ptrace_stop_current,
 };
@@ -209,6 +209,11 @@ pub fn new_user_task(name: &str, mut uctx: UserContext, set_child_tid: usize) ->
                 }
 
                 if !unblock_next_signal() {
+                    // POSIX timers are also driven by the alarm task, but polling
+                    // here closes the window where an expired timer is only noticed
+                    // after the current syscall returns to userspace.
+                    poll_process_timer(thr.proc_data.proc.pid());
+
                     let eintr_code = -(ax_errno::LinuxError::EINTR.code() as isize);
                     let restart = if is_syscall
                         && (uctx.retval() as isize) == eintr_code

os/arceos/api/arceos_api/src/imp/task.rs

@@ -1,10 +1,12 @@
+#[track_caller]
 pub fn ax_sleep_until(deadline: crate::time::AxTimeValue) {
     #[cfg(feature = "multitask")]
     ax_task::sleep_until(deadline);
     #[cfg(not(feature = "multitask"))]
     ax_hal::time::busy_wait_until(deadline);
 }
 
+#[track_caller]
 pub fn ax_yield_now() {
     #[cfg(feature = "multitask")]
     ax_task::yield_now();
@@ -16,6 +18,7 @@ pub fn ax_yield_now() {
     }
 }
 
+#[track_caller]
 pub fn ax_exit(_exit_code: i32) -> ! {
     #[cfg(feature = "multitask")]
     ax_task::exit(_exit_code);
@@ -78,6 +81,7 @@ cfg_task! {
         }
     }
 
+    #[track_caller]
     pub fn ax_wait_for_exit(task: AxTaskHandle) -> i32 {
         task.inner.join()
     }
@@ -93,6 +97,7 @@ cfg_task! {
         }
     }
 
+    #[track_caller]
     pub fn ax_set_current_affinity(cpumask: AxCpuMask) -> crate::AxResult {
         if ax_task::set_current_affinity(cpumask) {
             Ok(())
@@ -104,6 +109,7 @@ cfg_task! {
         }
     }
 
+    #[track_caller]
     pub fn ax_wait_queue_wait(wq: &AxWaitQueueHandle, timeout: Option<Duration>) -> bool {
         #[cfg(feature = "irq")]
         if let Some(dur) = timeout {
@@ -117,6 +123,7 @@ cfg_task! {
         false
     }
 
+    #[track_caller]
     pub fn ax_wait_queue_wait_until(
         wq: &AxWaitQueueHandle,
         until_condition: impl Fn() -> bool,

os/arceos/api/arceos_api/src/lib.rs

@@ -129,15 +129,18 @@ pub mod task {
         /// Current task is going to sleep, it will be woken up at the given deadline.
         ///
         /// If the feature `multitask` is not enabled, it uses busy-wait instead
+        #[track_caller]
         pub fn ax_sleep_until(deadline: crate::time::AxTimeValue);
 
         /// Current task gives up the CPU time voluntarily, and switches to another
         /// ready task.
         ///
         /// If the feature `multitask` is not enabled, it does nothing.
+        #[track_caller]
         pub fn ax_yield_now();
 
         /// Exits the current task with the given exit code.
+        #[track_caller]
         pub fn ax_exit(exit_code: i32) -> !;
     }
 
@@ -154,18 +157,22 @@ pub mod task {
         ) -> AxTaskHandle;
         /// Waits for the given task to exit, and returns its exit code (the
         /// argument of [`ax_exit`]).
+        #[track_caller]
         pub fn ax_wait_for_exit(task: AxTaskHandle) -> i32;
         /// Sets the priority of the current task.
         pub fn ax_set_current_priority(prio: isize) -> crate::AxResult;
         /// Sets the cpu affinity of the current task.
+        #[track_caller]
         pub fn ax_set_current_affinity(cpumask: AxCpuMask) -> crate::AxResult;
         /// Blocks the current task and put it into the wait queue, until
         /// other tasks notify the wait queue, or the given duration has
         /// elapsed (if specified).
+        #[track_caller]
         pub fn ax_wait_queue_wait(wq: &AxWaitQueueHandle, timeout: Option<core::time::Duration>) -> bool;
         /// Blocks the current task and put it into the wait queue, until the
         /// given condition becomes true, or the given duration has elapsed
         /// (if specified).
+        #[track_caller]
         pub fn ax_wait_queue_wait_until(
             wq: &AxWaitQueueHandle,
             until_condition: impl Fn() -> bool,

os/arceos/api/arceos_posix_api/src/imp/pthread/mod.rs

@@ -90,12 +90,14 @@ impl Pthread {
         unsafe { core::ptr::NonNull::new(Self::current_ptr()).map(|ptr| ptr.as_ref()) }
     }
 
+    #[track_caller]
     fn exit_current(retval: *mut c_void) -> ! {
         let thread = Self::current().expect("fail to get current thread");
         unsafe { *thread.retval.result.get() = retval };
         ax_task::exit(0);
     }
 
+    #[track_caller]
     fn join(ptr: ctypes::pthread_t) -> LinuxResult<*mut c_void> {
         if core::ptr::eq(ptr, Self::current_ptr() as _) {
             return Err(LinuxError::EDEADLK);
@@ -138,12 +140,14 @@ pub unsafe fn sys_pthread_create(
 }
 
 /// Exits the current thread. The value `retval` will be returned to the joiner.
+#[track_caller]
 pub fn sys_pthread_exit(retval: *mut c_void) -> ! {
     debug!("sys_pthread_exit <= {:#x}", retval as usize);
     Pthread::exit_current(retval);
 }
 
 /// Waits for the given thread to exit, and stores the return value in `retval`.
+#[track_caller]
 pub unsafe fn sys_pthread_join(thread: ctypes::pthread_t, retval: *mut *mut c_void) -> c_int {
     debug!("sys_pthread_join <= {:#x}", retval as usize);
     syscall_body!(sys_pthread_join, {

os/arceos/api/arceos_posix_api/src/imp/task.rs

@@ -4,6 +4,7 @@ use core::ffi::c_int;
 ///
 /// For single-threaded configuration (`multitask` feature is disabled), we just
 /// relax the CPU and wait for incoming interrupts.
+#[track_caller]
 pub fn sys_sched_yield() -> c_int {
     #[cfg(feature = "multitask")]
     ax_task::yield_now();
@@ -31,6 +32,7 @@ pub fn sys_getpid() -> c_int {
 }
 
 /// Exit current task
+#[track_caller]
 pub fn sys_exit(exit_code: c_int) -> ! {
     debug!("sys_exit <= {exit_code}");
     #[cfg(feature = "multitask")]

os/arceos/api/arceos_posix_api/src/imp/time.rs

@@ -63,6 +63,7 @@ pub unsafe fn sys_clock_gettime(clk: ctypes::clockid_t, ts: *mut ctypes::timespe
 /// Sleep some nanoseconds
 ///
 /// TODO: should be woken by signals, and set errno
+#[track_caller]
 pub unsafe fn sys_nanosleep(req: *const ctypes::timespec, rem: *mut ctypes::timespec) -> c_int {
     syscall_body!(sys_nanosleep, {
         unsafe {

os/arceos/modules/axtask/src/api.rs

@@ -231,6 +231,7 @@ pub fn set_priority(prio: isize) -> bool {
 /// Returns `true` if the affinity is set successfully.
 ///
 /// TODO: support set the affinity for other tasks.
+#[track_caller]
 pub fn set_current_affinity(cpumask: AxCpuMask) -> bool {
     might_sleep();
 
@@ -282,13 +283,15 @@ pub(crate) fn yield_now_unchecked() {
 /// Current task is going to sleep for the given duration.
 ///
 /// If the feature `irq` is not enabled, it uses busy-wait instead.
+#[track_caller]
 pub fn sleep(dur: core::time::Duration) {
     sleep_until(ax_hal::time::wall_time() + dur);
 }
 
 /// Current task is going to sleep, it will be woken up at the given deadline.
 ///
 /// If the feature `irq` is not enabled, it uses busy-wait instead.
+#[track_caller]
 pub fn sleep_until(deadline: ax_hal::time::TimeValue) {
     #[cfg(feature = "irq")]
     might_sleep();
@@ -299,6 +302,7 @@ pub fn sleep_until(deadline: ax_hal::time::TimeValue) {
 }
 
 /// Exits the current task.
+#[track_caller]
 pub fn exit(exit_code: i32) -> ! {
     might_sleep();
 
@@ -316,6 +320,10 @@ fn current_preempt_count() -> usize {
     }
 }
 
+fn current_task_id() -> Option<u64> {
+    current_may_uninit().map(|curr| curr.id().as_u64())
+}
+
 /// Returns whether the current context is atomic, meaning sleeping or
 /// rescheduling is not allowed.
 ///
@@ -343,9 +351,11 @@ pub fn might_sleep() {
     if in_atomic_context() {
         panic!(
             "sleeping or rescheduling is not allowed in atomic context: irq_enabled={}, \
-             preempt_count={}",
+             preempt_count={}, cpu_id={}, task_id={:?}",
             ax_hal::asm::irqs_enabled(),
-            current_preempt_count()
+            current_preempt_count(),
+            ax_hal::percpu::this_cpu_id(),
+            current_task_id()
         );
     }
 }

os/arceos/modules/axtask/src/future/mod.rs

@@ -55,6 +55,7 @@ impl Wake for AxWaker {
 /// from signal delivery), this function yields the CPU to allow signal
 /// processing on the return-to-userspace path. The future will be re-polled
 /// after the yield.
+#[track_caller]
 pub fn block_on<F: IntoFuture>(f: F) -> F::Output {
     crate::api::might_sleep();
 

os/arceos/modules/axtask/src/task.rs

@@ -202,6 +202,7 @@ impl TaskInner {
     /// Wait for the task to exit, and return the exit code.
     ///
     /// It will return immediately if the task has already exited (but not dropped).
+    #[track_caller]
     pub fn join(&self) -> i32 {
         crate::api::might_sleep();
         self.wait_for_exit

os/arceos/modules/axtask/src/wait_queue.rs

@@ -73,6 +73,7 @@ impl WaitQueue {
 
     /// Blocks the current task and put it into the wait queue, until other task
     /// notifies it.
+    #[track_caller]
     pub fn wait(&self) {
         crate::api::might_sleep();
         current_run_queue::<NoPreemptIrqSave>().blocked_resched(self.queue.lock());
@@ -84,6 +85,7 @@ impl WaitQueue {
     ///
     /// Note that even other tasks notify this task, it will not wake up until
     /// the condition becomes true.
+    #[track_caller]
     pub fn wait_until<F>(&self, condition: F)
     where
         F: Fn() -> bool,
@@ -106,6 +108,7 @@ impl WaitQueue {
     /// Blocks the current task and put it into the wait queue, until other tasks
     /// notify it, or the given duration has elapsed.
     #[cfg(feature = "irq")]
+    #[track_caller]
     pub fn wait_timeout(&self, dur: core::time::Duration) -> bool {
         crate::api::might_sleep();
         let mut rq = current_run_queue::<NoPreemptIrqSave>();
@@ -133,6 +136,7 @@ impl WaitQueue {
     /// Note that even other tasks notify this task, it will not wake up until
     /// the above conditions are met.
     #[cfg(feature = "irq")]
+    #[track_caller]
     pub fn wait_timeout_until<F>(&self, dur: core::time::Duration, condition: F) -> bool
     where
         F: Fn() -> bool,