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