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

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

98.5% Lines (64/65) 71.9% Functions (123/171) 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 5716 aligned_offset(std::size_t n) noexcept
153 {
154 5716 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 2858 operator new(std::size_t size)
167 {
168 2858 auto* mr = current_frame_allocator();
169
2/2
✓ Branch 0 taken 83 times.
✓ Branch 1 taken 2775 times.
 2858 if(!mr)
170 83 mr = std::pmr::get_default_resource();
171
172 // Allocate extra space for memory_resource pointer
173 2858 std::size_t ptr_offset = aligned_offset(size);
174 2858 std::size_t total = ptr_offset + sizeof(std::pmr::memory_resource*);
175 2858 void* raw = mr->allocate(total, alignof(std::max_align_t));
176
177 // Store the allocator pointer at the end
178 2858 auto* ptr_loc = reinterpret_cast<std::pmr::memory_resource**>(
179 static_cast<char*>(raw) + ptr_offset);
180 2858 *ptr_loc = mr;
181
182 2858 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 2858 operator delete(void* ptr, std::size_t size)
192 {
193 // Read the allocator pointer from the end of the frame
194 2858 std::size_t ptr_offset = aligned_offset(size);
195 2858 auto* ptr_loc = reinterpret_cast<std::pmr::memory_resource**>(
196 static_cast<char*>(ptr) + ptr_offset);
197 2858 auto* mr = *ptr_loc;
198
199 2858 std::size_t total = ptr_offset + sizeof(std::pmr::memory_resource*);
200 2858 mr->deallocate(ptr, total, alignof(std::max_align_t));
201 2858 }
202
203 2858 ~io_awaitable_support()
204 {
205
2/2
✓ Branch 1 taken 1 time.
✓ Branch 2 taken 2857 times.
 2858 if (cont_)
206 1 cont_.destroy();
207 2858 }
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 2845 set_frame_allocator(std::pmr::memory_resource* alloc) noexcept
218 {
219 2845 alloc_ = alloc;
220 2845 }
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 18937 frame_allocator() const noexcept
228 {
229 18937 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 2805 void set_continuation(coro cont, executor_ref caller_ex) noexcept
247 {
248 2805 cont_ = cont;
249 2805 caller_ex_ = caller_ex;
250 2805 }
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 first.
258
259 Call this from your `final_suspend` awaiter's `await_suspend`.
260
261 @return A coroutine handle for symmetric transfer.
262 */
263 2846 coro complete() const noexcept
264 {
265
2/2
✓ Branch 1 taken 42 times.
✓ Branch 2 taken 2804 times.
 2846 if(!cont_)
266 42 return std::noop_coroutine();
267
1/2
✓ Branch 1 taken 2804 times.
✗ Branch 2 not taken.
 2804 if(executor_ == caller_ex_)
268 2804 return std::exchange(cont_, nullptr);
269 ✗ return caller_ex_.dispatch(std::exchange(cont_, nullptr));
270 }
271
272 /** Store a stop token for later retrieval.
273
274 Call this from your coroutine type's `await_suspend`
275 overload to make the token available via
276 `co_await this_coro::stop_token`.
277
278 @param token The stop token to store.
279 */
280 2807 void set_stop_token(std::stop_token token) noexcept
281 {
282 2807 stop_token_ = token;
283 2807 }
284
285 /** Return the stored stop token.
286
287 @return The stop token, or a default-constructed token if none was set.
288 */
289 1803 std::stop_token const& stop_token() const noexcept
290 {
291 1803 return stop_token_;
292 }
293
294 /** Store an executor for later retrieval.
295
296 Call this from your coroutine type's `await_suspend`
297 overload to make the executor available via
298 `co_await this_coro::executor`.
299
300 @param ex The executor to store.
301 */
302 2807 void set_executor(executor_ref ex) noexcept
303 {
304 2807 executor_ = ex;
305 2807 }
306
307 /** Return the stored executor.
308
309 @return The executor, or a default-constructed executor_ref if none was set.
310 */
311 1803 executor_ref executor() const noexcept
312 {
313 1803 return executor_;
314 }
315
316 /** Transform an awaitable before co_await.
317
318 Override this in your derived promise type to customize how
319 awaitables are transformed. The default implementation passes
320 the awaitable through unchanged.
321
322 @param a The awaitable expression from `co_await a`.
323
324 @return The transformed awaitable.
325 */
326 template<typename A>
327 decltype(auto) transform_awaitable(A&& a)
328 {
329 return std::forward<A>(a);
330 }
331
332 /** Intercept co_await expressions.
333
334 This function handles @ref this_coro::stop_token_tag and
335 @ref this_coro::executor_tag specially, returning an awaiter that
336 yields the stored value. All other awaitables are delegated to
337 @ref transform_awaitable.
338
339 @param t The awaited expression.
340
341 @return An awaiter for the expression.
342 */
343 template<typename T>
344 6852 auto await_transform(T&& t)
345 {
346 if constexpr (std::is_same_v<std::decay_t<T>, this_coro::stop_token_tag>)
347 {
348 struct awaiter
349 {
350 std::stop_token token_;
351
352 14 bool await_ready() const noexcept
353 {
354 14 return true;
355 }
356
357 1 void await_suspend(coro) const noexcept
358 {
359 1 }
360
361 13 std::stop_token await_resume() const noexcept
362 {
363 13 return token_;
364 }
365 };
366 15 return awaiter{stop_token_};
367 }
368 else if constexpr (std::is_same_v<std::decay_t<T>, this_coro::executor_tag>)
369 {
370 struct awaiter
371 {
372 executor_ref executor_;
373
374 2 bool await_ready() const noexcept
375 {
376 2 return true;
377 }
378
379 1 void await_suspend(coro) const noexcept
380 {
381 1 }
382
383 1 executor_ref await_resume() const noexcept
384 {
385 1 return executor_;
386 }
387 };
388 3 return awaiter{executor_};
389 }
390 else
391 {
392 5606 return static_cast<Derived*>(this)->transform_awaitable(
393 6834 std::forward<T>(t));
394 }
395 }
396 };
397
398 } // namespace capy
399 } // namespace boost
400
401 #endif
402