Fix scheduler shrink issue causing I/O orphaned by retired workers

include/elio/runtime/scheduler.hpp

@@ -962,21 +962,44 @@ inline void worker_thread::run() {
     // This ensures shutdown() returns only when all submitted tasks have
     // fully completed (including coroutine cleanup and lambda destruction).
     //
+    // Retirement vs. shutdown: if the SCHEDULER is still running, this
+    // worker is being retired by a pool shrink (set_thread_count), not a
+    // full shutdown. In that case queued tasks must be redistributed to
+    // active workers rather than resumed here — resuming a queued coroutine
+    // on the retiring worker lets it start a fresh async I/O that binds to
+    // this worker's io_context (via worker_thread::current()), which stops
+    // being polled once this thread exits, orphaning the operation and
+    // hanging the coroutine. shutdown_force() flips the scheduler's
+    // running_ to false BEFORE joining workers, so a genuine shutdown still
+    // takes the run_task() path below and drains locally.
+    //
     // Note: pop() (steal-style CAS path) — NOT pop_local(false). Other
     // workers may still be inside try_steal()->steal_task() against this
     // worker's deque because they haven't yet observed our running_=false.
     // Chase-Lev requires either single-threaded owner pop OR concurrent
     // steal+pop synchronized via the seq_cst fence inside pop(); the
     // single-thread fast path of pop_local() races with those stealers.
+    const bool retiring = scheduler_->is_running();
     while (true) {
         drain_inbox();
         void* addr = queue_->pop();
         if (!addr) break;
 
         auto handle = std::coroutine_handle<>::from_address(addr);
         if (handle && !handle.done()) {
-            needs_sync_ = true;  // Conservatively ensure memory visibility for drained tasks
-            run_task(handle);
+            if (retiring) {
+                // Hand the task back to the scheduler; do_spawn() routes it to
+                // a surviving worker (num_threads_ was already lowered before
+                // this worker was stopped, so it never lands back here). The
+                // release fence inside do_spawn() carries frame visibility, so
+                // no needs_sync_ handshake is required on the receiving worker.
+                scheduler_->spawn(handle);
+            } else {
+                needs_sync_ = true;  // Conservatively ensure memory visibility for drained tasks
+                run_task(handle);
+            }
+        } else if (handle) {
+            handle.destroy();
         }
     }
 

tests/unit/test_scheduler.cpp

@@ -1,12 +1,15 @@
 #include <catch2/catch_test_macros.hpp>
 #include <elio/runtime/scheduler.hpp>
 #include <elio/coro/task.hpp>
+#include <elio/time/timer.hpp>
 #include <atomic>
 #include <chrono>
+#include <thread>
 #include "../test_main.cpp"  // For scaled timeouts
 
 using namespace elio::runtime;
 using namespace elio::coro;
+using namespace elio::time;
 using namespace elio::test;
 
 // Standalone task functions to avoid lambda capture lifetime issues
@@ -26,6 +29,27 @@ task<void> empty_task() {
     co_return;
 }
 
+// Occupies a worker thread with a plain (non-coroutine) sleep so the worker
+// cannot drain its inbox while this runs. Sets `running` once it starts so the
+// test can synchronize on the worker actually being busy. Does NOT register any
+// async I/O, so the retiring worker's io_context stays at pending_count()==0.
+task<void> occupy_worker_task(std::atomic<bool>* running,
+                              std::chrono::milliseconds busy_for) {
+    running->store(true, std::memory_order_release);
+    std::this_thread::sleep_for(busy_for);
+    co_return;
+}
+
+// Starts an async I/O (sleep_for goes through the current worker's io_context)
+// and only sets `done` after it resumes. If this coroutine is resumed on a
+// worker that is about to exit, the I/O binds to that worker's io_context and
+// is orphaned — `done` would then never become true.
+task<void> io_after_resume_task(std::chrono::milliseconds sleep_for_dur,
+                                std::atomic<bool>* done) {
+    co_await sleep_for(sleep_for_dur);
+    done->store(true, std::memory_order_release);
+}
+
 } // namespace
 
 TEST_CASE("Scheduler construction", "[scheduler]") {
@@ -243,12 +267,79 @@ TEST_CASE("Scheduler handles empty spawn", "[scheduler]") {
 
 TEST_CASE("Scheduler handles spawn before start", "[scheduler]") {
     scheduler sched(2);
-    
+
     // Should not crash, but task won't execute (scheduler not running)
     sched.go(empty_task);
-    
+
     // Now start - but the task was already queued
     sched.start();
     std::this_thread::sleep_for(scaled_ms(100));
     sched.shutdown();
 }
+
+// Regression test for the "shrink orphans I/O started by retired workers"
+// bug: when set_thread_count() retires a worker, any task still queued on
+// that worker must NOT start fresh async I/O on the retiring worker's
+// io_context (which stops being polled once the worker exits). The fix
+// redistributes such queued tasks to a surviving worker during the drain
+// phase instead of resuming them locally.
+//
+// Setup that deterministically forces the bug's preconditions:
+//   1. Pin an occupier task to worker 1 that busies the thread with a plain
+//      sleep. While worker 1 is inside run_task() it cannot drain its inbox.
+//   2. Once the occupier is confirmed running, pin an I/O-starting task to
+//      worker 1. Because worker 1 is busy, this task sits in worker 1's MPSC
+//      inbox. It cannot be stolen (stealing only touches the deque), so it
+//      stays pinned to the retiring worker.
+//   3. Call set_thread_count(1). The scheduler lowers the thread count, sees
+//      worker 1's io_context has 0 pending ops (the occupier does no async
+//      I/O and the I/O task hasn't run yet), then stop()s + joins worker 1.
+//   4. The join blocks until the occupier's sleep ends. Worker 1 then enters
+//      its drain phase, moves the I/O task from inbox to deque, and pops it.
+//
+// Without the fix: worker 1 resumes the I/O task locally, its co_await
+// sleep_for binds to worker 1's (soon-orphaned) io_context, worker 1 exits,
+// and the sleep never fires — `io_done` stays false.
+// With the fix: worker 1 redistributes the task to worker 0, which polls its
+// I/O to completion — `io_done` becomes true.
+TEST_CASE("Scheduler shrink does not orphan I/O from a retiring worker",
+          "[scheduler]") {
+    scheduler sched(2);
+    sched.start();
+
+    std::atomic<bool> occupier_running{false};
+    std::atomic<bool> io_done{false};
+
+    // (1) Occupy worker 1 for long enough that we can queue the I/O task and
+    // trigger the shrink while it is still busy.
+    sched.go_to(1, occupy_worker_task, &occupier_running, scaled_ms(300));
+
+    // (2) Wait until worker 1 is genuinely busy running the occupier.
+    auto deadline = std::chrono::steady_clock::now() + scaled_sec(5);
+    while (!occupier_running.load(std::memory_order_acquire) &&
+           std::chrono::steady_clock::now() < deadline) {
+        std::this_thread::sleep_for(std::chrono::milliseconds(1));
+    }
+    REQUIRE(occupier_running.load(std::memory_order_acquire));
+
+    // (3) Queue an I/O-starting task pinned to the (now busy) worker 1. It
+    // parks in worker 1's inbox because worker 1 can't drain while running the
+    // occupier.
+    sched.go_to(1, io_after_resume_task, scaled_ms(40), &io_done);
+
+    // (4) Shrink to a single worker. This retires worker 1 while the I/O task
+    // is still sitting in its inbox. Not called from a worker thread, so it is
+    // allowed to join.
+    sched.set_thread_count(1);
+    REQUIRE(sched.num_threads() == 1);
+
+    // The redistributed I/O task must complete on the surviving worker.
+    auto io_deadline = std::chrono::steady_clock::now() + scaled_sec(5);
+    while (!io_done.load(std::memory_order_acquire) &&
+           std::chrono::steady_clock::now() < io_deadline) {
+        std::this_thread::sleep_for(std::chrono::milliseconds(1));
+    }
+    REQUIRE(io_done.load(std::memory_order_acquire));
+
+    sched.shutdown();
+}