Skip to content

Commit a194534

Browse files
committed
Split lists of Waker and directly-registered Future callbacks
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.
1 parent 73da722 commit a194534

File tree

1 file changed

+12
-8
lines changed

1 file changed

+12
-8
lines changed

lightning/src/util/wakers.rs

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ impl Notifier {
6969
} else {
7070
let state = Arc::new(Mutex::new(FutureState {
7171
callbacks: Vec::new(),
72+
std_future_callbacks: Vec::new(),
7273
callbacks_with_state: Vec::new(),
7374
complete: lock.0,
7475
callbacks_made: false,
@@ -112,7 +113,8 @@ pub(crate) struct FutureState {
112113
// When we're tracking whether a callback counts as having woken the user's code, we check the
113114
// first bool - set to false if we're just calling a Waker, and true if we're calling an actual
114115
// user-provided function.
115-
callbacks: Vec<(bool, Box<dyn FutureCallback>)>,
116+
callbacks: Vec<Box<dyn FutureCallback>>,
117+
std_future_callbacks: Vec<StdWaker>,
116118
callbacks_with_state: Vec<(bool, Box<dyn Fn(&Arc<Mutex<FutureState>>) -> () + Send>)>,
117119
complete: bool,
118120
callbacks_made: bool,
@@ -121,9 +123,12 @@ pub(crate) struct FutureState {
121123
fn complete_future(this: &Arc<Mutex<FutureState>>) -> bool {
122124
let mut state_lock = this.lock().unwrap();
123125
let state = &mut *state_lock;
124-
for (counts_as_call, callback) in state.callbacks.drain(..) {
126+
for callback in state.callbacks.drain(..) {
125127
callback.call();
126-
state.callbacks_made |= counts_as_call;
128+
state.callbacks_made = true;
129+
}
130+
for waker in state.std_future_callbacks.drain(..) {
131+
waker.0.wake_by_ref();
127132
}
128133
for (counts_as_call, callback) in state.callbacks_with_state.drain(..) {
129134
(callback)(this);
@@ -153,7 +158,7 @@ impl Future {
153158
mem::drop(state);
154159
callback.call();
155160
} else {
156-
state.callbacks.push((true, callback));
161+
state.callbacks.push(callback);
157162
}
158163
}
159164

@@ -193,9 +198,6 @@ impl Future {
193198

194199
use core::task::Waker;
195200
struct StdWaker(pub Waker);
196-
impl FutureCallback for StdWaker {
197-
fn call(&self) { self.0.wake_by_ref() }
198-
}
199201

200202
/// This is not exported to bindings users as Rust Futures aren't usable in language bindings.
201203
impl<'a> StdFuture for Future {
@@ -208,7 +210,7 @@ impl<'a> StdFuture for Future {
208210
Poll::Ready(())
209211
} else {
210212
let waker = cx.waker().clone();
211-
state.callbacks.push((false, Box::new(StdWaker(waker))));
213+
state.std_future_callbacks.push(StdWaker(waker));
212214
Poll::Pending
213215
}
214216
}
@@ -455,6 +457,7 @@ mod tests {
455457
let future = Future {
456458
state: Arc::new(Mutex::new(FutureState {
457459
callbacks: Vec::new(),
460+
std_future_callbacks: Vec::new(),
458461
callbacks_with_state: Vec::new(),
459462
complete: false,
460463
callbacks_made: false,
@@ -514,6 +517,7 @@ mod tests {
514517
let mut future = Future {
515518
state: Arc::new(Mutex::new(FutureState {
516519
callbacks: Vec::new(),
520+
std_future_callbacks: Vec::new(),
517521
callbacks_with_state: Vec::new(),
518522
complete: false,
519523
callbacks_made: false,

0 commit comments

Comments
 (0)