_ 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 core::time::Duration;
  19 use sync::{Condvar, Mutex};
  20
  21 use prelude::{Box, Vec};
  22
  23 #[cfg(any(test, feature = "std"))]
  24 use std::time::Instant;
  25
  26 use core::future::Future as StdFuture;
  27 use core::task::{Context, Poll};
  28 use core::pin::Pin;
  29
  30 use prelude::*;
  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         condvar: Condvar,
  36 }
  37
  38 impl Notifier {
  39         pub(crate) fn new() -> Self {
  40                 Self {
  41                         notify_pending: Mutex::new((false, None)),
  42                         condvar: Condvar::new(),
  43                 }
  44         }
  45
  46         pub(crate) fn wait(&self) {
  47                 loop {
  48                         let mut guard = self.notify_pending.lock().unwrap();
  49                         if guard.0 {
  50                                 guard.0 = false;
  51                                 return;
  52                         }
  53                         guard = self.condvar.wait(guard).unwrap();
  54                         let result = guard.0;
  55                         if result {
  56                                 guard.0 = false;
  57                                 return
  58                         }
  59                 }
  60         }
  61
  62         #[cfg(any(test, feature = "std"))]
  63         pub(crate) fn wait_timeout(&self, max_wait: Duration) -> bool {
  64                 let current_time = Instant::now();
  65                 loop {
  66                         let mut guard = self.notify_pending.lock().unwrap();
  67                         if guard.0 {
  68                                 guard.0 = false;
  69                                 return true;
  70                         }
  71                         guard = self.condvar.wait_timeout(guard, max_wait).unwrap().0;
  72                         // Due to spurious wakeups that can happen on `wait_timeout`, here we need to check if the
  73                         // desired wait time has actually passed, and if not then restart the loop with a reduced wait
  74                         // time. Note that this logic can be highly simplified through the use of
  75                         // `Condvar::wait_while` and `Condvar::wait_timeout_while`, if and when our MSRV is raised to
  76                         // 1.42.0.
  77                         let elapsed = current_time.elapsed();
  78                         let result = guard.0;
  79                         if result || elapsed >= max_wait {
  80                                 guard.0 = false;
  81                                 return result;
  82                         }
  83                         match max_wait.checked_sub(elapsed) {
  84                                 None => return result,
  85                                 Some(_) => continue
  86                         }
  87                 }
  88         }
  89
  90         /// Wake waiters, tracking that wake needs to occur even if there are currently no waiters.
  91         pub(crate) fn notify(&self) {
  92                 let mut lock = self.notify_pending.lock().unwrap();
  93                 lock.0 = true;
  94                 if let Some(future_state) = lock.1.take() {
  95                         future_state.lock().unwrap().complete();
  96                 }
  97                 mem::drop(lock);
  98                 self.condvar.notify_all();
  99         }
 100
 101         /// Gets a [`Future`] that will get woken up with any waiters
 102         pub(crate) fn get_future(&self) -> Future {
 103                 let mut lock = self.notify_pending.lock().unwrap();
 104                 if lock.0 {
 105                         Future {
 106                                 state: Arc::new(Mutex::new(FutureState {
 107                                         callbacks: Vec::new(),
 108                                         complete: false,
 109                                 }))
 110                         }
 111                 } else if let Some(existing_state) = &lock.1 {
 112                         Future { state: Arc::clone(&existing_state) }
 113                 } else {
 114                         let state = Arc::new(Mutex::new(FutureState {
 115                                 callbacks: Vec::new(),
 116                                 complete: false,
 117                         }));
 118                         lock.1 = Some(Arc::clone(&state));
 119                         Future { state }
 120                 }
 121         }
 122
 123         #[cfg(any(test, feature = "_test_utils"))]
 124         pub fn notify_pending(&self) -> bool {
 125                 self.notify_pending.lock().unwrap().0
 126         }
 127 }
 128
 129 /// A callback which is called when a [`Future`] completes.
 130 ///
 131 /// Note that this MUST NOT call back into LDK directly, it must instead schedule actions to be
 132 /// taken later. Rust users should use the [`std::future::Future`] implementation for [`Future`]
 133 /// instead.
 134 ///
 135 /// Note that the [`std::future::Future`] implementation may only work for runtimes which schedule
 136 /// futures when they receive a wake, rather than immediately executing them.
 137 pub trait FutureCallback : Send {
 138         /// The method which is called.
 139         fn call(&self);
 140 }
 141
 142 impl<F: Fn() + Send> FutureCallback for F {
 143         fn call(&self) { (self)(); }
 144 }
 145
 146 pub(crate) struct FutureState {
 147         callbacks: Vec<Box<dyn FutureCallback>>,
 148         complete: bool,
 149 }
 150
 151 impl FutureState {
 152         fn complete(&mut self) {
 153                 for callback in self.callbacks.drain(..) {
 154                         callback.call();
 155                 }
 156                 self.complete = true;
 157         }
 158 }
 159
 160 /// A simple future which can complete once, and calls some callback(s) when it does so.
 161 pub struct Future {
 162         state: Arc<Mutex<FutureState>>,
 163 }
 164
 165 impl Future {
 166         /// Registers a callback to be called upon completion of this future. If the future has already
 167         /// completed, the callback will be called immediately.
 168         pub fn register_callback(&self, callback: Box<dyn FutureCallback>) {
 169                 let mut state = self.state.lock().unwrap();
 170                 if state.complete {
 171                         mem::drop(state);
 172                         callback.call();
 173                 } else {
 174                         state.callbacks.push(callback);
 175                 }
 176         }
 177 }
 178
 179 mod std_future {
 180         use core::task::Waker;
 181         pub struct StdWaker(pub Waker);
 182         impl super::FutureCallback for StdWaker {
 183                 fn call(&self) { self.0.wake_by_ref() }
 184         }
 185 }
 186
 187 /// (C-not exported) as Rust Futures aren't usable in language bindings.
 188 impl<'a> StdFuture for Future {
 189         type Output = ();
 190
 191         fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
 192                 let mut state = self.state.lock().unwrap();
 193                 if state.complete {
 194                         Poll::Ready(())
 195                 } else {
 196                         let waker = cx.waker().clone();
 197                         state.callbacks.push(Box::new(std_future::StdWaker(waker)));
 198                         Poll::Pending
 199                 }
 200         }
 201 }
 202
 203 #[cfg(test)]
 204 mod tests {
 205         use super::*;
 206         use core::sync::atomic::{AtomicBool, Ordering};
 207         use core::future::Future as FutureTrait;
 208         use core::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
 209
 210         #[cfg(feature = "std")]
 211         #[test]
 212         fn test_wait_timeout() {
 213                 use sync::Arc;
 214                 use std::thread;
 215
 216                 let persistence_notifier = Arc::new(Notifier::new());
 217                 let thread_notifier = Arc::clone(&persistence_notifier);
 218
 219                 let exit_thread = Arc::new(AtomicBool::new(false));
 220                 let exit_thread_clone = exit_thread.clone();
 221                 thread::spawn(move || {
 222                         loop {
 223                                 let mut lock = thread_notifier.notify_pending.lock().unwrap();
 224                                 lock.0 = true;
 225                                 thread_notifier.condvar.notify_all();
 226
 227                                 if exit_thread_clone.load(Ordering::SeqCst) {
 228                                         break
 229                                 }
 230                         }
 231                 });
 232
 233                 // Check that we can block indefinitely until updates are available.
 234                 let _ = persistence_notifier.wait();
 235
 236                 // Check that the Notifier will return after the given duration if updates are
 237                 // available.
 238                 loop {
 239                         if persistence_notifier.wait_timeout(Duration::from_millis(100)) {
 240                                 break
 241                         }
 242                 }
 243
 244                 exit_thread.store(true, Ordering::SeqCst);
 245
 246                 // Check that the Notifier will return after the given duration even if no updates
 247                 // are available.
 248                 loop {
 249                         if !persistence_notifier.wait_timeout(Duration::from_millis(100)) {
 250                                 break
 251                         }
 252                 }
 253         }
 254
 255         #[test]
 256         fn test_future_callbacks() {
 257                 let future = Future {
 258                         state: Arc::new(Mutex::new(FutureState {
 259                                 callbacks: Vec::new(),
 260                                 complete: false,
 261                         }))
 262                 };
 263                 let callback = Arc::new(AtomicBool::new(false));
 264                 let callback_ref = Arc::clone(&callback);
 265                 future.register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
 266
 267                 assert!(!callback.load(Ordering::SeqCst));
 268                 future.state.lock().unwrap().complete();
 269                 assert!(callback.load(Ordering::SeqCst));
 270                 future.state.lock().unwrap().complete();
 271         }
 272
 273         #[test]
 274         fn test_pre_completed_future_callbacks() {
 275                 let future = Future {
 276                         state: Arc::new(Mutex::new(FutureState {
 277                                 callbacks: Vec::new(),
 278                                 complete: false,
 279                         }))
 280                 };
 281                 future.state.lock().unwrap().complete();
 282
 283                 let callback = Arc::new(AtomicBool::new(false));
 284                 let callback_ref = Arc::clone(&callback);
 285                 future.register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
 286
 287                 assert!(callback.load(Ordering::SeqCst));
 288                 assert!(future.state.lock().unwrap().callbacks.is_empty());
 289         }
 290
 291         // Rather annoyingly, there's no safe way in Rust std to construct a Waker despite it being
 292         // totally possible to construct from a trait implementation (though somewhat less effecient
 293         // compared to a raw VTable). Instead, we have to write out a lot of boilerplate to build a
 294         // waker, which we do here with a trivial Arc<AtomicBool> data element to track woke-ness.
 295         const WAKER_V_TABLE: RawWakerVTable = RawWakerVTable::new(waker_clone, wake, wake_by_ref, drop);
 296         unsafe fn wake_by_ref(ptr: *const ()) { let p = ptr as *const Arc<AtomicBool>; assert!(!(*p).fetch_or(true, Ordering::SeqCst)); }
 297         unsafe fn drop(ptr: *const ()) { let p = ptr as *mut Arc<AtomicBool>; Box::from_raw(p); }
 298         unsafe fn wake(ptr: *const ()) { wake_by_ref(ptr); drop(ptr); }
 299         unsafe fn waker_clone(ptr: *const ()) -> RawWaker {
 300                 let p = ptr as *const Arc<AtomicBool>;
 301                 RawWaker::new(Box::into_raw(Box::new(Arc::clone(&*p))) as *const (), &WAKER_V_TABLE)
 302         }
 303
 304         fn create_waker() -> (Arc<AtomicBool>, Waker) {
 305                 let a = Arc::new(AtomicBool::new(false));
 306                 let waker = unsafe { Waker::from_raw(waker_clone((&a as *const Arc<AtomicBool>) as *const ())) };
 307                 (a, waker)
 308         }
 309
 310         #[test]
 311         fn test_future() {
 312                 let mut future = Future {
 313                         state: Arc::new(Mutex::new(FutureState {
 314                                 callbacks: Vec::new(),
 315                                 complete: false,
 316                         }))
 317                 };
 318                 let mut second_future = Future { state: Arc::clone(&future.state) };
 319
 320                 let (woken, waker) = create_waker();
 321                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Pending);
 322                 assert!(!woken.load(Ordering::SeqCst));
 323
 324                 let (second_woken, second_waker) = create_waker();
 325                 assert_eq!(Pin::new(&mut second_future).poll(&mut Context::from_waker(&second_waker)), Poll::Pending);
 326                 assert!(!second_woken.load(Ordering::SeqCst));
 327
 328                 future.state.lock().unwrap().complete();
 329                 assert!(woken.load(Ordering::SeqCst));
 330                 assert!(second_woken.load(Ordering::SeqCst));
 331                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(()));
 332                 assert_eq!(Pin::new(&mut second_future).poll(&mut Context::from_waker(&second_waker)), Poll::Ready(()));
 333         }
 334 }