fix(sched): closed the kill-before-block race in all blocking paths

kernel/sched/sched.cpp

@@ -186,10 +186,44 @@ bool is_kill_pending() {
 
 __PRIVILEGED_CODE void force_wake_for_kill(task* t) {
     __atomic_store_n(&t->kill_pending, 1, __ATOMIC_RELEASE);
+
+    // Pairs with block_task_interrupted_by_kill: the wake below sees BLOCKED,
+    // or the blocker's kill check sees kill_pending. Never neither.
+    __atomic_thread_fence(__ATOMIC_SEQ_CST);
     timer::cancel_sleep(t);
     wake(t);
 }
 
+/**
+ * @note Privilege: **required**
+ */
+__PRIVILEGED_CODE void prepare_to_block_task() {
+    __atomic_store_n(&current()->state, TASK_STATE_BLOCKED, __ATOMIC_RELEASE);
+}
+
+/**
+ * @note Privilege: **required**
+ */
+__PRIVILEGED_CODE bool block_task_interrupted_by_kill() {
+    // Fence pairs with the one in force_wake_for_kill.
+    __atomic_thread_fence(__ATOMIC_SEQ_CST);
+    return __atomic_load_n(&current()->kill_pending, __ATOMIC_ACQUIRE);
+}
+
+/**
+ * @note Privilege: **required**
+ */
+__PRIVILEGED_CODE void cancel_block_task() {
+    task* self = current();
+    uint32_t expected = TASK_STATE_BLOCKED;
+    if (!__atomic_compare_exchange_n(&self->state, &expected, TASK_STATE_RUNNING,
+                                      false, __ATOMIC_ACQ_REL, __ATOMIC_RELAXED)) {
+        // A wake already claimed this task READY and will requeue it once
+        // it is off-CPU. Yield so that handoff can complete.
+        yield();
+    }
+}
+
 /**
  * @note Privilege: **required**
  */
@@ -378,8 +412,17 @@ __PRIVILEGED_CODE void sleep_ns(uint64_t ns) {
     }
 
     uint64_t deadline = clock::now_ns() + ns;
-    self->state = TASK_STATE_BLOCKED;
+
+    prepare_to_block_task();
     timer::schedule_sleep(self, deadline);
+
+    if (block_task_interrupted_by_kill()) {
+        // Killed before or during sleep entry: do not serve the sleep.
+        timer::cancel_sleep(self);
+        cancel_block_task();
+        return;
+    }
+
     yield();
 }
 

kernel/sched/sched.h

@@ -110,12 +110,33 @@ __PRIVILEGED_CODE void wake(task* t);
 
 /**
  * @brief Mark a task for termination and wake it if blocked.
- * Sets kill_pending, cancels any timer sleep, and wakes via CAS.
- * Fire-and-forget: does not wait for the task to actually die.
+ * Fire-and-forget: the target is force-woken now or observes the kill
+ * at its next killable blocking attempt (sleep, futex, poll).
  * @note Privilege: **required**
  */
 __PRIVILEGED_CODE void force_wake_for_kill(task* t);
 
+/**
+ * @brief Publish intent to block: moves the current task to BLOCKED.
+ * Pair with block_task_interrupted_by_kill before yielding.
+ * @note Privilege: **required**
+ */
+__PRIVILEGED_CODE void prepare_to_block_task();
+
+/**
+ * @brief True if a kill arrived since prepare_to_block_task.
+ * On true, the caller must unwind its wait entry and call cancel_block_task.
+ * @note Privilege: **required**
+ */
+__PRIVILEGED_CODE bool block_task_interrupted_by_kill();
+
+/**
+ * @brief Revert an unfinished block after the caller unwound its entry.
+ * Yields once if a concurrent wake already claimed the task.
+ * @note Privilege: **required**
+ */
+__PRIVILEGED_CODE void cancel_block_task();
+
 /**
  * @brief Check if the current task has been marked for termination.
  * Safe to call from any kernel context with a valid per-CPU base.

kernel/sync/futex.cpp

@@ -73,7 +73,7 @@ __PRIVILEGED_CODE int32_t futex_wait(uintptr_t uaddr, uint32_t expected,
         return -11; // EAGAIN
     }
 
-    self->state = sched::TASK_STATE_BLOCKED;
+    sched::prepare_to_block_task();
     bucket->waiters.push_back(&waiter);
 
     if (timeout_ns > 0) {
@@ -82,6 +82,18 @@ __PRIVILEGED_CODE int32_t futex_wait(uintptr_t uaddr, uint32_t expected,
     }
 
     spin_unlock_irqrestore(bucket->lock, irq);
+
+    if (sched::block_task_interrupted_by_kill()) {
+        // Killed during futex entry: unwind waiter and timer, don't block.
+        timer::cancel_sleep(self);
+        irq = spin_lock_irqsave(bucket->lock);
+        if (waiter.link.is_linked()) {
+            bucket->waiters.remove(&waiter);
+        }
+        spin_unlock_irqrestore(bucket->lock, irq);
+        sched::cancel_block_task();
+        return -4; // EINTR
+    }
     sched::yield();
 
     // Cancel any outstanding timer to prevent spurious wakes of future

kernel/sync/poll.cpp

@@ -27,20 +27,6 @@ __PRIVILEGED_CODE void poll_subscribe(poll_table& pt, wait_queue& wq) {
     spin_unlock_irqrestore(wq.lock, irq);
 }
 
-__PRIVILEGED_CODE static void resolve_aborted_poll_block(sched::task* self) {
-    uint32_t expected = sched::TASK_STATE_BLOCKED;
-    if (__atomic_compare_exchange_n(&self->state, &expected,
-            sched::TASK_STATE_RUNNING, false,
-            __ATOMIC_ACQ_REL, __ATOMIC_ACQUIRE)) {
-        return;
-    }
-
-    // A concurrent wake already moved us to READY and queued a runqueue entry.
-    // Yield once so the scheduler can consume that entry instead of leaving
-    // sched_link live while we keep executing in C.
-    sched::yield();
-}
-
 __PRIVILEGED_CODE bool poll_wait(poll_table& pt, uint64_t timeout_ns) {
     if (__atomic_load_n(&pt.triggered, __ATOMIC_ACQUIRE)) {
         return true;
@@ -55,18 +41,19 @@ __PRIVILEGED_CODE bool poll_wait(poll_table& pt, uint64_t timeout_ns) {
         return false;
     }
 
-    self->state = sched::TASK_STATE_BLOCKED;
+    sched::prepare_to_block_task();
 
-    // Re-check after setting BLOCKED to close races where a source fires
-    // or kill arrives between the initial checks and the state transition.
+    if (sched::block_task_interrupted_by_kill()) {
+        sched::cancel_block_task();
+        return false;
+    }
+
+    // The kill check's fence also orders this re-check against the BLOCKED
+    // store, closing the race where a source fires during the transition.
     if (__atomic_load_n(&pt.triggered, __ATOMIC_ACQUIRE)) {
-        resolve_aborted_poll_block(self);
+        sched::cancel_block_task();
         return true;
     }
-    if (__atomic_load_n(&self->kill_pending, __ATOMIC_ACQUIRE)) {
-        resolve_aborted_poll_block(self);
-        return false;
-    }
 
     if (timeout_ns > 0) {
         uint64_t deadline = clock::now_ns() + timeout_ns;

kernel/sync/wait_queue.cpp

@@ -56,6 +56,8 @@ __PRIVILEGED_CODE static void notify_observers_and_unlock(
 }
 
 /**
+ * Not kill-interruptible on its own: a force-woken waiter re-blocks unless
+ * the caller loops on sched::is_kill_pending().
  * @note Privilege: **required**
  */
 __PRIVILEGED_CODE
@@ -69,7 +71,7 @@ irq_state wait(wait_queue& wq, spinlock& lock, irq_state saved) {
     }
 
     spin_lock(wq.lock);
-    self->state = sched::TASK_STATE_BLOCKED;
+    sched::prepare_to_block_task();
     wq.waiters.push_back(self);
     spin_unlock(wq.lock);
 

kernel/tests/framework/helpers.h

@@ -2,23 +2,25 @@
 #define STELLUX_TESTS_FRAMEWORK_HELPERS_H
 
 #include "common/types.h"
+#include "clock/clock.h"
 
 namespace test_helpers {
 
-constexpr uint64_t SPIN_TIMEOUT = 100000000;
+// Wall-clock bound: iteration counts vary ~100x across hosts and emulators.
+constexpr uint64_t SPIN_TIMEOUT_NS = 20000000000ULL; // 20s
 
 inline bool spin_wait(volatile uint32_t* flag) {
-    uint64_t spins = 0;
+    uint64_t deadline = clock::now_ns() + SPIN_TIMEOUT_NS;
     while (!__atomic_load_n(flag, __ATOMIC_ACQUIRE)) {
-        if (++spins > SPIN_TIMEOUT) return false;
+        if (clock::now_ns() > deadline) return false;
     }
     return true;
 }
 
 inline bool spin_wait_ge(volatile uint32_t* value, uint32_t target) {
-    uint64_t spins = 0;
+    uint64_t deadline = clock::now_ns() + SPIN_TIMEOUT_NS;
     while (__atomic_load_n(value, __ATOMIC_ACQUIRE) < target) {
-        if (++spins > SPIN_TIMEOUT) return false;
+        if (clock::now_ns() > deadline) return false;
     }
     return true;
 }

kernel/tests/sched/kill.test.cpp

@@ -13,6 +13,15 @@ using test_helpers::spin_wait;
 using test_helpers::spin_wait_ge;
 using test_helpers::brief_delay;
 
+// Deterministic handshake: wait until the task has actually blocked.
+static bool wait_until_blocked(sched::task* t) {
+    uint64_t deadline = clock::now_ns() + test_helpers::SPIN_TIMEOUT_NS;
+    while (__atomic_load_n(&t->state, __ATOMIC_ACQUIRE) != sched::TASK_STATE_BLOCKED) {
+        if (clock::now_ns() > deadline) return false;
+    }
+    return true;
+}
+
 TEST_SUITE(kill);
 
 // --- force_wake_kills_sleeping_task ---
@@ -52,7 +61,7 @@ TEST(kill, force_wake_kills_sleeping_task) {
         sched::enqueue(t);
     });
 
-    brief_delay();
+    ASSERT_TRUE(wait_until_blocked(t));
 
     RUN_ELEVATED({
         sched::force_wake_for_kill(t);
@@ -64,6 +73,28 @@ TEST(kill, force_wake_kills_sleeping_task) {
               static_cast<uint64_t>(2000000000)); // woke in < 2s, not 5s
 }
 
+// --- kill_before_sleep_aborts_sleep ---
+// A kill issued before the task ever runs must abort its sleep attempt
+// at the block commit point instead of being lost until natural expiry.
+
+TEST(kill, kill_before_sleep_aborts_sleep) {
+    g_sleep_kill_done = 0;
+    g_sleep_kill_was_pending = 0;
+    g_sleep_kill_elapsed_ns = 0;
+
+    RUN_ELEVATED({
+        sched::task* t = sched::create_kernel_task(sleep_kill_fn, nullptr, "kill_presleep");
+        ASSERT_NOT_NULL(t);
+        sched::force_wake_for_kill(t);
+        sched::enqueue(t);
+    });
+
+    ASSERT_TRUE(spin_wait(&g_sleep_kill_done));
+    EXPECT_EQ(__atomic_load_n(&g_sleep_kill_was_pending, __ATOMIC_ACQUIRE), 1u);
+    EXPECT_LT(__atomic_load_n(&g_sleep_kill_elapsed_ns, __ATOMIC_ACQUIRE),
+              static_cast<uint64_t>(2000000000)); // aborted, not slept for 5s
+}
+
 // --- force_wake_kills_blocked_on_wq ---
 // A task blocked on a wait queue is woken by force_wake_for_kill.
 // Verifies: task wakes, self-removes from wq, and sees kill_pending.
@@ -221,13 +252,15 @@ static volatile uint32_t g_ikp_before = 0xFF;
 static volatile uint32_t g_ikp_after = 0xFF;
 static volatile uint32_t g_ikp_done = 0;
 static volatile uint32_t g_ikp_flag_set = 0;
+static volatile uint32_t g_ikp_started = 0;
 
 static void ikp_fn(void*) {
     uint32_t before = 0;
     RUN_ELEVATED({
         before = sched::is_kill_pending() ? 1 : 0;
     });
     __atomic_store_n(&g_ikp_before, before, __ATOMIC_RELEASE);
+    __atomic_store_n(&g_ikp_started, 1, __ATOMIC_RELEASE);
 
     while (!__atomic_load_n(&g_ikp_flag_set, __ATOMIC_ACQUIRE)) {
         // busy wait for test driver to set kill_pending
@@ -247,6 +280,7 @@ TEST(kill, is_kill_pending_accessor) {
     g_ikp_after = 0xFF;
     g_ikp_done = 0;
     g_ikp_flag_set = 0;
+    g_ikp_started = 0;
 
     sched::task* t = nullptr;
     RUN_ELEVATED({
@@ -255,7 +289,7 @@ TEST(kill, is_kill_pending_accessor) {
         sched::enqueue(t);
     });
 
-    brief_delay();
+    ASSERT_TRUE(spin_wait(&g_ikp_started));
 
     RUN_ELEVATED({
         __atomic_store_n(&t->kill_pending, 1, __ATOMIC_RELEASE);