ex/io_awaitable_support.hpp

96.1% Lines (49/51) 100.0% Functions (239/239) 90.0% Branches (18/20)
ex/io_awaitable_support.hpp
Line Branch Hits 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
3/4
✓ Branch 0 taken 94 times.
✓ Branch 1 taken 3647 times.
✓ Branch 3 taken 94 times.
✗ Branch 4 not taken.
3741 static auto* const rmr = get_recycling_memory_resource();
157
158 3741 auto* mr = current_frame_allocator();
159
2/2
✓ Branch 0 taken 1994 times.
✓ Branch 1 taken 1747 times.
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
2/2
✓ Branch 0 taken 1734 times.
✓ Branch 1 taken 2007 times.
3741 if(mr == rmr)
165 raw = static_cast<recycling_memory_resource*>(mr)
166
1/1
✓ Branch 1 taken 1734 times.
1734 ->allocate_fast(total, alignof(std::max_align_t));
167 else
168
1/1
✓ Branch 1 taken 2007 times.
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
3/4
✓ Branch 0 taken 94 times.
✓ Branch 1 taken 3647 times.
✓ Branch 3 taken 94 times.
✗ Branch 4 not taken.
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
2/2
✓ Branch 0 taken 1734 times.
✓ Branch 1 taken 2007 times.
3741 if(mr == rmr)
188 static_cast<recycling_memory_resource*>(mr)
189
1/1
✓ Branch 1 taken 1734 times.
1734 ->deallocate_fast(ptr, total, alignof(std::max_align_t));
190 else
191
1/1
✓ Branch 1 taken 2007 times.
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
2/2
✓ Branch 3 taken 1 time.
✓ Branch 4 taken 3740 times.
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 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 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
346