diff --git a/AGENTS.md b/AGENTS.md index f18d628e..3029eabb 100644 --- a/AGENTS.md +++ b/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. diff --git a/CHANGELOG.md b/CHANGELOG.md index 6f226004..bd5e84c9 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/asyncband/src/oneshot/mod.rs b/asyncband/src/oneshot/mod.rs index fd7339cf..7859d7bc 100644 --- a/asyncband/src/oneshot/mod.rs +++ b/asyncband/src/oneshot/mod.rs @@ -302,6 +302,8 @@ impl Channel { /// /// 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 Channel { /// 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, + 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 { diff --git a/asyncband/src/oneshot/receiver.rs b/asyncband/src/oneshot/receiver.rs index 4b397079..9f0228b3 100644 --- a/asyncband/src/oneshot/receiver.rs +++ b/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`]. /// diff --git a/asyncband/src/oneshot/sender.rs b/asyncband/src/oneshot/sender.rs index af161d40..738c4dda 100644 --- a/asyncband/src/oneshot/sender.rs +++ b/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 { @@ -61,9 +61,9 @@ impl Sender { 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 Sender { // 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 Drop for Sender { // * 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 Drop for Sender { // 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 { diff --git a/asyncband/src/oneshot/tests.rs b/asyncband/src/oneshot/tests.rs index a3011ab1..4e293261 100644 --- a/asyncband/src/oneshot/tests.rs +++ b/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); } diff --git a/tests-integration/tests/oneshot_test/main.rs b/tests-integration/tests/oneshot_test/main.rs index 9fdd122f..69db4ddb 100644 --- a/tests-integration/tests/oneshot_test/main.rs +++ b/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::(); + 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::();