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