From ccf30c5cc300114de7d71ff887228068ba65db83 Mon Sep 17 00:00:00 2001 From: Alice Ryhl Date: Mon, 10 Oct 2022 10:02:46 +0200 Subject: [PATCH 1/8] macros: don't take ownership of futures in macros --- tokio/src/future/poll_fn.rs | 21 ++++++++++++++++----- tokio/src/macros/join.rs | 8 +++++++- tokio/src/macros/select.rs | 8 +++++++- tokio/src/macros/try_join.rs | 8 +++++++- 4 files changed, 37 insertions(+), 8 deletions(-) diff --git a/tokio/src/future/poll_fn.rs b/tokio/src/future/poll_fn.rs index 041e3d70a8d..ab1091d9289 100644 --- a/tokio/src/future/poll_fn.rs +++ b/tokio/src/future/poll_fn.rs @@ -6,20 +6,23 @@ use std::fmt; use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; +use std::marker::PhantomPinned; /// Future for the [`poll_fn`] function. pub struct PollFn { f: F, + _pinned: PhantomPinned, } -impl Unpin for PollFn {} - /// Creates a new future wrapping around a function returning [`Poll`]. pub fn poll_fn(f: F) -> PollFn where F: FnMut(&mut Context<'_>) -> Poll, { - PollFn { f } + PollFn { + f, + _pinned: PhantomPinned, + } } impl fmt::Debug for PollFn { @@ -34,7 +37,15 @@ where { type Output = T; - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - (self.f)(cx) + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + // SAFETY: We never construct a `Pin<&mut F>` anywhere, so accessing `f` + // mutably in an unpinned way is sound. + // + // This is, strictly speaking, not necessary. We could make `PollFn` + // unconditionally `Unpin` and avoid this unsafe. However, making this + // struct `!Unpin` mitigates the issues described here: + // + let me = unsafe { Pin::into_inner_unchecked(self) }; + (me.f)(cx) } } diff --git a/tokio/src/macros/join.rs b/tokio/src/macros/join.rs index 9697936c895..2c08bf08b9a 100644 --- a/tokio/src/macros/join.rs +++ b/tokio/src/macros/join.rs @@ -74,6 +74,12 @@ macro_rules! join { // the requirement of `Pin::new_unchecked` called below. let mut futures = ( $( maybe_done($e), )* ); + // This assignment makes sure that the `poll_fn` closure only has a + // reference to the futures, instead of taking ownership of them. This + // mitigates the issue described in + // + let mut futures = &mut futures; + // Each time the future created by poll_fn is polled, a different future will be polled first // to ensure every future passed to join! gets a chance to make progress even if // one of the futures consumes the whole budget. @@ -106,7 +112,7 @@ macro_rules! join { to_run -= 1; // Extract the future for this branch from the tuple. - let ( $($skip,)* fut, .. ) = &mut futures; + let ( $($skip,)* fut, .. ) = &mut *futures; // Safety: future is stored on the stack above // and never moved. diff --git a/tokio/src/macros/select.rs b/tokio/src/macros/select.rs index f38aee0f20b..fa68801fb94 100644 --- a/tokio/src/macros/select.rs +++ b/tokio/src/macros/select.rs @@ -462,6 +462,12 @@ macro_rules! select { // satisfy the requirement of `Pin::new_unchecked` called below. let mut futures = ( $( $fut , )+ ); + // This assignment makes sure that the `poll_fn` closure only has a + // reference to the futures, instead of taking ownership of them. + // This mitigates the issue described in + // + let mut futures = &mut *futures; + $crate::macros::support::poll_fn(|cx| { // Track if any branch returns pending. If no branch completes // **or** returns pending, this implies that all branches are @@ -497,7 +503,7 @@ macro_rules! select { // Extract the future for this branch from the // tuple - let ( $($skip,)* fut, .. ) = &mut futures; + let ( $($skip,)* fut, .. ) = &mut *futures; // Safety: future is stored on the stack above // and never moved. diff --git a/tokio/src/macros/try_join.rs b/tokio/src/macros/try_join.rs index c80395c9097..ca4b8f6dc10 100644 --- a/tokio/src/macros/try_join.rs +++ b/tokio/src/macros/try_join.rs @@ -120,6 +120,12 @@ macro_rules! try_join { // the requirement of `Pin::new_unchecked` called below. let mut futures = ( $( maybe_done($e), )* ); + // This assignment makes sure that the `poll_fn` closure only has a + // reference to the futures, instead of taking ownership of them. This + // mitigates the issue described in + // + let mut futures = &mut futures; + // Each time the future created by poll_fn is polled, a different future will be polled first // to ensure every future passed to join! gets a chance to make progress even if // one of the futures consumes the whole budget. @@ -152,7 +158,7 @@ macro_rules! try_join { to_run -= 1; // Extract the future for this branch from the tuple. - let ( $($skip,)* fut, .. ) = &mut futures; + let ( $($skip,)* fut, .. ) = &mut *futures; // Safety: future is stored on the stack above // and never moved. From 4a4c137cde03f2d2ea544b8a8994948002eee907 Mon Sep 17 00:00:00 2001 From: Alice Ryhl Date: Mon, 10 Oct 2022 10:13:08 +0200 Subject: [PATCH 2/8] fix --- tokio/src/future/poll_fn.rs | 2 +- tokio/src/macros/select.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tokio/src/future/poll_fn.rs b/tokio/src/future/poll_fn.rs index ab1091d9289..a21972a23d4 100644 --- a/tokio/src/future/poll_fn.rs +++ b/tokio/src/future/poll_fn.rs @@ -4,9 +4,9 @@ use std::fmt; use std::future::Future; +use std::marker::PhantomPinned; use std::pin::Pin; use std::task::{Context, Poll}; -use std::marker::PhantomPinned; /// Future for the [`poll_fn`] function. pub struct PollFn { diff --git a/tokio/src/macros/select.rs b/tokio/src/macros/select.rs index fa68801fb94..b79a4fcd1f8 100644 --- a/tokio/src/macros/select.rs +++ b/tokio/src/macros/select.rs @@ -466,7 +466,7 @@ macro_rules! select { // reference to the futures, instead of taking ownership of them. // This mitigates the issue described in // - let mut futures = &mut *futures; + let mut futures = &mut futures; $crate::macros::support::poll_fn(|cx| { // Track if any branch returns pending. If no branch completes From 12fd8ceb40f5703003f358ade35ce8c66a99dca4 Mon Sep 17 00:00:00 2001 From: Alice Ryhl Date: Mon, 10 Oct 2022 10:32:22 +0200 Subject: [PATCH 3/8] Fix size tests --- tokio/tests/macros_join.rs | 4 ++-- tokio/tests/macros_select.rs | 6 +++--- tokio/tests/macros_try_join.rs | 4 ++-- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/tokio/tests/macros_join.rs b/tokio/tests/macros_join.rs index 0efbd83c3ea..ac55707a42f 100644 --- a/tokio/tests/macros_join.rs +++ b/tokio/tests/macros_join.rs @@ -72,14 +72,14 @@ fn join_size() { let ready = future::ready(0i32); tokio::join!(ready) }; - assert_eq!(mem::size_of_val(&fut), 20); + assert_eq!(mem::size_of_val(&fut), 32); let fut = async { let ready1 = future::ready(0i32); let ready2 = future::ready(0i32); tokio::join!(ready1, ready2) }; - assert_eq!(mem::size_of_val(&fut), 32); + assert_eq!(mem::size_of_val(&fut), 48); } async fn non_cooperative_task(permits: Arc) -> usize { diff --git a/tokio/tests/macros_select.rs b/tokio/tests/macros_select.rs index e642a5d9ff4..57cd98d6379 100644 --- a/tokio/tests/macros_select.rs +++ b/tokio/tests/macros_select.rs @@ -219,7 +219,7 @@ async fn struct_size() { } }; - assert!(mem::size_of_val(&fut) <= 32); + assert_eq!(mem::size_of_val(&fut), 40); let fut = async { let ready1 = future::ready(0i32); @@ -231,7 +231,7 @@ async fn struct_size() { } }; - assert!(mem::size_of_val(&fut) <= 40); + assert_eq!(mem::size_of_val(&fut), 48); let fut = async { let ready1 = future::ready(0i32); @@ -245,7 +245,7 @@ async fn struct_size() { } }; - assert!(mem::size_of_val(&fut) <= 48); + assert_eq!(mem::size_of_val(&fut), 56); } #[maybe_tokio_test] diff --git a/tokio/tests/macros_try_join.rs b/tokio/tests/macros_try_join.rs index 681b1b7bc68..e9c468272f3 100644 --- a/tokio/tests/macros_try_join.rs +++ b/tokio/tests/macros_try_join.rs @@ -96,14 +96,14 @@ fn join_size() { let ready = future::ready(ok(0i32)); tokio::try_join!(ready) }; - assert_eq!(mem::size_of_val(&fut), 20); + assert_eq!(mem::size_of_val(&fut), 32); let fut = async { let ready1 = future::ready(ok(0i32)); let ready2 = future::ready(ok(0i32)); tokio::try_join!(ready1, ready2) }; - assert_eq!(mem::size_of_val(&fut), 32); + assert_eq!(mem::size_of_val(&fut), 48); } fn ok(val: T) -> Result { From f7346205fa73c11146317116521c78d7d51dde55 Mon Sep 17 00:00:00 2001 From: Alice Ryhl Date: Mon, 10 Oct 2022 10:46:57 +0200 Subject: [PATCH 4/8] Don't test macro sizes on 32-bit platforms --- tokio/tests/macros_join.rs | 1 + tokio/tests/macros_select.rs | 1 + tokio/tests/macros_try_join.rs | 1 + 3 files changed, 3 insertions(+) diff --git a/tokio/tests/macros_join.rs b/tokio/tests/macros_join.rs index ac55707a42f..9e8aacf1560 100644 --- a/tokio/tests/macros_join.rs +++ b/tokio/tests/macros_join.rs @@ -64,6 +64,7 @@ async fn two_await() { } #[test] +#[cfg(target_pointer_width = "64")] fn join_size() { use futures::future; use std::mem; diff --git a/tokio/tests/macros_select.rs b/tokio/tests/macros_select.rs index 57cd98d6379..e54e223ba5d 100644 --- a/tokio/tests/macros_select.rs +++ b/tokio/tests/macros_select.rs @@ -207,6 +207,7 @@ async fn nested() { } #[maybe_tokio_test] +#[cfg(target_pointer_width = "64")] async fn struct_size() { use futures::future; use std::mem; diff --git a/tokio/tests/macros_try_join.rs b/tokio/tests/macros_try_join.rs index e9c468272f3..209516bb998 100644 --- a/tokio/tests/macros_try_join.rs +++ b/tokio/tests/macros_try_join.rs @@ -88,6 +88,7 @@ async fn err_abort_early() { } #[test] +#[cfg(target_pointer_width = "64")] fn join_size() { use futures::future; use std::mem; From da99e6d7d8aaa1e6fce6527f85a64bbcc5945e7e Mon Sep 17 00:00:00 2001 From: Alice Ryhl Date: Mon, 10 Oct 2022 10:52:08 +0200 Subject: [PATCH 5/8] Fix unused import --- tokio/tests/macros_join.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/tokio/tests/macros_join.rs b/tokio/tests/macros_join.rs index 9e8aacf1560..c289c00679a 100644 --- a/tokio/tests/macros_join.rs +++ b/tokio/tests/macros_join.rs @@ -3,6 +3,7 @@ use std::sync::Arc; #[cfg(tokio_wasm_not_wasi)] +#[cfg(target_pointer_width = "64")] use wasm_bindgen_test::wasm_bindgen_test as test; #[cfg(tokio_wasm_not_wasi)] use wasm_bindgen_test::wasm_bindgen_test as maybe_tokio_test; From 0bae230ea2eb36b0ebd2d52c028b8219bfd3d3b8 Mon Sep 17 00:00:00 2001 From: Alice Ryhl Date: Mon, 10 Oct 2022 12:35:16 +0200 Subject: [PATCH 6/8] Add comments --- tokio/src/future/poll_fn.rs | 24 +++++++++++++----------- tokio/src/macros/join.rs | 4 ++++ tokio/src/macros/select.rs | 4 ++++ tokio/src/macros/try_join.rs | 4 ++++ tokio/tests/async_send_sync.rs | 16 ++++++++++++++++ 5 files changed, 41 insertions(+), 11 deletions(-) diff --git a/tokio/src/future/poll_fn.rs b/tokio/src/future/poll_fn.rs index a21972a23d4..9b09aee5ea7 100644 --- a/tokio/src/future/poll_fn.rs +++ b/tokio/src/future/poll_fn.rs @@ -4,14 +4,24 @@ use std::fmt; use std::future::Future; -use std::marker::PhantomPinned; use std::pin::Pin; use std::task::{Context, Poll}; +// This struct is intentionally `!Unpin` when `F` is `!Unpin`. This is to +// mitigate the issue where rust puts noalias on mutable references to the +// `PollFn` type if it is `Unpin`. If the closure has ownership of a future, +// then this "leaks" and the future is affected by noalias too, which we don't +// want. +// +// See this thread for more information: +// +// +// The fact that `PollFn` is not `Unpin` when it shouldn't be is tested in +// `tests/async_send_sync.rs`. + /// Future for the [`poll_fn`] function. pub struct PollFn { f: F, - _pinned: PhantomPinned, } /// Creates a new future wrapping around a function returning [`Poll`]. @@ -19,10 +29,7 @@ pub fn poll_fn(f: F) -> PollFn where F: FnMut(&mut Context<'_>) -> Poll, { - PollFn { - f, - _pinned: PhantomPinned, - } + PollFn { f } } impl fmt::Debug for PollFn { @@ -40,11 +47,6 @@ where fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { // SAFETY: We never construct a `Pin<&mut F>` anywhere, so accessing `f` // mutably in an unpinned way is sound. - // - // This is, strictly speaking, not necessary. We could make `PollFn` - // unconditionally `Unpin` and avoid this unsafe. However, making this - // struct `!Unpin` mitigates the issues described here: - // let me = unsafe { Pin::into_inner_unchecked(self) }; (me.f)(cx) } diff --git a/tokio/src/macros/join.rs b/tokio/src/macros/join.rs index 2c08bf08b9a..7e85203c0c4 100644 --- a/tokio/src/macros/join.rs +++ b/tokio/src/macros/join.rs @@ -72,6 +72,10 @@ macro_rules! join { // Safety: nothing must be moved out of `futures`. This is to satisfy // the requirement of `Pin::new_unchecked` called below. + // + // We can't use the `pin!` macro for this because `futures` is a tuple + // and the standard library provides no way to pin-project to the fields + // of a tuple. let mut futures = ( $( maybe_done($e), )* ); // This assignment makes sure that the `poll_fn` closure only has a diff --git a/tokio/src/macros/select.rs b/tokio/src/macros/select.rs index b79a4fcd1f8..7ba04a00768 100644 --- a/tokio/src/macros/select.rs +++ b/tokio/src/macros/select.rs @@ -460,6 +460,10 @@ macro_rules! select { let mut output = { // Safety: Nothing must be moved out of `futures`. This is to // satisfy the requirement of `Pin::new_unchecked` called below. + // + // We can't use the `pin!` macro for this because `futures` is a + // tuple and the standard library provides no way to pin-project to + // the fields of a tuple. let mut futures = ( $( $fut , )+ ); // This assignment makes sure that the `poll_fn` closure only has a diff --git a/tokio/src/macros/try_join.rs b/tokio/src/macros/try_join.rs index ca4b8f6dc10..597cd5df021 100644 --- a/tokio/src/macros/try_join.rs +++ b/tokio/src/macros/try_join.rs @@ -118,6 +118,10 @@ macro_rules! try_join { // Safety: nothing must be moved out of `futures`. This is to satisfy // the requirement of `Pin::new_unchecked` called below. + // + // We can't use the `pin!` macro for this because `futures` is a tuple + // and the standard library provides no way to pin-project to the fields + // of a tuple. let mut futures = ( $( maybe_done($e), )* ); // This assignment makes sure that the `poll_fn` closure only has a diff --git a/tokio/tests/async_send_sync.rs b/tokio/tests/async_send_sync.rs index ecfa9e27b85..e50ff3a010f 100644 --- a/tokio/tests/async_send_sync.rs +++ b/tokio/tests/async_send_sync.rs @@ -5,6 +5,7 @@ use std::cell::Cell; use std::future::Future; use std::io::SeekFrom; +use std::marker::PhantomPinned; use std::net::SocketAddr; use std::pin::Pin; use std::rc::Rc; @@ -136,6 +137,21 @@ macro_rules! cfg_not_wasi { } } +// Manualy re-implementation of `async_assert_fn` for `poll_fn`. The macro +// doesn't work for this particular case because constructing the closure +// is too complicated. +const _: fn() = || { + let pinned = std::marker::PhantomPinned; + let f = tokio::macros::support::poll_fn(move |_| { + // Use `pinned` to take ownership of it. + let _ = &pinned; + std::task::Poll::Pending::<()> + }); + require_send(&f); + require_sync(&f); + AmbiguousIfUnpin::some_item(&f); +}; + cfg_not_wasi! { mod fs { use super::*; From 8cf8a03a9beb32a7e5f39291366df1c69a22f6bf Mon Sep 17 00:00:00 2001 From: Alice Ryhl Date: Mon, 10 Oct 2022 12:40:25 +0200 Subject: [PATCH 7/8] Remove unused import --- tokio/tests/async_send_sync.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/tokio/tests/async_send_sync.rs b/tokio/tests/async_send_sync.rs index e50ff3a010f..2ef02c7e8d8 100644 --- a/tokio/tests/async_send_sync.rs +++ b/tokio/tests/async_send_sync.rs @@ -5,7 +5,6 @@ use std::cell::Cell; use std::future::Future; use std::io::SeekFrom; -use std::marker::PhantomPinned; use std::net::SocketAddr; use std::pin::Pin; use std::rc::Rc; From 624dff4529f89a2aea59adbd7bef45c6fe5b036a Mon Sep 17 00:00:00 2001 From: Alice Ryhl Date: Mon, 10 Oct 2022 12:54:11 +0200 Subject: [PATCH 8/8] Comment on why pin-project is not used --- tokio/src/future/poll_fn.rs | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tokio/src/future/poll_fn.rs b/tokio/src/future/poll_fn.rs index 9b09aee5ea7..074d9438eb4 100644 --- a/tokio/src/future/poll_fn.rs +++ b/tokio/src/future/poll_fn.rs @@ -45,8 +45,15 @@ where type Output = T; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - // SAFETY: We never construct a `Pin<&mut F>` anywhere, so accessing `f` + // Safety: We never construct a `Pin<&mut F>` anywhere, so accessing `f` // mutably in an unpinned way is sound. + // + // This use of unsafe cannot be replaced with the pin-project macro + // because: + // * If we put `#[pin]` on the field, then it gives us a `Pin<&mut F>`, + // which we can't use to call the closure. + // * If we don't put `#[pin]` on the field, then it makes `PollFn` be + // unconditionally `Unpin`, which we also don't want. let me = unsafe { Pin::into_inner_unchecked(self) }; (me.f)(cx) }