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