Root/libs / capy / include / boost / capy / ex / io_awaitable_support.hpp
Prev Index Next

libs/capy/include/boost/capy/ex/io_awaitable_support.hpp

97.0% Lines (64/66) 99.2% Functions (263/265) 87.5% Branches (7/8)
libs/capy/include/boost/capy/ex/io_awaitable_support.hpp
Line Branch Hits Source Code
1 //
2 // Copyright (c) 2025 Vinnie Falco (vinnie dot falco at gmail dot 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/coro.hpp>
15 #include <boost/capy/ex/executor_ref.hpp>
16 #include <boost/capy/ex/frame_allocator.hpp>
17 #include <boost/capy/ex/this_coro.hpp>
18
19 #include <coroutine>
20 #include <cstddef>
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. **Frame allocator storage** — The mixin stores the allocator pointer
36 for propagation to child tasks.
37
38 3. **Stop token storage** — The mixin stores the `std::stop_token`
39 that was passed when your coroutine was awaited.
40
41 4. **Stop token access** — Coroutine code can retrieve the token via
42 `co_await this_coro::stop_token`.
43
44 5. **Executor storage** — The mixin stores the `executor_ref`
45 that this coroutine is bound to.
46
47 6. **Executor access** — Coroutine code can retrieve the executor via
48 `co_await this_coro::executor`.
49
50 @tparam Derived The derived promise type (CRTP pattern).
51
52 @par Basic Usage
53
54 For coroutines that need to access their stop token or executor:
55
56 @code
57 struct my_task
58 {
59 struct promise_type : io_awaitable_support<promise_type>
60 {
61 my_task get_return_object();
62 std::suspend_always initial_suspend() noexcept;
63 std::suspend_always final_suspend() noexcept;
64 void return_void();
65 void unhandled_exception();
66 };
67
68 // ... awaitable interface ...
69 };
70
71 my_task example()
72 {
73 auto token = co_await this_coro::stop_token;
74 auto ex = co_await this_coro::executor;
75 // Use token and ex...
76 }
77 @endcode
78
79 @par Custom Awaitable Transformation
80
81 If your promise needs to transform awaitables (e.g., for affinity or
82 logging), override `transform_awaitable` instead of `await_transform`:
83
84 @code
85 struct promise_type : io_awaitable_support<promise_type>
86 {
87 template<typename A>
88 auto transform_awaitable(A&& a)
89 {
90 // Your custom transformation logic
91 return std::forward<A>(a);
92 }
93 };
94 @endcode
95
96 The mixin's `await_transform` intercepts @ref this_coro::stop_token_tag and
97 @ref this_coro::executor_tag, then delegates all other awaitables to your
98 `transform_awaitable`.
99
100 @par Making Your Coroutine an IoAwaitable
101
102 The mixin handles the "inside the coroutine" part—accessing the token
103 and executor. To receive these when your coroutine is awaited (satisfying
104 @ref IoAwaitable), implement the `await_suspend` overload on your
105 coroutine return type:
106
107 @code
108 struct my_task
109 {
110 struct promise_type : io_awaitable_support<promise_type> { ... };
111
112 std::coroutine_handle<promise_type> h_;
113
114 // IoAwaitable await_suspend receives and stores the token and executor
115 coro await_suspend(coro cont, executor_ref ex, std::stop_token token)
116 {
117 h_.promise().set_stop_token(token);
118 h_.promise().set_executor(ex);
119 // ... rest of suspend logic ...
120 }
121 };
122 @endcode
123
124 @par Thread Safety
125 The stop token and executor are stored during `await_suspend` and read
126 during `co_await this_coro::stop_token` or `co_await this_coro::executor`.
127 These occur on the same logical thread of execution, so no synchronization
128 is required.
129
130 @see this_coro::stop_token
131 @see this_coro::executor
132 @see IoAwaitable
133 */
134 template<typename Derived>
135 class io_awaitable_support
136 {
137 executor_ref executor_;
138 std::stop_token stop_token_;
139 std::pmr::memory_resource* alloc_ = nullptr;
140 executor_ref caller_ex_;
141 mutable coro cont_{nullptr};
142
143 public:
144 //----------------------------------------------------------
145 // Frame allocation support
146 //----------------------------------------------------------
147
148 private:
149 static constexpr std::size_t ptr_alignment = alignof(void*);
150
151 static std::size_t
152 6366 aligned_offset(std::size_t n) noexcept
153 {
154 6366 return (n + ptr_alignment - 1) & ~(ptr_alignment - 1);
155 }
156
157 public:
158 /** Allocate a coroutine frame.
159
160 Uses the thread-local frame allocator set by run_async.
161 Falls back to default memory resource if not set.
162 Stores the allocator pointer at the end of each frame for
163 correct deallocation even when TLS changes.
164 */
165 static void*
166 3183 operator new(std::size_t size)
167 {
168 3183 auto* mr = current_frame_allocator();
169
2/2
✓ Branch 0 taken 83 times.
✓ Branch 1 taken 3100 times.
 3183 if(!mr)
170 83 mr = std::pmr::get_default_resource();
171
172 // Allocate extra space for memory_resource pointer
173 3183 std::size_t ptr_offset = aligned_offset(size);
174 3183 std::size_t total = ptr_offset + sizeof(std::pmr::memory_resource*);
175 3183 void* raw = mr->allocate(total, alignof(std::max_align_t));
176
177 // Store the allocator pointer at the end
178 3183 auto* ptr_loc = reinterpret_cast<std::pmr::memory_resource**>(
179 static_cast<char*>(raw) + ptr_offset);
180 3183 *ptr_loc = mr;
181
182 3183 return raw;
183 }
184
185 /** Deallocate a coroutine frame.
186
187 Reads the allocator pointer stored at the end of the frame
188 to ensure correct deallocation regardless of current TLS.
189 */
190 static void
191 3183 operator delete(void* ptr, std::size_t size)
192 {
193 // Read the allocator pointer from the end of the frame
194 3183 std::size_t ptr_offset = aligned_offset(size);
195 3183 auto* ptr_loc = reinterpret_cast<std::pmr::memory_resource**>(
196 static_cast<char*>(ptr) + ptr_offset);
197 3183 auto* mr = *ptr_loc;
198
199 3183 std::size_t total = ptr_offset + sizeof(std::pmr::memory_resource*);
200 3183 mr->deallocate(ptr, total, alignof(std::max_align_t));
201 3183 }
202
203 3183 ~io_awaitable_support()
204 {
205
2/2
✓ Branch 1 taken 1 time.
✓ Branch 2 taken 3182 times.
 3183 if (cont_)
206 1 cont_.destroy();
207 3183 }
208
209 /** Store a frame allocator for later retrieval.
210
211 Call this from initial_suspend to capture the current
212 TLS allocator for propagation to child tasks.
213
214 @param alloc The allocator to store.
215 */
216 void
217 3179 set_frame_allocator(std::pmr::memory_resource* alloc) noexcept
218 {
219 3179 alloc_ = alloc;
220 3179 }
221
222 /** Return the stored frame allocator.
223
224 @return The allocator, or nullptr if none was set.
225 */
226 std::pmr::memory_resource*
227 20143 frame_allocator() const noexcept
228 {
229 20143 return alloc_;
230 }
231
232 //----------------------------------------------------------
233 // Continuation support
234 //----------------------------------------------------------
235
236 /** Store continuation and caller's executor for completion dispatch.
237
238 Call this from your coroutine type's `await_suspend` overload to
239 set up the completion path. On completion, the coroutine will
240 resume the continuation, dispatching through the caller's executor
241 if it differs from this coroutine's executor.
242
243 @param cont The continuation to resume on completion.
244 @param caller_ex The caller's executor for completion dispatch.
245 */
246 3129 void set_continuation(coro cont, executor_ref caller_ex) noexcept
247 {
248 3129 cont_ = cont;
249 3129 caller_ex_ = caller_ex;
250 3129 }
251
252 /** Return the handle to resume on completion with dispatch-awareness.
253
254 If no continuation was set, returns `std::noop_coroutine()`.
255 If the coroutine's executor matches the caller's executor, returns
256 the continuation directly for symmetric transfer.
257 Otherwise, dispatches through the caller's executor and returns
258 `std::noop_coroutine()`.
259
260 Call this from your `final_suspend` awaiter's `await_suspend`.
261
262 @return A coroutine handle for symmetric transfer.
263 */
264 3171 coro complete() const noexcept
265 {
266
2/2
✓ Branch 1 taken 43 times.
✓ Branch 2 taken 3128 times.
 3171 if(!cont_)
267 43 return std::noop_coroutine();
268
1/2
✓ Branch 1 taken 3128 times.
✗ Branch 2 not taken.
 3128 if(executor_ == caller_ex_)
269 3128 return std::exchange(cont_, nullptr);
270 ✗ caller_ex_.dispatch(std::exchange(cont_, nullptr));
271 ✗ return std::noop_coroutine();
272 }
273
274 /** Store a stop token for later retrieval.
275
276 Call this from your coroutine type's `await_suspend`
277 overload to make the token available via
278 `co_await this_coro::stop_token`.
279
280 @param token The stop token to store.
281 */
282 3131 void set_stop_token(std::stop_token token) noexcept
283 {
284 3131 stop_token_ = token;
285 3131 }
286
287 /** Return the stored stop token.
288
289 @return The stop token, or a default-constructed token if none was set.
290 */
291 1943 std::stop_token const& stop_token() const noexcept
292 {
293 1943 return stop_token_;
294 }
295
296 /** Store an executor for later retrieval.
297
298 Call this from your coroutine type's `await_suspend`
299 overload to make the executor available via
300 `co_await this_coro::executor`.
301
302 @param ex The executor to store.
303 */
304 3131 void set_executor(executor_ref ex) noexcept
305 {
306 3131 executor_ = ex;
307 3131 }
308
309 /** Return the stored executor.
310
311 @return The executor, or a default-constructed executor_ref if none was set.
312 */
313 1943 executor_ref executor() const noexcept
314 {
315 1943 return executor_;
316 }
317
318 /** Transform an awaitable before co_await.
319
320 Override this in your derived promise type to customize how
321 awaitables are transformed. The default implementation passes
322 the awaitable through unchanged.
323
324 @param a The awaitable expression from `co_await a`.
325
326 @return The transformed awaitable.
327 */
328 template<typename A>
329 decltype(auto) transform_awaitable(A&& a)
330 {
331 return std::forward<A>(a);
332 }
333
334 /** Intercept co_await expressions.
335
336 This function handles @ref this_coro::stop_token_tag and
337 @ref this_coro::executor_tag specially, returning an awaiter that
338 yields the stored value. All other awaitables are delegated to
339 @ref transform_awaitable.
340
341 @param t The awaited expression.
342
343 @return An awaiter for the expression.
344 */
345 template<typename T>
346 7013 auto await_transform(T&& t)
347 {
348 if constexpr (std::is_same_v<std::decay_t<T>, this_coro::stop_token_tag>)
349 {
350 struct awaiter
351 {
352 std::stop_token token_;
353
354 35 bool await_ready() const noexcept
355 {
356 35 return true;
357 }
358
359 1 void await_suspend(coro) const noexcept
360 {
361 1 }
362
363 34 std::stop_token await_resume() const noexcept
364 {
365 34 return token_;
366 }
367 };
368 36 return awaiter{stop_token_};
369 }
370 else if constexpr (std::is_same_v<std::decay_t<T>, this_coro::executor_tag>)
371 {
372 struct awaiter
373 {
374 executor_ref executor_;
375
376 2 bool await_ready() const noexcept
377 {
378 2 return true;
379 }
380
381 1 void await_suspend(coro) const noexcept
382 {
383 1 }
384
385 1 executor_ref await_resume() const noexcept
386 {
387 1 return executor_;
388 }
389 };
390 3 return awaiter{executor_};
391 }
392 else
393 {
394 5709 return static_cast<Derived*>(this)->transform_awaitable(
395 6974 std::forward<T>(t));
396 }
397 }
398 };
399
400 } // namespace capy
401 } // namespace boost
402
403 #endif
404