_ 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::{Condvar, Mutex, MutexGuard};
  19
  20 use crate::prelude::*;
  21
  22 #[cfg(any(test, feature = "std"))]
  23 use std::time::{Duration, Instant};
  24
  25 use core::future::Future as StdFuture;
  26 use core::task::{Context, Poll};
  27 use core::pin::Pin;
  28
  29
  30 /// Used to signal to one of many waiters that the condition they're waiting on has happened.
  31 pub(crate) struct Notifier {
  32         notify_pending: Mutex<(bool, Option<Arc<Mutex<FutureState>>>)>,
  33         condvar: Condvar,
  34 }
  35
  36 impl Notifier {
  37         pub(crate) fn new() -> Self {
  38                 Self {
  39                         notify_pending: Mutex::new((false, None)),
  40                         condvar: Condvar::new(),
  41                 }
  42         }
  43
  44         fn propagate_future_state_to_notify_flag(&self) -> MutexGuard<(bool, Option<Arc<Mutex<FutureState>>>)> {
  45                 let mut lock = self.notify_pending.lock().unwrap();
  46                 if let Some(existing_state) = &lock.1 {
  47                         if existing_state.lock().unwrap().callbacks_made {
  48                                 // If the existing `FutureState` has completed and actually made callbacks,
  49                                 // consider the notification flag to have been cleared and reset the future state.
  50                                 lock.1.take();
  51                                 lock.0 = false;
  52                         }
  53                 }
  54                 lock
  55         }
  56
  57         pub(crate) fn wait(&self) {
  58                 loop {
  59                         let mut guard = self.propagate_future_state_to_notify_flag();
  60                         if guard.0 {
  61                                 guard.0 = false;
  62                                 return;
  63                         }
  64                         guard = self.condvar.wait(guard).unwrap();
  65                         let result = guard.0;
  66                         if result {
  67                                 guard.0 = false;
  68                                 return
  69                         }
  70                 }
  71         }
  72
  73         #[cfg(any(test, feature = "std"))]
  74         pub(crate) fn wait_timeout(&self, max_wait: Duration) -> bool {
  75                 let current_time = Instant::now();
  76                 loop {
  77                         let mut guard = self.propagate_future_state_to_notify_flag();
  78                         if guard.0 {
  79                                 guard.0 = false;
  80                                 return true;
  81                         }
  82                         guard = self.condvar.wait_timeout(guard, max_wait).unwrap().0;
  83                         // Due to spurious wakeups that can happen on `wait_timeout`, here we need to check if the
  84                         // desired wait time has actually passed, and if not then restart the loop with a reduced wait
  85                         // time. Note that this logic can be highly simplified through the use of
  86                         // `Condvar::wait_while` and `Condvar::wait_timeout_while`, if and when our MSRV is raised to
  87                         // 1.42.0.
  88                         let elapsed = current_time.elapsed();
  89                         let result = guard.0;
  90                         if result || elapsed >= max_wait {
  91                                 guard.0 = false;
  92                                 return result;
  93                         }
  94                         match max_wait.checked_sub(elapsed) {
  95                                 None => return result,
  96                                 Some(_) => continue
  97                         }
  98                 }
  99         }
 100
 101         /// Wake waiters, tracking that wake needs to occur even if there are currently no waiters.
 102         pub(crate) fn notify(&self) {
 103                 let mut lock = self.notify_pending.lock().unwrap();
 104                 if let Some(future_state) = &lock.1 {
 105                         future_state.lock().unwrap().complete();
 106                 }
 107                 lock.0 = true;
 108                 mem::drop(lock);
 109                 self.condvar.notify_all();
 110         }
 111
 112         /// Gets a [`Future`] that will get woken up with any waiters
 113         pub(crate) fn get_future(&self) -> Future {
 114                 let mut lock = self.propagate_future_state_to_notify_flag();
 115                 if let Some(existing_state) = &lock.1 {
 116                         Future { state: Arc::clone(&existing_state) }
 117                 } else {
 118                         let state = Arc::new(Mutex::new(FutureState {
 119                                 callbacks: Vec::new(),
 120                                 complete: lock.0,
 121                                 callbacks_made: false,
 122                         }));
 123                         lock.1 = Some(Arc::clone(&state));
 124                         Future { state }
 125                 }
 126         }
 127
 128         #[cfg(any(test, feature = "_test_utils"))]
 129         pub fn notify_pending(&self) -> bool {
 130                 self.notify_pending.lock().unwrap().0
 131         }
 132 }
 133
 134 /// A callback which is called when a [`Future`] completes.
 135 ///
 136 /// Note that this MUST NOT call back into LDK directly, it must instead schedule actions to be
 137 /// taken later. Rust users should use the [`std::future::Future`] implementation for [`Future`]
 138 /// instead.
 139 ///
 140 /// Note that the [`std::future::Future`] implementation may only work for runtimes which schedule
 141 /// futures when they receive a wake, rather than immediately executing them.
 142 pub trait FutureCallback : Send {
 143         /// The method which is called.
 144         fn call(&self);
 145 }
 146
 147 impl<F: Fn() + Send> FutureCallback for F {
 148         fn call(&self) { (self)(); }
 149 }
 150
 151 pub(crate) struct FutureState {
 152         callbacks: Vec<Box<dyn FutureCallback>>,
 153         complete: bool,
 154         callbacks_made: bool,
 155 }
 156
 157 impl FutureState {
 158         fn complete(&mut self) {
 159                 for callback in self.callbacks.drain(..) {
 160                         callback.call();
 161                         self.callbacks_made = true;
 162                 }
 163                 self.complete = true;
 164         }
 165 }
 166
 167 /// A simple future which can complete once, and calls some callback(s) when it does so.
 168 pub struct Future {
 169         state: Arc<Mutex<FutureState>>,
 170 }
 171
 172 impl Future {
 173         /// Registers a callback to be called upon completion of this future. If the future has already
 174         /// completed, the callback will be called immediately.
 175         ///
 176         /// (C-not exported) use the bindings-only `register_callback_fn` instead
 177         pub fn register_callback(&self, callback: Box<dyn FutureCallback>) {
 178                 let mut state = self.state.lock().unwrap();
 179                 if state.complete {
 180                         state.callbacks_made = true;
 181                         mem::drop(state);
 182                         callback.call();
 183                 } else {
 184                         state.callbacks.push(callback);
 185                 }
 186         }
 187
 188         // C bindings don't (currently) know how to map `Box<dyn Trait>`, and while it could add the
 189         // following wrapper, doing it in the bindings is currently much more work than simply doing it
 190         // here.
 191         /// Registers a callback to be called upon completion of this future. If the future has already
 192         /// completed, the callback will be called immediately.
 193         #[cfg(c_bindings)]
 194         pub fn register_callback_fn<F: 'static + FutureCallback>(&self, callback: F) {
 195                 self.register_callback(Box::new(callback));
 196         }
 197 }
 198
 199 use core::task::Waker;
 200 struct StdWaker(pub Waker);
 201 impl FutureCallback for StdWaker {
 202         fn call(&self) { self.0.wake_by_ref() }
 203 }
 204
 205 /// (C-not exported) as Rust Futures aren't usable in language bindings.
 206 impl<'a> StdFuture for Future {
 207         type Output = ();
 208
 209         fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
 210                 let mut state = self.state.lock().unwrap();
 211                 if state.complete {
 212                         Poll::Ready(())
 213                 } else {
 214                         let waker = cx.waker().clone();
 215                         state.callbacks.push(Box::new(StdWaker(waker)));
 216                         Poll::Pending
 217                 }
 218         }
 219 }
 220
 221 #[cfg(test)]
 222 mod tests {
 223         use super::*;
 224         use core::sync::atomic::{AtomicBool, Ordering};
 225         use core::future::Future as FutureTrait;
 226         use core::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
 227
 228         #[test]
 229         fn notifier_pre_notified_future() {
 230                 // Previously, if we generated a future after a `Notifier` had been notified, the future
 231                 // would never complete. This tests this behavior, ensuring the future instead completes
 232                 // immediately.
 233                 let notifier = Notifier::new();
 234                 notifier.notify();
 235
 236                 let callback = Arc::new(AtomicBool::new(false));
 237                 let callback_ref = Arc::clone(&callback);
 238                 notifier.get_future().register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
 239                 assert!(callback.load(Ordering::SeqCst));
 240         }
 241
 242         #[test]
 243         fn notifier_future_completes_wake() {
 244                 // Previously, if we were only using the `Future` interface to learn when a `Notifier` has
 245                 // been notified, we'd never mark the notifier as not-awaiting-notify. This caused the
 246                 // `lightning-background-processor` to persist in a tight loop.
 247                 let notifier = Notifier::new();
 248
 249                 // First check the simple case, ensuring if we get notified a new future isn't woken until
 250                 // a second `notify`.
 251                 let callback = Arc::new(AtomicBool::new(false));
 252                 let callback_ref = Arc::clone(&callback);
 253                 notifier.get_future().register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
 254                 assert!(!callback.load(Ordering::SeqCst));
 255
 256                 notifier.notify();
 257                 assert!(callback.load(Ordering::SeqCst));
 258
 259                 let callback = Arc::new(AtomicBool::new(false));
 260                 let callback_ref = Arc::clone(&callback);
 261                 notifier.get_future().register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
 262                 assert!(!callback.load(Ordering::SeqCst));
 263
 264                 notifier.notify();
 265                 assert!(callback.load(Ordering::SeqCst));
 266
 267                 // Then check the case where the future is fetched before the notification, but a callback
 268                 // is only registered after the `notify`, ensuring that it is still sufficient to ensure we
 269                 // don't get an instant-wake when we get a new future.
 270                 let future = notifier.get_future();
 271                 notifier.notify();
 272
 273                 let callback = Arc::new(AtomicBool::new(false));
 274                 let callback_ref = Arc::clone(&callback);
 275                 future.register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
 276                 assert!(callback.load(Ordering::SeqCst));
 277
 278                 let callback = Arc::new(AtomicBool::new(false));
 279                 let callback_ref = Arc::clone(&callback);
 280                 notifier.get_future().register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
 281                 assert!(!callback.load(Ordering::SeqCst));
 282         }
 283
 284         #[test]
 285         fn new_future_wipes_notify_bit() {
 286                 // Previously, if we were only using the `Future` interface to learn when a `Notifier` has
 287                 // been notified, we'd never mark the notifier as not-awaiting-notify if a `Future` is
 288                 // fetched after the notify bit has been set.
 289                 let notifier = Notifier::new();
 290                 notifier.notify();
 291
 292                 let callback = Arc::new(AtomicBool::new(false));
 293                 let callback_ref = Arc::clone(&callback);
 294                 notifier.get_future().register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
 295                 assert!(callback.load(Ordering::SeqCst));
 296
 297                 let callback = Arc::new(AtomicBool::new(false));
 298                 let callback_ref = Arc::clone(&callback);
 299                 notifier.get_future().register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
 300                 assert!(!callback.load(Ordering::SeqCst));
 301
 302                 notifier.notify();
 303                 assert!(callback.load(Ordering::SeqCst));
 304         }
 305
 306         #[cfg(feature = "std")]
 307         #[test]
 308         fn test_wait_timeout() {
 309                 use crate::sync::Arc;
 310                 use std::thread;
 311
 312                 let persistence_notifier = Arc::new(Notifier::new());
 313                 let thread_notifier = Arc::clone(&persistence_notifier);
 314
 315                 let exit_thread = Arc::new(AtomicBool::new(false));
 316                 let exit_thread_clone = exit_thread.clone();
 317                 thread::spawn(move || {
 318                         loop {
 319                                 let mut lock = thread_notifier.notify_pending.lock().unwrap();
 320                                 lock.0 = true;
 321                                 thread_notifier.condvar.notify_all();
 322
 323                                 if exit_thread_clone.load(Ordering::SeqCst) {
 324                                         break
 325                                 }
 326                         }
 327                 });
 328
 329                 // Check that we can block indefinitely until updates are available.
 330                 let _ = persistence_notifier.wait();
 331
 332                 // Check that the Notifier will return after the given duration if updates are
 333                 // available.
 334                 loop {
 335                         if persistence_notifier.wait_timeout(Duration::from_millis(100)) {
 336                                 break
 337                         }
 338                 }
 339
 340                 exit_thread.store(true, Ordering::SeqCst);
 341
 342                 // Check that the Notifier will return after the given duration even if no updates
 343                 // are available.
 344                 loop {
 345                         if !persistence_notifier.wait_timeout(Duration::from_millis(100)) {
 346                                 break
 347                         }
 348                 }
 349         }
 350
 351         #[test]
 352         fn test_future_callbacks() {
 353                 let future = Future {
 354                         state: Arc::new(Mutex::new(FutureState {
 355                                 callbacks: Vec::new(),
 356                                 complete: false,
 357                                 callbacks_made: false,
 358                         }))
 359                 };
 360                 let callback = Arc::new(AtomicBool::new(false));
 361                 let callback_ref = Arc::clone(&callback);
 362                 future.register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
 363
 364                 assert!(!callback.load(Ordering::SeqCst));
 365                 future.state.lock().unwrap().complete();
 366                 assert!(callback.load(Ordering::SeqCst));
 367                 future.state.lock().unwrap().complete();
 368         }
 369
 370         #[test]
 371         fn test_pre_completed_future_callbacks() {
 372                 let future = Future {
 373                         state: Arc::new(Mutex::new(FutureState {
 374                                 callbacks: Vec::new(),
 375                                 complete: false,
 376                                 callbacks_made: false,
 377                         }))
 378                 };
 379                 future.state.lock().unwrap().complete();
 380
 381                 let callback = Arc::new(AtomicBool::new(false));
 382                 let callback_ref = Arc::clone(&callback);
 383                 future.register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
 384
 385                 assert!(callback.load(Ordering::SeqCst));
 386                 assert!(future.state.lock().unwrap().callbacks.is_empty());
 387         }
 388
 389         // Rather annoyingly, there's no safe way in Rust std to construct a Waker despite it being
 390         // totally possible to construct from a trait implementation (though somewhat less effecient
 391         // compared to a raw VTable). Instead, we have to write out a lot of boilerplate to build a
 392         // waker, which we do here with a trivial Arc<AtomicBool> data element to track woke-ness.
 393         const WAKER_V_TABLE: RawWakerVTable = RawWakerVTable::new(waker_clone, wake, wake_by_ref, drop);
 394         unsafe fn wake_by_ref(ptr: *const ()) { let p = ptr as *const Arc<AtomicBool>; assert!(!(*p).fetch_or(true, Ordering::SeqCst)); }
 395         unsafe fn drop(ptr: *const ()) { let p = ptr as *mut Arc<AtomicBool>; let _freed = Box::from_raw(p); }
 396         unsafe fn wake(ptr: *const ()) { wake_by_ref(ptr); drop(ptr); }
 397         unsafe fn waker_clone(ptr: *const ()) -> RawWaker {
 398                 let p = ptr as *const Arc<AtomicBool>;
 399                 RawWaker::new(Box::into_raw(Box::new(Arc::clone(&*p))) as *const (), &WAKER_V_TABLE)
 400         }
 401
 402         fn create_waker() -> (Arc<AtomicBool>, Waker) {
 403                 let a = Arc::new(AtomicBool::new(false));
 404                 let waker = unsafe { Waker::from_raw(waker_clone((&a as *const Arc<AtomicBool>) as *const ())) };
 405                 (a, waker)
 406         }
 407
 408         #[test]
 409         fn test_future() {
 410                 let mut future = Future {
 411                         state: Arc::new(Mutex::new(FutureState {
 412                                 callbacks: Vec::new(),
 413                                 complete: false,
 414                                 callbacks_made: false,
 415                         }))
 416                 };
 417                 let mut second_future = Future { state: Arc::clone(&future.state) };
 418
 419                 let (woken, waker) = create_waker();
 420                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Pending);
 421                 assert!(!woken.load(Ordering::SeqCst));
 422
 423                 let (second_woken, second_waker) = create_waker();
 424                 assert_eq!(Pin::new(&mut second_future).poll(&mut Context::from_waker(&second_waker)), Poll::Pending);
 425                 assert!(!second_woken.load(Ordering::SeqCst));
 426
 427                 future.state.lock().unwrap().complete();
 428                 assert!(woken.load(Ordering::SeqCst));
 429                 assert!(second_woken.load(Ordering::SeqCst));
 430                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(()));
 431                 assert_eq!(Pin::new(&mut second_future).poll(&mut Context::from_waker(&second_waker)), Poll::Ready(()));
 432         }
 433 }