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_allocator.hpp>
15 : #include <boost/capy/ex/io_env.hpp>
16 : #include <boost/capy/ex/recycling_memory_resource.hpp>
17 : #include <boost/capy/ex/this_coro.hpp>
18 :
19 : #include <coroutine>
20 : #include <cstddef>
21 : #include <cstring>
22 : #include <memory_resource>
23 : #include <stop_token>
24 : #include <type_traits>
25 :
26 : namespace boost {
27 : namespace capy {
28 :
29 : /** CRTP mixin that adds I/O awaitable support to a promise type.
30 :
31 : Inherit from this class to enable these capabilities in your coroutine:
32 :
33 : 1. **Frame allocation** — The mixin provides `operator new/delete` that
34 : use the thread-local frame allocator set by `run_async`.
35 :
36 : 2. **Environment storage** — The mixin stores a pointer to the `io_env`
37 : containing the executor, stop token, and allocator for this coroutine.
38 :
39 : 3. **Environment access** — Coroutine code can retrieve the environment
40 : via `co_await this_coro::environment`, or individual fields via
41 : `co_await this_coro::executor`, `co_await this_coro::stop_token`,
42 : and `co_await this_coro::allocator`.
43 :
44 : @tparam Derived The derived promise type (CRTP pattern).
45 :
46 : @par Basic Usage
47 :
48 : For coroutines that need to access their execution environment:
49 :
50 : @code
51 : struct my_task
52 : {
53 : struct promise_type : io_awaitable_promise_base<promise_type>
54 : {
55 : my_task get_return_object();
56 : std::suspend_always initial_suspend() noexcept;
57 : std::suspend_always final_suspend() noexcept;
58 : void return_void();
59 : void unhandled_exception();
60 : };
61 :
62 : // ... awaitable interface ...
63 : };
64 :
65 : my_task example()
66 : {
67 : auto env = co_await this_coro::environment;
68 : // Access env->executor, env->stop_token, env->allocator
69 :
70 : // Or use fine-grained accessors:
71 : auto ex = co_await this_coro::executor;
72 : auto token = co_await this_coro::stop_token;
73 : auto* alloc = co_await this_coro::allocator;
74 : }
75 : @endcode
76 :
77 : @par Custom Awaitable Transformation
78 :
79 : If your promise needs to transform awaitables (e.g., for affinity or
80 : logging), override `transform_awaitable` instead of `await_transform`:
81 :
82 : @code
83 : struct promise_type : io_awaitable_promise_base<promise_type>
84 : {
85 : template<typename A>
86 : auto transform_awaitable(A&& a)
87 : {
88 : // Your custom transformation logic
89 : return std::forward<A>(a);
90 : }
91 : };
92 : @endcode
93 :
94 : The mixin's `await_transform` intercepts @ref this_coro::environment_tag
95 : and the fine-grained tag types (@ref this_coro::executor_tag,
96 : @ref this_coro::stop_token_tag, @ref this_coro::allocator_tag),
97 : then delegates all other awaitables to your `transform_awaitable`.
98 :
99 : @par Making Your Coroutine an IoAwaitable
100 :
101 : The mixin handles the "inside the coroutine" part—accessing the
102 : environment. To receive the environment when your coroutine is awaited
103 : (satisfying @ref IoAwaitable), implement the `await_suspend` overload
104 : on your coroutine return type:
105 :
106 : @code
107 : struct my_task
108 : {
109 : struct promise_type : io_awaitable_promise_base<promise_type> { ... };
110 :
111 : std::coroutine_handle<promise_type> h_;
112 :
113 : // IoAwaitable await_suspend receives and stores the environment
114 : std::coroutine_handle<> await_suspend(std::coroutine_handle<> cont, io_env const* env)
115 : {
116 : h_.promise().set_environment(env);
117 : // ... rest of suspend logic ...
118 : }
119 : };
120 : @endcode
121 :
122 : @par Thread Safety
123 : The environment is stored during `await_suspend` and read during
124 : `co_await this_coro::environment`. These occur on the same logical
125 : thread of execution, so no synchronization is required.
126 :
127 : @see this_coro::environment, this_coro::executor,
128 : this_coro::stop_token, this_coro::allocator
129 : @see io_env
130 : @see IoAwaitable
131 : */
132 : template<typename Derived>
133 : class io_awaitable_promise_base
134 : {
135 : io_env const* env_ = nullptr;
136 : mutable std::coroutine_handle<> cont_{std::noop_coroutine()};
137 :
138 : public:
139 : /** Allocate a coroutine frame.
140 :
141 : Uses the thread-local frame allocator set by run_async.
142 : Falls back to default memory resource if not set.
143 : Stores the allocator pointer at the end of each frame for
144 : correct deallocation even when TLS changes. Uses memcpy
145 : to avoid alignment requirements on the trailing pointer.
146 : Bypasses virtual dispatch for the recycling allocator.
147 : */
148 3741 : static void* operator new(std::size_t size)
149 : {
150 3741 : static auto* const rmr = get_recycling_memory_resource();
151 :
152 3741 : auto* mr = get_current_frame_allocator();
153 3741 : if(!mr)
154 1980 : mr = std::pmr::get_default_resource();
155 :
156 3741 : auto total = size + sizeof(std::pmr::memory_resource*);
157 : void* raw;
158 3741 : if(mr == rmr)
159 : raw = static_cast<recycling_memory_resource*>(mr)
160 1748 : ->allocate_fast(total, alignof(std::max_align_t));
161 : else
162 1993 : raw = mr->allocate(total, alignof(std::max_align_t));
163 3741 : std::memcpy(static_cast<char*>(raw) + size, &mr, sizeof(mr));
164 3741 : return raw;
165 : }
166 :
167 : /** Deallocate a coroutine frame.
168 :
169 : Reads the allocator pointer stored at the end of the frame
170 : to ensure correct deallocation regardless of current TLS.
171 : Bypasses virtual dispatch for the recycling allocator.
172 : */
173 3741 : static void operator delete(void* ptr, std::size_t size) noexcept
174 : {
175 3741 : static auto* const rmr = get_recycling_memory_resource();
176 :
177 : std::pmr::memory_resource* mr;
178 3741 : std::memcpy(&mr, static_cast<char*>(ptr) + size, sizeof(mr));
179 3741 : auto total = size + sizeof(std::pmr::memory_resource*);
180 3741 : if(mr == rmr)
181 : static_cast<recycling_memory_resource*>(mr)
182 1748 : ->deallocate_fast(ptr, total, alignof(std::max_align_t));
183 : else
184 1993 : mr->deallocate(ptr, total, alignof(std::max_align_t));
185 3741 : }
186 :
187 3741 : ~io_awaitable_promise_base()
188 : {
189 : // Abnormal teardown: destroy orphaned continuation
190 3741 : if(cont_ != std::noop_coroutine())
191 1 : cont_.destroy();
192 3741 : }
193 :
194 : //----------------------------------------------------------
195 : // Continuation support
196 : //----------------------------------------------------------
197 :
198 : /** Store the continuation to resume on completion.
199 :
200 : Call this from your coroutine type's `await_suspend` overload
201 : to set up the completion path. The `final_suspend` awaiter
202 : returns this handle via unconditional symmetric transfer.
203 :
204 : @param cont The continuation to resume on completion.
205 : */
206 3660 : void set_continuation(std::coroutine_handle<> cont) noexcept
207 : {
208 3660 : cont_ = cont;
209 3660 : }
210 :
211 : /** Return and consume the stored continuation handle.
212 :
213 : Resets the stored handle to `noop_coroutine()` so the
214 : destructor will not double-destroy it.
215 :
216 : @return The continuation for symmetric transfer.
217 : */
218 3717 : std::coroutine_handle<> continuation() const noexcept
219 : {
220 3717 : return std::exchange(cont_, std::noop_coroutine());
221 : }
222 :
223 : //----------------------------------------------------------
224 : // Environment support
225 : //----------------------------------------------------------
226 :
227 : /** Store a pointer to the execution environment.
228 :
229 : Call this from your coroutine type's `await_suspend`
230 : overload to make the environment available via
231 : `co_await this_coro::environment`. The pointed-to
232 : `io_env` must outlive this coroutine.
233 :
234 : @param env The environment to store.
235 : */
236 3738 : void set_environment(io_env const* env) noexcept
237 : {
238 3738 : env_ = env;
239 3738 : }
240 :
241 : /** Return the stored execution environment.
242 :
243 : @return The environment.
244 : */
245 13096 : io_env const* environment() const noexcept
246 : {
247 13096 : BOOST_CAPY_ASSERT(env_);
248 13096 : return env_;
249 : }
250 :
251 : /** Transform an awaitable before co_await.
252 :
253 : Override this in your derived promise type to customize how
254 : awaitables are transformed. The default implementation passes
255 : the awaitable through unchanged.
256 :
257 : @param a The awaitable expression from `co_await a`.
258 :
259 : @return The transformed awaitable.
260 : */
261 : template<typename A>
262 : decltype(auto) transform_awaitable(A&& a)
263 : {
264 : return std::forward<A>(a);
265 : }
266 :
267 : /** Intercept co_await expressions.
268 :
269 : This function handles @ref this_coro::environment_tag and
270 : the fine-grained tags (@ref this_coro::executor_tag,
271 : @ref this_coro::stop_token_tag, @ref this_coro::allocator_tag)
272 : specially, returning an awaiter that yields the stored value.
273 : All other awaitables are delegated to @ref transform_awaitable.
274 :
275 : @param t The awaited expression.
276 :
277 : @return An awaiter for the expression.
278 : */
279 : template<typename T>
280 7326 : auto await_transform(T&& t)
281 : {
282 : using Tag = std::decay_t<T>;
283 :
284 : if constexpr (std::is_same_v<Tag, this_coro::environment_tag>)
285 : {
286 37 : BOOST_CAPY_ASSERT(env_);
287 : struct awaiter
288 : {
289 : io_env const* env_;
290 35 : bool await_ready() const noexcept { return true; }
291 2 : void await_suspend(std::coroutine_handle<>) const noexcept { }
292 34 : io_env const* await_resume() const noexcept { return env_; }
293 : };
294 37 : return awaiter{env_};
295 : }
296 : else if constexpr (std::is_same_v<Tag, this_coro::executor_tag>)
297 : {
298 3 : BOOST_CAPY_ASSERT(env_);
299 : struct awaiter
300 : {
301 : executor_ref executor_;
302 2 : bool await_ready() const noexcept { return true; }
303 : void await_suspend(std::coroutine_handle<>) const noexcept { }
304 2 : executor_ref await_resume() const noexcept { return executor_; }
305 : };
306 3 : return awaiter{env_->executor};
307 : }
308 : else if constexpr (std::is_same_v<Tag, this_coro::stop_token_tag>)
309 : {
310 7 : BOOST_CAPY_ASSERT(env_);
311 : struct awaiter
312 : {
313 : std::stop_token token_;
314 6 : bool await_ready() const noexcept { return true; }
315 0 : void await_suspend(std::coroutine_handle<>) const noexcept { }
316 6 : std::stop_token await_resume() const noexcept { return token_; }
317 : };
318 7 : return awaiter{env_->stop_token};
319 : }
320 : else if constexpr (std::is_same_v<Tag, this_coro::allocator_tag>)
321 : {
322 8 : BOOST_CAPY_ASSERT(env_);
323 : struct awaiter
324 : {
325 : std::pmr::memory_resource* allocator_;
326 6 : bool await_ready() const noexcept { return true; }
327 0 : void await_suspend(std::coroutine_handle<>) const noexcept { }
328 7 : std::pmr::memory_resource* await_resume() const noexcept { return allocator_; }
329 : };
330 8 : return awaiter{env_->allocator};
331 : }
332 : else
333 : {
334 5484 : return static_cast<Derived*>(this)->transform_awaitable(
335 7271 : std::forward<T>(t));
336 : }
337 : }
338 : };
339 :
340 : } // namespace capy
341 : } // namespace boost
342 :
343 : #endif
|
