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/tests/unit/test_scheduler.cpp b/tests/unit/test_scheduler.cpp index c91760c8..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 @@ -12,6 +14,7 @@ #include #include #include +#include #include #include #include "../test_main.cpp" // For scaled timeouts @@ -158,6 +161,297 @@ task cancellable_recv_on_worker( completed->store(true, std::memory_order_release); } +enum class spawn_affinity_mode { + inherit, + clear, + move_to_worker_one, +}; + +class scoped_socket_pair { +public: + scoped_socket_pair() noexcept { + valid_ = ::socketpair(AF_UNIX, + SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, + 0, fds_) == 0; + } + + ~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: + 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}; +}; + +task observe_spawn_affinity( + int fd, + spawn_affinity_mode mode, + elio::sync::channel* gate, + cancel_token cleanup_token, + spawn_affinity_observation* observed) { + auto* root_frame = promise_base::current_frame(); + if (!root_frame) { + throw std::logic_error("spawn affinity test requires a current frame"); + } + observed->context = root_frame->execution_context(); + observed->first_worker.store(current_worker_id(), std::memory_order_release); + 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->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, 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( + 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( + elio::sync::channel* gate, + std::atomic* sent) { + elio::runtime::detail::reject_next_schedule_for_test.store( + true, std::memory_order_release); + sent->store(gate->try_send(1), std::memory_order_release); + co_return; +} + +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) { + 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(); + + 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.reader(), mode, &gate, + cleanup_cancellation.get_token(), &observed)); + } else { + sched.go_to( + 0, observe_spawn_affinity, sockets.reader(), mode, &gate, + cleanup_cancellation.get_token(), &observed); + } + + 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); + 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); + + const bool io_staged = wait_for_scheduler_condition([&] { + return observed.io_staged.load(std::memory_order_acquire); + }); + + 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); + + const char byte = 'x'; + 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_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(shutdown_succeeded); +} + } // namespace TEST_CASE("Scheduler construction", "[scheduler]") { @@ -240,6 +534,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();