include/boost/capy/ex/io_awaitable_promise_base.hpp

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