diff --git a/include/elio/runtime/scheduler.hpp b/include/elio/runtime/scheduler.hpp index 911f706f..e951ed40 100644 --- a/include/elio/runtime/scheduler.hpp +++ b/include/elio/runtime/scheduler.hpp @@ -962,12 +962,24 @@ 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(); @@ -975,8 +987,19 @@ inline void worker_thread::run() { 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(); } } diff --git a/tests/unit/test_scheduler.cpp b/tests/unit/test_scheduler.cpp index a28fedde..56a599bc 100644 --- a/tests/unit/test_scheduler.cpp +++ b/tests/unit/test_scheduler.cpp @@ -1,12 +1,15 @@ #include #include #include +#include #include #include +#include #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 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 occupy_worker_task(std::atomic* 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 io_after_resume_task(std::chrono::milliseconds sleep_for_dur, + std::atomic* 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 occupier_running{false}; + std::atomic 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(); +}