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_IO_AWAITABLE_HPP
11 : #define BOOST_CAPY_IO_AWAITABLE_HPP
12 :
13 : #include <boost/capy/detail/config.hpp>
14 : #include <boost/capy/coro.hpp>
15 : #include <boost/capy/ex/executor_ref.hpp>
16 :
17 : #include <coroutine>
18 : #include <stop_token>
19 : #include <type_traits>
20 :
21 : namespace boost {
22 : namespace capy {
23 :
24 : /** Tag type for coroutine stop token retrieval.
25 :
26 : This tag is returned by @ref get_stop_token and intercepted by a
27 : promise type's `await_transform` to yield the coroutine's current
28 : stop token. The tag itself carries no data; it serves only as a
29 : sentinel for compile-time dispatch.
30 :
31 : @see get_stop_token
32 : @see io_awaitable_support
33 : */
34 : struct get_stop_token_tag {};
35 :
36 : /** Tag type for coroutine executor retrieval.
37 :
38 : This tag is returned by @ref get_executor and intercepted by a
39 : promise type's `await_transform` to yield the coroutine's current
40 : executor. The tag itself carries no data; it serves only as a
41 : sentinel for compile-time dispatch.
42 :
43 : @see get_executor
44 : @see io_awaitable_support
45 : */
46 : struct get_executor_tag {};
47 :
48 : /** Return a tag that yields the current stop token when awaited.
49 :
50 : Use `co_await get_stop_token()` inside a coroutine whose promise
51 : type supports stop token access (e.g., inherits from
52 : @ref io_awaitable_support). The returned stop token reflects whatever
53 : token was passed to this coroutine when it was awaited.
54 :
55 : @par Example
56 : @code
57 : task<void> cancellable_work()
58 : {
59 : auto token = co_await get_stop_token();
60 : for (int i = 0; i < 1000; ++i)
61 : {
62 : if (token.stop_requested())
63 : co_return; // Exit gracefully on cancellation
64 : co_await process_chunk(i);
65 : }
66 : }
67 : @endcode
68 :
69 : @par Behavior
70 : @li If no stop token was propagated, returns a default-constructed
71 : `std::stop_token` (where `stop_possible()` returns `false`).
72 : @li The returned token remains valid for the coroutine's lifetime.
73 : @li This operation never suspends; `await_ready()` always returns `true`.
74 :
75 : @return A tag that `await_transform` intercepts to return the stop token.
76 :
77 : @see get_stop_token_tag
78 : @see io_awaitable_support
79 : */
80 14 : inline get_stop_token_tag get_stop_token() noexcept
81 : {
82 14 : return {};
83 : }
84 :
85 : /** Return a tag that yields the current executor when awaited.
86 :
87 : Use `co_await get_executor()` inside a coroutine whose promise
88 : type supports executor access (e.g., inherits from
89 : @ref io_awaitable_support). The returned executor reflects the
90 : executor this coroutine is bound to.
91 :
92 : @par Example
93 : @code
94 : task<void> example()
95 : {
96 : executor_ref ex = co_await get_executor();
97 : // ex is the executor this coroutine is bound to
98 : }
99 : @endcode
100 :
101 : @par Behavior
102 : @li If no executor was set, returns a default-constructed
103 : `executor_ref` (where `operator bool()` returns `false`).
104 : @li This operation never suspends; `await_ready()` always returns `true`.
105 :
106 : @return A tag that `await_transform` intercepts to return the executor.
107 :
108 : @see get_executor_tag
109 : @see io_awaitable_support
110 : */
111 4 : inline get_executor_tag get_executor() noexcept
112 : {
113 4 : return {};
114 : }
115 :
116 : /** Concept for I/O awaitable types.
117 :
118 : An awaitable is an I/O awaitable if it participates in the I/O awaitable
119 : protocol by accepting an executor and a stop_token in its `await_suspend`
120 : method. This enables zero-overhead scheduler affinity and cancellation
121 : support.
122 :
123 : @tparam A The awaitable type.
124 : @tparam P The promise type (defaults to void).
125 :
126 : @par Requirements
127 : @li `A` must provide `await_suspend(std::coroutine_handle<P> h, Ex const& ex,
128 : std::stop_token token)`
129 : @li The awaitable must use the executor `ex` to resume the caller
130 : @li The awaitable should use the stop_token to support cancellation
131 :
132 : @par Example
133 : @code
134 : struct my_io_op
135 : {
136 : template<typename Executor>
137 : auto await_suspend(std::coroutine_handle<> h, Executor const& ex,
138 : std::stop_token token)
139 : {
140 : start_async([h, &ex, token] {
141 : if (token.stop_requested()) {
142 : // Handle cancellation
143 : }
144 : ex.dispatch(h); // Schedule resumption through executor
145 : });
146 : return std::noop_coroutine();
147 : }
148 : // ... await_ready, await_resume ...
149 : };
150 : @endcode
151 : */
152 : template<typename A, typename P = void>
153 : concept IoAwaitable =
154 : requires(
155 : A a,
156 : std::coroutine_handle<P> h,
157 : executor_ref ex,
158 : std::stop_token token)
159 : {
160 : a.await_suspend(h, ex, token);
161 : };
162 :
163 : /** CRTP mixin that adds I/O awaitable support to a promise type.
164 :
165 : Inherit from this class to enable these capabilities in your coroutine:
166 :
167 : 1. **Stop token storage** — The mixin stores the `std::stop_token`
168 : that was passed when your coroutine was awaited.
169 :
170 : 2. **Stop token access** — Coroutine code can retrieve the token via
171 : `co_await get_stop_token()`.
172 :
173 : 3. **Executor storage** — The mixin stores the `executor_ref`
174 : that this coroutine is bound to.
175 :
176 : 4. **Executor access** — Coroutine code can retrieve the executor via
177 : `co_await get_executor()`.
178 :
179 : @tparam Derived The derived promise type (CRTP pattern).
180 :
181 : @par Basic Usage
182 :
183 : For coroutines that need to access their stop token or executor:
184 :
185 : @code
186 : struct my_task
187 : {
188 : struct promise_type : io_awaitable_support<promise_type>
189 : {
190 : my_task get_return_object();
191 : std::suspend_always initial_suspend() noexcept;
192 : std::suspend_always final_suspend() noexcept;
193 : void return_void();
194 : void unhandled_exception();
195 : };
196 :
197 : // ... awaitable interface ...
198 : };
199 :
200 : my_task example()
201 : {
202 : auto token = co_await get_stop_token();
203 : auto ex = co_await get_executor();
204 : // Use token and ex...
205 : }
206 : @endcode
207 :
208 : @par Custom Awaitable Transformation
209 :
210 : If your promise needs to transform awaitables (e.g., for affinity or
211 : logging), override `transform_awaitable` instead of `await_transform`:
212 :
213 : @code
214 : struct promise_type : io_awaitable_support<promise_type>
215 : {
216 : template<typename A>
217 : auto transform_awaitable(A&& a)
218 : {
219 : // Your custom transformation logic
220 : return std::forward<A>(a);
221 : }
222 : };
223 : @endcode
224 :
225 : The mixin's `await_transform` intercepts @ref get_stop_token_tag and
226 : @ref get_executor_tag, then delegates all other awaitables to your
227 : `transform_awaitable`.
228 :
229 : @par Making Your Coroutine an IoAwaitable
230 :
231 : The mixin handles the "inside the coroutine" part—accessing the token
232 : and executor. To receive these when your coroutine is awaited (satisfying
233 : @ref IoAwaitable), implement the `await_suspend` overload on your
234 : coroutine return type:
235 :
236 : @code
237 : struct my_task
238 : {
239 : struct promise_type : io_awaitable_support<promise_type> { ... };
240 :
241 : std::coroutine_handle<promise_type> h_;
242 :
243 : // IoAwaitable await_suspend receives and stores the token and executor
244 : template<class Ex>
245 : coro await_suspend(coro cont, Ex const& ex, std::stop_token token)
246 : {
247 : h_.promise().set_stop_token(token);
248 : h_.promise().set_executor(ex);
249 : // ... rest of suspend logic ...
250 : }
251 : };
252 : @endcode
253 :
254 : @par Thread Safety
255 : The stop token and executor are stored during `await_suspend` and read
256 : during `co_await get_stop_token()` or `co_await get_executor()`. These
257 : occur on the same logical thread of execution, so no synchronization
258 : is required.
259 :
260 : @see get_stop_token
261 : @see get_executor
262 : @see IoAwaitable
263 : */
264 : template<typename Derived>
265 : class io_awaitable_support
266 : {
267 : std::stop_token stop_token_;
268 : executor_ref ex_;
269 :
270 : public:
271 : /** Store a stop token for later retrieval.
272 :
273 : Call this from your coroutine type's `await_suspend`
274 : overload to make the token available via `co_await get_stop_token()`.
275 :
276 : @param token The stop token to store.
277 : */
278 161 : void set_stop_token(std::stop_token token) noexcept
279 : {
280 161 : stop_token_ = token;
281 161 : }
282 :
283 : /** Return the stored stop token.
284 :
285 : @return The stop token, or a default-constructed token if none was set.
286 : */
287 66 : std::stop_token const& stop_token() const noexcept
288 : {
289 66 : return stop_token_;
290 : }
291 :
292 : /** Store an executor for later retrieval.
293 :
294 : Call this from your coroutine type's `await_suspend`
295 : overload to make the executor available via `co_await get_executor()`.
296 :
297 : @param ex The executor to store.
298 : */
299 184 : void set_executor(executor_ref ex) noexcept
300 : {
301 184 : ex_ = ex;
302 184 : }
303 :
304 : /** Return the stored executor.
305 :
306 : @return The executor, or a default-constructed executor_ref if none was set.
307 : */
308 66 : executor_ref executor() const noexcept
309 : {
310 66 : return ex_;
311 : }
312 :
313 : /** Transform an awaitable before co_await.
314 :
315 : Override this in your derived promise type to customize how
316 : awaitables are transformed. The default implementation passes
317 : the awaitable through unchanged.
318 :
319 : @param a The awaitable expression from `co_await a`.
320 :
321 : @return The transformed awaitable.
322 : */
323 : template<typename A>
324 : decltype(auto) transform_awaitable(A&& a)
325 : {
326 : return std::forward<A>(a);
327 : }
328 :
329 : /** Intercept co_await expressions.
330 :
331 : This function handles @ref get_stop_token_tag and @ref get_executor_tag
332 : specially, returning an awaiter that yields the stored value. All other
333 : awaitables are delegated to @ref transform_awaitable.
334 :
335 : @param t The awaited expression.
336 :
337 : @return An awaiter for the expression.
338 : */
339 : template<typename T>
340 81 : auto await_transform(T&& t)
341 : {
342 : if constexpr (std::is_same_v<std::decay_t<T>, get_stop_token_tag>)
343 : {
344 : struct awaiter
345 : {
346 : std::stop_token token_;
347 :
348 12 : bool await_ready() const noexcept
349 : {
350 12 : return true;
351 : }
352 :
353 1 : void await_suspend(coro) const noexcept
354 : {
355 1 : }
356 :
357 11 : std::stop_token await_resume() const noexcept
358 : {
359 11 : return token_;
360 : }
361 : };
362 13 : return awaiter{stop_token_};
363 : }
364 : else if constexpr (std::is_same_v<std::decay_t<T>, get_executor_tag>)
365 : {
366 : struct awaiter
367 : {
368 : executor_ref ex_;
369 :
370 2 : bool await_ready() const noexcept
371 : {
372 2 : return true;
373 : }
374 :
375 1 : void await_suspend(coro) const noexcept
376 : {
377 1 : }
378 :
379 1 : executor_ref await_resume() const noexcept
380 : {
381 1 : return ex_;
382 : }
383 : };
384 3 : return awaiter{ex_};
385 : }
386 : else
387 : {
388 25 : return static_cast<Derived*>(this)->transform_awaitable(
389 65 : std::forward<T>(t));
390 : }
391 : }
392 : };
393 :
394 : } // namespace capy
395 : } // namespace boost
396 :
397 : #endif
|
