Index/include / boost / capy / ex / io_awaitable_promise_base.hpp
Prev Index Next

include/boost/capy/ex/io_awaitable_promise_base.hpp

96.4% Lines (54/56) 100.0% Functions (241/241)
include/boost/capy/ex/io_awaitable_promise_base.hpp
Line 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 4719 static void* operator new(std::size_t size)
149 {
150 4719 static auto* const rmr = get_recycling_memory_resource();
151
152 4719 auto* mr = get_current_frame_allocator();
153 4719 if(!mr)
154 2718 mr = std::pmr::get_default_resource();
155
156 4719 auto total = size + sizeof(std::pmr::memory_resource*);
157 void* raw;
158 4719 if(mr == rmr)
159 raw = static_cast<recycling_memory_resource*>(mr)
160 1988 ->allocate_fast(total, alignof(std::max_align_t));
161 else
162 2731 raw = mr->allocate(total, alignof(std::max_align_t));
163 4719 std::memcpy(static_cast<char*>(raw) + size, &mr, sizeof(mr));
164 4719 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 4719 static void operator delete(void* ptr, std::size_t size) noexcept
174 {
175 4719 static auto* const rmr = get_recycling_memory_resource();
176
177 std::pmr::memory_resource* mr;
178 4719 std::memcpy(&mr, static_cast<char*>(ptr) + size, sizeof(mr));
179 4719 auto total = size + sizeof(std::pmr::memory_resource*);
180 4719 if(mr == rmr)
181 static_cast<recycling_memory_resource*>(mr)
182 1988 ->deallocate_fast(ptr, total, alignof(std::max_align_t));
183 else
184 2731 mr->deallocate(ptr, total, alignof(std::max_align_t));
185 4719 }
186
187 4719 ~io_awaitable_promise_base()
188 {
189 // Abnormal teardown: destroy orphaned continuation
190 4719 if(cont_ != std::noop_coroutine())
191 1 cont_.destroy();
192 4719 }
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 4638 void set_continuation(std::coroutine_handle<> cont) noexcept
207 {
208 4638 cont_ = cont;
209 4638 }
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 4695 std::coroutine_handle<> continuation() const noexcept
219 {
220 4695 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 4716 void set_environment(io_env const* env) noexcept
237 {
238 4716 env_ = env;
239 4716 }
240
241 /** Return the stored execution environment.
242
243 @return The environment.
244 */
245 15520 io_env const* environment() const noexcept
246 {
247 15520 BOOST_CAPY_ASSERT(env_);
248 15520 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 8659 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 6754 return static_cast<Derived*>(this)->transform_awaitable(
335 8604 std::forward<T>(t));
336 }
337 }
338 };
339
340 } // namespace capy
341 } // namespace boost
342
343 #endif
344