sync: drop rx waker when oneshot receiver is dropped (#7886)

tokio/src/sync/oneshot.rs

@@ -1391,6 +1391,19 @@ impl<T> Inner<T> {
             }
         }
 
+        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
     }
 

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::<i32>();
+
+        // 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();
+    });
+}