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

tokio/Cargo.toml

@@ -92,6 +92,7 @@ time = []
 tokio-macros = { version = "~2.2.0", path = "../tokio-macros", optional = true }
 
 pin-project-lite = "0.2.11"
+rustversion = { git = "https://github.com/wenym1/rustversion", rev = "0b11410" }
 
 # Everything else is optional...
 bytes = { version = "1.0.0", optional = true }
@@ -152,7 +153,8 @@ wasm-bindgen-test = "0.3.0"
 mio-aio = { version = "0.8.0", features = ["tokio"] }
 
 [target.'cfg(loom)'.dev-dependencies]
-loom = { version = "0.7", features = ["futures", "checkpoint"] }
+loom = { git = "https://github.com/wenym1/loom", rev = "4edf9435", features = ["futures", "checkpoint"] }
+#loom = { path = "/Users/william/repo/loom", features = ["futures", "checkpoint"] }
 
 [package.metadata.docs.rs]
 all-features = true

tokio/src/sync/oneshot.rs

@@ -627,7 +627,32 @@ impl<T> Sender<T> {
             )
         });
 
-        Ok(())
+        #[rustversion::since(1.70)]
+        fn consume_inner<T>(inner: Arc<Inner<T>>) -> Result<(), T> {
+            if let Some(inner) = Arc::into_inner(inner) {
+                if let Some(t) = inner.value.with_mut(|ptr| unsafe {
+                    // SAFETY: we have successfully returned with `Some`, which means we are the
+                    // only accessor to `ptr`.
+                    //
+                    // Note: value can be `None` even though we have previously set it as `Some`,
+                    // because the value may have been consumed by receiver before we reach here.
+                    (*ptr).take()
+                }) {
+                    Err(t)
+                } else {
+                    Ok(())
+                }
+            } else {
+                Ok(())
+            }
+        }
+
+        #[rustversion::before(1.70)]
+        fn consume_inner<T>(_inner: Arc<Inner<T>>) -> Result<(), T> {
+            Ok(())
+        }
+
+        consume_inner(inner)
     }
 
     /// Waits for the associated [`Receiver`] handle to close.

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