From a68675ec4756a409c1b522dd0ea7f136548213da Mon Sep 17 00:00:00 2001 From: Eric Niebler Date: Sat, 11 Apr 2026 10:35:07 -0700 Subject: [PATCH] fix how `stdexec::task` handles environments and queries a number of bugs and omissions meant that queries from a user-specified task environment were not getting propagated. fix the problem and add a test that pulls a query out of a custom task environment. fixes #1965 --- include/stdexec/__detail/__task.hpp | 59 +++++++++++++++++++---------- test/stdexec/types/test_task.cpp | 48 +++++++++++++++++++++-- 2 files changed, 83 insertions(+), 24 deletions(-) diff --git a/include/stdexec/__detail/__task.hpp b/include/stdexec/__detail/__task.hpp index 57c815558..a200cd999 100644 --- a/include/stdexec/__detail/__task.hpp +++ b/include/stdexec/__detail/__task.hpp @@ -394,10 +394,13 @@ namespace STDEXEC struct __opstate_base : private allocator_type { - template - constexpr explicit __opstate_base(task&& __task, _Env const & __env) noexcept + template + constexpr explicit __opstate_base(task&& __task, + _Env const & __env, + _OwnEnv const & __own_env) noexcept : allocator_type(__mk_alloc(__env)) , __sch_(__mk_sched(__env, __get_allocator())) + , __env_(__mk_env(__env, __own_env)) , __task_(static_cast(__task)) { auto& __promise = __task_.__coro_.promise(); @@ -434,19 +437,26 @@ namespace STDEXEC } start_scheduler_type __sch_; + _TaskEnv __env_; task __task_; __error_variant_t __errors_{__no_init}; }; + template + struct __own_env_box + { + __own_env_t<_Env> __own_env_; + }; + template struct STDEXEC_ATTRIBUTE(empty_bases) __awaiter final - : __opstate_base + : __own_env_box> + , __opstate_base , __stop_callback_box_t> { constexpr explicit __awaiter(task&& __task, _ParentPromise& __parent) noexcept - : __opstate_base(static_cast(__task), STDEXEC::get_env(__parent)) - , __own_env_(__mk_own_env(STDEXEC::get_env(__parent))) - , __env_(__mk_env(STDEXEC::get_env(__parent), __own_env_)) + : __awaiter::__own_env_box{__mk_own_env(STDEXEC::get_env(__parent))} + , __opstate_base(static_cast(__task), STDEXEC::get_env(__parent), this->__own_env_) , __parent_(__parent) {} @@ -501,10 +511,8 @@ namespace STDEXEC return __parent_.unhandled_stopped(); } - __own_env_t<_ParentPromise> __own_env_; - _TaskEnv __env_; - __std::coroutine_handle<> __continuation_; - _ParentPromise& __parent_; + __std::coroutine_handle<> __continuation_; + _ParentPromise& __parent_; }; struct __attrs @@ -542,21 +550,23 @@ namespace STDEXEC //////////////////////////////////////////////////////////////////////////////////////// // task::__opstate - template + template template - struct STDEXEC_ATTRIBUTE(empty_bases) task<_Ty, _Env>::__opstate final + struct STDEXEC_ATTRIBUTE(empty_bases) task<_Ty, _TaskEnv>::__opstate final : __rcvr_box<_Rcvr> // holds the receiver so that we can pass __opstate_base a reference to it - , __opstate_base + , __own_env_box> , __stop_callback_box_t> + , __opstate_base { public: using operation_state_concept = operation_state_tag; explicit __opstate(task&& __task, _Rcvr&& __rcvr) noexcept : __rcvr_box<_Rcvr>{static_cast<_Rcvr&&>(__rcvr)} - , __opstate_base(static_cast(__task), STDEXEC::get_env(this->__rcvr_)) - , __own_env_(__mk_own_env(STDEXEC::get_env(this->__rcvr_))) - , __env_(__mk_env(STDEXEC::get_env(this->__rcvr_), __own_env_)) + , __opstate::__own_env_box{__mk_own_env(STDEXEC::get_env(this->__rcvr_))} + , __opstate_base(static_cast(__task), + STDEXEC::get_env(this->__rcvr_), + this->__own_env_) {} void start() & noexcept @@ -634,15 +644,12 @@ namespace STDEXEC STDEXEC::set_stopped(static_cast<_Rcvr&&>(this->__rcvr_)); return std::noop_coroutine(); } - - __own_env_t<_Rcvr> __own_env_; - _Env __env_; }; //////////////////////////////////////////////////////////////////////////////////////// // task::promise_type - template - struct task<_Ty, _Env>::__promise : __task::__promise_base<_Ty> + template + struct task<_Ty, _TaskEnv>::__promise : __task::__promise_base<_Ty> { __promise() noexcept = default; @@ -814,6 +821,16 @@ namespace STDEXEC } } + template <__forwarding_query _Query, class... _Args> + requires __queryable_with<_TaskEnv, _Query, _Args...> + [[nodiscard]] + constexpr auto query(_Query __tag, _Args&&... __args) const + noexcept(__nothrow_queryable_with<_TaskEnv, _Query, _Args...>) + -> __query_result_t<_TaskEnv, _Query, _Args...> + { + return __query<_Query>()(__promise_->__state_->__env_, static_cast<_Args&&>(__args)...); + } + __promise const * __promise_; }; diff --git a/test/stdexec/types/test_task.cpp b/test/stdexec/types/test_task.cpp index 49bb83b00..481d2f5e3 100644 --- a/test/stdexec/types/test_task.cpp +++ b/test/stdexec/types/test_task.cpp @@ -79,7 +79,7 @@ namespace CHECK(i == 42); } - auto test_task_int_ref(int& i) -> ex::task + auto test_task_int_ref(int &i) -> ex::task { CHECK(get_id() == 0); co_await ex::schedule(ex::inline_scheduler{}); @@ -90,7 +90,7 @@ namespace TEST_CASE("test task", "[types][task]") { int value = 42; - auto t = test_task_int_ref(value) | ex::then([](int& i) { return std::ref(i); }); + auto t = test_task_int_ref(value) | ex::then([](int &i) { return std::ref(i); }); auto [i] = ex::sync_wait(std::move(t)).value(); STATIC_REQUIRE(std::same_as>); CHECK(&i.get() == &value); @@ -207,7 +207,7 @@ namespace template _Env> requires ex::__callable - explicit test_env2(_Env const & other) noexcept + explicit test_env2(_Env const &other) noexcept : sch(ex::get_scheduler(other)) {} @@ -278,6 +278,48 @@ namespace CHECK(i == 84'000'042); } + struct my_env + { + template + using env_type = my_env; + + template + requires std::invocable + && std::same_as, + ex::run_loop::scheduler> + explicit my_env(Env const &env) noexcept + : delegation_scheduler_(ex::get_delegation_scheduler(env)) + {} + + [[nodiscard]] + auto query(ex::get_delegation_scheduler_t) const noexcept + { + return delegation_scheduler_; + } + + ex::run_loop::scheduler delegation_scheduler_; + }; + + auto + test_task_provides_additional_queries_with_a_custom_env(ex::run_loop::scheduler sync_wt_dlgtn_sch) + -> ex::task + { + // Fetch sync_wait's run_loop scheduler from the environment. + ex::run_loop::scheduler tsk_dlgtn_sch = co_await ex::read_env(ex::get_delegation_scheduler); + CHECK(tsk_dlgtn_sch == sync_wt_dlgtn_sch); + co_return 13; + } + + TEST_CASE("task can provide additional queries through a custom environment", "[types][task]") + { + ex::sync_wait(ex::let_value(ex::read_env(ex::get_delegation_scheduler), + [](ex::run_loop::scheduler sync_wt_dlgtn_sch) + { + return test_task_provides_additional_queries_with_a_custom_env( + sync_wt_dlgtn_sch); + })); + } + // FUTURE TODO: add support so that `co_await sndr` can return a reference. // auto test_task_awaits_just_ref_sender() -> ex::task { // int value = 42;