X-Git-Url: http://git.bitcoin.ninja/index.cgi?a=blobdiff_plain;f=lightning%2Fsrc%2Futil%2Fwakers.rs;h=602c2ee04b7d0fb2b85554e8d5566aef8f3fd6d6;hb=ec3aa494953c7d720370119f673e8e6b3a2155d5;hp=fdbc22f116600b7162bdbdcafa4a6f314af944e1;hpb=32fdeb7b4ee12011f9503c7bdc6eb64423e6d232;p=rust-lightning diff --git a/lightning/src/util/wakers.rs b/lightning/src/util/wakers.rs index fdbc22f1..602c2ee0 100644 --- a/lightning/src/util/wakers.rs +++ b/lightning/src/util/wakers.rs @@ -15,12 +15,14 @@ use alloc::sync::Arc; use core::mem; -use crate::sync::{Condvar, Mutex, MutexGuard}; +use crate::sync::Mutex; use crate::prelude::*; -#[cfg(any(test, feature = "std"))] -use std::time::{Duration, Instant}; +#[cfg(feature = "std")] +use crate::sync::Condvar; +#[cfg(feature = "std")] +use std::time::Duration; use core::future::Future as StdFuture; use core::task::{Context, Poll}; @@ -30,74 +32,12 @@ use core::pin::Pin; /// Used to signal to one of many waiters that the condition they're waiting on has happened. pub(crate) struct Notifier { notify_pending: Mutex<(bool, Option>>)>, - condvar: Condvar, -} - -macro_rules! check_woken { - ($guard: expr, $retval: expr) => { { - if $guard.0 { - $guard.0 = false; - if $guard.1.as_ref().map(|l| l.lock().unwrap().complete).unwrap_or(false) { - // If we're about to return as woken, and the future state is marked complete, wipe - // the future state and let the next future wait until we get a new notify. - $guard.1.take(); - } - return $retval; - } - } } } impl Notifier { pub(crate) fn new() -> Self { Self { notify_pending: Mutex::new((false, None)), - condvar: Condvar::new(), - } - } - - fn propagate_future_state_to_notify_flag(&self) -> MutexGuard<(bool, Option>>)> { - let mut lock = self.notify_pending.lock().unwrap(); - if let Some(existing_state) = &lock.1 { - if existing_state.lock().unwrap().callbacks_made { - // If the existing `FutureState` has completed and actually made callbacks, - // consider the notification flag to have been cleared and reset the future state. - lock.1.take(); - lock.0 = false; - } - } - lock - } - - pub(crate) fn wait(&self) { - loop { - let mut guard = self.propagate_future_state_to_notify_flag(); - check_woken!(guard, ()); - guard = self.condvar.wait(guard).unwrap(); - check_woken!(guard, ()); - } - } - - #[cfg(any(test, feature = "std"))] - pub(crate) fn wait_timeout(&self, max_wait: Duration) -> bool { - let current_time = Instant::now(); - loop { - let mut guard = self.propagate_future_state_to_notify_flag(); - check_woken!(guard, true); - guard = self.condvar.wait_timeout(guard, max_wait).unwrap().0; - check_woken!(guard, true); - // Due to spurious wakeups that can happen on `wait_timeout`, here we need to check if the - // desired wait time has actually passed, and if not then restart the loop with a reduced wait - // time. Note that this logic can be highly simplified through the use of - // `Condvar::wait_while` and `Condvar::wait_timeout_while`, if and when our MSRV is raised to - // 1.42.0. - let elapsed = current_time.elapsed(); - if elapsed >= max_wait { - return false; - } - match max_wait.checked_sub(elapsed) { - None => return false, - Some(_) => continue - } } } @@ -105,16 +45,25 @@ impl Notifier { pub(crate) fn notify(&self) { let mut lock = self.notify_pending.lock().unwrap(); if let Some(future_state) = &lock.1 { - future_state.lock().unwrap().complete(); + if future_state.lock().unwrap().complete() { + lock.1 = None; + return; + } } lock.0 = true; - mem::drop(lock); - self.condvar.notify_all(); } /// Gets a [`Future`] that will get woken up with any waiters pub(crate) fn get_future(&self) -> Future { - let mut lock = self.propagate_future_state_to_notify_flag(); + let mut lock = self.notify_pending.lock().unwrap(); + if let Some(existing_state) = &lock.1 { + if existing_state.lock().unwrap().callbacks_made { + // If the existing `FutureState` has completed and actually made callbacks, + // consider the notification flag to have been cleared and reset the future state. + lock.1.take(); + lock.0 = false; + } + } if let Some(existing_state) = &lock.1 { Future { state: Arc::clone(&existing_state) } } else { @@ -134,6 +83,7 @@ impl Notifier { } } +macro_rules! define_callback { ($($bounds: path),*) => { /// A callback which is called when a [`Future`] completes. /// /// Note that this MUST NOT call back into LDK directly, it must instead schedule actions to be @@ -142,14 +92,20 @@ impl Notifier { /// /// Note that the [`std::future::Future`] implementation may only work for runtimes which schedule /// futures when they receive a wake, rather than immediately executing them. -pub trait FutureCallback : Send { +pub trait FutureCallback : $($bounds +)* { /// The method which is called. fn call(&self); } -impl FutureCallback for F { +impl FutureCallback for F { fn call(&self) { (self)(); } } +} } + +#[cfg(feature = "std")] +define_callback!(Send); +#[cfg(not(feature = "std"))] +define_callback!(); pub(crate) struct FutureState { // When we're tracking whether a callback counts as having woken the user's code, we check the @@ -161,16 +117,20 @@ pub(crate) struct FutureState { } impl FutureState { - fn complete(&mut self) { + fn complete(&mut self) -> bool { for (counts_as_call, callback) in self.callbacks.drain(..) { callback.call(); self.callbacks_made |= counts_as_call; } self.complete = true; + self.callbacks_made } } /// A simple future which can complete once, and calls some callback(s) when it does so. +/// +/// Clones can be made and all futures cloned from the same source will complete at the same time. +#[derive(Clone)] pub struct Future { state: Arc>, } @@ -179,7 +139,7 @@ impl Future { /// Registers a callback to be called upon completion of this future. If the future has already /// completed, the callback will be called immediately. /// - /// (C-not exported) use the bindings-only `register_callback_fn` instead + /// This is not exported to bindings users, use the bindings-only `register_callback_fn` instead pub fn register_callback(&self, callback: Box) { let mut state = self.state.lock().unwrap(); if state.complete { @@ -200,6 +160,29 @@ impl Future { pub fn register_callback_fn(&self, callback: F) { self.register_callback(Box::new(callback)); } + + /// Waits until this [`Future`] completes. + #[cfg(feature = "std")] + pub fn wait(self) { + Sleeper::from_single_future(self).wait(); + } + + /// Waits until this [`Future`] completes or the given amount of time has elapsed. + /// + /// Returns true if the [`Future`] completed, false if the time elapsed. + #[cfg(feature = "std")] + pub fn wait_timeout(self, max_wait: Duration) -> bool { + Sleeper::from_single_future(self).wait_timeout(max_wait) + } + + #[cfg(test)] + pub fn poll_is_complete(&self) -> bool { + let mut state = self.state.lock().unwrap(); + if state.complete { + state.callbacks_made = true; + true + } else { false } + } } use core::task::Waker; @@ -208,7 +191,7 @@ impl FutureCallback for StdWaker { fn call(&self) { self.0.wake_by_ref() } } -/// (C-not exported) as Rust Futures aren't usable in language bindings. +/// This is not exported to bindings users as Rust Futures aren't usable in language bindings. impl<'a> StdFuture for Future { type Output = (); @@ -225,6 +208,78 @@ impl<'a> StdFuture for Future { } } +/// A struct which can be used to select across many [`Future`]s at once without relying on a full +/// async context. +#[cfg(feature = "std")] +pub struct Sleeper { + notifiers: Vec>>, +} + +#[cfg(feature = "std")] +impl Sleeper { + /// Constructs a new sleeper from one future, allowing blocking on it. + pub fn from_single_future(future: Future) -> Self { + Self { notifiers: vec![future.state] } + } + /// Constructs a new sleeper from two futures, allowing blocking on both at once. + // Note that this is the common case - a ChannelManager and ChainMonitor. + pub fn from_two_futures(fut_a: Future, fut_b: Future) -> Self { + Self { notifiers: vec![fut_a.state, fut_b.state] } + } + /// Constructs a new sleeper on many futures, allowing blocking on all at once. + pub fn new(futures: Vec) -> Self { + Self { notifiers: futures.into_iter().map(|f| f.state).collect() } + } + /// Prepares to go into a wait loop body, creating a condition variable which we can block on + /// and an `Arc>>` which gets set to the waking `Future`'s state prior to the + /// condition variable being woken. + fn setup_wait(&self) -> (Arc, Arc>>>>) { + let cv = Arc::new(Condvar::new()); + let notified_fut_mtx = Arc::new(Mutex::new(None)); + { + for notifier_mtx in self.notifiers.iter() { + let cv_ref = Arc::clone(&cv); + let notified_fut_ref = Arc::clone(¬ified_fut_mtx); + let notifier_ref = Arc::clone(¬ifier_mtx); + let mut notifier = notifier_mtx.lock().unwrap(); + if notifier.complete { + *notified_fut_mtx.lock().unwrap() = Some(notifier_ref); + break; + } + notifier.callbacks.push((false, Box::new(move || { + *notified_fut_ref.lock().unwrap() = Some(Arc::clone(¬ifier_ref)); + cv_ref.notify_all(); + }))); + } + } + (cv, notified_fut_mtx) + } + + /// Wait until one of the [`Future`]s registered with this [`Sleeper`] has completed. + pub fn wait(&self) { + let (cv, notified_fut_mtx) = self.setup_wait(); + let notified_fut = cv.wait_while(notified_fut_mtx.lock().unwrap(), |fut_opt| fut_opt.is_none()) + .unwrap().take().expect("CV wait shouldn't have returned until the notifying future was set"); + notified_fut.lock().unwrap().callbacks_made = true; + } + + /// Wait until one of the [`Future`]s registered with this [`Sleeper`] has completed or the + /// given amount of time has elapsed. Returns true if a [`Future`] completed, false if the time + /// elapsed. + pub fn wait_timeout(&self, max_wait: Duration) -> bool { + let (cv, notified_fut_mtx) = self.setup_wait(); + let notified_fut = + match cv.wait_timeout_while(notified_fut_mtx.lock().unwrap(), max_wait, |fut_opt| fut_opt.is_none()) { + Ok((_, e)) if e.timed_out() => return false, + Ok((mut notified_fut, _)) => + notified_fut.take().expect("CV wait shouldn't have returned until the notifying future was set"), + Err(_) => panic!("Previous panic while a lock was held led to a lock panic"), + }; + notified_fut.lock().unwrap().callbacks_made = true; + true + } +} + #[cfg(test)] mod tests { use super::*; @@ -323,10 +378,7 @@ mod tests { let exit_thread_clone = exit_thread.clone(); thread::spawn(move || { loop { - let mut lock = thread_notifier.notify_pending.lock().unwrap(); - lock.0 = true; - thread_notifier.condvar.notify_all(); - + thread_notifier.notify(); if exit_thread_clone.load(Ordering::SeqCst) { break } @@ -334,12 +386,12 @@ mod tests { }); // Check that we can block indefinitely until updates are available. - let _ = persistence_notifier.wait(); + let _ = persistence_notifier.get_future().wait(); // Check that the Notifier will return after the given duration if updates are // available. loop { - if persistence_notifier.wait_timeout(Duration::from_millis(100)) { + if persistence_notifier.get_future().wait_timeout(Duration::from_millis(100)) { break } } @@ -349,7 +401,7 @@ mod tests { // Check that the Notifier will return after the given duration even if no updates // are available. loop { - if !persistence_notifier.wait_timeout(Duration::from_millis(100)) { + if !persistence_notifier.get_future().wait_timeout(Duration::from_millis(100)) { break } } @@ -439,6 +491,7 @@ mod tests { } #[test] + #[cfg(feature = "std")] fn test_dropped_future_doesnt_count() { // Tests that if a Future gets drop'd before it is poll()ed `Ready` it doesn't count as // having been woken, leaving the notify-required flag set. @@ -447,8 +500,8 @@ mod tests { // If we get a future and don't touch it we're definitely still notify-required. notifier.get_future(); - assert!(notifier.wait_timeout(Duration::from_millis(1))); - assert!(!notifier.wait_timeout(Duration::from_millis(1))); + assert!(notifier.get_future().wait_timeout(Duration::from_millis(1))); + assert!(!notifier.get_future().wait_timeout(Duration::from_millis(1))); // Even if we poll'd once but didn't observe a `Ready`, we should be notify-required. let mut future = notifier.get_future(); @@ -457,7 +510,7 @@ mod tests { notifier.notify(); assert!(woken.load(Ordering::SeqCst)); - assert!(notifier.wait_timeout(Duration::from_millis(1))); + assert!(notifier.get_future().wait_timeout(Duration::from_millis(1))); // However, once we do poll `Ready` it should wipe the notify-required flag. let mut future = notifier.get_future(); @@ -467,6 +520,128 @@ mod tests { notifier.notify(); assert!(woken.load(Ordering::SeqCst)); assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(())); - assert!(!notifier.wait_timeout(Duration::from_millis(1))); + assert!(!notifier.get_future().wait_timeout(Duration::from_millis(1))); + } + + #[test] + fn test_poll_post_notify_completes() { + // Tests that if we have a future state that has completed, and we haven't yet requested a + // new future, if we get a notify prior to requesting that second future it is generated + // pre-completed. + let notifier = Notifier::new(); + + notifier.notify(); + let mut future = notifier.get_future(); + let (woken, waker) = create_waker(); + assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(())); + assert!(!woken.load(Ordering::SeqCst)); + + notifier.notify(); + let mut future = notifier.get_future(); + let (woken, waker) = create_waker(); + assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(())); + assert!(!woken.load(Ordering::SeqCst)); + + let mut future = notifier.get_future(); + let (woken, waker) = create_waker(); + assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Pending); + assert!(!woken.load(Ordering::SeqCst)); + + notifier.notify(); + assert!(woken.load(Ordering::SeqCst)); + assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(())); + } + + #[test] + fn test_poll_post_notify_completes_initial_notified() { + // Identical to the previous test, but the first future completes via a wake rather than an + // immediate `Poll::Ready`. + let notifier = Notifier::new(); + + let mut future = notifier.get_future(); + let (woken, waker) = create_waker(); + assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Pending); + + notifier.notify(); + assert!(woken.load(Ordering::SeqCst)); + assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(())); + + notifier.notify(); + let mut future = notifier.get_future(); + let (woken, waker) = create_waker(); + assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(())); + assert!(!woken.load(Ordering::SeqCst)); + + let mut future = notifier.get_future(); + let (woken, waker) = create_waker(); + assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Pending); + assert!(!woken.load(Ordering::SeqCst)); + + notifier.notify(); + assert!(woken.load(Ordering::SeqCst)); + assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(())); + } + + #[test] + #[cfg(feature = "std")] + fn test_multi_future_sleep() { + // Tests the `Sleeper` with multiple futures. + let notifier_a = Notifier::new(); + let notifier_b = Notifier::new(); + + // Set both notifiers as woken without sleeping yet. + notifier_a.notify(); + notifier_b.notify(); + Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future()).wait(); + + // One future has woken us up, but the other should still have a pending notification. + Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future()).wait(); + + // However once we've slept twice, we should no longer have any pending notifications + assert!(!Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future()) + .wait_timeout(Duration::from_millis(10))); + + // Test ordering somewhat more. + notifier_a.notify(); + Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future()).wait(); + } + + #[test] + #[cfg(feature = "std")] + fn sleeper_with_pending_callbacks() { + // This is similar to the above `test_multi_future_sleep` test, but in addition registers + // "normal" callbacks which will cause the futures to assume notification has occurred, + // rather than waiting for a woken sleeper. + let notifier_a = Notifier::new(); + let notifier_b = Notifier::new(); + + // Set both notifiers as woken without sleeping yet. + notifier_a.notify(); + notifier_b.notify(); + + // After sleeping one future (not guaranteed which one, however) will have its notification + // bit cleared. + Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future()).wait(); + + // By registering a callback on the futures for both notifiers, one will complete + // immediately, but one will remain tied to the notifier, and will complete once the + // notifier is next woken, which will be considered the completion of the notification. + let callback_a = Arc::new(AtomicBool::new(false)); + let callback_b = Arc::new(AtomicBool::new(false)); + let callback_a_ref = Arc::clone(&callback_a); + let callback_b_ref = Arc::clone(&callback_b); + notifier_a.get_future().register_callback(Box::new(move || assert!(!callback_a_ref.fetch_or(true, Ordering::SeqCst)))); + notifier_b.get_future().register_callback(Box::new(move || assert!(!callback_b_ref.fetch_or(true, Ordering::SeqCst)))); + assert!(callback_a.load(Ordering::SeqCst) ^ callback_b.load(Ordering::SeqCst)); + + // If we now notify both notifiers again, the other callback will fire, completing the + // notification, and we'll be back to one pending notification. + notifier_a.notify(); + notifier_b.notify(); + + assert!(callback_a.load(Ordering::SeqCst) && callback_b.load(Ordering::SeqCst)); + Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future()).wait(); + assert!(!Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future()) + .wait_timeout(Duration::from_millis(10))); } }