Line data 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 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 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 0 : 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 0 : 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
|
