]> git.bitcoin.ninja Git - rust-lightning/blob - lightning/src/util/wakers.rs
Limit TLV stream decoding to type ranges
[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                 lock.0 = true;
92                 if let Some(future_state) = lock.1.take() {
93                         future_state.lock().unwrap().complete();
94                 }
95                 mem::drop(lock);
96                 self.condvar.notify_all();
97         }
98
99         /// Gets a [`Future`] that will get woken up with any waiters
100         pub(crate) fn get_future(&self) -> Future {
101                 let mut lock = self.notify_pending.lock().unwrap();
102                 if lock.0 {
103                         Future {
104                                 state: Arc::new(Mutex::new(FutureState {
105                                         callbacks: Vec::new(),
106                                         complete: true,
107                                 }))
108                         }
109                 } else if let Some(existing_state) = &lock.1 {
110                         Future { state: Arc::clone(&existing_state) }
111                 } else {
112                         let state = Arc::new(Mutex::new(FutureState {
113                                 callbacks: Vec::new(),
114                                 complete: false,
115                         }));
116                         lock.1 = Some(Arc::clone(&state));
117                         Future { state }
118                 }
119         }
120
121         #[cfg(any(test, feature = "_test_utils"))]
122         pub fn notify_pending(&self) -> bool {
123                 self.notify_pending.lock().unwrap().0
124         }
125 }
126
127 /// A callback which is called when a [`Future`] completes.
128 ///
129 /// Note that this MUST NOT call back into LDK directly, it must instead schedule actions to be
130 /// taken later. Rust users should use the [`std::future::Future`] implementation for [`Future`]
131 /// instead.
132 ///
133 /// Note that the [`std::future::Future`] implementation may only work for runtimes which schedule
134 /// futures when they receive a wake, rather than immediately executing them.
135 pub trait FutureCallback : Send {
136         /// The method which is called.
137         fn call(&self);
138 }
139
140 impl<F: Fn() + Send> FutureCallback for F {
141         fn call(&self) { (self)(); }
142 }
143
144 pub(crate) struct FutureState {
145         callbacks: Vec<Box<dyn FutureCallback>>,
146         complete: bool,
147 }
148
149 impl FutureState {
150         fn complete(&mut self) {
151                 for callback in self.callbacks.drain(..) {
152                         callback.call();
153                 }
154                 self.complete = true;
155         }
156 }
157
158 /// A simple future which can complete once, and calls some callback(s) when it does so.
159 pub struct Future {
160         state: Arc<Mutex<FutureState>>,
161 }
162
163 impl Future {
164         /// Registers a callback to be called upon completion of this future. If the future has already
165         /// completed, the callback will be called immediately.
166         ///
167         /// (C-not exported) use the bindings-only `register_callback_fn` instead
168         pub fn register_callback(&self, callback: Box<dyn FutureCallback>) {
169                 let mut state = self.state.lock().unwrap();
170                 if state.complete {
171                         mem::drop(state);
172                         callback.call();
173                 } else {
174                         state.callbacks.push(callback);
175                 }
176         }
177
178         // C bindings don't (currently) know how to map `Box<dyn Trait>`, and while it could add the
179         // following wrapper, doing it in the bindings is currently much more work than simply doing it
180         // here.
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         #[cfg(c_bindings)]
184         pub fn register_callback_fn<F: 'static + FutureCallback>(&self, callback: F) {
185                 self.register_callback(Box::new(callback));
186         }
187 }
188
189 mod std_future {
190         use core::task::Waker;
191         pub struct StdWaker(pub Waker);
192         impl super::FutureCallback for StdWaker {
193                 fn call(&self) { self.0.wake_by_ref() }
194         }
195 }
196
197 /// (C-not exported) as Rust Futures aren't usable in language bindings.
198 impl<'a> StdFuture for Future {
199         type Output = ();
200
201         fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
202                 let mut state = self.state.lock().unwrap();
203                 if state.complete {
204                         Poll::Ready(())
205                 } else {
206                         let waker = cx.waker().clone();
207                         state.callbacks.push(Box::new(std_future::StdWaker(waker)));
208                         Poll::Pending
209                 }
210         }
211 }
212
213 #[cfg(test)]
214 mod tests {
215         use super::*;
216         use core::sync::atomic::{AtomicBool, Ordering};
217         use core::future::Future as FutureTrait;
218         use core::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
219
220         #[test]
221         fn notifier_pre_notified_future() {
222                 // Previously, if we generated a future after a `Notifier` had been notified, the future
223                 // would never complete. This tests this behavior, ensuring the future instead completes
224                 // immediately.
225                 let notifier = Notifier::new();
226                 notifier.notify();
227
228                 let callback = Arc::new(AtomicBool::new(false));
229                 let callback_ref = Arc::clone(&callback);
230                 notifier.get_future().register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
231                 assert!(callback.load(Ordering::SeqCst));
232         }
233
234         #[cfg(feature = "std")]
235         #[test]
236         fn test_wait_timeout() {
237                 use crate::sync::Arc;
238                 use std::thread;
239
240                 let persistence_notifier = Arc::new(Notifier::new());
241                 let thread_notifier = Arc::clone(&persistence_notifier);
242
243                 let exit_thread = Arc::new(AtomicBool::new(false));
244                 let exit_thread_clone = exit_thread.clone();
245                 thread::spawn(move || {
246                         loop {
247                                 let mut lock = thread_notifier.notify_pending.lock().unwrap();
248                                 lock.0 = true;
249                                 thread_notifier.condvar.notify_all();
250
251                                 if exit_thread_clone.load(Ordering::SeqCst) {
252                                         break
253                                 }
254                         }
255                 });
256
257                 // Check that we can block indefinitely until updates are available.
258                 let _ = persistence_notifier.wait();
259
260                 // Check that the Notifier will return after the given duration if updates are
261                 // available.
262                 loop {
263                         if persistence_notifier.wait_timeout(Duration::from_millis(100)) {
264                                 break
265                         }
266                 }
267
268                 exit_thread.store(true, Ordering::SeqCst);
269
270                 // Check that the Notifier will return after the given duration even if no updates
271                 // are available.
272                 loop {
273                         if !persistence_notifier.wait_timeout(Duration::from_millis(100)) {
274                                 break
275                         }
276                 }
277         }
278
279         #[test]
280         fn test_future_callbacks() {
281                 let future = Future {
282                         state: Arc::new(Mutex::new(FutureState {
283                                 callbacks: Vec::new(),
284                                 complete: false,
285                         }))
286                 };
287                 let callback = Arc::new(AtomicBool::new(false));
288                 let callback_ref = Arc::clone(&callback);
289                 future.register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
290
291                 assert!(!callback.load(Ordering::SeqCst));
292                 future.state.lock().unwrap().complete();
293                 assert!(callback.load(Ordering::SeqCst));
294                 future.state.lock().unwrap().complete();
295         }
296
297         #[test]
298         fn test_pre_completed_future_callbacks() {
299                 let future = Future {
300                         state: Arc::new(Mutex::new(FutureState {
301                                 callbacks: Vec::new(),
302                                 complete: false,
303                         }))
304                 };
305                 future.state.lock().unwrap().complete();
306
307                 let callback = Arc::new(AtomicBool::new(false));
308                 let callback_ref = Arc::clone(&callback);
309                 future.register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst))));
310
311                 assert!(callback.load(Ordering::SeqCst));
312                 assert!(future.state.lock().unwrap().callbacks.is_empty());
313         }
314
315         // Rather annoyingly, there's no safe way in Rust std to construct a Waker despite it being
316         // totally possible to construct from a trait implementation (though somewhat less effecient
317         // compared to a raw VTable). Instead, we have to write out a lot of boilerplate to build a
318         // waker, which we do here with a trivial Arc<AtomicBool> data element to track woke-ness.
319         const WAKER_V_TABLE: RawWakerVTable = RawWakerVTable::new(waker_clone, wake, wake_by_ref, drop);
320         unsafe fn wake_by_ref(ptr: *const ()) { let p = ptr as *const Arc<AtomicBool>; assert!(!(*p).fetch_or(true, Ordering::SeqCst)); }
321         unsafe fn drop(ptr: *const ()) { let p = ptr as *mut Arc<AtomicBool>; let _freed = Box::from_raw(p); }
322         unsafe fn wake(ptr: *const ()) { wake_by_ref(ptr); drop(ptr); }
323         unsafe fn waker_clone(ptr: *const ()) -> RawWaker {
324                 let p = ptr as *const Arc<AtomicBool>;
325                 RawWaker::new(Box::into_raw(Box::new(Arc::clone(&*p))) as *const (), &WAKER_V_TABLE)
326         }
327
328         fn create_waker() -> (Arc<AtomicBool>, Waker) {
329                 let a = Arc::new(AtomicBool::new(false));
330                 let waker = unsafe { Waker::from_raw(waker_clone((&a as *const Arc<AtomicBool>) as *const ())) };
331                 (a, waker)
332         }
333
334         #[test]
335         fn test_future() {
336                 let mut future = Future {
337                         state: Arc::new(Mutex::new(FutureState {
338                                 callbacks: Vec::new(),
339                                 complete: false,
340                         }))
341                 };
342                 let mut second_future = Future { state: Arc::clone(&future.state) };
343
344                 let (woken, waker) = create_waker();
345                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Pending);
346                 assert!(!woken.load(Ordering::SeqCst));
347
348                 let (second_woken, second_waker) = create_waker();
349                 assert_eq!(Pin::new(&mut second_future).poll(&mut Context::from_waker(&second_waker)), Poll::Pending);
350                 assert!(!second_woken.load(Ordering::SeqCst));
351
352                 future.state.lock().unwrap().complete();
353                 assert!(woken.load(Ordering::SeqCst));
354                 assert!(second_woken.load(Ordering::SeqCst));
355                 assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(()));
356                 assert_eq!(Pin::new(&mut second_future).poll(&mut Context::from_waker(&second_waker)), Poll::Ready(()));
357         }
358 }