]> git.bitcoin.ninja Git - rust-lightning/blob - lightning/src/util/wakers.rs
Merge pull request #2226 from alecchendev/2023-04-persist-network-graph-on-rgs
[rust-lightning] / lightning / src / util / wakers.rs
1 // This file is Copyright its original authors, visible in version control
2 // history.
3 //
4 // This file is licensed under the Apache License, Version 2.0 <LICENSE-APACHE
5 // or http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
6 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your option.
7 // You may not use this file except in accordance with one or both of these
8 // licenses.
9
10 //! Utilities which allow users to block on some future notification from LDK. These are
11 //! specifically used by [`ChannelManager`] to allow waiting until the [`ChannelManager`] needs to
12 //! be re-persisted.
13 //!
14 //! [`ChannelManager`]: crate::ln::channelmanager::ChannelManager
15
16 use alloc::sync::Arc;
17 use core::mem;
18 use crate::sync::Mutex;
19
20 use crate::prelude::*;
21
22 #[cfg(feature = "std")]
23 use crate::sync::Condvar;
24 #[cfg(feature = "std")]
25 use std::time::Duration;
26
27 use core::future::Future as StdFuture;
28 use core::task::{Context, Poll};
29 use core::pin::Pin;
30
31
32 /// Used to signal to one of many waiters that the condition they're waiting on has happened.
33 pub(crate) struct Notifier {
34         notify_pending: Mutex<(bool, Option<Arc<Mutex<FutureState>>>)>,
35 }
36
37 impl Notifier {
38         pub(crate) fn new() -> Self {
39                 Self {
40                         notify_pending: Mutex::new((false, None)),
41                 }
42         }
43
44         /// Wake waiters, tracking that wake needs to occur even if there are currently no waiters.
45         pub(crate) fn notify(&self) {
46                 let mut lock = self.notify_pending.lock().unwrap();
47                 if let Some(future_state) = &lock.1 {
48                         if complete_future(future_state) {
49                                 lock.1 = None;
50                                 return;
51                         }
52                 }
53                 lock.0 = true;
54         }
55
56         /// Gets a [`Future`] that will get woken up with any waiters
57         pub(crate) fn get_future(&self) -> Future {
58                 let mut lock = self.notify_pending.lock().unwrap();
59                 if let Some(existing_state) = &lock.1 {
60                         if existing_state.lock().unwrap().callbacks_made {
61                                 // If the existing `FutureState` has completed and actually made callbacks,
62                                 // consider the notification flag to have been cleared and reset the future state.
63                                 lock.1.take();
64                                 lock.0 = false;
65                         }
66                 }
67                 if let Some(existing_state) = &lock.1 {
68                         Future { state: Arc::clone(&existing_state) }
69                 } else {
70                         let state = Arc::new(Mutex::new(FutureState {
71                                 callbacks: Vec::new(),
72                                 callbacks_with_state: Vec::new(),
73                                 complete: lock.0,
74                                 callbacks_made: false,
75                         }));
76                         lock.1 = Some(Arc::clone(&state));
77                         Future { state }
78                 }
79         }
80
81         #[cfg(any(test, feature = "_test_utils"))]
82         pub fn notify_pending(&self) -> bool {
83                 self.notify_pending.lock().unwrap().0
84         }
85 }
86
87 macro_rules! define_callback { ($($bounds: path),*) => {
88 /// A callback which is called when a [`Future`] completes.
89 ///
90 /// Note that this MUST NOT call back into LDK directly, it must instead schedule actions to be
91 /// taken later. Rust users should use the [`std::future::Future`] implementation for [`Future`]
92 /// instead.
93 ///
94 /// Note that the [`std::future::Future`] implementation may only work for runtimes which schedule
95 /// futures when they receive a wake, rather than immediately executing them.
96 pub trait FutureCallback : $($bounds +)* {
97         /// The method which is called.
98         fn call(&self);
99 }
100
101 impl<F: Fn() $(+ $bounds)*> FutureCallback for F {
102         fn call(&self) { (self)(); }
103 }
104 } }
105
106 #[cfg(feature = "std")]
107 define_callback!(Send);
108 #[cfg(not(feature = "std"))]
109 define_callback!();
110
111 pub(crate) struct FutureState {
112         // When we're tracking whether a callback counts as having woken the user's code, we check the
113         // first bool - set to false if we're just calling a Waker, and true if we're calling an actual
114         // user-provided function.
115         callbacks: Vec<(bool, Box<dyn FutureCallback>)>,
116         callbacks_with_state: Vec<(bool, Box<dyn Fn(&Arc<Mutex<FutureState>>) -> () + Send>)>,
117         complete: bool,
118         callbacks_made: bool,
119 }
120
121 fn complete_future(this: &Arc<Mutex<FutureState>>) -> bool {
122         let mut state_lock = this.lock().unwrap();
123         let state = &mut *state_lock;
124         for (counts_as_call, callback) in state.callbacks.drain(..) {
125                 callback.call();
126                 state.callbacks_made |= counts_as_call;
127         }
128         for (counts_as_call, callback) in state.callbacks_with_state.drain(..) {
129                 (callback)(this);
130                 state.callbacks_made |= counts_as_call;
131         }
132         state.complete = true;
133         state.callbacks_made
134 }
135
136 /// A simple future which can complete once, and calls some callback(s) when it does so.
137 ///
138 /// Clones can be made and all futures cloned from the same source will complete at the same time.
139 #[derive(Clone)]
140 pub struct Future {
141         state: Arc<Mutex<FutureState>>,
142 }
143
144 impl Future {
145         /// Registers a callback to be called upon completion of this future. If the future has already
146         /// completed, the callback will be called immediately.
147         ///
148         /// This is not exported to bindings users, use the bindings-only `register_callback_fn` instead
149         pub fn register_callback(&self, callback: Box<dyn FutureCallback>) {
150                 let mut state = self.state.lock().unwrap();
151                 if state.complete {
152                         state.callbacks_made = true;
153                         mem::drop(state);
154                         callback.call();
155                 } else {
156                         state.callbacks.push((true, callback));
157                 }
158         }
159
160         // C bindings don't (currently) know how to map `Box<dyn Trait>`, and while it could add the
161         // following wrapper, doing it in the bindings is currently much more work than simply doing it
162         // here.
163         /// Registers a callback to be called upon completion of this future. If the future has already
164         /// completed, the callback will be called immediately.
165         #[cfg(c_bindings)]
166         pub fn register_callback_fn<F: 'static + FutureCallback>(&self, callback: F) {
167                 self.register_callback(Box::new(callback));
168         }
169
170         /// Waits until this [`Future`] completes.
171         #[cfg(feature = "std")]
172         pub fn wait(self) {
173                 Sleeper::from_single_future(self).wait();
174         }
175
176         /// Waits until this [`Future`] completes or the given amount of time has elapsed.
177         ///
178         /// Returns true if the [`Future`] completed, false if the time elapsed.
179         #[cfg(feature = "std")]
180         pub fn wait_timeout(self, max_wait: Duration) -> bool {
181                 Sleeper::from_single_future(self).wait_timeout(max_wait)
182         }
183
184         #[cfg(test)]
185         pub fn poll_is_complete(&self) -> bool {
186                 let mut state = self.state.lock().unwrap();
187                 if state.complete {
188                         state.callbacks_made = true;
189                         true
190                 } else { false }
191         }
192 }
193
194 use core::task::Waker;
195 struct StdWaker(pub Waker);
196 impl FutureCallback for StdWaker {
197         fn call(&self) { self.0.wake_by_ref() }
198 }
199
200 /// This is not exported to bindings users as Rust Futures aren't usable in language bindings.
201 impl<'a> StdFuture for Future {
202         type Output = ();
203
204         fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
205                 let mut state = self.state.lock().unwrap();
206                 if state.complete {
207                         state.callbacks_made = true;
208                         Poll::Ready(())
209                 } else {
210                         let waker = cx.waker().clone();
211                         state.callbacks.push((false, Box::new(StdWaker(waker))));
212                         Poll::Pending
213                 }
214         }
215 }
216
217 /// A struct which can be used to select across many [`Future`]s at once without relying on a full
218 /// async context.
219 #[cfg(feature = "std")]
220 pub struct Sleeper {
221         notifiers: Vec<Arc<Mutex<FutureState>>>,
222 }
223
224 #[cfg(feature = "std")]
225 impl Sleeper {
226         /// Constructs a new sleeper from one future, allowing blocking on it.
227         pub fn from_single_future(future: Future) -> Self {
228                 Self { notifiers: vec![future.state] }
229         }
230         /// Constructs a new sleeper from two futures, allowing blocking on both at once.
231         // Note that this is the common case - a ChannelManager and ChainMonitor.
232         pub fn from_two_futures(fut_a: Future, fut_b: Future) -> Self {
233                 Self { notifiers: vec![fut_a.state, fut_b.state] }
234         }
235         /// Constructs a new sleeper on many futures, allowing blocking on all at once.
236         pub fn new(futures: Vec<Future>) -> Self {
237                 Self { notifiers: futures.into_iter().map(|f| f.state).collect() }
238         }
239         /// Prepares to go into a wait loop body, creating a condition variable which we can block on
240         /// and an `Arc<Mutex<Option<_>>>` which gets set to the waking `Future`'s state prior to the
241         /// condition variable being woken.
242         fn setup_wait(&self) -> (Arc<Condvar>, Arc<Mutex<Option<Arc<Mutex<FutureState>>>>>) {
243                 let cv = Arc::new(Condvar::new());
244                 let notified_fut_mtx = Arc::new(Mutex::new(None));
245                 {
246                         for notifier_mtx in self.notifiers.iter() {
247                                 let cv_ref = Arc::clone(&cv);
248                                 let notified_fut_ref = Arc::clone(&notified_fut_mtx);
249                                 let mut notifier = notifier_mtx.lock().unwrap();
250                                 if notifier.complete {
251                                         *notified_fut_mtx.lock().unwrap() = Some(Arc::clone(&notifier_mtx));
252                                         break;
253                                 }
254                                 notifier.callbacks_with_state.push((false, Box::new(move |notifier_ref| {
255                                         *notified_fut_ref.lock().unwrap() = Some(Arc::clone(notifier_ref));
256                                         cv_ref.notify_all();
257                                 })));
258                         }
259                 }
260                 (cv, notified_fut_mtx)
261         }
262
263         /// Wait until one of the [`Future`]s registered with this [`Sleeper`] has completed.
264         pub fn wait(&self) {
265                 let (cv, notified_fut_mtx) = self.setup_wait();
266                 let notified_fut = cv.wait_while(notified_fut_mtx.lock().unwrap(), |fut_opt| fut_opt.is_none())
267                         .unwrap().take().expect("CV wait shouldn't have returned until the notifying future was set");
268                 notified_fut.lock().unwrap().callbacks_made = true;
269         }
270
271         /// Wait until one of the [`Future`]s registered with this [`Sleeper`] has completed or the
272         /// given amount of time has elapsed. Returns true if a [`Future`] completed, false if the time
273         /// elapsed.
274         pub fn wait_timeout(&self, max_wait: Duration) -> bool {
275                 let (cv, notified_fut_mtx) = self.setup_wait();
276                 let notified_fut =
277                         match cv.wait_timeout_while(notified_fut_mtx.lock().unwrap(), max_wait, |fut_opt| fut_opt.is_none()) {
278                                 Ok((_, e)) if e.timed_out() => return false,
279                                 Ok((mut notified_fut, _)) =>
280                                         notified_fut.take().expect("CV wait shouldn't have returned until the notifying future was set"),
281                                 Err(_) => panic!("Previous panic while a lock was held led to a lock panic"),
282                         };
283                 notified_fut.lock().unwrap().callbacks_made = true;
284                 true
285         }
286 }
287
288 #[cfg(test)]
289 mod tests {
290         use super::*;
291         use core::sync::atomic::{AtomicBool, Ordering};
292         use core::future::Future as FutureTrait;
293         use core::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
294
295         #[test]
296         fn notifier_pre_notified_future() {
297                 // Previously, if we generated a future after a `Notifier` had been notified, the future
298                 // would never complete. This tests this behavior, ensuring the future instead completes
299                 // immediately.
300                 let notifier = Notifier::new();
301                 notifier.notify();
302
303                 let callback = Arc::new(AtomicBool::new(false));
304                 let callback_ref = Arc::clone(&callback);
305                 notifier.get_future().register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
306                 assert!(callback.load(Ordering::SeqCst));
307         }
308
309         #[test]
310         fn notifier_future_completes_wake() {
311                 // Previously, if we were only using the `Future` interface to learn when a `Notifier` has
312                 // been notified, we'd never mark the notifier as not-awaiting-notify. This caused the
313                 // `lightning-background-processor` to persist in a tight loop.
314                 let notifier = Notifier::new();
315
316                 // First check the simple case, ensuring if we get notified a new future isn't woken until
317                 // a second `notify`.
318                 let callback = Arc::new(AtomicBool::new(false));
319                 let callback_ref = Arc::clone(&callback);
320                 notifier.get_future().register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
321                 assert!(!callback.load(Ordering::SeqCst));
322
323                 notifier.notify();
324                 assert!(callback.load(Ordering::SeqCst));
325
326                 let callback = Arc::new(AtomicBool::new(false));
327                 let callback_ref = Arc::clone(&callback);
328                 notifier.get_future().register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
329                 assert!(!callback.load(Ordering::SeqCst));
330
331                 notifier.notify();
332                 assert!(callback.load(Ordering::SeqCst));
333
334                 // Then check the case where the future is fetched before the notification, but a callback
335                 // is only registered after the `notify`, ensuring that it is still sufficient to ensure we
336                 // don't get an instant-wake when we get a new future.
337                 let future = notifier.get_future();
338                 notifier.notify();
339
340                 let callback = Arc::new(AtomicBool::new(false));
341                 let callback_ref = Arc::clone(&callback);
342                 future.register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
343                 assert!(callback.load(Ordering::SeqCst));
344
345                 let callback = Arc::new(AtomicBool::new(false));
346                 let callback_ref = Arc::clone(&callback);
347                 notifier.get_future().register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
348                 assert!(!callback.load(Ordering::SeqCst));
349         }
350
351         #[test]
352         fn new_future_wipes_notify_bit() {
353                 // Previously, if we were only using the `Future` interface to learn when a `Notifier` has
354                 // been notified, we'd never mark the notifier as not-awaiting-notify if a `Future` is
355                 // fetched after the notify bit has been set.
356                 let notifier = Notifier::new();
357                 notifier.notify();
358
359                 let callback = Arc::new(AtomicBool::new(false));
360                 let callback_ref = Arc::clone(&callback);
361                 notifier.get_future().register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
362                 assert!(callback.load(Ordering::SeqCst));
363
364                 let callback = Arc::new(AtomicBool::new(false));
365                 let callback_ref = Arc::clone(&callback);
366                 notifier.get_future().register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
367                 assert!(!callback.load(Ordering::SeqCst));
368
369                 notifier.notify();
370                 assert!(callback.load(Ordering::SeqCst));
371         }
372
373         #[cfg(feature = "std")]
374         #[test]
375         fn test_wait_timeout() {
376                 use crate::sync::Arc;
377                 use std::thread;
378
379                 let persistence_notifier = Arc::new(Notifier::new());
380                 let thread_notifier = Arc::clone(&persistence_notifier);
381
382                 let exit_thread = Arc::new(AtomicBool::new(false));
383                 let exit_thread_clone = exit_thread.clone();
384                 thread::spawn(move || {
385                         loop {
386                                 thread_notifier.notify();
387                                 if exit_thread_clone.load(Ordering::SeqCst) {
388                                         break
389                                 }
390                         }
391                 });
392
393                 // Check that we can block indefinitely until updates are available.
394                 let _ = persistence_notifier.get_future().wait();
395
396                 // Check that the Notifier will return after the given duration if updates are
397                 // available.
398                 loop {
399                         if persistence_notifier.get_future().wait_timeout(Duration::from_millis(100)) {
400                                 break
401                         }
402                 }
403
404                 exit_thread.store(true, Ordering::SeqCst);
405
406                 // Check that the Notifier will return after the given duration even if no updates
407                 // are available.
408                 loop {
409                         if !persistence_notifier.get_future().wait_timeout(Duration::from_millis(100)) {
410                                 break
411                         }
412                 }
413         }
414
415         #[cfg(feature = "std")]
416         #[test]
417         fn test_state_drops() {
418                 // Previously, there was a leak if a `Notifier` was `drop`ed without ever being notified
419                 // but after having been slept-on. This tests for that leak.
420                 use crate::sync::Arc;
421                 use std::thread;
422
423                 let notifier_a = Arc::new(Notifier::new());
424                 let notifier_b = Arc::new(Notifier::new());
425
426                 let thread_notifier_a = Arc::clone(&notifier_a);
427
428                 let future_a = notifier_a.get_future();
429                 let future_state_a = Arc::downgrade(&future_a.state);
430
431                 let future_b = notifier_b.get_future();
432                 let future_state_b = Arc::downgrade(&future_b.state);
433
434                 let join_handle = thread::spawn(move || {
435                         // Let the other thread get to the wait point, then notify it.
436                         std::thread::sleep(Duration::from_millis(50));
437                         thread_notifier_a.notify();
438                 });
439
440                 // Wait on the other thread to finish its sleep, note that the leak only happened if we
441                 // actually have to sleep here, not if we immediately return.
442                 Sleeper::from_two_futures(future_a, future_b).wait();
443
444                 join_handle.join().unwrap();
445
446                 // then drop the notifiers and make sure the future states are gone.
447                 mem::drop(notifier_a);
448                 mem::drop(notifier_b);
449
450                 assert!(future_state_a.upgrade().is_none() && future_state_b.upgrade().is_none());
451         }
452
453         #[test]
454         fn test_future_callbacks() {
455                 let future = Future {
456                         state: Arc::new(Mutex::new(FutureState {
457                                 callbacks: Vec::new(),
458                                 callbacks_with_state: Vec::new(),
459                                 complete: false,
460                                 callbacks_made: false,
461                         }))
462                 };
463                 let callback = Arc::new(AtomicBool::new(false));
464                 let callback_ref = Arc::clone(&callback);
465                 future.register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
466
467                 assert!(!callback.load(Ordering::SeqCst));
468                 complete_future(&future.state);
469                 assert!(callback.load(Ordering::SeqCst));
470                 complete_future(&future.state);
471         }
472
473         #[test]
474         fn test_pre_completed_future_callbacks() {
475                 let future = Future {
476                         state: Arc::new(Mutex::new(FutureState {
477                                 callbacks: Vec::new(),
478                                 callbacks_with_state: Vec::new(),
479                                 complete: false,
480                                 callbacks_made: false,
481                         }))
482                 };
483                 complete_future(&future.state);
484
485                 let callback = Arc::new(AtomicBool::new(false));
486                 let callback_ref = Arc::clone(&callback);
487                 future.register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
488
489                 assert!(callback.load(Ordering::SeqCst));
490                 assert!(future.state.lock().unwrap().callbacks.is_empty());
491         }
492
493         // Rather annoyingly, there's no safe way in Rust std to construct a Waker despite it being
494         // totally possible to construct from a trait implementation (though somewhat less effecient
495         // compared to a raw VTable). Instead, we have to write out a lot of boilerplate to build a
496         // waker, which we do here with a trivial Arc<AtomicBool> data element to track woke-ness.
497         const WAKER_V_TABLE: RawWakerVTable = RawWakerVTable::new(waker_clone, wake, wake_by_ref, drop);
498         unsafe fn wake_by_ref(ptr: *const ()) { let p = ptr as *const Arc<AtomicBool>; assert!(!(*p).fetch_or(true, Ordering::SeqCst)); }
499         unsafe fn drop(ptr: *const ()) { let p = ptr as *mut Arc<AtomicBool>; let _freed = Box::from_raw(p); }
500         unsafe fn wake(ptr: *const ()) { wake_by_ref(ptr); drop(ptr); }
501         unsafe fn waker_clone(ptr: *const ()) -> RawWaker {
502                 let p = ptr as *const Arc<AtomicBool>;
503                 RawWaker::new(Box::into_raw(Box::new(Arc::clone(&*p))) as *const (), &WAKER_V_TABLE)
504         }
505
506         fn create_waker() -> (Arc<AtomicBool>, Waker) {
507                 let a = Arc::new(AtomicBool::new(false));
508                 let waker = unsafe { Waker::from_raw(waker_clone((&a as *const Arc<AtomicBool>) as *const ())) };
509                 (a, waker)
510         }
511
512         #[test]
513         fn test_future() {
514                 let mut future = Future {
515                         state: Arc::new(Mutex::new(FutureState {
516                                 callbacks: Vec::new(),
517                                 callbacks_with_state: Vec::new(),
518                                 complete: false,
519                                 callbacks_made: false,
520                         }))
521                 };
522                 let mut second_future = Future { state: Arc::clone(&future.state) };
523
524                 let (woken, waker) = create_waker();
525                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Pending);
526                 assert!(!woken.load(Ordering::SeqCst));
527
528                 let (second_woken, second_waker) = create_waker();
529                 assert_eq!(Pin::new(&mut second_future).poll(&mut Context::from_waker(&second_waker)), Poll::Pending);
530                 assert!(!second_woken.load(Ordering::SeqCst));
531
532                 complete_future(&future.state);
533                 assert!(woken.load(Ordering::SeqCst));
534                 assert!(second_woken.load(Ordering::SeqCst));
535                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(()));
536                 assert_eq!(Pin::new(&mut second_future).poll(&mut Context::from_waker(&second_waker)), Poll::Ready(()));
537         }
538
539         #[test]
540         #[cfg(feature = "std")]
541         fn test_dropped_future_doesnt_count() {
542                 // Tests that if a Future gets drop'd before it is poll()ed `Ready` it doesn't count as
543                 // having been woken, leaving the notify-required flag set.
544                 let notifier = Notifier::new();
545                 notifier.notify();
546
547                 // If we get a future and don't touch it we're definitely still notify-required.
548                 notifier.get_future();
549                 assert!(notifier.get_future().wait_timeout(Duration::from_millis(1)));
550                 assert!(!notifier.get_future().wait_timeout(Duration::from_millis(1)));
551
552                 // Even if we poll'd once but didn't observe a `Ready`, we should be notify-required.
553                 let mut future = notifier.get_future();
554                 let (woken, waker) = create_waker();
555                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Pending);
556
557                 notifier.notify();
558                 assert!(woken.load(Ordering::SeqCst));
559                 assert!(notifier.get_future().wait_timeout(Duration::from_millis(1)));
560
561                 // However, once we do poll `Ready` it should wipe the notify-required flag.
562                 let mut future = notifier.get_future();
563                 let (woken, waker) = create_waker();
564                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Pending);
565
566                 notifier.notify();
567                 assert!(woken.load(Ordering::SeqCst));
568                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(()));
569                 assert!(!notifier.get_future().wait_timeout(Duration::from_millis(1)));
570         }
571
572         #[test]
573         fn test_poll_post_notify_completes() {
574                 // Tests that if we have a future state that has completed, and we haven't yet requested a
575                 // new future, if we get a notify prior to requesting that second future it is generated
576                 // pre-completed.
577                 let notifier = Notifier::new();
578
579                 notifier.notify();
580                 let mut future = notifier.get_future();
581                 let (woken, waker) = create_waker();
582                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(()));
583                 assert!(!woken.load(Ordering::SeqCst));
584
585                 notifier.notify();
586                 let mut future = notifier.get_future();
587                 let (woken, waker) = create_waker();
588                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(()));
589                 assert!(!woken.load(Ordering::SeqCst));
590
591                 let mut future = notifier.get_future();
592                 let (woken, waker) = create_waker();
593                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Pending);
594                 assert!(!woken.load(Ordering::SeqCst));
595
596                 notifier.notify();
597                 assert!(woken.load(Ordering::SeqCst));
598                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(()));
599         }
600
601         #[test]
602         fn test_poll_post_notify_completes_initial_notified() {
603                 // Identical to the previous test, but the first future completes via a wake rather than an
604                 // immediate `Poll::Ready`.
605                 let notifier = Notifier::new();
606
607                 let mut future = notifier.get_future();
608                 let (woken, waker) = create_waker();
609                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Pending);
610
611                 notifier.notify();
612                 assert!(woken.load(Ordering::SeqCst));
613                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(()));
614
615                 notifier.notify();
616                 let mut future = notifier.get_future();
617                 let (woken, waker) = create_waker();
618                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(()));
619                 assert!(!woken.load(Ordering::SeqCst));
620
621                 let mut future = notifier.get_future();
622                 let (woken, waker) = create_waker();
623                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Pending);
624                 assert!(!woken.load(Ordering::SeqCst));
625
626                 notifier.notify();
627                 assert!(woken.load(Ordering::SeqCst));
628                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(()));
629         }
630
631         #[test]
632         #[cfg(feature = "std")]
633         fn test_multi_future_sleep() {
634                 // Tests the `Sleeper` with multiple futures.
635                 let notifier_a = Notifier::new();
636                 let notifier_b = Notifier::new();
637
638                 // Set both notifiers as woken without sleeping yet.
639                 notifier_a.notify();
640                 notifier_b.notify();
641                 Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future()).wait();
642
643                 // One future has woken us up, but the other should still have a pending notification.
644                 Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future()).wait();
645
646                 // However once we've slept twice, we should no longer have any pending notifications
647                 assert!(!Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future())
648                         .wait_timeout(Duration::from_millis(10)));
649
650                 // Test ordering somewhat more.
651                 notifier_a.notify();
652                 Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future()).wait();
653         }
654
655         #[test]
656         #[cfg(feature = "std")]
657         fn sleeper_with_pending_callbacks() {
658                 // This is similar to the above `test_multi_future_sleep` test, but in addition registers
659                 // "normal" callbacks which will cause the futures to assume notification has occurred,
660                 // rather than waiting for a woken sleeper.
661                 let notifier_a = Notifier::new();
662                 let notifier_b = Notifier::new();
663
664                 // Set both notifiers as woken without sleeping yet.
665                 notifier_a.notify();
666                 notifier_b.notify();
667
668                 // After sleeping one future (not guaranteed which one, however) will have its notification
669                 // bit cleared.
670                 Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future()).wait();
671
672                 // By registering a callback on the futures for both notifiers, one will complete
673                 // immediately, but one will remain tied to the notifier, and will complete once the
674                 // notifier is next woken, which will be considered the completion of the notification.
675                 let callback_a = Arc::new(AtomicBool::new(false));
676                 let callback_b = Arc::new(AtomicBool::new(false));
677                 let callback_a_ref = Arc::clone(&callback_a);
678                 let callback_b_ref = Arc::clone(&callback_b);
679                 notifier_a.get_future().register_callback(Box::new(move || assert!(!callback_a_ref.fetch_or(true, Ordering::SeqCst))));
680                 notifier_b.get_future().register_callback(Box::new(move || assert!(!callback_b_ref.fetch_or(true, Ordering::SeqCst))));
681                 assert!(callback_a.load(Ordering::SeqCst) ^ callback_b.load(Ordering::SeqCst));
682
683                 // If we now notify both notifiers again, the other callback will fire, completing the
684                 // notification, and we'll be back to one pending notification.
685                 notifier_a.notify();
686                 notifier_b.notify();
687
688                 assert!(callback_a.load(Ordering::SeqCst) && callback_b.load(Ordering::SeqCst));
689                 Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future()).wait();
690                 assert!(!Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future())
691                         .wait_timeout(Duration::from_millis(10)));
692         }
693 }