diff --git a/tokio/src/sync/oneshot.rs b/tokio/src/sync/oneshot.rs index 32e288442..7f9571c90 100644 --- a/tokio/src/sync/oneshot.rs +++ b/tokio/src/sync/oneshot.rs @@ -1391,6 +1391,19 @@ impl Inner { } } + if prev.is_rx_task_set() && !prev.is_complete() { + State::unset_rx_task(&self.state); + // SAFETY: The sender only accesses `rx_task` (via + // `wake_by_ref`) in `complete()` after successfully setting + // `VALUE_SENT`. But `set_complete` will not set `VALUE_SENT` + // if `CLOSED` is already set (its CAS loop breaks early). + // Since `prev` shows that `VALUE_SENT` was not set before we + // set `CLOSED`, the sender can no longer set `VALUE_SENT` and + // will never access `rx_task`. Therefore, we have exclusive + // access here. + unsafe { self.rx_task.drop_task() }; + } + prev } diff --git a/tokio/src/sync/tests/loom_oneshot.rs b/tokio/src/sync/tests/loom_oneshot.rs index 62e1251c2..552cfcc8c 100644 --- a/tokio/src/sync/tests/loom_oneshot.rs +++ b/tokio/src/sync/tests/loom_oneshot.rs @@ -3,6 +3,7 @@ use crate::sync::oneshot; use loom::future::block_on; use loom::thread; use std::future::poll_fn; +use std::pin::Pin; use std::task::Poll::{Pending, Ready}; #[test] @@ -87,7 +88,6 @@ fn recv_closed() { // TODO: Move this into `oneshot` proper. use std::future::Future; -use std::pin::Pin; use std::task::{Context, Poll}; struct OnClose<'a> { @@ -185,3 +185,29 @@ fn checking_tx_send_ok_not_drop() { tx_thread_join_handle.join().unwrap(); }); } + +#[test] +fn drop_rx_after_poll() { + // Test that rx_task is properly deallocated when the receiver is dropped + // after being polled (which sets rx_task), while the sender is concurrently + // trying to send. + loom::model(|| { + let (tx, mut rx) = oneshot::channel::(); + + // Poll once to set rx_task before entering the parallel part of the + // test. + let _ = block_on(poll_fn(|cx| { + let _ = Pin::new(&mut rx).poll(cx); + Ready(()) + })); + + // Drop the receiver concurrently with the sender trying to send. + let rx_thread = thread::spawn(move || { + drop(rx); + }); + + let _ = tx.send(1); + + rx_thread.join().unwrap(); + }); +}