From 804dbd6f8eb6ca65cc67a55048b4579ef1a457aa Mon Sep 17 00:00:00 2001 From: Carl Lerche Date: Thu, 10 Oct 2019 12:00:22 -0700 Subject: [PATCH] sync: fix mem leak in oneshot on task migration (#1648) When polling the task, the current waker is saved to the oneshot state. When the handle is migrated to a new task and polled again, the waker must be swaped from the old waker to the new waker. In some cases, there is a potential for the old waker to leak. This bug was caught by loom with the recently added memory leak detection. --- tokio-sync/src/oneshot.rs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tokio-sync/src/oneshot.rs b/tokio-sync/src/oneshot.rs index 04401b9fc..8ed1672fc 100644 --- a/tokio-sync/src/oneshot.rs +++ b/tokio-sync/src/oneshot.rs @@ -171,6 +171,8 @@ impl Sender { state = State::unset_tx_task(&inner.state); if state.is_closed() { + // Set the flag again so that the waker is released in drop + State::set_tx_task(&inner.state); return Ready(()); } else { unsafe { inner.drop_tx_task() }; @@ -359,6 +361,9 @@ impl Inner { // Unset the task state = State::unset_rx_task(&self.state); if state.is_complete() { + // Set the flag again so that the waker is released in drop + State::set_rx_task(&self.state); + return match unsafe { self.consume_value() } { Some(value) => Ready(Ok(value)), None => Ready(Err(RecvError(()))),