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.
This commit is contained in:
Matthew Leach
2025-12-24 22:05:57 +00:00
committed by Ashwin Naren
parent bd21276368
commit 69f50fef18
3 changed files with 53 additions and 4 deletions

View File

@@ -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",

View File

@@ -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;
}

View File

@@ -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;
}
_ => {}
}
}
}