From 69f50fef1804647645366182633b13645003ebbe Mon Sep 17 00:00:00 2001
From: Matthew Leach <dev@mattleach.net>
Date: Wed, 24 Dec 2025 22:05:57 +0000
Subject: [PATCH] sched: fix futures race-condition

When a future returns `Poll::Pending`, there is a window where, if a
waker is called, prior to the sched code setting the task's state to
`Sleeping`, the wake-up could be lost.  We get around this by
introducing a new state `Woken`.

A waker will set a `Running` task to this state. The sched code then
detects this and *does not* set the task's state to `Sleeping`, instead
it leaves it as running and attempts to re-schedule.
---
 src/process/mod.rs    |  2 ++
 src/sched/uspc_ret.rs | 40 ++++++++++++++++++++++++++++++++++++++--
 src/sched/waker.rs    | 15 +++++++++++++--
 3 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/src/process/mod.rs b/src/process/mod.rs
index f7a5255..cadefe0 100644
--- a/src/process/mod.rs
+++ b/src/process/mod.rs
@@ -117,6 +117,7 @@ impl TaskDescriptor {
 pub enum TaskState {
     Running,
     Runnable,
+    Woken,
     Stopped,
     Sleeping,
     Finished,
@@ -127,6 +128,7 @@ impl Display for TaskState {
         let state_str = match self {
             TaskState::Running => "R",
             TaskState::Runnable => "R",
+            TaskState::Woken => "W",
             TaskState::Stopped => "T",
             TaskState::Sleeping => "S",
             TaskState::Finished => "Z",
diff --git a/src/sched/uspc_ret.rs b/src/sched/uspc_ret.rs
index 8d601d9..8593e34 100644
--- a/src/sched/uspc_ret.rs
+++ b/src/sched/uspc_ret.rs
@@ -115,7 +115,25 @@ pub fn dispatch_userspace_task(ctx: *mut UserCtx) {
                         }
                         Poll::Pending => {
                             task.ctx.lock_save_irq().put_signal_work(signal_work);
-                            *task.state.lock_save_irq() = TaskState::Sleeping;
+                            let mut task_state = task.state.lock_save_irq();
+
+                            match *task_state {
+                                // The main path we expect to take to sleep the
+                                // task.
+                                TaskState::Running => *task_state = TaskState::Sleeping,
+                                // If we were woken between the future returning
+                                // `Poll::Pending` and acquiring the lock above,
+                                // the waker will have put us into this state.
+                                // Transition back to `Running` since we're
+                                // ready to progress with more work.
+                                TaskState::Woken => *task_state = TaskState::Running,
+                                // We should never get here for any other state.
+                                s => {
+                                    unreachable!(
+                                        "Unexpected task state {s:?} during signal task sleep"
+                                    );
+                                }
+                            }
 
                             state = State::PickNewTask;
                             continue;
@@ -166,7 +184,25 @@ pub fn dispatch_userspace_task(ctx: *mut UserCtx) {
                             // state to sleeping so it's not scheduled again and
                             // search for another task to execute.
                             task.ctx.lock_save_irq().put_kernel_work(kern_work);
-                            *task.state.lock_save_irq() = TaskState::Sleeping;
+                            let mut task_state = task.state.lock_save_irq();
+
+                            match *task_state {
+                                // The main path we expect to take to sleep the
+                                // task.
+                                TaskState::Running => *task_state = TaskState::Sleeping,
+                                // If we were woken between the future returning
+                                // `Poll::Pending` and acquiring the lock above,
+                                // the waker will have put us into this state.
+                                // Transition back to `Running` since we're
+                                // ready to progress with more work.
+                                TaskState::Woken => *task_state = TaskState::Running,
+                                // We should never get here for any other state.
+                                s => {
+                                    unreachable!(
+                                        "Unexpected task state {s:?} during kernel task sleep"
+                                    );
+                                }
+                            }
                             state = State::PickNewTask;
                             continue;
                         }
diff --git a/src/sched/waker.rs b/src/sched/waker.rs
index b9652a1..20dd21a 100644
--- a/src/sched/waker.rs
+++ b/src/sched/waker.rs
@@ -13,8 +13,19 @@ unsafe fn wake_waker(data: *const ()) {
         && let Some(proc) = proc.upgrade()
     {
         let mut state = proc.lock_save_irq();
-        if *state == TaskState::Sleeping {
-            *state = TaskState::Runnable;
+        match *state {
+            // If the task has been put to sleep, then wake it up.
+            TaskState::Sleeping => {
+                *state = TaskState::Runnable;
+            }
+            // If the task is running, mark it so it doesn't actually go to
+            // sleep when poll returns. This covers the small race-window
+            // between a future returning `Poll::Pending` and the sched setting
+            // the state to sleeping.
+            TaskState::Running => {
+                *state = TaskState::Woken;
+            }
+            _ => {}
         }
     }
 }