From 3873afcb752fa5e84868719e71d506ae8a6592ae Mon Sep 17 00:00:00 2001 From: Matt Corallo Date: Wed, 26 Apr 2023 04:29:59 +0000 Subject: [PATCH 1/2] Hold a reference to the `Arc` when completing futures This will allow us to pass in that state to the callbacks in the next commit. --- lightning/src/util/wakers.rs | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/lightning/src/util/wakers.rs b/lightning/src/util/wakers.rs index 602c2ee04b7..2dbea705a31 100644 --- a/lightning/src/util/wakers.rs +++ b/lightning/src/util/wakers.rs @@ -45,7 +45,7 @@ impl Notifier { pub(crate) fn notify(&self) { let mut lock = self.notify_pending.lock().unwrap(); if let Some(future_state) = &lock.1 { - if future_state.lock().unwrap().complete() { + if complete_future(future_state) { lock.1 = None; return; } @@ -116,15 +116,15 @@ pub(crate) struct FutureState { callbacks_made: bool, } -impl FutureState { - fn complete(&mut self) -> bool { - for (counts_as_call, callback) in self.callbacks.drain(..) { - callback.call(); - self.callbacks_made |= counts_as_call; - } - self.complete = true; - self.callbacks_made +fn complete_future(this: &Arc>) -> bool { + let mut state_lock = this.lock().unwrap(); + let state = &mut *state_lock; + for (counts_as_call, callback) in state.callbacks.drain(..) { + callback.call(); + state.callbacks_made |= counts_as_call; } + state.complete = true; + state.callbacks_made } /// A simple future which can complete once, and calls some callback(s) when it does so. @@ -421,9 +421,9 @@ mod tests { future.register_callback(Box::new(move || assert!(!callback_ref.fetch_or(true, Ordering::SeqCst)))); assert!(!callback.load(Ordering::SeqCst)); - future.state.lock().unwrap().complete(); + complete_future(&future.state); assert!(callback.load(Ordering::SeqCst)); - future.state.lock().unwrap().complete(); + complete_future(&future.state); } #[test] @@ -435,7 +435,7 @@ mod tests { callbacks_made: false, })) }; - future.state.lock().unwrap().complete(); + complete_future(&future.state); let callback = Arc::new(AtomicBool::new(false)); let callback_ref = Arc::clone(&callback); @@ -483,7 +483,7 @@ mod tests { assert_eq!(Pin::new(&mut second_future).poll(&mut Context::from_waker(&second_waker)), Poll::Pending); assert!(!second_woken.load(Ordering::SeqCst)); - future.state.lock().unwrap().complete(); + complete_future(&future.state); assert!(woken.load(Ordering::SeqCst)); assert!(second_woken.load(Ordering::SeqCst)); assert_eq!(Pin::new(&mut future).poll(&mut Context::from_waker(&waker)), Poll::Ready(())); From 7caa5840513332bab30b3f8db45c938c8e07011b Mon Sep 17 00:00:00 2001 From: Matt Corallo Date: Wed, 26 Apr 2023 05:01:13 +0000 Subject: [PATCH 2/2] Fix a leak in `FutureState` when a `Notifier` is dropped un-woken If a `Notifier` has an internal `FutureState` which gathers some sleeper callbacks, but is never actaully woken, those callbacks will leak due to a circular `Arc` reference when the `Notifier` is `drop`'d. Because `Notifier`s are rarely `drop`'d in production this isn't a huge deal, but shows up materially in bindings tests as they spawn many nodes over the course of a short test. Fixes #2232 --- lightning/src/util/wakers.rs | 54 +++++++++++++++++++++++++++++++++--- 1 file changed, 50 insertions(+), 4 deletions(-) diff --git a/lightning/src/util/wakers.rs b/lightning/src/util/wakers.rs index 2dbea705a31..37c036da959 100644 --- a/lightning/src/util/wakers.rs +++ b/lightning/src/util/wakers.rs @@ -69,6 +69,7 @@ impl Notifier { } else { let state = Arc::new(Mutex::new(FutureState { callbacks: Vec::new(), + callbacks_with_state: Vec::new(), complete: lock.0, callbacks_made: false, })); @@ -112,6 +113,7 @@ pub(crate) struct FutureState { // first bool - set to false if we're just calling a Waker, and true if we're calling an actual // user-provided function. callbacks: Vec<(bool, Box)>, + callbacks_with_state: Vec<(bool, Box>) -> () + Send>)>, complete: bool, callbacks_made: bool, } @@ -123,6 +125,10 @@ fn complete_future(this: &Arc>) -> bool { callback.call(); state.callbacks_made |= counts_as_call; } + for (counts_as_call, callback) in state.callbacks_with_state.drain(..) { + (callback)(this); + state.callbacks_made |= counts_as_call; + } state.complete = true; state.callbacks_made } @@ -240,14 +246,13 @@ impl Sleeper { for notifier_mtx in self.notifiers.iter() { let cv_ref = Arc::clone(&cv); let notified_fut_ref = Arc::clone(¬ified_fut_mtx); - let notifier_ref = Arc::clone(¬ifier_mtx); let mut notifier = notifier_mtx.lock().unwrap(); if notifier.complete { - *notified_fut_mtx.lock().unwrap() = Some(notifier_ref); + *notified_fut_mtx.lock().unwrap() = Some(Arc::clone(¬ifier_mtx)); break; } - notifier.callbacks.push((false, Box::new(move || { - *notified_fut_ref.lock().unwrap() = Some(Arc::clone(¬ifier_ref)); + notifier.callbacks_with_state.push((false, Box::new(move |notifier_ref| { + *notified_fut_ref.lock().unwrap() = Some(Arc::clone(notifier_ref)); cv_ref.notify_all(); }))); } @@ -407,11 +412,50 @@ mod tests { } } + #[cfg(feature = "std")] + #[test] + fn test_state_drops() { + // Previously, there was a leak if a `Notifier` was `drop`ed without ever being notified + // but after having been slept-on. This tests for that leak. + use crate::sync::Arc; + use std::thread; + + let notifier_a = Arc::new(Notifier::new()); + let notifier_b = Arc::new(Notifier::new()); + + let thread_notifier_a = Arc::clone(¬ifier_a); + + let future_a = notifier_a.get_future(); + let future_state_a = Arc::downgrade(&future_a.state); + + let future_b = notifier_b.get_future(); + let future_state_b = Arc::downgrade(&future_b.state); + + let join_handle = thread::spawn(move || { + // Let the other thread get to the wait point, then notify it. + std::thread::sleep(Duration::from_millis(50)); + thread_notifier_a.notify(); + }); + + // Wait on the other thread to finish its sleep, note that the leak only happened if we + // actually have to sleep here, not if we immediately return. + Sleeper::from_two_futures(future_a, future_b).wait(); + + join_handle.join().unwrap(); + + // then drop the notifiers and make sure the future states are gone. + mem::drop(notifier_a); + mem::drop(notifier_b); + + assert!(future_state_a.upgrade().is_none() && future_state_b.upgrade().is_none()); + } + #[test] fn test_future_callbacks() { let future = Future { state: Arc::new(Mutex::new(FutureState { callbacks: Vec::new(), + callbacks_with_state: Vec::new(), complete: false, callbacks_made: false, })) @@ -431,6 +475,7 @@ mod tests { let future = Future { state: Arc::new(Mutex::new(FutureState { callbacks: Vec::new(), + callbacks_with_state: Vec::new(), complete: false, callbacks_made: false, })) @@ -469,6 +514,7 @@ mod tests { let mut future = Future { state: Arc::new(Mutex::new(FutureState { callbacks: Vec::new(), + callbacks_with_state: Vec::new(), complete: false, callbacks_made: false, }))