TLA Line data Source code
1 : //
2 : // Copyright (c) 2025 Vinnie Falco (vinnie.falco@gmail.com)
3 : //
4 : // Distributed under the Boost Software License, Version 1.0. (See accompanying
5 : // file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
6 : //
7 : // Official repository: https://github.com/cppalliance/capy
8 : //
9 :
10 : #ifndef BOOST_CAPY_EX_IO_AWAITABLE_PROMISE_BASE_HPP
11 : #define BOOST_CAPY_EX_IO_AWAITABLE_PROMISE_BASE_HPP
12 :
13 : #include <boost/capy/detail/config.hpp>
14 : #include <boost/capy/ex/frame_alloc_mixin.hpp>
15 : #include <boost/capy/ex/frame_allocator.hpp>
16 : #include <boost/capy/ex/io_env.hpp>
17 : #include <boost/capy/ex/this_coro.hpp>
18 :
19 : #include <coroutine>
20 : #include <memory_resource>
21 : #include <stop_token>
22 : #include <type_traits>
23 :
24 : namespace boost {
25 : namespace capy {
26 :
27 : /** CRTP mixin that adds I/O awaitable support to a promise type.
28 :
29 : Inherit from this class to enable these capabilities in your coroutine:
30 :
31 : 1. **Frame allocation** — The mixin provides `operator new/delete` that
32 : use the thread-local frame allocator set by `run_async`.
33 :
34 : 2. **Environment storage** — The mixin stores a pointer to the `io_env`
35 : containing the executor, stop token, and allocator for this coroutine.
36 :
37 : 3. **Environment access** — Coroutine code can retrieve the environment
38 : via `co_await this_coro::environment`, or individual fields via
39 : `co_await this_coro::executor`, `co_await this_coro::stop_token`,
40 : and `co_await this_coro::frame_allocator`.
41 :
42 : @tparam Derived The derived promise type (CRTP pattern).
43 :
44 : @par Basic Usage
45 :
46 : For coroutines that need to access their execution environment:
47 :
48 : @code
49 : struct my_task
50 : {
51 : struct promise_type : io_awaitable_promise_base<promise_type>
52 : {
53 : my_task get_return_object();
54 : std::suspend_always initial_suspend() noexcept;
55 : std::suspend_always final_suspend() noexcept;
56 : void return_void();
57 : void unhandled_exception();
58 : };
59 :
60 : // ... awaitable interface ...
61 : };
62 :
63 : my_task example()
64 : {
65 : auto env = co_await this_coro::environment;
66 : // Access env->executor, env->stop_token, env->frame_allocator
67 :
68 : // Or use fine-grained accessors:
69 : auto ex = co_await this_coro::executor;
70 : auto token = co_await this_coro::stop_token;
71 : auto* alloc = co_await this_coro::frame_allocator;
72 : }
73 : @endcode
74 :
75 : @par Custom Awaitable Transformation
76 :
77 : If your promise needs to transform awaitables (e.g., for affinity or
78 : logging), override `transform_awaitable` instead of `await_transform`:
79 :
80 : @code
81 : struct promise_type : io_awaitable_promise_base<promise_type>
82 : {
83 : template<typename A>
84 : auto transform_awaitable(A&& a)
85 : {
86 : // Your custom transformation logic
87 : return std::forward<A>(a);
88 : }
89 : };
90 : @endcode
91 :
92 : The mixin's `await_transform` intercepts @ref this_coro::environment_tag
93 : and the fine-grained tag types (@ref this_coro::executor_tag,
94 : @ref this_coro::stop_token_tag, @ref this_coro::frame_allocator_tag),
95 : then delegates all other awaitables to your `transform_awaitable`.
96 :
97 : @par Making Your Coroutine an IoAwaitable
98 :
99 : The mixin handles the "inside the coroutine" part—accessing the
100 : environment. To receive the environment when your coroutine is awaited
101 : (satisfying @ref IoAwaitable), implement the `await_suspend` overload
102 : on your coroutine return type:
103 :
104 : @code
105 : struct my_task
106 : {
107 : struct promise_type : io_awaitable_promise_base<promise_type> { ... };
108 :
109 : std::coroutine_handle<promise_type> h_;
110 :
111 : // IoAwaitable await_suspend receives and stores the environment
112 : std::coroutine_handle<> await_suspend(std::coroutine_handle<> cont, io_env const* env)
113 : {
114 : h_.promise().set_environment(env);
115 : // ... rest of suspend logic ...
116 : }
117 : };
118 : @endcode
119 :
120 : @par Thread Safety
121 : The environment is stored during `await_suspend` and read during
122 : `co_await this_coro::environment`. These occur on the same logical
123 : thread of execution, so no synchronization is required.
124 :
125 : @see this_coro::environment, this_coro::executor,
126 : this_coro::stop_token, this_coro::frame_allocator
127 : @see io_env
128 : @see IoAwaitable
129 : */
130 : template<typename Derived>
131 : class io_awaitable_promise_base
132 : : public frame_alloc_mixin
133 : {
134 : io_env const* env_ = nullptr;
135 : mutable std::coroutine_handle<> cont_{std::noop_coroutine()};
136 :
137 : public:
138 HIT 5065 : ~io_awaitable_promise_base()
139 : {
140 : // Abnormal teardown: destroy an orphaned continuation, e.g.
141 : // a run_async trampoline when the task is destroyed before
142 : // reaching final_suspend. Callers must not destroy a task
143 : // via handle().destroy() while it is being awaited by a
144 : // parent coroutine: that puts cont_ under another owner
145 : // and would produce a double-destroy from this branch. See
146 : // task::handle() / quitter::handle() for the contract.
147 5065 : if(cont_ != std::noop_coroutine())
148 145 : cont_.destroy();
149 5065 : }
150 :
151 : //----------------------------------------------------------
152 : // Continuation support
153 : //----------------------------------------------------------
154 :
155 : /** Store the continuation to resume on completion.
156 :
157 : Call this from your coroutine type's `await_suspend` overload
158 : to set up the completion path. The `final_suspend` awaiter
159 : returns this handle via unconditional symmetric transfer.
160 :
161 : @param cont The continuation to resume on completion.
162 : */
163 4982 : void set_continuation(std::coroutine_handle<> cont) noexcept
164 : {
165 4982 : cont_ = cont;
166 4982 : }
167 :
168 : /** Return and consume the stored continuation handle.
169 :
170 : Resets the stored handle to `noop_coroutine()` so the
171 : destructor will not double-destroy it.
172 :
173 : @return The continuation for symmetric transfer.
174 : */
175 4897 : std::coroutine_handle<> continuation() const noexcept
176 : {
177 4897 : return std::exchange(cont_, std::noop_coroutine());
178 : }
179 :
180 : //----------------------------------------------------------
181 : // Environment support
182 : //----------------------------------------------------------
183 :
184 : /** Store a pointer to the execution environment.
185 :
186 : Call this from your coroutine type's `await_suspend`
187 : overload to make the environment available via
188 : `co_await this_coro::environment`. The pointed-to
189 : `io_env` must outlive this coroutine.
190 :
191 : @param env The environment to store.
192 : */
193 5062 : void set_environment(io_env const* env) noexcept
194 : {
195 5062 : env_ = env;
196 5062 : }
197 :
198 : /** Return the stored execution environment.
199 :
200 : @return The environment.
201 : */
202 16622 : io_env const* environment() const noexcept
203 : {
204 16622 : BOOST_CAPY_ASSERT(env_);
205 16622 : return env_;
206 : }
207 :
208 : /** Transform an awaitable before co_await.
209 :
210 : Override this in your derived promise type to customize how
211 : awaitables are transformed. The default implementation passes
212 : the awaitable through unchanged.
213 :
214 : @param a The awaitable expression from `co_await a`.
215 :
216 : @return The transformed awaitable.
217 : */
218 : template<typename A>
219 : decltype(auto) transform_awaitable(A&& a)
220 : {
221 : return std::forward<A>(a);
222 : }
223 :
224 : /** Intercept co_await expressions.
225 :
226 : This function handles @ref this_coro::environment_tag and
227 : the fine-grained tags (@ref this_coro::executor_tag,
228 : @ref this_coro::stop_token_tag, @ref this_coro::frame_allocator_tag)
229 : specially, returning an awaiter that yields the stored value.
230 : All other awaitables are delegated to @ref transform_awaitable.
231 :
232 : @param t The awaited expression.
233 :
234 : @return An awaiter for the expression.
235 : */
236 : template<typename T>
237 9225 : auto await_transform(T&& t)
238 : {
239 : using Tag = std::decay_t<T>;
240 :
241 : if constexpr (std::is_same_v<Tag, this_coro::environment_tag>)
242 : {
243 18 : BOOST_CAPY_ASSERT(env_);
244 : struct awaiter
245 : {
246 : io_env const* env_;
247 16 : bool await_ready() const noexcept { return true; }
248 2 : void await_suspend(std::coroutine_handle<>) const noexcept { }
249 15 : io_env const* await_resume() const noexcept { return env_; }
250 : };
251 18 : return awaiter{env_};
252 : }
253 : else if constexpr (std::is_same_v<Tag, this_coro::executor_tag>)
254 : {
255 4 : BOOST_CAPY_ASSERT(env_);
256 : struct awaiter
257 : {
258 : executor_ref executor_;
259 3 : bool await_ready() const noexcept { return true; }
260 MIS 0 : void await_suspend(std::coroutine_handle<>) const noexcept { }
261 HIT 3 : executor_ref await_resume() const noexcept { return executor_; }
262 : };
263 4 : return awaiter{env_->executor};
264 : }
265 : else if constexpr (std::is_same_v<Tag, this_coro::stop_token_tag>)
266 : {
267 15 : BOOST_CAPY_ASSERT(env_);
268 : struct awaiter
269 : {
270 : std::stop_token token_;
271 14 : bool await_ready() const noexcept { return true; }
272 MIS 0 : void await_suspend(std::coroutine_handle<>) const noexcept { }
273 HIT 14 : std::stop_token await_resume() const noexcept { return token_; }
274 : };
275 15 : return awaiter{env_->stop_token};
276 : }
277 : else if constexpr (std::is_same_v<Tag, this_coro::frame_allocator_tag>)
278 : {
279 8 : BOOST_CAPY_ASSERT(env_);
280 : struct awaiter
281 : {
282 : std::pmr::memory_resource* frame_allocator_;
283 6 : bool await_ready() const noexcept { return true; }
284 MIS 0 : void await_suspend(std::coroutine_handle<>) const noexcept { }
285 HIT 7 : std::pmr::memory_resource* await_resume() const noexcept { return frame_allocator_; }
286 : };
287 8 : return awaiter{env_->frame_allocator};
288 : }
289 : else
290 : {
291 7000 : return static_cast<Derived*>(this)->transform_awaitable(
292 9180 : std::forward<T>(t));
293 : }
294 : }
295 : };
296 :
297 : } // namespace capy
298 : } // namespace boost
299 :
300 : #endif
|
