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