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