diff --git a/tokio/src/sync/mpsc/chan.rs b/tokio/src/sync/mpsc/chan.rs index 0945d8805..f22880793 100644 --- a/tokio/src/sync/mpsc/chan.rs +++ b/tokio/src/sync/mpsc/chan.rs @@ -519,6 +519,11 @@ impl Drop for Rx { 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(); }); } diff --git a/tokio/tests/sync_mpsc.rs b/tokio/tests/sync_mpsc.rs index 93804581b..f91601313 100644 --- a/tokio/tests/sync_mpsc.rs +++ b/tokio/tests/sync_mpsc.rs @@ -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) {} + } + + // Create a dummy Arc 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) {}