mirror of
https://github.com/tokio-rs/tokio.git
synced 2026-08-07 00:00:09 +02:00
sync: reset Chan::rx_waker in chan::Rx's Drop impl (#8095)
I recently fixed a memory leak in an application where tokio's RawTask storage was being kept alive by a leaked Waker. The task itself was polling an `mpsc::Receiver` before then being aborted. During cleanup, all references to the RawTask were dropped, except for the one stored in `mpsc::chan::Chan::rx_waker`. While the `Receiver` was dropped as part of the task's future, one of the channel's `Sender`s was leaked outside the task. This meant the `Chan` was never dropped and its `rx_waker` contained the leaked Waker. I fixed the leak by properly cleaning up the `Sender`, but I also think keeping `rx_waker` around in this case is unnecessary. Once `chan::Rx` is dropped, it can't be polled anymore, so waking up the registered task will always be spurious. The commit includes a regression test to illustrate the problem that is fixed by removing the waker explicitly.
This commit is contained in:
@@ -519,6 +519,11 @@ impl<T, S: Semaphore> Drop for Rx<T, S> {
|
||||
sem: &self.inner.semaphore,
|
||||
};
|
||||
|
||||
// When Rx is dropped, there is nothing for a task to poll anymore.
|
||||
// This means we can drop our waker to potentially free up resources.
|
||||
// Do so before draining the channel where panics may occur.
|
||||
self.inner.rx_waker.take_waker();
|
||||
|
||||
guard.drain();
|
||||
});
|
||||
}
|
||||
|
||||
@@ -1631,4 +1631,38 @@ fn drop_all_elements_during_panic() {
|
||||
// `mpsc::Chan`'s drop is called, freeing the `Block` memory allocation.
|
||||
}
|
||||
|
||||
/// A task's wakers keep its memory allocation alive. This test ensures
|
||||
/// `mpsc::channel` releases the rx waker immediately upon drop to potentially
|
||||
/// free up resources.
|
||||
#[test]
|
||||
fn release_waker_on_rx_drop() {
|
||||
use std::task::{Context, Wake, Waker};
|
||||
|
||||
struct DummyWaker;
|
||||
impl Wake for DummyWaker {
|
||||
fn wake(self: Arc<Self>) {}
|
||||
}
|
||||
|
||||
// Create a dummy Arc<impl Wake> to count waker references
|
||||
let dummy = Arc::new(DummyWaker);
|
||||
let waker = Waker::from(dummy.clone());
|
||||
let mut cx = Context::from_waker(&waker);
|
||||
|
||||
// Baseline: 2 references (`dummy` and `waker`)
|
||||
assert_eq!(Arc::strong_count(&dummy), 2);
|
||||
|
||||
// Register `waker` in rx
|
||||
let (_tx, mut rx) = mpsc::channel::<()>(1);
|
||||
assert_pending!(rx.poll_recv(&mut cx));
|
||||
assert_eq!(Arc::strong_count(&dummy), 3);
|
||||
|
||||
// Drop rx, releasing the stored waker
|
||||
drop(rx);
|
||||
assert_eq!(
|
||||
Arc::strong_count(&dummy),
|
||||
2,
|
||||
"dropping rx did not drop its waker",
|
||||
);
|
||||
}
|
||||
|
||||
fn is_debug<T: fmt::Debug>(_: &T) {}
|
||||
|
||||
Reference in New Issue
Block a user