X-Git-Url: http://git.bitcoin.ninja/index.cgi?a=blobdiff_plain;f=lightning%2Fsrc%2Futil%2Fwakers.rs;h=37c036da9594747ea81b142e151ff960c2892dad;hb=bedc2c64fcfe5fa5f85ded630e9ed2eb3c3651eb;hp=fdbc22f116600b7162bdbdcafa4a6f314af944e1;hpb=8d93dba37031208878a6046110dc3f6e56796b65;p=rust-lightning diff --git a/lightning/src/util/wakers.rs b/lightning/src/util/wakers.rs index fdbc22f1..37c036da 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,21 +45,31 @@ 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 complete_future(future_state) { + 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 { let state = Arc::new(Mutex::new(FutureState { callbacks: Vec::new(), + callbacks_with_state: Vec::new(), complete: lock.0, callbacks_made: false, })); @@ -134,6 +84,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,35 +93,50 @@ 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 // first bool - set to false if we're just calling a Waker, and true if we're calling an actual // user-provided function. callbacks: Vec<(bool, Box)>, + callbacks_with_state: Vec<(bool, Box>) -> () + Send>)>, complete: bool, callbacks_made: bool, } -impl FutureState { - fn complete(&mut self) { - for (counts_as_call, callback) in self.callbacks.drain(..) { - callback.call(); - self.callbacks_made |= counts_as_call; - } - self.complete = true; +fn complete_future(this: &Arc>) -> bool { + let mut state_lock = this.lock().unwrap(); + let state = &mut *state_lock; + for (counts_as_call, callback) in state.callbacks.drain(..) { + callback.call(); + state.callbacks_made |= counts_as_call; + } + for (counts_as_call, callback) in state.callbacks_with_state.drain(..) { + (callback)(this); + state.callbacks_made |= counts_as_call; } + state.complete = true; + state.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 +145,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 +166,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 +197,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 +214,77 @@ 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 mut notifier = notifier_mtx.lock().unwrap(); + if notifier.complete { + *notified_fut_mtx.lock().unwrap() = Some(Arc::clone(¬ifier_mtx)); + break; + } + notifier.callbacks_with_state.push((false, Box::new(move |notifier_ref| { + *notified_fut_ref.lock().unwrap() = Some(Arc::clone(notifier_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 +383,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 +391,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,17 +406,56 @@ 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 } } } + #[cfg(feature = "std")] + #[test] + fn test_state_drops() { + // Previously, there was a leak if a `Notifier` was `drop`ed without ever being notified + // but after having been slept-on. This tests for that leak. + use crate::sync::Arc; + use std::thread; + + let notifier_a = Arc::new(Notifier::new()); + let notifier_b = Arc::new(Notifier::new()); + + let thread_notifier_a = Arc::clone(¬ifier_a); + + let future_a = notifier_a.get_future(); + let future_state_a = Arc::downgrade(&future_a.state); + + let future_b = notifier_b.get_future(); + let future_state_b = Arc::downgrade(&future_b.state); + + let join_handle = thread::spawn(move || { + // Let the other thread get to the wait point, then notify it. + std::thread::sleep(Duration::from_millis(50)); + thread_notifier_a.notify(); + }); + + // Wait on the other thread to finish its sleep, note that the leak only happened if we + // actually have to sleep here, not if we immediately return. + Sleeper::from_two_futures(future_a, future_b).wait(); + + join_handle.join().unwrap(); + + // then drop the notifiers and make sure the future states are gone. + mem::drop(notifier_a); + mem::drop(notifier_b); + + assert!(future_state_a.upgrade().is_none() && future_state_b.upgrade().is_none()); + } + #[test] fn test_future_callbacks() { let future = Future { state: Arc::new(Mutex::new(FutureState { callbacks: Vec::new(), + callbacks_with_state: Vec::new(), complete: false, callbacks_made: false, })) @@ -369,9 +465,9 @@ mod tests { future.register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst)))); assert!(!callback.load(Ordering::SeqCst)); - future.state.lock().unwrap().complete(); + complete_future(&future.state); assert!(callback.load(Ordering::SeqCst)); - future.state.lock().unwrap().complete(); + complete_future(&future.state); } #[test] @@ -379,11 +475,12 @@ mod tests { let future = Future { state: Arc::new(Mutex::new(FutureState { callbacks: Vec::new(), + callbacks_with_state: Vec::new(), complete: false, callbacks_made: false, })) }; - future.state.lock().unwrap().complete(); + complete_future(&future.state); let callback = Arc::new(AtomicBool::new(false)); let callback_ref = Arc::clone(&callback); @@ -417,6 +514,7 @@ mod tests { let mut future = Future { state: Arc::new(Mutex::new(FutureState { callbacks: Vec::new(), + callbacks_with_state: Vec::new(), complete: false, callbacks_made: false, })) @@ -431,7 +529,7 @@ mod tests { assert_eq!(Pin::new(&mut second_future).poll(&mut Context::from_waker(&second_waker)), Poll::Pending); assert!(!second_woken.load(Ordering::SeqCst)); - future.state.lock().unwrap().complete(); + complete_future(&future.state); assert!(woken.load(Ordering::SeqCst)); assert!(second_woken.load(Ordering::SeqCst)); assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(())); @@ -439,6 +537,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 +546,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 +556,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 +566,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))); } }