feat(oneshot channel): ensure msg won't be dropped on sender side when send returns ok

tokio/src/sync/oneshot.rs

@@ -1072,7 +1072,16 @@ impl<T> Receiver<T> {
 impl<T> Drop for Receiver<T> {
     fn drop(&mut self) {
         if let Some(inner) = self.inner.as_ref() {
-            inner.close();
+            let state = inner.close();
+
+            if state.is_complete() {
+                drop(unsafe {
+                    // SAFETY: we have ensured that the `VALUE_SENT` bit has been set,
+                    // so only the receiver can access the value.
+                    inner.consume_value()
+                });
+            }
+
             #[cfg(all(tokio_unstable, feature = "tracing"))]
             self.resource_span.in_scope(|| {
                 tracing::trace!(
@@ -1202,14 +1211,16 @@ impl<T> Inner<T> {
     }
 
     /// Called by `Receiver` to indicate that the value will never be received.
-    fn close(&self) {
+    fn close(&self) -> State {
         let prev = State::set_closed(&self.state);
 
         if prev.is_tx_task_set() && !prev.is_complete() {
             unsafe {
                 self.tx_task.with_task(Waker::wake_by_ref);
             }
         }
+
+        prev
     }
 
     /// Consumes the value. This function does not check `state`.
@@ -1248,6 +1259,16 @@ impl<T> Drop for Inner<T> {
                 self.tx_task.drop_task();
             }
         }
+
+        unsafe {
+            // SAFETY: we have `&mut self`, and therefore we have
+            // exclusive access to the value.
+            //
+            // Note: the assertion holds because if the value has been sent by sender,
+            // we must ensure that the value must have been consumed by the receiver before
+            // dropping the `Inner`.
+            debug_assert!(self.consume_value().is_none());
+        }
     }
 }
 

tokio/src/sync/tests/loom_oneshot.rs

@@ -1,4 +1,4 @@
-use crate::sync::oneshot;
+use crate::sync::{mpsc, oneshot};
 
 use futures::future::poll_fn;
 use loom::future::block_on;
@@ -86,6 +86,8 @@ fn recv_closed() {
 
 // TODO: Move this into `oneshot` proper.
 
+use crate::loom;
+use crate::sync::oneshot::{Receiver, Sender};
 use std::future::Future;
 use std::pin::Pin;
 use std::task::{Context, Poll};
@@ -138,3 +140,48 @@ fn changing_tx_task() {
         }
     });
 }
+
+#[test]
+fn checking_tx_send_ok_not_drop() {
+    use std::cell::RefCell;
+    loom::thread_local! {
+        static IS_RX: RefCell<bool> = RefCell::new(true);
+    }
+
+    struct Msg;
+
+    impl Drop for Msg {
+        fn drop(&mut self) {
+            IS_RX.with(|is_rx: &RefCell<_>| {
+                // On `tx.send(msg)` returning `Err(msg)`,
+                // we call `std::mem::forget(msg)`, so that
+                // `drop` is not expected to be called in the
+                // tx thread.
+                assert!(*is_rx.borrow());
+            });
+        }
+    }
+
+    loom::model(|| {
+        let (tx, rx) = oneshot::channel();
+
+        // tx thread
+        let tx_thread_join_handle = thread::spawn(move || {
+            // Ensure that `Msg::drop` in this thread will see is_rx == false
+            IS_RX.with(|is_rx: &RefCell<_>| {
+                *is_rx.borrow_mut() = false;
+            });
+            if let Err(msg) = tx.send(Msg) {
+                std::mem::forget(msg);
+            }
+        });
+
+        // rx thread
+        let rx_thread_join_handle = thread::spawn(move || {
+            drop(rx);
+        });
+
+        tx_thread_join_handle.join().unwrap();
+        rx_thread_join_handle.join().unwrap();
+    });
+}