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