Line data Source code
1 : //
2 : // Copyright (c) 2025 Vinnie Falco (vinnie dot falco at gmail dot 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/coro.hpp>
15 : #include <boost/capy/ex/executor_ref.hpp>
16 : #include <boost/capy/ex/frame_allocator.hpp>
17 : #include <boost/capy/ex/this_coro.hpp>
18 :
19 : #include <coroutine>
20 : #include <cstddef>
21 : #include <memory_resource>
22 : #include <stop_token>
23 : #include <type_traits>
24 :
25 : namespace boost {
26 : namespace capy {
27 :
28 : /** CRTP mixin that adds I/O awaitable support to a promise type.
29 :
30 : Inherit from this class to enable these capabilities in your coroutine:
31 :
32 : 1. **Frame allocation** — The mixin provides `operator new/delete` that
33 : use the thread-local frame allocator set by `run_async`.
34 :
35 : 2. **Frame allocator storage** — The mixin stores the allocator pointer
36 : for propagation to child tasks.
37 :
38 : 3. **Stop token storage** — The mixin stores the `std::stop_token`
39 : that was passed when your coroutine was awaited.
40 :
41 : 4. **Stop token access** — Coroutine code can retrieve the token via
42 : `co_await this_coro::stop_token`.
43 :
44 : 5. **Executor storage** — The mixin stores the `executor_ref`
45 : that this coroutine is bound to.
46 :
47 : 6. **Executor access** — Coroutine code can retrieve the executor via
48 : `co_await this_coro::executor`.
49 :
50 : @tparam Derived The derived promise type (CRTP pattern).
51 :
52 : @par Basic Usage
53 :
54 : For coroutines that need to access their stop token or executor:
55 :
56 : @code
57 : struct my_task
58 : {
59 : struct promise_type : io_awaitable_support<promise_type>
60 : {
61 : my_task get_return_object();
62 : std::suspend_always initial_suspend() noexcept;
63 : std::suspend_always final_suspend() noexcept;
64 : void return_void();
65 : void unhandled_exception();
66 : };
67 :
68 : // ... awaitable interface ...
69 : };
70 :
71 : my_task example()
72 : {
73 : auto token = co_await this_coro::stop_token;
74 : auto ex = co_await this_coro::executor;
75 : // Use token and ex...
76 : }
77 : @endcode
78 :
79 : @par Custom Awaitable Transformation
80 :
81 : If your promise needs to transform awaitables (e.g., for affinity or
82 : logging), override `transform_awaitable` instead of `await_transform`:
83 :
84 : @code
85 : struct promise_type : io_awaitable_support<promise_type>
86 : {
87 : template<typename A>
88 : auto transform_awaitable(A&& a)
89 : {
90 : // Your custom transformation logic
91 : return std::forward<A>(a);
92 : }
93 : };
94 : @endcode
95 :
96 : The mixin's `await_transform` intercepts @ref this_coro::stop_token_tag and
97 : @ref this_coro::executor_tag, then delegates all other awaitables to your
98 : `transform_awaitable`.
99 :
100 : @par Making Your Coroutine an IoAwaitable
101 :
102 : The mixin handles the "inside the coroutine" part—accessing the token
103 : and executor. To receive these when your coroutine is awaited (satisfying
104 : @ref IoAwaitable), implement the `await_suspend` overload on your
105 : coroutine return type:
106 :
107 : @code
108 : struct my_task
109 : {
110 : struct promise_type : io_awaitable_support<promise_type> { ... };
111 :
112 : std::coroutine_handle<promise_type> h_;
113 :
114 : // IoAwaitable await_suspend receives and stores the token and executor
115 : coro await_suspend(coro cont, executor_ref ex, std::stop_token token)
116 : {
117 : h_.promise().set_stop_token(token);
118 : h_.promise().set_executor(ex);
119 : // ... rest of suspend logic ...
120 : }
121 : };
122 : @endcode
123 :
124 : @par Thread Safety
125 : The stop token and executor are stored during `await_suspend` and read
126 : during `co_await this_coro::stop_token` or `co_await this_coro::executor`.
127 : These occur on the same logical thread of execution, so no synchronization
128 : is required.
129 :
130 : @see this_coro::stop_token
131 : @see this_coro::executor
132 : @see IoAwaitable
133 : */
134 : template<typename Derived>
135 : class io_awaitable_support
136 : {
137 : executor_ref executor_;
138 : std::stop_token stop_token_;
139 : std::pmr::memory_resource* alloc_ = nullptr;
140 : executor_ref caller_ex_;
141 : mutable coro cont_{nullptr};
142 :
143 : public:
144 : //----------------------------------------------------------
145 : // Frame allocation support
146 : //----------------------------------------------------------
147 :
148 : private:
149 : static constexpr std::size_t ptr_alignment = alignof(void*);
150 :
151 : static std::size_t
152 5716 : aligned_offset(std::size_t n) noexcept
153 : {
154 5716 : return (n + ptr_alignment - 1) & ~(ptr_alignment - 1);
155 : }
156 :
157 : public:
158 : /** Allocate a coroutine frame.
159 :
160 : Uses the thread-local frame allocator set by run_async.
161 : Falls back to default memory resource if not set.
162 : Stores the allocator pointer at the end of each frame for
163 : correct deallocation even when TLS changes.
164 : */
165 : static void*
166 2858 : operator new(std::size_t size)
167 : {
168 2858 : auto* mr = current_frame_allocator();
169 2858 : if(!mr)
170 83 : mr = std::pmr::get_default_resource();
171 :
172 : // Allocate extra space for memory_resource pointer
173 2858 : std::size_t ptr_offset = aligned_offset(size);
174 2858 : std::size_t total = ptr_offset + sizeof(std::pmr::memory_resource*);
175 2858 : void* raw = mr->allocate(total, alignof(std::max_align_t));
176 :
177 : // Store the allocator pointer at the end
178 2858 : auto* ptr_loc = reinterpret_cast<std::pmr::memory_resource**>(
179 : static_cast<char*>(raw) + ptr_offset);
180 2858 : *ptr_loc = mr;
181 :
182 2858 : return raw;
183 : }
184 :
185 : /** Deallocate a coroutine frame.
186 :
187 : Reads the allocator pointer stored at the end of the frame
188 : to ensure correct deallocation regardless of current TLS.
189 : */
190 : static void
191 2858 : operator delete(void* ptr, std::size_t size)
192 : {
193 : // Read the allocator pointer from the end of the frame
194 2858 : std::size_t ptr_offset = aligned_offset(size);
195 2858 : auto* ptr_loc = reinterpret_cast<std::pmr::memory_resource**>(
196 : static_cast<char*>(ptr) + ptr_offset);
197 2858 : auto* mr = *ptr_loc;
198 :
199 2858 : std::size_t total = ptr_offset + sizeof(std::pmr::memory_resource*);
200 2858 : mr->deallocate(ptr, total, alignof(std::max_align_t));
201 2858 : }
202 :
203 2858 : ~io_awaitable_support()
204 : {
205 2858 : if (cont_)
206 1 : cont_.destroy();
207 2858 : }
208 :
209 : /** Store a frame allocator for later retrieval.
210 :
211 : Call this from initial_suspend to capture the current
212 : TLS allocator for propagation to child tasks.
213 :
214 : @param alloc The allocator to store.
215 : */
216 : void
217 2845 : set_frame_allocator(std::pmr::memory_resource* alloc) noexcept
218 : {
219 2845 : alloc_ = alloc;
220 2845 : }
221 :
222 : /** Return the stored frame allocator.
223 :
224 : @return The allocator, or nullptr if none was set.
225 : */
226 : std::pmr::memory_resource*
227 18937 : frame_allocator() const noexcept
228 : {
229 18937 : return alloc_;
230 : }
231 :
232 : //----------------------------------------------------------
233 : // Continuation support
234 : //----------------------------------------------------------
235 :
236 : /** Store continuation and caller's executor for completion dispatch.
237 :
238 : Call this from your coroutine type's `await_suspend` overload to
239 : set up the completion path. On completion, the coroutine will
240 : resume the continuation, dispatching through the caller's executor
241 : if it differs from this coroutine's executor.
242 :
243 : @param cont The continuation to resume on completion.
244 : @param caller_ex The caller's executor for completion dispatch.
245 : */
246 2805 : void set_continuation(coro cont, executor_ref caller_ex) noexcept
247 : {
248 2805 : cont_ = cont;
249 2805 : caller_ex_ = caller_ex;
250 2805 : }
251 :
252 : /** Return the handle to resume on completion with dispatch-awareness.
253 :
254 : If no continuation was set, returns `std::noop_coroutine()`.
255 : If the coroutine's executor matches the caller's executor, returns
256 : the continuation directly for symmetric transfer.
257 : Otherwise, dispatches through the caller's executor first.
258 :
259 : Call this from your `final_suspend` awaiter's `await_suspend`.
260 :
261 : @return A coroutine handle for symmetric transfer.
262 : */
263 2846 : coro complete() const noexcept
264 : {
265 2846 : if(!cont_)
266 42 : return std::noop_coroutine();
267 2804 : if(executor_ == caller_ex_)
268 2804 : return std::exchange(cont_, nullptr);
269 0 : return caller_ex_.dispatch(std::exchange(cont_, nullptr));
270 : }
271 :
272 : /** Store a stop token for later retrieval.
273 :
274 : Call this from your coroutine type's `await_suspend`
275 : overload to make the token available via
276 : `co_await this_coro::stop_token`.
277 :
278 : @param token The stop token to store.
279 : */
280 2807 : void set_stop_token(std::stop_token token) noexcept
281 : {
282 2807 : stop_token_ = token;
283 2807 : }
284 :
285 : /** Return the stored stop token.
286 :
287 : @return The stop token, or a default-constructed token if none was set.
288 : */
289 1803 : std::stop_token const& stop_token() const noexcept
290 : {
291 1803 : return stop_token_;
292 : }
293 :
294 : /** Store an executor for later retrieval.
295 :
296 : Call this from your coroutine type's `await_suspend`
297 : overload to make the executor available via
298 : `co_await this_coro::executor`.
299 :
300 : @param ex The executor to store.
301 : */
302 2807 : void set_executor(executor_ref ex) noexcept
303 : {
304 2807 : executor_ = ex;
305 2807 : }
306 :
307 : /** Return the stored executor.
308 :
309 : @return The executor, or a default-constructed executor_ref if none was set.
310 : */
311 1803 : executor_ref executor() const noexcept
312 : {
313 1803 : return executor_;
314 : }
315 :
316 : /** Transform an awaitable before co_await.
317 :
318 : Override this in your derived promise type to customize how
319 : awaitables are transformed. The default implementation passes
320 : the awaitable through unchanged.
321 :
322 : @param a The awaitable expression from `co_await a`.
323 :
324 : @return The transformed awaitable.
325 : */
326 : template<typename A>
327 : decltype(auto) transform_awaitable(A&& a)
328 : {
329 : return std::forward<A>(a);
330 : }
331 :
332 : /** Intercept co_await expressions.
333 :
334 : This function handles @ref this_coro::stop_token_tag and
335 : @ref this_coro::executor_tag specially, returning an awaiter that
336 : yields the stored value. All other awaitables are delegated to
337 : @ref transform_awaitable.
338 :
339 : @param t The awaited expression.
340 :
341 : @return An awaiter for the expression.
342 : */
343 : template<typename T>
344 6852 : auto await_transform(T&& t)
345 : {
346 : if constexpr (std::is_same_v<std::decay_t<T>, this_coro::stop_token_tag>)
347 : {
348 : struct awaiter
349 : {
350 : std::stop_token token_;
351 :
352 14 : bool await_ready() const noexcept
353 : {
354 14 : return true;
355 : }
356 :
357 1 : void await_suspend(coro) const noexcept
358 : {
359 1 : }
360 :
361 13 : std::stop_token await_resume() const noexcept
362 : {
363 13 : return token_;
364 : }
365 : };
366 15 : return awaiter{stop_token_};
367 : }
368 : else if constexpr (std::is_same_v<std::decay_t<T>, this_coro::executor_tag>)
369 : {
370 : struct awaiter
371 : {
372 : executor_ref executor_;
373 :
374 2 : bool await_ready() const noexcept
375 : {
376 2 : return true;
377 : }
378 :
379 1 : void await_suspend(coro) const noexcept
380 : {
381 1 : }
382 :
383 1 : executor_ref await_resume() const noexcept
384 : {
385 1 : return executor_;
386 : }
387 : };
388 3 : return awaiter{executor_};
389 : }
390 : else
391 : {
392 5606 : return static_cast<Derived*>(this)->transform_awaitable(
393 6834 : std::forward<T>(t));
394 : }
395 : }
396 : };
397 :
398 : } // namespace capy
399 : } // namespace boost
400 :
401 : #endif
|
