_ Git - rust-lightning/blob - 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 use core::task::Waker;
 206 struct StdWaker(pub Waker);
 207
 208 /// This is not exported to bindings users as Rust Futures aren't usable in language bindings.
 209 impl<'a> StdFuture for Future {
 210         type Output = ();
 211
 212         fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
 213                 let mut state = self.state.lock().unwrap();
 214                 if state.complete {
 215                         state.callbacks_made = true;
 216                         Poll::Ready(())
 217                 } else {
 218                         let waker = cx.waker().clone();
 219                         state.std_future_callbacks.push((self.self_idx, StdWaker(waker)));
 220                         Poll::Pending
 221                 }
 222         }
 223 }
 224
 225 /// A struct which can be used to select across many [`Future`]s at once without relying on a full
 226 /// async context.
 227 #[cfg(feature = "std")]
 228 pub struct Sleeper {
 229         notifiers: Vec<Arc<Mutex<FutureState>>>,
 230 }
 231
 232 #[cfg(feature = "std")]
 233 impl Sleeper {
 234         /// Constructs a new sleeper from one future, allowing blocking on it.
 235         pub fn from_single_future(future: Future) -> Self {
 236                 Self { notifiers: vec![future.state] }
 237         }
 238         /// Constructs a new sleeper from two futures, allowing blocking on both at once.
 239         // Note that this is the common case - a ChannelManager and ChainMonitor.
 240         pub fn from_two_futures(fut_a: Future, fut_b: Future) -> Self {
 241                 Self { notifiers: vec![fut_a.state, fut_b.state] }
 242         }
 243         /// Constructs a new sleeper on many futures, allowing blocking on all at once.
 244         pub fn new(futures: Vec<Future>) -> Self {
 245                 Self { notifiers: futures.into_iter().map(|f| f.state).collect() }
 246         }
 247         /// Prepares to go into a wait loop body, creating a condition variable which we can block on
 248         /// and an `Arc<Mutex<Option<_>>>` which gets set to the waking `Future`'s state prior to the
 249         /// condition variable being woken.
 250         fn setup_wait(&self) -> (Arc<Condvar>, Arc<Mutex<Option<Arc<Mutex<FutureState>>>>>) {
 251                 let cv = Arc::new(Condvar::new());
 252                 let notified_fut_mtx = Arc::new(Mutex::new(None));
 253                 {
 254                         for notifier_mtx in self.notifiers.iter() {
 255                                 let cv_ref = Arc::clone(&cv);
 256                                 let notified_fut_ref = Arc::clone(&notified_fut_mtx);
 257                                 let mut notifier = notifier_mtx.lock().unwrap();
 258                                 if notifier.complete {
 259                                         *notified_fut_mtx.lock().unwrap() = Some(Arc::clone(&notifier_mtx));
 260                                         break;
 261                                 }
 262                                 notifier.callbacks_with_state.push(Box::new(move |notifier_ref| {
 263                                         *notified_fut_ref.lock().unwrap() = Some(Arc::clone(notifier_ref));
 264                                         cv_ref.notify_all();
 265                                 }));
 266                         }
 267                 }
 268                 (cv, notified_fut_mtx)
 269         }
 270
 271         /// Wait until one of the [`Future`]s registered with this [`Sleeper`] has completed.
 272         pub fn wait(&self) {
 273                 let (cv, notified_fut_mtx) = self.setup_wait();
 274                 let notified_fut = cv.wait_while(notified_fut_mtx.lock().unwrap(), |fut_opt| fut_opt.is_none())
 275                         .unwrap().take().expect("CV wait shouldn't have returned until the notifying future was set");
 276                 notified_fut.lock().unwrap().callbacks_made = true;
 277         }
 278
 279         /// Wait until one of the [`Future`]s registered with this [`Sleeper`] has completed or the
 280         /// given amount of time has elapsed. Returns true if a [`Future`] completed, false if the time
 281         /// elapsed.
 282         pub fn wait_timeout(&self, max_wait: Duration) -> bool {
 283                 let (cv, notified_fut_mtx) = self.setup_wait();
 284                 let notified_fut =
 285                         match cv.wait_timeout_while(notified_fut_mtx.lock().unwrap(), max_wait, |fut_opt| fut_opt.is_none()) {
 286                                 Ok((_, e)) if e.timed_out() => return false,
 287                                 Ok((mut notified_fut, _)) =>
 288                                         notified_fut.take().expect("CV wait shouldn't have returned until the notifying future was set"),
 289                                 Err(_) => panic!("Previous panic while a lock was held led to a lock panic"),
 290                         };
 291                 notified_fut.lock().unwrap().callbacks_made = true;
 292                 true
 293         }
 294 }
 295
 296 #[cfg(test)]
 297 mod tests {
 298         use super::*;
 299         use core::sync::atomic::{AtomicBool, Ordering};
 300         use core::future::Future as FutureTrait;
 301         use core::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
 302
 303         #[test]
 304         fn notifier_pre_notified_future() {
 305                 // Previously, if we generated a future after a `Notifier` had been notified, the future
 306                 // would never complete. This tests this behavior, ensuring the future instead completes
 307                 // immediately.
 308                 let notifier = Notifier::new();
 309                 notifier.notify();
 310
 311                 let callback = Arc::new(AtomicBool::new(false));
 312                 let callback_ref = Arc::clone(&callback);
 313                 notifier.get_future().register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
 314                 assert!(callback.load(Ordering::SeqCst));
 315         }
 316
 317         #[test]
 318         fn notifier_future_completes_wake() {
 319                 // Previously, if we were only using the `Future` interface to learn when a `Notifier` has
 320                 // been notified, we'd never mark the notifier as not-awaiting-notify. This caused the
 321                 // `lightning-background-processor` to persist in a tight loop.
 322                 let notifier = Notifier::new();
 323
 324                 // First check the simple case, ensuring if we get notified a new future isn't woken until
 325                 // a second `notify`.
 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                 let callback = Arc::new(AtomicBool::new(false));
 335                 let callback_ref = Arc::clone(&callback);
 336                 notifier.get_future().register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
 337                 assert!(!callback.load(Ordering::SeqCst));
 338
 339                 notifier.notify();
 340                 assert!(callback.load(Ordering::SeqCst));
 341
 342                 // Then check the case where the future is fetched before the notification, but a callback
 343                 // is only registered after the `notify`, ensuring that it is still sufficient to ensure we
 344                 // don't get an instant-wake when we get a new future.
 345                 let future = notifier.get_future();
 346                 notifier.notify();
 347
 348                 let callback = Arc::new(AtomicBool::new(false));
 349                 let callback_ref = Arc::clone(&callback);
 350                 future.register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
 351                 assert!(callback.load(Ordering::SeqCst));
 352
 353                 let callback = Arc::new(AtomicBool::new(false));
 354                 let callback_ref = Arc::clone(&callback);
 355                 notifier.get_future().register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
 356                 assert!(!callback.load(Ordering::SeqCst));
 357         }
 358
 359         #[test]
 360         fn new_future_wipes_notify_bit() {
 361                 // Previously, if we were only using the `Future` interface to learn when a `Notifier` has
 362                 // been notified, we'd never mark the notifier as not-awaiting-notify if a `Future` is
 363                 // fetched after the notify bit has been set.
 364                 let notifier = Notifier::new();
 365                 notifier.notify();
 366
 367                 let callback = Arc::new(AtomicBool::new(false));
 368                 let callback_ref = Arc::clone(&callback);
 369                 notifier.get_future().register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
 370                 assert!(callback.load(Ordering::SeqCst));
 371
 372                 let callback = Arc::new(AtomicBool::new(false));
 373                 let callback_ref = Arc::clone(&callback);
 374                 notifier.get_future().register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
 375                 assert!(!callback.load(Ordering::SeqCst));
 376
 377                 notifier.notify();
 378                 assert!(callback.load(Ordering::SeqCst));
 379         }
 380
 381         #[cfg(feature = "std")]
 382         #[test]
 383         fn test_wait_timeout() {
 384                 use crate::sync::Arc;
 385                 use std::thread;
 386
 387                 let persistence_notifier = Arc::new(Notifier::new());
 388                 let thread_notifier = Arc::clone(&persistence_notifier);
 389
 390                 let exit_thread = Arc::new(AtomicBool::new(false));
 391                 let exit_thread_clone = exit_thread.clone();
 392                 thread::spawn(move || {
 393                         loop {
 394                                 thread_notifier.notify();
 395                                 if exit_thread_clone.load(Ordering::SeqCst) {
 396                                         break
 397                                 }
 398                         }
 399                 });
 400
 401                 // Check that we can block indefinitely until updates are available.
 402                 let _ = persistence_notifier.get_future().wait();
 403
 404                 // Check that the Notifier will return after the given duration if updates are
 405                 // available.
 406                 loop {
 407                         if persistence_notifier.get_future().wait_timeout(Duration::from_millis(100)) {
 408                                 break
 409                         }
 410                 }
 411
 412                 exit_thread.store(true, Ordering::SeqCst);
 413
 414                 // Check that the Notifier will return after the given duration even if no updates
 415                 // are available.
 416                 loop {
 417                         if !persistence_notifier.get_future().wait_timeout(Duration::from_millis(100)) {
 418                                 break
 419                         }
 420                 }
 421         }
 422
 423         #[cfg(feature = "std")]
 424         #[test]
 425         fn test_state_drops() {
 426                 // Previously, there was a leak if a `Notifier` was `drop`ed without ever being notified
 427                 // but after having been slept-on. This tests for that leak.
 428                 use crate::sync::Arc;
 429                 use std::thread;
 430
 431                 let notifier_a = Arc::new(Notifier::new());
 432                 let notifier_b = Arc::new(Notifier::new());
 433
 434                 let thread_notifier_a = Arc::clone(&notifier_a);
 435
 436                 let future_a = notifier_a.get_future();
 437                 let future_state_a = Arc::downgrade(&future_a.state);
 438
 439                 let future_b = notifier_b.get_future();
 440                 let future_state_b = Arc::downgrade(&future_b.state);
 441
 442                 let join_handle = thread::spawn(move || {
 443                         // Let the other thread get to the wait point, then notify it.
 444                         std::thread::sleep(Duration::from_millis(50));
 445                         thread_notifier_a.notify();
 446                 });
 447
 448                 // Wait on the other thread to finish its sleep, note that the leak only happened if we
 449                 // actually have to sleep here, not if we immediately return.
 450                 Sleeper::from_two_futures(future_a, future_b).wait();
 451
 452                 join_handle.join().unwrap();
 453
 454                 // then drop the notifiers and make sure the future states are gone.
 455                 mem::drop(notifier_a);
 456                 mem::drop(notifier_b);
 457
 458                 assert!(future_state_a.upgrade().is_none() && future_state_b.upgrade().is_none());
 459         }
 460
 461         #[test]
 462         fn test_future_callbacks() {
 463                 let future = Future {
 464                         state: Arc::new(Mutex::new(FutureState {
 465                                 callbacks: Vec::new(),
 466                                 std_future_callbacks: Vec::new(),
 467                                 callbacks_with_state: Vec::new(),
 468                                 complete: false,
 469                                 callbacks_made: false,
 470                                 next_idx: 1,
 471                         })),
 472                         self_idx: 0,
 473                 };
 474                 let callback = Arc::new(AtomicBool::new(false));
 475                 let callback_ref = Arc::clone(&callback);
 476                 future.register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
 477
 478                 assert!(!callback.load(Ordering::SeqCst));
 479                 complete_future(&future.state);
 480                 assert!(callback.load(Ordering::SeqCst));
 481                 complete_future(&future.state);
 482         }
 483
 484         #[test]
 485         fn test_pre_completed_future_callbacks() {
 486                 let future = Future {
 487                         state: Arc::new(Mutex::new(FutureState {
 488                                 callbacks: Vec::new(),
 489                                 std_future_callbacks: Vec::new(),
 490                                 callbacks_with_state: Vec::new(),
 491                                 complete: false,
 492                                 callbacks_made: false,
 493                                 next_idx: 1,
 494                         })),
 495                         self_idx: 0,
 496                 };
 497                 complete_future(&future.state);
 498
 499                 let callback = Arc::new(AtomicBool::new(false));
 500                 let callback_ref = Arc::clone(&callback);
 501                 future.register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
 502
 503                 assert!(callback.load(Ordering::SeqCst));
 504                 assert!(future.state.lock().unwrap().callbacks.is_empty());
 505         }
 506
 507         // Rather annoyingly, there's no safe way in Rust std to construct a Waker despite it being
 508         // totally possible to construct from a trait implementation (though somewhat less efficient
 509         // compared to a raw VTable). Instead, we have to write out a lot of boilerplate to build a
 510         // waker, which we do here with a trivial Arc<AtomicBool> data element to track woke-ness.
 511         const WAKER_V_TABLE: RawWakerVTable = RawWakerVTable::new(waker_clone, wake, wake_by_ref, drop);
 512         unsafe fn wake_by_ref(ptr: *const ()) { let p = ptr as *const Arc<AtomicBool>; assert!(!(*p).fetch_or(true, Ordering::SeqCst)); }
 513         unsafe fn drop(ptr: *const ()) { let p = ptr as *mut Arc<AtomicBool>; let _freed = Box::from_raw(p); }
 514         unsafe fn wake(ptr: *const ()) { wake_by_ref(ptr); drop(ptr); }
 515         unsafe fn waker_clone(ptr: *const ()) -> RawWaker {
 516                 let p = ptr as *const Arc<AtomicBool>;
 517                 RawWaker::new(Box::into_raw(Box::new(Arc::clone(&*p))) as *const (), &WAKER_V_TABLE)
 518         }
 519
 520         fn create_waker() -> (Arc<AtomicBool>, Waker) {
 521                 let a = Arc::new(AtomicBool::new(false));
 522                 let waker = unsafe { Waker::from_raw(waker_clone((&a as *const Arc<AtomicBool>) as *const ())) };
 523                 (a, waker)
 524         }
 525
 526         #[test]
 527         fn test_future() {
 528                 let mut future = Future {
 529                         state: Arc::new(Mutex::new(FutureState {
 530                                 callbacks: Vec::new(),
 531                                 std_future_callbacks: Vec::new(),
 532                                 callbacks_with_state: Vec::new(),
 533                                 complete: false,
 534                                 callbacks_made: false,
 535                                 next_idx: 2,
 536                         })),
 537                         self_idx: 0,
 538                 };
 539                 let mut second_future = Future { state: Arc::clone(&future.state), self_idx: 1 };
 540
 541                 let (woken, waker) = create_waker();
 542                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Pending);
 543                 assert!(!woken.load(Ordering::SeqCst));
 544
 545                 let (second_woken, second_waker) = create_waker();
 546                 assert_eq!(Pin::new(&mut second_future).poll(&mut Context::from_waker(&second_waker)), Poll::Pending);
 547                 assert!(!second_woken.load(Ordering::SeqCst));
 548
 549                 complete_future(&future.state);
 550                 assert!(woken.load(Ordering::SeqCst));
 551                 assert!(second_woken.load(Ordering::SeqCst));
 552                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(()));
 553                 assert_eq!(Pin::new(&mut second_future).poll(&mut Context::from_waker(&second_waker)), Poll::Ready(()));
 554         }
 555
 556         #[test]
 557         #[cfg(feature = "std")]
 558         fn test_dropped_future_doesnt_count() {
 559                 // Tests that if a Future gets drop'd before it is poll()ed `Ready` it doesn't count as
 560                 // having been woken, leaving the notify-required flag set.
 561                 let notifier = Notifier::new();
 562                 notifier.notify();
 563
 564                 // If we get a future and don't touch it we're definitely still notify-required.
 565                 notifier.get_future();
 566                 assert!(notifier.get_future().wait_timeout(Duration::from_millis(1)));
 567                 assert!(!notifier.get_future().wait_timeout(Duration::from_millis(1)));
 568
 569                 // Even if we poll'd once but didn't observe a `Ready`, we should be notify-required.
 570                 let mut future = notifier.get_future();
 571                 let (woken, waker) = create_waker();
 572                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Pending);
 573
 574                 notifier.notify();
 575                 assert!(woken.load(Ordering::SeqCst));
 576                 assert!(notifier.get_future().wait_timeout(Duration::from_millis(1)));
 577
 578                 // However, once we do poll `Ready` it should wipe the notify-required flag.
 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_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(()));
 586                 assert!(!notifier.get_future().wait_timeout(Duration::from_millis(1)));
 587         }
 588
 589         #[test]
 590         fn test_poll_post_notify_completes() {
 591                 // Tests that if we have a future state that has completed, and we haven't yet requested a
 592                 // new future, if we get a notify prior to requesting that second future it is generated
 593                 // pre-completed.
 594                 let notifier = Notifier::new();
 595
 596                 notifier.notify();
 597                 let mut future = notifier.get_future();
 598                 let (woken, waker) = create_waker();
 599                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(()));
 600                 assert!(!woken.load(Ordering::SeqCst));
 601
 602                 notifier.notify();
 603                 let mut future = notifier.get_future();
 604                 let (woken, waker) = create_waker();
 605                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(()));
 606                 assert!(!woken.load(Ordering::SeqCst));
 607
 608                 let mut future = notifier.get_future();
 609                 let (woken, waker) = create_waker();
 610                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Pending);
 611                 assert!(!woken.load(Ordering::SeqCst));
 612
 613                 notifier.notify();
 614                 assert!(woken.load(Ordering::SeqCst));
 615                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(()));
 616         }
 617
 618         #[test]
 619         fn test_poll_post_notify_completes_initial_notified() {
 620                 // Identical to the previous test, but the first future completes via a wake rather than an
 621                 // immediate `Poll::Ready`.
 622                 let notifier = Notifier::new();
 623
 624                 let mut future = notifier.get_future();
 625                 let (woken, waker) = create_waker();
 626                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Pending);
 627
 628                 notifier.notify();
 629                 assert!(woken.load(Ordering::SeqCst));
 630                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(()));
 631
 632                 notifier.notify();
 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::Ready(()));
 636                 assert!(!woken.load(Ordering::SeqCst));
 637
 638                 let mut future = notifier.get_future();
 639                 let (woken, waker) = create_waker();
 640                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Pending);
 641                 assert!(!woken.load(Ordering::SeqCst));
 642
 643                 notifier.notify();
 644                 assert!(woken.load(Ordering::SeqCst));
 645                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(()));
 646         }
 647
 648         #[test]
 649         #[cfg(feature = "std")]
 650         fn test_multi_future_sleep() {
 651                 // Tests the `Sleeper` with multiple futures.
 652                 let notifier_a = Notifier::new();
 653                 let notifier_b = Notifier::new();
 654
 655                 // Set both notifiers as woken without sleeping yet.
 656                 notifier_a.notify();
 657                 notifier_b.notify();
 658                 Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future()).wait();
 659
 660                 // One future has woken us up, but the other should still have a pending notification.
 661                 Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future()).wait();
 662
 663                 // However once we've slept twice, we should no longer have any pending notifications
 664                 assert!(!Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future())
 665                         .wait_timeout(Duration::from_millis(10)));
 666
 667                 // Test ordering somewhat more.
 668                 notifier_a.notify();
 669                 Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future()).wait();
 670         }
 671
 672         #[test]
 673         #[cfg(feature = "std")]
 674         fn sleeper_with_pending_callbacks() {
 675                 // This is similar to the above `test_multi_future_sleep` test, but in addition registers
 676                 // "normal" callbacks which will cause the futures to assume notification has occurred,
 677                 // rather than waiting for a woken sleeper.
 678                 let notifier_a = Notifier::new();
 679                 let notifier_b = Notifier::new();
 680
 681                 // Set both notifiers as woken without sleeping yet.
 682                 notifier_a.notify();
 683                 notifier_b.notify();
 684
 685                 // After sleeping one future (not guaranteed which one, however) will have its notification
 686                 // bit cleared.
 687                 Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future()).wait();
 688
 689                 // By registering a callback on the futures for both notifiers, one will complete
 690                 // immediately, but one will remain tied to the notifier, and will complete once the
 691                 // notifier is next woken, which will be considered the completion of the notification.
 692                 let callback_a = Arc::new(AtomicBool::new(false));
 693                 let callback_b = Arc::new(AtomicBool::new(false));
 694                 let callback_a_ref = Arc::clone(&callback_a);
 695                 let callback_b_ref = Arc::clone(&callback_b);
 696                 notifier_a.get_future().register_callback(Box::new(move || assert!(!callback_a_ref.fetch_or(true, Ordering::SeqCst))));
 697                 notifier_b.get_future().register_callback(Box::new(move || assert!(!callback_b_ref.fetch_or(true, Ordering::SeqCst))));
 698                 assert!(callback_a.load(Ordering::SeqCst) ^ callback_b.load(Ordering::SeqCst));
 699
 700                 // If we now notify both notifiers again, the other callback will fire, completing the
 701                 // notification, and we'll be back to one pending notification.
 702                 notifier_a.notify();
 703                 notifier_b.notify();
 704
 705                 assert!(callback_a.load(Ordering::SeqCst) && callback_b.load(Ordering::SeqCst));
 706                 Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future()).wait();
 707                 assert!(!Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future())
 708                         .wait_timeout(Duration::from_millis(10)));
 709         }
 710 }