fix(oneshot): allow reclamation after sender handoff

AGENTS.md

@@ -31,6 +31,10 @@ Use the shared [license-audit skill](.agents/skills/license-audit/SKILL.md) for
 
 Declare restricted visibility at the module boundary and use `pub` for items in that module's API.
 
+Use imports for referenced symbols. Keep a short module qualifier or use an explicit alias when a bare name would obscure its origin or conflict with another symbol, such as `io::Error`.
+
+Start ordinary intra-crate imports at `crate`; avoid explicit `super::...` imports and qualified paths. Reserve `use super::*` for test modules.
+
 ## Waker Contract
 
 - Allow normal executor `Waker::clone` inside short state critical sections. Custom clone panic recovery and reentrancy are not general guarantees; do not require them in reviews unless an explicit local contract does.

CHANGELOG.md

@@ -30,6 +30,7 @@ All notable changes to this project will be documented in this file.
 ### Bug fixes
 
 * Prevent deadlocks in `singleflight::Group::work` and `try_work` when a duplicate key's destructor calls back into the same group.
+* Fix undefined behavior when sending or dropping a oneshot sender races with receiver completion or cancellation.
 
 ### Notable changes
 

asyncband/src/oneshot/mod.rs

@@ -302,6 +302,8 @@ impl<T> Channel<T> {
     ///
     /// Returns the waker and whether the receiver still owns allocation cleanup. If the receiver
     /// cancelled from `AWAKING`, this returns `false` and transfers cleanup to the caller.
+    /// The receiver may reclaim the allocation before this function returns, so this must take a
+    /// raw pointer rather than a reference that keeps the allocation borrowed for the entire call.
     ///
     /// # Safety
     ///
@@ -312,22 +314,32 @@ impl<T> Channel<T> {
     ///   method.
     /// * When publishing `MESSAGE`, the caller must own an initialized message that precedes the
     ///   release operation in this method.
+    /// * `channel_ptr` must point to that channel, and the caller must not access the channel after
+    ///   this call unless it returns `false` and transfers allocation cleanup to the caller.
     #[inline(always)]
-    unsafe fn finish_sender_awakening(&self, final_state: u8) -> (Waker, bool) {
+    unsafe fn finish_sender_awakening(
+        channel_ptr: NonNull<Self>,
+        final_state: u8,
+    ) -> (Waker, bool) {
         debug_assert!(matches!(final_state, MESSAGE | DISCONNECTED));
 
+        // SAFETY: The RECEIVING-to-AWAKING transition keeps the allocation live until this swap.
+        // Receiver cancellation can only transfer cleanup to this sender. This local reference
+        // must not be used after the swap, which can let the receiver reclaim the allocation.
+        let channel = unsafe { channel_ptr.as_ref() };
+
         // ORDERING: The caller's Release RMW read RECEIVING with a Relaxed load. Acquire
         // synchronizes that read with the receiver's Release publication before taking the waker.
         fence(Ordering::Acquire);
 
         // SAFETY: The caller's RECEIVING-to-AWAKING transition transferred exclusive ownership of
         // the initialized waker to the sender.
-        let waker = unsafe { self.take_waker() };
+        let waker = unsafe { channel.take_waker() };
 
         // ORDERING: Release publishes the message or disconnected state when this replaces
         // AWAKING. The RMW's load half is Relaxed; if it reads a receiver-written DISCONNECTED, the
         // conditional Acquire below completes the reverse allocation-ownership handoff.
-        let previous_state = self.state.swap(final_state, Ordering::Release);
+        let previous_state = channel.state.swap(final_state, Ordering::Release);
         if matches!(previous_state, AWAKING) {
             (waker, true)
         } else {

asyncband/src/oneshot/receiver.rs

@@ -31,16 +31,16 @@ use std::sync::atomic::fence;
 use std::task::Context;
 use std::task::Poll;
 
-use super::AWAKING;
-use super::Channel;
-use super::DISCONNECTED;
-use super::EMPTY;
-use super::MESSAGE;
-use super::RECEIVING;
+use crate::oneshot::AWAKING;
+use crate::oneshot::Channel;
+use crate::oneshot::DISCONNECTED;
+use crate::oneshot::EMPTY;
+use crate::oneshot::MESSAGE;
+use crate::oneshot::RECEIVING;
 #[cfg(doc)]
-use super::Sender;
-use super::deallocate_empty_channel;
-use super::drop_message_and_deallocate_channel;
+use crate::oneshot::Sender;
+use crate::oneshot::deallocate_empty_channel;
+use crate::oneshot::drop_message_and_deallocate_channel;
 
 /// Receives a value from the associated [`Sender`].
 ///

asyncband/src/oneshot/sender.rs

@@ -29,15 +29,15 @@ use std::ptr::NonNull;
 use std::sync::atomic::Ordering;
 use std::sync::atomic::fence;
 
-use super::Channel;
-use super::DISCONNECTED;
-use super::EMPTY;
-use super::MESSAGE;
-use super::RECEIVING;
+use crate::oneshot::Channel;
+use crate::oneshot::DISCONNECTED;
+use crate::oneshot::EMPTY;
+use crate::oneshot::MESSAGE;
+use crate::oneshot::RECEIVING;
 #[cfg(doc)]
-use super::Receiver;
-use super::deallocate_empty_channel;
-use super::drop_message_and_deallocate_channel;
+use crate::oneshot::Receiver;
+use crate::oneshot::deallocate_empty_channel;
+use crate::oneshot::drop_message_and_deallocate_channel;
 
 /// Sends a value to the associated [`Receiver`].
 pub struct Sender<T> {
@@ -61,9 +61,9 @@ impl<T> Sender<T> {
         let sender = ManuallyDrop::new(self);
         let channel_ptr = sender.channel_ptr;
 
-        // SAFETY: The channel exists on the heap for the entire duration of this method, and we
-        // only ever acquire shared references to it. Dropping the receiver does not immediately
-        // free the channel.
+        // SAFETY: The live sender keeps the allocation valid until it publishes MESSAGE. That
+        // publication can let the receiver reclaim it, so no channel access may follow unless
+        // receiver cancellation transfers allocation cleanup to this sender.
         let channel = unsafe { channel_ptr.as_ref() };
 
         // Write the message into the channel on the heap.
@@ -91,7 +91,7 @@ impl<T> Sender<T> {
                 // exclusive ownership of the published waker to this sender. The message was
                 // initialized before the RMW and is ready to publish.
                 let (waker, receiver_owns_allocation) =
-                    unsafe { channel.finish_sender_awakening(MESSAGE) };
+                    unsafe { Channel::finish_sender_awakening(channel_ptr, MESSAGE) };
                 if receiver_owns_allocation {
                     waker.wake();
                 } else {
@@ -159,6 +159,8 @@ impl<T> Drop for Sender<T> {
         // * If we are in the DISCONNECTED state, then the receiver either received the message,
         //   making this statement unreachable, or was dropped and observed that our side was still
         //   alive, and thus didn't free the channel.
+        // Publishing DISCONNECTED can let the receiver reclaim the allocation; no channel access
+        // may follow unless receiver cancellation transfers allocation cleanup to this sender.
         let channel = unsafe { self.channel_ptr.as_ref() };
 
         // Publish the channel's disconnected state directly, or begin awakening a receiving task:
@@ -180,7 +182,7 @@ impl<T> Drop for Sender<T> {
                 // exclusive ownership of the published waker to this sender. Sender::drop has no
                 // message to publish.
                 let (waker, receiver_owns_allocation) =
-                    unsafe { channel.finish_sender_awakening(DISCONNECTED) };
+                    unsafe { Channel::finish_sender_awakening(self.channel_ptr, DISCONNECTED) };
                 if receiver_owns_allocation {
                     waker.wake();
                 } else {

asyncband/src/oneshot/tests.rs

@@ -32,10 +32,14 @@ use std::task::Context;
 use std::task::Poll;
 use std::time::Duration;
 
-use self::support::DropProbe;
-use self::support::WakerProbe;
-use self::support::spawn_named;
 use crate::oneshot;
+use crate::oneshot::Channel;
+use crate::oneshot::MESSAGE;
+use crate::oneshot::RECEIVING;
+use crate::oneshot::drop_message_and_deallocate_channel;
+use crate::oneshot::tests::support::DropProbe;
+use crate::oneshot::tests::support::WakerProbe;
+use crate::oneshot::tests::support::spawn_named;
 
 // These tests stay next to the implementation because they inspect private state.
 
@@ -56,10 +60,7 @@ fn poll_returns_while_sender_owns_waker() {
     let channel = unsafe { channel_ptr.as_ref() };
     unsafe { channel.write_message(1234u128) };
     // Pause the synthetic sender in AWAKING immediately after it takes the stored waker.
-    assert_eq!(
-        channel.state.fetch_add(1, Ordering::Release),
-        super::RECEIVING
-    );
+    assert_eq!(channel.state.fetch_add(1, Ordering::Release), RECEIVING);
 
     let (started_tx, started_rx) = mpsc::sync_channel(0);
     let (result_tx, result_rx) = mpsc::sync_channel(1);
@@ -78,7 +79,7 @@ fn poll_returns_while_sender_owns_waker() {
     // SAFETY: fetch_add observed RECEIVING and changed it to AWAKING after the message and waker
     // were initialized.
     let (sender_waker, receiver_owns_allocation) =
-        unsafe { channel.finish_sender_awakening(super::MESSAGE) };
+        unsafe { Channel::finish_sender_awakening(channel_ptr, MESSAGE) };
     assert!(receiver_owns_allocation);
     sender_waker.wake();
     assert_eq!(stored_probe.wake_count(), 1);
@@ -124,10 +125,7 @@ fn drop_transfers_cleanup_while_sender_owns_waker() {
     let channel = unsafe { channel_ptr.as_ref() };
     unsafe { channel.write_message(message) };
     // Pause the synthetic sender in AWAKING immediately after it takes the stored waker.
-    assert_eq!(
-        channel.state.fetch_add(1, Ordering::Release),
-        super::RECEIVING
-    );
+    assert_eq!(channel.state.fetch_add(1, Ordering::Release), RECEIVING);
 
     let (started_tx, started_rx) = mpsc::sync_channel(0);
     let (done_tx, done_rx) = mpsc::sync_channel(1);
@@ -144,14 +142,14 @@ fn drop_transfers_cleanup_while_sender_owns_waker() {
     // SAFETY: fetch_add observed RECEIVING and changed it to AWAKING after the message and waker
     // were initialized.
     let (sender_waker, receiver_owns_allocation) =
-        unsafe { channel.finish_sender_awakening(super::MESSAGE) };
+        unsafe { Channel::finish_sender_awakening(channel_ptr, MESSAGE) };
     if receiver_owns_allocation {
         sender_waker.wake();
     } else {
         // SAFETY: Receiver cancellation transferred allocation cleanup to the synthetic sender, so
         // the original pointer provenance may be reclaimed as a Box. The message is initialized and
         // the waker was moved out above.
-        unsafe { super::drop_message_and_deallocate_channel(channel_ptr) };
+        unsafe { drop_message_and_deallocate_channel(channel_ptr) };
         drop(sender_waker);
     }
 

tests-integration/tests/oneshot_test/main.rs

@@ -300,6 +300,26 @@ fn poll_then_drop_receiver_during_send() {
     assert_eq!(message_drop_count.load(Ordering::Relaxed), 1);
 }
 
+#[test]
+fn poll_then_drop_receiver_during_sender_drop() {
+    let (sender, receiver) = oneshot::channel::<i32>();
+    let mut receiver = receiver.into_future();
+    let (waker, waker_probe) = WakerProbe::new();
+    let mut context = Context::from_waker(&waker);
+
+    assert!(matches!(
+        Pin::new(&mut receiver).poll(&mut context),
+        Poll::Pending
+    ));
+
+    let sender_thread = spawn_named("sender", move || drop(sender));
+    drop(receiver);
+    sender_thread.join().unwrap();
+
+    // Either endpoint may claim the stored waker; neither may leave it in the freed channel.
+    assert_eq!(WakerProbe::live_waker_count(&waker_probe), 1);
+}
+
 #[test]
 fn concurrent_send_and_try_recv_to_completion() {
     let (sender, receiver) = oneshot::channel::<i32>();