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_SUPPORT_HPP
11 : #define BOOST_CAPY_EX_IO_AWAITABLE_SUPPORT_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_support<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_support<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_support<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_support
134 : {
135 : io_env const* env_ = &detail::empty_io_env;
136 : mutable std::coroutine_handle<> cont_{std::noop_coroutine()};
137 :
138 : public:
139 : //----------------------------------------------------------
140 : // Frame allocation support
141 : //----------------------------------------------------------
142 :
143 : public:
144 : /** Allocate a coroutine frame.
145 :
146 : Uses the thread-local frame allocator set by run_async.
147 : Falls back to default memory resource if not set.
148 : Stores the allocator pointer at the end of each frame for
149 : correct deallocation even when TLS changes. Uses memcpy
150 : to avoid alignment requirements on the trailing pointer.
151 : Bypasses virtual dispatch for the recycling allocator.
152 : */
153 : static void*
154 3741 : operator new(std::size_t size)
155 : {
156 3741 : static auto* const rmr = get_recycling_memory_resource();
157 :
158 3741 : auto* mr = current_frame_allocator();
159 3741 : if(!mr)
160 1994 : mr = std::pmr::get_default_resource();
161 :
162 3741 : auto total = size + sizeof(std::pmr::memory_resource*);
163 : void* raw;
164 3741 : if(mr == rmr)
165 : raw = static_cast<recycling_memory_resource*>(mr)
166 1734 : ->allocate_fast(total, alignof(std::max_align_t));
167 : else
168 2007 : raw = mr->allocate(total, alignof(std::max_align_t));
169 3741 : std::memcpy(static_cast<char*>(raw) + size, &mr, sizeof(mr));
170 3741 : return raw;
171 : }
172 :
173 : /** Deallocate a coroutine frame.
174 :
175 : Reads the allocator pointer stored at the end of the frame
176 : to ensure correct deallocation regardless of current TLS.
177 : Bypasses virtual dispatch for the recycling allocator.
178 : */
179 : static void
180 3741 : operator delete(void* ptr, std::size_t size)
181 : {
182 3741 : static auto* const rmr = get_recycling_memory_resource();
183 :
184 : std::pmr::memory_resource* mr;
185 3741 : std::memcpy(&mr, static_cast<char*>(ptr) + size, sizeof(mr));
186 3741 : auto total = size + sizeof(std::pmr::memory_resource*);
187 3741 : if(mr == rmr)
188 : static_cast<recycling_memory_resource*>(mr)
189 1734 : ->deallocate_fast(ptr, total, alignof(std::max_align_t));
190 : else
191 2007 : mr->deallocate(ptr, total, alignof(std::max_align_t));
192 3741 : }
193 :
194 3741 : ~io_awaitable_support()
195 : {
196 : // Abnormal teardown: destroy orphaned continuation
197 3741 : if(cont_ != std::noop_coroutine())
198 1 : cont_.destroy();
199 3741 : }
200 :
201 : //----------------------------------------------------------
202 : // Continuation support
203 : //----------------------------------------------------------
204 :
205 : /** Store the continuation to resume on completion.
206 :
207 : Call this from your coroutine type's `await_suspend` overload
208 : to set up the completion path. The `final_suspend` awaiter
209 : returns this handle via unconditional symmetric transfer.
210 :
211 : @param cont The continuation to resume on completion.
212 : */
213 3635 : void set_continuation(std::coroutine_handle<> cont) noexcept
214 : {
215 3635 : cont_ = cont;
216 3635 : }
217 :
218 : /** Return and consume the stored continuation handle.
219 :
220 : Resets the stored handle to `noop_coroutine()` so the
221 : destructor will not double-destroy it.
222 :
223 : @return The continuation for symmetric transfer.
224 : */
225 3717 : std::coroutine_handle<> continuation() const noexcept
226 : {
227 3717 : return std::exchange(cont_, std::noop_coroutine());
228 : }
229 :
230 : //----------------------------------------------------------
231 : // Environment support
232 : //----------------------------------------------------------
233 :
234 : /** Store a pointer to the execution environment.
235 :
236 : Call this from your coroutine type's `await_suspend`
237 : overload to make the environment available via
238 : `co_await this_coro::environment`. The pointed-to
239 : `io_env` must outlive this coroutine.
240 :
241 : @param env The environment to store.
242 : */
243 3702 : void set_environment(io_env const* env) noexcept
244 : {
245 3702 : env_ = env;
246 3702 : }
247 :
248 : /** Return the stored execution environment.
249 :
250 : @return The environment.
251 : */
252 13096 : io_env const* environment() const noexcept
253 : {
254 13096 : return env_;
255 : }
256 :
257 : /** Transform an awaitable before co_await.
258 :
259 : Override this in your derived promise type to customize how
260 : awaitables are transformed. The default implementation passes
261 : the awaitable through unchanged.
262 :
263 : @param a The awaitable expression from `co_await a`.
264 :
265 : @return The transformed awaitable.
266 : */
267 : template<typename A>
268 : decltype(auto) transform_awaitable(A&& a)
269 : {
270 : return std::forward<A>(a);
271 : }
272 :
273 : /** Intercept co_await expressions.
274 :
275 : This function handles @ref this_coro::environment_tag and
276 : the fine-grained tags (@ref this_coro::executor_tag,
277 : @ref this_coro::stop_token_tag, @ref this_coro::allocator_tag)
278 : specially, returning an awaiter that yields the stored value.
279 : All other awaitables are delegated to @ref transform_awaitable.
280 :
281 : @param t The awaited expression.
282 :
283 : @return An awaiter for the expression.
284 : */
285 : template<typename T>
286 7326 : auto await_transform(T&& t)
287 : {
288 : using Tag = std::decay_t<T>;
289 :
290 : if constexpr (std::is_same_v<Tag, this_coro::environment_tag>)
291 : {
292 : struct awaiter
293 : {
294 : io_env const* env_;
295 35 : bool await_ready() const noexcept { return true; }
296 2 : void await_suspend(std::coroutine_handle<>) const noexcept { }
297 34 : io_env const* await_resume() const noexcept { return env_; }
298 : };
299 37 : return awaiter{env_};
300 : }
301 : else if constexpr (std::is_same_v<Tag, this_coro::executor_tag>)
302 : {
303 : struct awaiter
304 : {
305 : executor_ref executor_;
306 2 : bool await_ready() const noexcept { return true; }
307 : void await_suspend(std::coroutine_handle<>) const noexcept { }
308 2 : executor_ref await_resume() const noexcept { return executor_; }
309 : };
310 3 : return awaiter{env_->executor};
311 : }
312 : else if constexpr (std::is_same_v<Tag, this_coro::stop_token_tag>)
313 : {
314 : struct awaiter
315 : {
316 : std::stop_token token_;
317 6 : bool await_ready() const noexcept { return true; }
318 0 : void await_suspend(std::coroutine_handle<>) const noexcept { }
319 6 : std::stop_token await_resume() const noexcept { return token_; }
320 : };
321 7 : return awaiter{env_->stop_token};
322 : }
323 : else if constexpr (std::is_same_v<Tag, this_coro::allocator_tag>)
324 : {
325 : struct awaiter
326 : {
327 : std::pmr::memory_resource* allocator_;
328 6 : bool await_ready() const noexcept { return true; }
329 0 : void await_suspend(std::coroutine_handle<>) const noexcept { }
330 7 : std::pmr::memory_resource* await_resume() const noexcept { return allocator_; }
331 : };
332 8 : return awaiter{env_->allocator};
333 : }
334 : else
335 : {
336 5484 : return static_cast<Derived*>(this)->transform_awaitable(
337 7271 : std::forward<T>(t));
338 : }
339 : }
340 : };
341 :
342 : } // namespace capy
343 : } // namespace boost
344 :
345 : #endif
|
