Never store more than one StdWaker per live Future #2894
Conversation
WalkthroughThe update introduces enhancements to the handling of futures and wakers, focusing on optimizing memory usage and improving the efficiency of future polling. It involves changes in how futures are managed, specifically addressing issues with memory allocation by refining the callback and waker mechanisms. The modifications aim to prevent unnecessary memory growth and ensure more efficient future completion tracking, particularly in busy environments. Changes
Assessment against linked issues
Poem
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (invoked as PR comments)
Additionally, you can add CodeRabbit Configration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Review Status
Actionable comments generated: 0
Configuration used: CodeRabbit UI
Files selected for processing (2)
- lightning-background-processor/src/lib.rs (1 hunks)
- lightning/src/util/wakers.rs (15 hunks)
Additional comments: 8
lightning/src/util/wakers.rs (6)
- 59-70: The logic for managing
self_idxand resetting the future state whencallbacks_madeis true appears sound. However, ensure that the logic for resettinglock.0tofalseand handlingnext_idxis thoroughly tested, especially in concurrent scenarios where multiple futures might be interacting with the sameNotifier.- 78-82: Adding
std_future_callbacksto track standard wakers is a critical improvement. Ensure that the vector is efficiently managed, especially in terms of memory usage and performance when futures are frequently polled and woken.- 124-128: The management of
std_future_callbackswith a tuple of(usize, StdWaker)is appropriate for ensuring that only the most recent waker is used. However, consider the performance implications of frequently modifying this vector, especially in high-concurrency scenarios.- 207-208: The
Dropimplementation forFuturecorrectly ensures that thestd_future_callbacksvector does not retain entries for dropped futures. This is crucial for preventing memory leaks. Verify that this logic works as expected in scenarios where futures are dropped before being completed.- 225-226: The logic to retain only the last waker by removing the current
Future's waker before adding a new one is a good approach to prevent unnecessary wake-ups. Ensure this mechanism is thoroughly tested to confirm it behaves as expected under various polling scenarios.- 730-747: The test case
multi_poll_stores_single_wakereffectively demonstrates the intended behavior of storing only the last waker and cleaning up uponFuturedrop. Ensure comprehensive testing covers edge cases, such as rapid polling and dropping of futures in a multi-threaded environment.lightning-background-processor/src/lib.rs (2)
- 857-858: The changes involve passing references to futures obtained from
channel_manager.get_event_or_persistence_needed_future()andchain_monitor.get_update_future()withinSleeper::from_two_futures()instead of directly calling these functions. This adjustment is part of a broader strategy to manage futures and their lifecycles more efficiently.- 857-858: Ensure that the changes to pass references to futures instead of directly calling the functions are correctly implemented across all usages. This approach should improve efficiency by avoiding unnecessary calls and managing futures' lifecycles more effectively.
Codecov ReportAttention:
❗ Your organization needs to install the Codecov GitHub app to enable full functionality. Additional details and impacted files@@ Coverage Diff @@
## main #2894 +/- ##
==========================================
+ Coverage 89.13% 89.39% +0.26%
==========================================
Files 115 115
Lines 94179 96441 +2262
Branches 94179 96441 +2262
==========================================
+ Hits 83944 86211 +2267
- Misses 7761 7807 +46
+ Partials 2474 2423 -51 ☔ View full report in Codecov by Sentry. |
There was a problem hiding this comment.
LGTM
Grr, I recently even stumbled across the one line in the Future docs that mentions this but didn't think through the implications for LDK. FWIW, it would be nice if they would highlight the potential leakage and/or mention it somewhere besides this one line though.
Given how invasive the changes are, I think this could use a second reviewer.
| lock.1.take(); | ||
| lock.0 = false; | ||
| } else { | ||
| self_idx = locked.next_idx; |
There was a problem hiding this comment.
nit: 'fetch and add' could be a method so that (in future, no pun intended) we would never forget to increase the counter?
There was a problem hiding this comment.
I played with a constructor a bit trying to make it more robust but didn't really see a decent way to do it without double-locking everywhere. Just defining a method to fetch-and-increment the index doesn't seem like it'll actually prevent a bug cause we'll just forget to use it :)
TheBlueMatt
force-pushed
the
2024-02-future-poll-leak
branch
from
60bff36
to
5a5a8d1
Compare
There was a problem hiding this comment.
Review Status
Actionable comments generated: 0
Configuration used: CodeRabbit UI
Files selected for processing (2)
- lightning-background-processor/src/lib.rs (1 hunks)
- lightning/src/util/wakers.rs (15 hunks)
Files skipped from review as they are similar to previous changes (2)
- lightning-background-processor/src/lib.rs
- lightning/src/util/wakers.rs
|
Feel free to squash |
In the next commit we'll fix a memory leak due to keeping too many `std::task::Waker` callbacks in `FutureState` from redundant `poll` calls, but first we need to split handling of `StdWaker`-based future wake callbacks from normal ones, which we do here.
When an `std::future::Future` is `poll()`ed, we're only supposed to use the latest `Waker` provided. However, we currently push an `StdWaker` onto our callback list every time `poll` is called, waking every `Waker` but also using more and more memory until the `Future` itself is woken. Here we take a step towards fixing this by giving each `Future` a unique index and storing which `Future` an `StdWaker` came from in the callback list. This sets us up to deduplicate `StdWaker`s by `Future`s in the next commit.
When an `std::future::Future` is `poll()`ed, we're only supposed to use the latest `Waker` provided. However, we currently push an `StdWaker` onto our callback list every time `poll` is called, waking every `Waker` but also using more and more memory until the `Future` itself is woken. Here we fix this by removing any `StdWaker`s stored for a given `Future` when it is `drop`ped or prior to pushing a new `StdWaker` onto the list when `poll`ed. Sadly, the introduction of a `Drop` impl for `Future` means we can't trivially destructure the struct any longer, causing a few methods to need to take `Future`s by reference rather than ownership and `clone` a few `Arc`s. Fixes lightningdevkit#2874
TheBlueMatt
force-pushed
the
2024-02-future-poll-leak
branch
from
5a5a8d1
to
8157c01
Compare
|
Squashed without further changes, diff from yesterday: $ git diff-tree -U2 60bff36 8157c01e
diff --git a/lightning/src/util/wakers.rs b/lightning/src/util/wakers.rs
index 28cea2624..b2c9d21b9 100644
--- a/lightning/src/util/wakers.rs
+++ b/lightning/src/util/wakers.rs
@@ -118,10 +118,11 @@ define_callback!();
pub(crate) struct FutureState {
- // When we're tracking whether a callback counts as having woken the user's code, we check the
- // first bool - set to false if we're just calling a Waker, and true if we're calling an actual
- // user-provided function.
+ // `callbacks` count as having woken the users' code (as they go direct to the user), but
+ // `std_future_callbacks` and `callbacks_with_state` do not (as the first just wakes a future,
+ // we only count it after another `poll()` and the second wakes a `Sleeper` which handles
+ // setting `callbacks_made` itself).
callbacks: Vec<Box<dyn FutureCallback>>,
std_future_callbacks: Vec<(usize, StdWaker)>,
- callbacks_with_state: Vec<(bool, Box<dyn Fn(&Arc<Mutex<FutureState>>) -> () + Send>)>,
+ callbacks_with_state: Vec<Box<dyn Fn(&Arc<Mutex<FutureState>>) -> () + Send>>,
complete: bool,
callbacks_made: bool,
@@ -139,7 +140,6 @@ fn complete_future(this: &Arc<Mutex<FutureState>>) -> bool {
waker.0.wake_by_ref();
}
- for (counts_as_call, callback) in state.callbacks_with_state.drain(..) {
+ for callback in state.callbacks_with_state.drain(..) {
(callback)(this);
- state.callbacks_made |= counts_as_call;
}
state.complete = true;
@@ -267,8 +267,8 @@ impl Sleeper {
break;
}
- notifier.callbacks_with_state.push((false, Box::new(move |notifier_ref| {
+ notifier.callbacks_with_state.push(Box::new(move |notifier_ref| {
*notified_fut_ref.lock().unwrap() = Some(Arc::clone(notifier_ref));
cv_ref.notify_all();
- })));
+ }));
}
}
@@ -745,4 +745,15 @@ mod tests {
mem::drop(future_b);
assert_eq!(future_state.lock().unwrap().std_future_callbacks.len(), 0);
+
+ // Further, after polling a future twice, if the notifier is woken all Wakers are dropped.
+ let mut future_a = notifier.get_future();
+ assert_eq!(Pin::new(&mut future_a).poll(&mut Context::from_waker(&create_waker().1)), Poll::Pending);
+ assert_eq!(future_state.lock().unwrap().std_future_callbacks.len(), 1);
+ assert_eq!(Pin::new(&mut future_a).poll(&mut Context::from_waker(&create_waker().1)), Poll::Pending);
+ assert_eq!(future_state.lock().unwrap().std_future_callbacks.len(), 1);
+ notifier.notify();
+ assert_eq!(future_state.lock().unwrap().std_future_callbacks.len(), 0);
+ assert_eq!(Pin::new(&mut future_a).poll(&mut Context::from_waker(&create_waker().1)), Poll::Ready(()));
+ assert_eq!(future_state.lock().unwrap().std_future_callbacks.len(), 0);
}
} |
There was a problem hiding this comment.
Review Status
Actionable comments generated: 0
Configuration used: CodeRabbit UI
Files selected for processing (2)
- lightning-background-processor/src/lib.rs (1 hunks)
- lightning/src/util/wakers.rs (15 hunks)
Files skipped from review as they are similar to previous changes (2)
- lightning-background-processor/src/lib.rs
- lightning/src/util/wakers.rs
When an
std::future::Futureispoll()ed, we're only supposed touse the latest
Wakerprovided. However, we currently push anStdWakeronto our callback list every timepollis called,waking every
Wakerbut also using more and more memory until theFutureitself is woken.Here we fix this by removing any
StdWakers stored for a givenFuturewhen it isdropped or prior to pushing a newStdWakeronto the list when
polled.Sadly, the introduction of a
Dropimpl forFuturemeans wecan't trivially destructure the struct any longer, causing a few
methods to need to take
Futures by reference rather thanownership and
clonea fewArcs.Fixes #2874