TLA Line data Source code
1 : //
2 : // Copyright (c) 2026 Steve Gerbino
3 : // Copyright (c) 2026 Michael Vandeberg
4 : //
5 : // Distributed under the Boost Software License, Version 1.0. (See accompanying
6 : // file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
7 : //
8 : // Official repository: https://github.com/cppalliance/capy
9 : //
10 :
11 : #ifndef BOOST_CAPY_EX_ASYNC_WAKER_HPP
12 : #define BOOST_CAPY_EX_ASYNC_WAKER_HPP
13 :
14 : #include <boost/capy/detail/config.hpp>
15 : #include <boost/capy/continuation.hpp>
16 : #include <boost/capy/error.hpp>
17 : #include <boost/capy/ex/executor_ref.hpp>
18 : #include <boost/capy/ex/io_env.hpp>
19 : #include <boost/capy/io_result.hpp>
20 :
21 : #include <atomic>
22 : #include <coroutine>
23 : #include <new>
24 : #include <stop_token>
25 : #include <utility>
26 :
27 : /* async_waker implementation notes
28 : ===================================
29 :
30 : wake() must be callable from foreign threads (that is the whole
31 : point: the user's thread provides the timing). A waiter-side
32 : claimed_ flag is not enough there -- the
33 : waker has to dereference the waiter, and nothing would pin the
34 : waiter's frame between reading the pointer and claiming it.
35 :
36 : So the three-state st_ atomic is the single arbiter:
37 :
38 : empty --arm--> armed --wake/cancel CAS--> empty
39 : empty --wake--> token --wait consumes--> empty
40 :
41 : Whoever wins the armed->empty CAS owns the resume and may
42 : dereference waiter_: the frame cannot die underneath the
43 : winner because the coroutine only resumes when the winner
44 : posts it. The loser never touches the waiter. When the stop
45 : callback wins, a concurrent wake retries, finds empty, and
46 : latches a token -- a racing wakeup is deferred, never lost.
47 :
48 : Serialized resumption is required: await_suspend keeps
49 : writing after the publishing armed-CAS (the stop_cb
50 : placement-new and active_ = true), so a wake/cancel winner
51 : can post the continuation while that tail is still running.
52 : The posted resume must be ordered after await_suspend's
53 : return, which holds on a single-threaded executor (the one
54 : thread is still inside await_suspend) and on a strand (the
55 : resume is a later turn, synchronized with the current one).
56 : A raw multi-threaded executor lets another worker run
57 : await_resume against those in-flight writes. async_event and
58 : async_mutex make the same assumption; it is stated explicitly
59 : here because wake() invites foreign threads into the picture.
60 : */
61 :
62 : namespace boost {
63 : namespace capy {
64 :
65 : /** A single-slot waker that hands one wakeup to a waiting coroutine.
66 :
67 : This is the escape hatch for timing and other external events:
68 : the user provides the thread and the clock, capy provides the
69 : suspension point. One coroutine suspends in `wait()`; any
70 : thread wakes it with `wake()`.
71 :
72 : A wakeup with no waiter present is latched as a single pending
73 : token, and the next `wait()` consumes it immediately. This
74 : makes the wake-before-wait race benign without any lock
75 : protocol. Multiple wakes collapse into one token.
76 :
77 : @par Cancellation
78 :
79 : If the environment's stop token is triggered while suspended,
80 : the wait completes with `error::canceled`. A wake that loses
81 : the race against cancellation is latched for the next `wait()`
82 : rather than dropped.
83 :
84 : @par Zero Allocation
85 :
86 : No heap allocation occurs for wait or wake operations.
87 :
88 : @par Thread Safety
89 :
90 : Distinct objects: Safe.@n
91 : Shared objects: `wake()` may be called from any thread.
92 : `wait()` must only be awaited by one coroutine at a time. The
93 : executor must never run the coroutine's continuations
94 : concurrently: use a single-threaded executor, or a strand over
95 : a multi-threaded one. That is the same threading model as
96 : `async_event` and `async_mutex`. Awaiting `wait()` directly
97 : on a multi-threaded executor is undefined.
98 :
99 : This type is non-copyable and non-movable because a suspended
100 : waiter holds a pointer into the object.
101 :
102 : @par Example
103 : @code
104 : async_waker waker;
105 :
106 : // user-provided timing thread
107 : std::thread th([&waker] {
108 : std::this_thread::sleep_for(100ms);
109 : waker.wake();
110 : });
111 :
112 : task<> waiter() {
113 : auto [ec] = co_await waker.wait();
114 : // resumed on the executor after ~100ms
115 : }
116 : // ... th.join() after the pool drains
117 : @endcode
118 : */
119 : class async_waker
120 : {
121 : public:
122 : class wait_awaiter;
123 :
124 : private:
125 : static constexpr int state_empty = 0; // no token, no waiter
126 : static constexpr int state_token = 1; // latched wakeup
127 : static constexpr int state_armed = 2; // waiter suspended
128 :
129 : std::atomic<int> st_{state_empty};
130 : wait_awaiter* waiter_ = nullptr;
131 :
132 : public:
133 : /** Suspends the caller until `wake()` runs, or resumes it with `error::canceled` on a stop request.
134 : */
135 : class wait_awaiter
136 : {
137 : friend class async_waker;
138 :
139 : async_waker* waker_;
140 : continuation cont_;
141 : executor_ref ex_;
142 :
143 : // Declared before stop_cb_buf_: the callback accesses
144 : // these members, so they must still be alive if the
145 : // stop_cb_ destructor blocks.
146 : bool canceled_ = false;
147 : bool active_ = false;
148 : bool published_ = false;
149 :
150 : struct cancel_fn
151 : {
152 : wait_awaiter* self_;
153 :
154 HIT 15 : void operator()() const noexcept
155 : {
156 15 : int expected = state_armed;
157 30 : if(self_->waker_->st_.compare_exchange_strong(
158 : expected, state_empty,
159 : std::memory_order_acq_rel,
160 : std::memory_order_acquire))
161 : {
162 9 : self_->canceled_ = true;
163 9 : self_->ex_.post(self_->cont_);
164 : }
165 15 : }
166 : };
167 :
168 : using stop_cb_t = std::stop_callback<cancel_fn>;
169 :
170 : // Aligned storage for stop_cb_t. Declared last: its
171 : // destructor may block while the callback accesses the
172 : // members above.
173 : BOOST_CAPY_MSVC_WARNING_PUSH
174 : BOOST_CAPY_MSVC_WARNING_DISABLE(4324)
175 : alignas(stop_cb_t)
176 : unsigned char stop_cb_buf_[sizeof(stop_cb_t)];
177 : BOOST_CAPY_MSVC_WARNING_POP
178 :
179 33 : stop_cb_t& stop_cb_() noexcept
180 : {
181 33 : return *reinterpret_cast<stop_cb_t*>(stop_cb_buf_);
182 : }
183 :
184 : public:
185 : /** Destroy the awaiter, leaving the waker unable to reach it.
186 :
187 : Destroys the stop callback if one is registered. If the awaiter
188 : is still armed, it also returns the waker's slot to the empty
189 : state, so a later `wake()` cannot dereference a destroyed
190 : awaiter. That case means the frame is being torn down without
191 : ever being resumed; a wake arriving afterward latches a token
192 : instead.
193 : */
194 276 : ~wait_awaiter()
195 : {
196 276 : if(active_)
197 1 : stop_cb_().~stop_cb_t();
198 276 : if(published_)
199 : {
200 : // Destroyed while still armed (frame torn down
201 : // without resuming): deregister so a later
202 : // wake cannot touch the dead frame.
203 1 : int expected = state_armed;
204 1 : waker_->st_.compare_exchange_strong(
205 : expected, state_empty,
206 : std::memory_order_acq_rel,
207 : std::memory_order_acquire);
208 : }
209 276 : }
210 :
211 : /** Construct an awaiter for the given waker.
212 :
213 : @param waker The waker to wait on. It must outlive the awaiter.
214 : */
215 138 : explicit wait_awaiter(async_waker* waker) noexcept
216 138 : : waker_(waker)
217 : {
218 138 : }
219 :
220 : /** Construct by moving.
221 :
222 : The moved-from awaiter is left inert: its destructor no longer
223 : destroys the stop callback and no longer deregisters from the
224 : waker.
225 :
226 : @param o The awaiter to move from.
227 : */
228 138 : wait_awaiter(wait_awaiter&& o) noexcept
229 138 : : waker_(o.waker_)
230 138 : , cont_(o.cont_)
231 138 : , ex_(o.ex_)
232 138 : , canceled_(o.canceled_)
233 138 : , active_(std::exchange(o.active_, false))
234 138 : , published_(std::exchange(o.published_, false))
235 : {
236 138 : }
237 :
238 : /** Copy construction is disabled; an armed waiter is registered
239 : with the waker by address.
240 :
241 : @param other The awaiter that would be copied.
242 : */
243 : wait_awaiter(wait_awaiter const& other) = delete;
244 :
245 : /** Copy assignment is disabled; an armed waiter is registered
246 : with the waker by address.
247 :
248 : @param other The awaiter that would be assigned from.
249 :
250 : @return A reference to `*this`.
251 : */
252 : wait_awaiter& operator=(wait_awaiter const& other) = delete;
253 :
254 : /** Move assignment is disabled; an armed waiter is registered
255 : with the waker by address.
256 :
257 : @param other The awaiter that would be moved from.
258 :
259 : @return A reference to `*this`.
260 : */
261 : wait_awaiter& operator=(wait_awaiter&& other) = delete;
262 :
263 : /** Consume a latched token, completing synchronously.
264 :
265 : This is not a pure query: the check is a compare-exchange that
266 : takes the token. Calling it twice is not idempotent: the second
267 : call reports `false`, because the first already consumed the
268 : wakeup.
269 :
270 : @return `true` if a pending wakeup token was latched and has now
271 : been consumed, in which case the awaiting coroutine does not
272 : suspend; otherwise `false`.
273 : */
274 138 : bool await_ready() noexcept
275 : {
276 138 : int expected = state_token;
277 138 : return waker_->st_.compare_exchange_strong(
278 : expected, state_empty,
279 : std::memory_order_acq_rel,
280 138 : std::memory_order_acquire);
281 : }
282 :
283 : /** Arm the waker with the awaiting coroutine.
284 :
285 : This is the @ref IoAwaitable overload of `await_suspend`.
286 : Unlike `async_event` and `async_mutex`, it has three outcomes,
287 : because a `wake()` from another thread can land in the window
288 : between `await_ready` and this call.
289 :
290 : @li A stop request is already pending on `env->stop_token`: the
291 : awaiter records the cancellation and does not arm.
292 :
293 : @li The waker's slot is no longer empty. Under the single-waiter
294 : precondition that means a wakeup was latched after
295 : `await_ready` looked, so the token is consumed here instead
296 : and the wait succeeds.
297 :
298 : @li Otherwise the slot moves to the armed state, publishing this
299 : awaiter to the waker, and a stop callback is registered on
300 : `env->stop_token`. Whichever of `wake()` and that callback
301 : wins the armed-to-empty transition posts `h` through
302 : `env->executor`. The loser does nothing, and a losing
303 : `wake()` re-latches its token for the next `wait()`.
304 :
305 : @param h The awaiting coroutine, resumed when the waker fires
306 : or the wait is canceled.
307 :
308 : @param env The execution environment. Its executor posts the
309 : resumption and its stop token is watched for the duration of
310 : the wait. It must outlive the wait.
311 :
312 : @return `h` in the first two cases, which resumes the awaiting
313 : coroutine immediately; otherwise `std::noop_coroutine()`, which
314 : leaves the coroutine suspended and returns control to the
315 : resumer.
316 : */
317 : std::coroutine_handle<>
318 55 : await_suspend(
319 : std::coroutine_handle<> h,
320 : io_env const* env) noexcept
321 : {
322 55 : if(env->stop_token.stop_requested())
323 : {
324 22 : canceled_ = true;
325 22 : return h;
326 : }
327 33 : cont_.h = h;
328 33 : ex_ = env->executor;
329 33 : waker_->waiter_ = this;
330 :
331 33 : int expected = state_empty;
332 66 : if(!waker_->st_.compare_exchange_strong(
333 : expected, state_armed,
334 : std::memory_order_acq_rel,
335 : std::memory_order_acquire))
336 : {
337 : // Single-waiter precondition: a second concurrent
338 : // wait would find the slot armed.
339 MIS 0 : BOOST_CAPY_ASSERT(expected == state_token);
340 :
341 : // A wake latched between await_ready and here;
342 : // consume it and resume inline.
343 0 : waker_->st_.store(
344 : state_empty, std::memory_order_release);
345 0 : return h;
346 : }
347 HIT 33 : published_ = true;
348 :
349 99 : ::new(stop_cb_buf_) stop_cb_t(
350 33 : env->stop_token, cancel_fn{this});
351 33 : active_ = true;
352 33 : return std::noop_coroutine();
353 : }
354 :
355 : /** Complete the wait and report the outcome.
356 :
357 : Destroys the stop callback if one is registered and clears the
358 : armed bookkeeping, so the destructor does not deregister a slot
359 : the resumption already consumed.
360 :
361 : @return An empty `io_result<>` if the wait was woken, whether by
362 : `wake()` or by a token consumed inline. Otherwise one holding
363 : `error::canceled`, which means the stop token won the race.
364 : */
365 137 : io_result<> await_resume() noexcept
366 : {
367 137 : if(active_)
368 : {
369 32 : stop_cb_().~stop_cb_t();
370 32 : active_ = false;
371 : }
372 137 : published_ = false;
373 137 : if(canceled_)
374 30 : return {make_error_code(error::canceled)};
375 107 : return {{}};
376 : }
377 : };
378 :
379 : /// Construct with no token latched.
380 1 : async_waker() = default;
381 :
382 : /** Copy construction is disabled; an armed waiter points into the
383 : waker.
384 :
385 : @param other The waker that would be copied.
386 : */
387 : async_waker(async_waker const& other) = delete;
388 :
389 : /** Copy assignment is disabled; an armed waiter points into the waker.
390 :
391 : @param other The waker that would be assigned from.
392 :
393 : @return A reference to `*this`.
394 : */
395 : async_waker& operator=(async_waker const& other) = delete;
396 :
397 : /** Move construction is disabled; an armed waiter points into the
398 : waker.
399 :
400 : @param other The waker that would be moved from.
401 : */
402 : async_waker(async_waker&& other) = delete;
403 :
404 : /** Move assignment is disabled; an armed waiter points into the waker.
405 :
406 : @param other The waker that would be moved from.
407 :
408 : @return A reference to `*this`.
409 : */
410 : async_waker& operator=(async_waker&& other) = delete;
411 :
412 : /** Asynchronously wait until woken.
413 :
414 : If a token is latched, completes immediately and consumes
415 : it. Otherwise suspends until `wake()` or the stop token
416 : fires.
417 :
418 : @par Preconditions
419 : No other coroutine is currently waiting on this object.
420 :
421 : @return An awaitable that await-returns `io_result<>`;
422 : empty on wakeup, `error::canceled` if the stop
423 : token wins.
424 : */
425 138 : wait_awaiter wait() noexcept
426 : {
427 138 : return wait_awaiter{this};
428 : }
429 :
430 : /** Wake the waiter, or latch the wakeup if none waits.
431 :
432 : Callable from any thread. The waiter's resumption is
433 : posted through its executor; this call never resumes a
434 : coroutine inline. Multiple calls without an intervening
435 : `wait()` collapse into a single token.
436 : */
437 109 : void wake() noexcept
438 : {
439 : for(;;)
440 : {
441 109 : int s = st_.load(std::memory_order_acquire);
442 109 : if(s == state_token)
443 109 : return;
444 107 : if(s == state_empty)
445 : {
446 166 : if(st_.compare_exchange_weak(
447 : s, state_token,
448 : std::memory_order_acq_rel,
449 : std::memory_order_acquire))
450 83 : return;
451 MIS 0 : continue;
452 : }
453 : // armed: winning this CAS claims the waiter, whose
454 : // frame is pinned until we post its resumption.
455 HIT 48 : if(st_.compare_exchange_weak(
456 : s, state_empty,
457 : std::memory_order_acq_rel,
458 : std::memory_order_acquire))
459 : {
460 24 : auto* w = waiter_;
461 24 : w->ex_.post(w->cont_);
462 24 : return;
463 : }
464 MIS 0 : }
465 : }
466 : };
467 :
468 : } // namespace capy
469 : } // namespace boost
470 :
471 : #endif
|