From 04088e0f8f585476dab1c7945e73bd0433a6181b Mon Sep 17 00:00:00 2001 From: Coldwings Date: Wed, 22 Jul 2026 10:36:08 +0800 Subject: [PATCH 1/3] fix(runtime): propagate targeted spawn affinity --- include/elio/runtime/scheduler.hpp | 18 ++- tests/unit/test_scheduler.cpp | 217 +++++++++++++++++++++++++++++ 2 files changed, 233 insertions(+), 2 deletions(-) diff --git a/include/elio/runtime/scheduler.hpp b/include/elio/runtime/scheduler.hpp index 64a22e40..54f1bf2a 100644 --- a/include/elio/runtime/scheduler.hpp +++ b/include/elio/runtime/scheduler.hpp @@ -113,10 +113,22 @@ namespace detail { /// done in promise_base::~promise_base via the on_spawn_completion_ /// callback, which fires whenever the frame is destroyed regardless of /// whether the body ever ran. + template + void inherit_spawn_affinity(Task& task) noexcept { + auto* wrapper = coro::promise_base::current_frame(); + auto child = coro::detail::task_access::handle(task); + if (wrapper && wrapper->has_affinity() && child && + !child.promise().has_affinity()) { + child.promise().set_affinity(wrapper->affinity()); + } + } + template requires (std::invocable && is_task_v>) coro::task callable_wrapper_void(F f, Args... args) { - co_await std::invoke(std::move(f), std::move(args)...); + auto task = std::invoke(std::move(f), std::move(args)...); + inherit_spawn_affinity(task); + co_await task; } /// Wrapper coroutine for joinable spawn (go_joinable). @@ -124,7 +136,9 @@ namespace detail { template requires (std::invocable && is_task_v>) auto callable_wrapper(F f, Args... args) -> std::invoke_result_t { - co_return co_await std::invoke(std::move(f), std::move(args)...); + auto task = std::invoke(std::move(f), std::move(args)...); + inherit_spawn_affinity(task); + co_return co_await task; } } // namespace detail diff --git a/tests/unit/test_scheduler.cpp b/tests/unit/test_scheduler.cpp index c91760c8..6b8955e4 100644 --- a/tests/unit/test_scheduler.cpp +++ b/tests/unit/test_scheduler.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include "../test_main.cpp" // For scaled timeouts @@ -158,6 +159,200 @@ task cancellable_recv_on_worker( completed->store(true, std::memory_order_release); } +enum class spawn_affinity_mode { + inherit, + clear, + move_to_worker_one, +}; + +class manual_schedule_gate { +public: + class awaitable { + public: + explicit awaitable(manual_schedule_gate& gate) noexcept : gate_(gate) {} + + bool await_ready() const noexcept { return false; } + + bool await_suspend(std::coroutine_handle<> handle) noexcept { + gate_.handle_.store(handle.address(), std::memory_order_release); + gate_.published_.store(true, std::memory_order_release); + gate_.published_.notify_all(); + return true; + } + + void await_resume() const noexcept {} + + private: + manual_schedule_gate& gate_; + }; + + [[nodiscard]] awaitable wait() noexcept { return awaitable{*this}; } + + [[nodiscard]] bool published() const noexcept { + return published_.load(std::memory_order_acquire); + } + + void resume() noexcept { + void* address = handle_.exchange(nullptr, std::memory_order_acq_rel); + if (address) { + schedule_handle(std::coroutine_handle<>::from_address(address)); + } + } + +private: + std::atomic handle_{nullptr}; + std::atomic published_{false}; +}; + +struct spawn_affinity_observation { + std::shared_ptr context; + std::atomic first_worker{NO_AFFINITY}; + std::atomic resumed_worker{NO_AFFINITY}; + std::atomic completion_worker{NO_AFFINITY}; + std::atomic first_affinity{NO_AFFINITY}; + std::atomic resumed_affinity{NO_AFFINITY}; + std::atomic post_io_affinity{NO_AFFINITY}; + std::atomic io_staged{false}; + std::atomic recv_result{-1}; + std::atomic done{false}; +}; + +task observe_spawn_affinity( + int fd, + spawn_affinity_mode mode, + manual_schedule_gate* gate, + spawn_affinity_observation* observed) { + auto* frame = promise_base::current_frame(); + if (!frame) { + throw std::logic_error("spawn affinity test requires a current frame"); + } + observed->context = frame->execution_context(); + observed->first_worker.store(current_worker_id(), std::memory_order_release); + observed->first_affinity.store(frame->affinity(), std::memory_order_release); + + if (mode == spawn_affinity_mode::clear) { + co_await clear_affinity(); + } else if (mode == spawn_affinity_mode::move_to_worker_one) { + co_await set_affinity(1, false); + } + + co_await gate->wait(); + observed->resumed_worker.store( + current_worker_id(), std::memory_order_release); + observed->resumed_affinity.store( + frame->affinity(), std::memory_order_release); + + elio::io::detail::arm_next_cancellable_recv_staged_for_test( + observed->io_staged); + char byte = 0; + auto result = co_await elio::io::async_recv( + fd, &byte, 1, 0, cancel_token{}); + observed->recv_result.store(result.io.result, std::memory_order_release); + observed->completion_worker.store( + current_worker_id(), std::memory_order_release); + observed->post_io_affinity.store( + frame->effective_affinity(), std::memory_order_release); + observed->done.store(true, std::memory_order_release); + observed->done.notify_all(); +} + +task resume_spawn_affinity_probe(manual_schedule_gate* gate) { + elio::runtime::detail::reject_next_schedule_for_test.store( + true, std::memory_order_release); + gate->resume(); +} + +template +bool wait_for_scheduler_condition(Predicate&& predicate) { + const auto deadline = std::chrono::steady_clock::now() + scaled_sec(5); + while (!predicate() && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + return predicate(); +} + +void check_spawn_affinity_case(bool joinable, + spawn_affinity_mode mode, + size_t expected_worker, + size_t expected_affinity) { + int sockets[2] = {-1, -1}; + REQUIRE(::socketpair(AF_UNIX, + SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, + 0, sockets) == 0); + + elio::runtime::detail::reject_next_schedule_for_test.store( + false, std::memory_order_release); + scheduler sched(2); + sched.start(); + + manual_schedule_gate gate; + spawn_affinity_observation observed; + std::optional> joined; + if (joinable) { + joined.emplace(sched.go_joinable_to( + 0, observe_spawn_affinity, sockets[0], mode, &gate, &observed)); + } else { + sched.go_to( + 0, observe_spawn_affinity, sockets[0], mode, &gate, &observed); + } + + REQUIRE(wait_for_scheduler_condition([&] { return gate.published(); })); + REQUIRE(observed.context); + CHECK(observed.first_worker.load(std::memory_order_acquire) == 0); + CHECK(observed.first_affinity.load(std::memory_order_acquire) == 0); + + const auto local_fallbacks_before = + elio::runtime::detail::local_schedule_fallbacks_for_test.load( + std::memory_order_acquire); + auto notifier = sched.go_joinable_to( + 1, resume_spawn_affinity_probe, &gate); + notifier.wait_destroyed(); + notifier.await_resume(); + + REQUIRE(wait_for_scheduler_condition([&] { + return observed.io_staged.load(std::memory_order_acquire); + })); + CHECK(observed.resumed_worker.load(std::memory_order_acquire) == + expected_worker); + CHECK(observed.resumed_affinity.load(std::memory_order_acquire) == + expected_affinity); + + auto* owner = sched.get_worker(expected_worker); + REQUIRE(owner != nullptr); + CHECK(observed.context->is_io_pin_owner( + expected_worker, owner->io_context().generation())); + CHECK(observed.context->effective_affinity() == expected_worker); + + const auto local_fallbacks_after = + elio::runtime::detail::local_schedule_fallbacks_for_test.load( + std::memory_order_acquire); + if (mode == spawn_affinity_mode::inherit) { + CHECK(local_fallbacks_after == local_fallbacks_before); + } else { + CHECK(local_fallbacks_after == local_fallbacks_before + 1); + } + + const char byte = 'x'; + REQUIRE(::write(sockets[1], &byte, 1) == 1); + REQUIRE(wait_for_scheduler_condition([&] { + return observed.done.load(std::memory_order_acquire); + })); + + if (joined) { + joined->wait_destroyed(); + joined->await_resume(); + } + CHECK(observed.recv_result.load(std::memory_order_acquire) == 1); + CHECK(observed.completion_worker.load(std::memory_order_acquire) == + expected_worker); + CHECK(observed.post_io_affinity.load(std::memory_order_acquire) == + expected_affinity); + REQUIRE(sched.shutdown(scaled_sec(5))); + + ::close(sockets[0]); + ::close(sockets[1]); +} + } // namespace TEST_CASE("Scheduler construction", "[scheduler]") { @@ -240,6 +435,28 @@ TEST_CASE("Scheduler routes an I/O-pinned migration request to its owner", REQUIRE(sched.shutdown(scaled_sec(5))); } +TEST_CASE("go_to affinity reaches the returned task across suspension", + "[scheduler][task][io][affinity][regression]") { + check_spawn_affinity_case( + false, spawn_affinity_mode::inherit, 0, 0); +} + +TEST_CASE("go_joinable_to affinity reaches the returned task and remains overridable", + "[scheduler][task][io][affinity][join_handle][regression]") { + SECTION("inherit spawn affinity") { + check_spawn_affinity_case( + true, spawn_affinity_mode::inherit, 0, 0); + } + SECTION("clear spawn affinity") { + check_spawn_affinity_case( + true, spawn_affinity_mode::clear, 1, NO_AFFINITY); + } + SECTION("replace spawn affinity") { + check_spawn_affinity_case( + true, spawn_affinity_mode::move_to_worker_one, 1, 1); + } +} + TEST_CASE("Scheduler pause/resume", "[scheduler]") { scheduler sched(2); sched.start(); From 9fbcd6b4dbf15e21e9c3ab91003e04103acb19ef Mon Sep 17 00:00:00 2001 From: Coldwings Date: Wed, 22 Jul 2026 10:40:28 +0800 Subject: [PATCH 2/3] test(runtime): make affinity notifier a coroutine --- tests/unit/test_scheduler.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/test_scheduler.cpp b/tests/unit/test_scheduler.cpp index 6b8955e4..e79e3dea 100644 --- a/tests/unit/test_scheduler.cpp +++ b/tests/unit/test_scheduler.cpp @@ -260,6 +260,7 @@ task resume_spawn_affinity_probe(manual_schedule_gate* gate) { elio::runtime::detail::reject_next_schedule_for_test.store( true, std::memory_order_release); gate->resume(); + co_return; } template From 3514da610722270b0c3f597306a8469fa0d8ad2f Mon Sep 17 00:00:00 2001 From: Coldwings Date: Wed, 22 Jul 2026 11:50:55 +0800 Subject: [PATCH 3/3] fix(coro): propagate affinity through task chains --- include/elio/coro/task.hpp | 30 +++ include/elio/runtime/scheduler.hpp | 18 +- tests/unit/test_scheduler.cpp | 282 +++++++++++++++++++---------- 3 files changed, 222 insertions(+), 108 deletions(-) diff --git a/include/elio/coro/task.hpp b/include/elio/coro/task.hpp index 93059597..5f64272f 100644 --- a/include/elio/coro/task.hpp +++ b/include/elio/coro/task.hpp @@ -411,6 +411,7 @@ class task { struct promise_type : promise_base { std::optional value_; std::coroutine_handle<> continuation_; + promise_base* awaiter_promise_ = nullptr; bool detached_ = false; std::shared_ptr> join_state_; @@ -495,6 +496,14 @@ class task { if constexpr (std::is_convertible_v) { parent = std::addressof(awaiter.promise()); parent_token = parent->execution_context()->get_cancel_token(); + // Directly awaited Elio tasks form one logical vthread. Carry the + // current user affinity into a new child frame; await_resume() + // copies the child's final value back so explicit changes made by + // deeper frames remain visible to the rest of the chain. + handle_.promise().awaiter_promise_ = parent; + if (!handle_.promise().has_affinity() && parent->has_affinity()) { + handle_.promise().set_affinity(parent->affinity()); + } } handle_.promise().link_parent_cancellation(std::move(parent_token)); handle_.promise().continuation_ = awaiter; @@ -505,6 +514,13 @@ class task { assert(handle_ && "cannot resume an empty task"); auto& promise = handle_.promise(); auto exception = promise.exception(); + if (auto* awaiter = std::exchange(promise.awaiter_promise_, nullptr)) { + if (promise.has_affinity()) { + awaiter->set_affinity(promise.affinity()); + } else { + awaiter->clear_affinity(); + } + } promise.detach_from_parent(); if (exception) std::rethrow_exception(exception); return std::move(*promise.value_); @@ -523,6 +539,7 @@ class task { struct promise_type : promise_base { std::coroutine_handle<> continuation_; + promise_base* awaiter_promise_ = nullptr; bool detached_ = false; std::shared_ptr> join_state_; @@ -599,6 +616,12 @@ class task { if constexpr (std::is_convertible_v) { parent = std::addressof(awaiter.promise()); parent_token = parent->execution_context()->get_cancel_token(); + // Keep user affinity continuous across the logical vthread even + // though each lazy task owns a distinct execution context. + handle_.promise().awaiter_promise_ = parent; + if (!handle_.promise().has_affinity() && parent->has_affinity()) { + handle_.promise().set_affinity(parent->affinity()); + } } handle_.promise().link_parent_cancellation(std::move(parent_token)); handle_.promise().continuation_ = awaiter; @@ -609,6 +632,13 @@ class task { assert(handle_ && "cannot resume an empty task"); auto& promise = handle_.promise(); auto exception = promise.exception(); + if (auto* awaiter = std::exchange(promise.awaiter_promise_, nullptr)) { + if (promise.has_affinity()) { + awaiter->set_affinity(promise.affinity()); + } else { + awaiter->clear_affinity(); + } + } promise.detach_from_parent(); if (exception) std::rethrow_exception(exception); } diff --git a/include/elio/runtime/scheduler.hpp b/include/elio/runtime/scheduler.hpp index 54f1bf2a..64a22e40 100644 --- a/include/elio/runtime/scheduler.hpp +++ b/include/elio/runtime/scheduler.hpp @@ -113,22 +113,10 @@ namespace detail { /// done in promise_base::~promise_base via the on_spawn_completion_ /// callback, which fires whenever the frame is destroyed regardless of /// whether the body ever ran. - template - void inherit_spawn_affinity(Task& task) noexcept { - auto* wrapper = coro::promise_base::current_frame(); - auto child = coro::detail::task_access::handle(task); - if (wrapper && wrapper->has_affinity() && child && - !child.promise().has_affinity()) { - child.promise().set_affinity(wrapper->affinity()); - } - } - template requires (std::invocable && is_task_v>) coro::task callable_wrapper_void(F f, Args... args) { - auto task = std::invoke(std::move(f), std::move(args)...); - inherit_spawn_affinity(task); - co_await task; + co_await std::invoke(std::move(f), std::move(args)...); } /// Wrapper coroutine for joinable spawn (go_joinable). @@ -136,9 +124,7 @@ namespace detail { template requires (std::invocable && is_task_v>) auto callable_wrapper(F f, Args... args) -> std::invoke_result_t { - auto task = std::invoke(std::move(f), std::move(args)...); - inherit_spawn_affinity(task); - co_return co_await task; + co_return co_await std::invoke(std::move(f), std::move(args)...); } } // namespace detail diff --git a/tests/unit/test_scheduler.cpp b/tests/unit/test_scheduler.cpp index e79e3dea..b2870b0c 100644 --- a/tests/unit/test_scheduler.cpp +++ b/tests/unit/test_scheduler.cpp @@ -2,9 +2,11 @@ #include #include #include +#include #include #include #include +#include #include #include #include @@ -165,53 +167,45 @@ enum class spawn_affinity_mode { move_to_worker_one, }; -class manual_schedule_gate { +class scoped_socket_pair { public: - class awaitable { - public: - explicit awaitable(manual_schedule_gate& gate) noexcept : gate_(gate) {} - - bool await_ready() const noexcept { return false; } - - bool await_suspend(std::coroutine_handle<> handle) noexcept { - gate_.handle_.store(handle.address(), std::memory_order_release); - gate_.published_.store(true, std::memory_order_release); - gate_.published_.notify_all(); - return true; - } - - void await_resume() const noexcept {} - - private: - manual_schedule_gate& gate_; - }; - - [[nodiscard]] awaitable wait() noexcept { return awaitable{*this}; } - - [[nodiscard]] bool published() const noexcept { - return published_.load(std::memory_order_acquire); + scoped_socket_pair() noexcept { + valid_ = ::socketpair(AF_UNIX, + SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, + 0, fds_) == 0; } - void resume() noexcept { - void* address = handle_.exchange(nullptr, std::memory_order_acq_rel); - if (address) { - schedule_handle(std::coroutine_handle<>::from_address(address)); + ~scoped_socket_pair() { + for (int fd : fds_) { + if (fd >= 0) { + ::close(fd); + } } } + scoped_socket_pair(const scoped_socket_pair&) = delete; + scoped_socket_pair& operator=(const scoped_socket_pair&) = delete; + + [[nodiscard]] bool valid() const noexcept { return valid_; } + [[nodiscard]] int reader() const noexcept { return fds_[0]; } + [[nodiscard]] int writer() const noexcept { return fds_[1]; } + private: - std::atomic handle_{nullptr}; - std::atomic published_{false}; + int fds_[2] = {-1, -1}; + bool valid_ = false; }; struct spawn_affinity_observation { std::shared_ptr context; std::atomic first_worker{NO_AFFINITY}; std::atomic resumed_worker{NO_AFFINITY}; + std::atomic after_scope_worker{NO_AFFINITY}; std::atomic completion_worker{NO_AFFINITY}; std::atomic first_affinity{NO_AFFINITY}; std::atomic resumed_affinity{NO_AFFINITY}; + std::atomic after_scope_affinity{NO_AFFINITY}; std::atomic post_io_affinity{NO_AFFINITY}; + std::atomic gate_received{false}; std::atomic io_staged{false}; std::atomic recv_result{-1}; std::atomic done{false}; @@ -220,46 +214,67 @@ struct spawn_affinity_observation { task observe_spawn_affinity( int fd, spawn_affinity_mode mode, - manual_schedule_gate* gate, + elio::sync::channel* gate, + cancel_token cleanup_token, spawn_affinity_observation* observed) { - auto* frame = promise_base::current_frame(); - if (!frame) { + auto* root_frame = promise_base::current_frame(); + if (!root_frame) { throw std::logic_error("spawn affinity test requires a current frame"); } - observed->context = frame->execution_context(); + observed->context = root_frame->execution_context(); observed->first_worker.store(current_worker_id(), std::memory_order_release); - observed->first_affinity.store(frame->affinity(), std::memory_order_release); - - if (mode == spawn_affinity_mode::clear) { - co_await clear_affinity(); - } else if (mode == spawn_affinity_mode::move_to_worker_one) { - co_await set_affinity(1, false); - } - - co_await gate->wait(); - observed->resumed_worker.store( + observed->first_affinity.store( + root_frame->affinity(), std::memory_order_release); + + co_await task_scope( + [mode, gate, observed](task_group&) -> task { + auto* body_frame = promise_base::current_frame(); + if (!body_frame) { + throw std::logic_error( + "spawn affinity scope requires a current frame"); + } + + if (mode == spawn_affinity_mode::clear) { + co_await clear_affinity(); + } else if (mode == spawn_affinity_mode::move_to_worker_one) { + co_await set_affinity(1, false); + } + + auto value = co_await gate->recv(); + observed->gate_received.store( + value.has_value() && *value == 1, + std::memory_order_release); + observed->resumed_worker.store( + current_worker_id(), std::memory_order_release); + observed->resumed_affinity.store( + body_frame->affinity(), std::memory_order_release); + }); + + observed->after_scope_worker.store( current_worker_id(), std::memory_order_release); - observed->resumed_affinity.store( - frame->affinity(), std::memory_order_release); + observed->after_scope_affinity.store( + root_frame->affinity(), std::memory_order_release); elio::io::detail::arm_next_cancellable_recv_staged_for_test( observed->io_staged); char byte = 0; auto result = co_await elio::io::async_recv( - fd, &byte, 1, 0, cancel_token{}); + fd, &byte, 1, 0, std::move(cleanup_token)); observed->recv_result.store(result.io.result, std::memory_order_release); observed->completion_worker.store( current_worker_id(), std::memory_order_release); observed->post_io_affinity.store( - frame->effective_affinity(), std::memory_order_release); + root_frame->effective_affinity(), std::memory_order_release); observed->done.store(true, std::memory_order_release); observed->done.notify_all(); } -task resume_spawn_affinity_probe(manual_schedule_gate* gate) { +task resume_spawn_affinity_probe( + elio::sync::channel* gate, + std::atomic* sent) { elio::runtime::detail::reject_next_schedule_for_test.store( true, std::memory_order_release); - gate->resume(); + sent->store(gate->try_send(1), std::memory_order_release); co_return; } @@ -276,82 +291,165 @@ void check_spawn_affinity_case(bool joinable, spawn_affinity_mode mode, size_t expected_worker, size_t expected_affinity) { - int sockets[2] = {-1, -1}; - REQUIRE(::socketpair(AF_UNIX, - SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, - 0, sockets) == 0); + scoped_socket_pair sockets; + REQUIRE(sockets.valid()); elio::runtime::detail::reject_next_schedule_for_test.store( false, std::memory_order_release); + elio::sync::channel gate(1); + spawn_affinity_observation observed; + cancel_source cleanup_cancellation; + std::optional> joined; scheduler sched(2); sched.start(); - manual_schedule_gate gate; - spawn_affinity_observation observed; - std::optional> joined; + const auto recv_waits_before = + elio::sync::detail::bounded_recv_waits_for_test.load( + std::memory_order_acquire); if (joinable) { joined.emplace(sched.go_joinable_to( - 0, observe_spawn_affinity, sockets[0], mode, &gate, &observed)); + 0, observe_spawn_affinity, sockets.reader(), mode, &gate, + cleanup_cancellation.get_token(), &observed)); } else { sched.go_to( - 0, observe_spawn_affinity, sockets[0], mode, &gate, &observed); + 0, observe_spawn_affinity, sockets.reader(), mode, &gate, + cleanup_cancellation.get_token(), &observed); } - REQUIRE(wait_for_scheduler_condition([&] { return gate.published(); })); - REQUIRE(observed.context); - CHECK(observed.first_worker.load(std::memory_order_acquire) == 0); - CHECK(observed.first_affinity.load(std::memory_order_acquire) == 0); + const bool receiver_waited = wait_for_scheduler_condition([&] { + return elio::sync::detail::bounded_recv_waits_for_test.load( + std::memory_order_acquire) > recv_waits_before; + }); + + bool marker_completed = false; + bool marker_succeeded = false; + if (receiver_waited) { + auto marker = sched.go_joinable_to(0, empty_task); + marker_completed = wait_for_scheduler_condition([&] { + return marker.is_destroyed(); + }); + if (marker_completed) { + try { + marker.await_resume(); + marker_succeeded = true; + } catch (...) { + } + } + } const auto local_fallbacks_before = elio::runtime::detail::local_schedule_fallbacks_for_test.load( std::memory_order_acquire); - auto notifier = sched.go_joinable_to( - 1, resume_spawn_affinity_probe, &gate); - notifier.wait_destroyed(); - notifier.await_resume(); + std::atomic gate_sent{false}; + bool notifier_completed = false; + bool notifier_succeeded = false; + if (marker_succeeded) { + auto notifier = sched.go_joinable_to( + 1, resume_spawn_affinity_probe, &gate, &gate_sent); + notifier_completed = wait_for_scheduler_condition([&] { + return notifier.is_destroyed(); + }); + if (notifier_completed) { + try { + notifier.await_resume(); + notifier_succeeded = true; + } catch (...) { + } + } + } + elio::runtime::detail::reject_next_schedule_for_test.store( + false, std::memory_order_release); - REQUIRE(wait_for_scheduler_condition([&] { + const bool io_staged = wait_for_scheduler_condition([&] { return observed.io_staged.load(std::memory_order_acquire); - })); - CHECK(observed.resumed_worker.load(std::memory_order_acquire) == - expected_worker); - CHECK(observed.resumed_affinity.load(std::memory_order_acquire) == - expected_affinity); + }); - auto* owner = sched.get_worker(expected_worker); - REQUIRE(owner != nullptr); - CHECK(observed.context->is_io_pin_owner( - expected_worker, owner->io_context().generation())); - CHECK(observed.context->effective_affinity() == expected_worker); + bool io_owner_matches = false; + bool effective_io_affinity_matches = false; + if (io_staged && observed.context) { + auto* owner = sched.get_worker(expected_worker); + if (owner) { + io_owner_matches = observed.context->is_io_pin_owner( + expected_worker, owner->io_context().generation()); + effective_io_affinity_matches = + observed.context->effective_affinity() == expected_worker; + } + } const auto local_fallbacks_after = elio::runtime::detail::local_schedule_fallbacks_for_test.load( std::memory_order_acquire); - if (mode == spawn_affinity_mode::inherit) { - CHECK(local_fallbacks_after == local_fallbacks_before); - } else { - CHECK(local_fallbacks_after == local_fallbacks_before + 1); - } const char byte = 'x'; - REQUIRE(::write(sockets[1], &byte, 1) == 1); - REQUIRE(wait_for_scheduler_condition([&] { + const auto write_result = ::write(sockets.writer(), &byte, 1); + bool task_completed = wait_for_scheduler_condition([&] { return observed.done.load(std::memory_order_acquire); - })); + }); + if (!task_completed) { + gate.close(); + try { + cleanup_cancellation.cancel(); + } catch (...) { + } + task_completed = wait_for_scheduler_condition([&] { + return observed.done.load(std::memory_order_acquire); + }); + } + bool joined_destroyed = !joined.has_value(); + bool joined_succeeded = !joined.has_value(); if (joined) { - joined->wait_destroyed(); - joined->await_resume(); + joined_destroyed = wait_for_scheduler_condition([&] { + return joined->is_destroyed(); + }); + if (joined_destroyed) { + try { + joined->await_resume(); + joined_succeeded = true; + } catch (...) { + } + } + } + + gate.close(); + const bool shutdown_succeeded = sched.shutdown(scaled_sec(5)); + + REQUIRE(receiver_waited); + REQUIRE(marker_completed); + REQUIRE(marker_succeeded); + REQUIRE(notifier_completed); + REQUIRE(notifier_succeeded); + REQUIRE(gate_sent.load(std::memory_order_acquire)); + REQUIRE(io_staged); + REQUIRE(observed.context); + CHECK(observed.first_worker.load(std::memory_order_acquire) == 0); + CHECK(observed.first_affinity.load(std::memory_order_acquire) == 0); + CHECK(observed.gate_received.load(std::memory_order_acquire)); + CHECK(observed.resumed_worker.load(std::memory_order_acquire) == + expected_worker); + CHECK(observed.resumed_affinity.load(std::memory_order_acquire) == + expected_affinity); + CHECK(observed.after_scope_worker.load(std::memory_order_acquire) == + expected_worker); + CHECK(observed.after_scope_affinity.load(std::memory_order_acquire) == + expected_affinity); + CHECK(io_owner_matches); + CHECK(effective_io_affinity_matches); + if (mode == spawn_affinity_mode::inherit) { + CHECK(local_fallbacks_after == local_fallbacks_before); + } else { + CHECK(local_fallbacks_after == local_fallbacks_before + 1); } + CHECK(write_result == 1); + REQUIRE(task_completed); + REQUIRE(joined_destroyed); + REQUIRE(joined_succeeded); CHECK(observed.recv_result.load(std::memory_order_acquire) == 1); CHECK(observed.completion_worker.load(std::memory_order_acquire) == expected_worker); CHECK(observed.post_io_affinity.load(std::memory_order_acquire) == expected_affinity); - REQUIRE(sched.shutdown(scaled_sec(5))); - - ::close(sockets[0]); - ::close(sockets[1]); + REQUIRE(shutdown_succeeded); } } // namespace