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

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

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