diff options
author | Daniel Baumann <daniel.baumann@progress-linux.org> | 2024-05-30 03:57:31 +0000 |
---|---|---|
committer | Daniel Baumann <daniel.baumann@progress-linux.org> | 2024-05-30 03:57:31 +0000 |
commit | dc0db358abe19481e475e10c32149b53370f1a1c (patch) | |
tree | ab8ce99c4b255ce46f99ef402c27916055b899ee /vendor/tokio/src | |
parent | Releasing progress-linux version 1.71.1+dfsg1-2~progress7.99u1. (diff) | |
download | rustc-dc0db358abe19481e475e10c32149b53370f1a1c.tar.xz rustc-dc0db358abe19481e475e10c32149b53370f1a1c.zip |
Merging upstream version 1.72.1+dfsg1.
Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
Diffstat (limited to 'vendor/tokio/src')
297 files changed, 33090 insertions, 10007 deletions
diff --git a/vendor/tokio/src/blocking.rs b/vendor/tokio/src/blocking.rs index f88b1db11..f172399d5 100644 --- a/vendor/tokio/src/blocking.rs +++ b/vendor/tokio/src/blocking.rs @@ -1,5 +1,11 @@ cfg_rt! { pub(crate) use crate::runtime::spawn_blocking; + + cfg_fs! { + #[allow(unused_imports)] + pub(crate) use crate::runtime::spawn_mandatory_blocking; + } + pub(crate) use crate::task::JoinHandle; } @@ -16,7 +22,16 @@ cfg_not_rt! { { assert_send_sync::<JoinHandle<std::cell::Cell<()>>>(); panic!("requires the `rt` Tokio feature flag") + } + cfg_fs! { + pub(crate) fn spawn_mandatory_blocking<F, R>(_f: F) -> Option<JoinHandle<R>> + where + F: FnOnce() -> R + Send + 'static, + R: Send + 'static, + { + panic!("requires the `rt` Tokio feature flag") + } } pub(crate) struct JoinHandle<R> { diff --git a/vendor/tokio/src/doc/mod.rs b/vendor/tokio/src/doc/mod.rs index 12c224702..7992fc6bc 100644 --- a/vendor/tokio/src/doc/mod.rs +++ b/vendor/tokio/src/doc/mod.rs @@ -17,7 +17,7 @@ /// will ever accidentally use it. /// /// [`never` type]: https://doc.rust-lang.org/std/primitive.never.html +#[derive(Debug)] pub enum NotDefinedHere {} pub mod os; -pub mod winapi; diff --git a/vendor/tokio/src/doc/os.rs b/vendor/tokio/src/doc/os.rs index 0ddf86959..9404d4491 100644 --- a/vendor/tokio/src/doc/os.rs +++ b/vendor/tokio/src/doc/os.rs @@ -13,7 +13,7 @@ pub mod windows { /// See [std::os::windows::io::AsRawHandle](https://doc.rust-lang.org/std/os/windows/io/trait.AsRawHandle.html) pub trait AsRawHandle { - /// See [std::os::windows::io::FromRawHandle::from_raw_handle](https://doc.rust-lang.org/std/os/windows/io/trait.AsRawHandle.html#tymethod.as_raw_handle) + /// See [std::os::windows::io::AsRawHandle::as_raw_handle](https://doc.rust-lang.org/std/os/windows/io/trait.AsRawHandle.html#tymethod.as_raw_handle) fn as_raw_handle(&self) -> RawHandle; } @@ -22,5 +22,44 @@ pub mod windows { /// See [std::os::windows::io::FromRawHandle::from_raw_handle](https://doc.rust-lang.org/std/os/windows/io/trait.FromRawHandle.html#tymethod.from_raw_handle) unsafe fn from_raw_handle(handle: RawHandle) -> Self; } + + /// See [std::os::windows::io::RawSocket](https://doc.rust-lang.org/std/os/windows/io/type.RawSocket.html) + pub type RawSocket = crate::doc::NotDefinedHere; + + /// See [std::os::windows::io::AsRawSocket](https://doc.rust-lang.org/std/os/windows/io/trait.AsRawSocket.html) + pub trait AsRawSocket { + /// See [std::os::windows::io::AsRawSocket::as_raw_socket](https://doc.rust-lang.org/std/os/windows/io/trait.AsRawSocket.html#tymethod.as_raw_socket) + fn as_raw_socket(&self) -> RawSocket; + } + + /// See [std::os::windows::io::FromRawSocket](https://doc.rust-lang.org/std/os/windows/io/trait.FromRawSocket.html) + pub trait FromRawSocket { + /// See [std::os::windows::io::FromRawSocket::from_raw_socket](https://doc.rust-lang.org/std/os/windows/io/trait.FromRawSocket.html#tymethod.from_raw_socket) + unsafe fn from_raw_socket(sock: RawSocket) -> Self; + } + + /// See [std::os::windows::io::IntoRawSocket](https://doc.rust-lang.org/std/os/windows/io/trait.IntoRawSocket.html) + pub trait IntoRawSocket { + /// See [std::os::windows::io::IntoRawSocket::into_raw_socket](https://doc.rust-lang.org/std/os/windows/io/trait.IntoRawSocket.html#tymethod.into_raw_socket) + fn into_raw_socket(self) -> RawSocket; + } + + /// See [std::os::windows::io::BorrowedHandle](https://doc.rust-lang.org/std/os/windows/io/struct.BorrowedHandle.html) + pub type BorrowedHandle<'handle> = crate::doc::NotDefinedHere; + + /// See [std::os::windows::io::AsHandle](https://doc.rust-lang.org/std/os/windows/io/trait.AsHandle.html) + pub trait AsHandle { + /// See [std::os::windows::io::AsHandle::as_handle](https://doc.rust-lang.org/std/os/windows/io/trait.AsHandle.html#tymethod.as_handle) + fn as_handle(&self) -> BorrowedHandle<'_>; + } + + /// See [std::os::windows::io::BorrowedSocket](https://doc.rust-lang.org/std/os/windows/io/struct.BorrowedSocket.html) + pub type BorrowedSocket<'socket> = crate::doc::NotDefinedHere; + + /// See [std::os::windows::io::AsSocket](https://doc.rust-lang.org/std/os/windows/io/trait.AsSocket.html) + pub trait AsSocket { + /// See [std::os::windows::io::AsSocket::as_socket](https://doc.rust-lang.org/std/os/windows/io/trait.AsSocket.html#tymethod.as_socket) + fn as_socket(&self) -> BorrowedSocket<'_>; + } } } diff --git a/vendor/tokio/src/doc/winapi.rs b/vendor/tokio/src/doc/winapi.rs deleted file mode 100644 index be68749e0..000000000 --- a/vendor/tokio/src/doc/winapi.rs +++ /dev/null @@ -1,66 +0,0 @@ -//! See [winapi]. -//! -//! [winapi]: https://docs.rs/winapi - -/// See [winapi::shared](https://docs.rs/winapi/*/winapi/shared/index.html). -pub mod shared { - /// See [winapi::shared::winerror](https://docs.rs/winapi/*/winapi/shared/winerror/index.html). - #[allow(non_camel_case_types)] - pub mod winerror { - /// See [winapi::shared::winerror::ERROR_ACCESS_DENIED][winapi] - /// - /// [winapi]: https://docs.rs/winapi/*/winapi/shared/winerror/constant.ERROR_ACCESS_DENIED.html - pub type ERROR_ACCESS_DENIED = crate::doc::NotDefinedHere; - - /// See [winapi::shared::winerror::ERROR_PIPE_BUSY][winapi] - /// - /// [winapi]: https://docs.rs/winapi/*/winapi/shared/winerror/constant.ERROR_PIPE_BUSY.html - pub type ERROR_PIPE_BUSY = crate::doc::NotDefinedHere; - - /// See [winapi::shared::winerror::ERROR_MORE_DATA][winapi] - /// - /// [winapi]: https://docs.rs/winapi/*/winapi/shared/winerror/constant.ERROR_MORE_DATA.html - pub type ERROR_MORE_DATA = crate::doc::NotDefinedHere; - } -} - -/// See [winapi::um](https://docs.rs/winapi/*/winapi/um/index.html). -pub mod um { - /// See [winapi::um::winbase](https://docs.rs/winapi/*/winapi/um/winbase/index.html). - #[allow(non_camel_case_types)] - pub mod winbase { - /// See [winapi::um::winbase::PIPE_TYPE_MESSAGE][winapi] - /// - /// [winapi]: https://docs.rs/winapi/*/winapi/um/winbase/constant.PIPE_TYPE_MESSAGE.html - pub type PIPE_TYPE_MESSAGE = crate::doc::NotDefinedHere; - - /// See [winapi::um::winbase::PIPE_TYPE_BYTE][winapi] - /// - /// [winapi]: https://docs.rs/winapi/*/winapi/um/winbase/constant.PIPE_TYPE_BYTE.html - pub type PIPE_TYPE_BYTE = crate::doc::NotDefinedHere; - - /// See [winapi::um::winbase::PIPE_CLIENT_END][winapi] - /// - /// [winapi]: https://docs.rs/winapi/*/winapi/um/winbase/constant.PIPE_CLIENT_END.html - pub type PIPE_CLIENT_END = crate::doc::NotDefinedHere; - - /// See [winapi::um::winbase::PIPE_SERVER_END][winapi] - /// - /// [winapi]: https://docs.rs/winapi/*/winapi/um/winbase/constant.PIPE_SERVER_END.html - pub type PIPE_SERVER_END = crate::doc::NotDefinedHere; - - /// See [winapi::um::winbase::SECURITY_IDENTIFICATION][winapi] - /// - /// [winapi]: https://docs.rs/winapi/*/winapi/um/winbase/constant.SECURITY_IDENTIFICATION.html - pub type SECURITY_IDENTIFICATION = crate::doc::NotDefinedHere; - } - - /// See [winapi::um::minwinbase](https://docs.rs/winapi/*/winapi/um/minwinbase/index.html). - #[allow(non_camel_case_types)] - pub mod minwinbase { - /// See [winapi::um::minwinbase::SECURITY_ATTRIBUTES][winapi] - /// - /// [winapi]: https://docs.rs/winapi/*/winapi/um/minwinbase/constant.SECURITY_ATTRIBUTES.html - pub type SECURITY_ATTRIBUTES = crate::doc::NotDefinedHere; - } -} diff --git a/vendor/tokio/src/fs/create_dir.rs b/vendor/tokio/src/fs/create_dir.rs index e03b04dc4..411969500 100644 --- a/vendor/tokio/src/fs/create_dir.rs +++ b/vendor/tokio/src/fs/create_dir.rs @@ -3,7 +3,7 @@ use crate::fs::asyncify; use std::io; use std::path::Path; -/// Creates a new, empty directory at the provided path +/// Creates a new, empty directory at the provided path. /// /// This is an async version of [`std::fs::create_dir`][std] /// diff --git a/vendor/tokio/src/fs/dir_builder.rs b/vendor/tokio/src/fs/dir_builder.rs index b1849344b..97168bff7 100644 --- a/vendor/tokio/src/fs/dir_builder.rs +++ b/vendor/tokio/src/fs/dir_builder.rs @@ -14,7 +14,7 @@ pub struct DirBuilder { /// Indicates whether to create parent directories if they are missing. recursive: bool, - /// Set the Unix mode for newly created directories. + /// Sets the Unix mode for newly created directories. #[cfg(unix)] pub(super) mode: Option<u32>, } diff --git a/vendor/tokio/src/fs/file.rs b/vendor/tokio/src/fs/file.rs index 5286e6c5c..74f91958d 100644 --- a/vendor/tokio/src/fs/file.rs +++ b/vendor/tokio/src/fs/file.rs @@ -20,16 +20,16 @@ use std::task::Poll; use std::task::Poll::*; #[cfg(test)] -use super::mocks::spawn_blocking; -#[cfg(test)] use super::mocks::JoinHandle; #[cfg(test)] use super::mocks::MockFile as StdFile; -#[cfg(not(test))] -use crate::blocking::spawn_blocking; +#[cfg(test)] +use super::mocks::{spawn_blocking, spawn_mandatory_blocking}; #[cfg(not(test))] use crate::blocking::JoinHandle; #[cfg(not(test))] +use crate::blocking::{spawn_blocking, spawn_mandatory_blocking}; +#[cfg(not(test))] use std::fs::File as StdFile; /// A reference to an open file on the filesystem. @@ -74,7 +74,7 @@ use std::fs::File as StdFile; /// # } /// ``` /// -/// Read the contents of a file into a buffer +/// Read the contents of a file into a buffer: /// /// ```no_run /// use tokio::fs::File; @@ -383,7 +383,7 @@ impl File { asyncify(move || std.metadata()).await } - /// Create a new `File` instance that shares the same underlying file handle + /// Creates a new `File` instance that shares the same underlying file handle /// as the existing `File` instance. Reads, writes, and seeks will affect both /// File instances simultaneously. /// @@ -399,6 +399,7 @@ impl File { /// # } /// ``` pub async fn try_clone(&self) -> io::Result<File> { + self.inner.lock().await.complete_inflight().await; let std = self.std.clone(); let std_file = asyncify(move || std.try_clone()).await?; Ok(File::from_std(std_file)) @@ -498,6 +499,7 @@ impl AsyncRead for File { cx: &mut Context<'_>, dst: &mut ReadBuf<'_>, ) -> Poll<io::Result<()>> { + ready!(crate::trace::trace_leaf(cx)); let me = self.get_mut(); let inner = me.inner.get_mut(); @@ -565,34 +567,36 @@ impl AsyncSeek for File { let me = self.get_mut(); let inner = me.inner.get_mut(); - loop { - match inner.state { - Busy(_) => panic!("must wait for poll_complete before calling start_seek"), - Idle(ref mut buf_cell) => { - let mut buf = buf_cell.take().unwrap(); + match inner.state { + Busy(_) => Err(io::Error::new( + io::ErrorKind::Other, + "other file operation is pending, call poll_complete before start_seek", + )), + Idle(ref mut buf_cell) => { + let mut buf = buf_cell.take().unwrap(); - // Factor in any unread data from the buf - if !buf.is_empty() { - let n = buf.discard_read(); + // Factor in any unread data from the buf + if !buf.is_empty() { + let n = buf.discard_read(); - if let SeekFrom::Current(ref mut offset) = pos { - *offset += n; - } + if let SeekFrom::Current(ref mut offset) = pos { + *offset += n; } + } - let std = me.std.clone(); + let std = me.std.clone(); - inner.state = Busy(spawn_blocking(move || { - let res = (&*std).seek(pos); - (Operation::Seek(res), buf) - })); - return Ok(()); - } + inner.state = Busy(spawn_blocking(move || { + let res = (&*std).seek(pos); + (Operation::Seek(res), buf) + })); + Ok(()) } } } fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> { + ready!(crate::trace::trace_leaf(cx)); let inner = self.inner.get_mut(); loop { @@ -628,6 +632,7 @@ impl AsyncWrite for File { cx: &mut Context<'_>, src: &[u8], ) -> Poll<io::Result<usize>> { + ready!(crate::trace::trace_leaf(cx)); let me = self.get_mut(); let inner = me.inner.get_mut(); @@ -649,7 +654,7 @@ impl AsyncWrite for File { let n = buf.copy_from(src); let std = me.std.clone(); - inner.state = Busy(spawn_blocking(move || { + let blocking_task_join_handle = spawn_mandatory_blocking(move || { let res = if let Some(seek) = seek { (&*std).seek(seek).and_then(|_| buf.write_to(&mut &*std)) } else { @@ -657,7 +662,12 @@ impl AsyncWrite for File { }; (Operation::Write(res), buf) - })); + }) + .ok_or_else(|| { + io::Error::new(io::ErrorKind::Other, "background task failed") + })?; + + inner.state = Busy(blocking_task_join_handle); return Ready(Ok(n)); } @@ -689,11 +699,13 @@ impl AsyncWrite for File { } fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> { + ready!(crate::trace::trace_leaf(cx)); let inner = self.inner.get_mut(); inner.poll_flush(cx) } fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> { + ready!(crate::trace::trace_leaf(cx)); self.poll_flush(cx) } } @@ -719,6 +731,15 @@ impl std::os::unix::io::AsRawFd for File { } } +#[cfg(all(unix, not(tokio_no_as_fd)))] +impl std::os::unix::io::AsFd for File { + fn as_fd(&self) -> std::os::unix::io::BorrowedFd<'_> { + unsafe { + std::os::unix::io::BorrowedFd::borrow_raw(std::os::unix::io::AsRawFd::as_raw_fd(self)) + } + } +} + #[cfg(unix)] impl std::os::unix::io::FromRawFd for File { unsafe fn from_raw_fd(fd: std::os::unix::io::RawFd) -> Self { @@ -726,17 +747,32 @@ impl std::os::unix::io::FromRawFd for File { } } -#[cfg(windows)] -impl std::os::windows::io::AsRawHandle for File { - fn as_raw_handle(&self) -> std::os::windows::io::RawHandle { - self.std.as_raw_handle() +cfg_windows! { + use crate::os::windows::io::{AsRawHandle, FromRawHandle, RawHandle}; + #[cfg(not(tokio_no_as_fd))] + use crate::os::windows::io::{AsHandle, BorrowedHandle}; + + impl AsRawHandle for File { + fn as_raw_handle(&self) -> RawHandle { + self.std.as_raw_handle() + } } -} -#[cfg(windows)] -impl std::os::windows::io::FromRawHandle for File { - unsafe fn from_raw_handle(handle: std::os::windows::io::RawHandle) -> Self { - StdFile::from_raw_handle(handle).into() + #[cfg(not(tokio_no_as_fd))] + impl AsHandle for File { + fn as_handle(&self) -> BorrowedHandle<'_> { + unsafe { + BorrowedHandle::borrow_raw( + AsRawHandle::as_raw_handle(self), + ) + } + } + } + + impl FromRawHandle for File { + unsafe fn from_raw_handle(handle: RawHandle) -> Self { + StdFile::from_raw_handle(handle).into() + } } } @@ -744,8 +780,18 @@ impl Inner { async fn complete_inflight(&mut self) { use crate::future::poll_fn; - if let Err(e) = poll_fn(|cx| Pin::new(&mut *self).poll_flush(cx)).await { - self.last_write_err = Some(e.kind()); + poll_fn(|cx| self.poll_complete_inflight(cx)).await + } + + fn poll_complete_inflight(&mut self, cx: &mut Context<'_>) -> Poll<()> { + ready!(crate::trace::trace_leaf(cx)); + match self.poll_flush(cx) { + Poll::Ready(Err(e)) => { + self.last_write_err = Some(e.kind()); + Poll::Ready(()) + } + Poll::Ready(Ok(())) => Poll::Ready(()), + Poll::Pending => Poll::Pending, } } diff --git a/vendor/tokio/src/fs/file/tests.rs b/vendor/tokio/src/fs/file/tests.rs index 28b5ffe77..7c61b3c4b 100644 --- a/vendor/tokio/src/fs/file/tests.rs +++ b/vendor/tokio/src/fs/file/tests.rs @@ -228,14 +228,15 @@ fn flush_while_idle() { } #[test] +#[cfg_attr(miri, ignore)] // takes a really long time with miri fn read_with_buffer_larger_than_max() { // Chunks - let chunk_a = 16 * 1024; + let chunk_a = crate::io::blocking::MAX_BUF; let chunk_b = chunk_a * 2; let chunk_c = chunk_a * 3; let chunk_d = chunk_a * 4; - assert_eq!(chunk_d / 1024, 64); + assert_eq!(chunk_d / 1024 / 1024, 8); let mut data = vec![]; for i in 0..(chunk_d - 1) { @@ -299,14 +300,15 @@ fn read_with_buffer_larger_than_max() { } #[test] +#[cfg_attr(miri, ignore)] // takes a really long time with miri fn write_with_buffer_larger_than_max() { // Chunks - let chunk_a = 16 * 1024; + let chunk_a = crate::io::blocking::MAX_BUF; let chunk_b = chunk_a * 2; let chunk_c = chunk_a * 3; let chunk_d = chunk_a * 4; - assert_eq!(chunk_d / 1024, 64); + assert_eq!(chunk_d / 1024 / 1024, 8); let mut data = vec![]; for i in 0..(chunk_d - 1) { @@ -953,3 +955,24 @@ fn partial_read_set_len_ok() { assert_eq!(n, FOO.len()); assert_eq!(&buf[..n], FOO); } + +#[test] +fn busy_file_seek_error() { + let mut file = MockFile::default(); + let mut seq = Sequence::new(); + file.expect_inner_write() + .once() + .in_sequence(&mut seq) + .returning(|_| Err(io::ErrorKind::Other.into())); + + let mut file = crate::io::BufReader::new(File::from_std(file)); + { + let mut t = task::spawn(file.write(HELLO)); + assert_ready_ok!(t.poll()); + } + + pool::run_one(); + + let mut t = task::spawn(file.seek(SeekFrom::Start(0))); + assert_ready_err!(t.poll()); +} diff --git a/vendor/tokio/src/fs/mocks.rs b/vendor/tokio/src/fs/mocks.rs index 68ef4f3a7..aa01e2471 100644 --- a/vendor/tokio/src/fs/mocks.rs +++ b/vendor/tokio/src/fs/mocks.rs @@ -81,7 +81,7 @@ impl Write for &'_ MockFile { } } -thread_local! { +tokio_thread_local! { static QUEUE: RefCell<VecDeque<Box<dyn FnOnce() + Send>>> = RefCell::new(VecDeque::new()) } @@ -105,6 +105,21 @@ where JoinHandle { rx } } +pub(super) fn spawn_mandatory_blocking<F, R>(f: F) -> Option<JoinHandle<R>> +where + F: FnOnce() -> R + Send + 'static, + R: Send + 'static, +{ + let (tx, rx) = oneshot::channel(); + let task = Box::new(move || { + let _ = tx.send(f()); + }); + + QUEUE.with(|cell| cell.borrow_mut().push_back(task)); + + Some(JoinHandle { rx }) +} + impl<T> Future for JoinHandle<T> { type Output = Result<T, io::Error>; diff --git a/vendor/tokio/src/fs/mod.rs b/vendor/tokio/src/fs/mod.rs index ca0264b36..f6d9605fe 100644 --- a/vendor/tokio/src/fs/mod.rs +++ b/vendor/tokio/src/fs/mod.rs @@ -22,6 +22,24 @@ //! `std::io::ErrorKind::WouldBlock` if a *worker* thread can not be converted //! to a *backup* thread immediately. //! +//! **Warning**: These adapters may create a large number of temporary tasks, +//! especially when reading large files. When performing a lot of operations +//! in one batch, it may be significantly faster to use [`spawn_blocking`] +//! directly: +//! +//! ``` +//! use tokio::fs::File; +//! use std::io::{BufReader, BufRead}; +//! async fn count_lines(file: File) -> Result<usize, std::io::Error> { +//! let file = file.into_std().await; +//! tokio::task::spawn_blocking(move || { +//! let line_count = BufReader::new(file).lines().count(); +//! Ok(line_count) +//! }).await? +//! } +//! ``` +//! +//! [`spawn_blocking`]: fn@crate::task::spawn_blocking //! [`AsyncRead`]: trait@crate::io::AsyncRead mod canonicalize; @@ -84,6 +102,9 @@ pub use self::write::write; mod copy; pub use self::copy::copy; +mod try_exists; +pub use self::try_exists::try_exists; + #[cfg(test)] mod mocks; @@ -94,9 +115,7 @@ feature! { pub use self::symlink::symlink; } -feature! { - #![windows] - +cfg_windows! { mod symlink_dir; pub use self::symlink_dir::symlink_dir; diff --git a/vendor/tokio/src/fs/open_options.rs b/vendor/tokio/src/fs/open_options.rs index 3e73529ec..103510250 100644 --- a/vendor/tokio/src/fs/open_options.rs +++ b/vendor/tokio/src/fs/open_options.rs @@ -10,6 +10,11 @@ use mock_open_options::MockOpenOptions as StdOpenOptions; #[cfg(not(test))] use std::fs::OpenOptions as StdOpenOptions; +#[cfg(unix)] +use std::os::unix::fs::OpenOptionsExt; +#[cfg(windows)] +use std::os::windows::fs::OpenOptionsExt; + /// Options and flags which can be used to configure how a file is opened. /// /// This builder exposes the ability to configure how a [`File`] is opened and @@ -399,8 +404,6 @@ impl OpenOptions { feature! { #![unix] - use std::os::unix::fs::OpenOptionsExt; - impl OpenOptions { /// Sets the mode bits that a new file will be created with. /// @@ -430,7 +433,7 @@ feature! { self } - /// Pass custom flags to the `flags` argument of `open`. + /// Passes custom flags to the `flags` argument of `open`. /// /// The bits that define the access mode are masked out with `O_ACCMODE`, to /// ensure they do not interfere with the access mode set by Rusts options. @@ -464,11 +467,7 @@ feature! { } } -feature! { - #![windows] - - use std::os::windows::fs::OpenOptionsExt; - +cfg_windows! { impl OpenOptions { /// Overrides the `dwDesiredAccess` argument to the call to [`CreateFile`] /// with the specified value. @@ -542,7 +541,7 @@ feature! { /// # Examples /// /// ```no_run - /// use winapi::um::winbase::FILE_FLAG_DELETE_ON_CLOSE; + /// use windows_sys::Win32::Storage::FileSystem::FILE_FLAG_DELETE_ON_CLOSE; /// use tokio::fs::OpenOptions; /// /// # #[tokio::main] @@ -581,7 +580,7 @@ feature! { /// # Examples /// /// ```no_run - /// use winapi::um::winnt::FILE_ATTRIBUTE_HIDDEN; + /// use windows_sys::Win32::Storage::FileSystem::FILE_ATTRIBUTE_HIDDEN; /// use tokio::fs::OpenOptions; /// /// # #[tokio::main] @@ -624,7 +623,7 @@ feature! { /// # Examples /// /// ```no_run - /// use winapi::um::winbase::SECURITY_IDENTIFICATION; + /// use windows_sys::Win32::Storage::FileSystem::SECURITY_IDENTIFICATION; /// use tokio::fs::OpenOptions; /// /// # #[tokio::main] diff --git a/vendor/tokio/src/fs/open_options/mock_open_options.rs b/vendor/tokio/src/fs/open_options/mock_open_options.rs index cbbda0ec2..17b4a4864 100644 --- a/vendor/tokio/src/fs/open_options/mock_open_options.rs +++ b/vendor/tokio/src/fs/open_options/mock_open_options.rs @@ -1,3 +1,4 @@ +#![allow(unreachable_pub)] //! Mock version of std::fs::OpenOptions; use mockall::mock; diff --git a/vendor/tokio/src/fs/read.rs b/vendor/tokio/src/fs/read.rs index 2d80eb5bd..ada5ba391 100644 --- a/vendor/tokio/src/fs/read.rs +++ b/vendor/tokio/src/fs/read.rs @@ -13,8 +13,12 @@ use std::{io, path::Path}; /// buffer based on the file size when available, so it is generally faster than /// reading into a vector created with `Vec::new()`. /// +/// This operation is implemented by running the equivalent blocking operation +/// on a separate thread pool using [`spawn_blocking`]. +/// /// [`File::open`]: super::File::open /// [`read_to_end`]: crate::io::AsyncReadExt::read_to_end +/// [`spawn_blocking`]: crate::task::spawn_blocking /// /// # Errors /// diff --git a/vendor/tokio/src/fs/read_dir.rs b/vendor/tokio/src/fs/read_dir.rs index c1cb665de..def735b3c 100644 --- a/vendor/tokio/src/fs/read_dir.rs +++ b/vendor/tokio/src/fs/read_dir.rs @@ -1,5 +1,6 @@ use crate::fs::asyncify; +use std::collections::VecDeque; use std::ffi::OsString; use std::fs::{FileType, Metadata}; use std::future::Future; @@ -19,17 +20,29 @@ use crate::blocking::spawn_blocking; #[cfg(not(test))] use crate::blocking::JoinHandle; +const CHUNK_SIZE: usize = 32; + /// Returns a stream over the entries within a directory. /// /// This is an async version of [`std::fs::read_dir`](std::fs::read_dir) +/// +/// This operation is implemented by running the equivalent blocking +/// operation on a separate thread pool using [`spawn_blocking`]. +/// +/// [`spawn_blocking`]: crate::task::spawn_blocking pub async fn read_dir(path: impl AsRef<Path>) -> io::Result<ReadDir> { let path = path.as_ref().to_owned(); - let std = asyncify(|| std::fs::read_dir(path)).await?; + asyncify(|| -> io::Result<ReadDir> { + let mut std = std::fs::read_dir(path)?; + let mut buf = VecDeque::with_capacity(CHUNK_SIZE); + let remain = ReadDir::next_chunk(&mut buf, &mut std); - Ok(ReadDir(State::Idle(Some(std)))) + Ok(ReadDir(State::Idle(Some((buf, std, remain))))) + }) + .await } -/// Read the the entries in a directory. +/// Reads the entries in a directory. /// /// This struct is returned from the [`read_dir`] function of this module and /// will yield instances of [`DirEntry`]. Through a [`DirEntry`] information @@ -53,12 +66,16 @@ pub struct ReadDir(State); #[derive(Debug)] enum State { - Idle(Option<std::fs::ReadDir>), - Pending(JoinHandle<(Option<io::Result<std::fs::DirEntry>>, std::fs::ReadDir)>), + Idle(Option<(VecDeque<io::Result<DirEntry>>, std::fs::ReadDir, bool)>), + Pending(JoinHandle<(VecDeque<io::Result<DirEntry>>, std::fs::ReadDir, bool)>), } impl ReadDir { /// Returns the next entry in the directory stream. + /// + /// # Cancel safety + /// + /// This method is cancellation safe. pub async fn next_entry(&mut self) -> io::Result<Option<DirEntry>> { use crate::future::poll_fn; poll_fn(|cx| self.poll_next_entry(cx)).await @@ -85,28 +102,57 @@ impl ReadDir { pub fn poll_next_entry(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<Option<DirEntry>>> { loop { match self.0 { - State::Idle(ref mut std) => { - let mut std = std.take().unwrap(); + State::Idle(ref mut data) => { + let (buf, _, ref remain) = data.as_mut().unwrap(); + + if let Some(ent) = buf.pop_front() { + return Poll::Ready(ent.map(Some)); + } else if !remain { + return Poll::Ready(Ok(None)); + } + + let (mut buf, mut std, _) = data.take().unwrap(); self.0 = State::Pending(spawn_blocking(move || { - let ret = std.next(); - (ret, std) + let remain = ReadDir::next_chunk(&mut buf, &mut std); + (buf, std, remain) })); } State::Pending(ref mut rx) => { - let (ret, std) = ready!(Pin::new(rx).poll(cx))?; - self.0 = State::Idle(Some(std)); + self.0 = State::Idle(Some(ready!(Pin::new(rx).poll(cx))?)); + } + } + } + } - let ret = match ret { - Some(Ok(std)) => Ok(Some(DirEntry(Arc::new(std)))), - Some(Err(e)) => Err(e), - None => Ok(None), - }; + fn next_chunk(buf: &mut VecDeque<io::Result<DirEntry>>, std: &mut std::fs::ReadDir) -> bool { + for _ in 0..CHUNK_SIZE { + let ret = match std.next() { + Some(ret) => ret, + None => return false, + }; - return Poll::Ready(ret); - } + let success = ret.is_ok(); + + buf.push_back(ret.map(|std| DirEntry { + #[cfg(not(any( + target_os = "solaris", + target_os = "illumos", + target_os = "haiku", + target_os = "vxworks", + target_os = "nto", + target_os = "vita", + )))] + file_type: std.file_type().ok(), + std: Arc::new(std), + })); + + if !success { + break; } } + + true } } @@ -151,7 +197,18 @@ feature! { /// filesystem. Each entry can be inspected via methods to learn about the full /// path or possibly other metadata through per-platform extension traits. #[derive(Debug)] -pub struct DirEntry(Arc<std::fs::DirEntry>); +pub struct DirEntry { + #[cfg(not(any( + target_os = "solaris", + target_os = "illumos", + target_os = "haiku", + target_os = "vxworks", + target_os = "nto", + target_os = "vita", + )))] + file_type: Option<FileType>, + std: Arc<std::fs::DirEntry>, +} impl DirEntry { /// Returns the full path to the file that this entry represents. @@ -184,7 +241,7 @@ impl DirEntry { /// /// The exact text, of course, depends on what files you have in `.`. pub fn path(&self) -> PathBuf { - self.0.path() + self.std.path() } /// Returns the bare file name of this directory entry without any other @@ -205,7 +262,7 @@ impl DirEntry { /// # } /// ``` pub fn file_name(&self) -> OsString { - self.0.file_name() + self.std.file_name() } /// Returns the metadata for the file that this entry points at. @@ -239,7 +296,7 @@ impl DirEntry { /// # } /// ``` pub async fn metadata(&self) -> io::Result<Metadata> { - let std = self.0.clone(); + let std = self.std.clone(); asyncify(move || std.metadata()).await } @@ -274,13 +331,25 @@ impl DirEntry { /// # } /// ``` pub async fn file_type(&self) -> io::Result<FileType> { - let std = self.0.clone(); + #[cfg(not(any( + target_os = "solaris", + target_os = "illumos", + target_os = "haiku", + target_os = "vxworks", + target_os = "nto", + target_os = "vita", + )))] + if let Some(file_type) = self.file_type { + return Ok(file_type); + } + + let std = self.std.clone(); asyncify(move || std.file_type()).await } - /// Returns a reference to the underlying `std::fs::DirEntry` + /// Returns a reference to the underlying `std::fs::DirEntry`. #[cfg(unix)] pub(super) fn as_inner(&self) -> &std::fs::DirEntry { - &self.0 + &self.std } } diff --git a/vendor/tokio/src/fs/read_to_string.rs b/vendor/tokio/src/fs/read_to_string.rs index 4f37986d6..26228d98c 100644 --- a/vendor/tokio/src/fs/read_to_string.rs +++ b/vendor/tokio/src/fs/read_to_string.rs @@ -7,6 +7,10 @@ use std::{io, path::Path}; /// /// This is the async equivalent of [`std::fs::read_to_string`][std]. /// +/// This operation is implemented by running the equivalent blocking operation +/// on a separate thread pool using [`spawn_blocking`]. +/// +/// [`spawn_blocking`]: crate::task::spawn_blocking /// [std]: fn@std::fs::read_to_string /// /// # Examples diff --git a/vendor/tokio/src/fs/symlink_dir.rs b/vendor/tokio/src/fs/symlink_dir.rs index 736e762b4..6753c25eb 100644 --- a/vendor/tokio/src/fs/symlink_dir.rs +++ b/vendor/tokio/src/fs/symlink_dir.rs @@ -10,7 +10,7 @@ use std::path::Path; /// /// This is an async version of [`std::os::windows::fs::symlink_dir`][std] /// -/// [std]: std::os::windows::fs::symlink_dir +/// [std]: https://doc.rust-lang.org/std/os/windows/fs/fn.symlink_dir.html pub async fn symlink_dir(src: impl AsRef<Path>, dst: impl AsRef<Path>) -> io::Result<()> { let src = src.as_ref().to_owned(); let dst = dst.as_ref().to_owned(); diff --git a/vendor/tokio/src/fs/symlink_file.rs b/vendor/tokio/src/fs/symlink_file.rs index 07d8e6041..623352a1b 100644 --- a/vendor/tokio/src/fs/symlink_file.rs +++ b/vendor/tokio/src/fs/symlink_file.rs @@ -10,7 +10,7 @@ use std::path::Path; /// /// This is an async version of [`std::os::windows::fs::symlink_file`][std] /// -/// [std]: std::os::windows::fs::symlink_file +/// [std]: https://doc.rust-lang.org/std/os/windows/fs/fn.symlink_file.html pub async fn symlink_file(src: impl AsRef<Path>, dst: impl AsRef<Path>) -> io::Result<()> { let src = src.as_ref().to_owned(); let dst = dst.as_ref().to_owned(); diff --git a/vendor/tokio/src/fs/try_exists.rs b/vendor/tokio/src/fs/try_exists.rs new file mode 100644 index 000000000..069518bf9 --- /dev/null +++ b/vendor/tokio/src/fs/try_exists.rs @@ -0,0 +1,34 @@ +use crate::fs::asyncify; + +use std::io; +use std::path::Path; + +/// Returns `Ok(true)` if the path points at an existing entity. +/// +/// This function will traverse symbolic links to query information about the +/// destination file. In case of broken symbolic links this will return `Ok(false)`. +/// +/// This is the async equivalent of [`std::path::Path::try_exists`][std]. +/// +/// [std]: fn@std::path::Path::try_exists +/// +/// # Examples +/// +/// ```no_run +/// use tokio::fs; +/// +/// # async fn dox() -> std::io::Result<()> { +/// fs::try_exists("foo.txt").await?; +/// # Ok(()) +/// # } +/// ``` +pub async fn try_exists(path: impl AsRef<Path>) -> io::Result<bool> { + let path = path.as_ref().to_owned(); + // std's Path::try_exists is not available for current Rust min supported version. + // Current implementation is based on its internal implementation instead. + match asyncify(move || std::fs::metadata(path)).await { + Ok(_) => Ok(true), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(false), + Err(error) => Err(error), + } +} diff --git a/vendor/tokio/src/fs/write.rs b/vendor/tokio/src/fs/write.rs index 0ed908271..28606fb36 100644 --- a/vendor/tokio/src/fs/write.rs +++ b/vendor/tokio/src/fs/write.rs @@ -7,6 +7,10 @@ use std::{io, path::Path}; /// /// This is the async equivalent of [`std::fs::write`][std]. /// +/// This operation is implemented by running the equivalent blocking operation +/// on a separate thread pool using [`spawn_blocking`]. +/// +/// [`spawn_blocking`]: crate::task::spawn_blocking /// [std]: fn@std::fs::write /// /// # Examples diff --git a/vendor/tokio/src/future/block_on.rs b/vendor/tokio/src/future/block_on.rs index 91f9cc005..2c2ab3736 100644 --- a/vendor/tokio/src/future/block_on.rs +++ b/vendor/tokio/src/future/block_on.rs @@ -1,15 +1,22 @@ use std::future::Future; cfg_rt! { + #[track_caller] pub(crate) fn block_on<F: Future>(f: F) -> F::Output { - let mut e = crate::runtime::enter::enter(false); + let mut e = crate::runtime::context::try_enter_blocking_region().expect( + "Cannot block the current thread from within a runtime. This \ + happens because a function attempted to block the current \ + thread while the thread is being used to drive asynchronous \ + tasks." + ); e.block_on(f).unwrap() } } cfg_not_rt! { + #[track_caller] pub(crate) fn block_on<F: Future>(f: F) -> F::Output { - let mut park = crate::park::thread::CachedParkThread::new(); + let mut park = crate::runtime::park::CachedParkThread::new(); park.block_on(f).unwrap() } } diff --git a/vendor/tokio/src/future/maybe_done.rs b/vendor/tokio/src/future/maybe_done.rs index 1e083ad7f..486efbe01 100644 --- a/vendor/tokio/src/future/maybe_done.rs +++ b/vendor/tokio/src/future/maybe_done.rs @@ -1,4 +1,4 @@ -//! Definition of the MaybeDone combinator +//! Definition of the MaybeDone combinator. use std::future::Future; use std::mem; @@ -8,9 +8,9 @@ use std::task::{Context, Poll}; /// A future that may have completed. #[derive(Debug)] pub enum MaybeDone<Fut: Future> { - /// A not-yet-completed future + /// A not-yet-completed future. Future(Fut), - /// The output of the completed future + /// The output of the completed future. Done(Fut::Output), /// The empty variant after the result of a [`MaybeDone`] has been /// taken using the [`take_output`](MaybeDone::take_output) method. @@ -20,7 +20,7 @@ pub enum MaybeDone<Fut: Future> { // Safe because we never generate `Pin<&mut Fut::Output>` impl<Fut: Future + Unpin> Unpin for MaybeDone<Fut> {} -/// Wraps a future into a `MaybeDone` +/// Wraps a future into a `MaybeDone`. pub fn maybe_done<Fut: Future>(future: Fut) -> MaybeDone<Fut> { MaybeDone::Future(future) } diff --git a/vendor/tokio/src/future/mod.rs b/vendor/tokio/src/future/mod.rs index 96483acd7..084ddc571 100644 --- a/vendor/tokio/src/future/mod.rs +++ b/vendor/tokio/src/future/mod.rs @@ -8,11 +8,6 @@ pub(crate) mod maybe_done; mod poll_fn; pub use poll_fn::poll_fn; -cfg_not_loom! { - mod ready; - pub(crate) use ready::{ok, Ready}; -} - cfg_process! { mod try_join; pub(crate) use try_join::try_join3; diff --git a/vendor/tokio/src/future/poll_fn.rs b/vendor/tokio/src/future/poll_fn.rs index 0169bd5fc..074d9438e 100644 --- a/vendor/tokio/src/future/poll_fn.rs +++ b/vendor/tokio/src/future/poll_fn.rs @@ -1,19 +1,29 @@ #![allow(dead_code)] -//! Definition of the `PollFn` adapter combinator +//! Definition of the `PollFn` adapter combinator. use std::fmt; use std::future::Future; 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: +// <https://internals.rust-lang.org/t/surprising-soundness-trouble-around-pollfn/17484> +// +// 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: F, } -impl<F> Unpin for PollFn<F> {} - /// Creates a new future wrapping around a function returning [`Poll`]. pub fn poll_fn<T, F>(f: F) -> PollFn<F> where @@ -34,7 +44,17 @@ where { type Output = T; - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> { - (&mut self.f)(cx) + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> { + // 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) } } diff --git a/vendor/tokio/src/future/ready.rs b/vendor/tokio/src/future/ready.rs deleted file mode 100644 index de2d60c13..000000000 --- a/vendor/tokio/src/future/ready.rs +++ /dev/null @@ -1,27 +0,0 @@ -use std::future::Future; -use std::pin::Pin; -use std::task::{Context, Poll}; - -/// Future for the [`ok`](ok()) function. -/// -/// `pub` in order to use the future as an associated type in a sealed trait. -#[derive(Debug)] -// Used as an associated type in a "sealed" trait. -#[allow(unreachable_pub)] -pub struct Ready<T>(Option<T>); - -impl<T> Unpin for Ready<T> {} - -impl<T> Future for Ready<T> { - type Output = T; - - #[inline] - fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<T> { - Poll::Ready(self.0.take().unwrap()) - } -} - -/// Creates a future that is immediately ready with a success value. -pub(crate) fn ok<T, E>(t: T) -> Ready<Result<T, E>> { - Ready(Some(Ok(t))) -} diff --git a/vendor/tokio/src/fuzz.rs b/vendor/tokio/src/fuzz.rs new file mode 100644 index 000000000..d89718f0e --- /dev/null +++ b/vendor/tokio/src/fuzz.rs @@ -0,0 +1 @@ +pub use crate::util::linked_list::tests::fuzz_linked_list; diff --git a/vendor/tokio/src/io/async_fd.rs b/vendor/tokio/src/io/async_fd.rs index fa5bec530..0cde5a407 100644 --- a/vendor/tokio/src/io/async_fd.rs +++ b/vendor/tokio/src/io/async_fd.rs @@ -1,4 +1,6 @@ -use crate::io::driver::{Handle, Interest, ReadyEvent, Registration}; +use crate::io::{Interest, Ready}; +use crate::runtime::io::{ReadyEvent, Registration}; +use crate::runtime::scheduler; use mio::unix::SourceFd; use std::io; @@ -63,8 +65,8 @@ use std::{task::Context, task::Poll}; /// # Examples /// /// This example shows how to turn [`std::net::TcpStream`] asynchronous using -/// `AsyncFd`. It implements `read` as an async fn, and `AsyncWrite` as a trait -/// to show how to implement both approaches. +/// `AsyncFd`. It implements the read/write operations both as an `async fn` +/// and using the IO traits [`AsyncRead`] and [`AsyncWrite`]. /// /// ```no_run /// use futures::ready; @@ -72,7 +74,7 @@ use std::{task::Context, task::Poll}; /// use std::net::TcpStream; /// use std::pin::Pin; /// use std::task::{Context, Poll}; -/// use tokio::io::AsyncWrite; +/// use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; /// use tokio::io::unix::AsyncFd; /// /// pub struct AsyncTcpStream { @@ -81,6 +83,7 @@ use std::{task::Context, task::Poll}; /// /// impl AsyncTcpStream { /// pub fn new(tcp: TcpStream) -> io::Result<Self> { +/// tcp.set_nonblocking(true)?; /// Ok(Self { /// inner: AsyncFd::new(tcp)?, /// }) @@ -96,6 +99,39 @@ use std::{task::Context, task::Poll}; /// } /// } /// } +/// +/// pub async fn write(&self, buf: &[u8]) -> io::Result<usize> { +/// loop { +/// let mut guard = self.inner.writable().await?; +/// +/// match guard.try_io(|inner| inner.get_ref().write(buf)) { +/// Ok(result) => return result, +/// Err(_would_block) => continue, +/// } +/// } +/// } +/// } +/// +/// impl AsyncRead for AsyncTcpStream { +/// fn poll_read( +/// self: Pin<&mut Self>, +/// cx: &mut Context<'_>, +/// buf: &mut ReadBuf<'_> +/// ) -> Poll<io::Result<()>> { +/// loop { +/// let mut guard = ready!(self.inner.poll_read_ready(cx))?; +/// +/// let unfilled = buf.initialize_unfilled(); +/// match guard.try_io(|inner| inner.get_ref().read(unfilled)) { +/// Ok(Ok(len)) => { +/// buf.advance(len); +/// return Poll::Ready(Ok(())); +/// }, +/// Ok(Err(err)) => return Poll::Ready(Err(err)), +/// Err(_would_block) => continue, +/// } +/// } +/// } /// } /// /// impl AsyncWrite for AsyncTcpStream { @@ -136,6 +172,8 @@ use std::{task::Context, task::Poll}; /// [`writable`]: method@Self::writable /// [`AsyncFdReadyGuard`]: struct@self::AsyncFdReadyGuard /// [`TcpStream::poll_read_ready`]: struct@crate::net::TcpStream +/// [`AsyncRead`]: trait@crate::io::AsyncRead +/// [`AsyncWrite`]: trait@crate::io::AsyncWrite pub struct AsyncFd<T: AsRawFd> { registration: Registration, inner: Option<T>, @@ -163,35 +201,50 @@ pub struct AsyncFdReadyMutGuard<'a, T: AsRawFd> { event: Option<ReadyEvent>, } -const ALL_INTEREST: Interest = Interest::READABLE.add(Interest::WRITABLE); - impl<T: AsRawFd> AsyncFd<T> { - #[inline] /// Creates an AsyncFd backed by (and taking ownership of) an object /// implementing [`AsRawFd`]. The backing file descriptor is cached at the /// time of creation. /// + /// Only configures the [`Interest::READABLE`] and [`Interest::WRITABLE`] interests. For more + /// control, use [`AsyncFd::with_interest`]. + /// /// This method must be called in the context of a tokio runtime. + /// + /// # Panics + /// + /// This function panics if there is no current reactor set, or if the `rt` + /// feature flag is not enabled. + #[inline] + #[track_caller] pub fn new(inner: T) -> io::Result<Self> where T: AsRawFd, { - Self::with_interest(inner, ALL_INTEREST) + Self::with_interest(inner, Interest::READABLE | Interest::WRITABLE) } + /// Creates an AsyncFd backed by (and taking ownership of) an object + /// implementing [`AsRawFd`], with a specific [`Interest`]. The backing + /// file descriptor is cached at the time of creation. + /// + /// # Panics + /// + /// This function panics if there is no current reactor set, or if the `rt` + /// feature flag is not enabled. #[inline] - /// Creates new instance as `new` with additional ability to customize interest, - /// allowing to specify whether file descriptor will be polled for read, write or both. + #[track_caller] pub fn with_interest(inner: T, interest: Interest) -> io::Result<Self> where T: AsRawFd, { - Self::new_with_handle_and_interest(inner, Handle::current(), interest) + Self::new_with_handle_and_interest(inner, scheduler::Handle::current(), interest) } + #[track_caller] pub(crate) fn new_with_handle_and_interest( inner: T, - handle: Handle, + handle: scheduler::Handle, interest: Interest, ) -> io::Result<Self> { let fd = inner.as_raw_fd(); @@ -205,13 +258,13 @@ impl<T: AsRawFd> AsyncFd<T> { }) } - /// Returns a shared reference to the backing object of this [`AsyncFd`] + /// Returns a shared reference to the backing object of this [`AsyncFd`]. #[inline] pub fn get_ref(&self) -> &T { self.inner.as_ref().unwrap() } - /// Returns a mutable reference to the backing object of this [`AsyncFd`] + /// Returns a mutable reference to the backing object of this [`AsyncFd`]. #[inline] pub fn get_mut(&mut self) -> &mut T { self.inner.as_mut().unwrap() @@ -389,7 +442,96 @@ impl<T: AsRawFd> AsyncFd<T> { .into() } - async fn readiness(&self, interest: Interest) -> io::Result<AsyncFdReadyGuard<'_, T>> { + /// Waits for any of the requested ready states, returning a + /// [`AsyncFdReadyGuard`] that must be dropped to resume + /// polling for the requested ready states. + /// + /// The function may complete without the file descriptor being ready. This is a + /// false-positive and attempting an operation will return with + /// `io::ErrorKind::WouldBlock`. The function can also return with an empty + /// [`Ready`] set, so you should always check the returned value and possibly + /// wait again if the requested states are not set. + /// + /// When an IO operation does return `io::ErrorKind::WouldBlock`, the readiness must be cleared. + /// When a combined interest is used, it is important to clear only the readiness + /// that is actually observed to block. For instance when the combined + /// interest `Interest::READABLE | Interest::WRITABLE` is used, and a read blocks, only + /// read readiness should be cleared using the [`AsyncFdReadyGuard::clear_ready_matching`] method: + /// `guard.clear_ready_matching(Ready::READABLE)`. + /// Also clearing the write readiness in this case would be incorrect. The [`AsyncFdReadyGuard::clear_ready`] + /// method clears all readiness flags. + /// + /// This method takes `&self`, so it is possible to call this method + /// concurrently with other methods on this struct. This method only + /// provides shared access to the inner IO resource when handling the + /// [`AsyncFdReadyGuard`]. + /// + /// # Examples + /// + /// Concurrently read and write to a [`std::net::TcpStream`] on the same task without + /// splitting. + /// + /// ```no_run + /// use std::error::Error; + /// use std::io; + /// use std::io::{Read, Write}; + /// use std::net::TcpStream; + /// use tokio::io::unix::AsyncFd; + /// use tokio::io::{Interest, Ready}; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// let stream = TcpStream::connect("127.0.0.1:8080")?; + /// stream.set_nonblocking(true)?; + /// let stream = AsyncFd::new(stream)?; + /// + /// loop { + /// let mut guard = stream + /// .ready(Interest::READABLE | Interest::WRITABLE) + /// .await?; + /// + /// if guard.ready().is_readable() { + /// let mut data = vec![0; 1024]; + /// // Try to read data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match stream.get_ref().read(&mut data) { + /// Ok(n) => { + /// println!("read {} bytes", n); + /// } + /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + /// // a read has blocked, but a write might still succeed. + /// // clear only the read readiness. + /// guard.clear_ready_matching(Ready::READABLE); + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// + /// if guard.ready().is_writable() { + /// // Try to write data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match stream.get_ref().write(b"hello world") { + /// Ok(n) => { + /// println!("write {} bytes", n); + /// } + /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + /// // a write has blocked, but a read might still succeed. + /// // clear only the write readiness. + /// guard.clear_ready_matching(Ready::WRITABLE); + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// } + /// } + /// ``` + pub async fn ready(&self, interest: Interest) -> io::Result<AsyncFdReadyGuard<'_, T>> { let event = self.registration.readiness(interest).await?; Ok(AsyncFdReadyGuard { @@ -398,7 +540,94 @@ impl<T: AsRawFd> AsyncFd<T> { }) } - async fn readiness_mut( + /// Waits for any of the requested ready states, returning a + /// [`AsyncFdReadyMutGuard`] that must be dropped to resume + /// polling for the requested ready states. + /// + /// The function may complete without the file descriptor being ready. This is a + /// false-positive and attempting an operation will return with + /// `io::ErrorKind::WouldBlock`. The function can also return with an empty + /// [`Ready`] set, so you should always check the returned value and possibly + /// wait again if the requested states are not set. + /// + /// When an IO operation does return `io::ErrorKind::WouldBlock`, the readiness must be cleared. + /// When a combined interest is used, it is important to clear only the readiness + /// that is actually observed to block. For instance when the combined + /// interest `Interest::READABLE | Interest::WRITABLE` is used, and a read blocks, only + /// read readiness should be cleared using the [`AsyncFdReadyMutGuard::clear_ready_matching`] method: + /// `guard.clear_ready_matching(Ready::READABLE)`. + /// Also clearing the write readiness in this case would be incorrect. + /// The [`AsyncFdReadyMutGuard::clear_ready`] method clears all readiness flags. + /// + /// This method takes `&mut self`, so it is possible to access the inner IO + /// resource mutably when handling the [`AsyncFdReadyMutGuard`]. + /// + /// # Examples + /// + /// Concurrently read and write to a [`std::net::TcpStream`] on the same task without + /// splitting. + /// + /// ```no_run + /// use std::error::Error; + /// use std::io; + /// use std::io::{Read, Write}; + /// use std::net::TcpStream; + /// use tokio::io::unix::AsyncFd; + /// use tokio::io::{Interest, Ready}; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// let stream = TcpStream::connect("127.0.0.1:8080")?; + /// stream.set_nonblocking(true)?; + /// let mut stream = AsyncFd::new(stream)?; + /// + /// loop { + /// let mut guard = stream + /// .ready_mut(Interest::READABLE | Interest::WRITABLE) + /// .await?; + /// + /// if guard.ready().is_readable() { + /// let mut data = vec![0; 1024]; + /// // Try to read data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match guard.get_inner_mut().read(&mut data) { + /// Ok(n) => { + /// println!("read {} bytes", n); + /// } + /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + /// // a read has blocked, but a write might still succeed. + /// // clear only the read readiness. + /// guard.clear_ready_matching(Ready::READABLE); + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// + /// if guard.ready().is_writable() { + /// // Try to write data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match guard.get_inner_mut().write(b"hello world") { + /// Ok(n) => { + /// println!("write {} bytes", n); + /// } + /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + /// // a write has blocked, but a read might still succeed. + /// // clear only the write readiness. + /// guard.clear_ready_matching(Ready::WRITABLE); + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// } + /// } + /// ``` + pub async fn ready_mut( &mut self, interest: Interest, ) -> io::Result<AsyncFdReadyMutGuard<'_, T>> { @@ -420,7 +649,7 @@ impl<T: AsRawFd> AsyncFd<T> { /// [`AsyncFdReadyGuard`]. #[allow(clippy::needless_lifetimes)] // The lifetime improves rustdoc rendering. pub async fn readable<'a>(&'a self) -> io::Result<AsyncFdReadyGuard<'a, T>> { - self.readiness(Interest::READABLE).await + self.ready(Interest::READABLE).await } /// Waits for the file descriptor to become readable, returning a @@ -431,7 +660,7 @@ impl<T: AsRawFd> AsyncFd<T> { /// resource mutably when handling the [`AsyncFdReadyMutGuard`]. #[allow(clippy::needless_lifetimes)] // The lifetime improves rustdoc rendering. pub async fn readable_mut<'a>(&'a mut self) -> io::Result<AsyncFdReadyMutGuard<'a, T>> { - self.readiness_mut(Interest::READABLE).await + self.ready_mut(Interest::READABLE).await } /// Waits for the file descriptor to become writable, returning a @@ -444,7 +673,7 @@ impl<T: AsRawFd> AsyncFd<T> { /// [`AsyncFdReadyGuard`]. #[allow(clippy::needless_lifetimes)] // The lifetime improves rustdoc rendering. pub async fn writable<'a>(&'a self) -> io::Result<AsyncFdReadyGuard<'a, T>> { - self.readiness(Interest::WRITABLE).await + self.ready(Interest::WRITABLE).await } /// Waits for the file descriptor to become writable, returning a @@ -455,7 +684,110 @@ impl<T: AsRawFd> AsyncFd<T> { /// resource mutably when handling the [`AsyncFdReadyMutGuard`]. #[allow(clippy::needless_lifetimes)] // The lifetime improves rustdoc rendering. pub async fn writable_mut<'a>(&'a mut self) -> io::Result<AsyncFdReadyMutGuard<'a, T>> { - self.readiness_mut(Interest::WRITABLE).await + self.ready_mut(Interest::WRITABLE).await + } + + /// Reads or writes from the file descriptor using a user-provided IO operation. + /// + /// The `async_io` method is a convenience utility that waits for the file + /// descriptor to become ready, and then executes the provided IO operation. + /// Since file descriptors may be marked ready spuriously, the closure will + /// be called repeatedly until it returns something other than a + /// [`WouldBlock`] error. This is done using the following loop: + /// + /// ```no_run + /// # use std::io::{self, Result}; + /// # struct Dox<T> { inner: T } + /// # impl<T> Dox<T> { + /// # async fn writable(&self) -> Result<&Self> { + /// # Ok(self) + /// # } + /// # fn try_io<R>(&self, _: impl FnMut(&T) -> Result<R>) -> Result<Result<R>> { + /// # panic!() + /// # } + /// async fn async_io<R>(&self, mut f: impl FnMut(&T) -> io::Result<R>) -> io::Result<R> { + /// loop { + /// // or `readable` if called with the read interest. + /// let guard = self.writable().await?; + /// + /// match guard.try_io(&mut f) { + /// Ok(result) => return result, + /// Err(_would_block) => continue, + /// } + /// } + /// } + /// # } + /// ``` + /// + /// The closure should only return a [`WouldBlock`] error if it has performed + /// an IO operation on the file descriptor that failed due to the file descriptor not being + /// ready. Returning a [`WouldBlock`] error in any other situation will + /// incorrectly clear the readiness flag, which can cause the file descriptor to + /// behave incorrectly. + /// + /// The closure should not perform the IO operation using any of the methods + /// defined on the Tokio [`AsyncFd`] type, as this will mess with the + /// readiness flag and can cause the file descriptor to behave incorrectly. + /// + /// This method is not intended to be used with combined interests. + /// The closure should perform only one type of IO operation, so it should not + /// require more than one ready state. This method may panic or sleep forever + /// if it is called with a combined interest. + /// + /// # Examples + /// + /// This example sends some bytes on the inner [`std::net::UdpSocket`]. The `async_io` + /// method waits for readiness, and retries if the send operation does block. This example + /// is equivalent to the one given for [`try_io`]. + /// + /// ```no_run + /// use tokio::io::{Interest, unix::AsyncFd}; + /// + /// use std::io; + /// use std::net::UdpSocket; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let socket = UdpSocket::bind("0.0.0.0:8080")?; + /// socket.set_nonblocking(true)?; + /// let async_fd = AsyncFd::new(socket)?; + /// + /// let written = async_fd + /// .async_io(Interest::WRITABLE, |inner| inner.send(&[1, 2])) + /// .await?; + /// + /// println!("wrote {written} bytes"); + /// + /// Ok(()) + /// } + /// ``` + /// + /// [`try_io`]: AsyncFdReadyGuard::try_io + /// [`WouldBlock`]: std::io::ErrorKind::WouldBlock + pub async fn async_io<R>( + &self, + interest: Interest, + mut f: impl FnMut(&T) -> io::Result<R>, + ) -> io::Result<R> { + self.registration + .async_io(interest, || f(self.get_ref())) + .await + } + + /// Reads or writes from the file descriptor using a user-provided IO operation. + /// + /// The behavior is the same as [`async_io`], except that the closure can mutate the inner + /// value of the [`AsyncFd`]. + /// + /// [`async_io`]: AsyncFd::async_io + pub async fn async_io_mut<R>( + &mut self, + interest: Interest, + mut f: impl FnMut(&mut T) -> io::Result<R>, + ) -> io::Result<R> { + self.registration + .async_io(interest, || f(self.inner.as_mut().unwrap())) + .await } } @@ -465,6 +797,13 @@ impl<T: AsRawFd> AsRawFd for AsyncFd<T> { } } +#[cfg(not(tokio_no_as_fd))] +impl<T: AsRawFd> std::os::unix::io::AsFd for AsyncFd<T> { + fn as_fd(&self) -> std::os::unix::io::BorrowedFd<'_> { + unsafe { std::os::unix::io::BorrowedFd::borrow_raw(self.as_raw_fd()) } + } +} + impl<T: std::fmt::Debug + AsRawFd> std::fmt::Debug for AsyncFd<T> { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("AsyncFd") @@ -480,22 +819,117 @@ impl<T: AsRawFd> Drop for AsyncFd<T> { } impl<'a, Inner: AsRawFd> AsyncFdReadyGuard<'a, Inner> { - /// Indicates to tokio that the file descriptor is no longer ready. The - /// internal readiness flag will be cleared, and tokio will wait for the + /// Indicates to tokio that the file descriptor is no longer ready. All + /// internal readiness flags will be cleared, and tokio will wait for the /// next edge-triggered readiness notification from the OS. /// + /// This function is commonly used with guards returned by [`AsyncFd::readable`] and + /// [`AsyncFd::writable`]. + /// /// It is critical that this function not be called unless your code /// _actually observes_ that the file descriptor is _not_ ready. Do not call /// it simply because, for example, a read succeeded; it should be called /// when a read is observed to block. - /// - /// [`drop`]: method@std::mem::drop pub fn clear_ready(&mut self) { if let Some(event) = self.event.take() { self.async_fd.registration.clear_readiness(event); } } + /// Indicates to tokio that the file descriptor no longer has a specific readiness. + /// The internal readiness flag will be cleared, and tokio will wait for the + /// next edge-triggered readiness notification from the OS. + /// + /// This function is useful in combination with the [`AsyncFd::ready`] method when a + /// combined interest like `Interest::READABLE | Interest::WRITABLE` is used. + /// + /// It is critical that this function not be called unless your code + /// _actually observes_ that the file descriptor is _not_ ready for the provided `Ready`. + /// Do not call it simply because, for example, a read succeeded; it should be called + /// when a read is observed to block. Only clear the specific readiness that is observed to + /// block. For example when a read blocks when using a combined interest, + /// only clear `Ready::READABLE`. + /// + /// # Examples + /// + /// Concurrently read and write to a [`std::net::TcpStream`] on the same task without + /// splitting. + /// + /// ```no_run + /// use std::error::Error; + /// use std::io; + /// use std::io::{Read, Write}; + /// use std::net::TcpStream; + /// use tokio::io::unix::AsyncFd; + /// use tokio::io::{Interest, Ready}; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// let stream = TcpStream::connect("127.0.0.1:8080")?; + /// stream.set_nonblocking(true)?; + /// let stream = AsyncFd::new(stream)?; + /// + /// loop { + /// let mut guard = stream + /// .ready(Interest::READABLE | Interest::WRITABLE) + /// .await?; + /// + /// if guard.ready().is_readable() { + /// let mut data = vec![0; 1024]; + /// // Try to read data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match stream.get_ref().read(&mut data) { + /// Ok(n) => { + /// println!("read {} bytes", n); + /// } + /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + /// // a read has blocked, but a write might still succeed. + /// // clear only the read readiness. + /// guard.clear_ready_matching(Ready::READABLE); + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// + /// if guard.ready().is_writable() { + /// // Try to write data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match stream.get_ref().write(b"hello world") { + /// Ok(n) => { + /// println!("write {} bytes", n); + /// } + /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + /// // a write has blocked, but a read might still succeed. + /// // clear only the write readiness. + /// guard.clear_ready_matching(Ready::WRITABLE); + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// } + /// } + /// ``` + pub fn clear_ready_matching(&mut self, ready: Ready) { + if let Some(mut event) = self.event.take() { + self.async_fd + .registration + .clear_readiness(event.with_ready(ready)); + + // the event is no longer ready for the readiness that was just cleared + event.ready = event.ready - ready; + + if !event.ready.is_empty() { + self.event = Some(event); + } + } + } + /// This method should be invoked when you intentionally want to keep the /// ready flag asserted. /// @@ -505,6 +939,20 @@ impl<'a, Inner: AsRawFd> AsyncFdReadyGuard<'a, Inner> { // no-op } + /// Get the [`Ready`] value associated with this guard. + /// + /// This method will return the empty readiness state if + /// [`AsyncFdReadyGuard::clear_ready`] has been called on + /// the guard. + /// + /// [`Ready`]: crate::io::Ready + pub fn ready(&self) -> Ready { + match &self.event { + Some(event) => event.ready, + None => Ready::EMPTY, + } + } + /// Performs the provided IO operation. /// /// If `f` returns a [`WouldBlock`] error, the readiness state associated @@ -520,12 +968,49 @@ impl<'a, Inner: AsRawFd> AsyncFdReadyGuard<'a, Inner> { /// `AsyncFdReadyGuard` no longer expresses the readiness state that was queried to /// create this `AsyncFdReadyGuard`. /// + /// # Examples + /// + /// This example sends some bytes to the inner [`std::net::UdpSocket`]. Waiting + /// for write-readiness and retrying when the send operation does block are explicit. + /// This example can be written more succinctly using [`AsyncFd::async_io`]. + /// + /// ```no_run + /// use tokio::io::unix::AsyncFd; + /// + /// use std::io; + /// use std::net::UdpSocket; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let socket = UdpSocket::bind("0.0.0.0:8080")?; + /// socket.set_nonblocking(true)?; + /// let async_fd = AsyncFd::new(socket)?; + /// + /// let written = loop { + /// let mut guard = async_fd.writable().await?; + /// match guard.try_io(|inner| inner.get_ref().send(&[1, 2])) { + /// Ok(result) => { + /// break result?; + /// } + /// Err(_would_block) => { + /// // try_io already cleared the file descriptor's readiness state + /// continue; + /// } + /// } + /// }; + /// + /// println!("wrote {written} bytes"); + /// + /// Ok(()) + /// } + /// ``` + /// /// [`WouldBlock`]: std::io::ErrorKind::WouldBlock // Alias for old name in 0.x #[cfg_attr(docsrs, doc(alias = "with_io"))] pub fn try_io<R>( &mut self, - f: impl FnOnce(&AsyncFd<Inner>) -> io::Result<R>, + f: impl FnOnce(&'a AsyncFd<Inner>) -> io::Result<R>, ) -> Result<io::Result<R>, TryIoError> { let result = f(self.async_fd); @@ -542,33 +1027,128 @@ impl<'a, Inner: AsRawFd> AsyncFdReadyGuard<'a, Inner> { } /// Returns a shared reference to the inner [`AsyncFd`]. - pub fn get_ref(&self) -> &AsyncFd<Inner> { + pub fn get_ref(&self) -> &'a AsyncFd<Inner> { self.async_fd } /// Returns a shared reference to the backing object of the inner [`AsyncFd`]. - pub fn get_inner(&self) -> &Inner { + pub fn get_inner(&self) -> &'a Inner { self.get_ref().get_ref() } } impl<'a, Inner: AsRawFd> AsyncFdReadyMutGuard<'a, Inner> { - /// Indicates to tokio that the file descriptor is no longer ready. The - /// internal readiness flag will be cleared, and tokio will wait for the + /// Indicates to tokio that the file descriptor is no longer ready. All + /// internal readiness flags will be cleared, and tokio will wait for the /// next edge-triggered readiness notification from the OS. /// + /// This function is commonly used with guards returned by [`AsyncFd::readable_mut`] and + /// [`AsyncFd::writable_mut`]. + /// /// It is critical that this function not be called unless your code /// _actually observes_ that the file descriptor is _not_ ready. Do not call /// it simply because, for example, a read succeeded; it should be called /// when a read is observed to block. - /// - /// [`drop`]: method@std::mem::drop pub fn clear_ready(&mut self) { if let Some(event) = self.event.take() { self.async_fd.registration.clear_readiness(event); } } + /// Indicates to tokio that the file descriptor no longer has a specific readiness. + /// The internal readiness flag will be cleared, and tokio will wait for the + /// next edge-triggered readiness notification from the OS. + /// + /// This function is useful in combination with the [`AsyncFd::ready_mut`] method when a + /// combined interest like `Interest::READABLE | Interest::WRITABLE` is used. + /// + /// It is critical that this function not be called unless your code + /// _actually observes_ that the file descriptor is _not_ ready for the provided `Ready`. + /// Do not call it simply because, for example, a read succeeded; it should be called + /// when a read is observed to block. Only clear the specific readiness that is observed to + /// block. For example when a read blocks when using a combined interest, + /// only clear `Ready::READABLE`. + /// + /// # Examples + /// + /// Concurrently read and write to a [`std::net::TcpStream`] on the same task without + /// splitting. + /// + /// ```no_run + /// use std::error::Error; + /// use std::io; + /// use std::io::{Read, Write}; + /// use std::net::TcpStream; + /// use tokio::io::unix::AsyncFd; + /// use tokio::io::{Interest, Ready}; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// let stream = TcpStream::connect("127.0.0.1:8080")?; + /// stream.set_nonblocking(true)?; + /// let mut stream = AsyncFd::new(stream)?; + /// + /// loop { + /// let mut guard = stream + /// .ready_mut(Interest::READABLE | Interest::WRITABLE) + /// .await?; + /// + /// if guard.ready().is_readable() { + /// let mut data = vec![0; 1024]; + /// // Try to read data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match guard.get_inner_mut().read(&mut data) { + /// Ok(n) => { + /// println!("read {} bytes", n); + /// } + /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + /// // a read has blocked, but a write might still succeed. + /// // clear only the read readiness. + /// guard.clear_ready_matching(Ready::READABLE); + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// + /// if guard.ready().is_writable() { + /// // Try to write data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match guard.get_inner_mut().write(b"hello world") { + /// Ok(n) => { + /// println!("write {} bytes", n); + /// } + /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + /// // a write has blocked, but a read might still succeed. + /// // clear only the write readiness. + /// guard.clear_ready_matching(Ready::WRITABLE); + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// } + /// } + /// ``` + pub fn clear_ready_matching(&mut self, ready: Ready) { + if let Some(mut event) = self.event.take() { + self.async_fd + .registration + .clear_readiness(event.with_ready(ready)); + + // the event is no longer ready for the readiness that was just cleared + event.ready = event.ready - ready; + + if !event.ready.is_empty() { + self.event = Some(event); + } + } + } + /// This method should be invoked when you intentionally want to keep the /// ready flag asserted. /// @@ -578,6 +1158,20 @@ impl<'a, Inner: AsRawFd> AsyncFdReadyMutGuard<'a, Inner> { // no-op } + /// Get the [`Ready`] value associated with this guard. + /// + /// This method will return the empty readiness state if + /// [`AsyncFdReadyGuard::clear_ready`] has been called on + /// the guard. + /// + /// [`Ready`]: super::Ready + pub fn ready(&self) -> Ready { + match &self.event { + Some(event) => event.ready, + None => Ready::EMPTY, + } + } + /// Performs the provided IO operation. /// /// If `f` returns a [`WouldBlock`] error, the readiness state associated @@ -598,7 +1192,7 @@ impl<'a, Inner: AsRawFd> AsyncFdReadyMutGuard<'a, Inner> { &mut self, f: impl FnOnce(&mut AsyncFd<Inner>) -> io::Result<R>, ) -> Result<io::Result<R>, TryIoError> { - let result = f(&mut self.async_fd); + let result = f(self.async_fd); if let Err(e) = result.as_ref() { if e.kind() == io::ErrorKind::WouldBlock { diff --git a/vendor/tokio/src/io/blocking.rs b/vendor/tokio/src/io/blocking.rs index 94a3484ff..416573e97 100644 --- a/vendor/tokio/src/io/blocking.rs +++ b/vendor/tokio/src/io/blocking.rs @@ -16,7 +16,7 @@ use self::State::*; pub(crate) struct Blocking<T> { inner: Option<T>, state: State<T>, - /// `true` if the lower IO layer needs flushing + /// `true` if the lower IO layer needs flushing. need_flush: bool, } @@ -26,7 +26,7 @@ pub(crate) struct Buf { pos: usize, } -pub(crate) const MAX_BUF: usize = 16 * 1024; +pub(crate) const MAX_BUF: usize = 2 * 1024 * 1024; #[derive(Debug)] enum State<T> { @@ -34,8 +34,9 @@ enum State<T> { Busy(sys::Blocking<(io::Result<usize>, Buf, T)>), } -cfg_io_std! { +cfg_io_blocking! { impl<T> Blocking<T> { + #[cfg_attr(feature = "fs", allow(dead_code))] pub(crate) fn new(inner: T) -> Blocking<T> { Blocking { inner: Some(inner), @@ -175,7 +176,7 @@ where } } -/// Repeats operations that are interrupted +/// Repeats operations that are interrupted. macro_rules! uninterruptibly { ($e:expr) => {{ loop { diff --git a/vendor/tokio/src/io/bsd/poll_aio.rs b/vendor/tokio/src/io/bsd/poll_aio.rs new file mode 100644 index 000000000..6ac9e2880 --- /dev/null +++ b/vendor/tokio/src/io/bsd/poll_aio.rs @@ -0,0 +1,197 @@ +//! Use POSIX AIO futures with Tokio. + +use crate::io::interest::Interest; +use crate::runtime::io::{ReadyEvent, Registration}; +use crate::runtime::scheduler; +use mio::event::Source; +use mio::Registry; +use mio::Token; +use std::fmt; +use std::io; +use std::ops::{Deref, DerefMut}; +use std::os::unix::io::AsRawFd; +use std::os::unix::prelude::RawFd; +use std::task::{Context, Poll}; + +/// Like [`mio::event::Source`], but for POSIX AIO only. +/// +/// Tokio's consumer must pass an implementor of this trait to create a +/// [`Aio`] object. +pub trait AioSource { + /// Registers this AIO event source with Tokio's reactor. + fn register(&mut self, kq: RawFd, token: usize); + + /// Deregisters this AIO event source with Tokio's reactor. + fn deregister(&mut self); +} + +/// Wraps the user's AioSource in order to implement mio::event::Source, which +/// is what the rest of the crate wants. +struct MioSource<T>(T); + +impl<T: AioSource> Source for MioSource<T> { + fn register( + &mut self, + registry: &Registry, + token: Token, + interests: mio::Interest, + ) -> io::Result<()> { + assert!(interests.is_aio() || interests.is_lio()); + self.0.register(registry.as_raw_fd(), usize::from(token)); + Ok(()) + } + + fn deregister(&mut self, _registry: &Registry) -> io::Result<()> { + self.0.deregister(); + Ok(()) + } + + fn reregister( + &mut self, + registry: &Registry, + token: Token, + interests: mio::Interest, + ) -> io::Result<()> { + assert!(interests.is_aio() || interests.is_lio()); + self.0.register(registry.as_raw_fd(), usize::from(token)); + Ok(()) + } +} + +/// Associates a POSIX AIO control block with the reactor that drives it. +/// +/// `Aio`'s wrapped type must implement [`AioSource`] to be driven +/// by the reactor. +/// +/// The wrapped source may be accessed through the `Aio` via the `Deref` and +/// `DerefMut` traits. +/// +/// ## Clearing readiness +/// +/// If [`Aio::poll_ready`] returns ready, but the consumer determines that the +/// Source is not completely ready and must return to the Pending state, +/// [`Aio::clear_ready`] may be used. This can be useful with +/// [`lio_listio`], which may generate a kevent when only a portion of the +/// operations have completed. +/// +/// ## Platforms +/// +/// Only FreeBSD implements POSIX AIO with kqueue notification, so +/// `Aio` is only available for that operating system. +/// +/// [`lio_listio`]: https://pubs.opengroup.org/onlinepubs/9699919799/functions/lio_listio.html +// Note: Unlike every other kqueue event source, POSIX AIO registers events not +// via kevent(2) but when the aiocb is submitted to the kernel via aio_read, +// aio_write, etc. It needs the kqueue's file descriptor to do that. So +// AsyncFd can't be used for POSIX AIO. +// +// Note that Aio doesn't implement Drop. There's no need. Unlike other +// kqueue sources, simply dropping the object effectively deregisters it. +pub struct Aio<E> { + io: MioSource<E>, + registration: Registration, +} + +// ===== impl Aio ===== + +impl<E: AioSource> Aio<E> { + /// Creates a new `Aio` suitable for use with POSIX AIO functions. + /// + /// It will be associated with the default reactor. The runtime is usually + /// set implicitly when this function is called from a future driven by a + /// Tokio runtime, otherwise runtime can be set explicitly with + /// [`Runtime::enter`](crate::runtime::Runtime::enter) function. + pub fn new_for_aio(io: E) -> io::Result<Self> { + Self::new_with_interest(io, Interest::AIO) + } + + /// Creates a new `Aio` suitable for use with [`lio_listio`]. + /// + /// It will be associated with the default reactor. The runtime is usually + /// set implicitly when this function is called from a future driven by a + /// Tokio runtime, otherwise runtime can be set explicitly with + /// [`Runtime::enter`](crate::runtime::Runtime::enter) function. + /// + /// [`lio_listio`]: https://pubs.opengroup.org/onlinepubs/9699919799/functions/lio_listio.html + pub fn new_for_lio(io: E) -> io::Result<Self> { + Self::new_with_interest(io, Interest::LIO) + } + + fn new_with_interest(io: E, interest: Interest) -> io::Result<Self> { + let mut io = MioSource(io); + let handle = scheduler::Handle::current(); + let registration = Registration::new_with_interest_and_handle(&mut io, interest, handle)?; + Ok(Self { io, registration }) + } + + /// Indicates to Tokio that the source is no longer ready. The internal + /// readiness flag will be cleared, and tokio will wait for the next + /// edge-triggered readiness notification from the OS. + /// + /// It is critical that this method not be called unless your code + /// _actually observes_ that the source is _not_ ready. The OS must + /// deliver a subsequent notification, or this source will block + /// forever. It is equally critical that you `do` call this method if you + /// resubmit the same structure to the kernel and poll it again. + /// + /// This method is not very useful with AIO readiness, since each `aiocb` + /// structure is typically only used once. It's main use with + /// [`lio_listio`], which will sometimes send notification when only a + /// portion of its elements are complete. In that case, the caller must + /// call `clear_ready` before resubmitting it. + /// + /// [`lio_listio`]: https://pubs.opengroup.org/onlinepubs/9699919799/functions/lio_listio.html + pub fn clear_ready(&self, ev: AioEvent) { + self.registration.clear_readiness(ev.0) + } + + /// Destroy the [`Aio`] and return its inner source. + pub fn into_inner(self) -> E { + self.io.0 + } + + /// Polls for readiness. Either AIO or LIO counts. + /// + /// This method returns: + /// * `Poll::Pending` if the underlying operation is not complete, whether + /// or not it completed successfully. This will be true if the OS is + /// still processing it, or if it has not yet been submitted to the OS. + /// * `Poll::Ready(Ok(_))` if the underlying operation is complete. + /// * `Poll::Ready(Err(_))` if the reactor has been shutdown. This does + /// _not_ indicate that the underlying operation encountered an error. + /// + /// When the method returns `Poll::Pending`, the `Waker` in the provided `Context` + /// is scheduled to receive a wakeup when the underlying operation + /// completes. Note that on multiple calls to `poll_ready`, only the `Waker` from the + /// `Context` passed to the most recent call is scheduled to receive a wakeup. + pub fn poll_ready<'a>(&'a self, cx: &mut Context<'_>) -> Poll<io::Result<AioEvent>> { + let ev = ready!(self.registration.poll_read_ready(cx))?; + Poll::Ready(Ok(AioEvent(ev))) + } +} + +impl<E: AioSource> Deref for Aio<E> { + type Target = E; + + fn deref(&self) -> &E { + &self.io.0 + } +} + +impl<E: AioSource> DerefMut for Aio<E> { + fn deref_mut(&mut self) -> &mut E { + &mut self.io.0 + } +} + +impl<E: AioSource + fmt::Debug> fmt::Debug for Aio<E> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Aio").field("io", &self.io.0).finish() + } +} + +/// Opaque data returned by [`Aio::poll_ready`]. +/// +/// It can be fed back to [`Aio::clear_ready`]. +#[derive(Debug)] +pub struct AioEvent(ReadyEvent); diff --git a/vendor/tokio/src/io/driver/mod.rs b/vendor/tokio/src/io/driver/mod.rs deleted file mode 100644 index 3aa0cfbb2..000000000 --- a/vendor/tokio/src/io/driver/mod.rs +++ /dev/null @@ -1,353 +0,0 @@ -#![cfg_attr(not(feature = "rt"), allow(dead_code))] - -mod interest; -#[allow(unreachable_pub)] -pub use interest::Interest; - -mod ready; -#[allow(unreachable_pub)] -pub use ready::Ready; - -mod registration; -pub(crate) use registration::Registration; - -mod scheduled_io; -use scheduled_io::ScheduledIo; - -use crate::park::{Park, Unpark}; -use crate::util::slab::{self, Slab}; -use crate::{loom::sync::Mutex, util::bit}; - -use std::fmt; -use std::io; -use std::sync::{Arc, Weak}; -use std::time::Duration; - -/// I/O driver, backed by Mio -pub(crate) struct Driver { - /// Tracks the number of times `turn` is called. It is safe for this to wrap - /// as it is mostly used to determine when to call `compact()` - tick: u8, - - /// Reuse the `mio::Events` value across calls to poll. - events: Option<mio::Events>, - - /// Primary slab handle containing the state for each resource registered - /// with this driver. During Drop this is moved into the Inner structure, so - /// this is an Option to allow it to be vacated (until Drop this is always - /// Some) - resources: Option<Slab<ScheduledIo>>, - - /// The system event queue - poll: mio::Poll, - - /// State shared between the reactor and the handles. - inner: Arc<Inner>, -} - -/// A reference to an I/O driver -#[derive(Clone)] -pub(crate) struct Handle { - inner: Weak<Inner>, -} - -pub(crate) struct ReadyEvent { - tick: u8, - pub(crate) ready: Ready, -} - -pub(super) struct Inner { - /// Primary slab handle containing the state for each resource registered - /// with this driver. - /// - /// The ownership of this slab is moved into this structure during - /// `Driver::drop`, so that `Inner::drop` can notify all outstanding handles - /// without risking new ones being registered in the meantime. - resources: Mutex<Option<Slab<ScheduledIo>>>, - - /// Registers I/O resources - registry: mio::Registry, - - /// Allocates `ScheduledIo` handles when creating new resources. - pub(super) io_dispatch: slab::Allocator<ScheduledIo>, - - /// Used to wake up the reactor from a call to `turn` - waker: mio::Waker, -} - -#[derive(Debug, Eq, PartialEq, Clone, Copy)] -enum Direction { - Read, - Write, -} - -enum Tick { - Set(u8), - Clear(u8), -} - -// TODO: Don't use a fake token. Instead, reserve a slot entry for the wakeup -// token. -const TOKEN_WAKEUP: mio::Token = mio::Token(1 << 31); - -const ADDRESS: bit::Pack = bit::Pack::least_significant(24); - -// Packs the generation value in the `readiness` field. -// -// The generation prevents a race condition where a slab slot is reused for a -// new socket while the I/O driver is about to apply a readiness event. The -// generation value is checked when setting new readiness. If the generation do -// not match, then the readiness event is discarded. -const GENERATION: bit::Pack = ADDRESS.then(7); - -fn _assert_kinds() { - fn _assert<T: Send + Sync>() {} - - _assert::<Handle>(); -} - -// ===== impl Driver ===== - -impl Driver { - /// Creates a new event loop, returning any error that happened during the - /// creation. - pub(crate) fn new() -> io::Result<Driver> { - let poll = mio::Poll::new()?; - let waker = mio::Waker::new(poll.registry(), TOKEN_WAKEUP)?; - let registry = poll.registry().try_clone()?; - - let slab = Slab::new(); - let allocator = slab.allocator(); - - Ok(Driver { - tick: 0, - events: Some(mio::Events::with_capacity(1024)), - poll, - resources: Some(slab), - inner: Arc::new(Inner { - resources: Mutex::new(None), - registry, - io_dispatch: allocator, - waker, - }), - }) - } - - /// Returns a handle to this event loop which can be sent across threads - /// and can be used as a proxy to the event loop itself. - /// - /// Handles are cloneable and clones always refer to the same event loop. - /// This handle is typically passed into functions that create I/O objects - /// to bind them to this event loop. - pub(crate) fn handle(&self) -> Handle { - Handle { - inner: Arc::downgrade(&self.inner), - } - } - - fn turn(&mut self, max_wait: Option<Duration>) -> io::Result<()> { - // How often to call `compact()` on the resource slab - const COMPACT_INTERVAL: u8 = 255; - - self.tick = self.tick.wrapping_add(1); - - if self.tick == COMPACT_INTERVAL { - self.resources.as_mut().unwrap().compact() - } - - let mut events = self.events.take().expect("i/o driver event store missing"); - - // Block waiting for an event to happen, peeling out how many events - // happened. - match self.poll.poll(&mut events, max_wait) { - Ok(_) => {} - Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {} - Err(e) => return Err(e), - } - - // Process all the events that came in, dispatching appropriately - for event in events.iter() { - let token = event.token(); - - if token != TOKEN_WAKEUP { - self.dispatch(token, Ready::from_mio(event)); - } - } - - self.events = Some(events); - - Ok(()) - } - - fn dispatch(&mut self, token: mio::Token, ready: Ready) { - let addr = slab::Address::from_usize(ADDRESS.unpack(token.0)); - - let resources = self.resources.as_mut().unwrap(); - - let io = match resources.get(addr) { - Some(io) => io, - None => return, - }; - - let res = io.set_readiness(Some(token.0), Tick::Set(self.tick), |curr| curr | ready); - - if res.is_err() { - // token no longer valid! - return; - } - - io.wake(ready); - } -} - -impl Drop for Driver { - fn drop(&mut self) { - (*self.inner.resources.lock()) = self.resources.take(); - } -} - -impl Drop for Inner { - fn drop(&mut self) { - let resources = self.resources.lock().take(); - - if let Some(mut slab) = resources { - slab.for_each(|io| { - // If a task is waiting on the I/O resource, notify it. The task - // will then attempt to use the I/O resource and fail due to the - // driver being shutdown. - io.shutdown(); - }); - } - } -} - -impl Park for Driver { - type Unpark = Handle; - type Error = io::Error; - - fn unpark(&self) -> Self::Unpark { - self.handle() - } - - fn park(&mut self) -> io::Result<()> { - self.turn(None)?; - Ok(()) - } - - fn park_timeout(&mut self, duration: Duration) -> io::Result<()> { - self.turn(Some(duration))?; - Ok(()) - } - - fn shutdown(&mut self) {} -} - -impl fmt::Debug for Driver { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "Driver") - } -} - -// ===== impl Handle ===== - -cfg_rt! { - impl Handle { - /// Returns a handle to the current reactor - /// - /// # Panics - /// - /// This function panics if there is no current reactor set and `rt` feature - /// flag is not enabled. - pub(super) fn current() -> Self { - crate::runtime::context::io_handle().expect("A Tokio 1.x context was found, but IO is disabled. Call `enable_io` on the runtime builder to enable IO.") - } - } -} - -cfg_not_rt! { - impl Handle { - /// Returns a handle to the current reactor - /// - /// # Panics - /// - /// This function panics if there is no current reactor set, or if the `rt` - /// feature flag is not enabled. - pub(super) fn current() -> Self { - panic!("{}", crate::util::error::CONTEXT_MISSING_ERROR) - } - } -} - -impl Handle { - /// Forces a reactor blocked in a call to `turn` to wakeup, or otherwise - /// makes the next call to `turn` return immediately. - /// - /// This method is intended to be used in situations where a notification - /// needs to otherwise be sent to the main reactor. If the reactor is - /// currently blocked inside of `turn` then it will wake up and soon return - /// after this method has been called. If the reactor is not currently - /// blocked in `turn`, then the next call to `turn` will not block and - /// return immediately. - fn wakeup(&self) { - if let Some(inner) = self.inner() { - inner.waker.wake().expect("failed to wake I/O driver"); - } - } - - pub(super) fn inner(&self) -> Option<Arc<Inner>> { - self.inner.upgrade() - } -} - -impl Unpark for Handle { - fn unpark(&self) { - self.wakeup(); - } -} - -impl fmt::Debug for Handle { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "Handle") - } -} - -// ===== impl Inner ===== - -impl Inner { - /// Registers an I/O resource with the reactor for a given `mio::Ready` state. - /// - /// The registration token is returned. - pub(super) fn add_source( - &self, - source: &mut impl mio::event::Source, - interest: Interest, - ) -> io::Result<slab::Ref<ScheduledIo>> { - let (address, shared) = self.io_dispatch.allocate().ok_or_else(|| { - io::Error::new( - io::ErrorKind::Other, - "reactor at max registered I/O resources", - ) - })?; - - let token = GENERATION.pack(shared.generation(), ADDRESS.pack(address.as_usize(), 0)); - - self.registry - .register(source, mio::Token(token), interest.to_mio())?; - - Ok(shared) - } - - /// Deregisters an I/O resource from the reactor. - pub(super) fn deregister_source(&self, source: &mut impl mio::event::Source) -> io::Result<()> { - self.registry.deregister(source) - } -} - -impl Direction { - pub(super) fn mask(self) -> Ready { - match self { - Direction::Read => Ready::READABLE | Ready::READ_CLOSED, - Direction::Write => Ready::WRITABLE | Ready::WRITE_CLOSED, - } - } -} diff --git a/vendor/tokio/src/io/driver/platform.rs b/vendor/tokio/src/io/driver/platform.rs deleted file mode 100644 index 6b27988ce..000000000 --- a/vendor/tokio/src/io/driver/platform.rs +++ /dev/null @@ -1,44 +0,0 @@ -pub(crate) use self::sys::*; - -#[cfg(unix)] -mod sys { - use mio::unix::UnixReady; - use mio::Ready; - - pub(crate) fn hup() -> Ready { - UnixReady::hup().into() - } - - pub(crate) fn is_hup(ready: Ready) -> bool { - UnixReady::from(ready).is_hup() - } - - pub(crate) fn error() -> Ready { - UnixReady::error().into() - } - - pub(crate) fn is_error(ready: Ready) -> bool { - UnixReady::from(ready).is_error() - } -} - -#[cfg(windows)] -mod sys { - use mio::Ready; - - pub(crate) fn hup() -> Ready { - Ready::empty() - } - - pub(crate) fn is_hup(_: Ready) -> bool { - false - } - - pub(crate) fn error() -> Ready { - Ready::empty() - } - - pub(crate) fn is_error(_: Ready) -> bool { - false - } -} diff --git a/vendor/tokio/src/io/driver/interest.rs b/vendor/tokio/src/io/interest.rs index 36951cf5a..3a39cf761 100644 --- a/vendor/tokio/src/io/driver/interest.rs +++ b/vendor/tokio/src/io/interest.rs @@ -1,11 +1,11 @@ #![cfg_attr(not(feature = "net"), allow(dead_code, unreachable_pub))] -use crate::io::driver::Ready; +use crate::io::ready::Ready; use std::fmt; use std::ops; -/// Readiness event interest +/// Readiness event interest. /// /// Specifies the readiness events the caller is interested in when awaiting on /// I/O resource readiness states. @@ -14,16 +14,41 @@ use std::ops; pub struct Interest(mio::Interest); impl Interest { + // The non-FreeBSD definitions in this block are active only when + // building documentation. + cfg_aio! { + /// Interest for POSIX AIO. + #[cfg(target_os = "freebsd")] + pub const AIO: Interest = Interest(mio::Interest::AIO); + + /// Interest for POSIX AIO. + #[cfg(not(target_os = "freebsd"))] + pub const AIO: Interest = Interest(mio::Interest::READABLE); + + /// Interest for POSIX AIO lio_listio events. + #[cfg(target_os = "freebsd")] + pub const LIO: Interest = Interest(mio::Interest::LIO); + + /// Interest for POSIX AIO lio_listio events. + #[cfg(not(target_os = "freebsd"))] + pub const LIO: Interest = Interest(mio::Interest::READABLE); + } + /// Interest in all readable events. /// /// Readable interest includes read-closed events. pub const READABLE: Interest = Interest(mio::Interest::READABLE); - /// Interest in all writable events + /// Interest in all writable events. /// /// Writable interest includes write-closed events. pub const WRITABLE: Interest = Interest(mio::Interest::WRITABLE); + /// Returns a `Interest` set representing priority completion interests. + #[cfg(any(target_os = "linux", target_os = "android"))] + #[cfg_attr(docsrs, doc(cfg(any(target_os = "linux", target_os = "android"))))] + pub const PRIORITY: Interest = Interest(mio::Interest::PRIORITY); + /// Returns true if the value includes readable interest. /// /// # Examples @@ -58,6 +83,25 @@ impl Interest { self.0.is_writable() } + /// Returns true if the value includes priority interest. + /// + /// # Examples + /// + /// ``` + /// use tokio::io::Interest; + /// + /// assert!(!Interest::READABLE.is_priority()); + /// assert!(Interest::PRIORITY.is_priority()); + /// + /// let both = Interest::READABLE | Interest::PRIORITY; + /// assert!(both.is_priority()); + /// ``` + #[cfg(any(target_os = "linux", target_os = "android"))] + #[cfg_attr(docsrs, doc(cfg(any(target_os = "linux", target_os = "android"))))] + pub const fn is_priority(self) -> bool { + self.0.is_priority() + } + /// Add together two `Interest` values. /// /// This function works from a `const` context. @@ -80,10 +124,12 @@ impl Interest { self.0 } - pub(super) fn mask(self) -> Ready { + pub(crate) fn mask(self) -> Ready { match self { Interest::READABLE => Ready::READABLE | Ready::READ_CLOSED, Interest::WRITABLE => Ready::WRITABLE | Ready::WRITE_CLOSED, + #[cfg(any(target_os = "linux", target_os = "android"))] + Interest::PRIORITY => Ready::PRIORITY | Ready::READ_CLOSED, _ => Ready::EMPTY, } } diff --git a/vendor/tokio/src/io/mod.rs b/vendor/tokio/src/io/mod.rs index 14a4a6304..666d69cb4 100644 --- a/vendor/tokio/src/io/mod.rs +++ b/vendor/tokio/src/io/mod.rs @@ -1,5 +1,3 @@ -#![cfg_attr(loom, allow(dead_code, unreachable_pub))] - //! Traits, helpers, and type definitions for asynchronous I/O functionality. //! //! This module is the asynchronous version of `std::io`. Primarily, it @@ -132,19 +130,23 @@ //! other words, these types must never block the thread, and instead the //! current task is notified when the I/O resource is ready. //! -//! ## Conversion to and from Sink/Stream +//! ## Conversion to and from Stream/Sink +//! +//! It is often convenient to encapsulate the reading and writing of bytes in a +//! [`Stream`] or [`Sink`] of data. +//! +//! Tokio provides simple wrappers for converting [`AsyncRead`] to [`Stream`] +//! and vice-versa in the [tokio-util] crate, see [`ReaderStream`] and +//! [`StreamReader`]. //! -//! It is often convenient to encapsulate the reading and writing of -//! bytes and instead work with a [`Sink`] or [`Stream`] of some data -//! type that is encoded as bytes and/or decoded from bytes. Tokio -//! provides some utility traits in the [tokio-util] crate that -//! abstract the asynchronous buffering that is required and allows -//! you to write [`Encoder`] and [`Decoder`] functions working with a -//! buffer of bytes, and then use that ["codec"] to transform anything -//! that implements [`AsyncRead`] and [`AsyncWrite`] into a `Sink`/`Stream` of -//! your structured data. +//! There are also utility traits that abstract the asynchronous buffering +//! necessary to write your own adaptors for encoding and decoding bytes to/from +//! your structured data, allowing to transform something that implements +//! [`AsyncRead`]/[`AsyncWrite`] into a [`Stream`]/[`Sink`], see [`Decoder`] and +//! [`Encoder`] in the [tokio-util::codec] module. //! -//! [tokio-util]: https://docs.rs/tokio-util/0.6/tokio_util/codec/index.html +//! [tokio-util]: https://docs.rs/tokio-util +//! [tokio-util::codec]: https://docs.rs/tokio-util/latest/tokio_util/codec/index.html //! //! # Standard input and output //! @@ -169,9 +171,11 @@ //! [`AsyncWrite`]: trait@AsyncWrite //! [`AsyncReadExt`]: trait@AsyncReadExt //! [`AsyncWriteExt`]: trait@AsyncWriteExt -//! ["codec"]: https://docs.rs/tokio-util/0.6/tokio_util/codec/index.html -//! [`Encoder`]: https://docs.rs/tokio-util/0.6/tokio_util/codec/trait.Encoder.html -//! [`Decoder`]: https://docs.rs/tokio-util/0.6/tokio_util/codec/trait.Decoder.html +//! ["codec"]: https://docs.rs/tokio-util/latest/tokio_util/codec/index.html +//! [`Encoder`]: https://docs.rs/tokio-util/latest/tokio_util/codec/trait.Encoder.html +//! [`Decoder`]: https://docs.rs/tokio-util/latest/tokio_util/codec/trait.Decoder.html +//! [`ReaderStream`]: https://docs.rs/tokio-util/latest/tokio_util/io/struct.ReaderStream.html +//! [`StreamReader`]: https://docs.rs/tokio-util/latest/tokio_util/io/struct.StreamReader.html //! [`Error`]: struct@Error //! [`ErrorKind`]: enum@ErrorKind //! [`Result`]: type@Result @@ -180,6 +184,12 @@ //! [`Sink`]: https://docs.rs/futures/0.3/futures/sink/trait.Sink.html //! [`Stream`]: https://docs.rs/futures/0.3/futures/stream/trait.Stream.html //! [`Write`]: std::io::Write + +#![cfg_attr( + not(all(feature = "rt", feature = "net")), + allow(dead_code, unused_imports) +)] + cfg_io_blocking! { pub(crate) mod blocking; } @@ -205,18 +215,31 @@ pub use self::read_buf::ReadBuf; pub use std::io::{Error, ErrorKind, Result, SeekFrom}; cfg_io_driver_impl! { - pub(crate) mod driver; + pub(crate) mod interest; + pub(crate) mod ready; cfg_net! { - pub use driver::{Interest, Ready}; + pub use interest::Interest; + pub use ready::Ready; } + #[cfg_attr(tokio_wasi, allow(unused_imports))] mod poll_evented; #[cfg(not(loom))] + #[cfg_attr(tokio_wasi, allow(unused_imports))] pub(crate) use poll_evented::PollEvented; } +cfg_aio! { + /// BSD-specific I/O types. + pub mod bsd { + mod poll_aio; + + pub use poll_aio::{Aio, AioEvent, AioSource}; + } +} + cfg_net_unix! { mod async_fd; diff --git a/vendor/tokio/src/io/poll_evented.rs b/vendor/tokio/src/io/poll_evented.rs index a31e6db7b..875ad7e0d 100644 --- a/vendor/tokio/src/io/poll_evented.rs +++ b/vendor/tokio/src/io/poll_evented.rs @@ -1,16 +1,19 @@ -use crate::io::driver::{Handle, Interest, Registration}; +use crate::io::interest::Interest; +use crate::runtime::io::Registration; +use crate::runtime::scheduler; use mio::event::Source; use std::fmt; use std::io; use std::ops::Deref; +use std::panic::{RefUnwindSafe, UnwindSafe}; cfg_io_driver! { /// Associates an I/O resource that implements the [`std::io::Read`] and/or /// [`std::io::Write`] traits with the reactor that drives it. /// /// `PollEvented` uses [`Registration`] internally to take a type that - /// implements [`mio::event::Source`] as well as [`std::io::Read`] and or + /// implements [`mio::event::Source`] as well as [`std::io::Read`] and/or /// [`std::io::Write`] and associate it with a reactor that will drive it. /// /// Once the [`mio::event::Source`] type is wrapped by `PollEvented`, it can be @@ -40,13 +43,12 @@ cfg_io_driver! { /// [`poll_read_ready`] again will also indicate read readiness. /// /// When the operation is attempted and is unable to succeed due to the I/O - /// resource not being ready, the caller must call `clear_read_ready` or - /// `clear_write_ready`. This clears the readiness state until a new - /// readiness event is received. + /// resource not being ready, the caller must call [`clear_readiness`]. + /// This clears the readiness state until a new readiness event is received. /// /// This allows the caller to implement additional functions. For example, /// [`TcpListener`] implements poll_accept by using [`poll_read_ready`] and - /// `clear_read_ready`. + /// [`clear_readiness`]. /// /// ## Platform-specific events /// @@ -57,6 +59,7 @@ cfg_io_driver! { /// [`AsyncRead`]: crate::io::AsyncRead /// [`AsyncWrite`]: crate::io::AsyncWrite /// [`TcpListener`]: crate::net::TcpListener + /// [`clear_readiness`]: Registration::clear_readiness /// [`poll_read_ready`]: Registration::poll_read_ready /// [`poll_write_ready`]: Registration::poll_write_ready pub(crate) struct PollEvented<E: Source> { @@ -70,6 +73,9 @@ cfg_io_driver! { impl<E: Source> PollEvented<E> { /// Creates a new `PollEvented` associated with the default reactor. /// + /// The returned `PollEvented` has readable and writable interests. For more control, use + /// [`Self::new_with_interest`]. + /// /// # Panics /// /// This function panics if thread-local runtime is not set. @@ -77,6 +83,7 @@ impl<E: Source> PollEvented<E> { /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. + #[track_caller] #[cfg_attr(feature = "signal", allow(unused))] pub(crate) fn new(io: E) -> io::Result<Self> { PollEvented::new_with_interest(io, Interest::READABLE | Interest::WRITABLE) @@ -97,15 +104,17 @@ impl<E: Source> PollEvented<E> { /// a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) /// function. + #[track_caller] #[cfg_attr(feature = "signal", allow(unused))] pub(crate) fn new_with_interest(io: E, interest: Interest) -> io::Result<Self> { - Self::new_with_interest_and_handle(io, interest, Handle::current()) + Self::new_with_interest_and_handle(io, interest, scheduler::Handle::current()) } + #[track_caller] pub(crate) fn new_with_interest_and_handle( mut io: E, interest: Interest, - handle: Handle, + handle: scheduler::Handle, ) -> io::Result<Self> { let registration = Registration::new_with_interest_and_handle(&mut io, interest, handle)?; Ok(Self { @@ -114,17 +123,13 @@ impl<E: Source> PollEvented<E> { }) } - /// Returns a reference to the registration - #[cfg(any( - feature = "net", - all(unix, feature = "process"), - all(unix, feature = "signal"), - ))] + /// Returns a reference to the registration. + #[cfg(any(feature = "net"))] pub(crate) fn registration(&self) -> &Registration { &self.registration } - /// Deregister the inner io from the registration and returns a Result containing the inner io + /// Deregisters the inner io from the registration and returns a Result containing the inner io. #[cfg(any(feature = "net", feature = "process"))] pub(crate) fn into_inner(mut self) -> io::Result<E> { let mut inner = self.io.take().unwrap(); // As io shouldn't ever be None, just unwrap here. @@ -134,7 +139,7 @@ impl<E: Source> PollEvented<E> { } feature! { - #![any(feature = "net", feature = "process")] + #![any(feature = "net", all(unix, feature = "process"))] use crate::io::ReadBuf; use std::task::{Context, Poll}; @@ -151,16 +156,32 @@ feature! { { use std::io::Read; - let n = ready!(self.registration.poll_read_io(cx, || { + loop { + let evt = ready!(self.registration.poll_read_ready(cx))?; + let b = &mut *(buf.unfilled_mut() as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]); - self.io.as_ref().unwrap().read(b) - }))?; - - // Safety: We trust `TcpStream::read` to have filled up `n` bytes in the - // buffer. - buf.assume_init(n); - buf.advance(n); - Poll::Ready(Ok(())) + let len = b.len(); + + match self.io.as_ref().unwrap().read(b) { + Ok(n) => { + // if we read a partially full buffer, this is sufficient on unix to show + // that the socket buffer has been drained + if n > 0 && (!cfg!(windows) && n < len) { + self.registration.clear_readiness(evt); + } + + // Safety: We trust `TcpStream::read` to have filled up `n` bytes in the + // buffer. + buf.assume_init(n); + buf.advance(n); + return Poll::Ready(Ok(())); + }, + Err(e) if e.kind() == io::ErrorKind::WouldBlock => { + self.registration.clear_readiness(evt); + } + Err(e) => return Poll::Ready(Err(e)), + } + } } pub(crate) fn poll_write<'a>(&'a self, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> @@ -168,10 +189,29 @@ feature! { &'a E: io::Write + 'a, { use std::io::Write; - self.registration.poll_write_io(cx, || self.io.as_ref().unwrap().write(buf)) + + loop { + let evt = ready!(self.registration.poll_write_ready(cx))?; + + match self.io.as_ref().unwrap().write(buf) { + Ok(n) => { + // if we write only part of our buffer, this is sufficient on unix to show + // that the socket buffer is full + if n > 0 && (!cfg!(windows) && n < buf.len()) { + self.registration.clear_readiness(evt); + } + + return Poll::Ready(Ok(n)); + }, + Err(e) if e.kind() == io::ErrorKind::WouldBlock => { + self.registration.clear_readiness(evt); + } + Err(e) => return Poll::Ready(Err(e)), + } + } } - #[cfg(feature = "net")] + #[cfg(any(feature = "net", feature = "process"))] pub(crate) fn poll_write_vectored<'a>( &'a self, cx: &mut Context<'_>, @@ -186,6 +226,10 @@ feature! { } } +impl<E: Source> UnwindSafe for PollEvented<E> {} + +impl<E: Source> RefUnwindSafe for PollEvented<E> {} + impl<E: Source> Deref for PollEvented<E> { type Target = E; diff --git a/vendor/tokio/src/io/read_buf.rs b/vendor/tokio/src/io/read_buf.rs index ad58cbe75..283d96e30 100644 --- a/vendor/tokio/src/io/read_buf.rs +++ b/vendor/tokio/src/io/read_buf.rs @@ -1,9 +1,5 @@ -// This lint claims ugly casting is somehow safer than transmute, but there's -// no evidence that is the case. Shush. -#![allow(clippy::transmute_ptr_to_ptr)] - use std::fmt; -use std::mem::{self, MaybeUninit}; +use std::mem::MaybeUninit; /// A wrapper around a byte buffer that is incrementally filled and initialized. /// @@ -35,7 +31,7 @@ impl<'a> ReadBuf<'a> { #[inline] pub fn new(buf: &'a mut [u8]) -> ReadBuf<'a> { let initialized = buf.len(); - let buf = unsafe { mem::transmute::<&mut [u8], &mut [MaybeUninit<u8>]>(buf) }; + let buf = unsafe { slice_to_uninit_mut(buf) }; ReadBuf { buf, filled: 0, @@ -67,8 +63,7 @@ impl<'a> ReadBuf<'a> { let slice = &self.buf[..self.filled]; // safety: filled describes how far into the buffer that the // user has filled with bytes, so it's been initialized. - // TODO: This could use `MaybeUninit::slice_get_ref` when it is stable. - unsafe { mem::transmute::<&[MaybeUninit<u8>], &[u8]>(slice) } + unsafe { slice_assume_init(slice) } } /// Returns a mutable reference to the filled portion of the buffer. @@ -77,8 +72,7 @@ impl<'a> ReadBuf<'a> { let slice = &mut self.buf[..self.filled]; // safety: filled describes how far into the buffer that the // user has filled with bytes, so it's been initialized. - // TODO: This could use `MaybeUninit::slice_get_mut` when it is stable. - unsafe { mem::transmute::<&mut [MaybeUninit<u8>], &mut [u8]>(slice) } + unsafe { slice_assume_init_mut(slice) } } /// Returns a new `ReadBuf` comprised of the unfilled section up to `n`. @@ -97,8 +91,7 @@ impl<'a> ReadBuf<'a> { let slice = &self.buf[..self.initialized]; // safety: initialized describes how far into the buffer that the // user has at some point initialized with bytes. - // TODO: This could use `MaybeUninit::slice_get_ref` when it is stable. - unsafe { mem::transmute::<&[MaybeUninit<u8>], &[u8]>(slice) } + unsafe { slice_assume_init(slice) } } /// Returns a mutable reference to the initialized portion of the buffer. @@ -109,15 +102,14 @@ impl<'a> ReadBuf<'a> { let slice = &mut self.buf[..self.initialized]; // safety: initialized describes how far into the buffer that the // user has at some point initialized with bytes. - // TODO: This could use `MaybeUninit::slice_get_mut` when it is stable. - unsafe { mem::transmute::<&mut [MaybeUninit<u8>], &mut [u8]>(slice) } + unsafe { slice_assume_init_mut(slice) } } /// Returns a mutable reference to the entire buffer, without ensuring that it has been fully /// initialized. /// /// The elements between 0 and `self.filled().len()` are filled, and those between 0 and - /// `self.initialized().len()` are initialized (and so can be transmuted to a `&mut [u8]`). + /// `self.initialized().len()` are initialized (and so can be converted to a `&mut [u8]`). /// /// The caller of this method must ensure that these invariants are upheld. For example, if the /// caller initializes some of the uninitialized section of the buffer, it must call @@ -160,6 +152,7 @@ impl<'a> ReadBuf<'a> { /// /// Panics if `self.remaining()` is less than `n`. #[inline] + #[track_caller] pub fn initialize_unfilled_to(&mut self, n: usize) -> &mut [u8] { assert!(self.remaining() >= n, "n overflows remaining"); @@ -178,7 +171,7 @@ impl<'a> ReadBuf<'a> { let slice = &mut self.buf[self.filled..end]; // safety: just above, we checked that the end of the buf has // been initialized to some value. - unsafe { mem::transmute::<&mut [MaybeUninit<u8>], &mut [u8]>(slice) } + unsafe { slice_assume_init_mut(slice) } } /// Returns the number of bytes at the end of the slice that have not yet been filled. @@ -203,6 +196,7 @@ impl<'a> ReadBuf<'a> { /// /// Panics if the filled region of the buffer would become larger than the initialized region. #[inline] + #[track_caller] pub fn advance(&mut self, n: usize) { let new = self.filled.checked_add(n).expect("filled overflow"); self.set_filled(new); @@ -219,6 +213,7 @@ impl<'a> ReadBuf<'a> { /// /// Panics if the filled region of the buffer would become larger than the initialized region. #[inline] + #[track_caller] pub fn set_filled(&mut self, n: usize) { assert!( n <= self.initialized, @@ -249,6 +244,7 @@ impl<'a> ReadBuf<'a> { /// /// Panics if `self.remaining()` is less than `buf.len()`. #[inline] + #[track_caller] pub fn put_slice(&mut self, buf: &[u8]) { assert!( self.remaining() >= buf.len(), @@ -274,6 +270,33 @@ impl<'a> ReadBuf<'a> { } } +#[cfg(feature = "io-util")] +#[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] +unsafe impl<'a> bytes::BufMut for ReadBuf<'a> { + fn remaining_mut(&self) -> usize { + self.remaining() + } + + // SAFETY: The caller guarantees that at least `cnt` unfilled bytes have been initialized. + unsafe fn advance_mut(&mut self, cnt: usize) { + self.assume_init(cnt); + self.advance(cnt); + } + + fn chunk_mut(&mut self) -> &mut bytes::buf::UninitSlice { + // SAFETY: No region of `unfilled` will be deinitialized because it is + // exposed as an `UninitSlice`, whose API guarantees that the memory is + // never deinitialized. + let unfilled = unsafe { self.unfilled_mut() }; + let len = unfilled.len(); + let ptr = unfilled.as_mut_ptr() as *mut u8; + + // SAFETY: The pointer is valid for `len` bytes because it comes from a + // slice of that length. + unsafe { bytes::buf::UninitSlice::from_raw_parts_mut(ptr, len) } + } +} + impl fmt::Debug for ReadBuf<'_> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("ReadBuf") @@ -283,3 +306,17 @@ impl fmt::Debug for ReadBuf<'_> { .finish() } } + +unsafe fn slice_to_uninit_mut(slice: &mut [u8]) -> &mut [MaybeUninit<u8>] { + &mut *(slice as *mut [u8] as *mut [MaybeUninit<u8>]) +} + +// TODO: This could use `MaybeUninit::slice_assume_init` when it is stable. +unsafe fn slice_assume_init(slice: &[MaybeUninit<u8>]) -> &[u8] { + &*(slice as *const [MaybeUninit<u8>] as *const [u8]) +} + +// TODO: This could use `MaybeUninit::slice_assume_init_mut` when it is stable. +unsafe fn slice_assume_init_mut(slice: &mut [MaybeUninit<u8>]) -> &mut [u8] { + &mut *(slice as *mut [MaybeUninit<u8>] as *mut [u8]) +} diff --git a/vendor/tokio/src/io/driver/ready.rs b/vendor/tokio/src/io/ready.rs index 2ac01bdbe..33e2ef0b6 100644 --- a/vendor/tokio/src/io/driver/ready.rs +++ b/vendor/tokio/src/io/ready.rs @@ -7,12 +7,14 @@ const READABLE: usize = 0b0_01; const WRITABLE: usize = 0b0_10; const READ_CLOSED: usize = 0b0_0100; const WRITE_CLOSED: usize = 0b0_1000; +#[cfg(any(target_os = "linux", target_os = "android"))] +const PRIORITY: usize = 0b1_0000; /// Describes the readiness state of an I/O resources. /// /// `Ready` tracks which operation an I/O resource is ready to perform. #[cfg_attr(docsrs, doc(cfg(feature = "net")))] -#[derive(Clone, Copy, PartialEq, PartialOrd)] +#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] pub struct Ready(usize); impl Ready { @@ -31,13 +33,34 @@ impl Ready { /// Returns a `Ready` representing write closed readiness. pub const WRITE_CLOSED: Ready = Ready(WRITE_CLOSED); + /// Returns a `Ready` representing priority readiness. + #[cfg(any(target_os = "linux", target_os = "android"))] + #[cfg_attr(docsrs, doc(cfg(any(target_os = "linux", target_os = "android"))))] + pub const PRIORITY: Ready = Ready(PRIORITY); + + /// Returns a `Ready` representing readiness for all operations. + #[cfg(any(target_os = "linux", target_os = "android"))] + pub const ALL: Ready = Ready(READABLE | WRITABLE | READ_CLOSED | WRITE_CLOSED | PRIORITY); + /// Returns a `Ready` representing readiness for all operations. + #[cfg(not(any(target_os = "linux", target_os = "android")))] pub const ALL: Ready = Ready(READABLE | WRITABLE | READ_CLOSED | WRITE_CLOSED); // Must remain crate-private to avoid adding a public dependency on Mio. pub(crate) fn from_mio(event: &mio::event::Event) -> Ready { let mut ready = Ready::EMPTY; + #[cfg(all(target_os = "freebsd", feature = "net"))] + { + if event.is_aio() { + ready |= Ready::READABLE; + } + + if event.is_lio() { + ready |= Ready::READABLE; + } + } + if event.is_readable() { ready |= Ready::READABLE; } @@ -54,10 +77,17 @@ impl Ready { ready |= Ready::WRITE_CLOSED; } + #[cfg(any(target_os = "linux", target_os = "android"))] + { + if event.is_priority() { + ready |= Ready::PRIORITY; + } + } + ready } - /// Returns true if `Ready` is the empty set + /// Returns true if `Ready` is the empty set. /// /// # Examples /// @@ -71,7 +101,7 @@ impl Ready { self == Ready::EMPTY } - /// Returns `true` if the value includes `readable` + /// Returns `true` if the value includes `readable`. /// /// # Examples /// @@ -87,7 +117,7 @@ impl Ready { self.contains(Ready::READABLE) || self.is_read_closed() } - /// Returns `true` if the value includes writable `readiness` + /// Returns `true` if the value includes writable `readiness`. /// /// # Examples /// @@ -103,7 +133,7 @@ impl Ready { self.contains(Ready::WRITABLE) || self.is_write_closed() } - /// Returns `true` if the value includes read-closed `readiness` + /// Returns `true` if the value includes read-closed `readiness`. /// /// # Examples /// @@ -118,7 +148,7 @@ impl Ready { self.contains(Ready::READ_CLOSED) } - /// Returns `true` if the value includes write-closed `readiness` + /// Returns `true` if the value includes write-closed `readiness`. /// /// # Examples /// @@ -133,6 +163,23 @@ impl Ready { self.contains(Ready::WRITE_CLOSED) } + /// Returns `true` if the value includes priority `readiness`. + /// + /// # Examples + /// + /// ``` + /// use tokio::io::Ready; + /// + /// assert!(!Ready::EMPTY.is_priority()); + /// assert!(!Ready::WRITABLE.is_priority()); + /// assert!(Ready::PRIORITY.is_priority()); + /// ``` + #[cfg(any(target_os = "linux", target_os = "android"))] + #[cfg_attr(docsrs, doc(cfg(any(target_os = "linux", target_os = "android"))))] + pub fn is_priority(self) -> bool { + self.contains(Ready::PRIORITY) + } + /// Returns true if `self` is a superset of `other`. /// /// `other` may represent more than one readiness operations, in which case @@ -143,7 +190,7 @@ impl Ready { (self & other) == other } - /// Create a `Ready` instance using the given `usize` representation. + /// Creates a `Ready` instance using the given `usize` representation. /// /// The `usize` representation must have been obtained from a call to /// `Readiness::as_usize`. @@ -180,6 +227,12 @@ cfg_io_readiness! { ready |= Ready::WRITE_CLOSED; } + #[cfg(any(target_os = "linux", target_os = "android"))] + if interest.is_priority() { + ready |= Ready::PRIORITY; + ready |= Ready::READ_CLOSED; + } + ready } @@ -229,11 +282,16 @@ impl ops::Sub<Ready> for Ready { impl fmt::Debug for Ready { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - fmt.debug_struct("Ready") - .field("is_readable", &self.is_readable()) + let mut fmt = fmt.debug_struct("Ready"); + + fmt.field("is_readable", &self.is_readable()) .field("is_writable", &self.is_writable()) .field("is_read_closed", &self.is_read_closed()) - .field("is_write_closed", &self.is_write_closed()) - .finish() + .field("is_write_closed", &self.is_write_closed()); + + #[cfg(any(target_os = "linux", target_os = "android"))] + fmt.field("is_priority", &self.is_priority()); + + fmt.finish() } } diff --git a/vendor/tokio/src/io/split.rs b/vendor/tokio/src/io/split.rs index 732eb3b3a..f067b65a8 100644 --- a/vendor/tokio/src/io/split.rs +++ b/vendor/tokio/src/io/split.rs @@ -63,7 +63,7 @@ impl<T> ReadHalf<T> { /// Checks if this `ReadHalf` and some `WriteHalf` were split from the same /// stream. pub fn is_pair_of(&self, other: &WriteHalf<T>) -> bool { - other.is_pair_of(&self) + other.is_pair_of(self) } /// Reunites with a previously split `WriteHalf`. @@ -74,7 +74,11 @@ impl<T> ReadHalf<T> { /// same `split` operation this method will panic. /// This can be checked ahead of time by comparing the stream ID /// of the two halves. - pub fn unsplit(self, wr: WriteHalf<T>) -> T { + #[track_caller] + pub fn unsplit(self, wr: WriteHalf<T>) -> T + where + T: Unpin, + { if self.is_pair_of(&wr) { drop(wr); @@ -90,7 +94,7 @@ impl<T> ReadHalf<T> { } impl<T> WriteHalf<T> { - /// Check if this `WriteHalf` and some `ReadHalf` were split from the same + /// Checks if this `WriteHalf` and some `ReadHalf` were split from the same /// stream. pub fn is_pair_of(&self, other: &ReadHalf<T>) -> bool { Arc::ptr_eq(&self.inner, &other.inner) diff --git a/vendor/tokio/src/io/stderr.rs b/vendor/tokio/src/io/stderr.rs index 2f624fba9..2119c6d6c 100644 --- a/vendor/tokio/src/io/stderr.rs +++ b/vendor/tokio/src/io/stderr.rs @@ -74,16 +74,43 @@ cfg_io_std! { } #[cfg(unix)] -impl std::os::unix::io::AsRawFd for Stderr { - fn as_raw_fd(&self) -> std::os::unix::io::RawFd { - std::io::stderr().as_raw_fd() +mod sys { + #[cfg(not(tokio_no_as_fd))] + use std::os::unix::io::{AsFd, BorrowedFd}; + use std::os::unix::io::{AsRawFd, RawFd}; + + use super::Stderr; + + impl AsRawFd for Stderr { + fn as_raw_fd(&self) -> RawFd { + std::io::stderr().as_raw_fd() + } + } + + #[cfg(not(tokio_no_as_fd))] + impl AsFd for Stderr { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } } } -#[cfg(windows)] -impl std::os::windows::io::AsRawHandle for Stderr { - fn as_raw_handle(&self) -> std::os::windows::io::RawHandle { - std::io::stderr().as_raw_handle() +cfg_windows! { + #[cfg(not(tokio_no_as_fd))] + use crate::os::windows::io::{AsHandle, BorrowedHandle}; + use crate::os::windows::io::{AsRawHandle, RawHandle}; + + impl AsRawHandle for Stderr { + fn as_raw_handle(&self) -> RawHandle { + std::io::stderr().as_raw_handle() + } + } + + #[cfg(not(tokio_no_as_fd))] + impl AsHandle for Stderr { + fn as_handle(&self) -> BorrowedHandle<'_> { + unsafe { BorrowedHandle::borrow_raw(self.as_raw_handle()) } + } } } diff --git a/vendor/tokio/src/io/stdin.rs b/vendor/tokio/src/io/stdin.rs index c9578f17b..04cced425 100644 --- a/vendor/tokio/src/io/stdin.rs +++ b/vendor/tokio/src/io/stdin.rs @@ -49,16 +49,43 @@ cfg_io_std! { } #[cfg(unix)] -impl std::os::unix::io::AsRawFd for Stdin { - fn as_raw_fd(&self) -> std::os::unix::io::RawFd { - std::io::stdin().as_raw_fd() +mod sys { + #[cfg(not(tokio_no_as_fd))] + use std::os::unix::io::{AsFd, BorrowedFd}; + use std::os::unix::io::{AsRawFd, RawFd}; + + use super::Stdin; + + impl AsRawFd for Stdin { + fn as_raw_fd(&self) -> RawFd { + std::io::stdin().as_raw_fd() + } + } + + #[cfg(not(tokio_no_as_fd))] + impl AsFd for Stdin { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } } } -#[cfg(windows)] -impl std::os::windows::io::AsRawHandle for Stdin { - fn as_raw_handle(&self) -> std::os::windows::io::RawHandle { - std::io::stdin().as_raw_handle() +cfg_windows! { + #[cfg(not(tokio_no_as_fd))] + use crate::os::windows::io::{AsHandle, BorrowedHandle}; + use crate::os::windows::io::{AsRawHandle, RawHandle}; + + impl AsRawHandle for Stdin { + fn as_raw_handle(&self) -> RawHandle { + std::io::stdin().as_raw_handle() + } + } + + #[cfg(not(tokio_no_as_fd))] + impl AsHandle for Stdin { + fn as_handle(&self) -> BorrowedHandle<'_> { + unsafe { BorrowedHandle::borrow_raw(self.as_raw_handle()) } + } } } diff --git a/vendor/tokio/src/io/stdio_common.rs b/vendor/tokio/src/io/stdio_common.rs index 56c4520c6..06da761b8 100644 --- a/vendor/tokio/src/io/stdio_common.rs +++ b/vendor/tokio/src/io/stdio_common.rs @@ -7,7 +7,7 @@ use std::task::{Context, Poll}; /// if buffer contents seems to be utf8. Otherwise it only trims buffer down to MAX_BUF. /// That's why, wrapped writer will always receive well-formed utf-8 bytes. /// # Other platforms -/// passes data to `inner` as is +/// Passes data to `inner` as is. #[derive(Debug)] pub(crate) struct SplitByUtf8BoundaryIfWindows<W> { inner: W, @@ -42,7 +42,7 @@ where // for further code. Since `AsyncWrite` can always shrink // buffer at its discretion, excessive (i.e. in tests) shrinking // does not break correctness. - // 2. If buffer is small, it will not be shrinked. + // 2. If buffer is small, it will not be shrunk. // That's why, it's "textness" will not change, so we don't have // to fixup it. if cfg!(not(any(target_os = "windows", test))) || buf.len() <= crate::io::blocking::MAX_BUF @@ -108,14 +108,13 @@ where #[cfg(test)] #[cfg(not(loom))] mod tests { + use crate::io::blocking::MAX_BUF; use crate::io::AsyncWriteExt; use std::io; use std::pin::Pin; use std::task::Context; use std::task::Poll; - const MAX_BUF: usize = 16 * 1024; - struct TextMockWriter; impl crate::io::AsyncWrite for TextMockWriter { @@ -177,6 +176,7 @@ mod tests { } #[test] + #[cfg_attr(miri, ignore)] fn test_splitter() { let data = str::repeat("â–ˆ", MAX_BUF); let mut wr = super::SplitByUtf8BoundaryIfWindows::new(TextMockWriter); @@ -190,10 +190,11 @@ mod tests { } #[test] + #[cfg_attr(miri, ignore)] fn test_pseudo_text() { // In this test we write a piece of binary data, whose beginning is // text though. We then validate that even in this corner case buffer - // was not shrinked too much. + // was not shrunk too much. let checked_count = super::MAGIC_CONST * super::MAX_BYTES_PER_CHAR; let mut data: Vec<u8> = str::repeat("a", checked_count).into(); data.extend(std::iter::repeat(0b1010_1010).take(MAX_BUF - checked_count + 1)); @@ -212,7 +213,7 @@ mod tests { writer.write_history.iter().copied().sum::<usize>(), data.len() ); - // Check that at most MAX_BYTES_PER_CHAR + 1 (i.e. 5) bytes were shrinked + // Check that at most MAX_BYTES_PER_CHAR + 1 (i.e. 5) bytes were shrunk // from the buffer: one because it was outside of MAX_BUF boundary, and // up to one "utf8 code point". assert!(data.len() - writer.write_history[0] <= super::MAX_BYTES_PER_CHAR + 1); diff --git a/vendor/tokio/src/io/stdout.rs b/vendor/tokio/src/io/stdout.rs index a08ed01ee..f52bf5a4b 100644 --- a/vendor/tokio/src/io/stdout.rs +++ b/vendor/tokio/src/io/stdout.rs @@ -73,16 +73,43 @@ cfg_io_std! { } #[cfg(unix)] -impl std::os::unix::io::AsRawFd for Stdout { - fn as_raw_fd(&self) -> std::os::unix::io::RawFd { - std::io::stdout().as_raw_fd() +mod sys { + #[cfg(not(tokio_no_as_fd))] + use std::os::unix::io::{AsFd, BorrowedFd}; + use std::os::unix::io::{AsRawFd, RawFd}; + + use super::Stdout; + + impl AsRawFd for Stdout { + fn as_raw_fd(&self) -> RawFd { + std::io::stdout().as_raw_fd() + } + } + + #[cfg(not(tokio_no_as_fd))] + impl AsFd for Stdout { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } } } -#[cfg(windows)] -impl std::os::windows::io::AsRawHandle for Stdout { - fn as_raw_handle(&self) -> std::os::windows::io::RawHandle { - std::io::stdout().as_raw_handle() +cfg_windows! { + #[cfg(not(tokio_no_as_fd))] + use crate::os::windows::io::{AsHandle, BorrowedHandle}; + use crate::os::windows::io::{AsRawHandle, RawHandle}; + + impl AsRawHandle for Stdout { + fn as_raw_handle(&self) -> RawHandle { + std::io::stdout().as_raw_handle() + } + } + + #[cfg(not(tokio_no_as_fd))] + impl AsHandle for Stdout { + fn as_handle(&self) -> BorrowedHandle<'_> { + unsafe { BorrowedHandle::borrow_raw(self.as_raw_handle()) } + } } } diff --git a/vendor/tokio/src/io/util/async_buf_read_ext.rs b/vendor/tokio/src/io/util/async_buf_read_ext.rs index 233ac31c4..50ea2d91a 100644 --- a/vendor/tokio/src/io/util/async_buf_read_ext.rs +++ b/vendor/tokio/src/io/util/async_buf_read_ext.rs @@ -1,3 +1,4 @@ +use crate::io::util::fill_buf::{fill_buf, FillBuf}; use crate::io::util::lines::{lines, Lines}; use crate::io::util::read_line::{read_line, ReadLine}; use crate::io::util::read_until::{read_until, ReadUntil}; @@ -36,6 +37,18 @@ cfg_io_util! { /// [`fill_buf`]: AsyncBufRead::poll_fill_buf /// [`ErrorKind::Interrupted`]: std::io::ErrorKind::Interrupted /// + /// # Cancel safety + /// + /// If the method is used as the event in a + /// [`tokio::select!`](crate::select) statement and some other branch + /// completes first, then some data may have been partially read. Any + /// partially read bytes are appended to `buf`, and the method can be + /// called again to continue reading until `byte`. + /// + /// This method returns the total number of bytes read. If you cancel + /// the call to `read_until` and then call it again to continue reading, + /// the counter is reset. + /// /// # Examples /// /// [`std::io::Cursor`][`Cursor`] is a type that implements `BufRead`. In @@ -114,6 +127,30 @@ cfg_io_util! { /// /// [`read_until`]: AsyncBufReadExt::read_until /// + /// # Cancel safety + /// + /// This method is not cancellation safe. If the method is used as the + /// event in a [`tokio::select!`](crate::select) statement and some + /// other branch completes first, then some data may have been partially + /// read, and this data is lost. There are no guarantees regarding the + /// contents of `buf` when the call is cancelled. The current + /// implementation replaces `buf` with the empty string, but this may + /// change in the future. + /// + /// This function does not behave like [`read_until`] because of the + /// requirement that a string contains only valid utf-8. If you need a + /// cancellation safe `read_line`, there are three options: + /// + /// * Call [`read_until`] with a newline character and manually perform the utf-8 check. + /// * The stream returned by [`lines`] has a cancellation safe + /// [`next_line`] method. + /// * Use [`tokio_util::codec::LinesCodec`][LinesCodec]. + /// + /// [LinesCodec]: https://docs.rs/tokio-util/latest/tokio_util/codec/struct.LinesCodec.html + /// [`read_until`]: Self::read_until + /// [`lines`]: Self::lines + /// [`next_line`]: crate::io::Lines::next_line + /// /// # Examples /// /// [`std::io::Cursor`][`Cursor`] is a type that implements @@ -173,10 +210,11 @@ cfg_io_util! { /// [`BufRead::split`](std::io::BufRead::split). /// /// The stream returned from this function will yield instances of - /// [`io::Result`]`<`[`Vec<u8>`]`>`. Each vector returned will *not* have + /// [`io::Result`]`<`[`Option`]`<`[`Vec<u8>`]`>>`. Each vector returned will *not* have /// the delimiter byte at the end. /// /// [`io::Result`]: std::io::Result + /// [`Option`]: core::option::Option /// [`Vec<u8>`]: std::vec::Vec /// /// # Errors @@ -206,14 +244,68 @@ cfg_io_util! { split(self, byte) } + /// Returns the contents of the internal buffer, filling it with more + /// data from the inner reader if it is empty. + /// + /// This function is a lower-level call. It needs to be paired with the + /// [`consume`] method to function properly. When calling this method, + /// none of the contents will be "read" in the sense that later calling + /// `read` may return the same contents. As such, [`consume`] must be + /// called with the number of bytes that are consumed from this buffer + /// to ensure that the bytes are never returned twice. + /// + /// An empty buffer returned indicates that the stream has reached EOF. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn fill_buf(&mut self) -> io::Result<&[u8]>; + /// ``` + /// + /// # Errors + /// + /// This function will return an I/O error if the underlying reader was + /// read, but returned an error. + /// + /// [`consume`]: crate::io::AsyncBufReadExt::consume + fn fill_buf(&mut self) -> FillBuf<'_, Self> + where + Self: Unpin, + { + fill_buf(self) + } + + /// Tells this buffer that `amt` bytes have been consumed from the + /// buffer, so they should no longer be returned in calls to [`read`]. + /// + /// This function is a lower-level call. It needs to be paired with the + /// [`fill_buf`] method to function properly. This function does not + /// perform any I/O, it simply informs this object that some amount of + /// its buffer, returned from [`fill_buf`], has been consumed and should + /// no longer be returned. As such, this function may do odd things if + /// [`fill_buf`] isn't called before calling it. + /// + /// The `amt` must be less than the number of bytes in the buffer + /// returned by [`fill_buf`]. + /// + /// [`read`]: crate::io::AsyncReadExt::read + /// [`fill_buf`]: crate::io::AsyncBufReadExt::fill_buf + fn consume(&mut self, amt: usize) + where + Self: Unpin, + { + std::pin::Pin::new(self).consume(amt) + } + /// Returns a stream over the lines of this reader. /// This method is the async equivalent to [`BufRead::lines`](std::io::BufRead::lines). /// /// The stream returned from this function will yield instances of - /// [`io::Result`]`<`[`String`]`>`. Each string returned will *not* have a newline + /// [`io::Result`]`<`[`Option`]`<`[`String`]`>>`. Each string returned will *not* have a newline /// byte (the 0xA byte) or CRLF (0xD, 0xA bytes) at the end. /// /// [`io::Result`]: std::io::Result + /// [`Option`]: core::option::Option /// [`String`]: String /// /// # Errors diff --git a/vendor/tokio/src/io/util/async_read_ext.rs b/vendor/tokio/src/io/util/async_read_ext.rs index 10f641744..df5445c2c 100644 --- a/vendor/tokio/src/io/util/async_read_ext.rs +++ b/vendor/tokio/src/io/util/async_read_ext.rs @@ -2,6 +2,7 @@ use crate::io::util::chain::{chain, Chain}; use crate::io::util::read::{read, Read}; use crate::io::util::read_buf::{read_buf, ReadBuf}; use crate::io::util::read_exact::{read_exact, ReadExact}; +use crate::io::util::read_int::{ReadF32, ReadF32Le, ReadF64, ReadF64Le}; use crate::io::util::read_int::{ ReadI128, ReadI128Le, ReadI16, ReadI16Le, ReadI32, ReadI32Le, ReadI64, ReadI64Le, ReadI8, }; @@ -691,6 +692,82 @@ cfg_io_util! { /// ``` fn read_i128(&mut self) -> ReadI128; + /// Reads an 32-bit floating point type in big-endian order from the + /// underlying reader. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_f32(&mut self) -> io::Result<f32>; + /// ``` + /// + /// It is recommended to use a buffered reader to avoid excessive + /// syscalls. + /// + /// # Errors + /// + /// This method returns the same errors as [`AsyncReadExt::read_exact`]. + /// + /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact + /// + /// # Examples + /// + /// Read 32-bit floating point type from a `AsyncRead`: + /// + /// ```rust + /// use tokio::io::{self, AsyncReadExt}; + /// + /// use std::io::Cursor; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut reader = Cursor::new(vec![0xff, 0x7f, 0xff, 0xff]); + /// + /// assert_eq!(f32::MIN, reader.read_f32().await?); + /// Ok(()) + /// } + /// ``` + fn read_f32(&mut self) -> ReadF32; + + /// Reads an 64-bit floating point type in big-endian order from the + /// underlying reader. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_f64(&mut self) -> io::Result<f64>; + /// ``` + /// + /// It is recommended to use a buffered reader to avoid excessive + /// syscalls. + /// + /// # Errors + /// + /// This method returns the same errors as [`AsyncReadExt::read_exact`]. + /// + /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact + /// + /// # Examples + /// + /// Read 64-bit floating point type from a `AsyncRead`: + /// + /// ```rust + /// use tokio::io::{self, AsyncReadExt}; + /// + /// use std::io::Cursor; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut reader = Cursor::new(vec![ + /// 0xff, 0xef, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff + /// ]); + /// + /// assert_eq!(f64::MIN, reader.read_f64().await?); + /// Ok(()) + /// } + /// ``` + fn read_f64(&mut self) -> ReadF64; + /// Reads an unsigned 16-bit integer in little-endian order from the /// underlying reader. /// @@ -997,6 +1074,82 @@ cfg_io_util! { /// } /// ``` fn read_i128_le(&mut self) -> ReadI128Le; + + /// Reads an 32-bit floating point type in little-endian order from the + /// underlying reader. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_f32_le(&mut self) -> io::Result<f32>; + /// ``` + /// + /// It is recommended to use a buffered reader to avoid excessive + /// syscalls. + /// + /// # Errors + /// + /// This method returns the same errors as [`AsyncReadExt::read_exact`]. + /// + /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact + /// + /// # Examples + /// + /// Read 32-bit floating point type from a `AsyncRead`: + /// + /// ```rust + /// use tokio::io::{self, AsyncReadExt}; + /// + /// use std::io::Cursor; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut reader = Cursor::new(vec![0xff, 0xff, 0x7f, 0xff]); + /// + /// assert_eq!(f32::MIN, reader.read_f32_le().await?); + /// Ok(()) + /// } + /// ``` + fn read_f32_le(&mut self) -> ReadF32Le; + + /// Reads an 64-bit floating point type in little-endian order from the + /// underlying reader. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_f64_le(&mut self) -> io::Result<f64>; + /// ``` + /// + /// It is recommended to use a buffered reader to avoid excessive + /// syscalls. + /// + /// # Errors + /// + /// This method returns the same errors as [`AsyncReadExt::read_exact`]. + /// + /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact + /// + /// # Examples + /// + /// Read 64-bit floating point type from a `AsyncRead`: + /// + /// ```rust + /// use tokio::io::{self, AsyncReadExt}; + /// + /// use std::io::Cursor; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut reader = Cursor::new(vec![ + /// 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xef, 0xff + /// ]); + /// + /// assert_eq!(f64::MIN, reader.read_f64_le().await?); + /// Ok(()) + /// } + /// ``` + fn read_f64_le(&mut self) -> ReadF64Le; } /// Reads all bytes until EOF in this source, placing them into `buf`. diff --git a/vendor/tokio/src/io/util/async_seek_ext.rs b/vendor/tokio/src/io/util/async_seek_ext.rs index 297a4a61a..aadf3a76a 100644 --- a/vendor/tokio/src/io/util/async_seek_ext.rs +++ b/vendor/tokio/src/io/util/async_seek_ext.rs @@ -67,6 +67,16 @@ cfg_io_util! { seek(self, pos) } + /// Creates a future which will rewind to the beginning of the stream. + /// + /// This is convenience method, equivalent to `self.seek(SeekFrom::Start(0))`. + fn rewind(&mut self) -> Seek<'_, Self> + where + Self: Unpin, + { + self.seek(SeekFrom::Start(0)) + } + /// Creates a future which will return the current seek position from the /// start of the stream. /// diff --git a/vendor/tokio/src/io/util/async_write_ext.rs b/vendor/tokio/src/io/util/async_write_ext.rs index e214321b7..dfdde82f3 100644 --- a/vendor/tokio/src/io/util/async_write_ext.rs +++ b/vendor/tokio/src/io/util/async_write_ext.rs @@ -4,6 +4,7 @@ use crate::io::util::write::{write, Write}; use crate::io::util::write_all::{write_all, WriteAll}; use crate::io::util::write_all_buf::{write_all_buf, WriteAllBuf}; use crate::io::util::write_buf::{write_buf, WriteBuf}; +use crate::io::util::write_int::{WriteF32, WriteF32Le, WriteF64, WriteF64Le}; use crate::io::util::write_int::{ WriteI128, WriteI128Le, WriteI16, WriteI16Le, WriteI32, WriteI32Le, WriteI64, WriteI64Le, WriteI8, @@ -19,7 +20,7 @@ use std::io::IoSlice; use bytes::Buf; cfg_io_util! { - /// Defines numeric writer + /// Defines numeric writer. macro_rules! write_impl { ( $( @@ -255,7 +256,7 @@ cfg_io_util! { write_buf(self, src) } - /// Attempts to write an entire buffer into this writer + /// Attempts to write an entire buffer into this writer. /// /// Equivalent to: /// @@ -352,9 +353,9 @@ cfg_io_util! { /// /// #[tokio::main] /// async fn main() -> io::Result<()> { - /// let mut buffer = File::create("foo.txt").await?; + /// let mut file = File::create("foo.txt").await?; /// - /// buffer.write_all(b"some bytes").await?; + /// file.write_all(b"some bytes").await?; /// Ok(()) /// } /// ``` @@ -405,7 +406,7 @@ cfg_io_util! { /// ``` fn write_u8(&mut self, n: u8) -> WriteU8; - /// Writes an unsigned 8-bit integer to the underlying writer. + /// Writes a signed 8-bit integer to the underlying writer. /// /// Equivalent to: /// @@ -424,7 +425,7 @@ cfg_io_util! { /// /// # Examples /// - /// Write unsigned 8 bit integers to a `AsyncWrite`: + /// Write signed 8 bit integers to a `AsyncWrite`: /// /// ```rust /// use tokio::io::{self, AsyncWriteExt}; @@ -433,10 +434,10 @@ cfg_io_util! { /// async fn main() -> io::Result<()> { /// let mut writer = Vec::new(); /// - /// writer.write_u8(2).await?; - /// writer.write_u8(5).await?; + /// writer.write_i8(-2).await?; + /// writer.write_i8(126).await?; /// - /// assert_eq!(writer, b"\x02\x05"); + /// assert_eq!(writer, b"\xFE\x7E"); /// Ok(()) /// } /// ``` @@ -750,6 +751,81 @@ cfg_io_util! { /// ``` fn write_i128(&mut self, n: i128) -> WriteI128; + /// Writes an 32-bit floating point type in big-endian order to the + /// underlying writer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_f32(&mut self, n: f32) -> io::Result<()>; + /// ``` + /// + /// It is recommended to use a buffered writer to avoid excessive + /// syscalls. + /// + /// # Errors + /// + /// This method returns the same errors as [`AsyncWriteExt::write_all`]. + /// + /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all + /// + /// # Examples + /// + /// Write 32-bit floating point type to a `AsyncWrite`: + /// + /// ```rust + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut writer = Vec::new(); + /// + /// writer.write_f32(f32::MIN).await?; + /// + /// assert_eq!(writer, vec![0xff, 0x7f, 0xff, 0xff]); + /// Ok(()) + /// } + /// ``` + fn write_f32(&mut self, n: f32) -> WriteF32; + + /// Writes an 64-bit floating point type in big-endian order to the + /// underlying writer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_f64(&mut self, n: f64) -> io::Result<()>; + /// ``` + /// + /// It is recommended to use a buffered writer to avoid excessive + /// syscalls. + /// + /// # Errors + /// + /// This method returns the same errors as [`AsyncWriteExt::write_all`]. + /// + /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all + /// + /// # Examples + /// + /// Write 64-bit floating point type to a `AsyncWrite`: + /// + /// ```rust + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut writer = Vec::new(); + /// + /// writer.write_f64(f64::MIN).await?; + /// + /// assert_eq!(writer, vec![ + /// 0xff, 0xef, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff + /// ]); + /// Ok(()) + /// } + /// ``` + fn write_f64(&mut self, n: f64) -> WriteF64; /// Writes an unsigned 16-bit integer in little-endian order to the /// underlying writer. @@ -1058,6 +1134,82 @@ cfg_io_util! { /// } /// ``` fn write_i128_le(&mut self, n: i128) -> WriteI128Le; + + /// Writes an 32-bit floating point type in little-endian order to the + /// underlying writer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_f32_le(&mut self, n: f32) -> io::Result<()>; + /// ``` + /// + /// It is recommended to use a buffered writer to avoid excessive + /// syscalls. + /// + /// # Errors + /// + /// This method returns the same errors as [`AsyncWriteExt::write_all`]. + /// + /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all + /// + /// # Examples + /// + /// Write 32-bit floating point type to a `AsyncWrite`: + /// + /// ```rust + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut writer = Vec::new(); + /// + /// writer.write_f32_le(f32::MIN).await?; + /// + /// assert_eq!(writer, vec![0xff, 0xff, 0x7f, 0xff]); + /// Ok(()) + /// } + /// ``` + fn write_f32_le(&mut self, n: f32) -> WriteF32Le; + + /// Writes an 64-bit floating point type in little-endian order to the + /// underlying writer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_f64_le(&mut self, n: f64) -> io::Result<()>; + /// ``` + /// + /// It is recommended to use a buffered writer to avoid excessive + /// syscalls. + /// + /// # Errors + /// + /// This method returns the same errors as [`AsyncWriteExt::write_all`]. + /// + /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all + /// + /// # Examples + /// + /// Write 64-bit floating point type to a `AsyncWrite`: + /// + /// ```rust + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut writer = Vec::new(); + /// + /// writer.write_f64_le(f64::MIN).await?; + /// + /// assert_eq!(writer, vec![ + /// 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xef, 0xff + /// ]); + /// Ok(()) + /// } + /// ``` + fn write_f64_le(&mut self, n: f64) -> WriteF64Le; } /// Flushes this output stream, ensuring that all intermediately buffered diff --git a/vendor/tokio/src/io/util/buf_reader.rs b/vendor/tokio/src/io/util/buf_reader.rs index 7cfd46ce0..60879c0fd 100644 --- a/vendor/tokio/src/io/util/buf_reader.rs +++ b/vendor/tokio/src/io/util/buf_reader.rs @@ -155,7 +155,7 @@ pub(super) enum SeekState { Pending, } -/// Seek to an offset, in bytes, in the underlying reader. +/// Seeks to an offset, in bytes, in the underlying reader. /// /// The position used for seeking with `SeekFrom::Current(_)` is the /// position the underlying reader would be at if the `BufReader` had no @@ -204,7 +204,6 @@ impl<R: AsyncRead + AsyncSeek> AsyncSeek for BufReader<R> { self.as_mut() .get_pin_mut() .start_seek(SeekFrom::Current(offset))?; - self.as_mut().get_pin_mut().poll_complete(cx)? } else { // seek backwards by our remainder, and then by the offset self.as_mut() @@ -221,8 +220,8 @@ impl<R: AsyncRead + AsyncSeek> AsyncSeek for BufReader<R> { self.as_mut() .get_pin_mut() .start_seek(SeekFrom::Current(n))?; - self.as_mut().get_pin_mut().poll_complete(cx)? } + self.as_mut().get_pin_mut().poll_complete(cx)? } SeekState::PendingOverflowed(n) => { if self.as_mut().get_pin_mut().poll_complete(cx)?.is_pending() { diff --git a/vendor/tokio/src/io/util/copy.rs b/vendor/tokio/src/io/util/copy.rs index 3cd425b34..55861244f 100644 --- a/vendor/tokio/src/io/util/copy.rs +++ b/vendor/tokio/src/io/util/copy.rs @@ -8,6 +8,7 @@ use std::task::{Context, Poll}; #[derive(Debug)] pub(super) struct CopyBuffer { read_done: bool, + need_flush: bool, pos: usize, cap: usize, amt: u64, @@ -18,10 +19,56 @@ impl CopyBuffer { pub(super) fn new() -> Self { Self { read_done: false, + need_flush: false, pos: 0, cap: 0, amt: 0, - buf: vec![0; 2048].into_boxed_slice(), + buf: vec![0; super::DEFAULT_BUF_SIZE].into_boxed_slice(), + } + } + + fn poll_fill_buf<R>( + &mut self, + cx: &mut Context<'_>, + reader: Pin<&mut R>, + ) -> Poll<io::Result<()>> + where + R: AsyncRead + ?Sized, + { + let me = &mut *self; + let mut buf = ReadBuf::new(&mut me.buf); + buf.set_filled(me.cap); + + let res = reader.poll_read(cx, &mut buf); + if let Poll::Ready(Ok(_)) = res { + let filled_len = buf.filled().len(); + me.read_done = me.cap == filled_len; + me.cap = filled_len; + } + res + } + + fn poll_write_buf<R, W>( + &mut self, + cx: &mut Context<'_>, + mut reader: Pin<&mut R>, + mut writer: Pin<&mut W>, + ) -> Poll<io::Result<usize>> + where + R: AsyncRead + ?Sized, + W: AsyncWrite + ?Sized, + { + let me = &mut *self; + match writer.as_mut().poll_write(cx, &me.buf[me.pos..me.cap]) { + Poll::Pending => { + // Top up the buffer towards full if we can read a bit more + // data - this should improve the chances of a large write + if !me.read_done && me.cap < me.buf.len() { + ready!(me.poll_fill_buf(cx, reader.as_mut()))?; + } + Poll::Pending + } + res => res, } } @@ -39,22 +86,28 @@ impl CopyBuffer { // If our buffer is empty, then we need to read some data to // continue. if self.pos == self.cap && !self.read_done { - let me = &mut *self; - let mut buf = ReadBuf::new(&mut me.buf); - ready!(reader.as_mut().poll_read(cx, &mut buf))?; - let n = buf.filled().len(); - if n == 0 { - self.read_done = true; - } else { - self.pos = 0; - self.cap = n; + self.pos = 0; + self.cap = 0; + + match self.poll_fill_buf(cx, reader.as_mut()) { + Poll::Ready(Ok(_)) => (), + Poll::Ready(Err(err)) => return Poll::Ready(Err(err)), + Poll::Pending => { + // Try flushing when the reader has no progress to avoid deadlock + // when the reader depends on buffered writer. + if self.need_flush { + ready!(writer.as_mut().poll_flush(cx))?; + self.need_flush = false; + } + + return Poll::Pending; + } } } // If our buffer has some data, let's write it out! while self.pos < self.cap { - let me = &mut *self; - let i = ready!(writer.as_mut().poll_write(cx, &me.buf[me.pos..me.cap]))?; + let i = ready!(self.poll_write_buf(cx, reader.as_mut(), writer.as_mut()))?; if i == 0 { return Poll::Ready(Err(io::Error::new( io::ErrorKind::WriteZero, @@ -63,9 +116,18 @@ impl CopyBuffer { } else { self.pos += i; self.amt += i as u64; + self.need_flush = true; } } + // If pos larger than cap, this loop will never stop. + // In particular, user's wrong poll_write implementation returning + // incorrect written length may lead to thread blocking. + debug_assert!( + self.pos <= self.cap, + "writer returned length larger than input slice" + ); + // If we've written all the data and we've seen EOF, flush out the // data and finish the transfer. if self.pos == self.cap && self.read_done { @@ -91,14 +153,22 @@ cfg_io_util! { /// /// This function returns a future that will continuously read data from /// `reader` and then write it into `writer` in a streaming fashion until - /// `reader` returns EOF. + /// `reader` returns EOF or fails. /// /// On success, the total number of bytes that were copied from `reader` to /// `writer` is returned. /// /// This is an asynchronous version of [`std::io::copy`][std]. /// + /// A heap-allocated copy buffer with 8 KB is created to take data from the + /// reader to the writer, check [`copy_buf`] if you want an alternative for + /// [`AsyncBufRead`]. You can use `copy_buf` with [`BufReader`] to change the + /// buffer capacity. + /// /// [std]: std::io::copy + /// [`copy_buf`]: crate::io::copy_buf + /// [`AsyncBufRead`]: crate::io::AsyncBufRead + /// [`BufReader`]: crate::io::BufReader /// /// # Errors /// diff --git a/vendor/tokio/src/io/util/copy_bidirectional.rs b/vendor/tokio/src/io/util/copy_bidirectional.rs index c93060b36..e1a7db127 100644 --- a/vendor/tokio/src/io/util/copy_bidirectional.rs +++ b/vendor/tokio/src/io/util/copy_bidirectional.rs @@ -1,8 +1,8 @@ use super::copy::CopyBuffer; +use crate::future::poll_fn; use crate::io::{AsyncRead, AsyncWrite}; -use std::future::Future; use std::io; use std::pin::Pin; use std::task::{Context, Poll}; @@ -13,13 +13,6 @@ enum TransferState { Done(u64), } -struct CopyBidirectional<'a, A: ?Sized, B: ?Sized> { - a: &'a mut A, - b: &'a mut B, - a_to_b: TransferState, - b_to_a: TransferState, -} - fn transfer_one_direction<A, B>( cx: &mut Context<'_>, state: &mut TransferState, @@ -48,35 +41,6 @@ where } } } - -impl<'a, A, B> Future for CopyBidirectional<'a, A, B> -where - A: AsyncRead + AsyncWrite + Unpin + ?Sized, - B: AsyncRead + AsyncWrite + Unpin + ?Sized, -{ - type Output = io::Result<(u64, u64)>; - - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { - // Unpack self into mut refs to each field to avoid borrow check issues. - let CopyBidirectional { - a, - b, - a_to_b, - b_to_a, - } = &mut *self; - - let a_to_b = transfer_one_direction(cx, a_to_b, &mut *a, &mut *b)?; - let b_to_a = transfer_one_direction(cx, b_to_a, &mut *b, &mut *a)?; - - // It is not a problem if ready! returns early because transfer_one_direction for the - // other direction will keep returning TransferState::Done(count) in future calls to poll - let a_to_b = ready!(a_to_b); - let b_to_a = ready!(b_to_a); - - Poll::Ready(Ok((a_to_b, b_to_a))) - } -} - /// Copies data in both directions between `a` and `b`. /// /// This function returns a future that will read from both streams, @@ -110,11 +74,18 @@ where A: AsyncRead + AsyncWrite + Unpin + ?Sized, B: AsyncRead + AsyncWrite + Unpin + ?Sized, { - CopyBidirectional { - a, - b, - a_to_b: TransferState::Running(CopyBuffer::new()), - b_to_a: TransferState::Running(CopyBuffer::new()), - } + let mut a_to_b = TransferState::Running(CopyBuffer::new()); + let mut b_to_a = TransferState::Running(CopyBuffer::new()); + poll_fn(|cx| { + let a_to_b = transfer_one_direction(cx, &mut a_to_b, a, b)?; + let b_to_a = transfer_one_direction(cx, &mut b_to_a, b, a)?; + + // It is not a problem if ready! returns early because transfer_one_direction for the + // other direction will keep returning TransferState::Done(count) in future calls to poll + let a_to_b = ready!(a_to_b); + let b_to_a = ready!(b_to_a); + + Poll::Ready(Ok((a_to_b, b_to_a))) + }) .await } diff --git a/vendor/tokio/src/io/util/copy_buf.rs b/vendor/tokio/src/io/util/copy_buf.rs index 6831580b4..c23fc9a2b 100644 --- a/vendor/tokio/src/io/util/copy_buf.rs +++ b/vendor/tokio/src/io/util/copy_buf.rs @@ -24,11 +24,17 @@ cfg_io_util! { /// /// This function returns a future that will continuously read data from /// `reader` and then write it into `writer` in a streaming fashion until - /// `reader` returns EOF. + /// `reader` returns EOF or fails. /// /// On success, the total number of bytes that were copied from `reader` to /// `writer` is returned. /// + /// This is a [`tokio::io::copy`] alternative for [`AsyncBufRead`] readers + /// with no extra buffer allocation, since [`AsyncBufRead`] allow access + /// to the reader's inner buffer. + /// + /// [`tokio::io::copy`]: crate::io::copy + /// [`AsyncBufRead`]: crate::io::AsyncBufRead /// /// # Errors /// diff --git a/vendor/tokio/src/io/util/empty.rs b/vendor/tokio/src/io/util/empty.rs index f964d18e6..b96fabbaa 100644 --- a/vendor/tokio/src/io/util/empty.rs +++ b/vendor/tokio/src/io/util/empty.rs @@ -50,16 +50,20 @@ impl AsyncRead for Empty { #[inline] fn poll_read( self: Pin<&mut Self>, - _: &mut Context<'_>, + cx: &mut Context<'_>, _: &mut ReadBuf<'_>, ) -> Poll<io::Result<()>> { + ready!(crate::trace::trace_leaf(cx)); + ready!(poll_proceed_and_make_progress(cx)); Poll::Ready(Ok(())) } } impl AsyncBufRead for Empty { #[inline] - fn poll_fill_buf(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<&[u8]>> { + fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> { + ready!(crate::trace::trace_leaf(cx)); + ready!(poll_proceed_and_make_progress(cx)); Poll::Ready(Ok(&[])) } @@ -73,6 +77,20 @@ impl fmt::Debug for Empty { } } +cfg_coop! { + fn poll_proceed_and_make_progress(cx: &mut Context<'_>) -> Poll<()> { + let coop = ready!(crate::runtime::coop::poll_proceed(cx)); + coop.made_progress(); + Poll::Ready(()) + } +} + +cfg_not_coop! { + fn poll_proceed_and_make_progress(_: &mut Context<'_>) -> Poll<()> { + Poll::Ready(()) + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/vendor/tokio/src/io/util/fill_buf.rs b/vendor/tokio/src/io/util/fill_buf.rs new file mode 100644 index 000000000..bb07c766e --- /dev/null +++ b/vendor/tokio/src/io/util/fill_buf.rs @@ -0,0 +1,59 @@ +use crate::io::AsyncBufRead; + +use pin_project_lite::pin_project; +use std::future::Future; +use std::io; +use std::marker::PhantomPinned; +use std::pin::Pin; +use std::task::{Context, Poll}; + +pin_project! { + /// Future for the [`fill_buf`](crate::io::AsyncBufReadExt::fill_buf) method. + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct FillBuf<'a, R: ?Sized> { + reader: Option<&'a mut R>, + #[pin] + _pin: PhantomPinned, + } +} + +pub(crate) fn fill_buf<R>(reader: &mut R) -> FillBuf<'_, R> +where + R: AsyncBufRead + ?Sized + Unpin, +{ + FillBuf { + reader: Some(reader), + _pin: PhantomPinned, + } +} + +impl<'a, R: AsyncBufRead + ?Sized + Unpin> Future for FillBuf<'a, R> { + type Output = io::Result<&'a [u8]>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let me = self.project(); + + let reader = me.reader.take().expect("Polled after completion."); + match Pin::new(&mut *reader).poll_fill_buf(cx) { + Poll::Ready(Ok(slice)) => unsafe { + // Safety: This is necessary only due to a limitation in the + // borrow checker. Once Rust starts using the polonius borrow + // checker, this can be simplified. + // + // The safety of this transmute relies on the fact that the + // value of `reader` is `None` when we return in this branch. + // Otherwise the caller could poll us again after + // completion, and access the mutable reference while the + // returned immutable reference still exists. + let slice = std::mem::transmute::<&[u8], &'a [u8]>(slice); + Poll::Ready(Ok(slice)) + }, + Poll::Ready(Err(err)) => Poll::Ready(Err(err)), + Poll::Pending => { + *me.reader = Some(reader); + Poll::Pending + } + } + } +} diff --git a/vendor/tokio/src/io/util/lines.rs b/vendor/tokio/src/io/util/lines.rs index d02a4538a..717f633f9 100644 --- a/vendor/tokio/src/io/util/lines.rs +++ b/vendor/tokio/src/io/util/lines.rs @@ -8,7 +8,7 @@ use std::pin::Pin; use std::task::{Context, Poll}; pin_project! { - /// Read lines from an [`AsyncBufRead`]. + /// Reads lines from an [`AsyncBufRead`]. /// /// A `Lines` can be turned into a `Stream` with [`LinesStream`]. /// @@ -47,6 +47,10 @@ where { /// Returns the next line in the stream. /// + /// # Cancel safety + /// + /// This method is cancellation safe. + /// /// # Examples /// /// ``` @@ -68,12 +72,12 @@ where poll_fn(|cx| Pin::new(&mut *self).poll_next_line(cx)).await } - /// Obtain a mutable reference to the underlying reader + /// Obtains a mutable reference to the underlying reader. pub fn get_mut(&mut self) -> &mut R { &mut self.reader } - /// Obtain a reference to the underlying reader + /// Obtains a reference to the underlying reader. pub fn get_ref(&mut self) -> &R { &self.reader } @@ -102,11 +106,9 @@ where /// /// When the method returns `Poll::Pending`, the `Waker` in the provided /// `Context` is scheduled to receive a wakeup when more bytes become - /// available on the underlying IO resource. - /// - /// Note that on multiple calls to `poll_next_line`, only the `Waker` from - /// the `Context` passed to the most recent call is scheduled to receive a - /// wakeup. + /// available on the underlying IO resource. Note that on multiple calls to + /// `poll_next_line`, only the `Waker` from the `Context` passed to the most + /// recent call is scheduled to receive a wakeup. pub fn poll_next_line( self: Pin<&mut Self>, cx: &mut Context<'_>, diff --git a/vendor/tokio/src/io/util/mem.rs b/vendor/tokio/src/io/util/mem.rs index 4eefe7b26..5b404c219 100644 --- a/vendor/tokio/src/io/util/mem.rs +++ b/vendor/tokio/src/io/util/mem.rs @@ -177,10 +177,8 @@ impl Pipe { waker.wake(); } } -} -impl AsyncRead for Pipe { - fn poll_read( + fn poll_read_internal( mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &mut ReadBuf<'_>, @@ -204,10 +202,8 @@ impl AsyncRead for Pipe { Poll::Pending } } -} -impl AsyncWrite for Pipe { - fn poll_write( + fn poll_write_internal( mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &[u8], @@ -228,6 +224,66 @@ impl AsyncWrite for Pipe { } Poll::Ready(Ok(len)) } +} + +impl AsyncRead for Pipe { + cfg_coop! { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut task::Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<std::io::Result<()>> { + ready!(crate::trace::trace_leaf(cx)); + let coop = ready!(crate::runtime::coop::poll_proceed(cx)); + + let ret = self.poll_read_internal(cx, buf); + if ret.is_ready() { + coop.made_progress(); + } + ret + } + } + + cfg_not_coop! { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut task::Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<std::io::Result<()>> { + ready!(crate::trace::trace_leaf(cx)); + self.poll_read_internal(cx, buf) + } + } +} + +impl AsyncWrite for Pipe { + cfg_coop! { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut task::Context<'_>, + buf: &[u8], + ) -> Poll<std::io::Result<usize>> { + ready!(crate::trace::trace_leaf(cx)); + let coop = ready!(crate::runtime::coop::poll_proceed(cx)); + + let ret = self.poll_write_internal(cx, buf); + if ret.is_ready() { + coop.made_progress(); + } + ret + } + } + + cfg_not_coop! { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut task::Context<'_>, + buf: &[u8], + ) -> Poll<std::io::Result<usize>> { + ready!(crate::trace::trace_leaf(cx)); + self.poll_write_internal(cx, buf) + } + } fn poll_flush(self: Pin<&mut Self>, _: &mut task::Context<'_>) -> Poll<std::io::Result<()>> { Poll::Ready(Ok(())) diff --git a/vendor/tokio/src/io/util/mod.rs b/vendor/tokio/src/io/util/mod.rs index fd3dd0dfc..21199d0be 100644 --- a/vendor/tokio/src/io/util/mod.rs +++ b/vendor/tokio/src/io/util/mod.rs @@ -49,6 +49,7 @@ cfg_io_util! { mod read_exact; mod read_int; mod read_line; + mod fill_buf; mod read_to_end; mod vec_with_initialized; diff --git a/vendor/tokio/src/io/util/read.rs b/vendor/tokio/src/io/util/read.rs index edc9d5a9e..a1f9c8a05 100644 --- a/vendor/tokio/src/io/util/read.rs +++ b/vendor/tokio/src/io/util/read.rs @@ -48,7 +48,7 @@ where fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> { let me = self.project(); - let mut buf = ReadBuf::new(*me.buf); + let mut buf = ReadBuf::new(me.buf); ready!(Pin::new(me.reader).poll_read(cx, &mut buf))?; Poll::Ready(Ok(buf.filled().len())) } diff --git a/vendor/tokio/src/io/util/read_exact.rs b/vendor/tokio/src/io/util/read_exact.rs index 1e8150eb2..dbdd58bae 100644 --- a/vendor/tokio/src/io/util/read_exact.rs +++ b/vendor/tokio/src/io/util/read_exact.rs @@ -51,13 +51,13 @@ where type Output = io::Result<usize>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> { - let mut me = self.project(); + let me = self.project(); loop { // if our buffer is empty, then we need to read some data to continue. let rem = me.buf.remaining(); if rem != 0 { - ready!(Pin::new(&mut *me.reader).poll_read(cx, &mut me.buf))?; + ready!(Pin::new(&mut *me.reader).poll_read(cx, me.buf))?; if me.buf.remaining() == rem { return Err(eof()).into(); } diff --git a/vendor/tokio/src/io/util/read_int.rs b/vendor/tokio/src/io/util/read_int.rs index 5b9fb7bf7..164dcf596 100644 --- a/vendor/tokio/src/io/util/read_int.rs +++ b/vendor/tokio/src/io/util/read_int.rs @@ -142,6 +142,9 @@ reader!(ReadI32, i32, get_i32); reader!(ReadI64, i64, get_i64); reader!(ReadI128, i128, get_i128); +reader!(ReadF32, f32, get_f32); +reader!(ReadF64, f64, get_f64); + reader!(ReadU16Le, u16, get_u16_le); reader!(ReadU32Le, u32, get_u32_le); reader!(ReadU64Le, u64, get_u64_le); @@ -151,3 +154,6 @@ reader!(ReadI16Le, i16, get_i16_le); reader!(ReadI32Le, i32, get_i32_le); reader!(ReadI64Le, i64, get_i64_le); reader!(ReadI128Le, i128, get_i128_le); + +reader!(ReadF32Le, f32, get_f32_le); +reader!(ReadF64Le, f64, get_f64_le); diff --git a/vendor/tokio/src/io/util/read_to_end.rs b/vendor/tokio/src/io/util/read_to_end.rs index f4a564d7d..8edba2a17 100644 --- a/vendor/tokio/src/io/util/read_to_end.rs +++ b/vendor/tokio/src/io/util/read_to_end.rs @@ -1,11 +1,11 @@ use crate::io::util::vec_with_initialized::{into_read_buf_parts, VecU8, VecWithInitialized}; -use crate::io::AsyncRead; +use crate::io::{AsyncRead, ReadBuf}; use pin_project_lite::pin_project; use std::future::Future; use std::io; use std::marker::PhantomPinned; -use std::mem; +use std::mem::{self, MaybeUninit}; use std::pin::Pin; use std::task::{Context, Poll}; @@ -67,16 +67,47 @@ fn poll_read_to_end<V: VecU8, R: AsyncRead + ?Sized>( // has 4 bytes while still making large reads if the reader does have a ton // of data to return. Simply tacking on an extra DEFAULT_BUF_SIZE space every // time is 4,500 times (!) slower than this if the reader has a very small - // amount of data to return. - buf.reserve(32); + // amount of data to return. When the vector is full with its starting + // capacity, we first try to read into a small buffer to see if we reached + // an EOF. This only happens when the starting capacity is >= NUM_BYTES, since + // we allocate at least NUM_BYTES each time. This avoids the unnecessary + // allocation that we attempt before reading into the vector. + + const NUM_BYTES: usize = 32; + let try_small_read = buf.try_small_read_first(NUM_BYTES); // Get a ReadBuf into the vector. - let mut read_buf = buf.get_read_buf(); + let mut read_buf; + let poll_result; + + let n = if try_small_read { + // Read some bytes using a small read. + let mut small_buf: [MaybeUninit<u8>; NUM_BYTES] = [MaybeUninit::uninit(); NUM_BYTES]; + let mut small_read_buf = ReadBuf::uninit(&mut small_buf); + poll_result = read.poll_read(cx, &mut small_read_buf); + let to_write = small_read_buf.filled(); + + // Ensure we have enough space to fill our vector with what we read. + read_buf = buf.get_read_buf(); + if to_write.len() > read_buf.remaining() { + buf.reserve(NUM_BYTES); + read_buf = buf.get_read_buf(); + } + read_buf.put_slice(to_write); + + to_write.len() + } else { + // Ensure we have enough space for reading. + buf.reserve(NUM_BYTES); + read_buf = buf.get_read_buf(); + + // Read data directly into vector. + let filled_before = read_buf.filled().len(); + poll_result = read.poll_read(cx, &mut read_buf); - let filled_before = read_buf.filled().len(); - let poll_result = read.poll_read(cx, &mut read_buf); - let filled_after = read_buf.filled().len(); - let n = filled_after - filled_before; + // Compute the number of bytes read. + read_buf.filled().len() - filled_before + }; // Update the length of the vector using the result of poll_read. let read_buf_parts = into_read_buf_parts(read_buf); @@ -87,11 +118,11 @@ fn poll_read_to_end<V: VecU8, R: AsyncRead + ?Sized>( // In this case, nothing should have been read. However we still // update the vector in case the poll_read call initialized parts of // the vector's unused capacity. - debug_assert_eq!(filled_before, filled_after); + debug_assert_eq!(n, 0); Poll::Pending } Poll::Ready(Err(err)) => { - debug_assert_eq!(filled_before, filled_after); + debug_assert_eq!(n, 0); Poll::Ready(Err(err)) } Poll::Ready(Ok(())) => Poll::Ready(Ok(n)), diff --git a/vendor/tokio/src/io/util/read_until.rs b/vendor/tokio/src/io/util/read_until.rs index 90a0e8a18..fb6fb22d9 100644 --- a/vendor/tokio/src/io/util/read_until.rs +++ b/vendor/tokio/src/io/util/read_until.rs @@ -1,4 +1,5 @@ use crate::io::AsyncBufRead; +use crate::util::memchr; use pin_project_lite::pin_project; use std::future::Future; diff --git a/vendor/tokio/src/io/util/take.rs b/vendor/tokio/src/io/util/take.rs index b5e90c936..df2f61b9e 100644 --- a/vendor/tokio/src/io/util/take.rs +++ b/vendor/tokio/src/io/util/take.rs @@ -86,7 +86,11 @@ impl<R: AsyncRead> AsyncRead for Take<R> { let me = self.project(); let mut b = buf.take(*me.limit_ as usize); + + let buf_ptr = b.filled().as_ptr(); ready!(me.inner.poll_read(cx, &mut b))?; + assert_eq!(b.filled().as_ptr(), buf_ptr); + let n = b.filled().len(); // We need to update the original ReadBuf diff --git a/vendor/tokio/src/io/util/vec_with_initialized.rs b/vendor/tokio/src/io/util/vec_with_initialized.rs index 208cc939c..f527492dc 100644 --- a/vendor/tokio/src/io/util/vec_with_initialized.rs +++ b/vendor/tokio/src/io/util/vec_with_initialized.rs @@ -1,19 +1,18 @@ use crate::io::ReadBuf; use std::mem::MaybeUninit; -mod private { - pub trait Sealed {} - - impl Sealed for Vec<u8> {} - impl Sealed for &mut Vec<u8> {} -} +/// Something that looks like a `Vec<u8>`. +/// +/// # Safety +/// +/// The implementor must guarantee that the vector returned by the +/// `as_mut` and `as_mut` methods do not change from one call to +/// another. +pub(crate) unsafe trait VecU8: AsRef<Vec<u8>> + AsMut<Vec<u8>> {} -/// A sealed trait that constrains the generic type parameter in `VecWithInitialized<V>`. That struct's safety relies -/// on certain invariants upheld by `Vec<u8>`. -pub(crate) trait VecU8: AsMut<Vec<u8>> + private::Sealed {} +unsafe impl VecU8 for Vec<u8> {} +unsafe impl VecU8 for &mut Vec<u8> {} -impl VecU8 for Vec<u8> {} -impl VecU8 for &mut Vec<u8> {} /// This struct wraps a `Vec<u8>` or `&mut Vec<u8>`, combining it with a /// `num_initialized`, which keeps track of the number of initialized bytes /// in the unused capacity. @@ -29,6 +28,7 @@ pub(crate) struct VecWithInitialized<V> { // The number of initialized bytes in the vector. // Always between `vec.len()` and `vec.capacity()`. num_initialized: usize, + starting_capacity: usize, } impl VecWithInitialized<Vec<u8>> { @@ -48,6 +48,7 @@ where // to its length are initialized. Self { num_initialized: vec.as_mut().len(), + starting_capacity: vec.as_ref().capacity(), vec, } } @@ -64,8 +65,8 @@ where } #[cfg(feature = "io-util")] - pub(crate) fn is_empty(&mut self) -> bool { - self.vec.as_mut().is_empty() + pub(crate) fn is_empty(&self) -> bool { + self.vec.as_ref().is_empty() } pub(crate) fn get_read_buf<'a>(&'a mut self) -> ReadBuf<'a> { @@ -112,6 +113,15 @@ where vec.set_len(parts.len); } } + + // Returns a boolean telling the caller to try reading into a small local buffer first if true. + // Doing so would avoid overallocating when vec is filled to capacity and we reached EOF. + pub(crate) fn try_small_read_first(&self, num_bytes: usize) -> bool { + let vec = self.vec.as_ref(); + vec.capacity() - vec.len() < num_bytes + && self.starting_capacity == vec.capacity() + && self.starting_capacity >= num_bytes + } } pub(crate) struct ReadBufParts { diff --git a/vendor/tokio/src/io/util/write_all.rs b/vendor/tokio/src/io/util/write_all.rs index e59d41e4d..abd3e39d3 100644 --- a/vendor/tokio/src/io/util/write_all.rs +++ b/vendor/tokio/src/io/util/write_all.rs @@ -42,7 +42,7 @@ where while !me.buf.is_empty() { let n = ready!(Pin::new(&mut *me.writer).poll_write(cx, me.buf))?; { - let (_, rest) = mem::replace(&mut *me.buf, &[]).split_at(n); + let (_, rest) = mem::take(&mut *me.buf).split_at(n); *me.buf = rest; } if n == 0 { diff --git a/vendor/tokio/src/io/util/write_int.rs b/vendor/tokio/src/io/util/write_int.rs index 13bc191ed..63cd49126 100644 --- a/vendor/tokio/src/io/util/write_int.rs +++ b/vendor/tokio/src/io/util/write_int.rs @@ -135,6 +135,9 @@ writer!(WriteI32, i32, put_i32); writer!(WriteI64, i64, put_i64); writer!(WriteI128, i128, put_i128); +writer!(WriteF32, f32, put_f32); +writer!(WriteF64, f64, put_f64); + writer!(WriteU16Le, u16, put_u16_le); writer!(WriteU32Le, u32, put_u32_le); writer!(WriteU64Le, u64, put_u64_le); @@ -144,3 +147,6 @@ writer!(WriteI16Le, i16, put_i16_le); writer!(WriteI32Le, i32, put_i32_le); writer!(WriteI64Le, i64, put_i64_le); writer!(WriteI128Le, i128, put_i128_le); + +writer!(WriteF32Le, f32, put_f32_le); +writer!(WriteF64Le, f64, put_f64_le); diff --git a/vendor/tokio/src/lib.rs b/vendor/tokio/src/lib.rs index c74b9643f..c712945a4 100644 --- a/vendor/tokio/src/lib.rs +++ b/vendor/tokio/src/lib.rs @@ -1,7 +1,9 @@ #![allow( clippy::cognitive_complexity, clippy::large_enum_variant, - clippy::needless_doctest_main + clippy::module_inception, + clippy::needless_doctest_main, + clippy::declare_interior_mutable_const )] #![warn( missing_debug_implementations, @@ -10,12 +12,13 @@ unreachable_pub )] #![deny(unused_must_use)] -#![cfg_attr(docsrs, deny(broken_intra_doc_links))] #![doc(test( no_crate_inject, attr(deny(warnings, rust_2018_idioms), allow(dead_code, unused_variables)) ))] #![cfg_attr(docsrs, feature(doc_cfg))] +#![cfg_attr(docsrs, allow(unused_attributes))] +#![cfg_attr(loom, allow(dead_code, unreachable_pub))] //! A runtime for writing reliable network applications without compromising speed. //! @@ -114,7 +117,7 @@ //! The [`tokio::sync`] module contains synchronization primitives to use when //! needing to communicate or share data. These include: //! -//! * channels ([`oneshot`], [`mpsc`], and [`watch`]), for sending values +//! * channels ([`oneshot`], [`mpsc`], [`watch`], and [`broadcast`]), for sending values //! between tasks, //! * a non-blocking [`Mutex`], for controlling access to a shared, mutable //! value, @@ -130,6 +133,7 @@ //! [`oneshot`]: crate::sync::oneshot //! [`mpsc`]: crate::sync::mpsc //! [`watch`]: crate::sync::watch +//! [`broadcast`]: crate::sync::broadcast //! //! The [`tokio::time`] module provides utilities for tracking time and //! scheduling work. This includes functions for setting [timeouts][timeout] for @@ -151,7 +155,7 @@ //! provide the functionality you need. //! //! Using the runtime requires the "rt" or "rt-multi-thread" feature flags, to -//! enable the basic [single-threaded scheduler][rt] and the [thread-pool +//! enable the current-thread [single-threaded scheduler][rt] and the [multi-thread //! scheduler][rt-multi-thread], respectively. See the [`runtime` module //! documentation][rt-features] for details. In addition, the "macros" feature //! flag enables the `#[tokio::main]` and `#[tokio::test]` attributes. @@ -160,8 +164,8 @@ //! [`tokio::runtime`]: crate::runtime //! [`Builder`]: crate::runtime::Builder //! [`Runtime`]: crate::runtime::Runtime -//! [rt]: runtime/index.html#basic-scheduler -//! [rt-multi-thread]: runtime/index.html#threaded-scheduler +//! [rt]: runtime/index.html#current-thread-scheduler +//! [rt-multi-thread]: runtime/index.html#multi-thread-scheduler //! [rt-features]: runtime/index.html#runtime-scheduler //! //! ## CPU-bound tasks and blocking code @@ -170,12 +174,15 @@ //! swapping the currently running task on each thread. However, this kind of //! swapping can only happen at `.await` points, so code that spends a long time //! without reaching an `.await` will prevent other tasks from running. To -//! combat this, Tokio provides two kinds of threads: Core threads and blocking -//! threads. The core threads are where all asynchronous code runs, and Tokio -//! will by default spawn one for each CPU core. The blocking threads are -//! spawned on demand, can be used to run blocking code that would otherwise -//! block other tasks from running and are kept alive when not used for a certain -//! amount of time which can be configured with [`thread_keep_alive`]. +//! combat this, Tokio provides two kinds of threads: Core threads and blocking threads. +//! +//! The core threads are where all asynchronous code runs, and Tokio will by default +//! spawn one for each CPU core. You can use the environment variable `TOKIO_WORKER_THREADS` +//! to override the default value. +//! +//! The blocking threads are spawned on demand, can be used to run blocking code +//! that would otherwise block other tasks from running and are kept alive when +//! not used for a certain amount of time which can be configured with [`thread_keep_alive`]. //! Since it is not possible for Tokio to swap out blocking tasks, like it //! can do with asynchronous code, the upper limit on the number of blocking //! threads is very large. These limits can be configured on the [`Builder`]. @@ -204,9 +211,15 @@ //! ``` //! //! If your code is CPU-bound and you wish to limit the number of threads used -//! to run it, you should run it on another thread pool such as [rayon]. You -//! can use an [`oneshot`] channel to send the result back to Tokio when the -//! rayon task finishes. +//! to run it, you should use a separate thread pool dedicated to CPU bound tasks. +//! For example, you could consider using the [rayon] library for CPU-bound +//! tasks. It is also possible to create an extra Tokio runtime dedicated to +//! CPU-bound tasks, but if you do this, you should be careful that the extra +//! runtime runs _only_ CPU-bound tasks, as IO-bound tasks on that runtime +//! will behave poorly. +//! +//! Hint: If using rayon, you can use a [`oneshot`] channel to send the result back +//! to Tokio when the rayon task finishes. //! //! [rayon]: https://docs.rs/rayon //! [`oneshot`]: crate::sync::oneshot @@ -301,14 +314,15 @@ //! Beware though that this will pull in many extra dependencies that you may not //! need. //! -//! - `full`: Enables all Tokio public API features listed below except `test-util`. -//! - `rt`: Enables `tokio::spawn`, the basic (current thread) scheduler, +//! - `full`: Enables all features listed below except `test-util` and `tracing`. +//! - `rt`: Enables `tokio::spawn`, the current-thread scheduler, //! and non-scheduler utilities. //! - `rt-multi-thread`: Enables the heavier, multi-threaded, work-stealing scheduler. //! - `io-util`: Enables the IO based `Ext` traits. //! - `io-std`: Enable `Stdout`, `Stdin` and `Stderr` types. -//! - `net`: Enables `tokio::net` types such as `TcpStream`, `UnixStream` and `UdpSocket`, -//! as well as (on Unix-like systems) `AsyncFd` +//! - `net`: Enables `tokio::net` types such as `TcpStream`, `UnixStream` and +//! `UdpSocket`, as well as (on Unix-like systems) `AsyncFd` and (on +//! FreeBSD) `PollAio`. //! - `time`: Enables `tokio::time` types and allows the schedulers to enable //! the built in timer. //! - `process`: Enables `tokio::process` types. @@ -317,30 +331,176 @@ //! - `signal`: Enables all `tokio::signal` types. //! - `fs`: Enables `tokio::fs` types. //! - `test-util`: Enables testing based infrastructure for the Tokio runtime. +//! - `parking_lot`: As a potential optimization, use the _parking_lot_ crate's +//! synchronization primitives internally. Also, this +//! dependency is necessary to construct some of our primitives +//! in a const context. MSRV may increase according to the +//! _parking_lot_ release in use. //! //! _Note: `AsyncRead` and `AsyncWrite` traits do not require any features and are //! always available._ //! -//! ### Internal features +//! ### Unstable features //! -//! These features do not expose any new API, but influence internal -//! implementation aspects of Tokio, and can pull in additional -//! dependencies. +//! Some feature flags are only available when specifying the `tokio_unstable` flag: //! -//! - `parking_lot`: As a potential optimization, use the _parking_lot_ crate's -//! synchronization primitives internally. MSRV may increase according to the -//! _parking_lot_ release in use. +//! - `tracing`: Enables tracing events. //! -//! ### Unstable features +//! Likewise, some parts of the API are only available with the same flag: //! -//! These feature flags enable **unstable** features. The public API may break in 1.x -//! releases. To enable these features, the `--cfg tokio_unstable` must be passed to -//! `rustc` when compiling. This is easiest done using the `RUSTFLAGS` env variable: -//! `RUSTFLAGS="--cfg tokio_unstable"`. +//! - [`task::Builder`] +//! - Some methods on [`task::JoinSet`] +//! - [`runtime::RuntimeMetrics`] +//! - [`runtime::Builder::unhandled_panic`] +//! - [`task::Id`] //! -//! - `tracing`: Enables tracing events. +//! This flag enables **unstable** features. The public API of these features +//! may break in 1.x releases. To enable these features, the `--cfg +//! tokio_unstable` argument must be passed to `rustc` when compiling. This +//! serves to explicitly opt-in to features which may break semver conventions, +//! since Cargo [does not yet directly support such opt-ins][unstable features]. +//! +//! You can specify it in your project's `.cargo/config.toml` file: +//! +//! ```toml +//! [build] +//! rustflags = ["--cfg", "tokio_unstable"] +//! ``` +//! +//! Alternatively, you can specify it with an environment variable: +//! +//! ```sh +//! ## Many *nix shells: +//! export RUSTFLAGS="--cfg tokio_unstable" +//! cargo build +//! ``` +//! +//! ```powershell +//! ## Windows PowerShell: +//! $Env:RUSTFLAGS="--cfg tokio_unstable" +//! cargo build +//! ``` //! +//! [unstable features]: https://internals.rust-lang.org/t/feature-request-unstable-opt-in-non-transitive-crate-features/16193#why-not-a-crate-feature-2 //! [feature flags]: https://doc.rust-lang.org/cargo/reference/manifest.html#the-features-section +//! +//! ## Supported platforms +//! +//! Tokio currently guarantees support for the following platforms: +//! +//! * Linux +//! * Windows +//! * Android (API level 21) +//! * macOS +//! * iOS +//! * FreeBSD +//! +//! Tokio will continue to support these platforms in the future. However, +//! future releases may change requirements such as the minimum required libc +//! version on Linux, the API level on Android, or the supported FreeBSD +//! release. +//! +//! Beyond the above platforms, Tokio is intended to work on all platforms +//! supported by the mio crate. You can find a longer list [in mio's +//! documentation][mio-supported]. However, these additional platforms may +//! become unsupported in the future. +//! +//! Note that Wine is considered to be a different platform from Windows. See +//! mio's documentation for more information on Wine support. +//! +//! [mio-supported]: https://crates.io/crates/mio#platforms +//! +//! ### WASM support +//! +//! Tokio has some limited support for the WASM platform. Without the +//! `tokio_unstable` flag, the following features are supported: +//! +//! * `sync` +//! * `macros` +//! * `io-util` +//! * `rt` +//! * `time` +//! +//! Enabling any other feature (including `full`) will cause a compilation +//! failure. +//! +//! The `time` module will only work on WASM platforms that have support for +//! timers (e.g. wasm32-wasi). The timing functions will panic if used on a WASM +//! platform that does not support timers. +//! +//! Note also that if the runtime becomes indefinitely idle, it will panic +//! immediately instead of blocking forever. On platforms that don't support +//! time, this means that the runtime can never be idle in any way. +//! +//! ### Unstable WASM support +//! +//! Tokio also has unstable support for some additional WASM features. This +//! requires the use of the `tokio_unstable` flag. +//! +//! Using this flag enables the use of `tokio::net` on the wasm32-wasi target. +//! However, not all methods are available on the networking types as WASI +//! currently does not support the creation of new sockets from within WASM. +//! Because of this, sockets must currently be created via the `FromRawFd` +//! trait. + +// Test that pointer width is compatible. This asserts that e.g. usize is at +// least 32 bits, which a lot of components in Tokio currently assumes. +// +// TODO: improve once we have MSRV access to const eval to make more flexible. +#[cfg(not(any( + target_pointer_width = "32", + target_pointer_width = "64", + target_pointer_width = "128" +)))] +compile_error! { + "Tokio requires the platform pointer width to be 32, 64, or 128 bits" +} + +// Ensure that our build script has correctly set cfg flags for wasm. +// +// Each condition is written all(a, not(b)). This should be read as +// "if a, then we must also have b". +#[cfg(any( + all(target_arch = "wasm32", not(tokio_wasm)), + all(target_arch = "wasm64", not(tokio_wasm)), + all(target_family = "wasm", not(tokio_wasm)), + all(target_os = "wasi", not(tokio_wasm)), + all(target_os = "wasi", not(tokio_wasi)), + all(target_os = "wasi", tokio_wasm_not_wasi), + all(tokio_wasm, not(any(target_arch = "wasm32", target_arch = "wasm64"))), + all(tokio_wasm_not_wasi, not(tokio_wasm)), + all(tokio_wasi, not(tokio_wasm)) +))] +compile_error!("Tokio's build script has incorrectly detected wasm."); + +#[cfg(all( + not(tokio_unstable), + tokio_wasm, + any( + feature = "fs", + feature = "io-std", + feature = "net", + feature = "process", + feature = "rt-multi-thread", + feature = "signal" + ) +))] +compile_error!("Only features sync,macros,io-util,rt,time are supported on wasm."); + +#[cfg(all(not(tokio_unstable), tokio_taskdump))] +compile_error!("The `tokio_taskdump` feature requires `--cfg tokio_unstable`."); + +#[cfg(all( + tokio_taskdump, + not(all( + target_os = "linux", + any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64") + )) +))] +compile_error!( + "The `tokio_taskdump` feature is only currently supported on \ +linux, on `aarch64`, `x86` and `x86_64`." +); // Includes re-exports used by macros. // @@ -360,20 +520,25 @@ pub mod io; pub mod net; mod loom; -mod park; cfg_process! { pub mod process; } -#[cfg(any(feature = "net", feature = "fs", feature = "io-std"))] +#[cfg(any( + feature = "fs", + feature = "io-std", + feature = "net", + all(windows, feature = "process"), +))] mod blocking; cfg_rt! { pub mod runtime; } - -pub(crate) mod coop; +cfg_not_rt! { + pub(crate) mod runtime; +} cfg_signal! { pub mod signal; @@ -402,6 +567,40 @@ cfg_time! { pub mod time; } +mod trace { + use std::future::Future; + use std::pin::Pin; + use std::task::{Context, Poll}; + + cfg_taskdump! { + pub(crate) use crate::runtime::task::trace::trace_leaf; + } + + cfg_not_taskdump! { + #[inline(always)] + #[allow(dead_code)] + pub(crate) fn trace_leaf(_: &mut std::task::Context<'_>) -> std::task::Poll<()> { + std::task::Poll::Ready(()) + } + } + + #[cfg_attr(not(feature = "sync"), allow(dead_code))] + pub(crate) fn async_trace_leaf() -> impl Future<Output = ()> { + struct Trace; + + impl Future for Trace { + type Output = (); + + #[inline(always)] + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { + trace_leaf(cx) + } + } + + Trace + } +} + mod util; /// Due to the `Stream` trait's inclusion in `std` landing later than Tokio's 1.0 @@ -457,14 +656,6 @@ pub(crate) use self::doc::os; #[allow(unused)] pub(crate) use std::os; -#[cfg(docsrs)] -#[allow(unused)] -pub(crate) use self::doc::winapi; - -#[cfg(all(not(docsrs), windows, feature = "net"))] -#[allow(unused)] -pub(crate) use ::winapi; - cfg_macros! { /// Implementation detail of the `select!` macro. This macro is **not** /// intended to be used as part of the public API and is permitted to @@ -472,6 +663,12 @@ cfg_macros! { #[doc(hidden)] pub use tokio_macros::select_priv_declare_output_enum; + /// Implementation detail of the `select!` macro. This macro is **not** + /// intended to be used as part of the public API and is permitted to + /// change. + #[doc(hidden)] + pub use tokio_macros::select_priv_clean_pattern; + cfg_rt! { #[cfg(feature = "rt-multi-thread")] #[cfg(not(test))] // Work around for rust-lang/rust#62127 @@ -509,3 +706,7 @@ cfg_macros! { #[cfg(feature = "io-util")] #[cfg(test)] fn is_unpin<T: Unpin>() {} + +/// fuzz test (fuzz_linked_list) +#[cfg(fuzzing)] +pub mod fuzz; diff --git a/vendor/tokio/src/loom/mocked.rs b/vendor/tokio/src/loom/mocked.rs index 367d59b43..56dc1a063 100644 --- a/vendor/tokio/src/loom/mocked.rs +++ b/vendor/tokio/src/loom/mocked.rs @@ -25,6 +25,13 @@ pub(crate) mod sync { } } pub(crate) use loom::sync::*; + + pub(crate) mod atomic { + pub(crate) use loom::sync::atomic::*; + + // TODO: implement a loom version + pub(crate) type StaticAtomicU64 = std::sync::atomic::AtomicU64; + } } pub(crate) mod rand { @@ -38,3 +45,8 @@ pub(crate) mod sys { 2 } } + +pub(crate) mod thread { + pub use loom::lazy_static::AccessError; + pub use loom::thread::*; +} diff --git a/vendor/tokio/src/loom/std/atomic_ptr.rs b/vendor/tokio/src/loom/std/atomic_ptr.rs deleted file mode 100644 index 236645f03..000000000 --- a/vendor/tokio/src/loom/std/atomic_ptr.rs +++ /dev/null @@ -1,34 +0,0 @@ -use std::fmt; -use std::ops::{Deref, DerefMut}; - -/// `AtomicPtr` providing an additional `load_unsync` function. -pub(crate) struct AtomicPtr<T> { - inner: std::sync::atomic::AtomicPtr<T>, -} - -impl<T> AtomicPtr<T> { - pub(crate) fn new(ptr: *mut T) -> AtomicPtr<T> { - let inner = std::sync::atomic::AtomicPtr::new(ptr); - AtomicPtr { inner } - } -} - -impl<T> Deref for AtomicPtr<T> { - type Target = std::sync::atomic::AtomicPtr<T>; - - fn deref(&self) -> &Self::Target { - &self.inner - } -} - -impl<T> DerefMut for AtomicPtr<T> { - fn deref_mut(&mut self) -> &mut Self::Target { - &mut self.inner - } -} - -impl<T> fmt::Debug for AtomicPtr<T> { - fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - self.deref().fmt(fmt) - } -} diff --git a/vendor/tokio/src/loom/std/atomic_u16.rs b/vendor/tokio/src/loom/std/atomic_u16.rs index c1c531208..c9e105c19 100644 --- a/vendor/tokio/src/loom/std/atomic_u16.rs +++ b/vendor/tokio/src/loom/std/atomic_u16.rs @@ -2,7 +2,7 @@ use std::cell::UnsafeCell; use std::fmt; use std::ops::Deref; -/// `AtomicU16` providing an additional `load_unsync` function. +/// `AtomicU16` providing an additional `unsync_load` function. pub(crate) struct AtomicU16 { inner: UnsafeCell<std::sync::atomic::AtomicU16>, } @@ -23,7 +23,7 @@ impl AtomicU16 { /// All mutations must have happened before the unsynchronized load. /// Additionally, there must be no concurrent mutations. pub(crate) unsafe fn unsync_load(&self) -> u16 { - *(*self.inner.get()).get_mut() + core::ptr::read(self.inner.get() as *const u16) } } diff --git a/vendor/tokio/src/loom/std/atomic_u32.rs b/vendor/tokio/src/loom/std/atomic_u32.rs index 61f95fb30..ee0d2d380 100644 --- a/vendor/tokio/src/loom/std/atomic_u32.rs +++ b/vendor/tokio/src/loom/std/atomic_u32.rs @@ -2,7 +2,7 @@ use std::cell::UnsafeCell; use std::fmt; use std::ops::Deref; -/// `AtomicU32` providing an additional `load_unsync` function. +/// `AtomicU32` providing an additional `unsync_load` function. pub(crate) struct AtomicU32 { inner: UnsafeCell<std::sync::atomic::AtomicU32>, } @@ -15,6 +15,16 @@ impl AtomicU32 { let inner = UnsafeCell::new(std::sync::atomic::AtomicU32::new(val)); AtomicU32 { inner } } + + /// Performs an unsynchronized load. + /// + /// # Safety + /// + /// All mutations must have happened before the unsynchronized load. + /// Additionally, there must be no concurrent mutations. + pub(crate) unsafe fn unsync_load(&self) -> u32 { + core::ptr::read(self.inner.get() as *const u32) + } } impl Deref for AtomicU32 { diff --git a/vendor/tokio/src/loom/std/atomic_u64.rs b/vendor/tokio/src/loom/std/atomic_u64.rs index 7eb457a24..ce391be3e 100644 --- a/vendor/tokio/src/loom/std/atomic_u64.rs +++ b/vendor/tokio/src/loom/std/atomic_u64.rs @@ -2,74 +2,18 @@ //! re-export of `AtomicU64`. On 32 bit platforms, this is implemented using a //! `Mutex`. -pub(crate) use self::imp::AtomicU64; - // `AtomicU64` can only be used on targets with `target_has_atomic` is 64 or greater. // Once `cfg_target_has_atomic` feature is stable, we can replace it with // `#[cfg(target_has_atomic = "64")]`. // Refs: https://github.com/rust-lang/rust/tree/master/src/librustc_target -#[cfg(not(any(target_arch = "arm", target_arch = "mips", target_arch = "powerpc")))] -mod imp { - pub(crate) use std::sync::atomic::AtomicU64; +cfg_has_atomic_u64! { + #[path = "atomic_u64_native.rs"] + mod imp; } -#[cfg(any(target_arch = "arm", target_arch = "mips", target_arch = "powerpc"))] -mod imp { - use crate::loom::sync::Mutex; - use std::sync::atomic::Ordering; - - #[derive(Debug)] - pub(crate) struct AtomicU64 { - inner: Mutex<u64>, - } - - impl AtomicU64 { - pub(crate) fn new(val: u64) -> Self { - Self { - inner: Mutex::new(val), - } - } - - pub(crate) fn load(&self, _: Ordering) -> u64 { - *self.inner.lock() - } - - pub(crate) fn store(&self, val: u64, _: Ordering) { - *self.inner.lock() = val; - } - - pub(crate) fn fetch_or(&self, val: u64, _: Ordering) -> u64 { - let mut lock = self.inner.lock(); - let prev = *lock; - *lock = prev | val; - prev - } - - pub(crate) fn compare_exchange( - &self, - current: u64, - new: u64, - _success: Ordering, - _failure: Ordering, - ) -> Result<u64, u64> { - let mut lock = self.inner.lock(); - - if *lock == current { - *lock = new; - Ok(current) - } else { - Err(*lock) - } - } - - pub(crate) fn compare_exchange_weak( - &self, - current: u64, - new: u64, - success: Ordering, - failure: Ordering, - ) -> Result<u64, u64> { - self.compare_exchange(current, new, success, failure) - } - } +cfg_not_has_atomic_u64! { + #[path = "atomic_u64_as_mutex.rs"] + mod imp; } + +pub(crate) use imp::{AtomicU64, StaticAtomicU64}; diff --git a/vendor/tokio/src/loom/std/atomic_u64_as_mutex.rs b/vendor/tokio/src/loom/std/atomic_u64_as_mutex.rs new file mode 100644 index 000000000..9b3b6fac6 --- /dev/null +++ b/vendor/tokio/src/loom/std/atomic_u64_as_mutex.rs @@ -0,0 +1,76 @@ +use crate::loom::sync::Mutex; +use std::sync::atomic::Ordering; + +cfg_has_const_mutex_new! { + #[path = "atomic_u64_static_const_new.rs"] + mod static_macro; +} + +cfg_not_has_const_mutex_new! { + #[path = "atomic_u64_static_once_cell.rs"] + mod static_macro; +} + +pub(crate) use static_macro::StaticAtomicU64; + +#[derive(Debug)] +pub(crate) struct AtomicU64 { + inner: Mutex<u64>, +} + +impl AtomicU64 { + pub(crate) fn load(&self, _: Ordering) -> u64 { + *self.inner.lock() + } + + pub(crate) fn store(&self, val: u64, _: Ordering) { + *self.inner.lock() = val; + } + + pub(crate) fn fetch_add(&self, val: u64, _: Ordering) -> u64 { + let mut lock = self.inner.lock(); + let prev = *lock; + *lock = prev + val; + prev + } + + pub(crate) fn fetch_or(&self, val: u64, _: Ordering) -> u64 { + let mut lock = self.inner.lock(); + let prev = *lock; + *lock = prev | val; + prev + } + + pub(crate) fn compare_exchange( + &self, + current: u64, + new: u64, + _success: Ordering, + _failure: Ordering, + ) -> Result<u64, u64> { + let mut lock = self.inner.lock(); + + if *lock == current { + *lock = new; + Ok(current) + } else { + Err(*lock) + } + } + + pub(crate) fn compare_exchange_weak( + &self, + current: u64, + new: u64, + success: Ordering, + failure: Ordering, + ) -> Result<u64, u64> { + self.compare_exchange(current, new, success, failure) + } +} + +impl Default for AtomicU64 { + fn default() -> AtomicU64 { + AtomicU64::new(u64::default()) + } +} diff --git a/vendor/tokio/src/loom/std/atomic_u64_native.rs b/vendor/tokio/src/loom/std/atomic_u64_native.rs new file mode 100644 index 000000000..08adb2862 --- /dev/null +++ b/vendor/tokio/src/loom/std/atomic_u64_native.rs @@ -0,0 +1,4 @@ +pub(crate) use std::sync::atomic::{AtomicU64, Ordering}; + +/// Alias `AtomicU64` to `StaticAtomicU64` +pub(crate) type StaticAtomicU64 = AtomicU64; diff --git a/vendor/tokio/src/loom/std/atomic_u64_static_const_new.rs b/vendor/tokio/src/loom/std/atomic_u64_static_const_new.rs new file mode 100644 index 000000000..a4215342b --- /dev/null +++ b/vendor/tokio/src/loom/std/atomic_u64_static_const_new.rs @@ -0,0 +1,12 @@ +use super::AtomicU64; +use crate::loom::sync::Mutex; + +pub(crate) type StaticAtomicU64 = AtomicU64; + +impl AtomicU64 { + pub(crate) const fn new(val: u64) -> Self { + Self { + inner: Mutex::const_new(val), + } + } +} diff --git a/vendor/tokio/src/loom/std/atomic_u64_static_once_cell.rs b/vendor/tokio/src/loom/std/atomic_u64_static_once_cell.rs new file mode 100644 index 000000000..40c6172a5 --- /dev/null +++ b/vendor/tokio/src/loom/std/atomic_u64_static_once_cell.rs @@ -0,0 +1,57 @@ +use super::AtomicU64; +use crate::loom::sync::{atomic::Ordering, Mutex}; +use crate::util::once_cell::OnceCell; + +pub(crate) struct StaticAtomicU64 { + init: u64, + cell: OnceCell<Mutex<u64>>, +} + +impl AtomicU64 { + pub(crate) fn new(val: u64) -> Self { + Self { + inner: Mutex::new(val), + } + } +} + +impl StaticAtomicU64 { + pub(crate) const fn new(val: u64) -> StaticAtomicU64 { + StaticAtomicU64 { + init: val, + cell: OnceCell::new(), + } + } + + pub(crate) fn load(&self, order: Ordering) -> u64 { + *self.inner().lock() + } + + pub(crate) fn fetch_add(&self, val: u64, order: Ordering) -> u64 { + let mut lock = self.inner().lock(); + let prev = *lock; + *lock = prev + val; + prev + } + + pub(crate) fn compare_exchange_weak( + &self, + current: u64, + new: u64, + _success: Ordering, + _failure: Ordering, + ) -> Result<u64, u64> { + let mut lock = self.inner().lock(); + + if *lock == current { + *lock = new; + Ok(current) + } else { + Err(*lock) + } + } + + fn inner(&self) -> &Mutex<u64> { + self.cell.get(|| Mutex::new(self.init)) + } +} diff --git a/vendor/tokio/src/loom/std/atomic_u8.rs b/vendor/tokio/src/loom/std/atomic_u8.rs deleted file mode 100644 index 408aea338..000000000 --- a/vendor/tokio/src/loom/std/atomic_u8.rs +++ /dev/null @@ -1,34 +0,0 @@ -use std::cell::UnsafeCell; -use std::fmt; -use std::ops::Deref; - -/// `AtomicU8` providing an additional `load_unsync` function. -pub(crate) struct AtomicU8 { - inner: UnsafeCell<std::sync::atomic::AtomicU8>, -} - -unsafe impl Send for AtomicU8 {} -unsafe impl Sync for AtomicU8 {} - -impl AtomicU8 { - pub(crate) const fn new(val: u8) -> AtomicU8 { - let inner = UnsafeCell::new(std::sync::atomic::AtomicU8::new(val)); - AtomicU8 { inner } - } -} - -impl Deref for AtomicU8 { - type Target = std::sync::atomic::AtomicU8; - - fn deref(&self) -> &Self::Target { - // safety: it is always safe to access `&self` fns on the inner value as - // we never perform unsafe mutations. - unsafe { &*self.inner.get() } - } -} - -impl fmt::Debug for AtomicU8 { - fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - self.deref().fmt(fmt) - } -} diff --git a/vendor/tokio/src/loom/std/atomic_usize.rs b/vendor/tokio/src/loom/std/atomic_usize.rs index 0d5f36e43..c5503a2c1 100644 --- a/vendor/tokio/src/loom/std/atomic_usize.rs +++ b/vendor/tokio/src/loom/std/atomic_usize.rs @@ -2,7 +2,7 @@ use std::cell::UnsafeCell; use std::fmt; use std::ops; -/// `AtomicUsize` providing an additional `load_unsync` function. +/// `AtomicUsize` providing an additional `unsync_load` function. pub(crate) struct AtomicUsize { inner: UnsafeCell<std::sync::atomic::AtomicUsize>, } @@ -23,7 +23,7 @@ impl AtomicUsize { /// All mutations must have happened before the unsynchronized load. /// Additionally, there must be no concurrent mutations. pub(crate) unsafe fn unsync_load(&self) -> usize { - *(*self.inner.get()).get_mut() + core::ptr::read(self.inner.get() as *const usize) } pub(crate) fn with_mut<R>(&mut self, f: impl FnOnce(&mut usize) -> R) -> R { diff --git a/vendor/tokio/src/loom/std/barrier.rs b/vendor/tokio/src/loom/std/barrier.rs new file mode 100644 index 000000000..a3f0ca0ab --- /dev/null +++ b/vendor/tokio/src/loom/std/barrier.rs @@ -0,0 +1,217 @@ +//! A `Barrier` that provides `wait_timeout`. +//! +//! This implementation mirrors that of the Rust standard library. + +use crate::loom::sync::{Condvar, Mutex}; +use std::fmt; +use std::time::{Duration, Instant}; + +/// A barrier enables multiple threads to synchronize the beginning +/// of some computation. +/// +/// # Examples +/// +/// ``` +/// use std::sync::{Arc, Barrier}; +/// use std::thread; +/// +/// let mut handles = Vec::with_capacity(10); +/// let barrier = Arc::new(Barrier::new(10)); +/// for _ in 0..10 { +/// let c = Arc::clone(&barrier); +/// // The same messages will be printed together. +/// // You will NOT see any interleaving. +/// handles.push(thread::spawn(move|| { +/// println!("before wait"); +/// c.wait(); +/// println!("after wait"); +/// })); +/// } +/// // Wait for other threads to finish. +/// for handle in handles { +/// handle.join().unwrap(); +/// } +/// ``` +pub(crate) struct Barrier { + lock: Mutex<BarrierState>, + cvar: Condvar, + num_threads: usize, +} + +// The inner state of a double barrier +struct BarrierState { + count: usize, + generation_id: usize, +} + +/// A `BarrierWaitResult` is returned by [`Barrier::wait()`] when all threads +/// in the [`Barrier`] have rendezvoused. +/// +/// # Examples +/// +/// ``` +/// use std::sync::Barrier; +/// +/// let barrier = Barrier::new(1); +/// let barrier_wait_result = barrier.wait(); +/// ``` +pub(crate) struct BarrierWaitResult(bool); + +impl fmt::Debug for Barrier { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Barrier").finish_non_exhaustive() + } +} + +impl Barrier { + /// Creates a new barrier that can block a given number of threads. + /// + /// A barrier will block `n`-1 threads which call [`wait()`] and then wake + /// up all threads at once when the `n`th thread calls [`wait()`]. + /// + /// [`wait()`]: Barrier::wait + /// + /// # Examples + /// + /// ``` + /// use std::sync::Barrier; + /// + /// let barrier = Barrier::new(10); + /// ``` + #[must_use] + pub(crate) fn new(n: usize) -> Barrier { + Barrier { + lock: Mutex::new(BarrierState { + count: 0, + generation_id: 0, + }), + cvar: Condvar::new(), + num_threads: n, + } + } + + /// Blocks the current thread until all threads have rendezvoused here. + /// + /// Barriers are re-usable after all threads have rendezvoused once, and can + /// be used continuously. + /// + /// A single (arbitrary) thread will receive a [`BarrierWaitResult`] that + /// returns `true` from [`BarrierWaitResult::is_leader()`] when returning + /// from this function, and all other threads will receive a result that + /// will return `false` from [`BarrierWaitResult::is_leader()`]. + /// + /// # Examples + /// + /// ``` + /// use std::sync::{Arc, Barrier}; + /// use std::thread; + /// + /// let mut handles = Vec::with_capacity(10); + /// let barrier = Arc::new(Barrier::new(10)); + /// for _ in 0..10 { + /// let c = Arc::clone(&barrier); + /// // The same messages will be printed together. + /// // You will NOT see any interleaving. + /// handles.push(thread::spawn(move|| { + /// println!("before wait"); + /// c.wait(); + /// println!("after wait"); + /// })); + /// } + /// // Wait for other threads to finish. + /// for handle in handles { + /// handle.join().unwrap(); + /// } + /// ``` + pub(crate) fn wait(&self) -> BarrierWaitResult { + let mut lock = self.lock.lock(); + let local_gen = lock.generation_id; + lock.count += 1; + if lock.count < self.num_threads { + // We need a while loop to guard against spurious wakeups. + // https://en.wikipedia.org/wiki/Spurious_wakeup + while local_gen == lock.generation_id { + lock = self.cvar.wait(lock).unwrap(); + } + BarrierWaitResult(false) + } else { + lock.count = 0; + lock.generation_id = lock.generation_id.wrapping_add(1); + self.cvar.notify_all(); + BarrierWaitResult(true) + } + } + + /// Blocks the current thread until all threads have rendezvoused here for + /// at most `timeout` duration. + pub(crate) fn wait_timeout(&self, timeout: Duration) -> Option<BarrierWaitResult> { + // This implementation mirrors `wait`, but with each blocking operation + // replaced by a timeout-amenable alternative. + + let deadline = Instant::now() + timeout; + + // Acquire `self.lock` with at most `timeout` duration. + let mut lock = loop { + if let Some(guard) = self.lock.try_lock() { + break guard; + } else if Instant::now() > deadline { + return None; + } else { + std::thread::yield_now(); + } + }; + + // Shrink the `timeout` to account for the time taken to acquire `lock`. + let timeout = deadline.saturating_duration_since(Instant::now()); + + let local_gen = lock.generation_id; + lock.count += 1; + if lock.count < self.num_threads { + // We need a while loop to guard against spurious wakeups. + // https://en.wikipedia.org/wiki/Spurious_wakeup + while local_gen == lock.generation_id { + let (guard, timeout_result) = self.cvar.wait_timeout(lock, timeout).unwrap(); + lock = guard; + if timeout_result.timed_out() { + return None; + } + } + Some(BarrierWaitResult(false)) + } else { + lock.count = 0; + lock.generation_id = lock.generation_id.wrapping_add(1); + self.cvar.notify_all(); + Some(BarrierWaitResult(true)) + } + } +} + +impl fmt::Debug for BarrierWaitResult { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BarrierWaitResult") + .field("is_leader", &self.is_leader()) + .finish() + } +} + +impl BarrierWaitResult { + /// Returns `true` if this thread is the "leader thread" for the call to + /// [`Barrier::wait()`]. + /// + /// Only one thread will have `true` returned from their result, all other + /// threads will have `false` returned. + /// + /// # Examples + /// + /// ``` + /// use std::sync::Barrier; + /// + /// let barrier = Barrier::new(1); + /// let barrier_wait_result = barrier.wait(); + /// println!("{:?}", barrier_wait_result.is_leader()); + /// ``` + #[must_use] + pub(crate) fn is_leader(&self) -> bool { + self.0 + } +} diff --git a/vendor/tokio/src/loom/std/mod.rs b/vendor/tokio/src/loom/std/mod.rs index b29cbeeb8..0a732791f 100644 --- a/vendor/tokio/src/loom/std/mod.rs +++ b/vendor/tokio/src/loom/std/mod.rs @@ -1,11 +1,10 @@ #![cfg_attr(any(not(feature = "full"), loom), allow(unused_imports, dead_code))] -mod atomic_ptr; mod atomic_u16; mod atomic_u32; mod atomic_u64; -mod atomic_u8; mod atomic_usize; +mod barrier; mod mutex; #[cfg(feature = "parking_lot")] mod parking_lot; @@ -25,6 +24,10 @@ pub(crate) mod future { pub(crate) use crate::sync::AtomicWaker; } +pub(crate) mod hint { + pub(crate) use std::hint::spin_loop; +} + pub(crate) mod rand { use std::collections::hash_map::RandomState; use std::hash::{BuildHasher, Hash, Hasher}; @@ -67,24 +70,41 @@ pub(crate) mod sync { pub(crate) use crate::loom::std::mutex::Mutex; pub(crate) mod atomic { - pub(crate) use crate::loom::std::atomic_ptr::AtomicPtr; pub(crate) use crate::loom::std::atomic_u16::AtomicU16; pub(crate) use crate::loom::std::atomic_u32::AtomicU32; - pub(crate) use crate::loom::std::atomic_u64::AtomicU64; - pub(crate) use crate::loom::std::atomic_u8::AtomicU8; + pub(crate) use crate::loom::std::atomic_u64::{AtomicU64, StaticAtomicU64}; pub(crate) use crate::loom::std::atomic_usize::AtomicUsize; - pub(crate) use std::sync::atomic::{fence, AtomicBool, Ordering}; - // TODO: once we bump MSRV to 1.49+, use `hint::spin_loop` instead. - #[allow(deprecated)] - pub(crate) use std::sync::atomic::spin_loop_hint; + pub(crate) use std::sync::atomic::{fence, AtomicBool, AtomicPtr, AtomicU8, Ordering}; } + + pub(crate) use super::barrier::Barrier; } pub(crate) mod sys { #[cfg(feature = "rt-multi-thread")] pub(crate) fn num_cpus() -> usize { - usize::max(1, num_cpus::get()) + const ENV_WORKER_THREADS: &str = "TOKIO_WORKER_THREADS"; + + match std::env::var(ENV_WORKER_THREADS) { + Ok(s) => { + let n = s.parse().unwrap_or_else(|e| { + panic!( + "\"{}\" must be usize, error: {}, value: {}", + ENV_WORKER_THREADS, e, s + ) + }); + assert!(n > 0, "\"{}\" cannot be set to 0", ENV_WORKER_THREADS); + n + } + Err(std::env::VarError::NotPresent) => usize::max(1, num_cpus::get()), + Err(std::env::VarError::NotUnicode(e)) => { + panic!( + "\"{}\" must be valid unicode, error: {:?}", + ENV_WORKER_THREADS, e + ) + } + } } #[cfg(not(feature = "rt-multi-thread"))] @@ -93,4 +113,15 @@ pub(crate) mod sys { } } -pub(crate) use std::thread; +pub(crate) mod thread { + #[inline] + pub(crate) fn yield_now() { + std::hint::spin_loop(); + } + + #[allow(unused_imports)] + pub(crate) use std::thread::{ + current, panicking, park, park_timeout, sleep, spawn, AccessError, Builder, JoinHandle, + LocalKey, Result, Thread, ThreadId, + }; +} diff --git a/vendor/tokio/src/loom/std/mutex.rs b/vendor/tokio/src/loom/std/mutex.rs index bf14d6242..076f78611 100644 --- a/vendor/tokio/src/loom/std/mutex.rs +++ b/vendor/tokio/src/loom/std/mutex.rs @@ -1,7 +1,7 @@ use std::sync::{self, MutexGuard, TryLockError}; /// Adapter for `std::Mutex` that removes the poisoning aspects -// from its api +/// from its api. #[derive(Debug)] pub(crate) struct Mutex<T: ?Sized>(sync::Mutex<T>); @@ -13,6 +13,12 @@ impl<T> Mutex<T> { } #[inline] + #[cfg(not(tokio_no_const_mutex_new))] + pub(crate) const fn const_new(t: T) -> Mutex<T> { + Mutex(sync::Mutex::new(t)) + } + + #[inline] pub(crate) fn lock(&self) -> MutexGuard<'_, T> { match self.0.lock() { Ok(guard) => guard, diff --git a/vendor/tokio/src/loom/std/parking_lot.rs b/vendor/tokio/src/loom/std/parking_lot.rs index 8448bed53..e3af258d1 100644 --- a/vendor/tokio/src/loom/std/parking_lot.rs +++ b/vendor/tokio/src/loom/std/parking_lot.rs @@ -3,83 +3,143 @@ //! //! This can be extended to additional types/methods as required. +use std::fmt; +use std::marker::PhantomData; +use std::ops::{Deref, DerefMut}; use std::sync::LockResult; use std::time::Duration; +// All types in this file are marked with PhantomData to ensure that +// parking_lot's send_guard feature does not leak through and affect when Tokio +// types are Send. +// +// See <https://github.com/tokio-rs/tokio/pull/4359> for more info. + // Types that do not need wrapping -pub(crate) use parking_lot::{MutexGuard, RwLockReadGuard, RwLockWriteGuard, WaitTimeoutResult}; +pub(crate) use parking_lot::WaitTimeoutResult; + +#[derive(Debug)] +pub(crate) struct Mutex<T: ?Sized>(PhantomData<std::sync::Mutex<T>>, parking_lot::Mutex<T>); + +#[derive(Debug)] +pub(crate) struct RwLock<T>(PhantomData<std::sync::RwLock<T>>, parking_lot::RwLock<T>); + +#[derive(Debug)] +pub(crate) struct Condvar(PhantomData<std::sync::Condvar>, parking_lot::Condvar); -/// Adapter for `parking_lot::Mutex` to the `std::sync::Mutex` interface. #[derive(Debug)] -pub(crate) struct Mutex<T: ?Sized>(parking_lot::Mutex<T>); +pub(crate) struct MutexGuard<'a, T: ?Sized>( + PhantomData<std::sync::MutexGuard<'a, T>>, + parking_lot::MutexGuard<'a, T>, +); #[derive(Debug)] -pub(crate) struct RwLock<T>(parking_lot::RwLock<T>); +pub(crate) struct RwLockReadGuard<'a, T: ?Sized>( + PhantomData<std::sync::RwLockReadGuard<'a, T>>, + parking_lot::RwLockReadGuard<'a, T>, +); -/// Adapter for `parking_lot::Condvar` to the `std::sync::Condvar` interface. #[derive(Debug)] -pub(crate) struct Condvar(parking_lot::Condvar); +pub(crate) struct RwLockWriteGuard<'a, T: ?Sized>( + PhantomData<std::sync::RwLockWriteGuard<'a, T>>, + parking_lot::RwLockWriteGuard<'a, T>, +); impl<T> Mutex<T> { #[inline] pub(crate) fn new(t: T) -> Mutex<T> { - Mutex(parking_lot::Mutex::new(t)) + Mutex(PhantomData, parking_lot::Mutex::new(t)) } #[inline] - #[cfg(all(feature = "parking_lot", not(all(loom, test)),))] + #[cfg(all(feature = "parking_lot", not(all(loom, test))))] #[cfg_attr(docsrs, doc(cfg(all(feature = "parking_lot",))))] pub(crate) const fn const_new(t: T) -> Mutex<T> { - Mutex(parking_lot::const_mutex(t)) + Mutex(PhantomData, parking_lot::const_mutex(t)) } #[inline] pub(crate) fn lock(&self) -> MutexGuard<'_, T> { - self.0.lock() + MutexGuard(PhantomData, self.1.lock()) } #[inline] pub(crate) fn try_lock(&self) -> Option<MutexGuard<'_, T>> { - self.0.try_lock() + self.1 + .try_lock() + .map(|guard| MutexGuard(PhantomData, guard)) } #[inline] pub(crate) fn get_mut(&mut self) -> &mut T { - self.0.get_mut() + self.1.get_mut() } // Note: Additional methods `is_poisoned` and `into_inner`, can be // provided here as needed. } +impl<'a, T: ?Sized> Deref for MutexGuard<'a, T> { + type Target = T; + fn deref(&self) -> &T { + self.1.deref() + } +} + +impl<'a, T: ?Sized> DerefMut for MutexGuard<'a, T> { + fn deref_mut(&mut self) -> &mut T { + self.1.deref_mut() + } +} + impl<T> RwLock<T> { pub(crate) fn new(t: T) -> RwLock<T> { - RwLock(parking_lot::RwLock::new(t)) + RwLock(PhantomData, parking_lot::RwLock::new(t)) } pub(crate) fn read(&self) -> LockResult<RwLockReadGuard<'_, T>> { - Ok(self.0.read()) + Ok(RwLockReadGuard(PhantomData, self.1.read())) } pub(crate) fn write(&self) -> LockResult<RwLockWriteGuard<'_, T>> { - Ok(self.0.write()) + Ok(RwLockWriteGuard(PhantomData, self.1.write())) + } +} + +impl<'a, T: ?Sized> Deref for RwLockReadGuard<'a, T> { + type Target = T; + fn deref(&self) -> &T { + self.1.deref() + } +} + +impl<'a, T: ?Sized> Deref for RwLockWriteGuard<'a, T> { + type Target = T; + fn deref(&self) -> &T { + self.1.deref() + } +} + +impl<'a, T: ?Sized> DerefMut for RwLockWriteGuard<'a, T> { + fn deref_mut(&mut self) -> &mut T { + self.1.deref_mut() } } impl Condvar { #[inline] pub(crate) fn new() -> Condvar { - Condvar(parking_lot::Condvar::new()) + Condvar(PhantomData, parking_lot::Condvar::new()) } #[inline] pub(crate) fn notify_one(&self) { - self.0.notify_one(); + self.1.notify_one(); } #[inline] pub(crate) fn notify_all(&self) { - self.0.notify_all(); + self.1.notify_all(); } #[inline] @@ -87,7 +147,7 @@ impl Condvar { &self, mut guard: MutexGuard<'a, T>, ) -> LockResult<MutexGuard<'a, T>> { - self.0.wait(&mut guard); + self.1.wait(&mut guard.1); Ok(guard) } @@ -97,10 +157,28 @@ impl Condvar { mut guard: MutexGuard<'a, T>, timeout: Duration, ) -> LockResult<(MutexGuard<'a, T>, WaitTimeoutResult)> { - let wtr = self.0.wait_for(&mut guard, timeout); + let wtr = self.1.wait_for(&mut guard.1, timeout); Ok((guard, wtr)) } // Note: Additional methods `wait_timeout_ms`, `wait_timeout_until`, // `wait_until` can be provided here as needed. } + +impl<'a, T: ?Sized + fmt::Display> fmt::Display for MutexGuard<'a, T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Display::fmt(&self.1, f) + } +} + +impl<'a, T: ?Sized + fmt::Display> fmt::Display for RwLockReadGuard<'a, T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Display::fmt(&self.1, f) + } +} + +impl<'a, T: ?Sized + fmt::Display> fmt::Display for RwLockWriteGuard<'a, T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Display::fmt(&self.1, f) + } +} diff --git a/vendor/tokio/src/macros/addr_of.rs b/vendor/tokio/src/macros/addr_of.rs new file mode 100644 index 000000000..9d28158f2 --- /dev/null +++ b/vendor/tokio/src/macros/addr_of.rs @@ -0,0 +1,22 @@ +//! This module defines a macro that lets you go from a raw pointer to a struct +//! to a raw pointer to a field of the struct. + +macro_rules! generate_addr_of_methods { + ( + impl<$($gen:ident)*> $struct_name:ty {$( + $(#[$attrs:meta])* + $vis:vis unsafe fn $fn_name:ident(self: NonNull<Self>) -> NonNull<$field_type:ty> { + &self$(.$field_name:tt)+ + } + )*} + ) => { + impl<$($gen)*> $struct_name {$( + $(#[$attrs])* + $vis unsafe fn $fn_name(me: ::core::ptr::NonNull<Self>) -> ::core::ptr::NonNull<$field_type> { + let me = me.as_ptr(); + let field = ::std::ptr::addr_of_mut!((*me) $(.$field_name)+ ); + ::core::ptr::NonNull::new_unchecked(field) + } + )*} + }; +} diff --git a/vendor/tokio/src/macros/cfg.rs b/vendor/tokio/src/macros/cfg.rs index 1e77556d8..52ffc102b 100644 --- a/vendor/tokio/src/macros/cfg.rs +++ b/vendor/tokio/src/macros/cfg.rs @@ -13,7 +13,19 @@ macro_rules! feature { } } -/// Enables enter::block_on +/// Enables Windows-specific code. +/// Use this macro instead of `cfg(windows)` to generate docs properly. +macro_rules! cfg_windows { + ($($item:item)*) => { + $( + #[cfg(any(all(doc, docsrs), windows))] + #[cfg_attr(docsrs, doc(cfg(windows)))] + $item + )* + } +} + +/// Enables enter::block_on. macro_rules! cfg_block_on { ($($item:item)*) => { $( @@ -28,7 +40,7 @@ macro_rules! cfg_block_on { } } -/// Enables internal `AtomicWaker` impl +/// Enables internal `AtomicWaker` impl. macro_rules! cfg_atomic_waker_impl { ($($item:item)*) => { $( @@ -45,10 +57,23 @@ macro_rules! cfg_atomic_waker_impl { } } +macro_rules! cfg_aio { + ($($item:item)*) => { + $( + #[cfg(all(any(docsrs, target_os = "freebsd"), feature = "net"))] + #[cfg_attr(docsrs, + doc(cfg(all(target_os = "freebsd", feature = "net"))) + )] + $item + )* + } +} + macro_rules! cfg_fs { ($($item:item)*) => { $( #[cfg(feature = "fs")] + #[cfg(not(tokio_wasi))] #[cfg_attr(docsrs, doc(cfg(feature = "fs")))] $item )* @@ -57,7 +82,11 @@ macro_rules! cfg_fs { macro_rules! cfg_io_blocking { ($($item:item)*) => { - $( #[cfg(any(feature = "io-std", feature = "fs"))] $item )* + $( #[cfg(any( + feature = "io-std", + feature = "fs", + all(windows, feature = "process"), + ))] $item )* } } @@ -66,12 +95,12 @@ macro_rules! cfg_io_driver { $( #[cfg(any( feature = "net", - feature = "process", + all(unix, feature = "process"), all(unix, feature = "signal"), ))] #[cfg_attr(docsrs, doc(cfg(any( feature = "net", - feature = "process", + all(unix, feature = "process"), all(unix, feature = "signal"), ))))] $item @@ -84,7 +113,7 @@ macro_rules! cfg_io_driver_impl { $( #[cfg(any( feature = "net", - feature = "process", + all(unix, feature = "process"), all(unix, feature = "signal"), ))] $item @@ -97,7 +126,7 @@ macro_rules! cfg_not_io_driver { $( #[cfg(not(any( feature = "net", - feature = "process", + all(unix, feature = "process"), all(unix, feature = "signal"), )))] $item @@ -162,6 +191,43 @@ macro_rules! cfg_macros { } } +macro_rules! cfg_metrics { + ($($item:item)*) => { + $( + // For now, metrics is only disabled in loom tests. + // When stabilized, it might have a dedicated feature flag. + #[cfg(all(tokio_unstable, not(loom)))] + #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))] + $item + )* + } +} + +macro_rules! cfg_not_metrics { + ($($item:item)*) => { + $( + #[cfg(not(all(tokio_unstable, not(loom))))] + $item + )* + } +} + +macro_rules! cfg_not_rt_and_metrics_and_net { + ($($item:item)*) => { + $( #[cfg(not(all(feature = "net", feature = "rt", all(tokio_unstable, not(loom)))))]$item )* + } +} + +macro_rules! cfg_net_or_process { + ($($item:item)*) => { + $( + #[cfg(any(feature = "net", feature = "process"))] + #[cfg_attr(docsrs, doc(cfg(any(feature = "net", feature = "process"))))] + $item + )* + } +} + macro_rules! cfg_net { ($($item:item)*) => { $( @@ -176,7 +242,7 @@ macro_rules! cfg_net_unix { ($($item:item)*) => { $( #[cfg(all(unix, feature = "net"))] - #[cfg_attr(docsrs, doc(cfg(feature = "net")))] + #[cfg_attr(docsrs, doc(cfg(all(unix, feature = "net"))))] $item )* } @@ -185,7 +251,7 @@ macro_rules! cfg_net_unix { macro_rules! cfg_net_windows { ($($item:item)*) => { $( - #[cfg(all(any(docsrs, windows), feature = "net"))] + #[cfg(all(any(all(doc, docsrs), windows), feature = "net"))] #[cfg_attr(docsrs, doc(cfg(all(windows, feature = "net"))))] $item )* @@ -198,6 +264,7 @@ macro_rules! cfg_process { #[cfg(feature = "process")] #[cfg_attr(docsrs, doc(cfg(feature = "process")))] #[cfg(not(loom))] + #[cfg(not(tokio_wasi))] $item )* } @@ -226,6 +293,7 @@ macro_rules! cfg_signal { #[cfg(feature = "signal")] #[cfg_attr(docsrs, doc(cfg(feature = "signal")))] #[cfg(not(loom))] + #[cfg(not(tokio_wasi))] $item )* } @@ -241,6 +309,13 @@ macro_rules! cfg_signal_internal { } } +macro_rules! cfg_signal_internal_and_unix { + ($($item:item)*) => { + #[cfg(unix)] + cfg_signal_internal! { $($item)* } + } +} + macro_rules! cfg_not_signal_internal { ($($item:item)*) => { $( @@ -285,7 +360,7 @@ macro_rules! cfg_not_rt { macro_rules! cfg_rt_multi_thread { ($($item:item)*) => { $( - #[cfg(feature = "rt-multi-thread")] + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] #[cfg_attr(docsrs, doc(cfg(feature = "rt-multi-thread")))] $item )* @@ -298,6 +373,44 @@ macro_rules! cfg_not_rt_multi_thread { } } +macro_rules! cfg_taskdump { + ($($item:item)*) => { + $( + #[cfg(all( + tokio_unstable, + tokio_taskdump, + feature = "rt", + target_os = "linux", + any( + target_arch = "aarch64", + target_arch = "x86", + target_arch = "x86_64" + ) + ))] + $item + )* + }; +} + +macro_rules! cfg_not_taskdump { + ($($item:item)*) => { + $( + #[cfg(not(all( + tokio_unstable, + tokio_taskdump, + feature = "rt", + target_os = "linux", + any( + target_arch = "aarch64", + target_arch = "x86", + target_arch = "x86_64" + ) + )))] + $item + )* + }; +} + macro_rules! cfg_test_util { ($($item:item)*) => { $( @@ -334,10 +447,20 @@ macro_rules! cfg_trace { ($($item:item)*) => { $( #[cfg(all(tokio_unstable, feature = "tracing"))] - #[cfg_attr(docsrs, doc(cfg(feature = "tracing")))] + #[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "tracing"))))] $item )* - } + }; +} + +macro_rules! cfg_unstable { + ($($item:item)*) => { + $( + #[cfg(tokio_unstable)] + #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))] + $item + )* + }; } macro_rules! cfg_not_trace { @@ -384,3 +507,83 @@ macro_rules! cfg_not_coop { )* } } + +macro_rules! cfg_has_atomic_u64 { + ($($item:item)*) => { + $( + #[cfg_attr( + not(tokio_no_target_has_atomic), + cfg(all(target_has_atomic = "64", not(tokio_no_atomic_u64)) + ))] + #[cfg_attr( + tokio_no_target_has_atomic, + cfg(not(tokio_no_atomic_u64)) + )] + $item + )* + } +} + +macro_rules! cfg_not_has_atomic_u64 { + ($($item:item)*) => { + $( + #[cfg_attr( + not(tokio_no_target_has_atomic), + cfg(any(not(target_has_atomic = "64"), tokio_no_atomic_u64) + ))] + #[cfg_attr( + tokio_no_target_has_atomic, + cfg(tokio_no_atomic_u64) + )] + $item + )* + } +} + +macro_rules! cfg_has_const_mutex_new { + ($($item:item)*) => { + $( + #[cfg(all( + not(all(loom, test)), + any( + feature = "parking_lot", + not(tokio_no_const_mutex_new) + ) + ))] + $item + )* + } +} + +macro_rules! cfg_not_has_const_mutex_new { + ($($item:item)*) => { + $( + #[cfg(not(all( + not(all(loom, test)), + any( + feature = "parking_lot", + not(tokio_no_const_mutex_new) + ) + )))] + $item + )* + } +} + +macro_rules! cfg_not_wasi { + ($($item:item)*) => { + $( + #[cfg(not(tokio_wasi))] + $item + )* + } +} + +macro_rules! cfg_is_wasm_not_wasi { + ($($item:item)*) => { + $( + #[cfg(tokio_wasm_not_wasi)] + $item + )* + } +} diff --git a/vendor/tokio/src/macros/join.rs b/vendor/tokio/src/macros/join.rs index 5f37af510..8a0198600 100644 --- a/vendor/tokio/src/macros/join.rs +++ b/vendor/tokio/src/macros/join.rs @@ -1,4 +1,4 @@ -/// Wait on multiple concurrent branches, returning when **all** branches +/// Waits on multiple concurrent branches, returning when **all** branches /// complete. /// /// The `join!` macro must be used inside of async functions, closures, and @@ -12,7 +12,7 @@ /// for **all** branches complete regardless if any complete with `Err`. Use /// [`try_join!`] to return early when `Err` is encountered. /// -/// [`try_join!`]: macro@try_join +/// [`try_join!`]: crate::try_join /// /// # Notes /// @@ -60,6 +60,9 @@ macro_rules! join { // normalization is complete. ( $($count:tt)* ) + // The expression `0+1+1+ ... +1` equal to the number of branches. + ( $($total:tt)* ) + // Normalized join! branches $( ( $($skip:tt)* ) $e:expr, )* @@ -69,24 +72,66 @@ 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 + // reference to the futures, instead of taking ownership of them. This + // mitigates the issue described in + // <https://internals.rust-lang.org/t/surprising-soundness-trouble-around-pollfn/17484> + 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. + // + // This is number of futures that will be skipped in the first loop + // iteration the next time. + let mut skip_next_time: u32 = 0; + poll_fn(move |cx| { + const COUNT: u32 = $($total)*; + let mut is_pending = false; + let mut to_run = COUNT; + + // The number of futures that will be skipped in the first loop iteration. + let mut skip = skip_next_time; + + skip_next_time = if skip + 1 == COUNT { 0 } else { skip + 1 }; + + // This loop runs twice and the first `skip` futures + // are not polled in the first iteration. + loop { $( - // Extract the future for this branch from the tuple. - let ( $($skip,)* fut, .. ) = &mut futures; + if skip == 0 { + if to_run == 0 { + // Every future has been polled + break; + } + to_run -= 1; - // Safety: future is stored on the stack above - // and never moved. - let mut fut = unsafe { Pin::new_unchecked(fut) }; + // Extract the future for this branch from the tuple. + let ( $($skip,)* fut, .. ) = &mut *futures; - // Try polling - if fut.poll(cx).is_pending() { - is_pending = true; + // Safety: future is stored on the stack above + // and never moved. + let mut fut = unsafe { Pin::new_unchecked(fut) }; + + // Try polling + if fut.poll(cx).is_pending() { + is_pending = true; + } + } else { + // Future skipped, one less future to skip in the next iteration + skip -= 1; } )* + } if is_pending { Pending @@ -107,13 +152,15 @@ macro_rules! join { // ===== Normalize ===== - (@ { ( $($s:tt)* ) $($t:tt)* } $e:expr, $($r:tt)* ) => { - $crate::join!(@{ ($($s)* _) $($t)* ($($s)*) $e, } $($r)*) + (@ { ( $($s:tt)* ) ( $($n:tt)* ) $($t:tt)* } $e:expr, $($r:tt)* ) => { + $crate::join!(@{ ($($s)* _) ($($n)* + 1) $($t)* ($($s)*) $e, } $($r)*) }; // ===== Entry point ===== - ( $($e:expr),* $(,)?) => { - $crate::join!(@{ () } $($e,)*) + ( $($e:expr),+ $(,)?) => { + $crate::join!(@{ () (0) } $($e,)*) }; + + () => { async {}.await } } diff --git a/vendor/tokio/src/macros/loom.rs b/vendor/tokio/src/macros/loom.rs index d57d9fb0f..fa2653ce7 100644 --- a/vendor/tokio/src/macros/loom.rs +++ b/vendor/tokio/src/macros/loom.rs @@ -1,11 +1,7 @@ macro_rules! if_loom { ($($t:tt)*) => {{ #[cfg(loom)] - const LOOM: bool = true; - #[cfg(not(loom))] - const LOOM: bool = false; - - if LOOM { + { $($t)* } }} diff --git a/vendor/tokio/src/macros/mod.rs b/vendor/tokio/src/macros/mod.rs index b0af52152..82f42dbff 100644 --- a/vendor/tokio/src/macros/mod.rs +++ b/vendor/tokio/src/macros/mod.rs @@ -16,8 +16,12 @@ mod ready; mod thread_local; #[macro_use] -#[cfg(feature = "rt")] -pub(crate) mod scoped_tls; +mod addr_of; + +cfg_trace! { + #[macro_use] + mod trace; +} cfg_macros! { #[macro_use] diff --git a/vendor/tokio/src/macros/scoped_tls.rs b/vendor/tokio/src/macros/scoped_tls.rs deleted file mode 100644 index a00aae2fb..000000000 --- a/vendor/tokio/src/macros/scoped_tls.rs +++ /dev/null @@ -1,77 +0,0 @@ -use crate::loom::thread::LocalKey; - -use std::cell::Cell; -use std::marker; - -/// Set a reference as a thread-local -macro_rules! scoped_thread_local { - ($(#[$attrs:meta])* $vis:vis static $name:ident: $ty:ty) => ( - $(#[$attrs])* - $vis static $name: $crate::macros::scoped_tls::ScopedKey<$ty> - = $crate::macros::scoped_tls::ScopedKey { - inner: { - thread_local!(static FOO: ::std::cell::Cell<*const ()> = { - std::cell::Cell::new(::std::ptr::null()) - }); - &FOO - }, - _marker: ::std::marker::PhantomData, - }; - ) -} - -/// Type representing a thread local storage key corresponding to a reference -/// to the type parameter `T`. -pub(crate) struct ScopedKey<T> { - pub(crate) inner: &'static LocalKey<Cell<*const ()>>, - pub(crate) _marker: marker::PhantomData<T>, -} - -unsafe impl<T> Sync for ScopedKey<T> {} - -impl<T> ScopedKey<T> { - /// Inserts a value into this scoped thread local storage slot for a - /// duration of a closure. - pub(crate) fn set<F, R>(&'static self, t: &T, f: F) -> R - where - F: FnOnce() -> R, - { - struct Reset { - key: &'static LocalKey<Cell<*const ()>>, - val: *const (), - } - - impl Drop for Reset { - fn drop(&mut self) { - self.key.with(|c| c.set(self.val)); - } - } - - let prev = self.inner.with(|c| { - let prev = c.get(); - c.set(t as *const _ as *const ()); - prev - }); - - let _reset = Reset { - key: self.inner, - val: prev, - }; - - f() - } - - /// Gets a value out of this scoped variable. - pub(crate) fn with<F, R>(&'static self, f: F) -> R - where - F: FnOnce(Option<&T>) -> R, - { - let val = self.inner.with(|c| c.get()); - - if val.is_null() { - f(None) - } else { - unsafe { f(Some(&*(val as *const T))) } - } - } -} diff --git a/vendor/tokio/src/macros/select.rs b/vendor/tokio/src/macros/select.rs index a90ee9eb5..31c9b3ac2 100644 --- a/vendor/tokio/src/macros/select.rs +++ b/vendor/tokio/src/macros/select.rs @@ -1,4 +1,4 @@ -/// Wait on multiple concurrent branches, returning when the **first** branch +/// Waits on multiple concurrent branches, returning when the **first** branch /// completes, cancelling the remaining branches. /// /// The `select!` macro must be used inside of async functions, closures, and @@ -101,6 +101,7 @@ /// * [`tokio::sync::watch::Receiver::changed`](crate::sync::watch::Receiver::changed) /// * [`tokio::net::TcpListener::accept`](crate::net::TcpListener::accept) /// * [`tokio::net::UnixListener::accept`](crate::net::UnixListener::accept) +/// * [`tokio::signal::unix::Signal::recv`](crate::signal::unix::Signal::recv) /// * [`tokio::io::AsyncReadExt::read`](crate::io::AsyncReadExt::read) on any `AsyncRead` /// * [`tokio::io::AsyncReadExt::read_buf`](crate::io::AsyncReadExt::read_buf) on any `AsyncRead` /// * [`tokio::io::AsyncWriteExt::write`](crate::io::AsyncWriteExt::write) on any `AsyncWrite` @@ -130,6 +131,13 @@ /// correctly even if it is restarted while waiting at an `.await`, then it is /// cancellation safe. /// +/// Cancellation safety can be defined in the following way: If you have a +/// future that has not yet completed, then it must be a no-op to drop that +/// future and recreate it. This definition is motivated by the situation where +/// a `select!` is used in a loop. Without this guarantee, you would lose your +/// progress when another branch completes and you restart the `select!` by +/// going around the loop. +/// /// Be aware that cancelling something that is not cancellation safe is not /// necessarily wrong. For example, if you are cancelling a task because the /// application is shutting down, then you probably don't care that partially @@ -429,7 +437,8 @@ macro_rules! select { // // This module is defined within a scope and should not leak out of this // macro. - mod util { + #[doc(hidden)] + mod __tokio_select_util { // Generate an enum with one variant per select branch $crate::select_priv_declare_output_enum!( ( $($count)* ) ); } @@ -442,13 +451,13 @@ macro_rules! select { const BRANCHES: u32 = $crate::count!( $($count)* ); - let mut disabled: util::Mask = Default::default(); + let mut disabled: __tokio_select_util::Mask = Default::default(); // First, invoke all the pre-conditions. For any that return true, // set the appropriate bit in `disabled`. $( if !$c { - let mask: util::Mask = 1 << $crate::count!( $($skip)* ); + let mask: __tokio_select_util::Mask = 1 << $crate::count!( $($skip)* ); disabled |= mask; } )* @@ -458,8 +467,18 @@ 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 + // reference to the futures, instead of taking ownership of them. + // This mitigates the issue described in + // <https://internals.rust-lang.org/t/surprising-soundness-trouble-around-pollfn/17484> + 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 @@ -495,14 +514,14 @@ 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. let mut fut = unsafe { Pin::new_unchecked(fut) }; // Try polling it - let out = match fut.poll(cx) { + let out = match Future::poll(fut, cx) { Ready(out) => out, Pending => { // Track that at least one future is @@ -520,12 +539,12 @@ macro_rules! select { #[allow(unused_variables)] #[allow(unused_mut)] match &out { - $bind => {} + $crate::select_priv_clean_pattern!($bind) => {} _ => continue, } // The select is complete, return the value - return Ready($crate::select_variant!(util::Out, ($($skip)*))(out)); + return Ready($crate::select_variant!(__tokio_select_util::Out, ($($skip)*))(out)); } )* _ => unreachable!("reaching this means there probably is an off by one bug"), @@ -536,16 +555,16 @@ macro_rules! select { Pending } else { // All branches have been disabled. - Ready(util::Out::Disabled) + Ready(__tokio_select_util::Out::Disabled) } }).await }; match output { $( - $crate::select_variant!(util::Out, ($($skip)*) ($bind)) => $handle, + $crate::select_variant!(__tokio_select_util::Out, ($($skip)*) ($bind)) => $handle, )* - util::Out::Disabled => $else, + __tokio_select_util::Out::Disabled => $else, _ => unreachable!("failed to match bind"), } }}; @@ -801,6 +820,9 @@ macro_rules! count { (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 63 }; + (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { + 64 + }; } #[macro_export] diff --git a/vendor/tokio/src/macros/support.rs b/vendor/tokio/src/macros/support.rs index 7f11bc680..10526bcbc 100644 --- a/vendor/tokio/src/macros/support.rs +++ b/vendor/tokio/src/macros/support.rs @@ -1,7 +1,11 @@ cfg_macros! { pub use crate::future::poll_fn; pub use crate::future::maybe_done::maybe_done; - pub use crate::util::thread_rng_n; + + #[doc(hidden)] + pub fn thread_rng_n(n: u32) -> u32 { + crate::runtime::context::thread_rng_n(n) + } } pub use std::future::Future; diff --git a/vendor/tokio/src/macros/thread_local.rs b/vendor/tokio/src/macros/thread_local.rs index d84894735..74be99abf 100644 --- a/vendor/tokio/src/macros/thread_local.rs +++ b/vendor/tokio/src/macros/thread_local.rs @@ -1,4 +1,32 @@ #[cfg(all(loom, test))] -macro_rules! thread_local { +macro_rules! tokio_thread_local { + ($(#[$attrs:meta])* $vis:vis static $name:ident: $ty:ty = const { $expr:expr } $(;)?) => { + loom::thread_local! { + $(#[$attrs])* + $vis static $name: $ty = $expr; + } + }; + ($($tts:tt)+) => { loom::thread_local!{ $($tts)+ } } } + +#[cfg(not(tokio_no_const_thread_local))] +#[cfg(not(all(loom, test)))] +macro_rules! tokio_thread_local { + ($($tts:tt)+) => { + ::std::thread_local!{ $($tts)+ } + } +} + +#[cfg(tokio_no_const_thread_local)] +#[cfg(not(all(loom, test)))] +macro_rules! tokio_thread_local { + ($(#[$attrs:meta])* $vis:vis static $name:ident: $ty:ty = const { $expr:expr } $(;)?) => { + ::std::thread_local! { + $(#[$attrs])* + $vis static $name: $ty = $expr; + } + }; + + ($($tts:tt)+) => { ::std::thread_local!{ $($tts)+ } } +} diff --git a/vendor/tokio/src/macros/trace.rs b/vendor/tokio/src/macros/trace.rs new file mode 100644 index 000000000..80a257e18 --- /dev/null +++ b/vendor/tokio/src/macros/trace.rs @@ -0,0 +1,26 @@ +cfg_trace! { + macro_rules! trace_op { + ($name:expr, $readiness:literal) => { + tracing::trace!( + target: "runtime::resource::poll_op", + op_name = $name, + is_ready = $readiness + ); + } + } + + macro_rules! trace_poll_op { + ($name:expr, $poll:expr $(,)*) => { + match $poll { + std::task::Poll::Ready(t) => { + trace_op!($name, true); + std::task::Poll::Ready(t) + } + std::task::Poll::Pending => { + trace_op!($name, false); + return std::task::Poll::Pending; + } + } + }; + } +} diff --git a/vendor/tokio/src/macros/try_join.rs b/vendor/tokio/src/macros/try_join.rs index fa5850ef0..7b1237092 100644 --- a/vendor/tokio/src/macros/try_join.rs +++ b/vendor/tokio/src/macros/try_join.rs @@ -1,4 +1,4 @@ -/// Wait on multiple concurrent branches, returning when **all** branches +/// Waits on multiple concurrent branches, returning when **all** branches /// complete with `Ok(_)` or on the first `Err(_)`. /// /// The `try_join!` macro must be used inside of async functions, closures, and @@ -59,6 +59,45 @@ /// } /// } /// ``` +/// +/// Using `try_join!` with spawned tasks. +/// +/// ``` +/// use tokio::task::JoinHandle; +/// +/// async fn do_stuff_async() -> Result<(), &'static str> { +/// // async work +/// # Err("failed") +/// } +/// +/// async fn more_async_work() -> Result<(), &'static str> { +/// // more here +/// # Ok(()) +/// } +/// +/// async fn flatten<T>(handle: JoinHandle<Result<T, &'static str>>) -> Result<T, &'static str> { +/// match handle.await { +/// Ok(Ok(result)) => Ok(result), +/// Ok(Err(err)) => Err(err), +/// Err(err) => Err("handling failed"), +/// } +/// } +/// +/// #[tokio::main] +/// async fn main() { +/// let handle1 = tokio::spawn(do_stuff_async()); +/// let handle2 = tokio::spawn(more_async_work()); +/// match tokio::try_join!(flatten(handle1), flatten(handle2)) { +/// Ok(val) => { +/// // do something with the values +/// } +/// Err(err) => { +/// println!("Failed with {}.", err); +/// # assert_eq!(err, "failed"); +/// } +/// } +/// } +/// ``` #[macro_export] #[cfg_attr(docsrs, doc(cfg(feature = "macros")))] macro_rules! try_join { @@ -67,6 +106,9 @@ macro_rules! try_join { // normalization is complete. ( $($count:tt)* ) + // The expression `0+1+1+ ... +1` equal to the number of branches. + ( $($total:tt)* ) + // Normalized try_join! branches $( ( $($skip:tt)* ) $e:expr, )* @@ -76,26 +118,68 @@ 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 + // reference to the futures, instead of taking ownership of them. This + // mitigates the issue described in + // <https://internals.rust-lang.org/t/surprising-soundness-trouble-around-pollfn/17484> + 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. + // + // This is number of futures that will be skipped in the first loop + // iteration the next time. + let mut skip_next_time: u32 = 0; + poll_fn(move |cx| { + const COUNT: u32 = $($total)*; + let mut is_pending = false; + let mut to_run = COUNT; + + // The number of futures that will be skipped in the first loop iteration + let mut skip = skip_next_time; + + skip_next_time = if skip + 1 == COUNT { 0 } else { skip + 1 }; + + // This loop runs twice and the first `skip` futures + // are not polled in the first iteration. + loop { $( - // Extract the future for this branch from the tuple. - let ( $($skip,)* fut, .. ) = &mut futures; - - // Safety: future is stored on the stack above - // and never moved. - let mut fut = unsafe { Pin::new_unchecked(fut) }; - - // Try polling - if fut.as_mut().poll(cx).is_pending() { - is_pending = true; - } else if fut.as_mut().output_mut().expect("expected completed future").is_err() { - return Ready(Err(fut.take_output().expect("expected completed future").err().unwrap())) + if skip == 0 { + if to_run == 0 { + // Every future has been polled + break; + } + to_run -= 1; + + // Extract the future for this branch from the tuple. + let ( $($skip,)* fut, .. ) = &mut *futures; + + // Safety: future is stored on the stack above + // and never moved. + let mut fut = unsafe { Pin::new_unchecked(fut) }; + + // Try polling + if fut.as_mut().poll(cx).is_pending() { + is_pending = true; + } else if fut.as_mut().output_mut().expect("expected completed future").is_err() { + return Ready(Err(fut.take_output().expect("expected completed future").err().unwrap())) + } + } else { + // Future skipped, one less future to skip in the next iteration + skip -= 1; } )* + } if is_pending { Pending @@ -120,13 +204,15 @@ macro_rules! try_join { // ===== Normalize ===== - (@ { ( $($s:tt)* ) $($t:tt)* } $e:expr, $($r:tt)* ) => { - $crate::try_join!(@{ ($($s)* _) $($t)* ($($s)*) $e, } $($r)*) + (@ { ( $($s:tt)* ) ( $($n:tt)* ) $($t:tt)* } $e:expr, $($r:tt)* ) => { + $crate::try_join!(@{ ($($s)* _) ($($n)* + 1) $($t)* ($($s)*) $e, } $($r)*) }; // ===== Entry point ===== - ( $($e:expr),* $(,)?) => { - $crate::try_join!(@{ () } $($e,)*) + ( $($e:expr),+ $(,)?) => { + $crate::try_join!(@{ () (0) } $($e,)*) }; + + () => { async { Ok(()) }.await } } diff --git a/vendor/tokio/src/net/addr.rs b/vendor/tokio/src/net/addr.rs index ec4fa198e..fb8248fb3 100644 --- a/vendor/tokio/src/net/addr.rs +++ b/vendor/tokio/src/net/addr.rs @@ -1,5 +1,4 @@ -use crate::future; - +use std::future; use std::io; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6}; @@ -56,7 +55,7 @@ impl sealed::ToSocketAddrsPriv for SocketAddr { fn to_socket_addrs(&self, _: sealed::Internal) -> Self::Future { let iter = Some(*self).into_iter(); - future::ok(iter) + future::ready(Ok(iter)) } } @@ -96,7 +95,7 @@ impl sealed::ToSocketAddrsPriv for (IpAddr, u16) { fn to_socket_addrs(&self, _: sealed::Internal) -> Self::Future { let iter = Some(SocketAddr::from(*self)).into_iter(); - future::ok(iter) + future::ready(Ok(iter)) } } @@ -137,8 +136,23 @@ impl sealed::ToSocketAddrsPriv for &[SocketAddr] { type Future = ReadyFuture<Self::Iter>; fn to_socket_addrs(&self, _: sealed::Internal) -> Self::Future { - let iter = self.to_vec().into_iter(); - future::ok(iter) + #[inline] + fn slice_to_vec(addrs: &[SocketAddr]) -> Vec<SocketAddr> { + addrs.to_vec() + } + + // This uses a helper method because clippy doesn't like the `to_vec()` + // call here (it will allocate, whereas `self.iter().copied()` would + // not), but it's actually necessary in order to ensure that the + // returned iterator is valid for the `'static` lifetime, which the + // borrowed `slice::Iter` iterator would not be. + // + // Note that we can't actually add an `allow` attribute for + // `clippy::unnecessary_to_owned` here, as Tokio's CI runs clippy lints + // on Rust 1.52 to avoid breaking LTS releases of Tokio. Users of newer + // Rust versions who see this lint should just ignore it. + let iter = slice_to_vec(self).into_iter(); + future::ready(Ok(iter)) } } @@ -230,7 +244,7 @@ cfg_net! { type Future = <str as sealed::ToSocketAddrsPriv>::Future; fn to_socket_addrs(&self, _: sealed::Internal) -> Self::Future { - (&self[..]).to_socket_addrs(sealed::Internal) + self[..].to_socket_addrs(sealed::Internal) } } } diff --git a/vendor/tokio/src/net/mod.rs b/vendor/tokio/src/net/mod.rs index 0b8c1ecd1..2d317a8a2 100644 --- a/vendor/tokio/src/net/mod.rs +++ b/vendor/tokio/src/net/mod.rs @@ -23,8 +23,10 @@ //! [`UnixDatagram`]: UnixDatagram mod addr; -#[cfg(feature = "net")] -pub(crate) use addr::to_socket_addrs; +cfg_not_wasi! { + #[cfg(feature = "net")] + pub(crate) use addr::to_socket_addrs; +} pub use addr::ToSocketAddrs; cfg_net! { @@ -33,11 +35,13 @@ cfg_net! { pub mod tcp; pub use tcp::listener::TcpListener; - pub use tcp::socket::TcpSocket; pub use tcp::stream::TcpStream; + cfg_not_wasi! { + pub use tcp::socket::TcpSocket; - mod udp; - pub use udp::UdpSocket; + mod udp; + pub use udp::UdpSocket; + } } cfg_net_unix! { diff --git a/vendor/tokio/src/net/tcp/listener.rs b/vendor/tokio/src/net/tcp/listener.rs index 86f0ec1d2..34e393c89 100644 --- a/vendor/tokio/src/net/tcp/listener.rs +++ b/vendor/tokio/src/net/tcp/listener.rs @@ -1,8 +1,10 @@ use crate::io::{Interest, PollEvented}; use crate::net::tcp::TcpStream; -use crate::net::{to_socket_addrs, ToSocketAddrs}; -use std::convert::TryFrom; +cfg_not_wasi! { + use crate::net::{to_socket_addrs, ToSocketAddrs}; +} + use std::fmt; use std::io; use std::net::{self, SocketAddr}; @@ -55,68 +57,70 @@ cfg_net! { } impl TcpListener { - /// Creates a new TcpListener, which will be bound to the specified address. - /// - /// The returned listener is ready for accepting connections. - /// - /// Binding with a port number of 0 will request that the OS assigns a port - /// to this listener. The port allocated can be queried via the `local_addr` - /// method. - /// - /// The address type can be any implementor of the [`ToSocketAddrs`] trait. - /// If `addr` yields multiple addresses, bind will be attempted with each of - /// the addresses until one succeeds and returns the listener. If none of - /// the addresses succeed in creating a listener, the error returned from - /// the last attempt (the last address) is returned. - /// - /// This function sets the `SO_REUSEADDR` option on the socket. - /// - /// To configure the socket before binding, you can use the [`TcpSocket`] - /// type. - /// - /// [`ToSocketAddrs`]: trait@crate::net::ToSocketAddrs - /// [`TcpSocket`]: struct@crate::net::TcpSocket - /// - /// # Examples - /// - /// ```no_run - /// use tokio::net::TcpListener; - /// - /// use std::io; - /// - /// #[tokio::main] - /// async fn main() -> io::Result<()> { - /// let listener = TcpListener::bind("127.0.0.1:2345").await?; - /// - /// // use the listener - /// - /// # let _ = listener; - /// Ok(()) - /// } - /// ``` - pub async fn bind<A: ToSocketAddrs>(addr: A) -> io::Result<TcpListener> { - let addrs = to_socket_addrs(addr).await?; + cfg_not_wasi! { + /// Creates a new TcpListener, which will be bound to the specified address. + /// + /// The returned listener is ready for accepting connections. + /// + /// Binding with a port number of 0 will request that the OS assigns a port + /// to this listener. The port allocated can be queried via the `local_addr` + /// method. + /// + /// The address type can be any implementor of the [`ToSocketAddrs`] trait. + /// If `addr` yields multiple addresses, bind will be attempted with each of + /// the addresses until one succeeds and returns the listener. If none of + /// the addresses succeed in creating a listener, the error returned from + /// the last attempt (the last address) is returned. + /// + /// This function sets the `SO_REUSEADDR` option on the socket. + /// + /// To configure the socket before binding, you can use the [`TcpSocket`] + /// type. + /// + /// [`ToSocketAddrs`]: trait@crate::net::ToSocketAddrs + /// [`TcpSocket`]: struct@crate::net::TcpSocket + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpListener; + /// + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let listener = TcpListener::bind("127.0.0.1:2345").await?; + /// + /// // use the listener + /// + /// # let _ = listener; + /// Ok(()) + /// } + /// ``` + pub async fn bind<A: ToSocketAddrs>(addr: A) -> io::Result<TcpListener> { + let addrs = to_socket_addrs(addr).await?; - let mut last_err = None; + let mut last_err = None; - for addr in addrs { - match TcpListener::bind_addr(addr) { - Ok(listener) => return Ok(listener), - Err(e) => last_err = Some(e), + for addr in addrs { + match TcpListener::bind_addr(addr) { + Ok(listener) => return Ok(listener), + Err(e) => last_err = Some(e), + } } - } - Err(last_err.unwrap_or_else(|| { - io::Error::new( - io::ErrorKind::InvalidInput, - "could not resolve to any address", - ) - })) - } + Err(last_err.unwrap_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidInput, + "could not resolve to any address", + ) + })) + } - fn bind_addr(addr: SocketAddr) -> io::Result<TcpListener> { - let listener = mio::net::TcpListener::bind(addr)?; - TcpListener::new(listener) + fn bind_addr(addr: SocketAddr) -> io::Result<TcpListener> { + let listener = mio::net::TcpListener::bind(addr)?; + TcpListener::new(listener) + } } /// Accepts a new incoming connection from this listener. @@ -190,15 +194,22 @@ impl TcpListener { /// Creates new `TcpListener` from a `std::net::TcpListener`. /// /// This function is intended to be used to wrap a TCP listener from the - /// standard library in the Tokio equivalent. The conversion assumes nothing - /// about the underlying listener; it is left up to the user to set it in - /// non-blocking mode. + /// standard library in the Tokio equivalent. /// /// This API is typically paired with the `socket2` crate and the `Socket` /// type to build up and customize a listener before it's shipped off to the /// backing event loop. This allows configuration of options like /// `SO_REUSEPORT`, binding to multiple addresses, etc. /// + /// # Notes + /// + /// The caller is responsible for ensuring that the listener is in + /// non-blocking mode. Otherwise all I/O operations on the listener + /// will block the thread, which will cause unexpected behavior. + /// Non-blocking mode can be set using [`set_nonblocking`]. + /// + /// [`set_nonblocking`]: std::net::TcpListener::set_nonblocking + /// /// # Examples /// /// ```rust,no_run @@ -216,18 +227,20 @@ impl TcpListener { /// /// # Panics /// - /// This function panics if thread-local runtime is not set. + /// This function panics if it is not called from within a runtime with + /// IO enabled. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. + #[track_caller] pub fn from_std(listener: net::TcpListener) -> io::Result<TcpListener> { let io = mio::net::TcpListener::from_std(listener); let io = PollEvented::new(io)?; Ok(TcpListener { io }) } - /// Turn a [`tokio::net::TcpListener`] into a [`std::net::TcpListener`]. + /// Turns a [`tokio::net::TcpListener`] into a [`std::net::TcpListener`]. /// /// The returned [`std::net::TcpListener`] will have nonblocking mode set as /// `true`. Use [`set_nonblocking`] to change the blocking mode if needed. @@ -267,11 +280,22 @@ impl TcpListener { .map(|io| io.into_raw_socket()) .map(|raw_socket| unsafe { std::net::TcpListener::from_raw_socket(raw_socket) }) } + + #[cfg(tokio_wasi)] + { + use std::os::wasi::io::{FromRawFd, IntoRawFd}; + self.io + .into_inner() + .map(|io| io.into_raw_fd()) + .map(|raw_fd| unsafe { std::net::TcpListener::from_raw_fd(raw_fd) }) + } } - pub(crate) fn new(listener: mio::net::TcpListener) -> io::Result<TcpListener> { - let io = PollEvented::new(listener)?; - Ok(TcpListener { io }) + cfg_not_wasi! { + pub(crate) fn new(listener: mio::net::TcpListener) -> io::Result<TcpListener> { + let io = PollEvented::new(listener)?; + Ok(TcpListener { io }) + } } /// Returns the local address that this listener is bound to. @@ -382,16 +406,51 @@ mod sys { self.io.as_raw_fd() } } + + #[cfg(not(tokio_no_as_fd))] + impl AsFd for TcpListener { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } + } } -#[cfg(windows)] -mod sys { - use super::TcpListener; - use std::os::windows::prelude::*; +cfg_unstable! { + #[cfg(tokio_wasi)] + mod sys { + use super::TcpListener; + use std::os::wasi::prelude::*; + + impl AsRawFd for TcpListener { + fn as_raw_fd(&self) -> RawFd { + self.io.as_raw_fd() + } + } + + #[cfg(not(tokio_no_as_fd))] + impl AsFd for TcpListener { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } + } + } +} + +cfg_windows! { + use crate::os::windows::io::{AsRawSocket, RawSocket}; + #[cfg(not(tokio_no_as_fd))] + use crate::os::windows::io::{AsSocket, BorrowedSocket}; impl AsRawSocket for TcpListener { fn as_raw_socket(&self) -> RawSocket { self.io.as_raw_socket() } } + + #[cfg(not(tokio_no_as_fd))] + impl AsSocket for TcpListener { + fn as_socket(&self) -> BorrowedSocket<'_> { + unsafe { BorrowedSocket::borrow_raw(self.as_raw_socket()) } + } + } } diff --git a/vendor/tokio/src/net/tcp/mod.rs b/vendor/tokio/src/net/tcp/mod.rs index 7f0f6d914..734eabe6d 100644 --- a/vendor/tokio/src/net/tcp/mod.rs +++ b/vendor/tokio/src/net/tcp/mod.rs @@ -1,8 +1,10 @@ -//! TCP utility types +//! TCP utility types. pub(crate) mod listener; -pub(crate) mod socket; +cfg_not_wasi! { + pub(crate) mod socket; +} mod split; pub use split::{ReadHalf, WriteHalf}; diff --git a/vendor/tokio/src/net/tcp/socket.rs b/vendor/tokio/src/net/tcp/socket.rs index 02cb6377e..df792f9a6 100644 --- a/vendor/tokio/src/net/tcp/socket.rs +++ b/vendor/tokio/src/net/tcp/socket.rs @@ -4,10 +4,17 @@ use std::fmt; use std::io; use std::net::SocketAddr; +#[cfg(all(unix, not(tokio_no_as_fd)))] +use std::os::unix::io::{AsFd, BorrowedFd}; #[cfg(unix)] use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd}; -#[cfg(windows)] -use std::os::windows::io::{AsRawSocket, FromRawSocket, IntoRawSocket, RawSocket}; +use std::time::Duration; + +cfg_windows! { + use crate::os::windows::io::{AsRawSocket, FromRawSocket, IntoRawSocket, RawSocket}; + #[cfg(not(tokio_no_as_fd))] + use crate::os::windows::io::{AsSocket, BorrowedSocket}; +} cfg_net! { /// A TCP socket that has not yet been converted to a `TcpStream` or @@ -81,13 +88,14 @@ cfg_net! { /// [`AsRawFd`]: https://doc.rust-lang.org/std/os/unix/io/trait.AsRawFd.html /// [`AsRawSocket`]: https://doc.rust-lang.org/std/os/windows/io/trait.AsRawSocket.html /// [`socket2`]: https://docs.rs/socket2/ + #[cfg_attr(docsrs, doc(alias = "connect_std"))] pub struct TcpSocket { - inner: mio::net::TcpSocket, + inner: socket2::Socket, } } impl TcpSocket { - /// Create a new socket configured for IPv4. + /// Creates a new socket configured for IPv4. /// /// Calls `socket(2)` with `AF_INET` and `SOCK_STREAM`. /// @@ -117,11 +125,10 @@ impl TcpSocket { /// } /// ``` pub fn new_v4() -> io::Result<TcpSocket> { - let inner = mio::net::TcpSocket::new_v4()?; - Ok(TcpSocket { inner }) + TcpSocket::new(socket2::Domain::IPV4) } - /// Create a new socket configured for IPv6. + /// Creates a new socket configured for IPv6. /// /// Calls `socket(2)` with `AF_INET6` and `SOCK_STREAM`. /// @@ -151,11 +158,38 @@ impl TcpSocket { /// } /// ``` pub fn new_v6() -> io::Result<TcpSocket> { - let inner = mio::net::TcpSocket::new_v6()?; + TcpSocket::new(socket2::Domain::IPV6) + } + + fn new(domain: socket2::Domain) -> io::Result<TcpSocket> { + let ty = socket2::Type::STREAM; + #[cfg(any( + target_os = "android", + target_os = "dragonfly", + target_os = "freebsd", + target_os = "fuchsia", + target_os = "illumos", + target_os = "linux", + target_os = "netbsd", + target_os = "openbsd" + ))] + let ty = ty.nonblocking(); + let inner = socket2::Socket::new(domain, ty, Some(socket2::Protocol::TCP))?; + #[cfg(not(any( + target_os = "android", + target_os = "dragonfly", + target_os = "freebsd", + target_os = "fuchsia", + target_os = "illumos", + target_os = "linux", + target_os = "netbsd", + target_os = "openbsd" + )))] + inner.set_nonblocking(true)?; Ok(TcpSocket { inner }) } - /// Allow the socket to bind to an in-use address. + /// Allows the socket to bind to an in-use address. /// /// Behavior is platform specific. Refer to the target platform's /// documentation for more details. @@ -182,10 +216,10 @@ impl TcpSocket { /// } /// ``` pub fn set_reuseaddr(&self, reuseaddr: bool) -> io::Result<()> { - self.inner.set_reuseaddr(reuseaddr) + self.inner.set_reuse_address(reuseaddr) } - /// Retrieves the value set for `SO_REUSEADDR` on this socket + /// Retrieves the value set for `SO_REUSEADDR` on this socket. /// /// # Examples /// @@ -208,10 +242,10 @@ impl TcpSocket { /// } /// ``` pub fn reuseaddr(&self) -> io::Result<bool> { - self.inner.get_reuseaddr() + self.inner.reuse_address() } - /// Allow the socket to bind to an in-use port. Only available for unix systems + /// Allows the socket to bind to an in-use port. Only available for unix systems /// (excluding Solaris & Illumos). /// /// Behavior is platform specific. Refer to the target platform's @@ -242,10 +276,10 @@ impl TcpSocket { doc(cfg(all(unix, not(target_os = "solaris"), not(target_os = "illumos")))) )] pub fn set_reuseport(&self, reuseport: bool) -> io::Result<()> { - self.inner.set_reuseport(reuseport) + self.inner.set_reuse_port(reuseport) } - /// Allow the socket to bind to an in-use port. Only available for unix systems + /// Allows the socket to bind to an in-use port. Only available for unix systems /// (excluding Solaris & Illumos). /// /// Behavior is platform specific. Refer to the target platform's @@ -277,14 +311,14 @@ impl TcpSocket { doc(cfg(all(unix, not(target_os = "solaris"), not(target_os = "illumos")))) )] pub fn reuseport(&self) -> io::Result<bool> { - self.inner.get_reuseport() + self.inner.reuse_port() } /// Sets the size of the TCP send buffer on this socket. /// /// On most operating systems, this sets the `SO_SNDBUF` socket option. pub fn set_send_buffer_size(&self, size: u32) -> io::Result<()> { - self.inner.set_send_buffer_size(size) + self.inner.set_send_buffer_size(size as usize) } /// Returns the size of the TCP send buffer for this socket. @@ -311,14 +345,14 @@ impl TcpSocket { /// /// [`set_send_buffer_size`]: #method.set_send_buffer_size pub fn send_buffer_size(&self) -> io::Result<u32> { - self.inner.get_send_buffer_size() + self.inner.send_buffer_size().map(|n| n as u32) } /// Sets the size of the TCP receive buffer on this socket. /// /// On most operating systems, this sets the `SO_RCVBUF` socket option. pub fn set_recv_buffer_size(&self, size: u32) -> io::Result<()> { - self.inner.set_recv_buffer_size(size) + self.inner.set_recv_buffer_size(size as usize) } /// Returns the size of the TCP receive buffer for this socket. @@ -345,10 +379,160 @@ impl TcpSocket { /// /// [`set_recv_buffer_size`]: #method.set_recv_buffer_size pub fn recv_buffer_size(&self) -> io::Result<u32> { - self.inner.get_recv_buffer_size() + self.inner.recv_buffer_size().map(|n| n as u32) + } + + /// Sets the linger duration of this socket by setting the SO_LINGER option. + /// + /// This option controls the action taken when a stream has unsent messages and the stream is + /// closed. If SO_LINGER is set, the system shall block the process until it can transmit the + /// data or until the time expires. + /// + /// If SO_LINGER is not specified, and the socket is closed, the system handles the call in a + /// way that allows the process to continue as quickly as possible. + pub fn set_linger(&self, dur: Option<Duration>) -> io::Result<()> { + self.inner.set_linger(dur) + } + + /// Reads the linger duration for this socket by getting the `SO_LINGER` + /// option. + /// + /// For more information about this option, see [`set_linger`]. + /// + /// [`set_linger`]: TcpSocket::set_linger + pub fn linger(&self) -> io::Result<Option<Duration>> { + self.inner.linger() + } + + /// Sets the value of the `TCP_NODELAY` option on this socket. + /// + /// If set, this option disables the Nagle algorithm. This means that segments are always + /// sent as soon as possible, even if there is only a small amount of data. When not set, + /// data is buffered until there is a sufficient amount to send out, thereby avoiding + /// the frequent sending of small packets. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpSocket; + /// + /// # async fn dox() -> Result<(), Box<dyn std::error::Error>> { + /// let socket = TcpSocket::new_v4()?; + /// + /// println!("{:?}", socket.nodelay()?); + /// # Ok(()) + /// # } + /// ``` + pub fn set_nodelay(&self, nodelay: bool) -> io::Result<()> { + self.inner.set_nodelay(nodelay) + } + + /// Gets the value of the `TCP_NODELAY` option on this socket. + /// + /// For more information about this option, see [`set_nodelay`]. + /// + /// [`set_nodelay`]: TcpSocket::set_nodelay + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpSocket; + /// + /// # async fn dox() -> Result<(), Box<dyn std::error::Error>> { + /// let stream = TcpSocket::new_v4()?; + /// + /// stream.set_nodelay(true)?; + /// # Ok(()) + /// # } + /// ``` + pub fn nodelay(&self) -> io::Result<bool> { + self.inner.nodelay() } - /// Get the local address of this socket. + /// Gets the value of the `IP_TOS` option for this socket. + /// + /// For more information about this option, see [`set_tos`]. + /// + /// **NOTE:** On Windows, `IP_TOS` is only supported on [Windows 8+ or + /// Windows Server 2012+.](https://docs.microsoft.com/en-us/windows/win32/winsock/ipproto-ip-socket-options) + /// + /// [`set_tos`]: Self::set_tos + // https://docs.rs/socket2/0.4.2/src/socket2/socket.rs.html#1178 + #[cfg(not(any( + target_os = "fuchsia", + target_os = "redox", + target_os = "solaris", + target_os = "illumos", + )))] + #[cfg_attr( + docsrs, + doc(cfg(not(any( + target_os = "fuchsia", + target_os = "redox", + target_os = "solaris", + target_os = "illumos", + )))) + )] + pub fn tos(&self) -> io::Result<u32> { + self.inner.tos() + } + + /// Sets the value for the `IP_TOS` option on this socket. + /// + /// This value sets the type-of-service field that is used in every packet + /// sent from this socket. + /// + /// **NOTE:** On Windows, `IP_TOS` is only supported on [Windows 8+ or + /// Windows Server 2012+.](https://docs.microsoft.com/en-us/windows/win32/winsock/ipproto-ip-socket-options) + // https://docs.rs/socket2/0.4.2/src/socket2/socket.rs.html#1178 + #[cfg(not(any( + target_os = "fuchsia", + target_os = "redox", + target_os = "solaris", + target_os = "illumos", + )))] + #[cfg_attr( + docsrs, + doc(cfg(not(any( + target_os = "fuchsia", + target_os = "redox", + target_os = "solaris", + target_os = "illumos", + )))) + )] + pub fn set_tos(&self, tos: u32) -> io::Result<()> { + self.inner.set_tos(tos) + } + + /// Gets the value for the `SO_BINDTODEVICE` option on this socket + /// + /// This value gets the socket binded device's interface name. + #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux",))] + #[cfg_attr( + docsrs, + doc(cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux",))) + )] + pub fn device(&self) -> io::Result<Option<Vec<u8>>> { + self.inner.device() + } + + /// Sets the value for the `SO_BINDTODEVICE` option on this socket + /// + /// If a socket is bound to an interface, only packets received from that + /// particular interface are processed by the socket. Note that this only + /// works for some socket types, particularly `AF_INET` sockets. + /// + /// If `interface` is `None` or an empty string it removes the binding. + #[cfg(all(any(target_os = "android", target_os = "fuchsia", target_os = "linux")))] + #[cfg_attr( + docsrs, + doc(cfg(all(any(target_os = "android", target_os = "fuchsia", target_os = "linux")))) + )] + pub fn bind_device(&self, interface: Option<&[u8]>) -> io::Result<()> { + self.inner.bind_device(interface) + } + + /// Gets the local address of this socket. /// /// Will fail on windows if called before `bind`. /// @@ -371,10 +555,15 @@ impl TcpSocket { /// } /// ``` pub fn local_addr(&self) -> io::Result<SocketAddr> { - self.inner.get_localaddr() + self.inner.local_addr().and_then(convert_address) + } + + /// Returns the value of the `SO_ERROR` option. + pub fn take_error(&self) -> io::Result<Option<io::Error>> { + self.inner.take_error() } - /// Bind the socket to the given address. + /// Binds the socket to the given address. /// /// This calls the `bind(2)` operating-system function. Behavior is /// platform specific. Refer to the target platform's documentation for more @@ -403,10 +592,10 @@ impl TcpSocket { /// } /// ``` pub fn bind(&self, addr: SocketAddr) -> io::Result<()> { - self.inner.bind(addr) + self.inner.bind(&addr.into()) } - /// Establish a TCP connection with a peer at the specified socket address. + /// Establishes a TCP connection with a peer at the specified socket address. /// /// The `TcpSocket` is consumed. Once the connection is established, a /// connected [`TcpStream`] is returned. If the connection fails, the @@ -439,11 +628,36 @@ impl TcpSocket { /// } /// ``` pub async fn connect(self, addr: SocketAddr) -> io::Result<TcpStream> { - let mio = self.inner.connect(addr)?; + if let Err(err) = self.inner.connect(&addr.into()) { + #[cfg(unix)] + if err.raw_os_error() != Some(libc::EINPROGRESS) { + return Err(err); + } + #[cfg(windows)] + if err.kind() != io::ErrorKind::WouldBlock { + return Err(err); + } + } + #[cfg(unix)] + let mio = { + use std::os::unix::io::{FromRawFd, IntoRawFd}; + + let raw_fd = self.inner.into_raw_fd(); + unsafe { mio::net::TcpStream::from_raw_fd(raw_fd) } + }; + + #[cfg(windows)] + let mio = { + use std::os::windows::io::{FromRawSocket, IntoRawSocket}; + + let raw_socket = self.inner.into_raw_socket(); + unsafe { mio::net::TcpStream::from_raw_socket(raw_socket) } + }; + TcpStream::connect_mio(mio).await } - /// Convert the socket into a `TcpListener`. + /// Converts the socket into a `TcpListener`. /// /// `backlog` defines the maximum number of pending connections are queued /// by the operating system at any given time. Connection are removed from @@ -479,7 +693,23 @@ impl TcpSocket { /// } /// ``` pub fn listen(self, backlog: u32) -> io::Result<TcpListener> { - let mio = self.inner.listen(backlog)?; + self.inner.listen(backlog as i32)?; + #[cfg(unix)] + let mio = { + use std::os::unix::io::{FromRawFd, IntoRawFd}; + + let raw_fd = self.inner.into_raw_fd(); + unsafe { mio::net::TcpListener::from_raw_fd(raw_fd) } + }; + + #[cfg(windows)] + let mio = { + use std::os::windows::io::{FromRawSocket, IntoRawSocket}; + + let raw_socket = self.inner.into_raw_socket(); + unsafe { mio::net::TcpListener::from_raw_socket(raw_socket) } + }; + TcpListener::new(mio) } @@ -491,6 +721,15 @@ impl TcpSocket { /// [`std::net::TcpStream`]: struct@std::net::TcpStream /// [`socket2`]: https://docs.rs/socket2/ /// + /// # Notes + /// + /// The caller is responsible for ensuring that the socket is in + /// non-blocking mode. Otherwise all I/O operations on the socket + /// will block the thread, which will cause unexpected behavior. + /// Non-blocking mode can be set using [`set_nonblocking`]. + /// + /// [`set_nonblocking`]: std::net::TcpStream::set_nonblocking + /// /// # Examples /// /// ``` @@ -499,8 +738,8 @@ impl TcpSocket { /// /// #[tokio::main] /// async fn main() -> std::io::Result<()> { - /// /// let socket2_socket = Socket::new(Domain::IPV4, Type::STREAM, None)?; + /// socket2_socket.set_nonblocking(true)?; /// /// let socket = TcpSocket::from_std_stream(socket2_socket.into()); /// @@ -526,6 +765,16 @@ impl TcpSocket { } } +fn convert_address(address: socket2::SockAddr) -> io::Result<SocketAddr> { + match address.as_socket() { + Some(address) => Ok(address), + None => Err(io::Error::new( + io::ErrorKind::InvalidInput, + "invalid address family (not IPv4 or IPv6)", + )), + } +} + impl fmt::Debug for TcpSocket { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { self.inner.fmt(fmt) @@ -539,6 +788,13 @@ impl AsRawFd for TcpSocket { } } +#[cfg(all(unix, not(tokio_no_as_fd)))] +impl AsFd for TcpSocket { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } +} + #[cfg(unix)] impl FromRawFd for TcpSocket { /// Converts a `RawFd` to a `TcpSocket`. @@ -548,7 +804,7 @@ impl FromRawFd for TcpSocket { /// The caller is responsible for ensuring that the socket is in /// non-blocking mode. unsafe fn from_raw_fd(fd: RawFd) -> TcpSocket { - let inner = mio::net::TcpSocket::from_raw_fd(fd); + let inner = socket2::Socket::from_raw_fd(fd); TcpSocket { inner } } } @@ -560,30 +816,36 @@ impl IntoRawFd for TcpSocket { } } -#[cfg(windows)] -impl IntoRawSocket for TcpSocket { - fn into_raw_socket(self) -> RawSocket { - self.inner.into_raw_socket() +cfg_windows! { + impl IntoRawSocket for TcpSocket { + fn into_raw_socket(self) -> RawSocket { + self.inner.into_raw_socket() + } + } + + impl AsRawSocket for TcpSocket { + fn as_raw_socket(&self) -> RawSocket { + self.inner.as_raw_socket() + } } -} -#[cfg(windows)] -impl AsRawSocket for TcpSocket { - fn as_raw_socket(&self) -> RawSocket { - self.inner.as_raw_socket() + #[cfg(not(tokio_no_as_fd))] + impl AsSocket for TcpSocket { + fn as_socket(&self) -> BorrowedSocket<'_> { + unsafe { BorrowedSocket::borrow_raw(self.as_raw_socket()) } + } } -} -#[cfg(windows)] -impl FromRawSocket for TcpSocket { - /// Converts a `RawSocket` to a `TcpStream`. - /// - /// # Notes - /// - /// The caller is responsible for ensuring that the socket is in - /// non-blocking mode. - unsafe fn from_raw_socket(socket: RawSocket) -> TcpSocket { - let inner = mio::net::TcpSocket::from_raw_socket(socket); - TcpSocket { inner } + impl FromRawSocket for TcpSocket { + /// Converts a `RawSocket` to a `TcpStream`. + /// + /// # Notes + /// + /// The caller is responsible for ensuring that the socket is in + /// non-blocking mode. + unsafe fn from_raw_socket(socket: RawSocket) -> TcpSocket { + let inner = socket2::Socket::from_raw_socket(socket); + TcpSocket { inner } + } } } diff --git a/vendor/tokio/src/net/tcp/split.rs b/vendor/tokio/src/net/tcp/split.rs index 8ae70ce13..1a1e22253 100644 --- a/vendor/tokio/src/net/tcp/split.rs +++ b/vendor/tokio/src/net/tcp/split.rs @@ -9,14 +9,18 @@ //! level. use crate::future::poll_fn; -use crate::io::{AsyncRead, AsyncWrite, ReadBuf}; +use crate::io::{AsyncRead, AsyncWrite, Interest, ReadBuf, Ready}; use crate::net::TcpStream; use std::io; -use std::net::Shutdown; +use std::net::{Shutdown, SocketAddr}; use std::pin::Pin; use std::task::{Context, Poll}; +cfg_io_util! { + use bytes::BufMut; +} + /// Borrowed read half of a [`TcpStream`], created by [`split`]. /// /// Reading from a `ReadHalf` is usually done using the convenience methods found on the @@ -49,7 +53,7 @@ pub(crate) fn split(stream: &mut TcpStream) -> (ReadHalf<'_>, WriteHalf<'_>) { } impl ReadHalf<'_> { - /// Attempt to receive data on the socket, without removing that data from + /// Attempts to receive data on the socket, without removing that data from /// the queue, registering the current task for wakeup if data is not yet /// available. /// @@ -134,6 +138,233 @@ impl ReadHalf<'_> { let mut buf = ReadBuf::new(buf); poll_fn(|cx| self.poll_peek(cx, &mut buf)).await } + + /// Waits for any of the requested ready states. + /// + /// This function is usually paired with [`try_read()`]. It can be used instead + /// of [`readable()`] to check the returned ready set for [`Ready::READABLE`] + /// and [`Ready::READ_CLOSED`] events. + /// + /// The function may complete without the socket being ready. This is a + /// false-positive and attempting an operation will return with + /// `io::ErrorKind::WouldBlock`. The function can also return with an empty + /// [`Ready`] set, so you should always check the returned value and possibly + /// wait again if the requested states are not set. + /// + /// This function is equivalent to [`TcpStream::ready`]. + /// + /// [`try_read()`]: Self::try_read + /// [`readable()`]: Self::readable + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to read or write that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn ready(&self, interest: Interest) -> io::Result<Ready> { + self.0.ready(interest).await + } + + /// Waits for the socket to become readable. + /// + /// This function is equivalent to `ready(Interest::READABLE)` and is usually + /// paired with `try_read()`. + /// + /// This function is also equivalent to [`TcpStream::ready`]. + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to read that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn readable(&self) -> io::Result<()> { + self.0.readable().await + } + + /// Tries to read data from the stream into the provided buffer, returning how + /// many bytes were read. + /// + /// Receives any pending data from the socket but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read()` is non-blocking, the buffer does not have to be stored by + /// the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: Self::readable() + /// [`ready()`]: Self::ready() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. If `n` is `0`, then it can indicate one of two scenarios: + /// + /// 1. The stream's read half is closed and will no longer yield data. + /// 2. The specified buffer was 0 bytes in length. + /// + /// If the stream is not ready to read data, + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_read(&self, buf: &mut [u8]) -> io::Result<usize> { + self.0.try_read(buf) + } + + /// Tries to read data from the stream into the provided buffers, returning + /// how many bytes were read. + /// + /// Data is copied to fill each buffer in order, with the final buffer + /// written to possibly being only partially filled. This method behaves + /// equivalently to a single call to [`try_read()`] with concatenated + /// buffers. + /// + /// Receives any pending data from the socket but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read_vectored()` is non-blocking, the buffer does not have to be + /// stored by the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`try_read()`]: Self::try_read() + /// [`readable()`]: Self::readable() + /// [`ready()`]: Self::ready() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. `Ok(0)` indicates the stream's read half is closed + /// and will no longer yield data. If the stream is not ready to read data + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_read_vectored(&self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result<usize> { + self.0.try_read_vectored(bufs) + } + + cfg_io_util! { + /// Tries to read data from the stream into the provided buffer, advancing the + /// buffer's internal cursor, returning how many bytes were read. + /// + /// Receives any pending data from the socket but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read_buf()` is non-blocking, the buffer does not have to be stored by + /// the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: Self::readable() + /// [`ready()`]: Self::ready() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. `Ok(0)` indicates the stream's read half is closed + /// and will no longer yield data. If the stream is not ready to read data + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_read_buf<B: BufMut>(&self, buf: &mut B) -> io::Result<usize> { + self.0.try_read_buf(buf) + } + } + + /// Returns the remote address that this stream is connected to. + pub fn peer_addr(&self) -> io::Result<SocketAddr> { + self.0.peer_addr() + } + + /// Returns the local address that this stream is bound to. + pub fn local_addr(&self) -> io::Result<SocketAddr> { + self.0.local_addr() + } +} + +impl WriteHalf<'_> { + /// Waits for any of the requested ready states. + /// + /// This function is usually paired with [`try_write()`]. It can be used instead + /// of [`writable()`] to check the returned ready set for [`Ready::WRITABLE`] + /// and [`Ready::WRITE_CLOSED`] events. + /// + /// The function may complete without the socket being ready. This is a + /// false-positive and attempting an operation will return with + /// `io::ErrorKind::WouldBlock`. The function can also return with an empty + /// [`Ready`] set, so you should always check the returned value and possibly + /// wait again if the requested states are not set. + /// + /// This function is equivalent to [`TcpStream::ready`]. + /// + /// [`try_write()`]: Self::try_write + /// [`writable()`]: Self::writable + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to read or write that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn ready(&self, interest: Interest) -> io::Result<Ready> { + self.0.ready(interest).await + } + + /// Waits for the socket to become writable. + /// + /// This function is equivalent to `ready(Interest::WRITABLE)` and is usually + /// paired with `try_write()`. + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to write that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn writable(&self) -> io::Result<()> { + self.0.writable().await + } + + /// Tries to write a buffer to the stream, returning how many bytes were + /// written. + /// + /// The function will attempt to write the entire contents of `buf`, but + /// only part of the buffer may be written. + /// + /// This function is usually paired with `writable()`. + /// + /// # Return + /// + /// If data is successfully written, `Ok(n)` is returned, where `n` is the + /// number of bytes written. If the stream is not ready to write data, + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_write(&self, buf: &[u8]) -> io::Result<usize> { + self.0.try_write(buf) + } + + /// Tries to write several buffers to the stream, returning how many bytes + /// were written. + /// + /// Data is written from each buffer in order, with the final buffer read + /// from possible being only partially consumed. This method behaves + /// equivalently to a single call to [`try_write()`] with concatenated + /// buffers. + /// + /// This function is usually paired with `writable()`. + /// + /// [`try_write()`]: Self::try_write() + /// + /// # Return + /// + /// If data is successfully written, `Ok(n)` is returned, where `n` is the + /// number of bytes written. If the stream is not ready to write data, + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_write_vectored(&self, bufs: &[io::IoSlice<'_>]) -> io::Result<usize> { + self.0.try_write_vectored(bufs) + } + + /// Returns the remote address that this stream is connected to. + pub fn peer_addr(&self) -> io::Result<SocketAddr> { + self.0.peer_addr() + } + + /// Returns the local address that this stream is bound to. + pub fn local_addr(&self) -> io::Result<SocketAddr> { + self.0.local_addr() + } } impl AsyncRead for ReadHalf<'_> { diff --git a/vendor/tokio/src/net/tcp/split_owned.rs b/vendor/tokio/src/net/tcp/split_owned.rs index 1bcb4f2ea..6771d6497 100644 --- a/vendor/tokio/src/net/tcp/split_owned.rs +++ b/vendor/tokio/src/net/tcp/split_owned.rs @@ -9,16 +9,20 @@ //! level. use crate::future::poll_fn; -use crate::io::{AsyncRead, AsyncWrite, ReadBuf}; +use crate::io::{AsyncRead, AsyncWrite, Interest, ReadBuf, Ready}; use crate::net::TcpStream; use std::error::Error; -use std::net::Shutdown; +use std::net::{Shutdown, SocketAddr}; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; use std::{fmt, io}; +cfg_io_util! { + use bytes::BufMut; +} + /// Owned read half of a [`TcpStream`], created by [`into_split`]. /// /// Reading from an `OwnedReadHalf` is usually done using the convenience methods found @@ -189,6 +193,141 @@ impl OwnedReadHalf { let mut buf = ReadBuf::new(buf); poll_fn(|cx| self.poll_peek(cx, &mut buf)).await } + + /// Waits for any of the requested ready states. + /// + /// This function is usually paired with [`try_read()`]. It can be used instead + /// of [`readable()`] to check the returned ready set for [`Ready::READABLE`] + /// and [`Ready::READ_CLOSED`] events. + /// + /// The function may complete without the socket being ready. This is a + /// false-positive and attempting an operation will return with + /// `io::ErrorKind::WouldBlock`. The function can also return with an empty + /// [`Ready`] set, so you should always check the returned value and possibly + /// wait again if the requested states are not set. + /// + /// This function is equivalent to [`TcpStream::ready`]. + /// + /// [`try_read()`]: Self::try_read + /// [`readable()`]: Self::readable + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to read or write that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn ready(&self, interest: Interest) -> io::Result<Ready> { + self.inner.ready(interest).await + } + + /// Waits for the socket to become readable. + /// + /// This function is equivalent to `ready(Interest::READABLE)` and is usually + /// paired with `try_read()`. + /// + /// This function is also equivalent to [`TcpStream::ready`]. + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to read that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn readable(&self) -> io::Result<()> { + self.inner.readable().await + } + + /// Tries to read data from the stream into the provided buffer, returning how + /// many bytes were read. + /// + /// Receives any pending data from the socket but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read()` is non-blocking, the buffer does not have to be stored by + /// the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: Self::readable() + /// [`ready()`]: Self::ready() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. If `n` is `0`, then it can indicate one of two scenarios: + /// + /// 1. The stream's read half is closed and will no longer yield data. + /// 2. The specified buffer was 0 bytes in length. + /// + /// If the stream is not ready to read data, + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_read(&self, buf: &mut [u8]) -> io::Result<usize> { + self.inner.try_read(buf) + } + + /// Tries to read data from the stream into the provided buffers, returning + /// how many bytes were read. + /// + /// Data is copied to fill each buffer in order, with the final buffer + /// written to possibly being only partially filled. This method behaves + /// equivalently to a single call to [`try_read()`] with concatenated + /// buffers. + /// + /// Receives any pending data from the socket but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read_vectored()` is non-blocking, the buffer does not have to be + /// stored by the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`try_read()`]: Self::try_read() + /// [`readable()`]: Self::readable() + /// [`ready()`]: Self::ready() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. `Ok(0)` indicates the stream's read half is closed + /// and will no longer yield data. If the stream is not ready to read data + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_read_vectored(&self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result<usize> { + self.inner.try_read_vectored(bufs) + } + + cfg_io_util! { + /// Tries to read data from the stream into the provided buffer, advancing the + /// buffer's internal cursor, returning how many bytes were read. + /// + /// Receives any pending data from the socket but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read_buf()` is non-blocking, the buffer does not have to be stored by + /// the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: Self::readable() + /// [`ready()`]: Self::ready() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. `Ok(0)` indicates the stream's read half is closed + /// and will no longer yield data. If the stream is not ready to read data + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_read_buf<B: BufMut>(&self, buf: &mut B) -> io::Result<usize> { + self.inner.try_read_buf(buf) + } + } + + /// Returns the remote address that this stream is connected to. + pub fn peer_addr(&self) -> io::Result<SocketAddr> { + self.inner.peer_addr() + } + + /// Returns the local address that this stream is bound to. + pub fn local_addr(&self) -> io::Result<SocketAddr> { + self.inner.local_addr() + } } impl AsyncRead for OwnedReadHalf { @@ -211,13 +350,103 @@ impl OwnedWriteHalf { reunite(other, self) } - /// Destroy the write half, but don't close the write half of the stream + /// Destroys the write half, but don't close the write half of the stream /// until the read half is dropped. If the read half has already been /// dropped, this closes the stream. pub fn forget(mut self) { self.shutdown_on_drop = false; drop(self); } + + /// Waits for any of the requested ready states. + /// + /// This function is usually paired with [`try_write()`]. It can be used instead + /// of [`writable()`] to check the returned ready set for [`Ready::WRITABLE`] + /// and [`Ready::WRITE_CLOSED`] events. + /// + /// The function may complete without the socket being ready. This is a + /// false-positive and attempting an operation will return with + /// `io::ErrorKind::WouldBlock`. The function can also return with an empty + /// [`Ready`] set, so you should always check the returned value and possibly + /// wait again if the requested states are not set. + /// + /// This function is equivalent to [`TcpStream::ready`]. + /// + /// [`try_write()`]: Self::try_write + /// [`writable()`]: Self::writable + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to read or write that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn ready(&self, interest: Interest) -> io::Result<Ready> { + self.inner.ready(interest).await + } + + /// Waits for the socket to become writable. + /// + /// This function is equivalent to `ready(Interest::WRITABLE)` and is usually + /// paired with `try_write()`. + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to write that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn writable(&self) -> io::Result<()> { + self.inner.writable().await + } + + /// Tries to write a buffer to the stream, returning how many bytes were + /// written. + /// + /// The function will attempt to write the entire contents of `buf`, but + /// only part of the buffer may be written. + /// + /// This function is usually paired with `writable()`. + /// + /// # Return + /// + /// If data is successfully written, `Ok(n)` is returned, where `n` is the + /// number of bytes written. If the stream is not ready to write data, + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_write(&self, buf: &[u8]) -> io::Result<usize> { + self.inner.try_write(buf) + } + + /// Tries to write several buffers to the stream, returning how many bytes + /// were written. + /// + /// Data is written from each buffer in order, with the final buffer read + /// from possible being only partially consumed. This method behaves + /// equivalently to a single call to [`try_write()`] with concatenated + /// buffers. + /// + /// This function is usually paired with `writable()`. + /// + /// [`try_write()`]: Self::try_write() + /// + /// # Return + /// + /// If data is successfully written, `Ok(n)` is returned, where `n` is the + /// number of bytes written. If the stream is not ready to write data, + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_write_vectored(&self, bufs: &[io::IoSlice<'_>]) -> io::Result<usize> { + self.inner.try_write_vectored(bufs) + } + + /// Returns the remote address that this stream is connected to. + pub fn peer_addr(&self) -> io::Result<SocketAddr> { + self.inner.peer_addr() + } + + /// Returns the local address that this stream is bound to. + pub fn local_addr(&self) -> io::Result<SocketAddr> { + self.inner.local_addr() + } } impl Drop for OwnedWriteHalf { @@ -267,12 +496,12 @@ impl AsyncWrite for OwnedWriteHalf { impl AsRef<TcpStream> for OwnedReadHalf { fn as_ref(&self) -> &TcpStream { - &*self.inner + &self.inner } } impl AsRef<TcpStream> for OwnedWriteHalf { fn as_ref(&self) -> &TcpStream { - &*self.inner + &self.inner } } diff --git a/vendor/tokio/src/net/tcp/stream.rs b/vendor/tokio/src/net/tcp/stream.rs index 0277a360d..b7104d78b 100644 --- a/vendor/tokio/src/net/tcp/stream.rs +++ b/vendor/tokio/src/net/tcp/stream.rs @@ -1,16 +1,18 @@ -use crate::future::poll_fn; +cfg_not_wasi! { + use crate::future::poll_fn; + use crate::net::{to_socket_addrs, ToSocketAddrs}; + use std::time::Duration; +} + use crate::io::{AsyncRead, AsyncWrite, Interest, PollEvented, ReadBuf, Ready}; use crate::net::tcp::split::{split, ReadHalf, WriteHalf}; use crate::net::tcp::split_owned::{split_owned, OwnedReadHalf, OwnedWriteHalf}; -use crate::net::{to_socket_addrs, ToSocketAddrs}; -use std::convert::TryFrom; use std::fmt; use std::io; use std::net::{Shutdown, SocketAddr}; use std::pin::Pin; use std::task::{Context, Poll}; -use std::time::Duration; cfg_io_util! { use bytes::BufMut; @@ -70,86 +72,88 @@ cfg_net! { } impl TcpStream { - /// Opens a TCP connection to a remote host. - /// - /// `addr` is an address of the remote host. Anything which implements the - /// [`ToSocketAddrs`] trait can be supplied as the address. If `addr` - /// yields multiple addresses, connect will be attempted with each of the - /// addresses until a connection is successful. If none of the addresses - /// result in a successful connection, the error returned from the last - /// connection attempt (the last address) is returned. - /// - /// To configure the socket before connecting, you can use the [`TcpSocket`] - /// type. - /// - /// [`ToSocketAddrs`]: trait@crate::net::ToSocketAddrs - /// [`TcpSocket`]: struct@crate::net::TcpSocket - /// - /// # Examples - /// - /// ```no_run - /// use tokio::net::TcpStream; - /// use tokio::io::AsyncWriteExt; - /// use std::error::Error; - /// - /// #[tokio::main] - /// async fn main() -> Result<(), Box<dyn Error>> { - /// // Connect to a peer - /// let mut stream = TcpStream::connect("127.0.0.1:8080").await?; - /// - /// // Write some data. - /// stream.write_all(b"hello world!").await?; - /// - /// Ok(()) - /// } - /// ``` - /// - /// The [`write_all`] method is defined on the [`AsyncWriteExt`] trait. - /// - /// [`write_all`]: fn@crate::io::AsyncWriteExt::write_all - /// [`AsyncWriteExt`]: trait@crate::io::AsyncWriteExt - pub async fn connect<A: ToSocketAddrs>(addr: A) -> io::Result<TcpStream> { - let addrs = to_socket_addrs(addr).await?; + cfg_not_wasi! { + /// Opens a TCP connection to a remote host. + /// + /// `addr` is an address of the remote host. Anything which implements the + /// [`ToSocketAddrs`] trait can be supplied as the address. If `addr` + /// yields multiple addresses, connect will be attempted with each of the + /// addresses until a connection is successful. If none of the addresses + /// result in a successful connection, the error returned from the last + /// connection attempt (the last address) is returned. + /// + /// To configure the socket before connecting, you can use the [`TcpSocket`] + /// type. + /// + /// [`ToSocketAddrs`]: trait@crate::net::ToSocketAddrs + /// [`TcpSocket`]: struct@crate::net::TcpSocket + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpStream; + /// use tokio::io::AsyncWriteExt; + /// use std::error::Error; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// // Connect to a peer + /// let mut stream = TcpStream::connect("127.0.0.1:8080").await?; + /// + /// // Write some data. + /// stream.write_all(b"hello world!").await?; + /// + /// Ok(()) + /// } + /// ``` + /// + /// The [`write_all`] method is defined on the [`AsyncWriteExt`] trait. + /// + /// [`write_all`]: fn@crate::io::AsyncWriteExt::write_all + /// [`AsyncWriteExt`]: trait@crate::io::AsyncWriteExt + pub async fn connect<A: ToSocketAddrs>(addr: A) -> io::Result<TcpStream> { + let addrs = to_socket_addrs(addr).await?; - let mut last_err = None; + let mut last_err = None; - for addr in addrs { - match TcpStream::connect_addr(addr).await { - Ok(stream) => return Ok(stream), - Err(e) => last_err = Some(e), + for addr in addrs { + match TcpStream::connect_addr(addr).await { + Ok(stream) => return Ok(stream), + Err(e) => last_err = Some(e), + } } + + Err(last_err.unwrap_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidInput, + "could not resolve to any address", + ) + })) } - Err(last_err.unwrap_or_else(|| { - io::Error::new( - io::ErrorKind::InvalidInput, - "could not resolve to any address", - ) - })) - } + /// Establishes a connection to the specified `addr`. + async fn connect_addr(addr: SocketAddr) -> io::Result<TcpStream> { + let sys = mio::net::TcpStream::connect(addr)?; + TcpStream::connect_mio(sys).await + } - /// Establishes a connection to the specified `addr`. - async fn connect_addr(addr: SocketAddr) -> io::Result<TcpStream> { - let sys = mio::net::TcpStream::connect(addr)?; - TcpStream::connect_mio(sys).await - } + pub(crate) async fn connect_mio(sys: mio::net::TcpStream) -> io::Result<TcpStream> { + let stream = TcpStream::new(sys)?; - pub(crate) async fn connect_mio(sys: mio::net::TcpStream) -> io::Result<TcpStream> { - let stream = TcpStream::new(sys)?; + // Once we've connected, wait for the stream to be writable as + // that's when the actual connection has been initiated. Once we're + // writable we check for `take_socket_error` to see if the connect + // actually hit an error or not. + // + // If all that succeeded then we ship everything on up. + poll_fn(|cx| stream.io.registration().poll_write_ready(cx)).await?; - // Once we've connected, wait for the stream to be writable as - // that's when the actual connection has been initiated. Once we're - // writable we check for `take_socket_error` to see if the connect - // actually hit an error or not. - // - // If all that succeeded then we ship everything on up. - poll_fn(|cx| stream.io.registration().poll_write_ready(cx)).await?; + if let Some(e) = stream.io.take_error()? { + return Err(e); + } - if let Some(e) = stream.io.take_error()? { - return Err(e); + Ok(stream) } - - Ok(stream) } pub(crate) fn new(connected: mio::net::TcpStream) -> io::Result<TcpStream> { @@ -160,9 +164,16 @@ impl TcpStream { /// Creates new `TcpStream` from a `std::net::TcpStream`. /// /// This function is intended to be used to wrap a TCP stream from the - /// standard library in the Tokio equivalent. The conversion assumes nothing - /// about the underlying stream; it is left up to the user to set it in - /// non-blocking mode. + /// standard library in the Tokio equivalent. + /// + /// # Notes + /// + /// The caller is responsible for ensuring that the stream is in + /// non-blocking mode. Otherwise all I/O operations on the stream + /// will block the thread, which will cause unexpected behavior. + /// Non-blocking mode can be set using [`set_nonblocking`]. + /// + /// [`set_nonblocking`]: std::net::TcpStream::set_nonblocking /// /// # Examples /// @@ -181,18 +192,20 @@ impl TcpStream { /// /// # Panics /// - /// This function panics if thread-local runtime is not set. + /// This function panics if it is not called from within a runtime with + /// IO enabled. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. + #[track_caller] pub fn from_std(stream: std::net::TcpStream) -> io::Result<TcpStream> { let io = mio::net::TcpStream::from_std(stream); let io = PollEvented::new(io)?; Ok(TcpStream { io }) } - /// Turn a [`tokio::net::TcpStream`] into a [`std::net::TcpStream`]. + /// Turns a [`tokio::net::TcpStream`] into a [`std::net::TcpStream`]. /// /// The returned [`std::net::TcpStream`] will have nonblocking mode set as `true`. /// Use [`set_nonblocking`] to change the blocking mode if needed. @@ -244,6 +257,15 @@ impl TcpStream { .map(|io| io.into_raw_socket()) .map(|raw_socket| unsafe { std::net::TcpStream::from_raw_socket(raw_socket) }) } + + #[cfg(tokio_wasi)] + { + use std::os::wasi::io::{FromRawFd, IntoRawFd}; + self.io + .into_inner() + .map(|io| io.into_raw_fd()) + .map(|raw_fd| unsafe { std::net::TcpStream::from_raw_fd(raw_fd) }) + } } /// Returns the local address that this stream is bound to. @@ -264,6 +286,11 @@ impl TcpStream { self.io.local_addr() } + /// Returns the value of the `SO_ERROR` option. + pub fn take_error(&self) -> io::Result<Option<io::Error>> { + self.io.take_error() + } + /// Returns the remote address that this stream is connected to. /// /// # Examples @@ -350,12 +377,18 @@ impl TcpStream { } } - /// Wait for any of the requested ready states. + /// Waits for any of the requested ready states. /// /// This function is usually paired with `try_read()` or `try_write()`. It /// can be used to concurrently read / write to the same socket on a single /// task without splitting the socket. /// + /// The function may complete without the socket being ready. This is a + /// false-positive and attempting an operation will return with + /// `io::ErrorKind::WouldBlock`. The function can also return with an empty + /// [`Ready`] set, so you should always check the returned value and possibly + /// wait again if the requested states are not set. + /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method @@ -387,7 +420,7 @@ impl TcpStream { /// // if the readiness event is a false positive. /// match stream.try_read(&mut data) { /// Ok(n) => { - /// println!("read {} bytes", n); + /// println!("read {} bytes", n); /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; @@ -422,7 +455,7 @@ impl TcpStream { Ok(event.ready) } - /// Wait for the socket to become readable. + /// Waits for the socket to become readable. /// /// This function is equivalent to `ready(Interest::READABLE)` and is usually /// paired with `try_read()`. @@ -510,7 +543,7 @@ impl TcpStream { self.io.registration().poll_read_ready(cx).map_ok(|_| ()) } - /// Try to read data from the stream into the provided buffer, returning how + /// Tries to read data from the stream into the provided buffer, returning how /// many bytes were read. /// /// Receives any pending data from the socket but does not wait for new data @@ -526,8 +559,12 @@ impl TcpStream { /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the - /// number of bytes read. `Ok(0)` indicates the stream's read half is closed - /// and will no longer yield data. If the stream is not ready to read data + /// number of bytes read. If `n` is `0`, then it can indicate one of two scenarios: + /// + /// 1. The stream's read half is closed and will no longer yield data. + /// 2. The specified buffer was 0 bytes in length. + /// + /// If the stream is not ready to read data, /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples @@ -577,7 +614,7 @@ impl TcpStream { .try_io(Interest::READABLE, || (&*self.io).read(buf)) } - /// Try to read data from the stream into the provided buffers, returning + /// Tries to read data from the stream into the provided buffers, returning /// how many bytes were read. /// /// Data is copied to fill each buffer in order, with the final buffer @@ -656,7 +693,7 @@ impl TcpStream { } cfg_io_util! { - /// Try to read data from the stream into the provided buffer, advancing the + /// Tries to read data from the stream into the provided buffer, advancing the /// buffer's internal cursor, returning how many bytes were read. /// /// Receives any pending data from the socket but does not wait for new data @@ -734,7 +771,7 @@ impl TcpStream { } } - /// Wait for the socket to become writable. + /// Waits for the socket to become writable. /// /// This function is equivalent to `ready(Interest::WRITABLE)` and is usually /// paired with `try_write()`. @@ -874,7 +911,7 @@ impl TcpStream { .try_io(Interest::WRITABLE, || (&*self.io).write(buf)) } - /// Try to write several buffers to the stream, returning how many bytes + /// Tries to write several buffers to the stream, returning how many bytes /// were written. /// /// Data is written from each buffer in order, with the final buffer read @@ -936,6 +973,84 @@ impl TcpStream { .try_io(Interest::WRITABLE, || (&*self.io).write_vectored(bufs)) } + /// Tries to read or write from the socket using a user-provided IO operation. + /// + /// If the socket is ready, the provided closure is called. The closure + /// should attempt to perform IO operation on the socket by manually + /// calling the appropriate syscall. If the operation fails because the + /// socket is not actually ready, then the closure should return a + /// `WouldBlock` error and the readiness flag is cleared. The return value + /// of the closure is then returned by `try_io`. + /// + /// If the socket is not ready, then the closure is not called + /// and a `WouldBlock` error is returned. + /// + /// The closure should only return a `WouldBlock` error if it has performed + /// an IO operation on the socket that failed due to the socket not being + /// ready. Returning a `WouldBlock` error in any other situation will + /// incorrectly clear the readiness flag, which can cause the socket to + /// behave incorrectly. + /// + /// The closure should not perform the IO operation using any of the methods + /// defined on the Tokio `TcpStream` type, as this will mess with the + /// readiness flag and can cause the socket to behave incorrectly. + /// + /// This method is not intended to be used with combined interests. + /// The closure should perform only one type of IO operation, so it should not + /// require more than one ready state. This method may panic or sleep forever + /// if it is called with a combined interest. + /// + /// Usually, [`readable()`], [`writable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: TcpStream::readable() + /// [`writable()`]: TcpStream::writable() + /// [`ready()`]: TcpStream::ready() + pub fn try_io<R>( + &self, + interest: Interest, + f: impl FnOnce() -> io::Result<R>, + ) -> io::Result<R> { + self.io + .registration() + .try_io(interest, || self.io.try_io(f)) + } + + /// Reads or writes from the socket using a user-provided IO operation. + /// + /// The readiness of the socket is awaited and when the socket is ready, + /// the provided closure is called. The closure should attempt to perform + /// IO operation on the socket by manually calling the appropriate syscall. + /// If the operation fails because the socket is not actually ready, + /// then the closure should return a `WouldBlock` error. In such case the + /// readiness flag is cleared and the socket readiness is awaited again. + /// This loop is repeated until the closure returns an `Ok` or an error + /// other than `WouldBlock`. + /// + /// The closure should only return a `WouldBlock` error if it has performed + /// an IO operation on the socket that failed due to the socket not being + /// ready. Returning a `WouldBlock` error in any other situation will + /// incorrectly clear the readiness flag, which can cause the socket to + /// behave incorrectly. + /// + /// The closure should not perform the IO operation using any of the methods + /// defined on the Tokio `TcpStream` type, as this will mess with the + /// readiness flag and can cause the socket to behave incorrectly. + /// + /// This method is not intended to be used with combined interests. + /// The closure should perform only one type of IO operation, so it should not + /// require more than one ready state. This method may panic or sleep forever + /// if it is called with a combined interest. + pub async fn async_io<R>( + &self, + interest: Interest, + mut f: impl FnMut() -> io::Result<R>, + ) -> io::Result<R> { + self.io + .registration() + .async_io(interest, || self.io.try_io(&mut f)) + .await + } + /// Receives data on the socket from the remote address to which it is /// connected, without removing that data from the queue. On success, /// returns the number of bytes peeked. @@ -1035,69 +1150,53 @@ impl TcpStream { self.io.set_nodelay(nodelay) } - /// Reads the linger duration for this socket by getting the `SO_LINGER` - /// option. - /// - /// For more information about this option, see [`set_linger`]. - /// - /// [`set_linger`]: TcpStream::set_linger - /// - /// # Examples - /// - /// ```no_run - /// use tokio::net::TcpStream; - /// - /// # async fn dox() -> Result<(), Box<dyn std::error::Error>> { - /// let stream = TcpStream::connect("127.0.0.1:8080").await?; - /// - /// println!("{:?}", stream.linger()?); - /// # Ok(()) - /// # } - /// ``` - pub fn linger(&self) -> io::Result<Option<Duration>> { - let mio_socket = std::mem::ManuallyDrop::new(self.to_mio()); - - mio_socket.get_linger() - } - - /// Sets the linger duration of this socket by setting the SO_LINGER option. - /// - /// This option controls the action taken when a stream has unsent messages and the stream is - /// closed. If SO_LINGER is set, the system shall block the process until it can transmit the - /// data or until the time expires. - /// - /// If SO_LINGER is not specified, and the stream is closed, the system handles the call in a - /// way that allows the process to continue as quickly as possible. - /// - /// # Examples - /// - /// ```no_run - /// use tokio::net::TcpStream; - /// - /// # async fn dox() -> Result<(), Box<dyn std::error::Error>> { - /// let stream = TcpStream::connect("127.0.0.1:8080").await?; - /// - /// stream.set_linger(None)?; - /// # Ok(()) - /// # } - /// ``` - pub fn set_linger(&self, dur: Option<Duration>) -> io::Result<()> { - let mio_socket = std::mem::ManuallyDrop::new(self.to_mio()); - - mio_socket.set_linger(dur) - } - - fn to_mio(&self) -> mio::net::TcpSocket { - #[cfg(windows)] - { - use std::os::windows::io::{AsRawSocket, FromRawSocket}; - unsafe { mio::net::TcpSocket::from_raw_socket(self.as_raw_socket()) } + cfg_not_wasi! { + /// Reads the linger duration for this socket by getting the `SO_LINGER` + /// option. + /// + /// For more information about this option, see [`set_linger`]. + /// + /// [`set_linger`]: TcpStream::set_linger + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpStream; + /// + /// # async fn dox() -> Result<(), Box<dyn std::error::Error>> { + /// let stream = TcpStream::connect("127.0.0.1:8080").await?; + /// + /// println!("{:?}", stream.linger()?); + /// # Ok(()) + /// # } + /// ``` + pub fn linger(&self) -> io::Result<Option<Duration>> { + socket2::SockRef::from(self).linger() } - #[cfg(unix)] - { - use std::os::unix::io::{AsRawFd, FromRawFd}; - unsafe { mio::net::TcpSocket::from_raw_fd(self.as_raw_fd()) } + /// Sets the linger duration of this socket by setting the SO_LINGER option. + /// + /// This option controls the action taken when a stream has unsent messages and the stream is + /// closed. If SO_LINGER is set, the system shall block the process until it can transmit the + /// data or until the time expires. + /// + /// If SO_LINGER is not specified, and the stream is closed, the system handles the call in a + /// way that allows the process to continue as quickly as possible. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpStream; + /// + /// # async fn dox() -> Result<(), Box<dyn std::error::Error>> { + /// let stream = TcpStream::connect("127.0.0.1:8080").await?; + /// + /// stream.set_linger(None)?; + /// # Ok(()) + /// # } + /// ``` + pub fn set_linger(&self, dur: Option<Duration>) -> io::Result<()> { + socket2::SockRef::from(self).set_linger(dur) } } @@ -1278,16 +1377,49 @@ mod sys { self.io.as_raw_fd() } } + + #[cfg(not(tokio_no_as_fd))] + impl AsFd for TcpStream { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } + } } -#[cfg(windows)] -mod sys { - use super::TcpStream; - use std::os::windows::prelude::*; +cfg_windows! { + use crate::os::windows::io::{AsRawSocket, RawSocket}; + #[cfg(not(tokio_no_as_fd))] + use crate::os::windows::io::{AsSocket, BorrowedSocket}; impl AsRawSocket for TcpStream { fn as_raw_socket(&self) -> RawSocket { self.io.as_raw_socket() } } + + #[cfg(not(tokio_no_as_fd))] + impl AsSocket for TcpStream { + fn as_socket(&self) -> BorrowedSocket<'_> { + unsafe { BorrowedSocket::borrow_raw(self.as_raw_socket()) } + } + } +} + +#[cfg(all(tokio_unstable, tokio_wasi))] +mod sys { + use super::TcpStream; + use std::os::wasi::prelude::*; + + impl AsRawFd for TcpStream { + fn as_raw_fd(&self) -> RawFd { + self.io.as_raw_fd() + } + } + + #[cfg(not(tokio_no_as_fd))] + impl AsFd for TcpStream { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } + } } diff --git a/vendor/tokio/src/net/udp.rs b/vendor/tokio/src/net/udp.rs index 301a85cc0..8da377e16 100644 --- a/vendor/tokio/src/net/udp.rs +++ b/vendor/tokio/src/net/udp.rs @@ -1,7 +1,6 @@ use crate::io::{Interest, PollEvented, ReadBuf, Ready}; use crate::net::{to_socket_addrs, ToSocketAddrs}; -use std::convert::TryFrom; use std::fmt; use std::io; use std::net::{self, Ipv4Addr, Ipv6Addr, SocketAddr}; @@ -12,7 +11,7 @@ cfg_io_util! { } cfg_net! { - /// A UDP socket + /// A UDP socket. /// /// UDP is "connectionless", unlike TCP. Meaning, regardless of what address you've bound to, a `UdpSocket` /// is free to communicate with many different remotes. In tokio there are basically two main ways to use `UdpSocket`: @@ -128,6 +127,10 @@ impl UdpSocket { /// This function will create a new UDP socket and attempt to bind it to /// the `addr` provided. /// + /// Binding with a port number of 0 will request that the OS assigns a port + /// to this listener. The port allocated can be queried via the `local_addr` + /// method. + /// /// # Example /// /// ```no_run @@ -166,6 +169,7 @@ impl UdpSocket { UdpSocket::new(sys) } + #[track_caller] fn new(socket: mio::net::UdpSocket) -> io::Result<UdpSocket> { let io = PollEvented::new(socket)?; Ok(UdpSocket { io }) @@ -174,14 +178,21 @@ impl UdpSocket { /// Creates new `UdpSocket` from a previously bound `std::net::UdpSocket`. /// /// This function is intended to be used to wrap a UDP socket from the - /// standard library in the Tokio equivalent. The conversion assumes nothing - /// about the underlying socket; it is left up to the user to set it in - /// non-blocking mode. + /// standard library in the Tokio equivalent. /// /// This can be used in conjunction with socket2's `Socket` interface to /// configure a socket before it's handed off, such as setting options like /// `reuse_address` or binding to multiple addresses. /// + /// # Notes + /// + /// The caller is responsible for ensuring that the socket is in + /// non-blocking mode. Otherwise all I/O operations on the socket + /// will block the thread, which will cause unexpected behavior. + /// Non-blocking mode can be set using [`set_nonblocking`]. + /// + /// [`set_nonblocking`]: std::net::UdpSocket::set_nonblocking + /// /// # Panics /// /// This function panics if thread-local runtime is not set. @@ -206,12 +217,13 @@ impl UdpSocket { /// # Ok(()) /// # } /// ``` + #[track_caller] pub fn from_std(socket: net::UdpSocket) -> io::Result<UdpSocket> { let io = mio::net::UdpSocket::from_std(socket); UdpSocket::new(io) } - /// Turn a [`tokio::net::UdpSocket`] into a [`std::net::UdpSocket`]. + /// Turns a [`tokio::net::UdpSocket`] into a [`std::net::UdpSocket`]. /// /// The returned [`std::net::UdpSocket`] will have nonblocking mode set as /// `true`. Use [`set_nonblocking`] to change the blocking mode if needed. @@ -253,6 +265,10 @@ impl UdpSocket { } } + fn as_socket(&self) -> socket2::SockRef<'_> { + socket2::SockRef::from(self) + } + /// Returns the local address that this socket is bound to. /// /// # Example @@ -274,6 +290,28 @@ impl UdpSocket { self.io.local_addr() } + /// Returns the socket address of the remote peer this socket was connected to. + /// + /// # Example + /// + /// ``` + /// use tokio::net::UdpSocket; + /// + /// # use std::{io, net::SocketAddr}; + /// # #[tokio::main] + /// # async fn main() -> io::Result<()> { + /// let addr = "0.0.0.0:8080".parse::<SocketAddr>().unwrap(); + /// let peer = "127.0.0.1:11100".parse::<SocketAddr>().unwrap(); + /// let sock = UdpSocket::bind(addr).await?; + /// sock.connect(peer).await?; + /// assert_eq!(peer, sock.peer_addr()?); + /// # Ok(()) + /// # } + /// ``` + pub fn peer_addr(&self) -> io::Result<SocketAddr> { + self.io.peer_addr() + } + /// Connects the UDP socket setting the default destination for send() and /// limiting packets that are read via recv from the address specified in /// `addr`. @@ -317,7 +355,7 @@ impl UdpSocket { })) } - /// Wait for any of the requested ready states. + /// Waits for any of the requested ready states. /// /// This function is usually paired with `try_recv()` or `try_send()`. It /// can be used to concurrently recv / send to the same socket on a single @@ -325,7 +363,9 @@ impl UdpSocket { /// /// The function may complete without the socket being ready. This is a /// false-positive and attempting an operation will return with - /// `io::ErrorKind::WouldBlock`. + /// `io::ErrorKind::WouldBlock`. The function can also return with an empty + /// [`Ready`] set, so you should always check the returned value and possibly + /// wait again if the requested states are not set. /// /// # Cancel safety /// @@ -388,7 +428,7 @@ impl UdpSocket { Ok(event.ready) } - /// Wait for the socket to become writable. + /// Waits for the socket to become writable. /// /// This function is equivalent to `ready(Interest::WRITABLE)` and is /// usually paired with `try_send()` or `try_send_to()`. @@ -443,6 +483,39 @@ impl UdpSocket { Ok(()) } + /// Polls for write/send readiness. + /// + /// If the udp stream is not currently ready for sending, this method will + /// store a clone of the `Waker` from the provided `Context`. When the udp + /// stream becomes ready for sending, `Waker::wake` will be called on the + /// waker. + /// + /// Note that on multiple calls to `poll_send_ready` or `poll_send`, only + /// the `Waker` from the `Context` passed to the most recent call is + /// scheduled to receive a wakeup. (However, `poll_recv_ready` retains a + /// second, independent waker.) + /// + /// This function is intended for cases where creating and pinning a future + /// via [`writable`] is not feasible. Where possible, using [`writable`] is + /// preferred, as this supports polling from multiple tasks at once. + /// + /// # Return value + /// + /// The function returns: + /// + /// * `Poll::Pending` if the udp stream is not ready for writing. + /// * `Poll::Ready(Ok(()))` if the udp stream is ready for writing. + /// * `Poll::Ready(Err(e))` if an error is encountered. + /// + /// # Errors + /// + /// This function may encounter any standard I/O error except `WouldBlock`. + /// + /// [`writable`]: method@Self::writable + pub fn poll_send_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + self.io.registration().poll_write_ready(cx).map_ok(|_| ()) + } + /// Sends data on the socket to the remote address that the socket is /// connected to. /// @@ -516,7 +589,7 @@ impl UdpSocket { .poll_write_io(cx, || self.io.send(buf)) } - /// Try to send data on the socket to the remote address to which it is + /// Tries to send data on the socket to the remote address to which it is /// connected. /// /// When the socket buffer is full, `Err(io::ErrorKind::WouldBlock)` is @@ -570,7 +643,7 @@ impl UdpSocket { .try_io(Interest::WRITABLE, || self.io.send(buf)) } - /// Wait for the socket to become readable. + /// Waits for the socket to become readable. /// /// This function is equivalent to `ready(Interest::READABLE)` and is usually /// paired with `try_recv()`. @@ -630,6 +703,39 @@ impl UdpSocket { Ok(()) } + /// Polls for read/receive readiness. + /// + /// If the udp stream is not currently ready for receiving, this method will + /// store a clone of the `Waker` from the provided `Context`. When the udp + /// socket becomes ready for reading, `Waker::wake` will be called on the + /// waker. + /// + /// Note that on multiple calls to `poll_recv_ready`, `poll_recv` or + /// `poll_peek`, only the `Waker` from the `Context` passed to the most + /// recent call is scheduled to receive a wakeup. (However, + /// `poll_send_ready` retains a second, independent waker.) + /// + /// This function is intended for cases where creating and pinning a future + /// via [`readable`] is not feasible. Where possible, using [`readable`] is + /// preferred, as this supports polling from multiple tasks at once. + /// + /// # Return value + /// + /// The function returns: + /// + /// * `Poll::Pending` if the udp stream is not ready for reading. + /// * `Poll::Ready(Ok(()))` if the udp stream is ready for reading. + /// * `Poll::Ready(Err(e))` if an error is encountered. + /// + /// # Errors + /// + /// This function may encounter any standard I/O error except `WouldBlock`. + /// + /// [`readable`]: method@Self::readable + pub fn poll_recv_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + self.io.registration().poll_read_ready(cx).map_ok(|_| ()) + } + /// Receives a single datagram message on the socket from the remote address /// to which it is connected. On success, returns the number of bytes read. /// @@ -642,7 +748,7 @@ impl UdpSocket { /// /// # Cancel safety /// - /// This method is cancel safe. If `recv_from` is used as the event in a + /// This method is cancel safe. If `recv` is used as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, it is guaranteed that no messages were received on this /// socket. @@ -715,11 +821,11 @@ impl UdpSocket { Poll::Ready(Ok(())) } - /// Try to receive a single datagram message on the socket from the remote + /// Tries to receive a single datagram message on the socket from the remote /// address to which it is connected. On success, returns the number of /// bytes read. /// - /// The function must be called with valid byte array buf of sufficient size + /// This method must be called with valid byte array buf of sufficient size /// to hold the message bytes. If a message is too long to fit in the /// supplied buffer, excess bytes may be discarded. /// @@ -772,13 +878,15 @@ impl UdpSocket { } cfg_io_util! { - /// Try to receive data from the stream into the provided buffer, advancing the + /// Tries to receive data from the stream into the provided buffer, advancing the /// buffer's internal cursor, returning how many bytes were read. /// - /// The function must be called with valid byte array buf of sufficient size + /// This method must be called with valid byte array buf of sufficient size /// to hold the message bytes. If a message is too long to fit in the /// supplied buffer, excess bytes may be discarded. /// + /// This method can be used even if `buf` is uninitialized. + /// /// When there is no pending data, `Err(io::ErrorKind::WouldBlock)` is /// returned. This function is usually paired with `readable()`. /// @@ -825,10 +933,10 @@ impl UdpSocket { let dst = unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]) }; + let n = (*self.io).recv(dst)?; + // Safety: We trust `UdpSocket::recv` to have filled up `n` bytes in the // buffer. - let n = (&*self.io).recv(dst)?; - unsafe { buf.advance_mut(n); } @@ -837,16 +945,75 @@ impl UdpSocket { }) } - /// Try to receive a single datagram message on the socket. On success, + /// Receives a single datagram message on the socket from the remote address + /// to which it is connected, advancing the buffer's internal cursor, + /// returning how many bytes were read. + /// + /// This method must be called with valid byte array buf of sufficient size + /// to hold the message bytes. If a message is too long to fit in the + /// supplied buffer, excess bytes may be discarded. + /// + /// This method can be used even if `buf` is uninitialized. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::UdpSocket; + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// // Connect to a peer + /// let socket = UdpSocket::bind("127.0.0.1:8080").await?; + /// socket.connect("127.0.0.1:8081").await?; + /// + /// let mut buf = Vec::with_capacity(512); + /// let len = socket.recv_buf(&mut buf).await?; + /// + /// println!("received {} bytes {:?}", len, &buf[..len]); + /// + /// Ok(()) + /// } + /// ``` + pub async fn recv_buf<B: BufMut>(&self, buf: &mut B) -> io::Result<usize> { + self.io.registration().async_io(Interest::READABLE, || { + let dst = buf.chunk_mut(); + let dst = + unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]) }; + + let n = (*self.io).recv(dst)?; + + // Safety: We trust `UdpSocket::recv` to have filled up `n` bytes in the + // buffer. + unsafe { + buf.advance_mut(n); + } + + Ok(n) + }).await + } + + /// Tries to receive a single datagram message on the socket. On success, /// returns the number of bytes read and the origin. /// - /// The function must be called with valid byte array buf of sufficient size + /// This method must be called with valid byte array buf of sufficient size /// to hold the message bytes. If a message is too long to fit in the /// supplied buffer, excess bytes may be discarded. /// + /// This method can be used even if `buf` is uninitialized. + /// /// When there is no pending data, `Err(io::ErrorKind::WouldBlock)` is /// returned. This function is usually paired with `readable()`. /// + /// # Notes + /// Note that the socket address **cannot** be implicitly trusted, because it is relatively + /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. + /// Because UDP is stateless and does not validate the origin of a packet, + /// the attacker does not need to be able to intercept traffic in order to interfere. + /// It is important to be aware of this when designing your application-level protocol. + /// + /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection + /// /// # Examples /// /// ```no_run @@ -889,10 +1056,10 @@ impl UdpSocket { let dst = unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]) }; + let (n, addr) = (*self.io).recv_from(dst)?; + // Safety: We trust `UdpSocket::recv_from` to have filled up `n` bytes in the // buffer. - let (n, addr) = (&*self.io).recv_from(dst)?; - unsafe { buf.advance_mut(n); } @@ -900,6 +1067,62 @@ impl UdpSocket { Ok((n, addr)) }) } + + /// Receives a single datagram message on the socket, advancing the + /// buffer's internal cursor, returning how many bytes were read and the origin. + /// + /// This method must be called with valid byte array buf of sufficient size + /// to hold the message bytes. If a message is too long to fit in the + /// supplied buffer, excess bytes may be discarded. + /// + /// This method can be used even if `buf` is uninitialized. + /// + /// # Notes + /// Note that the socket address **cannot** be implicitly trusted, because it is relatively + /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. + /// Because UDP is stateless and does not validate the origin of a packet, + /// the attacker does not need to be able to intercept traffic in order to interfere. + /// It is important to be aware of this when designing your application-level protocol. + /// + /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::UdpSocket; + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// // Connect to a peer + /// let socket = UdpSocket::bind("127.0.0.1:8080").await?; + /// socket.connect("127.0.0.1:8081").await?; + /// + /// let mut buf = Vec::with_capacity(512); + /// let (len, addr) = socket.recv_buf_from(&mut buf).await?; + /// + /// println!("received {:?} bytes from {:?}", len, addr); + /// + /// Ok(()) + /// } + /// ``` + pub async fn recv_buf_from<B: BufMut>(&self, buf: &mut B) -> io::Result<(usize, SocketAddr)> { + self.io.registration().async_io(Interest::READABLE, || { + let dst = buf.chunk_mut(); + let dst = + unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]) }; + + let (n, addr) = (*self.io).recv_from(dst)?; + + // Safety: We trust `UdpSocket::recv_from` to have filled up `n` bytes in the + // buffer. + unsafe { + buf.advance_mut(n); + } + + Ok((n,addr)) + }).await + } } /// Sends data on the socket to the given address. On success, returns the @@ -978,7 +1201,7 @@ impl UdpSocket { .poll_write_io(cx, || self.io.send_to(buf, target)) } - /// Try to send data on the socket to the given address, but if the send is + /// Tries to send data on the socket to the given address, but if the send is /// blocked this will return right away. /// /// This function is usually paired with `writable()`. @@ -1070,6 +1293,15 @@ impl UdpSocket { /// Ok(()) /// } /// ``` + /// + /// # Notes + /// Note that the socket address **cannot** be implicitly trusted, because it is relatively + /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. + /// Because UDP is stateless and does not validate the origin of a packet, + /// the attacker does not need to be able to intercept traffic in order to interfere. + /// It is important to be aware of this when designing your application-level protocol. + /// + /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection pub async fn recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> { self.io .registration() @@ -1094,6 +1326,15 @@ impl UdpSocket { /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. + /// + /// # Notes + /// Note that the socket address **cannot** be implicitly trusted, because it is relatively + /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. + /// Because UDP is stateless and does not validate the origin of a packet, + /// the attacker does not need to be able to intercept traffic in order to interfere. + /// It is important to be aware of this when designing your application-level protocol. + /// + /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection pub fn poll_recv_from( &self, cx: &mut Context<'_>, @@ -1116,16 +1357,26 @@ impl UdpSocket { Poll::Ready(Ok(addr)) } - /// Try to receive a single datagram message on the socket. On success, + /// Tries to receive a single datagram message on the socket. On success, /// returns the number of bytes read and the origin. /// - /// The function must be called with valid byte array buf of sufficient size + /// This method must be called with valid byte array buf of sufficient size /// to hold the message bytes. If a message is too long to fit in the /// supplied buffer, excess bytes may be discarded. /// /// When there is no pending data, `Err(io::ErrorKind::WouldBlock)` is /// returned. This function is usually paired with `readable()`. /// + /// # Notes + /// + /// Note that the socket address **cannot** be implicitly trusted, because it is relatively + /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. + /// Because UDP is stateless and does not validate the origin of a packet, + /// the attacker does not need to be able to intercept traffic in order to interfere. + /// It is important to be aware of this when designing your application-level protocol. + /// + /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection + /// /// # Examples /// /// ```no_run @@ -1170,6 +1421,84 @@ impl UdpSocket { .try_io(Interest::READABLE, || self.io.recv_from(buf)) } + /// Tries to read or write from the socket using a user-provided IO operation. + /// + /// If the socket is ready, the provided closure is called. The closure + /// should attempt to perform IO operation on the socket by manually + /// calling the appropriate syscall. If the operation fails because the + /// socket is not actually ready, then the closure should return a + /// `WouldBlock` error and the readiness flag is cleared. The return value + /// of the closure is then returned by `try_io`. + /// + /// If the socket is not ready, then the closure is not called + /// and a `WouldBlock` error is returned. + /// + /// The closure should only return a `WouldBlock` error if it has performed + /// an IO operation on the socket that failed due to the socket not being + /// ready. Returning a `WouldBlock` error in any other situation will + /// incorrectly clear the readiness flag, which can cause the socket to + /// behave incorrectly. + /// + /// The closure should not perform the IO operation using any of the methods + /// defined on the Tokio `UdpSocket` type, as this will mess with the + /// readiness flag and can cause the socket to behave incorrectly. + /// + /// This method is not intended to be used with combined interests. + /// The closure should perform only one type of IO operation, so it should not + /// require more than one ready state. This method may panic or sleep forever + /// if it is called with a combined interest. + /// + /// Usually, [`readable()`], [`writable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: UdpSocket::readable() + /// [`writable()`]: UdpSocket::writable() + /// [`ready()`]: UdpSocket::ready() + pub fn try_io<R>( + &self, + interest: Interest, + f: impl FnOnce() -> io::Result<R>, + ) -> io::Result<R> { + self.io + .registration() + .try_io(interest, || self.io.try_io(f)) + } + + /// Reads or writes from the socket using a user-provided IO operation. + /// + /// The readiness of the socket is awaited and when the socket is ready, + /// the provided closure is called. The closure should attempt to perform + /// IO operation on the socket by manually calling the appropriate syscall. + /// If the operation fails because the socket is not actually ready, + /// then the closure should return a `WouldBlock` error. In such case the + /// readiness flag is cleared and the socket readiness is awaited again. + /// This loop is repeated until the closure returns an `Ok` or an error + /// other than `WouldBlock`. + /// + /// The closure should only return a `WouldBlock` error if it has performed + /// an IO operation on the socket that failed due to the socket not being + /// ready. Returning a `WouldBlock` error in any other situation will + /// incorrectly clear the readiness flag, which can cause the socket to + /// behave incorrectly. + /// + /// The closure should not perform the IO operation using any of the methods + /// defined on the Tokio `UdpSocket` type, as this will mess with the + /// readiness flag and can cause the socket to behave incorrectly. + /// + /// This method is not intended to be used with combined interests. + /// The closure should perform only one type of IO operation, so it should not + /// require more than one ready state. This method may panic or sleep forever + /// if it is called with a combined interest. + pub async fn async_io<R>( + &self, + interest: Interest, + mut f: impl FnMut() -> io::Result<R>, + ) -> io::Result<R> { + self.io + .registration() + .async_io(interest, || self.io.try_io(&mut f)) + .await + } + /// Receives data from the socket, without removing it from the input queue. /// On success, returns the number of bytes read and the address from whence /// the data came. @@ -1182,6 +1511,17 @@ impl UdpSocket { /// Make sure to always use a sufficiently large buffer to hold the /// maximum UDP packet size, which can be up to 65536 bytes in size. /// + /// MacOS will return an error if you pass a zero-sized buffer. + /// + /// If you're merely interested in learning the sender of the data at the head of the queue, + /// try [`peek_sender`]. + /// + /// Note that the socket address **cannot** be implicitly trusted, because it is relatively + /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. + /// Because UDP is stateless and does not validate the origin of a packet, + /// the attacker does not need to be able to intercept traffic in order to interfere. + /// It is important to be aware of this when designing your application-level protocol. + /// /// # Examples /// /// ```no_run @@ -1200,6 +1540,9 @@ impl UdpSocket { /// Ok(()) /// } /// ``` + /// + /// [`peek_sender`]: method@Self::peek_sender + /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection pub async fn peek_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> { self.io .registration() @@ -1208,7 +1551,7 @@ impl UdpSocket { } /// Receives data from the socket, without removing it from the input queue. - /// On success, returns the number of bytes read. + /// On success, returns the sending address of the datagram. /// /// # Notes /// @@ -1222,6 +1565,17 @@ impl UdpSocket { /// Make sure to always use a sufficiently large buffer to hold the /// maximum UDP packet size, which can be up to 65536 bytes in size. /// + /// MacOS will return an error if you pass a zero-sized buffer. + /// + /// If you're merely interested in learning the sender of the data at the head of the queue, + /// try [`poll_peek_sender`]. + /// + /// Note that the socket address **cannot** be implicitly trusted, because it is relatively + /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. + /// Because UDP is stateless and does not validate the origin of a packet, + /// the attacker does not need to be able to intercept traffic in order to interfere. + /// It is important to be aware of this when designing your application-level protocol. + /// /// # Return value /// /// The function returns: @@ -1233,6 +1587,9 @@ impl UdpSocket { /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. + /// + /// [`poll_peek_sender`]: method@Self::poll_peek_sender + /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection pub fn poll_peek_from( &self, cx: &mut Context<'_>, @@ -1255,6 +1612,117 @@ impl UdpSocket { Poll::Ready(Ok(addr)) } + /// Tries to receive data on the socket without removing it from the input queue. + /// On success, returns the number of bytes read and the sending address of the + /// datagram. + /// + /// When there is no pending data, `Err(io::ErrorKind::WouldBlock)` is + /// returned. This function is usually paired with `readable()`. + /// + /// # Notes + /// + /// On Windows, if the data is larger than the buffer specified, the buffer + /// is filled with the first part of the data, and peek returns the error + /// WSAEMSGSIZE(10040). The excess data is lost. + /// Make sure to always use a sufficiently large buffer to hold the + /// maximum UDP packet size, which can be up to 65536 bytes in size. + /// + /// MacOS will return an error if you pass a zero-sized buffer. + /// + /// If you're merely interested in learning the sender of the data at the head of the queue, + /// try [`try_peek_sender`]. + /// + /// Note that the socket address **cannot** be implicitly trusted, because it is relatively + /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. + /// Because UDP is stateless and does not validate the origin of a packet, + /// the attacker does not need to be able to intercept traffic in order to interfere. + /// It is important to be aware of this when designing your application-level protocol. + /// + /// [`try_peek_sender`]: method@Self::try_peek_sender + /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection + pub fn try_peek_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> { + self.io + .registration() + .try_io(Interest::READABLE, || self.io.peek_from(buf)) + } + + /// Retrieve the sender of the data at the head of the input queue, waiting if empty. + /// + /// This is equivalent to calling [`peek_from`] with a zero-sized buffer, + /// but suppresses the `WSAEMSGSIZE` error on Windows and the "invalid argument" error on macOS. + /// + /// Note that the socket address **cannot** be implicitly trusted, because it is relatively + /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. + /// Because UDP is stateless and does not validate the origin of a packet, + /// the attacker does not need to be able to intercept traffic in order to interfere. + /// It is important to be aware of this when designing your application-level protocol. + /// + /// [`peek_from`]: method@Self::peek_from + /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection + pub async fn peek_sender(&self) -> io::Result<SocketAddr> { + self.io + .registration() + .async_io(Interest::READABLE, || self.peek_sender_inner()) + .await + } + + /// Retrieve the sender of the data at the head of the input queue, + /// scheduling a wakeup if empty. + /// + /// This is equivalent to calling [`poll_peek_from`] with a zero-sized buffer, + /// but suppresses the `WSAEMSGSIZE` error on Windows and the "invalid argument" error on macOS. + /// + /// # Notes + /// + /// Note that on multiple calls to a `poll_*` method in the recv direction, only the + /// `Waker` from the `Context` passed to the most recent call will be scheduled to + /// receive a wakeup. + /// + /// Note that the socket address **cannot** be implicitly trusted, because it is relatively + /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. + /// Because UDP is stateless and does not validate the origin of a packet, + /// the attacker does not need to be able to intercept traffic in order to interfere. + /// It is important to be aware of this when designing your application-level protocol. + /// + /// [`poll_peek_from`]: method@Self::poll_peek_from + /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection + pub fn poll_peek_sender(&self, cx: &mut Context<'_>) -> Poll<io::Result<SocketAddr>> { + self.io + .registration() + .poll_read_io(cx, || self.peek_sender_inner()) + } + + /// Try to retrieve the sender of the data at the head of the input queue. + /// + /// When there is no pending data, `Err(io::ErrorKind::WouldBlock)` is + /// returned. This function is usually paired with `readable()`. + /// + /// Note that the socket address **cannot** be implicitly trusted, because it is relatively + /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. + /// Because UDP is stateless and does not validate the origin of a packet, + /// the attacker does not need to be able to intercept traffic in order to interfere. + /// It is important to be aware of this when designing your application-level protocol. + /// + /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection + pub fn try_peek_sender(&self) -> io::Result<SocketAddr> { + self.io + .registration() + .try_io(Interest::READABLE, || self.peek_sender_inner()) + } + + #[inline] + fn peek_sender_inner(&self) -> io::Result<SocketAddr> { + self.io.try_io(|| { + self.as_socket() + .peek_sender()? + // May be `None` if the platform doesn't populate the sender for some reason. + // In testing, that only occurred on macOS if you pass a zero-sized buffer, + // but the implementation of `Socket::peek_sender()` covers that. + .as_socket() + .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "sender not available")) + }) + } + /// Gets the value of the `SO_BROADCAST` option for this socket. /// /// For more information about this option, see [`set_broadcast`]. @@ -1379,6 +1847,89 @@ impl UdpSocket { self.io.set_ttl(ttl) } + /// Gets the value of the `IP_TOS` option for this socket. + /// + /// For more information about this option, see [`set_tos`]. + /// + /// **NOTE:** On Windows, `IP_TOS` is only supported on [Windows 8+ or + /// Windows Server 2012+.](https://docs.microsoft.com/en-us/windows/win32/winsock/ipproto-ip-socket-options) + /// + /// [`set_tos`]: Self::set_tos + // https://docs.rs/socket2/0.4.2/src/socket2/socket.rs.html#1178 + #[cfg(not(any( + target_os = "fuchsia", + target_os = "redox", + target_os = "solaris", + target_os = "illumos", + )))] + #[cfg_attr( + docsrs, + doc(cfg(not(any( + target_os = "fuchsia", + target_os = "redox", + target_os = "solaris", + target_os = "illumos", + )))) + )] + pub fn tos(&self) -> io::Result<u32> { + self.as_socket().tos() + } + + /// Sets the value for the `IP_TOS` option on this socket. + /// + /// This value sets the type-of-service field that is used in every packet + /// sent from this socket. + /// + /// **NOTE:** On Windows, `IP_TOS` is only supported on [Windows 8+ or + /// Windows Server 2012+.](https://docs.microsoft.com/en-us/windows/win32/winsock/ipproto-ip-socket-options) + // https://docs.rs/socket2/0.4.2/src/socket2/socket.rs.html#1178 + #[cfg(not(any( + target_os = "fuchsia", + target_os = "redox", + target_os = "solaris", + target_os = "illumos", + )))] + #[cfg_attr( + docsrs, + doc(cfg(not(any( + target_os = "fuchsia", + target_os = "redox", + target_os = "solaris", + target_os = "illumos", + )))) + )] + pub fn set_tos(&self, tos: u32) -> io::Result<()> { + self.as_socket().set_tos(tos) + } + + /// Gets the value for the `SO_BINDTODEVICE` option on this socket + /// + /// This value gets the socket-bound device's interface name. + #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux",))] + #[cfg_attr( + docsrs, + doc(cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux",))) + )] + pub fn device(&self) -> io::Result<Option<Vec<u8>>> { + self.as_socket().device() + } + + /// Sets the value for the `SO_BINDTODEVICE` option on this socket + /// + /// If a socket is bound to an interface, only packets received from that + /// particular interface are processed by the socket. Note that this only + /// works for some socket types, particularly `AF_INET` sockets. + /// + /// If `interface` is `None` or an empty string it removes the binding. + #[cfg(all(any(target_os = "android", target_os = "fuchsia", target_os = "linux")))] + #[cfg_attr( + docsrs, + doc(cfg(all(any(target_os = "android", target_os = "fuchsia", target_os = "linux")))) + )] + pub fn bind_device(&self, interface: Option<&[u8]>) -> io::Result<()> { + self.as_socket().bind_device(interface) + } + /// Executes an operation of the `IP_ADD_MEMBERSHIP` type. /// /// This function specifies a new multicast group for this socket to join. @@ -1459,7 +2010,7 @@ impl fmt::Debug for UdpSocket { } } -#[cfg(all(unix))] +#[cfg(unix)] mod sys { use super::UdpSocket; use std::os::unix::prelude::*; @@ -1469,16 +2020,30 @@ mod sys { self.io.as_raw_fd() } } + + #[cfg(not(tokio_no_as_fd))] + impl AsFd for UdpSocket { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } + } } -#[cfg(windows)] -mod sys { - use super::UdpSocket; - use std::os::windows::prelude::*; +cfg_windows! { + use crate::os::windows::io::{AsRawSocket, RawSocket}; + #[cfg(not(tokio_no_as_fd))] + use crate::os::windows::io::{AsSocket, BorrowedSocket}; impl AsRawSocket for UdpSocket { fn as_raw_socket(&self) -> RawSocket { self.io.as_raw_socket() } } + + #[cfg(not(tokio_no_as_fd))] + impl AsSocket for UdpSocket { + fn as_socket(&self) -> BorrowedSocket<'_> { + unsafe { BorrowedSocket::borrow_raw(self.as_raw_socket()) } + } + } } diff --git a/vendor/tokio/src/net/unix/datagram/socket.rs b/vendor/tokio/src/net/unix/datagram/socket.rs index 2d2177803..c97d101aa 100644 --- a/vendor/tokio/src/net/unix/datagram/socket.rs +++ b/vendor/tokio/src/net/unix/datagram/socket.rs @@ -1,10 +1,11 @@ use crate::io::{Interest, PollEvented, ReadBuf, Ready}; use crate::net::unix::SocketAddr; -use std::convert::TryFrom; use std::fmt; use std::io; use std::net::Shutdown; +#[cfg(not(tokio_no_as_fd))] +use std::os::unix::io::{AsFd, BorrowedFd}; use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd}; use std::os::unix::net; use std::path::Path; @@ -90,13 +91,14 @@ cfg_net_unix! { /// # Ok(()) /// # } /// ``` + #[cfg_attr(docsrs, doc(alias = "uds"))] pub struct UnixDatagram { io: PollEvented<mio::net::UnixDatagram>, } } impl UnixDatagram { - /// Wait for any of the requested ready states. + /// Waits for any of the requested ready states. /// /// This function is usually paired with `try_recv()` or `try_send()`. It /// can be used to concurrently recv / send to the same socket on a single @@ -104,7 +106,9 @@ impl UnixDatagram { /// /// The function may complete without the socket being ready. This is a /// false-positive and attempting an operation will return with - /// `io::ErrorKind::WouldBlock`. + /// `io::ErrorKind::WouldBlock`. The function can also return with an empty + /// [`Ready`] set, so you should always check the returned value and possibly + /// wait again if the requested states are not set. /// /// # Cancel safety /// @@ -169,7 +173,7 @@ impl UnixDatagram { Ok(event.ready) } - /// Wait for the socket to become writable. + /// Waits for the socket to become writable. /// /// This function is equivalent to `ready(Interest::WRITABLE)` and is /// usually paired with `try_send()` or `try_send_to()`. @@ -226,7 +230,40 @@ impl UnixDatagram { Ok(()) } - /// Wait for the socket to become readable. + /// Polls for write/send readiness. + /// + /// If the socket is not currently ready for sending, this method will + /// store a clone of the `Waker` from the provided `Context`. When the socket + /// becomes ready for sending, `Waker::wake` will be called on the + /// waker. + /// + /// Note that on multiple calls to `poll_send_ready` or `poll_send`, only + /// the `Waker` from the `Context` passed to the most recent call is + /// scheduled to receive a wakeup. (However, `poll_recv_ready` retains a + /// second, independent waker.) + /// + /// This function is intended for cases where creating and pinning a future + /// via [`writable`] is not feasible. Where possible, using [`writable`] is + /// preferred, as this supports polling from multiple tasks at once. + /// + /// # Return value + /// + /// The function returns: + /// + /// * `Poll::Pending` if the socket is not ready for writing. + /// * `Poll::Ready(Ok(()))` if the socket is ready for writing. + /// * `Poll::Ready(Err(e))` if an error is encountered. + /// + /// # Errors + /// + /// This function may encounter any standard I/O error except `WouldBlock`. + /// + /// [`writable`]: method@Self::writable + pub fn poll_send_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + self.io.registration().poll_write_ready(cx).map_ok(|_| ()) + } + + /// Waits for the socket to become readable. /// /// This function is equivalent to `ready(Interest::READABLE)` and is usually /// paired with `try_recv()`. @@ -289,6 +326,39 @@ impl UnixDatagram { Ok(()) } + /// Polls for read/receive readiness. + /// + /// If the socket is not currently ready for receiving, this method will + /// store a clone of the `Waker` from the provided `Context`. When the + /// socket becomes ready for reading, `Waker::wake` will be called on the + /// waker. + /// + /// Note that on multiple calls to `poll_recv_ready`, `poll_recv` or + /// `poll_peek`, only the `Waker` from the `Context` passed to the most + /// recent call is scheduled to receive a wakeup. (However, + /// `poll_send_ready` retains a second, independent waker.) + /// + /// This function is intended for cases where creating and pinning a future + /// via [`readable`] is not feasible. Where possible, using [`readable`] is + /// preferred, as this supports polling from multiple tasks at once. + /// + /// # Return value + /// + /// The function returns: + /// + /// * `Poll::Pending` if the socket is not ready for reading. + /// * `Poll::Ready(Ok(()))` if the socket is ready for reading. + /// * `Poll::Ready(Err(e))` if an error is encountered. + /// + /// # Errors + /// + /// This function may encounter any standard I/O error except `WouldBlock`. + /// + /// [`readable`]: method@Self::readable + pub fn poll_recv_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + self.io.registration().poll_read_ready(cx).map_ok(|_| ()) + } + /// Creates a new `UnixDatagram` bound to the specified path. /// /// # Examples @@ -358,13 +428,21 @@ impl UnixDatagram { /// Creates new `UnixDatagram` from a `std::os::unix::net::UnixDatagram`. /// /// This function is intended to be used to wrap a UnixDatagram from the - /// standard library in the Tokio equivalent. The conversion assumes - /// nothing about the underlying datagram; it is left up to the user to set - /// it in non-blocking mode. + /// standard library in the Tokio equivalent. + /// + /// # Notes + /// + /// The caller is responsible for ensuring that the socker is in + /// non-blocking mode. Otherwise all I/O operations on the socket + /// will block the thread, which will cause unexpected behavior. + /// Non-blocking mode can be set using [`set_nonblocking`]. + /// + /// [`set_nonblocking`]: std::os::unix::net::UnixDatagram::set_nonblocking /// /// # Panics /// - /// This function panics if thread-local runtime is not set. + /// This function panics if it is not called from within a runtime with + /// IO enabled. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a Tokio runtime, otherwise runtime can be set @@ -391,30 +469,29 @@ impl UnixDatagram { /// # Ok(()) /// # } /// ``` + #[track_caller] pub fn from_std(datagram: net::UnixDatagram) -> io::Result<UnixDatagram> { let socket = mio::net::UnixDatagram::from_std(datagram); let io = PollEvented::new(socket)?; Ok(UnixDatagram { io }) } - /// Turn a [`tokio::net::UnixDatagram`] into a [`std::os::unix::net::UnixDatagram`]. + /// Turns a [`tokio::net::UnixDatagram`] into a [`std::os::unix::net::UnixDatagram`]. /// /// The returned [`std::os::unix::net::UnixDatagram`] will have nonblocking - /// mode set as `true`. Use [`set_nonblocking`] to change the blocking mode + /// mode set as `true`. Use [`set_nonblocking`] to change the blocking mode /// if needed. /// /// # Examples /// /// ```rust,no_run - /// use std::error::Error; - /// - /// #[tokio::main] - /// async fn main() -> Result<(), Box<dyn Error>> { - /// let tokio_socket = tokio::net::UnixDatagram::bind("127.0.0.1:0")?; - /// let std_socket = tokio_socket.into_std()?; - /// std_socket.set_nonblocking(false)?; - /// Ok(()) - /// } + /// # use std::error::Error; + /// # async fn dox() -> Result<(), Box<dyn Error>> { + /// let tokio_socket = tokio::net::UnixDatagram::bind("/path/to/the/socket")?; + /// let std_socket = tokio_socket.into_std()?; + /// std_socket.set_nonblocking(false)?; + /// # Ok(()) + /// # } /// ``` /// /// [`tokio::net::UnixDatagram`]: UnixDatagram @@ -548,7 +625,7 @@ impl UnixDatagram { .await } - /// Try to send a datagram to the peer without waiting. + /// Tries to send a datagram to the peer without waiting. /// /// # Examples /// @@ -592,7 +669,7 @@ impl UnixDatagram { .try_io(Interest::WRITABLE, || self.io.send(buf)) } - /// Try to send a datagram to the peer without waiting. + /// Tries to send a datagram to the peer without waiting. /// /// # Examples /// @@ -678,7 +755,7 @@ impl UnixDatagram { .await } - /// Try to receive a datagram from the peer without waiting. + /// Tries to receive a datagram from the peer without waiting. /// /// # Examples /// @@ -729,7 +806,9 @@ impl UnixDatagram { } cfg_io_util! { - /// Try to receive data from the socket without waiting. + /// Tries to receive data from the socket without waiting. + /// + /// This method can be used even if `buf` is uninitialized. /// /// # Examples /// @@ -778,7 +857,7 @@ impl UnixDatagram { // Safety: We trust `UnixDatagram::recv_from` to have filled up `n` bytes in the // buffer. - let (n, addr) = (&*self.io).recv_from(dst)?; + let (n, addr) = (*self.io).recv_from(dst)?; unsafe { buf.advance_mut(n); @@ -790,9 +869,64 @@ impl UnixDatagram { Ok((n, SocketAddr(addr))) } - /// Try to read data from the stream into the provided buffer, advancing the + /// Receives from the socket, advances the + /// buffer's internal cursor and returns how many bytes were read and the origin. + /// + /// This method can be used even if `buf` is uninitialized. + /// + /// # Examples + /// ``` + /// # use std::error::Error; + /// # #[tokio::main] + /// # async fn main() -> Result<(), Box<dyn Error>> { + /// use tokio::net::UnixDatagram; + /// use tempfile::tempdir; + /// + /// // We use a temporary directory so that the socket + /// // files left by the bound sockets will get cleaned up. + /// let tmp = tempdir()?; + /// + /// // Bind each socket to a filesystem path + /// let tx_path = tmp.path().join("tx"); + /// let tx = UnixDatagram::bind(&tx_path)?; + /// let rx_path = tmp.path().join("rx"); + /// let rx = UnixDatagram::bind(&rx_path)?; + /// + /// let bytes = b"hello world"; + /// tx.send_to(bytes, &rx_path).await?; + /// + /// let mut buf = Vec::with_capacity(24); + /// let (size, addr) = rx.recv_buf_from(&mut buf).await?; + /// + /// let dgram = &buf[..size]; + /// assert_eq!(dgram, bytes); + /// assert_eq!(addr.as_pathname().unwrap(), &tx_path); + /// + /// # Ok(()) + /// # } + /// ``` + pub async fn recv_buf_from<B: BufMut>(&self, buf: &mut B) -> io::Result<(usize, SocketAddr)> { + self.io.registration().async_io(Interest::READABLE, || { + let dst = buf.chunk_mut(); + let dst = + unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]) }; + + // Safety: We trust `UnixDatagram::recv_from` to have filled up `n` bytes in the + // buffer. + let (n, addr) = (*self.io).recv_from(dst)?; + + unsafe { + buf.advance_mut(n); + } + Ok((n,SocketAddr(addr))) + }).await + } + + /// Tries to read data from the stream into the provided buffer, advancing the /// buffer's internal cursor, returning how many bytes were read. /// + /// This method can be used even if `buf` is uninitialized. + /// /// # Examples /// /// ```no_run @@ -841,7 +975,7 @@ impl UnixDatagram { // Safety: We trust `UnixDatagram::recv` to have filled up `n` bytes in the // buffer. - let n = (&*self.io).recv(dst)?; + let n = (*self.io).recv(dst)?; unsafe { buf.advance_mut(n); @@ -850,6 +984,52 @@ impl UnixDatagram { Ok(n) }) } + + /// Receives data from the socket from the address to which it is connected, + /// advancing the buffer's internal cursor, returning how many bytes were read. + /// + /// This method can be used even if `buf` is uninitialized. + /// + /// # Examples + /// ``` + /// # use std::error::Error; + /// # #[tokio::main] + /// # async fn main() -> Result<(), Box<dyn Error>> { + /// use tokio::net::UnixDatagram; + /// + /// // Create the pair of sockets + /// let (sock1, sock2) = UnixDatagram::pair()?; + /// + /// // Since the sockets are paired, the paired send/recv + /// // functions can be used + /// let bytes = b"hello world"; + /// sock1.send(bytes).await?; + /// + /// let mut buff = Vec::with_capacity(24); + /// let size = sock2.recv_buf(&mut buff).await?; + /// + /// let dgram = &buff[..size]; + /// assert_eq!(dgram, bytes); + /// + /// # Ok(()) + /// # } + /// ``` + pub async fn recv_buf<B: BufMut>(&self, buf: &mut B) -> io::Result<usize> { + self.io.registration().async_io(Interest::READABLE, || { + let dst = buf.chunk_mut(); + let dst = + unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]) }; + + // Safety: We trust `UnixDatagram::recv_from` to have filled up `n` bytes in the + // buffer. + let n = (*self.io).recv(dst)?; + + unsafe { + buf.advance_mut(n); + } + Ok(n) + }).await + } } /// Sends data on the socket to the specified address. @@ -1091,7 +1271,7 @@ impl UnixDatagram { Poll::Ready(Ok(())) } - /// Try to receive data from the socket without waiting. + /// Tries to receive data from the socket without waiting. /// /// # Examples /// @@ -1143,6 +1323,84 @@ impl UnixDatagram { Ok((n, SocketAddr(addr))) } + /// Tries to read or write from the socket using a user-provided IO operation. + /// + /// If the socket is ready, the provided closure is called. The closure + /// should attempt to perform IO operation on the socket by manually + /// calling the appropriate syscall. If the operation fails because the + /// socket is not actually ready, then the closure should return a + /// `WouldBlock` error and the readiness flag is cleared. The return value + /// of the closure is then returned by `try_io`. + /// + /// If the socket is not ready, then the closure is not called + /// and a `WouldBlock` error is returned. + /// + /// The closure should only return a `WouldBlock` error if it has performed + /// an IO operation on the socket that failed due to the socket not being + /// ready. Returning a `WouldBlock` error in any other situation will + /// incorrectly clear the readiness flag, which can cause the socket to + /// behave incorrectly. + /// + /// The closure should not perform the IO operation using any of the methods + /// defined on the Tokio `UnixDatagram` type, as this will mess with the + /// readiness flag and can cause the socket to behave incorrectly. + /// + /// This method is not intended to be used with combined interests. + /// The closure should perform only one type of IO operation, so it should not + /// require more than one ready state. This method may panic or sleep forever + /// if it is called with a combined interest. + /// + /// Usually, [`readable()`], [`writable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: UnixDatagram::readable() + /// [`writable()`]: UnixDatagram::writable() + /// [`ready()`]: UnixDatagram::ready() + pub fn try_io<R>( + &self, + interest: Interest, + f: impl FnOnce() -> io::Result<R>, + ) -> io::Result<R> { + self.io + .registration() + .try_io(interest, || self.io.try_io(f)) + } + + /// Reads or writes from the socket using a user-provided IO operation. + /// + /// The readiness of the socket is awaited and when the socket is ready, + /// the provided closure is called. The closure should attempt to perform + /// IO operation on the socket by manually calling the appropriate syscall. + /// If the operation fails because the socket is not actually ready, + /// then the closure should return a `WouldBlock` error. In such case the + /// readiness flag is cleared and the socket readiness is awaited again. + /// This loop is repeated until the closure returns an `Ok` or an error + /// other than `WouldBlock`. + /// + /// The closure should only return a `WouldBlock` error if it has performed + /// an IO operation on the socket that failed due to the socket not being + /// ready. Returning a `WouldBlock` error in any other situation will + /// incorrectly clear the readiness flag, which can cause the socket to + /// behave incorrectly. + /// + /// The closure should not perform the IO operation using any of the methods + /// defined on the Tokio `UnixDatagram` type, as this will mess with the + /// readiness flag and can cause the socket to behave incorrectly. + /// + /// This method is not intended to be used with combined interests. + /// The closure should perform only one type of IO operation, so it should not + /// require more than one ready state. This method may panic or sleep forever + /// if it is called with a combined interest. + pub async fn async_io<R>( + &self, + interest: Interest, + mut f: impl FnMut() -> io::Result<R>, + ) -> io::Result<R> { + self.io + .registration() + .async_io(interest, || self.io.try_io(&mut f)) + .await + } + /// Returns the local address that this socket is bound to. /// /// # Examples @@ -1319,3 +1577,10 @@ impl AsRawFd for UnixDatagram { self.io.as_raw_fd() } } + +#[cfg(not(tokio_no_as_fd))] +impl AsFd for UnixDatagram { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } +} diff --git a/vendor/tokio/src/net/unix/listener.rs b/vendor/tokio/src/net/unix/listener.rs index efb9503d4..37a5a4ae4 100644 --- a/vendor/tokio/src/net/unix/listener.rs +++ b/vendor/tokio/src/net/unix/listener.rs @@ -1,9 +1,10 @@ use crate::io::{Interest, PollEvented}; use crate::net::unix::{SocketAddr, UnixStream}; -use std::convert::TryFrom; use std::fmt; use std::io; +#[cfg(not(tokio_no_as_fd))] +use std::os::unix::io::{AsFd, BorrowedFd}; use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd}; use std::os::unix::net; use std::path::Path; @@ -44,6 +45,7 @@ cfg_net_unix! { /// } /// } /// ``` + #[cfg_attr(docsrs, doc(alias = "uds"))] pub struct UnixListener { io: PollEvented<mio::net::UnixListener>, } @@ -54,11 +56,13 @@ impl UnixListener { /// /// # Panics /// - /// This function panics if thread-local runtime is not set. + /// This function panics if it is not called from within a runtime with + /// IO enabled. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. + #[track_caller] pub fn bind<P>(path: P) -> io::Result<UnixListener> where P: AsRef<Path>, @@ -71,40 +75,62 @@ impl UnixListener { /// Creates new `UnixListener` from a `std::os::unix::net::UnixListener `. /// /// This function is intended to be used to wrap a UnixListener from the - /// standard library in the Tokio equivalent. The conversion assumes - /// nothing about the underlying listener; it is left up to the user to set - /// it in non-blocking mode. + /// standard library in the Tokio equivalent. + /// + /// # Notes + /// + /// The caller is responsible for ensuring that the listener is in + /// non-blocking mode. Otherwise all I/O operations on the listener + /// will block the thread, which will cause unexpected behavior. + /// Non-blocking mode can be set using [`set_nonblocking`]. + /// + /// [`set_nonblocking`]: std::os::unix::net::UnixListener::set_nonblocking + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::UnixListener; + /// use std::os::unix::net::UnixListener as StdUnixListener; + /// # use std::error::Error; + /// + /// # async fn dox() -> Result<(), Box<dyn Error>> { + /// let std_listener = StdUnixListener::bind("/path/to/the/socket")?; + /// std_listener.set_nonblocking(true)?; + /// let listener = UnixListener::from_std(std_listener)?; + /// # Ok(()) + /// # } + /// ``` /// /// # Panics /// - /// This function panics if thread-local runtime is not set. + /// This function panics if it is not called from within a runtime with + /// IO enabled. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. + #[track_caller] pub fn from_std(listener: net::UnixListener) -> io::Result<UnixListener> { let listener = mio::net::UnixListener::from_std(listener); let io = PollEvented::new(listener)?; Ok(UnixListener { io }) } - /// Turn a [`tokio::net::UnixListener`] into a [`std::os::unix::net::UnixListener`]. + /// Turns a [`tokio::net::UnixListener`] into a [`std::os::unix::net::UnixListener`]. /// /// The returned [`std::os::unix::net::UnixListener`] will have nonblocking mode - /// set as `true`. Use [`set_nonblocking`] to change the blocking mode if needed. + /// set as `true`. Use [`set_nonblocking`] to change the blocking mode if needed. /// /// # Examples /// /// ```rust,no_run - /// use std::error::Error; - /// - /// #[tokio::main] - /// async fn main() -> Result<(), Box<dyn Error>> { - /// let tokio_listener = tokio::net::UnixListener::bind("127.0.0.1:0")?; - /// let std_listener = tokio_listener.into_std()?; - /// std_listener.set_nonblocking(false)?; - /// Ok(()) - /// } + /// # use std::error::Error; + /// # async fn dox() -> Result<(), Box<dyn Error>> { + /// let tokio_listener = tokio::net::UnixListener::bind("/path/to/the/socket")?; + /// let std_listener = tokio_listener.into_std()?; + /// std_listener.set_nonblocking(false)?; + /// # Ok(()) + /// # } /// ``` /// /// [`tokio::net::UnixListener`]: UnixListener @@ -184,3 +210,10 @@ impl AsRawFd for UnixListener { self.io.as_raw_fd() } } + +#[cfg(not(tokio_no_as_fd))] +impl AsFd for UnixListener { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } +} diff --git a/vendor/tokio/src/net/unix/mod.rs b/vendor/tokio/src/net/unix/mod.rs index c3046f17e..a49b70af3 100644 --- a/vendor/tokio/src/net/unix/mod.rs +++ b/vendor/tokio/src/net/unix/mod.rs @@ -1,5 +1,4 @@ -//! Unix domain socket utility types - +//! Unix specific network types. // This module does not currently provide any public API, but it was // unintentionally defined as a public module. Hide it from the documentation // instead of changing it to a private module to avoid breakage. @@ -22,3 +21,17 @@ pub(crate) use stream::UnixStream; mod ucred; pub use ucred::UCred; + +pub mod pipe; + +/// A type representing process and process group IDs. +#[allow(non_camel_case_types)] +pub type uid_t = u32; + +/// A type representing user ID. +#[allow(non_camel_case_types)] +pub type gid_t = u32; + +/// A type representing group ID. +#[allow(non_camel_case_types)] +pub type pid_t = i32; diff --git a/vendor/tokio/src/net/unix/pipe.rs b/vendor/tokio/src/net/unix/pipe.rs new file mode 100644 index 000000000..188ac68b1 --- /dev/null +++ b/vendor/tokio/src/net/unix/pipe.rs @@ -0,0 +1,1222 @@ +//! Unix pipe types. + +use crate::io::interest::Interest; +use crate::io::{AsyncRead, AsyncWrite, PollEvented, ReadBuf, Ready}; + +use mio::unix::pipe as mio_pipe; +use std::fs::File; +use std::io::{self, Read, Write}; +use std::os::unix::fs::{FileTypeExt, OpenOptionsExt}; +#[cfg(not(tokio_no_as_fd))] +use std::os::unix::io::{AsFd, BorrowedFd}; +use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd}; +use std::path::Path; +use std::pin::Pin; +use std::task::{Context, Poll}; + +cfg_io_util! { + use bytes::BufMut; +} + +/// Options and flags which can be used to configure how a FIFO file is opened. +/// +/// This builder allows configuring how to create a pipe end from a FIFO file. +/// Generally speaking, when using `OpenOptions`, you'll first call [`new`], +/// then chain calls to methods to set each option, then call either +/// [`open_receiver`] or [`open_sender`], passing the path of the FIFO file you +/// are trying to open. This will give you a [`io::Result`][result] with a pipe +/// end inside that you can further operate on. +/// +/// [`new`]: OpenOptions::new +/// [`open_receiver`]: OpenOptions::open_receiver +/// [`open_sender`]: OpenOptions::open_sender +/// [result]: std::io::Result +/// +/// # Examples +/// +/// Opening a pair of pipe ends from a FIFO file: +/// +/// ```no_run +/// use tokio::net::unix::pipe; +/// # use std::error::Error; +/// +/// const FIFO_NAME: &str = "path/to/a/fifo"; +/// +/// # async fn dox() -> Result<(), Box<dyn Error>> { +/// let rx = pipe::OpenOptions::new().open_receiver(FIFO_NAME)?; +/// let tx = pipe::OpenOptions::new().open_sender(FIFO_NAME)?; +/// # Ok(()) +/// # } +/// ``` +/// +/// Opening a [`Sender`] on Linux when you are sure the file is a FIFO: +/// +/// ```ignore +/// use tokio::net::unix::pipe; +/// use nix::{unistd::mkfifo, sys::stat::Mode}; +/// # use std::error::Error; +/// +/// // Our program has exclusive access to this path. +/// const FIFO_NAME: &str = "path/to/a/new/fifo"; +/// +/// # async fn dox() -> Result<(), Box<dyn Error>> { +/// mkfifo(FIFO_NAME, Mode::S_IRWXU)?; +/// let tx = pipe::OpenOptions::new() +/// .read_write(true) +/// .unchecked(true) +/// .open_sender(FIFO_NAME)?; +/// # Ok(()) +/// # } +/// ``` +#[derive(Clone, Debug)] +pub struct OpenOptions { + #[cfg(target_os = "linux")] + read_write: bool, + unchecked: bool, +} + +impl OpenOptions { + /// Creates a blank new set of options ready for configuration. + /// + /// All options are initially set to `false`. + pub fn new() -> OpenOptions { + OpenOptions { + #[cfg(target_os = "linux")] + read_write: false, + unchecked: false, + } + } + + /// Sets the option for read-write access. + /// + /// This option, when true, will indicate that a FIFO file will be opened + /// in read-write access mode. This operation is not defined by the POSIX + /// standard and is only guaranteed to work on Linux. + /// + /// # Examples + /// + /// Opening a [`Sender`] even if there are no open reading ends: + /// + /// ```ignore + /// use tokio::net::unix::pipe; + /// + /// let tx = pipe::OpenOptions::new() + /// .read_write(true) + /// .open_sender("path/to/a/fifo"); + /// ``` + /// + /// Opening a resilient [`Receiver`] i.e. a reading pipe end which will not + /// fail with [`UnexpectedEof`] during reading if all writing ends of the + /// pipe close the FIFO file. + /// + /// [`UnexpectedEof`]: std::io::ErrorKind::UnexpectedEof + /// + /// ```ignore + /// use tokio::net::unix::pipe; + /// + /// let tx = pipe::OpenOptions::new() + /// .read_write(true) + /// .open_receiver("path/to/a/fifo"); + /// ``` + #[cfg(target_os = "linux")] + #[cfg_attr(docsrs, doc(cfg(target_os = "linux")))] + pub fn read_write(&mut self, value: bool) -> &mut Self { + self.read_write = value; + self + } + + /// Sets the option to skip the check for FIFO file type. + /// + /// By default, [`open_receiver`] and [`open_sender`] functions will check + /// if the opened file is a FIFO file. Set this option to `true` if you are + /// sure the file is a FIFO file. + /// + /// [`open_receiver`]: OpenOptions::open_receiver + /// [`open_sender`]: OpenOptions::open_sender + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::unix::pipe; + /// use nix::{unistd::mkfifo, sys::stat::Mode}; + /// # use std::error::Error; + /// + /// // Our program has exclusive access to this path. + /// const FIFO_NAME: &str = "path/to/a/new/fifo"; + /// + /// # async fn dox() -> Result<(), Box<dyn Error>> { + /// mkfifo(FIFO_NAME, Mode::S_IRWXU)?; + /// let rx = pipe::OpenOptions::new() + /// .unchecked(true) + /// .open_receiver(FIFO_NAME)?; + /// # Ok(()) + /// # } + /// ``` + pub fn unchecked(&mut self, value: bool) -> &mut Self { + self.unchecked = value; + self + } + + /// Creates a [`Receiver`] from a FIFO file with the options specified by `self`. + /// + /// This function will open the FIFO file at the specified path, possibly + /// check if it is a pipe, and associate the pipe with the default event + /// loop for reading. + /// + /// # Errors + /// + /// If the file type check fails, this function will fail with `io::ErrorKind::InvalidInput`. + /// This function may also fail with other standard OS errors. + /// + /// # Panics + /// + /// This function panics if it is not called from within a runtime with + /// IO enabled. + /// + /// The runtime is usually set implicitly when this function is called + /// from a future driven by a tokio runtime, otherwise runtime can be set + /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. + pub fn open_receiver<P: AsRef<Path>>(&self, path: P) -> io::Result<Receiver> { + let file = self.open(path.as_ref(), PipeEnd::Receiver)?; + Receiver::from_file_unchecked(file) + } + + /// Creates a [`Sender`] from a FIFO file with the options specified by `self`. + /// + /// This function will open the FIFO file at the specified path, possibly + /// check if it is a pipe, and associate the pipe with the default event + /// loop for writing. + /// + /// # Errors + /// + /// If the file type check fails, this function will fail with `io::ErrorKind::InvalidInput`. + /// If the file is not opened in read-write access mode and the file is not + /// currently open for reading, this function will fail with `ENXIO`. + /// This function may also fail with other standard OS errors. + /// + /// # Panics + /// + /// This function panics if it is not called from within a runtime with + /// IO enabled. + /// + /// The runtime is usually set implicitly when this function is called + /// from a future driven by a tokio runtime, otherwise runtime can be set + /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. + pub fn open_sender<P: AsRef<Path>>(&self, path: P) -> io::Result<Sender> { + let file = self.open(path.as_ref(), PipeEnd::Sender)?; + Sender::from_file_unchecked(file) + } + + fn open(&self, path: &Path, pipe_end: PipeEnd) -> io::Result<File> { + let mut options = std::fs::OpenOptions::new(); + options + .read(pipe_end == PipeEnd::Receiver) + .write(pipe_end == PipeEnd::Sender) + .custom_flags(libc::O_NONBLOCK); + + #[cfg(target_os = "linux")] + if self.read_write { + options.read(true).write(true); + } + + let file = options.open(path)?; + + if !self.unchecked && !is_fifo(&file)? { + return Err(io::Error::new(io::ErrorKind::InvalidInput, "not a pipe")); + } + + Ok(file) + } +} + +impl Default for OpenOptions { + fn default() -> OpenOptions { + OpenOptions::new() + } +} + +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +enum PipeEnd { + Sender, + Receiver, +} + +/// Writing end of a Unix pipe. +/// +/// It can be constructed from a FIFO file with [`OpenOptions::open_sender`]. +/// +/// Opening a named pipe for writing involves a few steps. +/// Call to [`OpenOptions::open_sender`] might fail with an error indicating +/// different things: +/// +/// * [`io::ErrorKind::NotFound`] - There is no file at the specified path. +/// * [`io::ErrorKind::InvalidInput`] - The file exists, but it is not a FIFO. +/// * [`ENXIO`] - The file is a FIFO, but no process has it open for reading. +/// Sleep for a while and try again. +/// * Other OS errors not specific to opening FIFO files. +/// +/// Opening a `Sender` from a FIFO file should look like this: +/// +/// ```no_run +/// use tokio::net::unix::pipe; +/// use tokio::time::{self, Duration}; +/// +/// const FIFO_NAME: &str = "path/to/a/fifo"; +/// +/// # async fn dox() -> Result<(), Box<dyn std::error::Error>> { +/// // Wait for a reader to open the file. +/// let tx = loop { +/// match pipe::OpenOptions::new().open_sender(FIFO_NAME) { +/// Ok(tx) => break tx, +/// Err(e) if e.raw_os_error() == Some(libc::ENXIO) => {}, +/// Err(e) => return Err(e.into()), +/// } +/// +/// time::sleep(Duration::from_millis(50)).await; +/// }; +/// # Ok(()) +/// # } +/// ``` +/// +/// On Linux, it is possible to create a `Sender` without waiting in a sleeping +/// loop. This is done by opening a named pipe in read-write access mode with +/// `OpenOptions::read_write`. This way, a `Sender` can at the same time hold +/// both a writing end and a reading end, and the latter allows to open a FIFO +/// without [`ENXIO`] error since the pipe is open for reading as well. +/// +/// `Sender` cannot be used to read from a pipe, so in practice the read access +/// is only used when a FIFO is opened. However, using a `Sender` in read-write +/// mode **may lead to lost data**, because written data will be dropped by the +/// system as soon as all pipe ends are closed. To avoid lost data you have to +/// make sure that a reading end has been opened before dropping a `Sender`. +/// +/// Note that using read-write access mode with FIFO files is not defined by +/// the POSIX standard and it is only guaranteed to work on Linux. +/// +/// ```ignore +/// use tokio::io::AsyncWriteExt; +/// use tokio::net::unix::pipe; +/// +/// const FIFO_NAME: &str = "path/to/a/fifo"; +/// +/// # async fn dox() -> Result<(), Box<dyn std::error::Error>> { +/// let mut tx = pipe::OpenOptions::new() +/// .read_write(true) +/// .open_sender(FIFO_NAME)?; +/// +/// // Asynchronously write to the pipe before a reader. +/// tx.write_all(b"hello world").await?; +/// # Ok(()) +/// # } +/// ``` +/// +/// [`ENXIO`]: https://docs.rs/libc/latest/libc/constant.ENXIO.html +#[derive(Debug)] +pub struct Sender { + io: PollEvented<mio_pipe::Sender>, +} + +impl Sender { + fn from_mio(mio_tx: mio_pipe::Sender) -> io::Result<Sender> { + let io = PollEvented::new_with_interest(mio_tx, Interest::WRITABLE)?; + Ok(Sender { io }) + } + + /// Creates a new `Sender` from a [`File`]. + /// + /// This function is intended to construct a pipe from a [`File`] representing + /// a special FIFO file. It will check if the file is a pipe and has write access, + /// set it in non-blocking mode and perform the conversion. + /// + /// # Errors + /// + /// Fails with `io::ErrorKind::InvalidInput` if the file is not a pipe or it + /// does not have write access. Also fails with any standard OS error if it occurs. + /// + /// # Panics + /// + /// This function panics if it is not called from within a runtime with + /// IO enabled. + /// + /// The runtime is usually set implicitly when this function is called + /// from a future driven by a tokio runtime, otherwise runtime can be set + /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. + pub fn from_file(mut file: File) -> io::Result<Sender> { + if !is_fifo(&file)? { + return Err(io::Error::new(io::ErrorKind::InvalidInput, "not a pipe")); + } + + let flags = get_file_flags(&file)?; + if has_write_access(flags) { + set_nonblocking(&mut file, flags)?; + Sender::from_file_unchecked(file) + } else { + Err(io::Error::new( + io::ErrorKind::InvalidInput, + "not in O_WRONLY or O_RDWR access mode", + )) + } + } + + /// Creates a new `Sender` from a [`File`] without checking pipe properties. + /// + /// This function is intended to construct a pipe from a File representing + /// a special FIFO file. The conversion assumes nothing about the underlying + /// file; it is left up to the user to make sure it is opened with write access, + /// represents a pipe and is set in non-blocking mode. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::unix::pipe; + /// use std::fs::OpenOptions; + /// use std::os::unix::fs::{FileTypeExt, OpenOptionsExt}; + /// # use std::error::Error; + /// + /// const FIFO_NAME: &str = "path/to/a/fifo"; + /// + /// # async fn dox() -> Result<(), Box<dyn Error>> { + /// let file = OpenOptions::new() + /// .write(true) + /// .custom_flags(libc::O_NONBLOCK) + /// .open(FIFO_NAME)?; + /// if file.metadata()?.file_type().is_fifo() { + /// let tx = pipe::Sender::from_file_unchecked(file)?; + /// /* use the Sender */ + /// } + /// # Ok(()) + /// # } + /// ``` + /// + /// # Panics + /// + /// This function panics if it is not called from within a runtime with + /// IO enabled. + /// + /// The runtime is usually set implicitly when this function is called + /// from a future driven by a tokio runtime, otherwise runtime can be set + /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. + pub fn from_file_unchecked(file: File) -> io::Result<Sender> { + let raw_fd = file.into_raw_fd(); + let mio_tx = unsafe { mio_pipe::Sender::from_raw_fd(raw_fd) }; + Sender::from_mio(mio_tx) + } + + /// Waits for any of the requested ready states. + /// + /// This function can be used instead of [`writable()`] to check the returned + /// ready set for [`Ready::WRITABLE`] and [`Ready::WRITE_CLOSED`] events. + /// + /// The function may complete without the pipe being ready. This is a + /// false-positive and attempting an operation will return with + /// `io::ErrorKind::WouldBlock`. The function can also return with an empty + /// [`Ready`] set, so you should always check the returned value and possibly + /// wait again if the requested states are not set. + /// + /// [`writable()`]: Self::writable + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to write that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn ready(&self, interest: Interest) -> io::Result<Ready> { + let event = self.io.registration().readiness(interest).await?; + Ok(event.ready) + } + + /// Waits for the pipe to become writable. + /// + /// This function is equivalent to `ready(Interest::WRITABLE)` and is usually + /// paired with [`try_write()`]. + /// + /// [`try_write()`]: Self::try_write + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::unix::pipe; + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// // Open a writing end of a fifo + /// let tx = pipe::OpenOptions::new().open_sender("path/to/a/fifo")?; + /// + /// loop { + /// // Wait for the pipe to be writable + /// tx.writable().await?; + /// + /// // Try to write data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match tx.try_write(b"hello world") { + /// Ok(n) => { + /// break; + /// } + /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// + /// Ok(()) + /// } + /// ``` + pub async fn writable(&self) -> io::Result<()> { + self.ready(Interest::WRITABLE).await?; + Ok(()) + } + + /// Polls for write readiness. + /// + /// If the pipe is not currently ready for writing, this method will + /// store a clone of the `Waker` from the provided `Context`. When the pipe + /// becomes ready for writing, `Waker::wake` will be called on the waker. + /// + /// Note that on multiple calls to `poll_write_ready` or `poll_write`, only + /// the `Waker` from the `Context` passed to the most recent call is + /// scheduled to receive a wakeup. + /// + /// This function is intended for cases where creating and pinning a future + /// via [`writable`] is not feasible. Where possible, using [`writable`] is + /// preferred, as this supports polling from multiple tasks at once. + /// + /// [`writable`]: Self::writable + /// + /// # Return value + /// + /// The function returns: + /// + /// * `Poll::Pending` if the pipe is not ready for writing. + /// * `Poll::Ready(Ok(()))` if the pipe is ready for writing. + /// * `Poll::Ready(Err(e))` if an error is encountered. + /// + /// # Errors + /// + /// This function may encounter any standard I/O error except `WouldBlock`. + pub fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + self.io.registration().poll_write_ready(cx).map_ok(|_| ()) + } + + /// Tries to write a buffer to the pipe, returning how many bytes were + /// written. + /// + /// The function will attempt to write the entire contents of `buf`, but + /// only part of the buffer may be written. If the length of `buf` is not + /// greater than `PIPE_BUF` (an OS constant, 4096 under Linux), then the + /// write is guaranteed to be atomic, i.e. either the entire content of + /// `buf` will be written or this method will fail with `WouldBlock`. There + /// is no such guarantee if `buf` is larger than `PIPE_BUF`. + /// + /// This function is usually paired with [`writable`]. + /// + /// [`writable`]: Self::writable + /// + /// # Return + /// + /// If data is successfully written, `Ok(n)` is returned, where `n` is the + /// number of bytes written. If the pipe is not ready to write data, + /// `Err(io::ErrorKind::WouldBlock)` is returned. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::unix::pipe; + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// // Open a writing end of a fifo + /// let tx = pipe::OpenOptions::new().open_sender("path/to/a/fifo")?; + /// + /// loop { + /// // Wait for the pipe to be writable + /// tx.writable().await?; + /// + /// // Try to write data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match tx.try_write(b"hello world") { + /// Ok(n) => { + /// break; + /// } + /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// + /// Ok(()) + /// } + /// ``` + pub fn try_write(&self, buf: &[u8]) -> io::Result<usize> { + self.io + .registration() + .try_io(Interest::WRITABLE, || (&*self.io).write(buf)) + } + + /// Tries to write several buffers to the pipe, returning how many bytes + /// were written. + /// + /// Data is written from each buffer in order, with the final buffer read + /// from possible being only partially consumed. This method behaves + /// equivalently to a single call to [`try_write()`] with concatenated + /// buffers. + /// + /// If the total length of buffers is not greater than `PIPE_BUF` (an OS + /// constant, 4096 under Linux), then the write is guaranteed to be atomic, + /// i.e. either the entire contents of buffers will be written or this + /// method will fail with `WouldBlock`. There is no such guarantee if the + /// total length of buffers is greater than `PIPE_BUF`. + /// + /// This function is usually paired with [`writable`]. + /// + /// [`try_write()`]: Self::try_write() + /// [`writable`]: Self::writable + /// + /// # Return + /// + /// If data is successfully written, `Ok(n)` is returned, where `n` is the + /// number of bytes written. If the pipe is not ready to write data, + /// `Err(io::ErrorKind::WouldBlock)` is returned. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::unix::pipe; + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// // Open a writing end of a fifo + /// let tx = pipe::OpenOptions::new().open_sender("path/to/a/fifo")?; + /// + /// let bufs = [io::IoSlice::new(b"hello "), io::IoSlice::new(b"world")]; + /// + /// loop { + /// // Wait for the pipe to be writable + /// tx.writable().await?; + /// + /// // Try to write data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match tx.try_write_vectored(&bufs) { + /// Ok(n) => { + /// break; + /// } + /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// + /// Ok(()) + /// } + /// ``` + pub fn try_write_vectored(&self, buf: &[io::IoSlice<'_>]) -> io::Result<usize> { + self.io + .registration() + .try_io(Interest::WRITABLE, || (&*self.io).write_vectored(buf)) + } +} + +impl AsyncWrite for Sender { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + self.io.poll_write(cx, buf) + } + + fn poll_write_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[io::IoSlice<'_>], + ) -> Poll<io::Result<usize>> { + self.io.poll_write_vectored(cx, bufs) + } + + fn is_write_vectored(&self) -> bool { + true + } + + fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> { + Poll::Ready(Ok(())) + } +} + +impl AsRawFd for Sender { + fn as_raw_fd(&self) -> RawFd { + self.io.as_raw_fd() + } +} + +#[cfg(not(tokio_no_as_fd))] +impl AsFd for Sender { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } +} + +/// Reading end of a Unix pipe. +/// +/// It can be constructed from a FIFO file with [`OpenOptions::open_receiver`]. +/// +/// # Examples +/// +/// Receiving messages from a named pipe in a loop: +/// +/// ```no_run +/// use tokio::net::unix::pipe; +/// use tokio::io::{self, AsyncReadExt}; +/// +/// const FIFO_NAME: &str = "path/to/a/fifo"; +/// +/// # async fn dox() -> Result<(), Box<dyn std::error::Error>> { +/// let mut rx = pipe::OpenOptions::new().open_receiver(FIFO_NAME)?; +/// loop { +/// let mut msg = vec![0; 256]; +/// match rx.read_exact(&mut msg).await { +/// Ok(_) => { +/// /* handle the message */ +/// } +/// Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => { +/// // Writing end has been closed, we should reopen the pipe. +/// rx = pipe::OpenOptions::new().open_receiver(FIFO_NAME)?; +/// } +/// Err(e) => return Err(e.into()), +/// } +/// } +/// # } +/// ``` +/// +/// On Linux, you can use a `Receiver` in read-write access mode to implement +/// resilient reading from a named pipe. Unlike `Receiver` opened in read-only +/// mode, read from a pipe in read-write mode will not fail with `UnexpectedEof` +/// when the writing end is closed. This way, a `Receiver` can asynchronously +/// wait for the next writer to open the pipe. +/// +/// You should not use functions waiting for EOF such as [`read_to_end`] with +/// a `Receiver` in read-write access mode, since it **may wait forever**. +/// `Receiver` in this mode also holds an open writing end, which prevents +/// receiving EOF. +/// +/// To set the read-write access mode you can use `OpenOptions::read_write`. +/// Note that using read-write access mode with FIFO files is not defined by +/// the POSIX standard and it is only guaranteed to work on Linux. +/// +/// ```ignore +/// use tokio::net::unix::pipe; +/// use tokio::io::AsyncReadExt; +/// # use std::error::Error; +/// +/// const FIFO_NAME: &str = "path/to/a/fifo"; +/// +/// # async fn dox() -> Result<(), Box<dyn Error>> { +/// let mut rx = pipe::OpenOptions::new() +/// .read_write(true) +/// .open_receiver(FIFO_NAME)?; +/// loop { +/// let mut msg = vec![0; 256]; +/// rx.read_exact(&mut msg).await?; +/// /* handle the message */ +/// } +/// # } +/// ``` +/// +/// [`read_to_end`]: crate::io::AsyncReadExt::read_to_end +#[derive(Debug)] +pub struct Receiver { + io: PollEvented<mio_pipe::Receiver>, +} + +impl Receiver { + fn from_mio(mio_rx: mio_pipe::Receiver) -> io::Result<Receiver> { + let io = PollEvented::new_with_interest(mio_rx, Interest::READABLE)?; + Ok(Receiver { io }) + } + + /// Creates a new `Receiver` from a [`File`]. + /// + /// This function is intended to construct a pipe from a [`File`] representing + /// a special FIFO file. It will check if the file is a pipe and has read access, + /// set it in non-blocking mode and perform the conversion. + /// + /// # Errors + /// + /// Fails with `io::ErrorKind::InvalidInput` if the file is not a pipe or it + /// does not have read access. Also fails with any standard OS error if it occurs. + /// + /// # Panics + /// + /// This function panics if it is not called from within a runtime with + /// IO enabled. + /// + /// The runtime is usually set implicitly when this function is called + /// from a future driven by a tokio runtime, otherwise runtime can be set + /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. + pub fn from_file(mut file: File) -> io::Result<Receiver> { + if !is_fifo(&file)? { + return Err(io::Error::new(io::ErrorKind::InvalidInput, "not a pipe")); + } + + let flags = get_file_flags(&file)?; + if has_read_access(flags) { + set_nonblocking(&mut file, flags)?; + Receiver::from_file_unchecked(file) + } else { + Err(io::Error::new( + io::ErrorKind::InvalidInput, + "not in O_RDONLY or O_RDWR access mode", + )) + } + } + + /// Creates a new `Receiver` from a [`File`] without checking pipe properties. + /// + /// This function is intended to construct a pipe from a File representing + /// a special FIFO file. The conversion assumes nothing about the underlying + /// file; it is left up to the user to make sure it is opened with read access, + /// represents a pipe and is set in non-blocking mode. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::unix::pipe; + /// use std::fs::OpenOptions; + /// use std::os::unix::fs::{FileTypeExt, OpenOptionsExt}; + /// # use std::error::Error; + /// + /// const FIFO_NAME: &str = "path/to/a/fifo"; + /// + /// # async fn dox() -> Result<(), Box<dyn Error>> { + /// let file = OpenOptions::new() + /// .read(true) + /// .custom_flags(libc::O_NONBLOCK) + /// .open(FIFO_NAME)?; + /// if file.metadata()?.file_type().is_fifo() { + /// let rx = pipe::Receiver::from_file_unchecked(file)?; + /// /* use the Receiver */ + /// } + /// # Ok(()) + /// # } + /// ``` + /// + /// # Panics + /// + /// This function panics if it is not called from within a runtime with + /// IO enabled. + /// + /// The runtime is usually set implicitly when this function is called + /// from a future driven by a tokio runtime, otherwise runtime can be set + /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. + pub fn from_file_unchecked(file: File) -> io::Result<Receiver> { + let raw_fd = file.into_raw_fd(); + let mio_rx = unsafe { mio_pipe::Receiver::from_raw_fd(raw_fd) }; + Receiver::from_mio(mio_rx) + } + + /// Waits for any of the requested ready states. + /// + /// This function can be used instead of [`readable()`] to check the returned + /// ready set for [`Ready::READABLE`] and [`Ready::READ_CLOSED`] events. + /// + /// The function may complete without the pipe being ready. This is a + /// false-positive and attempting an operation will return with + /// `io::ErrorKind::WouldBlock`. The function can also return with an empty + /// [`Ready`] set, so you should always check the returned value and possibly + /// wait again if the requested states are not set. + /// + /// [`readable()`]: Self::readable + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to read that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn ready(&self, interest: Interest) -> io::Result<Ready> { + let event = self.io.registration().readiness(interest).await?; + Ok(event.ready) + } + + /// Waits for the pipe to become readable. + /// + /// This function is equivalent to `ready(Interest::READABLE)` and is usually + /// paired with [`try_read()`]. + /// + /// [`try_read()`]: Self::try_read() + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::unix::pipe; + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// // Open a reading end of a fifo + /// let rx = pipe::OpenOptions::new().open_receiver("path/to/a/fifo")?; + /// + /// let mut msg = vec![0; 1024]; + /// + /// loop { + /// // Wait for the pipe to be readable + /// rx.readable().await?; + /// + /// // Try to read data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match rx.try_read(&mut msg) { + /// Ok(n) => { + /// msg.truncate(n); + /// break; + /// } + /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// + /// println!("GOT = {:?}", msg); + /// Ok(()) + /// } + /// ``` + pub async fn readable(&self) -> io::Result<()> { + self.ready(Interest::READABLE).await?; + Ok(()) + } + + /// Polls for read readiness. + /// + /// If the pipe is not currently ready for reading, this method will + /// store a clone of the `Waker` from the provided `Context`. When the pipe + /// becomes ready for reading, `Waker::wake` will be called on the waker. + /// + /// Note that on multiple calls to `poll_read_ready` or `poll_read`, only + /// the `Waker` from the `Context` passed to the most recent call is + /// scheduled to receive a wakeup. + /// + /// This function is intended for cases where creating and pinning a future + /// via [`readable`] is not feasible. Where possible, using [`readable`] is + /// preferred, as this supports polling from multiple tasks at once. + /// + /// [`readable`]: Self::readable + /// + /// # Return value + /// + /// The function returns: + /// + /// * `Poll::Pending` if the pipe is not ready for reading. + /// * `Poll::Ready(Ok(()))` if the pipe is ready for reading. + /// * `Poll::Ready(Err(e))` if an error is encountered. + /// + /// # Errors + /// + /// This function may encounter any standard I/O error except `WouldBlock`. + pub fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + self.io.registration().poll_read_ready(cx).map_ok(|_| ()) + } + + /// Tries to read data from the pipe into the provided buffer, returning how + /// many bytes were read. + /// + /// Reads any pending data from the pipe but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read()` is non-blocking, the buffer does not have to be stored by + /// the async task and can exist entirely on the stack. + /// + /// Usually [`readable()`] is used with this function. + /// + /// [`readable()`]: Self::readable() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. If `n` is `0`, then it can indicate one of two scenarios: + /// + /// 1. The pipe's writing end is closed and will no longer write data. + /// 2. The specified buffer was 0 bytes in length. + /// + /// If the pipe is not ready to read data, + /// `Err(io::ErrorKind::WouldBlock)` is returned. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::unix::pipe; + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// // Open a reading end of a fifo + /// let rx = pipe::OpenOptions::new().open_receiver("path/to/a/fifo")?; + /// + /// let mut msg = vec![0; 1024]; + /// + /// loop { + /// // Wait for the pipe to be readable + /// rx.readable().await?; + /// + /// // Try to read data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match rx.try_read(&mut msg) { + /// Ok(n) => { + /// msg.truncate(n); + /// break; + /// } + /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// + /// println!("GOT = {:?}", msg); + /// Ok(()) + /// } + /// ``` + pub fn try_read(&self, buf: &mut [u8]) -> io::Result<usize> { + self.io + .registration() + .try_io(Interest::READABLE, || (&*self.io).read(buf)) + } + + /// Tries to read data from the pipe into the provided buffers, returning + /// how many bytes were read. + /// + /// Data is copied to fill each buffer in order, with the final buffer + /// written to possibly being only partially filled. This method behaves + /// equivalently to a single call to [`try_read()`] with concatenated + /// buffers. + /// + /// Reads any pending data from the pipe but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read_vectored()` is non-blocking, the buffer does not have to be + /// stored by the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] is used with this function. + /// + /// [`try_read()`]: Self::try_read() + /// [`readable()`]: Self::readable() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. `Ok(0)` indicates the pipe's writing end is + /// closed and will no longer write data. If the pipe is not ready to read + /// data `Err(io::ErrorKind::WouldBlock)` is returned. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::unix::pipe; + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// // Open a reading end of a fifo + /// let rx = pipe::OpenOptions::new().open_receiver("path/to/a/fifo")?; + /// + /// loop { + /// // Wait for the pipe to be readable + /// rx.readable().await?; + /// + /// // Creating the buffer **after** the `await` prevents it from + /// // being stored in the async task. + /// let mut buf_a = [0; 512]; + /// let mut buf_b = [0; 1024]; + /// let mut bufs = [ + /// io::IoSliceMut::new(&mut buf_a), + /// io::IoSliceMut::new(&mut buf_b), + /// ]; + /// + /// // Try to read data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match rx.try_read_vectored(&mut bufs) { + /// Ok(0) => break, + /// Ok(n) => { + /// println!("read {} bytes", n); + /// } + /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// + /// Ok(()) + /// } + /// ``` + pub fn try_read_vectored(&self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result<usize> { + self.io + .registration() + .try_io(Interest::READABLE, || (&*self.io).read_vectored(bufs)) + } + + cfg_io_util! { + /// Tries to read data from the pipe into the provided buffer, advancing the + /// buffer's internal cursor, returning how many bytes were read. + /// + /// Reads any pending data from the pipe but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read_buf()` is non-blocking, the buffer does not have to be stored by + /// the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: Self::readable + /// [`ready()`]: Self::ready + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. `Ok(0)` indicates the pipe's writing end is + /// closed and will no longer write data. If the pipe is not ready to read + /// data `Err(io::ErrorKind::WouldBlock)` is returned. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::unix::pipe; + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// // Open a reading end of a fifo + /// let rx = pipe::OpenOptions::new().open_receiver("path/to/a/fifo")?; + /// + /// loop { + /// // Wait for the pipe to be readable + /// rx.readable().await?; + /// + /// let mut buf = Vec::with_capacity(4096); + /// + /// // Try to read data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match rx.try_read_buf(&mut buf) { + /// Ok(0) => break, + /// Ok(n) => { + /// println!("read {} bytes", n); + /// } + /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// + /// Ok(()) + /// } + /// ``` + pub fn try_read_buf<B: BufMut>(&self, buf: &mut B) -> io::Result<usize> { + self.io.registration().try_io(Interest::READABLE, || { + use std::io::Read; + + let dst = buf.chunk_mut(); + let dst = + unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]) }; + + // Safety: `mio_pipe::Receiver` uses a `std::fs::File` underneath, + // which correctly handles reads into uninitialized memory. + let n = (&*self.io).read(dst)?; + + unsafe { + buf.advance_mut(n); + } + + Ok(n) + }) + } + } +} + +impl AsyncRead for Receiver { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<()>> { + // Safety: `mio_pipe::Receiver` uses a `std::fs::File` underneath, + // which correctly handles reads into uninitialized memory. + unsafe { self.io.poll_read(cx, buf) } + } +} + +impl AsRawFd for Receiver { + fn as_raw_fd(&self) -> RawFd { + self.io.as_raw_fd() + } +} + +#[cfg(not(tokio_no_as_fd))] +impl AsFd for Receiver { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } +} + +/// Checks if file is a FIFO +fn is_fifo(file: &File) -> io::Result<bool> { + Ok(file.metadata()?.file_type().is_fifo()) +} + +/// Gets file descriptor's flags by fcntl. +fn get_file_flags(file: &File) -> io::Result<libc::c_int> { + let fd = file.as_raw_fd(); + let flags = unsafe { libc::fcntl(fd, libc::F_GETFL) }; + if flags < 0 { + Err(io::Error::last_os_error()) + } else { + Ok(flags) + } +} + +/// Checks for O_RDONLY or O_RDWR access mode. +fn has_read_access(flags: libc::c_int) -> bool { + let mode = flags & libc::O_ACCMODE; + mode == libc::O_RDONLY || mode == libc::O_RDWR +} + +/// Checks for O_WRONLY or O_RDWR access mode. +fn has_write_access(flags: libc::c_int) -> bool { + let mode = flags & libc::O_ACCMODE; + mode == libc::O_WRONLY || mode == libc::O_RDWR +} + +/// Sets file's flags with O_NONBLOCK by fcntl. +fn set_nonblocking(file: &mut File, current_flags: libc::c_int) -> io::Result<()> { + let fd = file.as_raw_fd(); + + let flags = current_flags | libc::O_NONBLOCK; + + if flags != current_flags { + let ret = unsafe { libc::fcntl(fd, libc::F_SETFL, flags) }; + if ret < 0 { + return Err(io::Error::last_os_error()); + } + } + + Ok(()) +} diff --git a/vendor/tokio/src/net/unix/split.rs b/vendor/tokio/src/net/unix/split.rs index 97214f7a7..6fc7067a7 100644 --- a/vendor/tokio/src/net/unix/split.rs +++ b/vendor/tokio/src/net/unix/split.rs @@ -8,14 +8,19 @@ //! split has no associated overhead and enforces all invariants at the type //! level. -use crate::io::{AsyncRead, AsyncWrite, ReadBuf}; +use crate::io::{AsyncRead, AsyncWrite, Interest, ReadBuf, Ready}; use crate::net::UnixStream; +use crate::net::unix::SocketAddr; use std::io; use std::net::Shutdown; use std::pin::Pin; use std::task::{Context, Poll}; +cfg_io_util! { + use bytes::BufMut; +} + /// Borrowed read half of a [`UnixStream`], created by [`split`]. /// /// Reading from a `ReadHalf` is usually done using the convenience methods found on the @@ -47,6 +52,232 @@ pub(crate) fn split(stream: &mut UnixStream) -> (ReadHalf<'_>, WriteHalf<'_>) { (ReadHalf(stream), WriteHalf(stream)) } +impl ReadHalf<'_> { + /// Wait for any of the requested ready states. + /// + /// This function is usually paired with [`try_read()`]. It can be used instead + /// of [`readable()`] to check the returned ready set for [`Ready::READABLE`] + /// and [`Ready::READ_CLOSED`] events. + /// + /// The function may complete without the socket being ready. This is a + /// false-positive and attempting an operation will return with + /// `io::ErrorKind::WouldBlock`. The function can also return with an empty + /// [`Ready`] set, so you should always check the returned value and possibly + /// wait again if the requested states are not set. + /// + /// This function is equivalent to [`UnixStream::ready`]. + /// + /// [`try_read()`]: Self::try_read + /// [`readable()`]: Self::readable + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to read or write that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn ready(&self, interest: Interest) -> io::Result<Ready> { + self.0.ready(interest).await + } + + /// Waits for the socket to become readable. + /// + /// This function is equivalent to `ready(Interest::READABLE)` and is usually + /// paired with `try_read()`. + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to read that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn readable(&self) -> io::Result<()> { + self.0.readable().await + } + + /// Tries to read data from the stream into the provided buffer, returning how + /// many bytes were read. + /// + /// Receives any pending data from the socket but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read()` is non-blocking, the buffer does not have to be stored by + /// the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: Self::readable() + /// [`ready()`]: Self::ready() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. If `n` is `0`, then it can indicate one of two scenarios: + /// + /// 1. The stream's read half is closed and will no longer yield data. + /// 2. The specified buffer was 0 bytes in length. + /// + /// If the stream is not ready to read data, + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_read(&self, buf: &mut [u8]) -> io::Result<usize> { + self.0.try_read(buf) + } + + cfg_io_util! { + /// Tries to read data from the stream into the provided buffer, advancing the + /// buffer's internal cursor, returning how many bytes were read. + /// + /// Receives any pending data from the socket but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read_buf()` is non-blocking, the buffer does not have to be stored by + /// the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: Self::readable() + /// [`ready()`]: Self::ready() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. `Ok(0)` indicates the stream's read half is closed + /// and will no longer yield data. If the stream is not ready to read data + pub fn try_read_buf<B: BufMut>(&self, buf: &mut B) -> io::Result<usize> { + self.0.try_read_buf(buf) + } + } + + /// Tries to read data from the stream into the provided buffers, returning + /// how many bytes were read. + /// + /// Data is copied to fill each buffer in order, with the final buffer + /// written to possibly being only partially filled. This method behaves + /// equivalently to a single call to [`try_read()`] with concatenated + /// buffers. + /// + /// Receives any pending data from the socket but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read_vectored()` is non-blocking, the buffer does not have to be + /// stored by the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`try_read()`]: Self::try_read() + /// [`readable()`]: Self::readable() + /// [`ready()`]: Self::ready() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. `Ok(0)` indicates the stream's read half is closed + /// and will no longer yield data. If the stream is not ready to read data + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_read_vectored(&self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result<usize> { + self.0.try_read_vectored(bufs) + } + + /// Returns the socket address of the remote half of this connection. + pub fn peer_addr(&self) -> io::Result<SocketAddr> { + self.0.peer_addr() + } + + /// Returns the socket address of the local half of this connection. + pub fn local_addr(&self) -> io::Result<SocketAddr> { + self.0.local_addr() + } +} + +impl WriteHalf<'_> { + /// Waits for any of the requested ready states. + /// + /// This function is usually paired with [`try_write()`]. It can be used instead + /// of [`writable()`] to check the returned ready set for [`Ready::WRITABLE`] + /// and [`Ready::WRITE_CLOSED`] events. + /// + /// The function may complete without the socket being ready. This is a + /// false-positive and attempting an operation will return with + /// `io::ErrorKind::WouldBlock`. The function can also return with an empty + /// [`Ready`] set, so you should always check the returned value and possibly + /// wait again if the requested states are not set. + /// + /// This function is equivalent to [`UnixStream::ready`]. + /// + /// [`try_write()`]: Self::try_write + /// [`writable()`]: Self::writable + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to read or write that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn ready(&self, interest: Interest) -> io::Result<Ready> { + self.0.ready(interest).await + } + + /// Waits for the socket to become writable. + /// + /// This function is equivalent to `ready(Interest::WRITABLE)` and is usually + /// paired with `try_write()`. + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to write that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn writable(&self) -> io::Result<()> { + self.0.writable().await + } + + /// Tries to write a buffer to the stream, returning how many bytes were + /// written. + /// + /// The function will attempt to write the entire contents of `buf`, but + /// only part of the buffer may be written. + /// + /// This function is usually paired with `writable()`. + /// + /// # Return + /// + /// If data is successfully written, `Ok(n)` is returned, where `n` is the + /// number of bytes written. If the stream is not ready to write data, + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_write(&self, buf: &[u8]) -> io::Result<usize> { + self.0.try_write(buf) + } + + /// Tries to write several buffers to the stream, returning how many bytes + /// were written. + /// + /// Data is written from each buffer in order, with the final buffer read + /// from possible being only partially consumed. This method behaves + /// equivalently to a single call to [`try_write()`] with concatenated + /// buffers. + /// + /// This function is usually paired with `writable()`. + /// + /// [`try_write()`]: Self::try_write() + /// + /// # Return + /// + /// If data is successfully written, `Ok(n)` is returned, where `n` is the + /// number of bytes written. If the stream is not ready to write data, + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_write_vectored(&self, buf: &[io::IoSlice<'_>]) -> io::Result<usize> { + self.0.try_write_vectored(buf) + } + + /// Returns the socket address of the remote half of this connection. + pub fn peer_addr(&self) -> io::Result<SocketAddr> { + self.0.peer_addr() + } + + /// Returns the socket address of the local half of this connection. + pub fn local_addr(&self) -> io::Result<SocketAddr> { + self.0.local_addr() + } +} + impl AsyncRead for ReadHalf<'_> { fn poll_read( self: Pin<&mut Self>, diff --git a/vendor/tokio/src/net/unix/split_owned.rs b/vendor/tokio/src/net/unix/split_owned.rs index 3d6ac6a7e..d1c6f93de 100644 --- a/vendor/tokio/src/net/unix/split_owned.rs +++ b/vendor/tokio/src/net/unix/split_owned.rs @@ -8,9 +8,10 @@ //! split has no associated overhead and enforces all invariants at the type //! level. -use crate::io::{AsyncRead, AsyncWrite, ReadBuf}; +use crate::io::{AsyncRead, AsyncWrite, Interest, ReadBuf, Ready}; use crate::net::UnixStream; +use crate::net::unix::SocketAddr; use std::error::Error; use std::net::Shutdown; use std::pin::Pin; @@ -18,6 +19,10 @@ use std::sync::Arc; use std::task::{Context, Poll}; use std::{fmt, io}; +cfg_io_util! { + use bytes::BufMut; +} + /// Owned read half of a [`UnixStream`], created by [`into_split`]. /// /// Reading from an `OwnedReadHalf` is usually done using the convenience methods found @@ -102,6 +107,139 @@ impl OwnedReadHalf { pub fn reunite(self, other: OwnedWriteHalf) -> Result<UnixStream, ReuniteError> { reunite(self, other) } + + /// Waits for any of the requested ready states. + /// + /// This function is usually paired with [`try_read()`]. It can be used instead + /// of [`readable()`] to check the returned ready set for [`Ready::READABLE`] + /// and [`Ready::READ_CLOSED`] events. + /// + /// The function may complete without the socket being ready. This is a + /// false-positive and attempting an operation will return with + /// `io::ErrorKind::WouldBlock`. The function can also return with an empty + /// [`Ready`] set, so you should always check the returned value and possibly + /// wait again if the requested states are not set. + /// + /// This function is equivalent to [`UnixStream::ready`]. + /// + /// [`try_read()`]: Self::try_read + /// [`readable()`]: Self::readable + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to read or write that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn ready(&self, interest: Interest) -> io::Result<Ready> { + self.inner.ready(interest).await + } + + /// Waits for the socket to become readable. + /// + /// This function is equivalent to `ready(Interest::READABLE)` and is usually + /// paired with `try_read()`. + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to read that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn readable(&self) -> io::Result<()> { + self.inner.readable().await + } + + /// Tries to read data from the stream into the provided buffer, returning how + /// many bytes were read. + /// + /// Receives any pending data from the socket but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read()` is non-blocking, the buffer does not have to be stored by + /// the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: Self::readable() + /// [`ready()`]: Self::ready() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. If `n` is `0`, then it can indicate one of two scenarios: + /// + /// 1. The stream's read half is closed and will no longer yield data. + /// 2. The specified buffer was 0 bytes in length. + /// + /// If the stream is not ready to read data, + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_read(&self, buf: &mut [u8]) -> io::Result<usize> { + self.inner.try_read(buf) + } + + cfg_io_util! { + /// Tries to read data from the stream into the provided buffer, advancing the + /// buffer's internal cursor, returning how many bytes were read. + /// + /// Receives any pending data from the socket but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read_buf()` is non-blocking, the buffer does not have to be stored by + /// the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: Self::readable() + /// [`ready()`]: Self::ready() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. `Ok(0)` indicates the stream's read half is closed + /// and will no longer yield data. If the stream is not ready to read data + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_read_buf<B: BufMut>(&self, buf: &mut B) -> io::Result<usize> { + self.inner.try_read_buf(buf) + } + } + + /// Tries to read data from the stream into the provided buffers, returning + /// how many bytes were read. + /// + /// Data is copied to fill each buffer in order, with the final buffer + /// written to possibly being only partially filled. This method behaves + /// equivalently to a single call to [`try_read()`] with concatenated + /// buffers. + /// + /// Receives any pending data from the socket but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read_vectored()` is non-blocking, the buffer does not have to be + /// stored by the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`try_read()`]: Self::try_read() + /// [`readable()`]: Self::readable() + /// [`ready()`]: Self::ready() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. `Ok(0)` indicates the stream's read half is closed + /// and will no longer yield data. If the stream is not ready to read data + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_read_vectored(&self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result<usize> { + self.inner.try_read_vectored(bufs) + } + + /// Returns the socket address of the remote half of this connection. + pub fn peer_addr(&self) -> io::Result<SocketAddr> { + self.inner.peer_addr() + } + + /// Returns the socket address of the local half of this connection. + pub fn local_addr(&self) -> io::Result<SocketAddr> { + self.inner.local_addr() + } } impl AsyncRead for OwnedReadHalf { @@ -124,13 +262,103 @@ impl OwnedWriteHalf { reunite(other, self) } - /// Destroy the write half, but don't close the write half of the stream + /// Destroys the write half, but don't close the write half of the stream /// until the read half is dropped. If the read half has already been /// dropped, this closes the stream. pub fn forget(mut self) { self.shutdown_on_drop = false; drop(self); } + + /// Waits for any of the requested ready states. + /// + /// This function is usually paired with [`try_write()`]. It can be used instead + /// of [`writable()`] to check the returned ready set for [`Ready::WRITABLE`] + /// and [`Ready::WRITE_CLOSED`] events. + /// + /// The function may complete without the socket being ready. This is a + /// false-positive and attempting an operation will return with + /// `io::ErrorKind::WouldBlock`. The function can also return with an empty + /// [`Ready`] set, so you should always check the returned value and possibly + /// wait again if the requested states are not set. + /// + /// This function is equivalent to [`UnixStream::ready`]. + /// + /// [`try_write()`]: Self::try_write + /// [`writable()`]: Self::writable + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to read or write that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn ready(&self, interest: Interest) -> io::Result<Ready> { + self.inner.ready(interest).await + } + + /// Waits for the socket to become writable. + /// + /// This function is equivalent to `ready(Interest::WRITABLE)` and is usually + /// paired with `try_write()`. + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to write that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn writable(&self) -> io::Result<()> { + self.inner.writable().await + } + + /// Tries to write a buffer to the stream, returning how many bytes were + /// written. + /// + /// The function will attempt to write the entire contents of `buf`, but + /// only part of the buffer may be written. + /// + /// This function is usually paired with `writable()`. + /// + /// # Return + /// + /// If data is successfully written, `Ok(n)` is returned, where `n` is the + /// number of bytes written. If the stream is not ready to write data, + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_write(&self, buf: &[u8]) -> io::Result<usize> { + self.inner.try_write(buf) + } + + /// Tries to write several buffers to the stream, returning how many bytes + /// were written. + /// + /// Data is written from each buffer in order, with the final buffer read + /// from possible being only partially consumed. This method behaves + /// equivalently to a single call to [`try_write()`] with concatenated + /// buffers. + /// + /// This function is usually paired with `writable()`. + /// + /// [`try_write()`]: Self::try_write() + /// + /// # Return + /// + /// If data is successfully written, `Ok(n)` is returned, where `n` is the + /// number of bytes written. If the stream is not ready to write data, + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_write_vectored(&self, buf: &[io::IoSlice<'_>]) -> io::Result<usize> { + self.inner.try_write_vectored(buf) + } + + /// Returns the socket address of the remote half of this connection. + pub fn peer_addr(&self) -> io::Result<SocketAddr> { + self.inner.peer_addr() + } + + /// Returns the socket address of the local half of this connection. + pub fn local_addr(&self) -> io::Result<SocketAddr> { + self.inner.local_addr() + } } impl Drop for OwnedWriteHalf { @@ -180,12 +408,12 @@ impl AsyncWrite for OwnedWriteHalf { impl AsRef<UnixStream> for OwnedReadHalf { fn as_ref(&self) -> &UnixStream { - &*self.inner + &self.inner } } impl AsRef<UnixStream> for OwnedWriteHalf { fn as_ref(&self) -> &UnixStream { - &*self.inner + &self.inner } } diff --git a/vendor/tokio/src/net/unix/stream.rs b/vendor/tokio/src/net/unix/stream.rs index 4baac6062..201645ba1 100644 --- a/vendor/tokio/src/net/unix/stream.rs +++ b/vendor/tokio/src/net/unix/stream.rs @@ -5,10 +5,11 @@ use crate::net::unix::split_owned::{split_owned, OwnedReadHalf, OwnedWriteHalf}; use crate::net::unix::ucred::{self, UCred}; use crate::net::unix::SocketAddr; -use std::convert::TryFrom; use std::fmt; use std::io::{self, Read, Write}; use std::net::Shutdown; +#[cfg(not(tokio_no_as_fd))] +use std::os::unix::io::{AsFd, BorrowedFd}; use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd}; use std::os::unix::net; use std::path::Path; @@ -22,8 +23,8 @@ cfg_io_util! { cfg_net_unix! { /// A structure representing a connected Unix socket. /// - /// This socket can be connected directly with `UnixStream::connect` or accepted - /// from a listener with `UnixListener::incoming`. Additionally, a pair of + /// This socket can be connected directly with [`UnixStream::connect`] or accepted + /// from a listener with [`UnixListener::accept`]. Additionally, a pair of /// anonymous Unix sockets can be created with `UnixStream::pair`. /// /// To shut down the stream in the write direction, you can call the @@ -32,6 +33,8 @@ cfg_net_unix! { /// the stream in one direction. /// /// [`shutdown()`]: fn@crate::io::AsyncWriteExt::shutdown + /// [`UnixListener::accept`]: crate::net::UnixListener::accept + #[cfg_attr(docsrs, doc(alias = "uds"))] pub struct UnixStream { io: PollEvented<mio::net::UnixStream>, } @@ -59,12 +62,18 @@ impl UnixStream { Ok(stream) } - /// Wait for any of the requested ready states. + /// Waits for any of the requested ready states. /// /// This function is usually paired with `try_read()` or `try_write()`. It /// can be used to concurrently read / write to the same socket on a single /// task without splitting the socket. /// + /// The function may complete without the socket being ready. This is a + /// false-positive and attempting an operation will return with + /// `io::ErrorKind::WouldBlock`. The function can also return with an empty + /// [`Ready`] set, so you should always check the returned value and possibly + /// wait again if the requested states are not set. + /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method @@ -133,7 +142,7 @@ impl UnixStream { Ok(event.ready) } - /// Wait for the socket to become readable. + /// Waits for the socket to become readable. /// /// This function is equivalent to `ready(Interest::READABLE)` and is usually /// paired with `try_read()`. @@ -239,8 +248,12 @@ impl UnixStream { /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the - /// number of bytes read. `Ok(0)` indicates the stream's read half is closed - /// and will no longer yield data. If the stream is not ready to read data + /// number of bytes read. If `n` is `0`, then it can indicate one of two scenarios: + /// + /// 1. The stream's read half is closed and will no longer yield data. + /// 2. The specified buffer was 0 bytes in length. + /// + /// If the stream is not ready to read data, /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples @@ -290,7 +303,7 @@ impl UnixStream { .try_io(Interest::READABLE, || (&*self.io).read(buf)) } - /// Try to read data from the stream into the provided buffers, returning + /// Tries to read data from the stream into the provided buffers, returning /// how many bytes were read. /// /// Data is copied to fill each buffer in order, with the final buffer @@ -369,7 +382,7 @@ impl UnixStream { } cfg_io_util! { - /// Try to read data from the stream into the provided buffer, advancing the + /// Tries to read data from the stream into the provided buffer, advancing the /// buffer's internal cursor, returning how many bytes were read. /// /// Receives any pending data from the socket but does not wait for new data @@ -449,7 +462,7 @@ impl UnixStream { } } - /// Wait for the socket to become writable. + /// Waits for the socket to become writable. /// /// This function is equivalent to `ready(Interest::WRITABLE)` and is usually /// paired with `try_write()`. @@ -535,7 +548,7 @@ impl UnixStream { self.io.registration().poll_write_ready(cx).map_ok(|_| ()) } - /// Try to write a buffer to the stream, returning how many bytes were + /// Tries to write a buffer to the stream, returning how many bytes were /// written. /// /// The function will attempt to write the entire contents of `buf`, but @@ -591,7 +604,7 @@ impl UnixStream { .try_io(Interest::WRITABLE, || (&*self.io).write(buf)) } - /// Try to write several buffers to the stream, returning how many bytes + /// Tries to write several buffers to the stream, returning how many bytes /// were written. /// /// Data is written from each buffer in order, with the final buffer read @@ -653,20 +666,122 @@ impl UnixStream { .try_io(Interest::WRITABLE, || (&*self.io).write_vectored(buf)) } + /// Tries to read or write from the socket using a user-provided IO operation. + /// + /// If the socket is ready, the provided closure is called. The closure + /// should attempt to perform IO operation on the socket by manually + /// calling the appropriate syscall. If the operation fails because the + /// socket is not actually ready, then the closure should return a + /// `WouldBlock` error and the readiness flag is cleared. The return value + /// of the closure is then returned by `try_io`. + /// + /// If the socket is not ready, then the closure is not called + /// and a `WouldBlock` error is returned. + /// + /// The closure should only return a `WouldBlock` error if it has performed + /// an IO operation on the socket that failed due to the socket not being + /// ready. Returning a `WouldBlock` error in any other situation will + /// incorrectly clear the readiness flag, which can cause the socket to + /// behave incorrectly. + /// + /// The closure should not perform the IO operation using any of the methods + /// defined on the Tokio `UnixStream` type, as this will mess with the + /// readiness flag and can cause the socket to behave incorrectly. + /// + /// This method is not intended to be used with combined interests. + /// The closure should perform only one type of IO operation, so it should not + /// require more than one ready state. This method may panic or sleep forever + /// if it is called with a combined interest. + /// + /// Usually, [`readable()`], [`writable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: UnixStream::readable() + /// [`writable()`]: UnixStream::writable() + /// [`ready()`]: UnixStream::ready() + pub fn try_io<R>( + &self, + interest: Interest, + f: impl FnOnce() -> io::Result<R>, + ) -> io::Result<R> { + self.io + .registration() + .try_io(interest, || self.io.try_io(f)) + } + + /// Reads or writes from the socket using a user-provided IO operation. + /// + /// The readiness of the socket is awaited and when the socket is ready, + /// the provided closure is called. The closure should attempt to perform + /// IO operation on the socket by manually calling the appropriate syscall. + /// If the operation fails because the socket is not actually ready, + /// then the closure should return a `WouldBlock` error. In such case the + /// readiness flag is cleared and the socket readiness is awaited again. + /// This loop is repeated until the closure returns an `Ok` or an error + /// other than `WouldBlock`. + /// + /// The closure should only return a `WouldBlock` error if it has performed + /// an IO operation on the socket that failed due to the socket not being + /// ready. Returning a `WouldBlock` error in any other situation will + /// incorrectly clear the readiness flag, which can cause the socket to + /// behave incorrectly. + /// + /// The closure should not perform the IO operation using any of the methods + /// defined on the Tokio `UnixStream` type, as this will mess with the + /// readiness flag and can cause the socket to behave incorrectly. + /// + /// This method is not intended to be used with combined interests. + /// The closure should perform only one type of IO operation, so it should not + /// require more than one ready state. This method may panic or sleep forever + /// if it is called with a combined interest. + pub async fn async_io<R>( + &self, + interest: Interest, + mut f: impl FnMut() -> io::Result<R>, + ) -> io::Result<R> { + self.io + .registration() + .async_io(interest, || self.io.try_io(&mut f)) + .await + } + /// Creates new `UnixStream` from a `std::os::unix::net::UnixStream`. /// /// This function is intended to be used to wrap a UnixStream from the - /// standard library in the Tokio equivalent. The conversion assumes - /// nothing about the underlying stream; it is left up to the user to set - /// it in non-blocking mode. + /// standard library in the Tokio equivalent. + /// + /// # Notes + /// + /// The caller is responsible for ensuring that the stream is in + /// non-blocking mode. Otherwise all I/O operations on the stream + /// will block the thread, which will cause unexpected behavior. + /// Non-blocking mode can be set using [`set_nonblocking`]. + /// + /// [`set_nonblocking`]: std::os::unix::net::UnixStream::set_nonblocking + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::UnixStream; + /// use std::os::unix::net::UnixStream as StdUnixStream; + /// # use std::error::Error; + /// + /// # async fn dox() -> Result<(), Box<dyn Error>> { + /// let std_stream = StdUnixStream::connect("/path/to/the/socket")?; + /// std_stream.set_nonblocking(true)?; + /// let stream = UnixStream::from_std(std_stream)?; + /// # Ok(()) + /// # } + /// ``` /// /// # Panics /// - /// This function panics if thread-local runtime is not set. + /// This function panics if it is not called from within a runtime with + /// IO enabled. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. + #[track_caller] pub fn from_std(stream: net::UnixStream) -> io::Result<UnixStream> { let stream = mio::net::UnixStream::from_std(stream); let io = PollEvented::new(stream)?; @@ -674,7 +789,7 @@ impl UnixStream { Ok(UnixStream { io }) } - /// Turn a [`tokio::net::UnixStream`] into a [`std::os::unix::net::UnixStream`]. + /// Turns a [`tokio::net::UnixStream`] into a [`std::os::unix::net::UnixStream`]. /// /// The returned [`std::os::unix::net::UnixStream`] will have nonblocking /// mode set as `true`. Use [`set_nonblocking`] to change the blocking @@ -738,11 +853,41 @@ impl UnixStream { } /// Returns the socket address of the local half of this connection. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::UnixStream; + /// + /// # async fn dox() -> Result<(), Box<dyn std::error::Error>> { + /// let dir = tempfile::tempdir().unwrap(); + /// let bind_path = dir.path().join("bind_path"); + /// let stream = UnixStream::connect(bind_path).await?; + /// + /// println!("{:?}", stream.local_addr()?); + /// # Ok(()) + /// # } + /// ``` pub fn local_addr(&self) -> io::Result<SocketAddr> { self.io.local_addr().map(SocketAddr) } /// Returns the socket address of the remote half of this connection. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::UnixStream; + /// + /// # async fn dox() -> Result<(), Box<dyn std::error::Error>> { + /// let dir = tempfile::tempdir().unwrap(); + /// let bind_path = dir.path().join("bind_path"); + /// let stream = UnixStream::connect(bind_path).await?; + /// + /// println!("{:?}", stream.peer_addr()?); + /// # Ok(()) + /// # } + /// ``` pub fn peer_addr(&self) -> io::Result<SocketAddr> { self.io.peer_addr().map(SocketAddr) } @@ -769,7 +914,7 @@ impl UnixStream { // These lifetime markers also appear in the generated documentation, and make // it more clear that this is a *borrowed* split. #[allow(clippy::needless_lifetimes)] - /// Split a `UnixStream` into a read half and a write half, which can be used + /// Splits a `UnixStream` into a read half and a write half, which can be used /// to read and write the stream concurrently. /// /// This method is more efficient than [`into_split`], but the halves cannot be @@ -893,3 +1038,10 @@ impl AsRawFd for UnixStream { self.io.as_raw_fd() } } + +#[cfg(not(tokio_no_as_fd))] +impl AsFd for UnixStream { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } +} diff --git a/vendor/tokio/src/net/unix/ucred.rs b/vendor/tokio/src/net/unix/ucred.rs index 6310183d7..edfab08ab 100644 --- a/vendor/tokio/src/net/unix/ucred.rs +++ b/vendor/tokio/src/net/unix/ucred.rs @@ -1,24 +1,24 @@ -use libc::{gid_t, pid_t, uid_t}; +use crate::net::unix; -/// Credentials of a process +/// Credentials of a process. #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)] pub struct UCred { - /// PID (process ID) of the process - pid: Option<pid_t>, - /// UID (user ID) of the process - uid: uid_t, - /// GID (group ID) of the process - gid: gid_t, + /// PID (process ID) of the process. + pid: Option<unix::pid_t>, + /// UID (user ID) of the process. + uid: unix::uid_t, + /// GID (group ID) of the process. + gid: unix::gid_t, } impl UCred { /// Gets UID (user ID) of the process. - pub fn uid(&self) -> uid_t { + pub fn uid(&self) -> unix::uid_t { self.uid } /// Gets GID (group ID) of the process. - pub fn gid(&self) -> gid_t { + pub fn gid(&self) -> unix::gid_t { self.gid } @@ -26,20 +26,23 @@ impl UCred { /// /// This is only implemented under Linux, Android, iOS, macOS, Solaris and /// Illumos. On other platforms this will always return `None`. - pub fn pid(&self) -> Option<pid_t> { + pub fn pid(&self) -> Option<unix::pid_t> { self.pid } } -#[cfg(any(target_os = "linux", target_os = "android"))] -pub(crate) use self::impl_linux::get_peer_cred; - #[cfg(any( - target_os = "dragonfly", - target_os = "freebsd", - target_os = "netbsd", + target_os = "linux", + target_os = "redox", + target_os = "android", target_os = "openbsd" ))] +pub(crate) use self::impl_linux::get_peer_cred; + +#[cfg(any(target_os = "netbsd"))] +pub(crate) use self::impl_netbsd::get_peer_cred; + +#[cfg(any(target_os = "dragonfly", target_os = "freebsd"))] pub(crate) use self::impl_bsd::get_peer_cred; #[cfg(any(target_os = "macos", target_os = "ios"))] @@ -48,13 +51,24 @@ pub(crate) use self::impl_macos::get_peer_cred; #[cfg(any(target_os = "solaris", target_os = "illumos"))] pub(crate) use self::impl_solaris::get_peer_cred; -#[cfg(any(target_os = "linux", target_os = "android"))] +#[cfg(target_os = "aix")] +pub(crate) use self::impl_aix::get_peer_cred; + +#[cfg(any( + target_os = "linux", + target_os = "redox", + target_os = "android", + target_os = "openbsd" +))] pub(crate) mod impl_linux { - use crate::net::unix::UnixStream; + use crate::net::unix::{self, UnixStream}; use libc::{c_void, getsockopt, socklen_t, SOL_SOCKET, SO_PEERCRED}; use std::{io, mem}; + #[cfg(target_os = "openbsd")] + use libc::sockpeercred as ucred; + #[cfg(any(target_os = "linux", target_os = "redox", target_os = "android"))] use libc::ucred; pub(crate) fn get_peer_cred(sock: &UnixStream) -> io::Result<super::UCred> { @@ -86,9 +100,9 @@ pub(crate) mod impl_linux { ); if ret == 0 && ucred_size as usize == mem::size_of::<ucred>() { Ok(super::UCred { - uid: ucred.uid, - gid: ucred.gid, - pid: Some(ucred.pid), + uid: ucred.uid as unix::uid_t, + gid: ucred.gid as unix::gid_t, + pid: Some(ucred.pid as unix::pid_t), }) } else { Err(io::Error::last_os_error()) @@ -97,14 +111,51 @@ pub(crate) mod impl_linux { } } -#[cfg(any( - target_os = "dragonfly", - target_os = "freebsd", - target_os = "netbsd", - target_os = "openbsd" -))] +#[cfg(any(target_os = "netbsd"))] +pub(crate) mod impl_netbsd { + use crate::net::unix::{self, UnixStream}; + + use libc::{c_void, getsockopt, socklen_t, unpcbid, LOCAL_PEEREID, SOL_SOCKET}; + use std::io; + use std::mem::size_of; + use std::os::unix::io::AsRawFd; + + pub(crate) fn get_peer_cred(sock: &UnixStream) -> io::Result<super::UCred> { + unsafe { + let raw_fd = sock.as_raw_fd(); + + let mut unpcbid = unpcbid { + unp_pid: 0, + unp_euid: 0, + unp_egid: 0, + }; + + let unpcbid_size = size_of::<unpcbid>(); + let mut unpcbid_size = unpcbid_size as socklen_t; + + let ret = getsockopt( + raw_fd, + SOL_SOCKET, + LOCAL_PEEREID, + &mut unpcbid as *mut unpcbid as *mut c_void, + &mut unpcbid_size, + ); + if ret == 0 && unpcbid_size as usize == size_of::<unpcbid>() { + Ok(super::UCred { + uid: unpcbid.unp_euid as unix::uid_t, + gid: unpcbid.unp_egid as unix::gid_t, + pid: Some(unpcbid.unp_pid as unix::pid_t), + }) + } else { + Err(io::Error::last_os_error()) + } + } + } +} + +#[cfg(any(target_os = "dragonfly", target_os = "freebsd"))] pub(crate) mod impl_bsd { - use crate::net::unix::UnixStream; + use crate::net::unix::{self, UnixStream}; use libc::getpeereid; use std::io; @@ -122,8 +173,8 @@ pub(crate) mod impl_bsd { if ret == 0 { Ok(super::UCred { - uid: uid.assume_init(), - gid: gid.assume_init(), + uid: uid.assume_init() as unix::uid_t, + gid: gid.assume_init() as unix::gid_t, pid: None, }) } else { @@ -135,7 +186,7 @@ pub(crate) mod impl_bsd { #[cfg(any(target_os = "macos", target_os = "ios"))] pub(crate) mod impl_macos { - use crate::net::unix::UnixStream; + use crate::net::unix::{self, UnixStream}; use libc::{c_void, getpeereid, getsockopt, pid_t, LOCAL_PEEREPID, SOL_LOCAL}; use std::io; @@ -169,9 +220,9 @@ pub(crate) mod impl_macos { if ret == 0 { Ok(super::UCred { - uid: uid.assume_init(), - gid: gid.assume_init(), - pid: Some(pid.assume_init()), + uid: uid.assume_init() as unix::uid_t, + gid: gid.assume_init() as unix::gid_t, + pid: Some(pid.assume_init() as unix::pid_t), }) } else { Err(io::Error::last_os_error()) @@ -182,7 +233,7 @@ pub(crate) mod impl_macos { #[cfg(any(target_os = "solaris", target_os = "illumos"))] pub(crate) mod impl_solaris { - use crate::net::unix::UnixStream; + use crate::net::unix::{self, UnixStream}; use std::io; use std::os::unix::io::AsRawFd; use std::ptr; @@ -202,9 +253,37 @@ pub(crate) mod impl_solaris { libc::ucred_free(cred); Ok(super::UCred { - uid, - gid, - pid: Some(pid), + uid: uid as unix::uid_t, + gid: gid as unix::gid_t, + pid: Some(pid as unix::pid_t), + }) + } else { + Err(io::Error::last_os_error()) + } + } + } +} + +#[cfg(target_os = "aix")] +pub(crate) mod impl_aix { + use crate::net::unix::UnixStream; + use std::io; + use std::os::unix::io::AsRawFd; + + pub(crate) fn get_peer_cred(sock: &UnixStream) -> io::Result<super::UCred> { + unsafe { + let raw_fd = sock.as_raw_fd(); + + let mut uid = std::mem::MaybeUninit::uninit(); + let mut gid = std::mem::MaybeUninit::uninit(); + + let ret = libc::getpeereid(raw_fd, uid.as_mut_ptr(), gid.as_mut_ptr()); + + if ret == 0 { + Ok(super::UCred { + uid: uid.assume_init(), + gid: gid.assume_init(), + pid: None, }) } else { Err(io::Error::last_os_error()) diff --git a/vendor/tokio/src/net/windows/named_pipe.rs b/vendor/tokio/src/net/windows/named_pipe.rs index b9f7d49d7..91f016267 100644 --- a/vendor/tokio/src/net/windows/named_pipe.rs +++ b/vendor/tokio/src/net/windows/named_pipe.rs @@ -10,27 +10,30 @@ use std::ptr; use std::task::{Context, Poll}; use crate::io::{AsyncRead, AsyncWrite, Interest, PollEvented, ReadBuf, Ready}; +#[cfg(not(tokio_no_as_fd))] +use crate::os::windows::io::{AsHandle, BorrowedHandle}; use crate::os::windows::io::{AsRawHandle, FromRawHandle, RawHandle}; +cfg_io_util! { + use bytes::BufMut; +} + // Hide imports which are not used when generating documentation. #[cfg(not(docsrs))] mod doc { pub(super) use crate::os::windows::ffi::OsStrExt; - pub(super) use crate::winapi::shared::minwindef::{DWORD, FALSE}; - pub(super) use crate::winapi::um::fileapi; - pub(super) use crate::winapi::um::handleapi; - pub(super) use crate::winapi::um::namedpipeapi; - pub(super) use crate::winapi::um::winbase; - pub(super) use crate::winapi::um::winnt; - + pub(super) mod windows_sys { + pub(crate) use windows_sys::{ + Win32::Foundation::*, Win32::Storage::FileSystem::*, Win32::System::Pipes::*, + Win32::System::SystemServices::*, + }; + } pub(super) use mio::windows as mio_windows; } // NB: none of these shows up in public API, so don't document them. #[cfg(docsrs)] mod doc { - pub type DWORD = crate::doc::NotDefinedHere; - pub(super) mod mio_windows { pub type NamedPipe = crate::doc::NotDefinedHere; } @@ -97,7 +100,6 @@ use self::doc::*; /// # Ok(()) } /// ``` /// -/// [`ERROR_PIPE_BUSY`]: crate::winapi::shared::winerror::ERROR_PIPE_BUSY /// [Windows named pipe]: https://docs.microsoft.com/en-us/windows/win32/ipc/named-pipes #[derive(Debug)] pub struct NamedPipeServer { @@ -105,7 +107,7 @@ pub struct NamedPipeServer { } impl NamedPipeServer { - /// Construct a new named pipe server from the specified raw handle. + /// Constructs a new named pipe server from the specified raw handle. /// /// This function will consume ownership of the handle given, passing /// responsibility for closing the handle to the returned object. @@ -188,17 +190,15 @@ impl NamedPipeServer { /// # Ok(()) } /// ``` pub async fn connect(&self) -> io::Result<()> { - loop { - match self.io.connect() { - Ok(()) => break, - Err(e) if e.kind() == io::ErrorKind::WouldBlock => { - self.io.registration().readiness(Interest::WRITABLE).await?; - } - Err(e) => return Err(e), + match self.io.connect() { + Err(e) if e.kind() == io::ErrorKind::WouldBlock => { + self.io + .registration() + .async_io(Interest::WRITABLE, || self.io.connect()) + .await } + x => x, } - - Ok(()) } /// Disconnects the server end of a named pipe instance from a client @@ -207,7 +207,7 @@ impl NamedPipeServer { /// ``` /// use tokio::io::AsyncWriteExt; /// use tokio::net::windows::named_pipe::{ClientOptions, ServerOptions}; - /// use winapi::shared::winerror; + /// use windows_sys::Win32::Foundation::ERROR_PIPE_NOT_CONNECTED; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-disconnect"; /// @@ -227,19 +227,25 @@ impl NamedPipeServer { /// // Write fails with an OS-specific error after client has been /// // disconnected. /// let e = client.write(b"ping").await.unwrap_err(); - /// assert_eq!(e.raw_os_error(), Some(winerror::ERROR_PIPE_NOT_CONNECTED as i32)); + /// assert_eq!(e.raw_os_error(), Some(ERROR_PIPE_NOT_CONNECTED as i32)); /// # Ok(()) } /// ``` pub fn disconnect(&self) -> io::Result<()> { self.io.disconnect() } - /// Wait for any of the requested ready states. + /// Waits for any of the requested ready states. /// /// This function is usually paired with `try_read()` or `try_write()`. It /// can be used to concurrently read / write to the same pipe on a single /// task without splitting the pipe. /// + /// The function may complete without the pipe being ready. This is a + /// false-positive and attempting an operation will return with + /// `io::ErrorKind::WouldBlock`. The function can also return with an empty + /// [`Ready`] set, so you should always check the returned value and possibly + /// wait again if the requested states are not set. + /// /// # Examples /// /// Concurrently read and write to the pipe on the same task without @@ -301,7 +307,7 @@ impl NamedPipeServer { Ok(event.ready) } - /// Wait for the pipe to become readable. + /// Waits for the pipe to become readable. /// /// This function is equivalent to `ready(Interest::READABLE)` and is usually /// paired with `try_read()`. @@ -383,7 +389,7 @@ impl NamedPipeServer { self.io.registration().poll_read_ready(cx).map_ok(|_| ()) } - /// Try to read data from the pipe into the provided buffer, returning how + /// Tries to read data from the pipe into the provided buffer, returning how /// many bytes were read. /// /// Receives any pending data from the pipe but does not wait for new data @@ -399,8 +405,12 @@ impl NamedPipeServer { /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the - /// number of bytes read. `Ok(0)` indicates the pipe's read half is closed - /// and will no longer yield data. If the pipe is not ready to read data + /// number of bytes read. If `n` is `0`, then it can indicate one of two scenarios: + /// + /// 1. The pipe's read half is closed and will no longer yield data. + /// 2. The specified buffer was 0 bytes in length. + /// + /// If the pipe is not ready to read data, /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples @@ -450,7 +460,7 @@ impl NamedPipeServer { .try_io(Interest::READABLE, || (&*self.io).read(buf)) } - /// Try to read data from the pipe into the provided buffers, returning + /// Tries to read data from the pipe into the provided buffers, returning /// how many bytes were read. /// /// Data is copied to fill each buffer in order, with the final buffer @@ -528,7 +538,87 @@ impl NamedPipeServer { .try_io(Interest::READABLE, || (&*self.io).read_vectored(bufs)) } - /// Wait for the pipe to become writable. + cfg_io_util! { + /// Tries to read data from the stream into the provided buffer, advancing the + /// buffer's internal cursor, returning how many bytes were read. + /// + /// Receives any pending data from the pipe but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read_buf()` is non-blocking, the buffer does not have to be stored by + /// the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: NamedPipeServer::readable() + /// [`ready()`]: NamedPipeServer::ready() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. `Ok(0)` indicates the stream's read half is closed + /// and will no longer yield data. If the stream is not ready to read data + /// `Err(io::ErrorKind::WouldBlock)` is returned. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::windows::named_pipe; + /// use std::error::Error; + /// use std::io; + /// + /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-client-readable"; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// let server = named_pipe::ServerOptions::new().create(PIPE_NAME)?; + /// + /// loop { + /// // Wait for the pipe to be readable + /// server.readable().await?; + /// + /// let mut buf = Vec::with_capacity(4096); + /// + /// // Try to read data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match server.try_read_buf(&mut buf) { + /// Ok(0) => break, + /// Ok(n) => { + /// println!("read {} bytes", n); + /// } + /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// + /// Ok(()) + /// } + /// ``` + pub fn try_read_buf<B: BufMut>(&self, buf: &mut B) -> io::Result<usize> { + self.io.registration().try_io(Interest::READABLE, || { + use std::io::Read; + + let dst = buf.chunk_mut(); + let dst = + unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]) }; + + // Safety: We trust `NamedPipeServer::read` to have filled up `n` bytes in the + // buffer. + let n = (&*self.io).read(dst)?; + + unsafe { + buf.advance_mut(n); + } + + Ok(n) + }) + } + } + + /// Waits for the pipe to become writable. /// /// This function is equivalent to `ready(Interest::WRITABLE)` and is usually /// paired with `try_write()`. @@ -606,7 +696,7 @@ impl NamedPipeServer { self.io.registration().poll_write_ready(cx).map_ok(|_| ()) } - /// Try to write a buffer to the pipe, returning how many bytes were + /// Tries to write a buffer to the pipe, returning how many bytes were /// written. /// /// The function will attempt to write the entire contents of `buf`, but @@ -662,7 +752,7 @@ impl NamedPipeServer { .try_io(Interest::WRITABLE, || (&*self.io).write(buf)) } - /// Try to write several buffers to the pipe, returning how many bytes + /// Tries to write several buffers to the pipe, returning how many bytes /// were written. /// /// Data is written from each buffer in order, with the final buffer read @@ -723,6 +813,79 @@ impl NamedPipeServer { .registration() .try_io(Interest::WRITABLE, || (&*self.io).write_vectored(buf)) } + + /// Tries to read or write from the pipe using a user-provided IO operation. + /// + /// If the pipe is ready, the provided closure is called. The closure + /// should attempt to perform IO operation from the pipe by manually + /// calling the appropriate syscall. If the operation fails because the + /// pipe is not actually ready, then the closure should return a + /// `WouldBlock` error and the readiness flag is cleared. The return value + /// of the closure is then returned by `try_io`. + /// + /// If the pipe is not ready, then the closure is not called + /// and a `WouldBlock` error is returned. + /// + /// The closure should only return a `WouldBlock` error if it has performed + /// an IO operation on the pipe that failed due to the pipe not being + /// ready. Returning a `WouldBlock` error in any other situation will + /// incorrectly clear the readiness flag, which can cause the pipe to + /// behave incorrectly. + /// + /// The closure should not perform the IO operation using any of the + /// methods defined on the Tokio `NamedPipeServer` type, as this will mess with + /// the readiness flag and can cause the pipe to behave incorrectly. + /// + /// This method is not intended to be used with combined interests. + /// The closure should perform only one type of IO operation, so it should not + /// require more than one ready state. This method may panic or sleep forever + /// if it is called with a combined interest. + /// + /// Usually, [`readable()`], [`writable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: NamedPipeServer::readable() + /// [`writable()`]: NamedPipeServer::writable() + /// [`ready()`]: NamedPipeServer::ready() + pub fn try_io<R>( + &self, + interest: Interest, + f: impl FnOnce() -> io::Result<R>, + ) -> io::Result<R> { + self.io.registration().try_io(interest, f) + } + + /// Reads or writes from the pipe using a user-provided IO operation. + /// + /// The readiness of the pipe is awaited and when the pipe is ready, + /// the provided closure is called. The closure should attempt to perform + /// IO operation on the pipe by manually calling the appropriate syscall. + /// If the operation fails because the pipe is not actually ready, + /// then the closure should return a `WouldBlock` error. In such case the + /// readiness flag is cleared and the pipe readiness is awaited again. + /// This loop is repeated until the closure returns an `Ok` or an error + /// other than `WouldBlock`. + /// + /// The closure should only return a `WouldBlock` error if it has performed + /// an IO operation on the pipe that failed due to the pipe not being + /// ready. Returning a `WouldBlock` error in any other situation will + /// incorrectly clear the readiness flag, which can cause the pipe to + /// behave incorrectly. + /// + /// The closure should not perform the IO operation using any of the methods + /// defined on the Tokio `NamedPipeServer` type, as this will mess with the + /// readiness flag and can cause the pipe to behave incorrectly. + /// + /// This method is not intended to be used with combined interests. + /// The closure should perform only one type of IO operation, so it should not + /// require more than one ready state. This method may panic or sleep forever + /// if it is called with a combined interest. + pub async fn async_io<R>( + &self, + interest: Interest, + f: impl FnMut() -> io::Result<R>, + ) -> io::Result<R> { + self.io.registration().async_io(interest, f).await + } } impl AsyncRead for NamedPipeServer { @@ -767,6 +930,13 @@ impl AsRawHandle for NamedPipeServer { } } +#[cfg(not(tokio_no_as_fd))] +impl AsHandle for NamedPipeServer { + fn as_handle(&self) -> BorrowedHandle<'_> { + unsafe { BorrowedHandle::borrow_raw(self.as_raw_handle()) } + } +} + /// A [Windows named pipe] client. /// /// Constructed using [`ClientOptions::open`]. @@ -784,7 +954,7 @@ impl AsRawHandle for NamedPipeServer { /// use std::time::Duration; /// use tokio::net::windows::named_pipe::ClientOptions; /// use tokio::time; -/// use winapi::shared::winerror; +/// use windows_sys::Win32::Foundation::ERROR_PIPE_BUSY; /// /// const PIPE_NAME: &str = r"\\.\pipe\named-pipe-idiomatic-client"; /// @@ -792,7 +962,7 @@ impl AsRawHandle for NamedPipeServer { /// let client = loop { /// match ClientOptions::new().open(PIPE_NAME) { /// Ok(client) => break client, -/// Err(e) if e.raw_os_error() == Some(winerror::ERROR_PIPE_BUSY as i32) => (), +/// Err(e) if e.raw_os_error() == Some(ERROR_PIPE_BUSY as i32) => (), /// Err(e) => return Err(e), /// } /// @@ -803,7 +973,7 @@ impl AsRawHandle for NamedPipeServer { /// # Ok(()) } /// ``` /// -/// [`ERROR_PIPE_BUSY`]: crate::winapi::shared::winerror::ERROR_PIPE_BUSY +/// [`ERROR_PIPE_BUSY`]: https://docs.rs/windows-sys/latest/windows_sys/Win32/Foundation/constant.ERROR_PIPE_BUSY.html /// [Windows named pipe]: https://docs.microsoft.com/en-us/windows/win32/ipc/named-pipes #[derive(Debug)] pub struct NamedPipeClient { @@ -811,7 +981,7 @@ pub struct NamedPipeClient { } impl NamedPipeClient { - /// Construct a new named pipe client from the specified raw handle. + /// Constructs a new named pipe client from the specified raw handle. /// /// This function will consume ownership of the handle given, passing /// responsibility for closing the handle to the returned object. @@ -861,12 +1031,18 @@ impl NamedPipeClient { unsafe { named_pipe_info(self.io.as_raw_handle()) } } - /// Wait for any of the requested ready states. + /// Waits for any of the requested ready states. /// /// This function is usually paired with `try_read()` or `try_write()`. It /// can be used to concurrently read / write to the same pipe on a single /// task without splitting the pipe. /// + /// The function may complete without the pipe being ready. This is a + /// false-positive and attempting an operation will return with + /// `io::ErrorKind::WouldBlock`. The function can also return with an empty + /// [`Ready`] set, so you should always check the returned value and possibly + /// wait again if the requested states are not set. + /// /// # Examples /// /// Concurrently read and write to the pipe on the same task without @@ -927,7 +1103,7 @@ impl NamedPipeClient { Ok(event.ready) } - /// Wait for the pipe to become readable. + /// Waits for the pipe to become readable. /// /// This function is equivalent to `ready(Interest::READABLE)` and is usually /// paired with `try_read()`. @@ -1008,7 +1184,7 @@ impl NamedPipeClient { self.io.registration().poll_read_ready(cx).map_ok(|_| ()) } - /// Try to read data from the pipe into the provided buffer, returning how + /// Tries to read data from the pipe into the provided buffer, returning how /// many bytes were read. /// /// Receives any pending data from the pipe but does not wait for new data @@ -1024,8 +1200,12 @@ impl NamedPipeClient { /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the - /// number of bytes read. `Ok(0)` indicates the pipe's read half is closed - /// and will no longer yield data. If the pipe is not ready to read data + /// number of bytes read. If `n` is `0`, then it can indicate one of two scenarios: + /// + /// 1. The pipe's read half is closed and will no longer yield data. + /// 2. The specified buffer was 0 bytes in length. + /// + /// If the pipe is not ready to read data, /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples @@ -1074,7 +1254,7 @@ impl NamedPipeClient { .try_io(Interest::READABLE, || (&*self.io).read(buf)) } - /// Try to read data from the pipe into the provided buffers, returning + /// Tries to read data from the pipe into the provided buffers, returning /// how many bytes were read. /// /// Data is copied to fill each buffer in order, with the final buffer @@ -1151,7 +1331,87 @@ impl NamedPipeClient { .try_io(Interest::READABLE, || (&*self.io).read_vectored(bufs)) } - /// Wait for the pipe to become writable. + cfg_io_util! { + /// Tries to read data from the stream into the provided buffer, advancing the + /// buffer's internal cursor, returning how many bytes were read. + /// + /// Receives any pending data from the pipe but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read_buf()` is non-blocking, the buffer does not have to be stored by + /// the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: NamedPipeClient::readable() + /// [`ready()`]: NamedPipeClient::ready() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. `Ok(0)` indicates the stream's read half is closed + /// and will no longer yield data. If the stream is not ready to read data + /// `Err(io::ErrorKind::WouldBlock)` is returned. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::windows::named_pipe; + /// use std::error::Error; + /// use std::io; + /// + /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-client-readable"; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// let client = named_pipe::ClientOptions::new().open(PIPE_NAME)?; + /// + /// loop { + /// // Wait for the pipe to be readable + /// client.readable().await?; + /// + /// let mut buf = Vec::with_capacity(4096); + /// + /// // Try to read data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match client.try_read_buf(&mut buf) { + /// Ok(0) => break, + /// Ok(n) => { + /// println!("read {} bytes", n); + /// } + /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// + /// Ok(()) + /// } + /// ``` + pub fn try_read_buf<B: BufMut>(&self, buf: &mut B) -> io::Result<usize> { + self.io.registration().try_io(Interest::READABLE, || { + use std::io::Read; + + let dst = buf.chunk_mut(); + let dst = + unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]) }; + + // Safety: We trust `NamedPipeClient::read` to have filled up `n` bytes in the + // buffer. + let n = (&*self.io).read(dst)?; + + unsafe { + buf.advance_mut(n); + } + + Ok(n) + }) + } + } + + /// Waits for the pipe to become writable. /// /// This function is equivalent to `ready(Interest::WRITABLE)` and is usually /// paired with `try_write()`. @@ -1228,7 +1488,7 @@ impl NamedPipeClient { self.io.registration().poll_write_ready(cx).map_ok(|_| ()) } - /// Try to write a buffer to the pipe, returning how many bytes were + /// Tries to write a buffer to the pipe, returning how many bytes were /// written. /// /// The function will attempt to write the entire contents of `buf`, but @@ -1283,7 +1543,7 @@ impl NamedPipeClient { .try_io(Interest::WRITABLE, || (&*self.io).write(buf)) } - /// Try to write several buffers to the pipe, returning how many bytes + /// Tries to write several buffers to the pipe, returning how many bytes /// were written. /// /// Data is written from each buffer in order, with the final buffer read @@ -1343,6 +1603,79 @@ impl NamedPipeClient { .registration() .try_io(Interest::WRITABLE, || (&*self.io).write_vectored(buf)) } + + /// Tries to read or write from the pipe using a user-provided IO operation. + /// + /// If the pipe is ready, the provided closure is called. The closure + /// should attempt to perform IO operation from the pipe by manually + /// calling the appropriate syscall. If the operation fails because the + /// pipe is not actually ready, then the closure should return a + /// `WouldBlock` error and the readiness flag is cleared. The return value + /// of the closure is then returned by `try_io`. + /// + /// If the pipe is not ready, then the closure is not called + /// and a `WouldBlock` error is returned. + /// + /// The closure should only return a `WouldBlock` error if it has performed + /// an IO operation on the pipe that failed due to the pipe not being + /// ready. Returning a `WouldBlock` error in any other situation will + /// incorrectly clear the readiness flag, which can cause the pipe to + /// behave incorrectly. + /// + /// The closure should not perform the IO operation using any of the methods + /// defined on the Tokio `NamedPipeClient` type, as this will mess with the + /// readiness flag and can cause the pipe to behave incorrectly. + /// + /// This method is not intended to be used with combined interests. + /// The closure should perform only one type of IO operation, so it should not + /// require more than one ready state. This method may panic or sleep forever + /// if it is called with a combined interest. + /// + /// Usually, [`readable()`], [`writable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: NamedPipeClient::readable() + /// [`writable()`]: NamedPipeClient::writable() + /// [`ready()`]: NamedPipeClient::ready() + pub fn try_io<R>( + &self, + interest: Interest, + f: impl FnOnce() -> io::Result<R>, + ) -> io::Result<R> { + self.io.registration().try_io(interest, f) + } + + /// Reads or writes from the pipe using a user-provided IO operation. + /// + /// The readiness of the pipe is awaited and when the pipe is ready, + /// the provided closure is called. The closure should attempt to perform + /// IO operation on the pipe by manually calling the appropriate syscall. + /// If the operation fails because the pipe is not actually ready, + /// then the closure should return a `WouldBlock` error. In such case the + /// readiness flag is cleared and the pipe readiness is awaited again. + /// This loop is repeated until the closure returns an `Ok` or an error + /// other than `WouldBlock`. + /// + /// The closure should only return a `WouldBlock` error if it has performed + /// an IO operation on the pipe that failed due to the pipe not being + /// ready. Returning a `WouldBlock` error in any other situation will + /// incorrectly clear the readiness flag, which can cause the pipe to + /// behave incorrectly. + /// + /// The closure should not perform the IO operation using any of the methods + /// defined on the Tokio `NamedPipeClient` type, as this will mess with the + /// readiness flag and can cause the pipe to behave incorrectly. + /// + /// This method is not intended to be used with combined interests. + /// The closure should perform only one type of IO operation, so it should not + /// require more than one ready state. This method may panic or sleep forever + /// if it is called with a combined interest. + pub async fn async_io<R>( + &self, + interest: Interest, + f: impl FnMut() -> io::Result<R>, + ) -> io::Result<R> { + self.io.registration().async_io(interest, f).await + } } impl AsyncRead for NamedPipeClient { @@ -1387,17 +1720,11 @@ impl AsRawHandle for NamedPipeClient { } } -// Helper to set a boolean flag as a bitfield. -macro_rules! bool_flag { - ($f:expr, $t:expr, $flag:expr) => {{ - let current = $f; - - if $t { - $f = current | $flag; - } else { - $f = current & !$flag; - }; - }}; +#[cfg(not(tokio_no_as_fd))] +impl AsHandle for NamedPipeClient { + fn as_handle(&self) -> BorrowedHandle<'_> { + unsafe { BorrowedHandle::borrow_raw(self.as_raw_handle()) } + } } /// A builder structure for construct a named pipe with named pipe-specific @@ -1407,12 +1734,21 @@ macro_rules! bool_flag { /// See [`ServerOptions::create`]. #[derive(Debug, Clone)] pub struct ServerOptions { - open_mode: DWORD, - pipe_mode: DWORD, - max_instances: DWORD, - out_buffer_size: DWORD, - in_buffer_size: DWORD, - default_timeout: DWORD, + // dwOpenMode + access_inbound: bool, + access_outbound: bool, + first_pipe_instance: bool, + write_dac: bool, + write_owner: bool, + access_system_security: bool, + // dwPipeMode + pipe_mode: PipeMode, + reject_remote_clients: bool, + // other options + max_instances: u32, + out_buffer_size: u32, + in_buffer_size: u32, + default_timeout: u32, } impl ServerOptions { @@ -1429,9 +1765,15 @@ impl ServerOptions { /// ``` pub fn new() -> ServerOptions { ServerOptions { - open_mode: winbase::PIPE_ACCESS_DUPLEX | winbase::FILE_FLAG_OVERLAPPED, - pipe_mode: winbase::PIPE_TYPE_BYTE | winbase::PIPE_REJECT_REMOTE_CLIENTS, - max_instances: winbase::PIPE_UNLIMITED_INSTANCES, + access_inbound: true, + access_outbound: true, + first_pipe_instance: false, + write_dac: false, + write_owner: false, + access_system_security: false, + pipe_mode: PipeMode::Byte, + reject_remote_clients: true, + max_instances: windows_sys::PIPE_UNLIMITED_INSTANCES, out_buffer_size: 65536, in_buffer_size: 65536, default_timeout: 0, @@ -1443,15 +1785,11 @@ impl ServerOptions { /// The default pipe mode is [`PipeMode::Byte`]. See [`PipeMode`] for /// documentation of what each mode means. /// - /// This corresponding to specifying [`dwPipeMode`]. + /// This corresponds to specifying `PIPE_TYPE_` and `PIPE_READMODE_` in [`dwPipeMode`]. /// /// [`dwPipeMode`]: https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-createnamedpipea pub fn pipe_mode(&mut self, pipe_mode: PipeMode) -> &mut Self { - self.pipe_mode = match pipe_mode { - PipeMode::Byte => winbase::PIPE_TYPE_BYTE, - PipeMode::Message => winbase::PIPE_TYPE_MESSAGE, - }; - + self.pipe_mode = pipe_mode; self } @@ -1547,7 +1885,7 @@ impl ServerOptions { /// # Ok(()) } /// ``` pub fn access_inbound(&mut self, allowed: bool) -> &mut Self { - bool_flag!(self.open_mode, allowed, winbase::PIPE_ACCESS_INBOUND); + self.access_inbound = allowed; self } @@ -1645,7 +1983,7 @@ impl ServerOptions { /// # Ok(()) } /// ``` pub fn access_outbound(&mut self, allowed: bool) -> &mut Self { - bool_flag!(self.open_mode, allowed, winbase::PIPE_ACCESS_OUTBOUND); + self.access_outbound = allowed; self } @@ -1713,11 +2051,109 @@ impl ServerOptions { /// [`create`]: ServerOptions::create /// [`FILE_FLAG_FIRST_PIPE_INSTANCE`]: https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-createnamedpipea#pipe_first_pipe_instance pub fn first_pipe_instance(&mut self, first: bool) -> &mut Self { - bool_flag!( - self.open_mode, - first, - winbase::FILE_FLAG_FIRST_PIPE_INSTANCE - ); + self.first_pipe_instance = first; + self + } + + /// Requests permission to modify the pipe's discretionary access control list. + /// + /// This corresponds to setting [`WRITE_DAC`] in dwOpenMode. + /// + /// # Examples + /// + /// ``` + /// use std::{io, os::windows::prelude::AsRawHandle, ptr}; + // + /// use tokio::net::windows::named_pipe::ServerOptions; + /// use windows_sys::{ + /// Win32::Foundation::ERROR_SUCCESS, + /// Win32::Security::DACL_SECURITY_INFORMATION, + /// Win32::Security::Authorization::{SetSecurityInfo, SE_KERNEL_OBJECT}, + /// }; + /// + /// const PIPE_NAME: &str = r"\\.\pipe\write_dac_pipe"; + /// + /// # #[tokio::main] async fn main() -> io::Result<()> { + /// let mut pipe_template = ServerOptions::new(); + /// pipe_template.write_dac(true); + /// let pipe = pipe_template.create(PIPE_NAME)?; + /// + /// unsafe { + /// assert_eq!( + /// ERROR_SUCCESS, + /// SetSecurityInfo( + /// pipe.as_raw_handle() as _, + /// SE_KERNEL_OBJECT, + /// DACL_SECURITY_INFORMATION, + /// ptr::null_mut(), + /// ptr::null_mut(), + /// ptr::null_mut(), + /// ptr::null_mut(), + /// ) + /// ); + /// } + /// + /// # Ok(()) } + /// ``` + /// + /// ``` + /// use std::{io, os::windows::prelude::AsRawHandle, ptr}; + // + /// use tokio::net::windows::named_pipe::ServerOptions; + /// use windows_sys::{ + /// Win32::Foundation::ERROR_ACCESS_DENIED, + /// Win32::Security::DACL_SECURITY_INFORMATION, + /// Win32::Security::Authorization::{SetSecurityInfo, SE_KERNEL_OBJECT}, + /// }; + /// + /// const PIPE_NAME: &str = r"\\.\pipe\write_dac_pipe_fail"; + /// + /// # #[tokio::main] async fn main() -> io::Result<()> { + /// let mut pipe_template = ServerOptions::new(); + /// pipe_template.write_dac(false); + /// let pipe = pipe_template.create(PIPE_NAME)?; + /// + /// unsafe { + /// assert_eq!( + /// ERROR_ACCESS_DENIED, + /// SetSecurityInfo( + /// pipe.as_raw_handle() as _, + /// SE_KERNEL_OBJECT, + /// DACL_SECURITY_INFORMATION, + /// ptr::null_mut(), + /// ptr::null_mut(), + /// ptr::null_mut(), + /// ptr::null_mut(), + /// ) + /// ); + /// } + /// + /// # Ok(()) } + /// ``` + /// + /// [`WRITE_DAC`]: https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-createnamedpipea + pub fn write_dac(&mut self, requested: bool) -> &mut Self { + self.write_dac = requested; + self + } + + /// Requests permission to modify the pipe's owner. + /// + /// This corresponds to setting [`WRITE_OWNER`] in dwOpenMode. + /// + /// [`WRITE_OWNER`]: https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-createnamedpipea + pub fn write_owner(&mut self, requested: bool) -> &mut Self { + self.write_owner = requested; + self + } + + /// Requests permission to modify the pipe's system access control list. + /// + /// This corresponds to setting [`ACCESS_SYSTEM_SECURITY`] in dwOpenMode. + /// + /// [`ACCESS_SYSTEM_SECURITY`]: https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-createnamedpipea + pub fn access_system_security(&mut self, requested: bool) -> &mut Self { + self.access_system_security = requested; self } @@ -1728,7 +2164,7 @@ impl ServerOptions { /// /// [`PIPE_REJECT_REMOTE_CLIENTS`]: https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-createnamedpipea#pipe_reject_remote_clients pub fn reject_remote_clients(&mut self, reject: bool) -> &mut Self { - bool_flag!(self.pipe_mode, reject, winbase::PIPE_REJECT_REMOTE_CLIENTS); + self.reject_remote_clients = reject; self } @@ -1750,7 +2186,7 @@ impl ServerOptions { /// ``` /// use std::io; /// use tokio::net::windows::named_pipe::{ServerOptions, ClientOptions}; - /// use winapi::shared::winerror; + /// use windows_sys::Win32::Foundation::ERROR_PIPE_BUSY; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-max-instances"; /// @@ -1766,11 +2202,11 @@ impl ServerOptions { /// /// // Too many servers! /// let e = server.create(PIPE_NAME).unwrap_err(); - /// assert_eq!(e.raw_os_error(), Some(winerror::ERROR_PIPE_BUSY as i32)); + /// assert_eq!(e.raw_os_error(), Some(ERROR_PIPE_BUSY as i32)); /// /// // Still too many servers even if we specify a higher value! /// let e = server.max_instances(100).create(PIPE_NAME).unwrap_err(); - /// assert_eq!(e.raw_os_error(), Some(winerror::ERROR_PIPE_BUSY as i32)); + /// assert_eq!(e.raw_os_error(), Some(ERROR_PIPE_BUSY as i32)); /// # Ok(()) } /// ``` /// @@ -1786,9 +2222,10 @@ impl ServerOptions { /// let builder = ServerOptions::new().max_instances(255); /// # Ok(()) } /// ``` + #[track_caller] pub fn max_instances(&mut self, instances: usize) -> &mut Self { assert!(instances < 255, "cannot specify more than 254 instances"); - self.max_instances = instances as DWORD; + self.max_instances = instances as u32; self } @@ -1798,7 +2235,7 @@ impl ServerOptions { /// /// [`nOutBufferSize`]: https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-createnamedpipea pub fn out_buffer_size(&mut self, buffer: u32) -> &mut Self { - self.out_buffer_size = buffer as DWORD; + self.out_buffer_size = buffer; self } @@ -1808,11 +2245,11 @@ impl ServerOptions { /// /// [`nInBufferSize`]: https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-createnamedpipea pub fn in_buffer_size(&mut self, buffer: u32) -> &mut Self { - self.in_buffer_size = buffer as DWORD; + self.in_buffer_size = buffer; self } - /// Create the named pipe identified by `addr` for use as a server. + /// Creates the named pipe identified by `addr` for use as a server. /// /// This uses the [`CreateNamedPipe`] function. /// @@ -1843,7 +2280,7 @@ impl ServerOptions { unsafe { self.create_with_security_attributes_raw(addr, ptr::null_mut()) } } - /// Create the named pipe identified by `addr` for use as a server. + /// Creates the named pipe identified by `addr` for use as a server. /// /// This is the same as [`create`] except that it supports providing the raw /// pointer to a structure of [`SECURITY_ATTRIBUTES`] which will be passed @@ -1865,7 +2302,7 @@ impl ServerOptions { /// /// [`create`]: ServerOptions::create /// [`CreateFile`]: https://docs.microsoft.com/en-us/windows/win32/api/fileapi/nf-fileapi-createfilew - /// [`SECURITY_ATTRIBUTES`]: crate::winapi::um::minwinbase::SECURITY_ATTRIBUTES + /// [`SECURITY_ATTRIBUTES`]: https://docs.rs/windows-sys/latest/windows_sys/Win32/Security/struct.SECURITY_ATTRIBUTES.html pub unsafe fn create_with_security_attributes_raw( &self, addr: impl AsRef<OsStr>, @@ -1873,10 +2310,46 @@ impl ServerOptions { ) -> io::Result<NamedPipeServer> { let addr = encode_addr(addr); - let h = namedpipeapi::CreateNamedPipeW( + let pipe_mode = { + let mut mode = if matches!(self.pipe_mode, PipeMode::Message) { + windows_sys::PIPE_TYPE_MESSAGE | windows_sys::PIPE_READMODE_MESSAGE + } else { + windows_sys::PIPE_TYPE_BYTE | windows_sys::PIPE_READMODE_BYTE + }; + if self.reject_remote_clients { + mode |= windows_sys::PIPE_REJECT_REMOTE_CLIENTS; + } else { + mode |= windows_sys::PIPE_ACCEPT_REMOTE_CLIENTS; + } + mode + }; + let open_mode = { + let mut mode = windows_sys::FILE_FLAG_OVERLAPPED; + if self.access_inbound { + mode |= windows_sys::PIPE_ACCESS_INBOUND; + } + if self.access_outbound { + mode |= windows_sys::PIPE_ACCESS_OUTBOUND; + } + if self.first_pipe_instance { + mode |= windows_sys::FILE_FLAG_FIRST_PIPE_INSTANCE; + } + if self.write_dac { + mode |= windows_sys::WRITE_DAC; + } + if self.write_owner { + mode |= windows_sys::WRITE_OWNER; + } + if self.access_system_security { + mode |= windows_sys::ACCESS_SYSTEM_SECURITY; + } + mode + }; + + let h = windows_sys::CreateNamedPipeW( addr.as_ptr(), - self.open_mode, - self.pipe_mode, + open_mode, + pipe_mode, self.max_instances, self.out_buffer_size, self.in_buffer_size, @@ -1884,11 +2357,11 @@ impl ServerOptions { attrs as *mut _, ); - if h == handleapi::INVALID_HANDLE_VALUE { + if h == windows_sys::INVALID_HANDLE_VALUE { return Err(io::Error::last_os_error()); } - NamedPipeServer::from_raw_handle(h) + NamedPipeServer::from_raw_handle(h as _) } } @@ -1898,8 +2371,10 @@ impl ServerOptions { /// See [`ClientOptions::open`]. #[derive(Debug, Clone)] pub struct ClientOptions { - desired_access: DWORD, - security_qos_flags: DWORD, + generic_read: bool, + generic_write: bool, + security_qos_flags: u32, + pipe_mode: PipeMode, } impl ClientOptions { @@ -1918,8 +2393,11 @@ impl ClientOptions { /// ``` pub fn new() -> Self { Self { - desired_access: winnt::GENERIC_READ | winnt::GENERIC_WRITE, - security_qos_flags: winbase::SECURITY_IDENTIFICATION | winbase::SECURITY_SQOS_PRESENT, + generic_read: true, + generic_write: true, + security_qos_flags: windows_sys::SECURITY_IDENTIFICATION + | windows_sys::SECURITY_SQOS_PRESENT, + pipe_mode: PipeMode::Byte, } } @@ -1930,7 +2408,7 @@ impl ClientOptions { /// [`GENERIC_READ`]: https://docs.microsoft.com/en-us/windows/win32/secauthz/generic-access-rights /// [`CreateFile`]: https://docs.microsoft.com/en-us/windows/win32/api/fileapi/nf-fileapi-createfilew pub fn read(&mut self, allowed: bool) -> &mut Self { - bool_flag!(self.desired_access, allowed, winnt::GENERIC_READ); + self.generic_read = allowed; self } @@ -1941,7 +2419,7 @@ impl ClientOptions { /// [`GENERIC_WRITE`]: https://docs.microsoft.com/en-us/windows/win32/secauthz/generic-access-rights /// [`CreateFile`]: https://docs.microsoft.com/en-us/windows/win32/api/fileapi/nf-fileapi-createfilew pub fn write(&mut self, allowed: bool) -> &mut Self { - bool_flag!(self.desired_access, allowed, winnt::GENERIC_WRITE); + self.generic_write = allowed; self } @@ -1964,15 +2442,24 @@ impl ClientOptions { /// automatically when using this method. /// /// [`CreateFile`]: https://docs.microsoft.com/en-us/windows/win32/api/fileapi/nf-fileapi-createfilea - /// [`SECURITY_IDENTIFICATION`]: crate::winapi::um::winbase::SECURITY_IDENTIFICATION + /// [`SECURITY_IDENTIFICATION`]: https://docs.rs/windows-sys/latest/windows_sys/Win32/Storage/FileSystem/constant.SECURITY_IDENTIFICATION.html /// [Impersonation Levels]: https://docs.microsoft.com/en-us/windows/win32/api/winnt/ne-winnt-security_impersonation_level pub fn security_qos_flags(&mut self, flags: u32) -> &mut Self { // See: https://github.com/rust-lang/rust/pull/58216 - self.security_qos_flags = flags | winbase::SECURITY_SQOS_PRESENT; + self.security_qos_flags = flags | windows_sys::SECURITY_SQOS_PRESENT; + self + } + + /// The pipe mode. + /// + /// The default pipe mode is [`PipeMode::Byte`]. See [`PipeMode`] for + /// documentation of what each mode means. + pub fn pipe_mode(&mut self, pipe_mode: PipeMode) -> &mut Self { + self.pipe_mode = pipe_mode; self } - /// Open the named pipe identified by `addr`. + /// Opens the named pipe identified by `addr`. /// /// This opens the client using [`CreateFile`] with the /// `dwCreationDisposition` option set to `OPEN_EXISTING`. @@ -1993,8 +2480,7 @@ impl ClientOptions { /// but the server is not currently waiting for a connection. Please see the /// examples for how to check for this error. /// - /// [`ERROR_PIPE_BUSY`]: crate::winapi::shared::winerror::ERROR_PIPE_BUSY - /// [`winapi`]: crate::winapi + /// [`ERROR_PIPE_BUSY`]: https://docs.rs/windows-sys/latest/windows_sys/Win32/Foundation/constant.ERROR_PIPE_BUSY.html /// [enabled I/O]: crate::runtime::Builder::enable_io /// [Tokio Runtime]: crate::runtime::Runtime /// @@ -2005,7 +2491,7 @@ impl ClientOptions { /// use std::time::Duration; /// use tokio::net::windows::named_pipe::ClientOptions; /// use tokio::time; - /// use winapi::shared::winerror; + /// use windows_sys::Win32::Foundation::ERROR_PIPE_BUSY; /// /// const PIPE_NAME: &str = r"\\.\pipe\mynamedpipe"; /// @@ -2013,7 +2499,7 @@ impl ClientOptions { /// let client = loop { /// match ClientOptions::new().open(PIPE_NAME) { /// Ok(client) => break client, - /// Err(e) if e.raw_os_error() == Some(winerror::ERROR_PIPE_BUSY as i32) => (), + /// Err(e) if e.raw_os_error() == Some(ERROR_PIPE_BUSY as i32) => (), /// Err(e) => return Err(e), /// } /// @@ -2029,7 +2515,7 @@ impl ClientOptions { unsafe { self.open_with_security_attributes_raw(addr, ptr::null_mut()) } } - /// Open the named pipe identified by `addr`. + /// Opens the named pipe identified by `addr`. /// /// This is the same as [`open`] except that it supports providing the raw /// pointer to a structure of [`SECURITY_ATTRIBUTES`] which will be passed @@ -2043,7 +2529,7 @@ impl ClientOptions { /// /// [`open`]: ClientOptions::open /// [`CreateFile`]: https://docs.microsoft.com/en-us/windows/win32/api/fileapi/nf-fileapi-createfilew - /// [`SECURITY_ATTRIBUTES`]: crate::winapi::um::minwinbase::SECURITY_ATTRIBUTES + /// [`SECURITY_ATTRIBUTES`]: https://docs.rs/windows-sys/latest/windows_sys/Win32/Security/struct.SECURITY_ATTRIBUTES.html pub unsafe fn open_with_security_attributes_raw( &self, addr: impl AsRef<OsStr>, @@ -2051,29 +2537,50 @@ impl ClientOptions { ) -> io::Result<NamedPipeClient> { let addr = encode_addr(addr); + let desired_access = { + let mut access = 0; + if self.generic_read { + access |= windows_sys::GENERIC_READ; + } + if self.generic_write { + access |= windows_sys::GENERIC_WRITE; + } + access + }; + // NB: We could use a platform specialized `OpenOptions` here, but since - // we have access to winapi it ultimately doesn't hurt to use + // we have access to windows_sys it ultimately doesn't hurt to use // `CreateFile` explicitly since it allows the use of our already // well-structured wide `addr` to pass into CreateFileW. - let h = fileapi::CreateFileW( + let h = windows_sys::CreateFileW( addr.as_ptr(), - self.desired_access, + desired_access, 0, attrs as *mut _, - fileapi::OPEN_EXISTING, + windows_sys::OPEN_EXISTING, self.get_flags(), - ptr::null_mut(), + 0, ); - if h == handleapi::INVALID_HANDLE_VALUE { + if h == windows_sys::INVALID_HANDLE_VALUE { return Err(io::Error::last_os_error()); } - NamedPipeClient::from_raw_handle(h) + if matches!(self.pipe_mode, PipeMode::Message) { + let mode = windows_sys::PIPE_READMODE_MESSAGE; + let result = + windows_sys::SetNamedPipeHandleState(h, &mode, ptr::null_mut(), ptr::null_mut()); + + if result == 0 { + return Err(io::Error::last_os_error()); + } + } + + NamedPipeClient::from_raw_handle(h as _) } fn get_flags(&self) -> u32 { - self.security_qos_flags | winbase::FILE_FLAG_OVERLAPPED + self.security_qos_flags | windows_sys::FILE_FLAG_OVERLAPPED } } @@ -2086,16 +2593,19 @@ pub enum PipeMode { /// Data is written to the pipe as a stream of bytes. The pipe does not /// distinguish bytes written during different write operations. /// - /// Corresponds to [`PIPE_TYPE_BYTE`][crate::winapi::um::winbase::PIPE_TYPE_BYTE]. + /// Corresponds to [`PIPE_TYPE_BYTE`]. + /// + /// [`PIPE_TYPE_BYTE`]: https://docs.rs/windows-sys/latest/windows_sys/Win32/System/Pipes/constant.PIPE_TYPE_BYTE.html Byte, /// Data is written to the pipe as a stream of messages. The pipe treats the /// bytes written during each write operation as a message unit. Any reading /// on a named pipe returns [`ERROR_MORE_DATA`] when a message is not read /// completely. /// - /// Corresponds to [`PIPE_TYPE_MESSAGE`][crate::winapi::um::winbase::PIPE_TYPE_MESSAGE]. + /// Corresponds to [`PIPE_TYPE_MESSAGE`]. /// - /// [`ERROR_MORE_DATA`]: crate::winapi::shared::winerror::ERROR_MORE_DATA + /// [`ERROR_MORE_DATA`]: https://docs.rs/windows-sys/latest/windows_sys/Win32/Foundation/constant.ERROR_MORE_DATA.html + /// [`PIPE_TYPE_MESSAGE`]: https://docs.rs/windows-sys/latest/windows_sys/Win32/System/Pipes/constant.PIPE_TYPE_MESSAGE.html Message, } @@ -2105,11 +2615,15 @@ pub enum PipeMode { pub enum PipeEnd { /// The named pipe refers to the client end of a named pipe instance. /// - /// Corresponds to [`PIPE_CLIENT_END`][crate::winapi::um::winbase::PIPE_CLIENT_END]. + /// Corresponds to [`PIPE_CLIENT_END`]. + /// + /// [`PIPE_CLIENT_END`]: https://docs.rs/windows-sys/latest/windows_sys/Win32/System/Pipes/constant.PIPE_CLIENT_END.html Client, /// The named pipe refers to the server end of a named pipe instance. /// - /// Corresponds to [`PIPE_SERVER_END`][crate::winapi::um::winbase::PIPE_SERVER_END]. + /// Corresponds to [`PIPE_SERVER_END`]. + /// + /// [`PIPE_SERVER_END`]: https://docs.rs/windows-sys/latest/windows_sys/Win32/System/Pipes/constant.PIPE_SERVER_END.html Server, } @@ -2131,7 +2645,7 @@ pub struct PipeInfo { pub in_buffer_size: u32, } -/// Encode an address so that it is a null-terminated wide string. +/// Encodes an address so that it is a null-terminated wide string. fn encode_addr(addr: impl AsRef<OsStr>) -> Box<[u16]> { let len = addr.as_ref().encode_wide().count(); let mut vec = Vec::with_capacity(len + 1); @@ -2147,26 +2661,26 @@ unsafe fn named_pipe_info(handle: RawHandle) -> io::Result<PipeInfo> { let mut in_buffer_size = 0; let mut max_instances = 0; - let result = namedpipeapi::GetNamedPipeInfo( - handle, + let result = windows_sys::GetNamedPipeInfo( + handle as _, &mut flags, &mut out_buffer_size, &mut in_buffer_size, &mut max_instances, ); - if result == FALSE { + if result == 0 { return Err(io::Error::last_os_error()); } let mut end = PipeEnd::Client; let mut mode = PipeMode::Byte; - if flags & winbase::PIPE_SERVER_END != 0 { + if flags & windows_sys::PIPE_SERVER_END != 0 { end = PipeEnd::Server; } - if flags & winbase::PIPE_TYPE_MESSAGE != 0 { + if flags & windows_sys::PIPE_TYPE_MESSAGE != 0 { mode = PipeMode::Message; } diff --git a/vendor/tokio/src/park/either.rs b/vendor/tokio/src/park/either.rs deleted file mode 100644 index ee02ec158..000000000 --- a/vendor/tokio/src/park/either.rs +++ /dev/null @@ -1,74 +0,0 @@ -#![cfg_attr(not(feature = "full"), allow(dead_code))] - -use crate::park::{Park, Unpark}; - -use std::fmt; -use std::time::Duration; - -pub(crate) enum Either<A, B> { - A(A), - B(B), -} - -impl<A, B> Park for Either<A, B> -where - A: Park, - B: Park, -{ - type Unpark = Either<A::Unpark, B::Unpark>; - type Error = Either<A::Error, B::Error>; - - fn unpark(&self) -> Self::Unpark { - match self { - Either::A(a) => Either::A(a.unpark()), - Either::B(b) => Either::B(b.unpark()), - } - } - - fn park(&mut self) -> Result<(), Self::Error> { - match self { - Either::A(a) => a.park().map_err(Either::A), - Either::B(b) => b.park().map_err(Either::B), - } - } - - fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> { - match self { - Either::A(a) => a.park_timeout(duration).map_err(Either::A), - Either::B(b) => b.park_timeout(duration).map_err(Either::B), - } - } - - fn shutdown(&mut self) { - match self { - Either::A(a) => a.shutdown(), - Either::B(b) => b.shutdown(), - } - } -} - -impl<A, B> Unpark for Either<A, B> -where - A: Unpark, - B: Unpark, -{ - fn unpark(&self) { - match self { - Either::A(a) => a.unpark(), - Either::B(b) => b.unpark(), - } - } -} - -impl<A, B> fmt::Debug for Either<A, B> -where - A: fmt::Debug, - B: fmt::Debug, -{ - fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Either::A(a) => a.fmt(fmt), - Either::B(b) => b.fmt(fmt), - } - } -} diff --git a/vendor/tokio/src/park/mod.rs b/vendor/tokio/src/park/mod.rs deleted file mode 100644 index edd937100..000000000 --- a/vendor/tokio/src/park/mod.rs +++ /dev/null @@ -1,117 +0,0 @@ -//! Abstraction over blocking and unblocking the current thread. -//! -//! Provides an abstraction over blocking the current thread. This is similar to -//! the park / unpark constructs provided by `std` but made generic. This allows -//! embedding custom functionality to perform when the thread is blocked. -//! -//! A blocked `Park` instance is unblocked by calling `unpark` on its -//! `Unpark` handle. -//! -//! The `ParkThread` struct implements `Park` using `thread::park` to put the -//! thread to sleep. The Tokio reactor also implements park, but uses -//! `mio::Poll` to block the thread instead. -//! -//! The `Park` trait is composable. A timer implementation might decorate a -//! `Park` implementation by checking if any timeouts have elapsed after the -//! inner `Park` implementation unblocks. -//! -//! # Model -//! -//! Conceptually, each `Park` instance has an associated token, which is -//! initially not present: -//! -//! * The `park` method blocks the current thread unless or until the token is -//! available, at which point it atomically consumes the token. -//! * The `unpark` method atomically makes the token available if it wasn't -//! already. -//! -//! Some things to note: -//! -//! * If `unpark` is called before `park`, the next call to `park` will -//! **not** block the thread. -//! * **Spurious** wakeups are permitted, i.e., the `park` method may unblock -//! even if `unpark` was not called. -//! * `park_timeout` does the same as `park` but allows specifying a maximum -//! time to block the thread for. - -cfg_rt! { - pub(crate) mod either; -} - -#[cfg(any(feature = "rt", feature = "sync"))] -pub(crate) mod thread; - -use std::fmt::Debug; -use std::sync::Arc; -use std::time::Duration; - -/// Block the current thread. -pub(crate) trait Park { - /// Unpark handle type for the `Park` implementation. - type Unpark: Unpark; - - /// Error returned by `park` - type Error: Debug; - - /// Gets a new `Unpark` handle associated with this `Park` instance. - fn unpark(&self) -> Self::Unpark; - - /// Blocks the current thread unless or until the token is available. - /// - /// A call to `park` does not guarantee that the thread will remain blocked - /// forever, and callers should be prepared for this possibility. This - /// function may wakeup spuriously for any reason. - /// - /// # Panics - /// - /// This function **should** not panic, but ultimately, panics are left as - /// an implementation detail. Refer to the documentation for the specific - /// `Park` implementation - fn park(&mut self) -> Result<(), Self::Error>; - - /// Parks the current thread for at most `duration`. - /// - /// This function is the same as `park` but allows specifying a maximum time - /// to block the thread for. - /// - /// Same as `park`, there is no guarantee that the thread will remain - /// blocked for any amount of time. Spurious wakeups are permitted for any - /// reason. - /// - /// # Panics - /// - /// This function **should** not panic, but ultimately, panics are left as - /// an implementation detail. Refer to the documentation for the specific - /// `Park` implementation - fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error>; - - /// Release all resources holded by the parker for proper leak-free shutdown - fn shutdown(&mut self); -} - -/// Unblock a thread blocked by the associated `Park` instance. -pub(crate) trait Unpark: Sync + Send + 'static { - /// Unblocks a thread that is blocked by the associated `Park` handle. - /// - /// Calling `unpark` atomically makes available the unpark token, if it is - /// not already available. - /// - /// # Panics - /// - /// This function **should** not panic, but ultimately, panics are left as - /// an implementation detail. Refer to the documentation for the specific - /// `Unpark` implementation - fn unpark(&self); -} - -impl Unpark for Box<dyn Unpark> { - fn unpark(&self) { - (**self).unpark() - } -} - -impl Unpark for Arc<dyn Unpark> { - fn unpark(&self) { - (**self).unpark() - } -} diff --git a/vendor/tokio/src/park/thread.rs b/vendor/tokio/src/park/thread.rs deleted file mode 100644 index 2725e4563..000000000 --- a/vendor/tokio/src/park/thread.rs +++ /dev/null @@ -1,346 +0,0 @@ -#![cfg_attr(not(feature = "full"), allow(dead_code))] - -use crate::loom::sync::atomic::AtomicUsize; -use crate::loom::sync::{Arc, Condvar, Mutex}; -use crate::park::{Park, Unpark}; - -use std::sync::atomic::Ordering::SeqCst; -use std::time::Duration; - -#[derive(Debug)] -pub(crate) struct ParkThread { - inner: Arc<Inner>, -} - -pub(crate) type ParkError = (); - -/// Unblocks a thread that was blocked by `ParkThread`. -#[derive(Clone, Debug)] -pub(crate) struct UnparkThread { - inner: Arc<Inner>, -} - -#[derive(Debug)] -struct Inner { - state: AtomicUsize, - mutex: Mutex<()>, - condvar: Condvar, -} - -const EMPTY: usize = 0; -const PARKED: usize = 1; -const NOTIFIED: usize = 2; - -thread_local! { - static CURRENT_PARKER: ParkThread = ParkThread::new(); -} - -// ==== impl ParkThread ==== - -impl ParkThread { - pub(crate) fn new() -> Self { - Self { - inner: Arc::new(Inner { - state: AtomicUsize::new(EMPTY), - mutex: Mutex::new(()), - condvar: Condvar::new(), - }), - } - } -} - -impl Park for ParkThread { - type Unpark = UnparkThread; - type Error = ParkError; - - fn unpark(&self) -> Self::Unpark { - let inner = self.inner.clone(); - UnparkThread { inner } - } - - fn park(&mut self) -> Result<(), Self::Error> { - self.inner.park(); - Ok(()) - } - - fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> { - self.inner.park_timeout(duration); - Ok(()) - } - - fn shutdown(&mut self) { - self.inner.shutdown(); - } -} - -// ==== impl Inner ==== - -impl Inner { - /// Park the current thread for at most `dur`. - fn park(&self) { - // If we were previously notified then we consume this notification and - // return quickly. - if self - .state - .compare_exchange(NOTIFIED, EMPTY, SeqCst, SeqCst) - .is_ok() - { - return; - } - - // Otherwise we need to coordinate going to sleep - let mut m = self.mutex.lock(); - - match self.state.compare_exchange(EMPTY, PARKED, SeqCst, SeqCst) { - Ok(_) => {} - Err(NOTIFIED) => { - // We must read here, even though we know it will be `NOTIFIED`. - // This is because `unpark` may have been called again since we read - // `NOTIFIED` in the `compare_exchange` above. We must perform an - // acquire operation that synchronizes with that `unpark` to observe - // any writes it made before the call to unpark. To do that we must - // read from the write it made to `state`. - let old = self.state.swap(EMPTY, SeqCst); - debug_assert_eq!(old, NOTIFIED, "park state changed unexpectedly"); - - return; - } - Err(actual) => panic!("inconsistent park state; actual = {}", actual), - } - - loop { - m = self.condvar.wait(m).unwrap(); - - if self - .state - .compare_exchange(NOTIFIED, EMPTY, SeqCst, SeqCst) - .is_ok() - { - // got a notification - return; - } - - // spurious wakeup, go back to sleep - } - } - - fn park_timeout(&self, dur: Duration) { - // Like `park` above we have a fast path for an already-notified thread, - // and afterwards we start coordinating for a sleep. Return quickly. - if self - .state - .compare_exchange(NOTIFIED, EMPTY, SeqCst, SeqCst) - .is_ok() - { - return; - } - - if dur == Duration::from_millis(0) { - return; - } - - let m = self.mutex.lock(); - - match self.state.compare_exchange(EMPTY, PARKED, SeqCst, SeqCst) { - Ok(_) => {} - Err(NOTIFIED) => { - // We must read again here, see `park`. - let old = self.state.swap(EMPTY, SeqCst); - debug_assert_eq!(old, NOTIFIED, "park state changed unexpectedly"); - - return; - } - Err(actual) => panic!("inconsistent park_timeout state; actual = {}", actual), - } - - // Wait with a timeout, and if we spuriously wake up or otherwise wake up - // from a notification, we just want to unconditionally set the state back to - // empty, either consuming a notification or un-flagging ourselves as - // parked. - let (_m, _result) = self.condvar.wait_timeout(m, dur).unwrap(); - - match self.state.swap(EMPTY, SeqCst) { - NOTIFIED => {} // got a notification, hurray! - PARKED => {} // no notification, alas - n => panic!("inconsistent park_timeout state: {}", n), - } - } - - fn unpark(&self) { - // To ensure the unparked thread will observe any writes we made before - // this call, we must perform a release operation that `park` can - // synchronize with. To do that we must write `NOTIFIED` even if `state` - // is already `NOTIFIED`. That is why this must be a swap rather than a - // compare-and-swap that returns if it reads `NOTIFIED` on failure. - match self.state.swap(NOTIFIED, SeqCst) { - EMPTY => return, // no one was waiting - NOTIFIED => return, // already unparked - PARKED => {} // gotta go wake someone up - _ => panic!("inconsistent state in unpark"), - } - - // There is a period between when the parked thread sets `state` to - // `PARKED` (or last checked `state` in the case of a spurious wake - // up) and when it actually waits on `cvar`. If we were to notify - // during this period it would be ignored and then when the parked - // thread went to sleep it would never wake up. Fortunately, it has - // `lock` locked at this stage so we can acquire `lock` to wait until - // it is ready to receive the notification. - // - // Releasing `lock` before the call to `notify_one` means that when the - // parked thread wakes it doesn't get woken only to have to wait for us - // to release `lock`. - drop(self.mutex.lock()); - - self.condvar.notify_one() - } - - fn shutdown(&self) { - self.condvar.notify_all(); - } -} - -impl Default for ParkThread { - fn default() -> Self { - Self::new() - } -} - -// ===== impl UnparkThread ===== - -impl Unpark for UnparkThread { - fn unpark(&self) { - self.inner.unpark(); - } -} - -use std::future::Future; -use std::marker::PhantomData; -use std::mem; -use std::rc::Rc; -use std::task::{RawWaker, RawWakerVTable, Waker}; - -/// Blocks the current thread using a condition variable. -#[derive(Debug)] -pub(crate) struct CachedParkThread { - _anchor: PhantomData<Rc<()>>, -} - -impl CachedParkThread { - /// Create a new `ParkThread` handle for the current thread. - /// - /// This type cannot be moved to other threads, so it should be created on - /// the thread that the caller intends to park. - pub(crate) fn new() -> CachedParkThread { - CachedParkThread { - _anchor: PhantomData, - } - } - - pub(crate) fn get_unpark(&self) -> Result<UnparkThread, ParkError> { - self.with_current(|park_thread| park_thread.unpark()) - } - - /// Get a reference to the `ParkThread` handle for this thread. - fn with_current<F, R>(&self, f: F) -> Result<R, ParkError> - where - F: FnOnce(&ParkThread) -> R, - { - CURRENT_PARKER.try_with(|inner| f(inner)).map_err(|_| ()) - } - - pub(crate) fn block_on<F: Future>(&mut self, f: F) -> Result<F::Output, ParkError> { - use std::task::Context; - use std::task::Poll::Ready; - - // `get_unpark()` should not return a Result - let waker = self.get_unpark()?.into_waker(); - let mut cx = Context::from_waker(&waker); - - pin!(f); - - loop { - if let Ready(v) = crate::coop::budget(|| f.as_mut().poll(&mut cx)) { - return Ok(v); - } - - self.park()?; - } - } -} - -impl Park for CachedParkThread { - type Unpark = UnparkThread; - type Error = ParkError; - - fn unpark(&self) -> Self::Unpark { - self.get_unpark().unwrap() - } - - fn park(&mut self) -> Result<(), Self::Error> { - self.with_current(|park_thread| park_thread.inner.park())?; - Ok(()) - } - - fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> { - self.with_current(|park_thread| park_thread.inner.park_timeout(duration))?; - Ok(()) - } - - fn shutdown(&mut self) { - let _ = self.with_current(|park_thread| park_thread.inner.shutdown()); - } -} - -impl UnparkThread { - pub(crate) fn into_waker(self) -> Waker { - unsafe { - let raw = unparker_to_raw_waker(self.inner); - Waker::from_raw(raw) - } - } -} - -impl Inner { - #[allow(clippy::wrong_self_convention)] - fn into_raw(this: Arc<Inner>) -> *const () { - Arc::into_raw(this) as *const () - } - - unsafe fn from_raw(ptr: *const ()) -> Arc<Inner> { - Arc::from_raw(ptr as *const Inner) - } -} - -unsafe fn unparker_to_raw_waker(unparker: Arc<Inner>) -> RawWaker { - RawWaker::new( - Inner::into_raw(unparker), - &RawWakerVTable::new(clone, wake, wake_by_ref, drop_waker), - ) -} - -unsafe fn clone(raw: *const ()) -> RawWaker { - let unparker = Inner::from_raw(raw); - - // Increment the ref count - mem::forget(unparker.clone()); - - unparker_to_raw_waker(unparker) -} - -unsafe fn drop_waker(raw: *const ()) { - let _ = Inner::from_raw(raw); -} - -unsafe fn wake(raw: *const ()) { - let unparker = Inner::from_raw(raw); - unparker.unpark(); -} - -unsafe fn wake_by_ref(raw: *const ()) { - let unparker = Inner::from_raw(raw); - unparker.unpark(); - - // We don't actually own a reference to the unparker - mem::forget(unparker); -} diff --git a/vendor/tokio/src/process/mod.rs b/vendor/tokio/src/process/mod.rs index 96ceb6d8d..d68a68995 100644 --- a/vendor/tokio/src/process/mod.rs +++ b/vendor/tokio/src/process/mod.rs @@ -97,13 +97,65 @@ //! } //! ``` //! +//! Here is another example using `sort` writing into the child process +//! standard input, capturing the output of the sorted text. +//! +//! ```no_run +//! use tokio::io::AsyncWriteExt; +//! use tokio::process::Command; +//! +//! use std::process::Stdio; +//! +//! #[tokio::main] +//! async fn main() -> Result<(), Box<dyn std::error::Error>> { +//! let mut cmd = Command::new("sort"); +//! +//! // Specifying that we want pipe both the output and the input. +//! // Similarly to capturing the output, by configuring the pipe +//! // to stdin it can now be used as an asynchronous writer. +//! cmd.stdout(Stdio::piped()); +//! cmd.stdin(Stdio::piped()); +//! +//! let mut child = cmd.spawn().expect("failed to spawn command"); +//! +//! // These are the animals we want to sort +//! let animals: &[&str] = &["dog", "bird", "frog", "cat", "fish"]; +//! +//! let mut stdin = child +//! .stdin +//! .take() +//! .expect("child did not have a handle to stdin"); +//! +//! // Write our animals to the child process +//! // Note that the behavior of `sort` is to buffer _all input_ before writing any output. +//! // In the general sense, it is recommended to write to the child in a separate task as +//! // awaiting its exit (or output) to avoid deadlocks (for example, the child tries to write +//! // some output but gets stuck waiting on the parent to read from it, meanwhile the parent +//! // is stuck waiting to write its input completely before reading the output). +//! stdin +//! .write(animals.join("\n").as_bytes()) +//! .await +//! .expect("could not write to stdin"); +//! +//! // We drop the handle here which signals EOF to the child process. +//! // This tells the child process that it there is no more data on the pipe. +//! drop(stdin); +//! +//! let op = child.wait_with_output().await?; +//! +//! // Results should come back in sorted order +//! assert_eq!(op.stdout, "bird\ncat\ndog\nfish\nfrog\n".as_bytes()); +//! +//! Ok(()) +//! } +//! ``` +//! //! With some coordination, we can also pipe the output of one command into //! another. //! //! ```no_run //! use tokio::join; //! use tokio::process::Command; -//! use std::convert::TryInto; //! use std::process::Stdio; //! //! #[tokio::main] @@ -164,7 +216,7 @@ //! from being spawned. //! //! The tokio runtime will, on a best-effort basis, attempt to reap and clean up -//! any process which it has spawned. No additional guarantees are made with regards +//! any process which it has spawned. No additional guarantees are made with regard to //! how quickly or how often this procedure will take place. //! //! It is recommended to avoid dropping a [`Child`] process handle before it has been @@ -192,20 +244,26 @@ mod kill; use crate::io::{AsyncRead, AsyncWrite, ReadBuf}; use crate::process::kill::Kill; -use std::convert::TryInto; use std::ffi::OsStr; use std::future::Future; use std::io; -#[cfg(unix)] -use std::os::unix::process::CommandExt; -#[cfg(windows)] -use std::os::windows::process::CommandExt; use std::path::Path; use std::pin::Pin; use std::process::{Command as StdCommand, ExitStatus, Output, Stdio}; use std::task::Context; use std::task::Poll; +#[cfg(unix)] +use std::os::unix::process::CommandExt; +#[cfg(windows)] +use std::os::windows::process::CommandExt; + +cfg_windows! { + use crate::os::windows::io::{AsRawHandle, RawHandle}; + #[cfg(not(tokio_no_as_fd))] + use crate::os::windows::io::{AsHandle, BorrowedHandle}; +} + /// This structure mimics the API of [`std::process::Command`] found in the standard library, but /// replaces functions that create a process with an asynchronous variant. The main provided /// asynchronous functions are [spawn](Command::spawn), [status](Command::status), and @@ -223,9 +281,9 @@ pub struct Command { pub(crate) struct SpawnedChild { child: imp::Child, - stdin: Option<imp::ChildStdin>, - stdout: Option<imp::ChildStdout>, - stderr: Option<imp::ChildStderr>, + stdin: Option<imp::ChildStdio>, + stdout: Option<imp::ChildStdio>, + stderr: Option<imp::ChildStdio>, } impl Command { @@ -254,7 +312,8 @@ impl Command { /// /// ```no_run /// use tokio::process::Command; - /// let command = Command::new("sh"); + /// let mut command = Command::new("sh"); + /// # let _ = command.output(); // assert borrow checker /// ``` /// /// [rust-lang/rust#37519]: https://github.com/rust-lang/rust/issues/37519 @@ -262,21 +321,31 @@ impl Command { Self::from(StdCommand::new(program)) } + /// Cheaply convert to a `&std::process::Command` for places where the type from the standard + /// library is expected. + pub fn as_std(&self) -> &StdCommand { + &self.std + } + /// Adds an argument to pass to the program. /// /// Only one argument can be passed per use. So instead of: /// /// ```no_run - /// tokio::process::Command::new("sh") - /// .arg("-C /path/to/repo"); + /// let mut command = tokio::process::Command::new("sh"); + /// command.arg("-C /path/to/repo"); + /// + /// # let _ = command.output(); // assert borrow checker /// ``` /// /// usage would be: /// /// ```no_run - /// tokio::process::Command::new("sh") - /// .arg("-C") - /// .arg("/path/to/repo"); + /// let mut command = tokio::process::Command::new("sh"); + /// command.arg("-C"); + /// command.arg("/path/to/repo"); + /// + /// # let _ = command.output(); // assert borrow checker /// ``` /// /// To pass multiple arguments see [`args`]. @@ -288,11 +357,15 @@ impl Command { /// Basic usage: /// /// ```no_run + /// # async fn test() { // allow using await /// use tokio::process::Command; /// - /// let command = Command::new("ls") + /// let output = Command::new("ls") /// .arg("-l") - /// .arg("-a"); + /// .arg("-a") + /// .output().await.unwrap(); + /// # } + /// /// ``` pub fn arg<S: AsRef<OsStr>>(&mut self, arg: S) -> &mut Command { self.std.arg(arg); @@ -310,10 +383,13 @@ impl Command { /// Basic usage: /// /// ```no_run + /// # async fn test() { // allow using await /// use tokio::process::Command; /// - /// let command = Command::new("ls") - /// .args(&["-l", "-a"]); + /// let output = Command::new("ls") + /// .args(&["-l", "-a"]) + /// .output().await.unwrap(); + /// # } /// ``` pub fn args<I, S>(&mut self, args: I) -> &mut Command where @@ -324,6 +400,22 @@ impl Command { self } + /// Append literal text to the command line without any quoting or escaping. + /// + /// This is useful for passing arguments to `cmd.exe /c`, which doesn't follow + /// `CommandLineToArgvW` escaping rules. + /// + /// **Note**: This is an [unstable API][unstable] but will be stabilised once + /// tokio's MSRV is sufficiently new. See [the documentation on + /// unstable features][unstable] for details about using unstable features. + #[cfg(windows)] + #[cfg(tokio_unstable)] + #[cfg_attr(docsrs, doc(cfg(all(windows, tokio_unstable))))] + pub fn raw_arg<S: AsRef<OsStr>>(&mut self, text_to_append_as_is: S) -> &mut Command { + self.std.raw_arg(text_to_append_as_is); + self + } + /// Inserts or updates an environment variable mapping. /// /// Note that environment variable names are case-insensitive (but case-preserving) on Windows, @@ -334,10 +426,13 @@ impl Command { /// Basic usage: /// /// ```no_run + /// # async fn test() { // allow using await /// use tokio::process::Command; /// - /// let command = Command::new("ls") - /// .env("PATH", "/bin"); + /// let output = Command::new("ls") + /// .env("PATH", "/bin") + /// .output().await.unwrap(); + /// # } /// ``` pub fn env<K, V>(&mut self, key: K, val: V) -> &mut Command where @@ -355,6 +450,7 @@ impl Command { /// Basic usage: /// /// ```no_run + /// # async fn test() { // allow using await /// use tokio::process::Command; /// use std::process::{Stdio}; /// use std::env; @@ -365,11 +461,13 @@ impl Command { /// k == "TERM" || k == "TZ" || k == "LANG" || k == "PATH" /// ).collect(); /// - /// let command = Command::new("printenv") + /// let output = Command::new("printenv") /// .stdin(Stdio::null()) /// .stdout(Stdio::inherit()) /// .env_clear() - /// .envs(&filtered_env); + /// .envs(&filtered_env) + /// .output().await.unwrap(); + /// # } /// ``` pub fn envs<I, K, V>(&mut self, vars: I) -> &mut Command where @@ -388,10 +486,13 @@ impl Command { /// Basic usage: /// /// ```no_run + /// # async fn test() { // allow using await /// use tokio::process::Command; /// - /// let command = Command::new("ls") - /// .env_remove("PATH"); + /// let output = Command::new("ls") + /// .env_remove("PATH") + /// .output().await.unwrap(); + /// # } /// ``` pub fn env_remove<K: AsRef<OsStr>>(&mut self, key: K) -> &mut Command { self.std.env_remove(key); @@ -405,10 +506,13 @@ impl Command { /// Basic usage: /// /// ```no_run + /// # async fn test() { // allow using await /// use tokio::process::Command; /// - /// let command = Command::new("ls") - /// .env_clear(); + /// let output = Command::new("ls") + /// .env_clear() + /// .output().await.unwrap(); + /// # } /// ``` pub fn env_clear(&mut self) -> &mut Command { self.std.env_clear(); @@ -432,10 +536,13 @@ impl Command { /// Basic usage: /// /// ```no_run + /// # async fn test() { // allow using await /// use tokio::process::Command; /// - /// let command = Command::new("ls") - /// .current_dir("/bin"); + /// let output = Command::new("ls") + /// .current_dir("/bin") + /// .output().await.unwrap(); + /// # } /// ``` pub fn current_dir<P: AsRef<Path>>(&mut self, dir: P) -> &mut Command { self.std.current_dir(dir); @@ -455,11 +562,14 @@ impl Command { /// Basic usage: /// /// ```no_run + /// # async fn test() { // allow using await /// use std::process::{Stdio}; /// use tokio::process::Command; /// - /// let command = Command::new("ls") - /// .stdin(Stdio::null()); + /// let output = Command::new("ls") + /// .stdin(Stdio::null()) + /// .output().await.unwrap(); + /// # } /// ``` pub fn stdin<T: Into<Stdio>>(&mut self, cfg: T) -> &mut Command { self.std.stdin(cfg); @@ -479,11 +589,14 @@ impl Command { /// Basic usage: /// /// ```no_run + /// # async fn test() { // allow using await /// use tokio::process::Command; /// use std::process::Stdio; /// - /// let command = Command::new("ls") - /// .stdout(Stdio::null()); + /// let output = Command::new("ls") + /// .stdout(Stdio::null()) + /// .output().await.unwrap(); + /// # } /// ``` pub fn stdout<T: Into<Stdio>>(&mut self, cfg: T) -> &mut Command { self.std.stdout(cfg); @@ -503,11 +616,14 @@ impl Command { /// Basic usage: /// /// ```no_run + /// # async fn test() { // allow using await /// use tokio::process::Command; /// use std::process::{Stdio}; /// - /// let command = Command::new("ls") - /// .stderr(Stdio::null()); + /// let output = Command::new("ls") + /// .stderr(Stdio::null()) + /// .output().await.unwrap(); + /// # } /// ``` pub fn stderr<T: Into<Stdio>>(&mut self, cfg: T) -> &mut Command { self.std.stderr(cfg); @@ -534,7 +650,7 @@ impl Command { /// operation, the resulting zombie process cannot be `.await`ed inside of the /// destructor to avoid blocking other tasks. The tokio runtime will, on a /// best-effort basis, attempt to reap and clean up such processes in the - /// background, but makes no additional guarantees are made with regards + /// background, but no additional guarantees are made with regard to /// how quickly or how often this procedure will take place. /// /// If stronger guarantees are required, it is recommended to avoid dropping @@ -545,21 +661,23 @@ impl Command { self } - /// Sets the [process creation flags][1] to be passed to `CreateProcess`. - /// - /// These will always be ORed with `CREATE_UNICODE_ENVIRONMENT`. - /// - /// [1]: https://msdn.microsoft.com/en-us/library/windows/desktop/ms684863(v=vs.85).aspx - #[cfg(windows)] - pub fn creation_flags(&mut self, flags: u32) -> &mut Command { - self.std.creation_flags(flags); - self + cfg_windows! { + /// Sets the [process creation flags][1] to be passed to `CreateProcess`. + /// + /// These will always be ORed with `CREATE_UNICODE_ENVIRONMENT`. + /// + /// [1]: https://msdn.microsoft.com/en-us/library/windows/desktop/ms684863(v=vs.85).aspx + pub fn creation_flags(&mut self, flags: u32) -> &mut Command { + self.std.creation_flags(flags); + self + } } /// Sets the child process's user ID. This translates to a /// `setuid` call in the child process. Failure in the `setuid` /// call will cause the spawn to fail. #[cfg(unix)] + #[cfg_attr(docsrs, doc(cfg(unix)))] pub fn uid(&mut self, id: u32) -> &mut Command { self.std.uid(id); self @@ -568,11 +686,26 @@ impl Command { /// Similar to `uid` but sets the group ID of the child process. This has /// the same semantics as the `uid` field. #[cfg(unix)] + #[cfg_attr(docsrs, doc(cfg(unix)))] pub fn gid(&mut self, id: u32) -> &mut Command { self.std.gid(id); self } + /// Sets executable argument. + /// + /// Set the first process argument, `argv[0]`, to something other than the + /// default executable path. + #[cfg(unix)] + #[cfg_attr(docsrs, doc(cfg(unix)))] + pub fn arg0<S>(&mut self, arg: S) -> &mut Command + where + S: AsRef<OsStr>, + { + self.std.arg0(arg); + self + } + /// Schedules a closure to be run just before the `exec` function is /// invoked. /// @@ -603,6 +736,7 @@ impl Command { /// working directory have successfully been changed, so output to these /// locations may not appear where intended. #[cfg(unix)] + #[cfg_attr(docsrs, doc(cfg(unix)))] pub unsafe fn pre_exec<F>(&mut self, f: F) -> &mut Command where F: FnMut() -> io::Result<()> + Send + Sync + 'static, @@ -611,6 +745,39 @@ impl Command { self } + /// Sets the process group ID (PGID) of the child process. Equivalent to a + /// setpgid call in the child process, but may be more efficient. + /// + /// Process groups determine which processes receive signals. + /// + /// **Note**: This is an [unstable API][unstable] but will be stabilised once + /// tokio's MSRV is sufficiently new. See [the documentation on + /// unstable features][unstable] for details about using unstable features. + /// + /// If you want similar behaviour without using this unstable feature you can + /// create a [`std::process::Command`] and convert that into a + /// [`tokio::process::Command`] using the `From` trait. + /// + /// [unstable]: crate#unstable-features + /// [`tokio::process::Command`]: crate::process::Command + /// + /// ```no_run + /// # async fn test() { // allow using await + /// use tokio::process::Command; + /// + /// let output = Command::new("ls") + /// .process_group(0) + /// .output().await.unwrap(); + /// # } + /// ``` + #[cfg(unix)] + #[cfg(tokio_unstable)] + #[cfg_attr(docsrs, doc(cfg(all(unix, tokio_unstable))))] + pub fn process_group(&mut self, pgroup: i32) -> &mut Command { + self.std.process_group(pgroup); + self + } + /// Executes the command as a child process, returning a handle to it. /// /// By default, stdin, stdout and stderr are inherited from the parent. @@ -663,7 +830,7 @@ impl Command { /// from being spawned. /// /// The tokio runtime will, on a best-effort basis, attempt to reap and clean up - /// any process which it has spawned. No additional guarantees are made with regards + /// any process which it has spawned. No additional guarantees are made with regard to /// how quickly or how often this procedure will take place. /// /// It is recommended to avoid dropping a [`Child`] process handle before it has been @@ -844,8 +1011,9 @@ where type Output = Result<T, E>; fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + ready!(crate::trace::trace_leaf(cx)); // Keep track of task budget - let coop = ready!(crate::coop::poll_proceed(cx)); + let coop = ready!(crate::runtime::coop::poll_proceed(cx)); let ret = Pin::new(&mut self.inner).poll(cx); @@ -934,6 +1102,17 @@ impl Child { } } + cfg_windows! { + /// Extracts the raw handle of the process associated with this child while + /// it is still running. Returns `None` if the child has exited. + pub fn raw_handle(&self) -> Option<RawHandle> { + match &self.child { + FusedChild::Child(c) => Some(c.inner.as_raw_handle()), + FusedChild::Done(_) => None, + } + } + } + /// Attempts to force the child to exit, but does not wait for the request /// to take effect. /// @@ -1094,19 +1273,26 @@ impl Child { pub async fn wait_with_output(mut self) -> io::Result<Output> { use crate::future::try_join3; - async fn read_to_end<A: AsyncRead + Unpin>(io: Option<A>) -> io::Result<Vec<u8>> { + async fn read_to_end<A: AsyncRead + Unpin>(io: &mut Option<A>) -> io::Result<Vec<u8>> { let mut vec = Vec::new(); - if let Some(mut io) = io { - crate::io::util::read_to_end(&mut io, &mut vec).await?; + if let Some(io) = io.as_mut() { + crate::io::util::read_to_end(io, &mut vec).await?; } Ok(vec) } - let stdout_fut = read_to_end(self.stdout.take()); - let stderr_fut = read_to_end(self.stderr.take()); + let mut stdout_pipe = self.stdout.take(); + let mut stderr_pipe = self.stderr.take(); + + let stdout_fut = read_to_end(&mut stdout_pipe); + let stderr_fut = read_to_end(&mut stderr_pipe); let (status, stdout, stderr) = try_join3(self.wait(), stdout_fut, stderr_fut).await?; + // Drop happens after `try_join` due to <https://github.com/tokio-rs/tokio/issues/4309> + drop(stdout_pipe); + drop(stderr_pipe); + Ok(Output { status, stdout, @@ -1121,7 +1307,7 @@ impl Child { /// handle of a child process asynchronously. #[derive(Debug)] pub struct ChildStdin { - inner: imp::ChildStdin, + inner: imp::ChildStdio, } /// The standard output stream for spawned children. @@ -1130,7 +1316,7 @@ pub struct ChildStdin { /// handle of a child process asynchronously. #[derive(Debug)] pub struct ChildStdout { - inner: imp::ChildStdout, + inner: imp::ChildStdio, } /// The standard error stream for spawned children. @@ -1139,46 +1325,101 @@ pub struct ChildStdout { /// handle of a child process asynchronously. #[derive(Debug)] pub struct ChildStderr { - inner: imp::ChildStderr, + inner: imp::ChildStdio, +} + +impl ChildStdin { + /// Creates an asynchronous `ChildStdin` from a synchronous one. + /// + /// # Errors + /// + /// This method may fail if an error is encountered when setting the pipe to + /// non-blocking mode, or when registering the pipe with the runtime's IO + /// driver. + pub fn from_std(inner: std::process::ChildStdin) -> io::Result<Self> { + Ok(Self { + inner: imp::stdio(inner)?, + }) + } +} + +impl ChildStdout { + /// Creates an asynchronous `ChildStdout` from a synchronous one. + /// + /// # Errors + /// + /// This method may fail if an error is encountered when setting the pipe to + /// non-blocking mode, or when registering the pipe with the runtime's IO + /// driver. + pub fn from_std(inner: std::process::ChildStdout) -> io::Result<Self> { + Ok(Self { + inner: imp::stdio(inner)?, + }) + } +} + +impl ChildStderr { + /// Creates an asynchronous `ChildStderr` from a synchronous one. + /// + /// # Errors + /// + /// This method may fail if an error is encountered when setting the pipe to + /// non-blocking mode, or when registering the pipe with the runtime's IO + /// driver. + pub fn from_std(inner: std::process::ChildStderr) -> io::Result<Self> { + Ok(Self { + inner: imp::stdio(inner)?, + }) + } } impl AsyncWrite for ChildStdin { fn poll_write( - self: Pin<&mut Self>, + mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll<io::Result<usize>> { - self.inner.poll_write(cx, buf) + Pin::new(&mut self.inner).poll_write(cx, buf) + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + Pin::new(&mut self.inner).poll_flush(cx) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + Pin::new(&mut self.inner).poll_shutdown(cx) } - fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> { - Poll::Ready(Ok(())) + fn poll_write_vectored( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[io::IoSlice<'_>], + ) -> Poll<Result<usize, io::Error>> { + Pin::new(&mut self.inner).poll_write_vectored(cx, bufs) } - fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> { - Poll::Ready(Ok(())) + fn is_write_vectored(&self) -> bool { + self.inner.is_write_vectored() } } impl AsyncRead for ChildStdout { fn poll_read( - self: Pin<&mut Self>, + mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll<io::Result<()>> { - // Safety: pipes support reading into uninitialized memory - unsafe { self.inner.poll_read(cx, buf) } + Pin::new(&mut self.inner).poll_read(cx, buf) } } impl AsyncRead for ChildStderr { fn poll_read( - self: Pin<&mut Self>, + mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll<io::Result<()>> { - // Safety: pipes support reading into uninitialized memory - unsafe { self.inner.poll_read(cx, buf) } + Pin::new(&mut self.inner).poll_read(cx, buf) } } @@ -1208,6 +1449,8 @@ impl TryInto<Stdio> for ChildStderr { #[cfg(unix)] mod sys { + #[cfg(not(tokio_no_as_fd))] + use std::os::unix::io::{AsFd, BorrowedFd}; use std::os::unix::io::{AsRawFd, RawFd}; use super::{ChildStderr, ChildStdin, ChildStdout}; @@ -1218,42 +1461,79 @@ mod sys { } } + #[cfg(not(tokio_no_as_fd))] + impl AsFd for ChildStdin { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } + } + impl AsRawFd for ChildStdout { fn as_raw_fd(&self) -> RawFd { self.inner.as_raw_fd() } } + #[cfg(not(tokio_no_as_fd))] + impl AsFd for ChildStdout { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } + } + impl AsRawFd for ChildStderr { fn as_raw_fd(&self) -> RawFd { self.inner.as_raw_fd() } } -} -#[cfg(windows)] -mod sys { - use std::os::windows::io::{AsRawHandle, RawHandle}; - - use super::{ChildStderr, ChildStdin, ChildStdout}; + #[cfg(not(tokio_no_as_fd))] + impl AsFd for ChildStderr { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } + } +} +cfg_windows! { impl AsRawHandle for ChildStdin { fn as_raw_handle(&self) -> RawHandle { self.inner.as_raw_handle() } } + #[cfg(not(tokio_no_as_fd))] + impl AsHandle for ChildStdin { + fn as_handle(&self) -> BorrowedHandle<'_> { + unsafe { BorrowedHandle::borrow_raw(self.as_raw_handle()) } + } + } + impl AsRawHandle for ChildStdout { fn as_raw_handle(&self) -> RawHandle { self.inner.as_raw_handle() } } + #[cfg(not(tokio_no_as_fd))] + impl AsHandle for ChildStdout { + fn as_handle(&self) -> BorrowedHandle<'_> { + unsafe { BorrowedHandle::borrow_raw(self.as_raw_handle()) } + } + } + impl AsRawHandle for ChildStderr { fn as_raw_handle(&self) -> RawHandle { self.inner.as_raw_handle() } } + + #[cfg(not(tokio_no_as_fd))] + impl AsHandle for ChildStderr { + fn as_handle(&self) -> BorrowedHandle<'_> { + unsafe { BorrowedHandle::borrow_raw(self.as_raw_handle()) } + } + } } #[cfg(all(test, not(loom)))] diff --git a/vendor/tokio/src/process/unix/mod.rs b/vendor/tokio/src/process/unix/mod.rs index fab63dd3d..41857f49c 100644 --- a/vendor/tokio/src/process/unix/mod.rs +++ b/vendor/tokio/src/process/unix/mod.rs @@ -1,4 +1,4 @@ -//! Unix handling of child processes +//! Unix handling of child processes. //! //! Right now the only "fancy" thing about this is how we implement the //! `Future` implementation on `Child` to get the exit status. Unix offers @@ -21,27 +21,26 @@ //! processes in general aren't scalable (e.g. millions) so it shouldn't be that //! bad in theory... -pub(crate) mod driver; - pub(crate) mod orphan; use orphan::{OrphanQueue, OrphanQueueImpl, Wait}; mod reap; use reap::Reaper; -use crate::io::PollEvented; +use crate::io::{AsyncRead, AsyncWrite, PollEvented, ReadBuf}; use crate::process::kill::Kill; use crate::process::SpawnedChild; -use crate::signal::unix::driver::Handle as SignalHandle; +use crate::runtime::signal::Handle as SignalHandle; use crate::signal::unix::{signal, Signal, SignalKind}; use mio::event::Source; use mio::unix::SourceFd; -use once_cell::sync::Lazy; use std::fmt; use std::fs::File; use std::future::Future; use std::io; +#[cfg(not(tokio_no_as_fd))] +use std::os::unix::io::{AsFd, BorrowedFd}; use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd}; use std::pin::Pin; use std::process::{Child as StdChild, ExitStatus, Stdio}; @@ -64,25 +63,41 @@ impl Kill for StdChild { } } -static ORPHAN_QUEUE: Lazy<OrphanQueueImpl<StdChild>> = Lazy::new(OrphanQueueImpl::new); +cfg_not_has_const_mutex_new! { + fn get_orphan_queue() -> &'static OrphanQueueImpl<StdChild> { + use crate::util::once_cell::OnceCell; + + static ORPHAN_QUEUE: OnceCell<OrphanQueueImpl<StdChild>> = OnceCell::new(); + + ORPHAN_QUEUE.get(OrphanQueueImpl::new) + } +} + +cfg_has_const_mutex_new! { + fn get_orphan_queue() -> &'static OrphanQueueImpl<StdChild> { + static ORPHAN_QUEUE: OrphanQueueImpl<StdChild> = OrphanQueueImpl::new(); + + &ORPHAN_QUEUE + } +} pub(crate) struct GlobalOrphanQueue; impl fmt::Debug for GlobalOrphanQueue { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - ORPHAN_QUEUE.fmt(fmt) + get_orphan_queue().fmt(fmt) } } impl GlobalOrphanQueue { - fn reap_orphans(handle: &SignalHandle) { - ORPHAN_QUEUE.reap_orphans(handle) + pub(crate) fn reap_orphans(handle: &SignalHandle) { + get_orphan_queue().reap_orphans(handle) } } impl OrphanQueue<StdChild> for GlobalOrphanQueue { fn push_orphan(&self, orphan: StdChild) { - ORPHAN_QUEUE.push_orphan(orphan) + get_orphan_queue().push_orphan(orphan) } } @@ -101,9 +116,9 @@ impl fmt::Debug for Child { pub(crate) fn spawn_child(cmd: &mut std::process::Command) -> io::Result<SpawnedChild> { let mut child = cmd.spawn()?; - let stdin = stdio(child.stdin.take())?; - let stdout = stdio(child.stdout.take())?; - let stderr = stdio(child.stderr.take())?; + let stdin = child.stdin.take().map(stdio).transpose()?; + let stdout = child.stdout.take().map(stdio).transpose()?; + let stderr = child.stderr.take().map(stdio).transpose()?; let signal = signal(SignalKind::child())?; @@ -143,7 +158,7 @@ impl Future for Child { #[derive(Debug)] pub(crate) struct Pipe { - // Actually a pipe and not a File. However, we are reusing `File` to get + // Actually a pipe is not a File. However, we are reusing `File` to get // close on drop. This is a similar trick as `mio`. fd: File, } @@ -169,6 +184,10 @@ impl<'a> io::Write for &'a Pipe { fn flush(&mut self) -> io::Result<()> { (&self.fd).flush() } + + fn write_vectored(&mut self, bufs: &[io::IoSlice<'_>]) -> io::Result<usize> { + (&self.fd).write_vectored(bufs) + } } impl AsRawFd for Pipe { @@ -177,8 +196,15 @@ impl AsRawFd for Pipe { } } -pub(crate) fn convert_to_stdio(io: PollEvented<Pipe>) -> io::Result<Stdio> { - let mut fd = io.into_inner()?.fd; +#[cfg(not(tokio_no_as_fd))] +impl AsFd for Pipe { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } +} + +pub(crate) fn convert_to_stdio(io: ChildStdio) -> io::Result<Stdio> { + let mut fd = io.inner.into_inner()?.fd; // Ensure that the fd to be inherited is set to *blocking* mode, as this // is the default that virtually all programs expect to have. Those @@ -213,9 +239,69 @@ impl Source for Pipe { } } -pub(crate) type ChildStdin = PollEvented<Pipe>; -pub(crate) type ChildStdout = PollEvented<Pipe>; -pub(crate) type ChildStderr = PollEvented<Pipe>; +pub(crate) struct ChildStdio { + inner: PollEvented<Pipe>, +} + +impl fmt::Debug for ChildStdio { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + self.inner.fmt(fmt) + } +} + +impl AsRawFd for ChildStdio { + fn as_raw_fd(&self) -> RawFd { + self.inner.as_raw_fd() + } +} + +#[cfg(not(tokio_no_as_fd))] +impl AsFd for ChildStdio { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } +} + +impl AsyncWrite for ChildStdio { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + self.inner.poll_write(cx, buf) + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> { + Poll::Ready(Ok(())) + } + + fn poll_write_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[io::IoSlice<'_>], + ) -> Poll<Result<usize, io::Error>> { + self.inner.poll_write_vectored(cx, bufs) + } + + fn is_write_vectored(&self) -> bool { + true + } +} + +impl AsyncRead for ChildStdio { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<()>> { + // Safety: pipes support reading into uninitialized memory + unsafe { self.inner.poll_read(cx, buf) } + } +} fn set_nonblocking<T: AsRawFd>(fd: &mut T, nonblocking: bool) -> io::Result<()> { unsafe { @@ -240,18 +326,13 @@ fn set_nonblocking<T: AsRawFd>(fd: &mut T, nonblocking: bool) -> io::Result<()> Ok(()) } -fn stdio<T>(option: Option<T>) -> io::Result<Option<PollEvented<Pipe>>> +pub(super) fn stdio<T>(io: T) -> io::Result<ChildStdio> where T: IntoRawFd, { - let io = match option { - Some(io) => io, - None => return Ok(None), - }; - // Set the fd to nonblocking before we pass it to the event loop let mut pipe = Pipe::from(io); set_nonblocking(&mut pipe, true)?; - Ok(Some(PollEvented::new(pipe)?)) + PollEvented::new(pipe).map(|inner| ChildStdio { inner }) } diff --git a/vendor/tokio/src/process/unix/orphan.rs b/vendor/tokio/src/process/unix/orphan.rs index 07f0dcfe4..340719603 100644 --- a/vendor/tokio/src/process/unix/orphan.rs +++ b/vendor/tokio/src/process/unix/orphan.rs @@ -1,5 +1,5 @@ use crate::loom::sync::{Mutex, MutexGuard}; -use crate::signal::unix::driver::Handle as SignalHandle; +use crate::runtime::signal::Handle as SignalHandle; use crate::signal::unix::{signal_with_handle, SignalKind}; use crate::sync::watch; use std::io; @@ -43,10 +43,21 @@ pub(crate) struct OrphanQueueImpl<T> { } impl<T> OrphanQueueImpl<T> { - pub(crate) fn new() -> Self { - Self { - sigchild: Mutex::new(None), - queue: Mutex::new(Vec::new()), + cfg_not_has_const_mutex_new! { + pub(crate) fn new() -> Self { + Self { + sigchild: Mutex::new(None), + queue: Mutex::new(Vec::new()), + } + } + } + + cfg_has_const_mutex_new! { + pub(crate) const fn new() -> Self { + Self { + sigchild: Mutex::const_new(None), + queue: Mutex::const_new(Vec::new()), + } } } @@ -87,7 +98,7 @@ impl<T> OrphanQueueImpl<T> { // means that the signal driver isn't running, in // which case there isn't anything we can // register/initialize here, so we can try again later - if let Ok(sigchild) = signal_with_handle(SignalKind::child(), &handle) { + if let Ok(sigchild) = signal_with_handle(SignalKind::child(), handle) { *sigchild_guard = Some(sigchild); drain_orphan_queue(queue); } @@ -120,8 +131,8 @@ where #[cfg(all(test, not(loom)))] pub(crate) mod test { use super::*; - use crate::io::driver::Driver as IoDriver; - use crate::signal::unix::driver::{Driver as SignalDriver, Handle as SignalHandle}; + use crate::runtime::io::Driver as IoDriver; + use crate::runtime::signal::{Driver as SignalDriver, Handle as SignalHandle}; use crate::sync::watch; use std::cell::{Cell, RefCell}; use std::io; @@ -280,9 +291,11 @@ pub(crate) mod test { drop(signal_guard); } + #[cfg_attr(miri, ignore)] // Miri does not support epoll. #[test] fn does_not_register_signal_if_queue_empty() { - let signal_driver = IoDriver::new().and_then(SignalDriver::new).unwrap(); + let (io_driver, io_handle) = IoDriver::new(1024).unwrap(); + let signal_driver = SignalDriver::new(io_driver, &io_handle).unwrap(); let handle = signal_driver.handle(); let orphanage = OrphanQueueImpl::new(); diff --git a/vendor/tokio/src/process/windows.rs b/vendor/tokio/src/process/windows.rs index 7237525da..eb412cdc5 100644 --- a/vendor/tokio/src/process/windows.rs +++ b/vendor/tokio/src/process/windows.rs @@ -15,29 +15,30 @@ //! `RegisterWaitForSingleObject` and then wait on the other end of the oneshot //! from then on out. -use crate::io::PollEvented; +use crate::io::{blocking::Blocking, AsyncRead, AsyncWrite, ReadBuf}; use crate::process::kill::Kill; use crate::process::SpawnedChild; use crate::sync::oneshot; -use mio::windows::NamedPipe; use std::fmt; +use std::fs::File as StdFile; use std::future::Future; use std::io; -use std::os::windows::prelude::{AsRawHandle, FromRawHandle, IntoRawHandle}; +use std::os::windows::prelude::{AsRawHandle, IntoRawHandle, RawHandle}; use std::pin::Pin; use std::process::Stdio; use std::process::{Child as StdChild, Command as StdCommand, ExitStatus}; -use std::ptr; -use std::task::Context; -use std::task::Poll; -use winapi::shared::minwindef::{DWORD, FALSE}; -use winapi::um::handleapi::{DuplicateHandle, INVALID_HANDLE_VALUE}; -use winapi::um::processthreadsapi::GetCurrentProcess; -use winapi::um::threadpoollegacyapiset::UnregisterWaitEx; -use winapi::um::winbase::{RegisterWaitForSingleObject, INFINITE}; -use winapi::um::winnt::{ - BOOLEAN, DUPLICATE_SAME_ACCESS, HANDLE, PVOID, WT_EXECUTEINWAITTHREAD, WT_EXECUTEONLYONCE, +use std::sync::Arc; +use std::task::{Context, Poll}; + +use windows_sys::{ + Win32::Foundation::{ + DuplicateHandle, BOOLEAN, DUPLICATE_SAME_ACCESS, HANDLE, INVALID_HANDLE_VALUE, + }, + Win32::System::Threading::{ + GetCurrentProcess, RegisterWaitForSingleObject, UnregisterWaitEx, INFINITE, + WT_EXECUTEINWAITTHREAD, WT_EXECUTEONLYONCE, + }, }; #[must_use = "futures do nothing unless polled"] @@ -67,9 +68,9 @@ unsafe impl Send for Waiting {} pub(crate) fn spawn_child(cmd: &mut StdCommand) -> io::Result<SpawnedChild> { let mut child = cmd.spawn()?; - let stdin = stdio(child.stdin.take()); - let stdout = stdio(child.stdout.take()); - let stderr = stdio(child.stderr.take()); + let stdin = child.stdin.take().map(stdio).transpose()?; + let stdout = child.stdout.take().map(stdio).transpose()?; + let stderr = child.stderr.take().map(stdio).transpose()?; Ok(SpawnedChild { child: Child { @@ -119,11 +120,11 @@ impl Future for Child { } let (tx, rx) = oneshot::channel(); let ptr = Box::into_raw(Box::new(Some(tx))); - let mut wait_object = ptr::null_mut(); + let mut wait_object = 0; let rc = unsafe { RegisterWaitForSingleObject( &mut wait_object, - inner.child.as_raw_handle(), + inner.child.as_raw_handle() as _, Some(callback), ptr as *mut _, INFINITE, @@ -144,6 +145,12 @@ impl Future for Child { } } +impl AsRawHandle for Child { + fn as_raw_handle(&self) -> RawHandle { + self.child.as_raw_handle() + } +} + impl Drop for Waiting { fn drop(&mut self) { unsafe { @@ -156,43 +163,106 @@ impl Drop for Waiting { } } -unsafe extern "system" fn callback(ptr: PVOID, _timer_fired: BOOLEAN) { +unsafe extern "system" fn callback(ptr: *mut std::ffi::c_void, _timer_fired: BOOLEAN) { let complete = &mut *(ptr as *mut Option<oneshot::Sender<()>>); let _ = complete.take().unwrap().send(()); } -pub(crate) type ChildStdin = PollEvented<NamedPipe>; -pub(crate) type ChildStdout = PollEvented<NamedPipe>; -pub(crate) type ChildStderr = PollEvented<NamedPipe>; +#[derive(Debug)] +struct ArcFile(Arc<StdFile>); + +impl io::Read for ArcFile { + fn read(&mut self, bytes: &mut [u8]) -> io::Result<usize> { + (&*self.0).read(bytes) + } +} + +impl io::Write for ArcFile { + fn write(&mut self, bytes: &[u8]) -> io::Result<usize> { + (&*self.0).write(bytes) + } + + fn flush(&mut self) -> io::Result<()> { + (&*self.0).flush() + } +} -fn stdio<T>(option: Option<T>) -> Option<PollEvented<NamedPipe>> +#[derive(Debug)] +pub(crate) struct ChildStdio { + // Used for accessing the raw handle, even if the io version is busy + raw: Arc<StdFile>, + // For doing I/O operations asynchronously + io: Blocking<ArcFile>, +} + +impl AsRawHandle for ChildStdio { + fn as_raw_handle(&self) -> RawHandle { + self.raw.as_raw_handle() + } +} + +impl AsyncRead for ChildStdio { + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<()>> { + Pin::new(&mut self.io).poll_read(cx, buf) + } +} + +impl AsyncWrite for ChildStdio { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + Pin::new(&mut self.io).poll_write(cx, buf) + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + Pin::new(&mut self.io).poll_flush(cx) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + Pin::new(&mut self.io).poll_shutdown(cx) + } +} + +pub(super) fn stdio<T>(io: T) -> io::Result<ChildStdio> where T: IntoRawHandle, { - let io = match option { - Some(io) => io, - None => return None, - }; - let pipe = unsafe { NamedPipe::from_raw_handle(io.into_raw_handle()) }; - PollEvented::new(pipe).ok() + use std::os::windows::prelude::FromRawHandle; + + let raw = Arc::new(unsafe { StdFile::from_raw_handle(io.into_raw_handle()) }); + let io = Blocking::new(ArcFile(raw.clone())); + Ok(ChildStdio { raw, io }) +} + +pub(crate) fn convert_to_stdio(child_stdio: ChildStdio) -> io::Result<Stdio> { + let ChildStdio { raw, io } = child_stdio; + drop(io); // Try to drop the Arc count here + + Arc::try_unwrap(raw) + .or_else(|raw| duplicate_handle(&*raw)) + .map(Stdio::from) } -pub(crate) fn convert_to_stdio(io: PollEvented<NamedPipe>) -> io::Result<Stdio> { - let named_pipe = io.into_inner()?; +fn duplicate_handle<T: AsRawHandle>(io: &T) -> io::Result<StdFile> { + use std::os::windows::prelude::FromRawHandle; - // Mio does not implement `IntoRawHandle` for `NamedPipe`, so we'll manually - // duplicate the handle here... unsafe { let mut dup_handle = INVALID_HANDLE_VALUE; let cur_proc = GetCurrentProcess(); let status = DuplicateHandle( cur_proc, - named_pipe.as_raw_handle(), + io.as_raw_handle() as _, cur_proc, &mut dup_handle, - 0 as DWORD, - FALSE, + 0, + 0, DUPLICATE_SAME_ACCESS, ); @@ -200,6 +270,6 @@ pub(crate) fn convert_to_stdio(io: PollEvented<NamedPipe>) -> io::Result<Stdio> return Err(io::Error::last_os_error()); } - Ok(Stdio::from_raw_handle(dup_handle)) + Ok(StdFile::from_raw_handle(dup_handle as _)) } } diff --git a/vendor/tokio/src/runtime/basic_scheduler.rs b/vendor/tokio/src/runtime/basic_scheduler.rs deleted file mode 100644 index 13dfb6973..000000000 --- a/vendor/tokio/src/runtime/basic_scheduler.rs +++ /dev/null @@ -1,534 +0,0 @@ -use crate::future::poll_fn; -use crate::loom::sync::atomic::AtomicBool; -use crate::loom::sync::Mutex; -use crate::park::{Park, Unpark}; -use crate::runtime::task::{self, JoinHandle, Schedule, Task}; -use crate::sync::notify::Notify; -use crate::util::linked_list::{Link, LinkedList}; -use crate::util::{waker_ref, Wake, WakerRef}; - -use std::cell::RefCell; -use std::collections::VecDeque; -use std::fmt; -use std::future::Future; -use std::ptr::NonNull; -use std::sync::atomic::Ordering::{AcqRel, Acquire, Release}; -use std::sync::Arc; -use std::task::Poll::{Pending, Ready}; -use std::time::Duration; - -/// Executes tasks on the current thread -pub(crate) struct BasicScheduler<P: Park> { - /// Inner state guarded by a mutex that is shared - /// between all `block_on` calls. - inner: Mutex<Option<Inner<P>>>, - - /// Notifier for waking up other threads to steal the - /// parker. - notify: Notify, - - /// Sendable task spawner - spawner: Spawner, -} - -/// The inner scheduler that owns the task queue and the main parker P. -struct Inner<P: Park> { - /// Scheduler run queue - /// - /// When the scheduler is executed, the queue is removed from `self` and - /// moved into `Context`. - /// - /// This indirection is to allow `BasicScheduler` to be `Send`. - tasks: Option<Tasks>, - - /// Sendable task spawner - spawner: Spawner, - - /// Current tick - tick: u8, - - /// Thread park handle - park: P, -} - -#[derive(Clone)] -pub(crate) struct Spawner { - shared: Arc<Shared>, -} - -struct Tasks { - /// Collection of all active tasks spawned onto this executor. - owned: LinkedList<Task<Arc<Shared>>, <Task<Arc<Shared>> as Link>::Target>, - - /// Local run queue. - /// - /// Tasks notified from the current thread are pushed into this queue. - queue: VecDeque<task::Notified<Arc<Shared>>>, -} - -/// A remote scheduler entry. -/// -/// These are filled in by remote threads sending instructions to the scheduler. -enum Entry { - /// A remote thread wants to spawn a task. - Schedule(task::Notified<Arc<Shared>>), - /// A remote thread wants a task to be released by the scheduler. We only - /// have access to its header. - Release(NonNull<task::Header>), -} - -// Safety: Used correctly, the task header is "thread safe". Ultimately the task -// is owned by the current thread executor, for which this instruction is being -// sent. -unsafe impl Send for Entry {} - -/// Scheduler state shared between threads. -struct Shared { - /// Remote run queue. None if the `Runtime` has been dropped. - queue: Mutex<Option<VecDeque<Entry>>>, - - /// Unpark the blocked thread. - unpark: Box<dyn Unpark>, - - /// Indicates whether the blocked on thread was woken. - woken: AtomicBool, -} - -/// Thread-local context. -struct Context { - /// Shared scheduler state - shared: Arc<Shared>, - - /// Local queue - tasks: RefCell<Tasks>, -} - -/// Initial queue capacity. -const INITIAL_CAPACITY: usize = 64; - -/// Max number of tasks to poll per tick. -#[cfg(loom)] -const MAX_TASKS_PER_TICK: usize = 4; -#[cfg(not(loom))] -const MAX_TASKS_PER_TICK: usize = 61; - -/// How often to check the remote queue first. -const REMOTE_FIRST_INTERVAL: u8 = 31; - -// Tracks the current BasicScheduler. -scoped_thread_local!(static CURRENT: Context); - -impl<P: Park> BasicScheduler<P> { - pub(crate) fn new(park: P) -> BasicScheduler<P> { - let unpark = Box::new(park.unpark()); - - let spawner = Spawner { - shared: Arc::new(Shared { - queue: Mutex::new(Some(VecDeque::with_capacity(INITIAL_CAPACITY))), - unpark: unpark as Box<dyn Unpark>, - woken: AtomicBool::new(false), - }), - }; - - let inner = Mutex::new(Some(Inner { - tasks: Some(Tasks { - owned: LinkedList::new(), - queue: VecDeque::with_capacity(INITIAL_CAPACITY), - }), - spawner: spawner.clone(), - tick: 0, - park, - })); - - BasicScheduler { - inner, - notify: Notify::new(), - spawner, - } - } - - pub(crate) fn spawner(&self) -> &Spawner { - &self.spawner - } - - pub(crate) fn block_on<F: Future>(&self, future: F) -> F::Output { - pin!(future); - - // Attempt to steal the dedicated parker and block_on the future if we can there, - // otherwise, lets select on a notification that the parker is available - // or the future is complete. - loop { - if let Some(inner) = &mut self.take_inner() { - return inner.block_on(future); - } else { - let mut enter = crate::runtime::enter(false); - - let notified = self.notify.notified(); - pin!(notified); - - if let Some(out) = enter - .block_on(poll_fn(|cx| { - if notified.as_mut().poll(cx).is_ready() { - return Ready(None); - } - - if let Ready(out) = future.as_mut().poll(cx) { - return Ready(Some(out)); - } - - Pending - })) - .expect("Failed to `Enter::block_on`") - { - return out; - } - } - } - } - - fn take_inner(&self) -> Option<InnerGuard<'_, P>> { - let inner = self.inner.lock().take()?; - - Some(InnerGuard { - inner: Some(inner), - basic_scheduler: &self, - }) - } -} - -impl<P: Park> Inner<P> { - /// Block on the future provided and drive the runtime's driver. - fn block_on<F: Future>(&mut self, future: F) -> F::Output { - enter(self, |scheduler, context| { - let _enter = crate::runtime::enter(false); - let waker = scheduler.spawner.waker_ref(); - let mut cx = std::task::Context::from_waker(&waker); - let mut polled = false; - - pin!(future); - - 'outer: loop { - if scheduler.spawner.was_woken() || !polled { - polled = true; - if let Ready(v) = crate::coop::budget(|| future.as_mut().poll(&mut cx)) { - return v; - } - } - - for _ in 0..MAX_TASKS_PER_TICK { - // Get and increment the current tick - let tick = scheduler.tick; - scheduler.tick = scheduler.tick.wrapping_add(1); - - let entry = if tick % REMOTE_FIRST_INTERVAL == 0 { - scheduler.spawner.pop().or_else(|| { - context - .tasks - .borrow_mut() - .queue - .pop_front() - .map(Entry::Schedule) - }) - } else { - context - .tasks - .borrow_mut() - .queue - .pop_front() - .map(Entry::Schedule) - .or_else(|| scheduler.spawner.pop()) - }; - - let entry = match entry { - Some(entry) => entry, - None => { - // Park until the thread is signaled - scheduler.park.park().expect("failed to park"); - - // Try polling the `block_on` future next - continue 'outer; - } - }; - - match entry { - Entry::Schedule(task) => crate::coop::budget(|| task.run()), - Entry::Release(ptr) => { - // Safety: the task header is only legally provided - // internally in the header, so we know that it is a - // valid (or in particular *allocated*) header that - // is part of the linked list. - unsafe { - let removed = context.tasks.borrow_mut().owned.remove(ptr); - - // TODO: This seems like it should hold, because - // there doesn't seem to be an avenue for anyone - // else to fiddle with the owned tasks - // collection *after* a remote thread has marked - // it as released, and at that point, the only - // location at which it can be removed is here - // or in the Drop implementation of the - // scheduler. - debug_assert!(removed.is_some()); - } - } - } - } - - // Yield to the park, this drives the timer and pulls any pending - // I/O events. - scheduler - .park - .park_timeout(Duration::from_millis(0)) - .expect("failed to park"); - } - }) - } -} - -/// Enter the scheduler context. This sets the queue and other necessary -/// scheduler state in the thread-local -fn enter<F, R, P>(scheduler: &mut Inner<P>, f: F) -> R -where - F: FnOnce(&mut Inner<P>, &Context) -> R, - P: Park, -{ - // Ensures the run queue is placed back in the `BasicScheduler` instance - // once `block_on` returns.` - struct Guard<'a, P: Park> { - context: Option<Context>, - scheduler: &'a mut Inner<P>, - } - - impl<P: Park> Drop for Guard<'_, P> { - fn drop(&mut self) { - let Context { tasks, .. } = self.context.take().expect("context missing"); - self.scheduler.tasks = Some(tasks.into_inner()); - } - } - - // Remove `tasks` from `self` and place it in a `Context`. - let tasks = scheduler.tasks.take().expect("invalid state"); - - let guard = Guard { - context: Some(Context { - shared: scheduler.spawner.shared.clone(), - tasks: RefCell::new(tasks), - }), - scheduler, - }; - - let context = guard.context.as_ref().unwrap(); - let scheduler = &mut *guard.scheduler; - - CURRENT.set(context, || f(scheduler, context)) -} - -impl<P: Park> Drop for BasicScheduler<P> { - fn drop(&mut self) { - // Avoid a double panic if we are currently panicking and - // the lock may be poisoned. - - let mut inner = match self.inner.lock().take() { - Some(inner) => inner, - None if std::thread::panicking() => return, - None => panic!("Oh no! We never placed the Inner state back, this is a bug!"), - }; - - enter(&mut inner, |scheduler, context| { - // Loop required here to ensure borrow is dropped between iterations - #[allow(clippy::while_let_loop)] - loop { - let task = match context.tasks.borrow_mut().owned.pop_back() { - Some(task) => task, - None => break, - }; - - task.shutdown(); - } - - // Drain local queue - for task in context.tasks.borrow_mut().queue.drain(..) { - task.shutdown(); - } - - // Drain remote queue and set it to None - let mut remote_queue = scheduler.spawner.shared.queue.lock(); - - // Using `Option::take` to replace the shared queue with `None`. - if let Some(remote_queue) = remote_queue.take() { - for entry in remote_queue { - match entry { - Entry::Schedule(task) => { - task.shutdown(); - } - Entry::Release(..) => { - // Do nothing, each entry in the linked list was *just* - // dropped by the scheduler above. - } - } - } - } - // By dropping the mutex lock after the full duration of the above loop, - // any thread that sees the queue in the `None` state is guaranteed that - // the runtime has fully shut down. - // - // The assert below is unrelated to this mutex. - drop(remote_queue); - - assert!(context.tasks.borrow().owned.is_empty()); - }); - } -} - -impl<P: Park> fmt::Debug for BasicScheduler<P> { - fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - fmt.debug_struct("BasicScheduler").finish() - } -} - -// ===== impl Spawner ===== - -impl Spawner { - /// Spawns a future onto the thread pool - pub(crate) fn spawn<F>(&self, future: F) -> JoinHandle<F::Output> - where - F: crate::future::Future + Send + 'static, - F::Output: Send + 'static, - { - let (task, handle) = task::joinable(future); - self.shared.schedule(task); - handle - } - - fn pop(&self) -> Option<Entry> { - match self.shared.queue.lock().as_mut() { - Some(queue) => queue.pop_front(), - None => None, - } - } - - fn waker_ref(&self) -> WakerRef<'_> { - // clear the woken bit - self.shared.woken.swap(false, AcqRel); - waker_ref(&self.shared) - } - - fn was_woken(&self) -> bool { - self.shared.woken.load(Acquire) - } -} - -impl fmt::Debug for Spawner { - fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - fmt.debug_struct("Spawner").finish() - } -} - -// ===== impl Shared ===== - -impl Schedule for Arc<Shared> { - fn bind(task: Task<Self>) -> Arc<Shared> { - CURRENT.with(|maybe_cx| { - let cx = maybe_cx.expect("scheduler context missing"); - cx.tasks.borrow_mut().owned.push_front(task); - cx.shared.clone() - }) - } - - fn release(&self, task: &Task<Self>) -> Option<Task<Self>> { - CURRENT.with(|maybe_cx| { - let ptr = NonNull::from(task.header()); - - if let Some(cx) = maybe_cx { - // safety: the task is inserted in the list in `bind`. - unsafe { cx.tasks.borrow_mut().owned.remove(ptr) } - } else { - // By sending an `Entry::Release` to the runtime, we ask the - // runtime to remove this task from the linked list in - // `Tasks::owned`. - // - // If the queue is `None`, then the task was already removed - // from that list in the destructor of `BasicScheduler`. We do - // not do anything in this case for the same reason that - // `Entry::Release` messages are ignored in the remote queue - // drain loop of `BasicScheduler`'s destructor. - if let Some(queue) = self.queue.lock().as_mut() { - queue.push_back(Entry::Release(ptr)); - } - - self.unpark.unpark(); - // Returning `None` here prevents the task plumbing from being - // freed. It is then up to the scheduler through the queue we - // just added to, or its Drop impl to free the task. - None - } - }) - } - - fn schedule(&self, task: task::Notified<Self>) { - CURRENT.with(|maybe_cx| match maybe_cx { - Some(cx) if Arc::ptr_eq(self, &cx.shared) => { - cx.tasks.borrow_mut().queue.push_back(task); - } - _ => { - let mut guard = self.queue.lock(); - if let Some(queue) = guard.as_mut() { - queue.push_back(Entry::Schedule(task)); - drop(guard); - self.unpark.unpark(); - } else { - // The runtime has shut down. We drop the new task - // immediately. - drop(guard); - task.shutdown(); - } - } - }); - } -} - -impl Wake for Shared { - fn wake(self: Arc<Self>) { - Wake::wake_by_ref(&self) - } - - /// Wake by reference - fn wake_by_ref(arc_self: &Arc<Self>) { - arc_self.woken.store(true, Release); - arc_self.unpark.unpark(); - } -} - -// ===== InnerGuard ===== - -/// Used to ensure we always place the Inner value -/// back into its slot in `BasicScheduler`, even if the -/// future panics. -struct InnerGuard<'a, P: Park> { - inner: Option<Inner<P>>, - basic_scheduler: &'a BasicScheduler<P>, -} - -impl<P: Park> InnerGuard<'_, P> { - fn block_on<F: Future>(&mut self, future: F) -> F::Output { - // The only time inner gets set to `None` is if we have dropped - // already so this unwrap is safe. - self.inner.as_mut().unwrap().block_on(future) - } -} - -impl<P: Park> Drop for InnerGuard<'_, P> { - fn drop(&mut self) { - if let Some(scheduler) = self.inner.take() { - let mut lock = self.basic_scheduler.inner.lock(); - - // Replace old scheduler back into the state to allow - // other threads to pick it up and drive it. - lock.replace(scheduler); - - // Wake up other possible threads that could steal - // the dedicated parker P. - self.basic_scheduler.notify.notify_one() - } - } -} diff --git a/vendor/tokio/src/runtime/blocking/mod.rs b/vendor/tokio/src/runtime/blocking/mod.rs index fece3c279..c42924be7 100644 --- a/vendor/tokio/src/runtime/blocking/mod.rs +++ b/vendor/tokio/src/runtime/blocking/mod.rs @@ -6,37 +6,21 @@ mod pool; pub(crate) use pool::{spawn_blocking, BlockingPool, Spawner}; +cfg_fs! { + pub(crate) use pool::spawn_mandatory_blocking; +} + +cfg_trace! { + pub(crate) use pool::Mandatory; +} + mod schedule; mod shutdown; -pub(crate) mod task; +mod task; +pub(crate) use task::BlockingTask; use crate::runtime::Builder; pub(crate) fn create_blocking_pool(builder: &Builder, thread_cap: usize) -> BlockingPool { BlockingPool::new(builder, thread_cap) } - -/* -cfg_not_blocking_impl! { - use crate::runtime::Builder; - use std::time::Duration; - - #[derive(Debug, Clone)] - pub(crate) struct BlockingPool {} - - pub(crate) use BlockingPool as Spawner; - - pub(crate) fn create_blocking_pool(_builder: &Builder, _thread_cap: usize) -> BlockingPool { - BlockingPool {} - } - - impl BlockingPool { - pub(crate) fn spawner(&self) -> &BlockingPool { - self - } - - pub(crate) fn shutdown(&mut self, _duration: Option<Duration>) { - } - } -} -*/ diff --git a/vendor/tokio/src/runtime/blocking/pool.rs b/vendor/tokio/src/runtime/blocking/pool.rs index b7d725128..a23b0a0d2 100644 --- a/vendor/tokio/src/runtime/blocking/pool.rs +++ b/vendor/tokio/src/runtime/blocking/pool.rs @@ -2,16 +2,16 @@ use crate::loom::sync::{Arc, Condvar, Mutex}; use crate::loom::thread; -use crate::runtime::blocking::schedule::NoopSchedule; -use crate::runtime::blocking::shutdown; +use crate::runtime::blocking::schedule::BlockingSchedule; +use crate::runtime::blocking::{shutdown, BlockingTask}; use crate::runtime::builder::ThreadNameFn; -use crate::runtime::context; use crate::runtime::task::{self, JoinHandle}; use crate::runtime::{Builder, Callback, Handle}; -use crate::util::error::CONTEXT_MISSING_ERROR; use std::collections::{HashMap, VecDeque}; use std::fmt; +use std::io; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; pub(crate) struct BlockingPool { @@ -24,36 +24,84 @@ pub(crate) struct Spawner { inner: Arc<Inner>, } +#[derive(Default)] +pub(crate) struct SpawnerMetrics { + num_threads: AtomicUsize, + num_idle_threads: AtomicUsize, + queue_depth: AtomicUsize, +} + +impl SpawnerMetrics { + fn num_threads(&self) -> usize { + self.num_threads.load(Ordering::Relaxed) + } + + fn num_idle_threads(&self) -> usize { + self.num_idle_threads.load(Ordering::Relaxed) + } + + cfg_metrics! { + fn queue_depth(&self) -> usize { + self.queue_depth.load(Ordering::Relaxed) + } + } + + fn inc_num_threads(&self) { + self.num_threads.fetch_add(1, Ordering::Relaxed); + } + + fn dec_num_threads(&self) { + self.num_threads.fetch_sub(1, Ordering::Relaxed); + } + + fn inc_num_idle_threads(&self) { + self.num_idle_threads.fetch_add(1, Ordering::Relaxed); + } + + fn dec_num_idle_threads(&self) -> usize { + self.num_idle_threads.fetch_sub(1, Ordering::Relaxed) + } + + fn inc_queue_depth(&self) { + self.queue_depth.fetch_add(1, Ordering::Relaxed); + } + + fn dec_queue_depth(&self) { + self.queue_depth.fetch_sub(1, Ordering::Relaxed); + } +} + struct Inner { - /// State shared between worker threads + /// State shared between worker threads. shared: Mutex<Shared>, /// Pool threads wait on this. condvar: Condvar, - /// Spawned threads use this name + /// Spawned threads use this name. thread_name: ThreadNameFn, - /// Spawned thread stack size + /// Spawned thread stack size. stack_size: Option<usize>, - /// Call after a thread starts + /// Call after a thread starts. after_start: Option<Callback>, - /// Call before a thread stops + /// Call before a thread stops. before_stop: Option<Callback>, - // Maximum number of threads + // Maximum number of threads. thread_cap: usize, - // Customizable wait timeout + // Customizable wait timeout. keep_alive: Duration, + + // Metrics about the pool. + metrics: SpawnerMetrics, } struct Shared { queue: VecDeque<Task>, - num_th: usize, - num_idle: u32, num_notify: u32, shutdown: bool, shutdown_tx: Option<shutdown::Sender>, @@ -67,24 +115,93 @@ struct Shared { /// calling shutdown handles joining on these. worker_threads: HashMap<usize, thread::JoinHandle<()>>, /// This is a counter used to iterate worker_threads in a consistent order (for loom's - /// benefit) + /// benefit). worker_thread_index: usize, } -type Task = task::Notified<NoopSchedule>; +pub(crate) struct Task { + task: task::UnownedTask<BlockingSchedule>, + mandatory: Mandatory, +} + +#[derive(PartialEq, Eq)] +pub(crate) enum Mandatory { + #[cfg_attr(not(fs), allow(dead_code))] + Mandatory, + NonMandatory, +} + +pub(crate) enum SpawnError { + /// Pool is shutting down and the task was not scheduled + ShuttingDown, + /// There are no worker threads available to take the task + /// and the OS failed to spawn a new one + NoThreads(io::Error), +} + +impl From<SpawnError> for io::Error { + fn from(e: SpawnError) -> Self { + match e { + SpawnError::ShuttingDown => { + io::Error::new(io::ErrorKind::Other, "blocking pool shutting down") + } + SpawnError::NoThreads(e) => e, + } + } +} + +impl Task { + pub(crate) fn new(task: task::UnownedTask<BlockingSchedule>, mandatory: Mandatory) -> Task { + Task { task, mandatory } + } + + fn run(self) { + self.task.run(); + } + + fn shutdown_or_run_if_mandatory(self) { + match self.mandatory { + Mandatory::NonMandatory => self.task.shutdown(), + Mandatory::Mandatory => self.task.run(), + } + } +} const KEEP_ALIVE: Duration = Duration::from_secs(10); -/// Run the provided function on an executor dedicated to blocking operations. +/// Runs the provided function on an executor dedicated to blocking operations. +/// Tasks will be scheduled as non-mandatory, meaning they may not get executed +/// in case of runtime shutdown. +#[track_caller] +#[cfg_attr(tokio_wasi, allow(dead_code))] pub(crate) fn spawn_blocking<F, R>(func: F) -> JoinHandle<R> where F: FnOnce() -> R + Send + 'static, R: Send + 'static, { - let rt = context::current().expect(CONTEXT_MISSING_ERROR); + let rt = Handle::current(); rt.spawn_blocking(func) } +cfg_fs! { + #[cfg_attr(any( + all(loom, not(test)), // the function is covered by loom tests + test + ), allow(dead_code))] + /// Runs the provided function on an executor dedicated to blocking + /// operations. Tasks will be scheduled as mandatory, meaning they are + /// guaranteed to run unless a shutdown is already taking place. In case a + /// shutdown is already taking place, `None` will be returned. + pub(crate) fn spawn_mandatory_blocking<F, R>(func: F) -> Option<JoinHandle<R>> + where + F: FnOnce() -> R + Send + 'static, + R: Send + 'static, + { + let rt = Handle::current(); + rt.inner.blocking_spawner().spawn_mandatory_blocking(&rt, func) + } +} + // ===== impl BlockingPool ===== impl BlockingPool { @@ -97,8 +214,6 @@ impl BlockingPool { inner: Arc::new(Inner { shared: Mutex::new(Shared { queue: VecDeque::new(), - num_th: 0, - num_idle: 0, num_notify: 0, shutdown: false, shutdown_tx: Some(shutdown_tx), @@ -113,6 +228,7 @@ impl BlockingPool { before_stop: builder.before_stop.clone(), thread_cap, keep_alive, + metrics: Default::default(), }), }, shutdown_rx, @@ -172,53 +288,164 @@ impl fmt::Debug for BlockingPool { // ===== impl Spawner ===== impl Spawner { - pub(crate) fn spawn(&self, task: Task, rt: &Handle) -> Result<(), ()> { - let shutdown_tx = { - let mut shared = self.inner.shared.lock(); - - if shared.shutdown { - // Shutdown the task - task.shutdown(); - - // no need to even push this task; it would never get picked up - return Err(()); + #[track_caller] + pub(crate) fn spawn_blocking<F, R>(&self, rt: &Handle, func: F) -> JoinHandle<R> + where + F: FnOnce() -> R + Send + 'static, + R: Send + 'static, + { + let (join_handle, spawn_result) = + if cfg!(debug_assertions) && std::mem::size_of::<F>() > 2048 { + self.spawn_blocking_inner(Box::new(func), Mandatory::NonMandatory, None, rt) + } else { + self.spawn_blocking_inner(func, Mandatory::NonMandatory, None, rt) + }; + + match spawn_result { + Ok(()) => join_handle, + // Compat: do not panic here, return the join_handle even though it will never resolve + Err(SpawnError::ShuttingDown) => join_handle, + Err(SpawnError::NoThreads(e)) => { + panic!("OS can't spawn worker thread: {}", e) } + } + } - shared.queue.push_back(task); - - if shared.num_idle == 0 { - // No threads are able to process the task. - - if shared.num_th == self.inner.thread_cap { - // At max number of threads - None - } else { - shared.num_th += 1; - assert!(shared.shutdown_tx.is_some()); - shared.shutdown_tx.clone() - } + cfg_fs! { + #[track_caller] + #[cfg_attr(any( + all(loom, not(test)), // the function is covered by loom tests + test + ), allow(dead_code))] + pub(crate) fn spawn_mandatory_blocking<F, R>(&self, rt: &Handle, func: F) -> Option<JoinHandle<R>> + where + F: FnOnce() -> R + Send + 'static, + R: Send + 'static, + { + let (join_handle, spawn_result) = if cfg!(debug_assertions) && std::mem::size_of::<F>() > 2048 { + self.spawn_blocking_inner( + Box::new(func), + Mandatory::Mandatory, + None, + rt, + ) + } else { + self.spawn_blocking_inner( + func, + Mandatory::Mandatory, + None, + rt, + ) + }; + + if spawn_result.is_ok() { + Some(join_handle) } else { - // Notify an idle worker thread. The notification counter - // is used to count the needed amount of notifications - // exactly. Thread libraries may generate spurious - // wakeups, this counter is used to keep us in a - // consistent state. - shared.num_idle -= 1; - shared.num_notify += 1; - self.inner.condvar.notify_one(); None } + } + } + + #[track_caller] + pub(crate) fn spawn_blocking_inner<F, R>( + &self, + func: F, + is_mandatory: Mandatory, + name: Option<&str>, + rt: &Handle, + ) -> (JoinHandle<R>, Result<(), SpawnError>) + where + F: FnOnce() -> R + Send + 'static, + R: Send + 'static, + { + let fut = BlockingTask::new(func); + let id = task::Id::next(); + #[cfg(all(tokio_unstable, feature = "tracing"))] + let fut = { + use tracing::Instrument; + let location = std::panic::Location::caller(); + let span = tracing::trace_span!( + target: "tokio::task::blocking", + "runtime.spawn", + kind = %"blocking", + task.name = %name.unwrap_or_default(), + task.id = id.as_u64(), + "fn" = %std::any::type_name::<F>(), + loc.file = location.file(), + loc.line = location.line(), + loc.col = location.column(), + ); + fut.instrument(span) }; - if let Some(shutdown_tx) = shutdown_tx { - let mut shared = self.inner.shared.lock(); + #[cfg(not(all(tokio_unstable, feature = "tracing")))] + let _ = name; - let id = shared.worker_thread_index; - shared.worker_thread_index += 1; + let (task, handle) = task::unowned(fut, BlockingSchedule::new(rt), id); - let handle = self.spawn_thread(shutdown_tx, rt, id); + let spawned = self.spawn_task(Task::new(task, is_mandatory), rt); + (handle, spawned) + } - shared.worker_threads.insert(id, handle); + fn spawn_task(&self, task: Task, rt: &Handle) -> Result<(), SpawnError> { + let mut shared = self.inner.shared.lock(); + + if shared.shutdown { + // Shutdown the task: it's fine to shutdown this task (even if + // mandatory) because it was scheduled after the shutdown of the + // runtime began. + task.task.shutdown(); + + // no need to even push this task; it would never get picked up + return Err(SpawnError::ShuttingDown); + } + + shared.queue.push_back(task); + self.inner.metrics.inc_queue_depth(); + + if self.inner.metrics.num_idle_threads() == 0 { + // No threads are able to process the task. + + if self.inner.metrics.num_threads() == self.inner.thread_cap { + // At max number of threads + } else { + assert!(shared.shutdown_tx.is_some()); + let shutdown_tx = shared.shutdown_tx.clone(); + + if let Some(shutdown_tx) = shutdown_tx { + let id = shared.worker_thread_index; + + match self.spawn_thread(shutdown_tx, rt, id) { + Ok(handle) => { + self.inner.metrics.inc_num_threads(); + shared.worker_thread_index += 1; + shared.worker_threads.insert(id, handle); + } + Err(ref e) + if is_temporary_os_thread_error(e) + && self.inner.metrics.num_threads() > 0 => + { + // OS temporarily failed to spawn a new thread. + // The task will be picked up eventually by a currently + // busy thread. + } + Err(e) => { + // The OS refused to spawn the thread and there is no thread + // to pick up the task that has just been pushed to the queue. + return Err(SpawnError::NoThreads(e)); + } + } + } + } + } else { + // Notify an idle worker thread. The notification counter + // is used to count the needed amount of notifications + // exactly. Thread libraries may generate spurious + // wakeups, this counter is used to keep us in a + // consistent state. + self.inner.metrics.dec_num_idle_threads(); + shared.num_notify += 1; + self.inner.condvar.notify_one(); } Ok(()) @@ -229,7 +456,7 @@ impl Spawner { shutdown_tx: shutdown::Sender, rt: &Handle, id: usize, - ) -> thread::JoinHandle<()> { + ) -> std::io::Result<thread::JoinHandle<()>> { let mut builder = thread::Builder::new().name((self.inner.thread_name)()); if let Some(stack_size) = self.inner.stack_size { @@ -238,17 +465,37 @@ impl Spawner { let rt = rt.clone(); - builder - .spawn(move || { - // Only the reference should be moved into the closure - let _enter = crate::runtime::context::enter(rt.clone()); - rt.blocking_spawner.inner.run(id); - drop(shutdown_tx); - }) - .unwrap() + builder.spawn(move || { + // Only the reference should be moved into the closure + let _enter = rt.enter(); + rt.inner.blocking_spawner().inner.run(id); + drop(shutdown_tx); + }) + } +} + +cfg_metrics! { + impl Spawner { + pub(crate) fn num_threads(&self) -> usize { + self.inner.metrics.num_threads() + } + + pub(crate) fn num_idle_threads(&self) -> usize { + self.inner.metrics.num_idle_threads() + } + + pub(crate) fn queue_depth(&self) -> usize { + self.inner.metrics.queue_depth() + } } } +// Tells whether the error when spawning a thread is temporary. +#[inline] +fn is_temporary_os_thread_error(error: &std::io::Error) -> bool { + matches!(error.kind(), std::io::ErrorKind::WouldBlock) +} + impl Inner { fn run(&self, worker_thread_id: usize) { if let Some(f) = &self.after_start { @@ -261,6 +508,7 @@ impl Inner { 'main: loop { // BUSY while let Some(task) = shared.queue.pop_front() { + self.metrics.dec_queue_depth(); drop(shared); task.run(); @@ -268,7 +516,7 @@ impl Inner { } // IDLE - shared.num_idle += 1; + self.metrics.inc_num_idle_threads(); while !shared.shutdown { let lock_result = self.condvar.wait_timeout(shared, self.keep_alive).unwrap(); @@ -302,8 +550,10 @@ impl Inner { if shared.shutdown { // Drain the queue while let Some(task) = shared.queue.pop_front() { + self.metrics.dec_queue_depth(); drop(shared); - task.shutdown(); + + task.shutdown_or_run_if_mandatory(); shared = self.shared.lock(); } @@ -311,7 +561,7 @@ impl Inner { // Work was produced, and we "took" it (by decrementing num_notify). // This means that num_idle was decremented once for our wakeup. // But, since we are exiting, we need to "undo" that, as we'll stay idle. - shared.num_idle += 1; + self.metrics.inc_num_idle_threads(); // NOTE: Technically we should also do num_notify++ and notify again, // but since we're shutting down anyway, that won't be necessary. break; @@ -319,17 +569,17 @@ impl Inner { } // Thread exit - shared.num_th -= 1; + self.metrics.dec_num_threads(); // num_idle should now be tracked exactly, panic // with a descriptive message if it is not the // case. - shared.num_idle = shared - .num_idle - .checked_sub(1) - .expect("num_idle underflowed on thread exit"); + let prev_idle = self.metrics.dec_num_idle_threads(); + if prev_idle < self.metrics.num_idle_threads() { + panic!("num_idle_threads underflowed on thread exit") + } - if shared.shutdown && shared.num_th == 0 { + if shared.shutdown && self.metrics.num_threads() == 0 { self.condvar.notify_one(); } diff --git a/vendor/tokio/src/runtime/blocking/schedule.rs b/vendor/tokio/src/runtime/blocking/schedule.rs index 4e044ab29..edf775be8 100644 --- a/vendor/tokio/src/runtime/blocking/schedule.rs +++ b/vendor/tokio/src/runtime/blocking/schedule.rs @@ -1,20 +1,52 @@ +#[cfg(feature = "test-util")] +use crate::runtime::scheduler; use crate::runtime::task::{self, Task}; +use crate::runtime::Handle; -/// `task::Schedule` implementation that does nothing. This is unique to the -/// blocking scheduler as tasks scheduled are not really futures but blocking -/// operations. +/// `task::Schedule` implementation that does nothing (except some bookkeeping +/// in test-util builds). This is unique to the blocking scheduler as tasks +/// scheduled are not really futures but blocking operations. /// /// We avoid storing the task by forgetting it in `bind` and re-materializing it -/// in `release. -pub(crate) struct NoopSchedule; +/// in `release`. +pub(crate) struct BlockingSchedule { + #[cfg(feature = "test-util")] + handle: Handle, +} -impl task::Schedule for NoopSchedule { - fn bind(_task: Task<Self>) -> NoopSchedule { - // Do nothing w/ the task - NoopSchedule +impl BlockingSchedule { + #[cfg_attr(not(feature = "test-util"), allow(unused_variables))] + pub(crate) fn new(handle: &Handle) -> Self { + #[cfg(feature = "test-util")] + { + match &handle.inner { + scheduler::Handle::CurrentThread(handle) => { + handle.driver.clock.inhibit_auto_advance(); + } + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + scheduler::Handle::MultiThread(_) => {} + } + } + BlockingSchedule { + #[cfg(feature = "test-util")] + handle: handle.clone(), + } } +} +impl task::Schedule for BlockingSchedule { fn release(&self, _task: &Task<Self>) -> Option<Task<Self>> { + #[cfg(feature = "test-util")] + { + match &self.handle.inner { + scheduler::Handle::CurrentThread(handle) => { + handle.driver.clock.allow_auto_advance(); + handle.driver.unpark(); + } + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + scheduler::Handle::MultiThread(_) => {} + } + } None } diff --git a/vendor/tokio/src/runtime/blocking/shutdown.rs b/vendor/tokio/src/runtime/blocking/shutdown.rs index 0cf22859b..fe5abae07 100644 --- a/vendor/tokio/src/runtime/blocking/shutdown.rs +++ b/vendor/tokio/src/runtime/blocking/shutdown.rs @@ -10,7 +10,7 @@ use std::time::Duration; #[derive(Debug, Clone)] pub(super) struct Sender { - tx: Arc<oneshot::Sender<()>>, + _tx: Arc<oneshot::Sender<()>>, } #[derive(Debug)] @@ -20,7 +20,7 @@ pub(super) struct Receiver { pub(super) fn channel() -> (Sender, Receiver) { let (tx, rx) = oneshot::channel(); - let tx = Sender { tx: Arc::new(tx) }; + let tx = Sender { _tx: Arc::new(tx) }; let rx = Receiver { rx }; (tx, rx) @@ -35,13 +35,13 @@ impl Receiver { /// /// If the timeout has elapsed, it returns `false`, otherwise it returns `true`. pub(crate) fn wait(&mut self, timeout: Option<Duration>) -> bool { - use crate::runtime::enter::try_enter; + use crate::runtime::context::try_enter_blocking_region; if timeout == Some(Duration::from_nanos(0)) { return false; } - let mut e = match try_enter(false) { + let mut e = match try_enter_blocking_region() { Some(enter) => enter, _ => { if std::thread::panicking() { diff --git a/vendor/tokio/src/runtime/blocking/task.rs b/vendor/tokio/src/runtime/blocking/task.rs index ee2d8d6d6..c44617540 100644 --- a/vendor/tokio/src/runtime/blocking/task.rs +++ b/vendor/tokio/src/runtime/blocking/task.rs @@ -2,13 +2,13 @@ use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; -/// Converts a function to a future that completes on poll +/// Converts a function to a future that completes on poll. pub(crate) struct BlockingTask<T> { func: Option<T>, } impl<T> BlockingTask<T> { - /// Initializes a new blocking task from the given function + /// Initializes a new blocking task from the given function. pub(crate) fn new(func: T) -> BlockingTask<T> { BlockingTask { func: Some(func) } } @@ -37,7 +37,7 @@ where // currently goes through Task::poll(), and so is subject to budgeting. That isn't really // what we want; a blocking task may itself want to run tasks (it might be a Worker!), so // we want it to start without any budgeting. - crate::coop::stop(); + crate::runtime::coop::stop(); Poll::Ready(func()) } diff --git a/vendor/tokio/src/runtime/builder.rs b/vendor/tokio/src/runtime/builder.rs index 51bf8c843..af9e0e172 100644 --- a/vendor/tokio/src/runtime/builder.rs +++ b/vendor/tokio/src/runtime/builder.rs @@ -1,5 +1,6 @@ use crate::runtime::handle::Handle; -use crate::runtime::{blocking, driver, Callback, Runtime, Spawner}; +use crate::runtime::{blocking, driver, Callback, HistogramBuilder, Runtime}; +use crate::util::rand::{RngSeed, RngSeedGenerator}; use std::fmt; use std::io; @@ -43,6 +44,7 @@ pub struct Builder { /// Whether or not to enable the I/O driver enable_io: bool, + nevents: usize, /// Whether or not to enable the time driver enable_time: bool, @@ -70,15 +72,132 @@ pub struct Builder { /// To run before each worker thread stops pub(super) before_stop: Option<Callback>, + /// To run before each worker thread is parked. + pub(super) before_park: Option<Callback>, + + /// To run after each thread is unparked. + pub(super) after_unpark: Option<Callback>, + /// Customizable keep alive timeout for BlockingPool pub(super) keep_alive: Option<Duration>, + + /// How many ticks before pulling a task from the global/remote queue? + /// + /// When `None`, the value is unspecified and behavior details are left to + /// the scheduler. Each scheduler flavor could choose to either pick its own + /// default value or use some other strategy to decide when to poll from the + /// global queue. For example, the multi-threaded scheduler uses a + /// self-tuning strategy based on mean task poll times. + pub(super) global_queue_interval: Option<u32>, + + /// How many ticks before yielding to the driver for timer and I/O events? + pub(super) event_interval: u32, + + /// When true, the multi-threade scheduler LIFO slot should not be used. + /// + /// This option should only be exposed as unstable. + pub(super) disable_lifo_slot: bool, + + /// Specify a random number generator seed to provide deterministic results + pub(super) seed_generator: RngSeedGenerator, + + /// When true, enables task poll count histogram instrumentation. + pub(super) metrics_poll_count_histogram_enable: bool, + + /// Configures the task poll count histogram + pub(super) metrics_poll_count_histogram: HistogramBuilder, + + #[cfg(tokio_unstable)] + pub(super) unhandled_panic: UnhandledPanic, +} + +cfg_unstable! { + /// How the runtime should respond to unhandled panics. + /// + /// Instances of `UnhandledPanic` are passed to `Builder::unhandled_panic` + /// to configure the runtime behavior when a spawned task panics. + /// + /// See [`Builder::unhandled_panic`] for more details. + #[derive(Debug, Clone)] + #[non_exhaustive] + pub enum UnhandledPanic { + /// The runtime should ignore panics on spawned tasks. + /// + /// The panic is forwarded to the task's [`JoinHandle`] and all spawned + /// tasks continue running normally. + /// + /// This is the default behavior. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::{self, UnhandledPanic}; + /// + /// # pub fn main() { + /// let rt = runtime::Builder::new_current_thread() + /// .unhandled_panic(UnhandledPanic::Ignore) + /// .build() + /// .unwrap(); + /// + /// let task1 = rt.spawn(async { panic!("boom"); }); + /// let task2 = rt.spawn(async { + /// // This task completes normally + /// "done" + /// }); + /// + /// rt.block_on(async { + /// // The panic on the first task is forwarded to the `JoinHandle` + /// assert!(task1.await.is_err()); + /// + /// // The second task completes normally + /// assert!(task2.await.is_ok()); + /// }) + /// # } + /// ``` + /// + /// [`JoinHandle`]: struct@crate::task::JoinHandle + Ignore, + + /// The runtime should immediately shutdown if a spawned task panics. + /// + /// The runtime will immediately shutdown even if the panicked task's + /// [`JoinHandle`] is still available. All further spawned tasks will be + /// immediately dropped and call to [`Runtime::block_on`] will panic. + /// + /// # Examples + /// + /// ```should_panic + /// use tokio::runtime::{self, UnhandledPanic}; + /// + /// # pub fn main() { + /// let rt = runtime::Builder::new_current_thread() + /// .unhandled_panic(UnhandledPanic::ShutdownRuntime) + /// .build() + /// .unwrap(); + /// + /// rt.spawn(async { panic!("boom"); }); + /// rt.spawn(async { + /// // This task never completes. + /// }); + /// + /// rt.block_on(async { + /// // Do some work + /// # loop { tokio::task::yield_now().await; } + /// }) + /// # } + /// ``` + /// + /// [`JoinHandle`]: struct@crate::task::JoinHandle + ShutdownRuntime, + } } pub(crate) type ThreadNameFn = std::sync::Arc<dyn Fn() -> String + Send + Sync + 'static>; +#[derive(Clone, Copy)] pub(crate) enum Kind { CurrentThread, - #[cfg(feature = "rt-multi-thread")] + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] MultiThread, } @@ -92,28 +211,38 @@ impl Builder { /// /// [`LocalSet`]: crate::task::LocalSet pub fn new_current_thread() -> Builder { - Builder::new(Kind::CurrentThread) + #[cfg(loom)] + const EVENT_INTERVAL: u32 = 4; + // The number `61` is fairly arbitrary. I believe this value was copied from golang. + #[cfg(not(loom))] + const EVENT_INTERVAL: u32 = 61; + + Builder::new(Kind::CurrentThread, EVENT_INTERVAL) } - /// Returns a new builder with the multi thread scheduler selected. - /// - /// Configuration methods can be chained on the return value. - #[cfg(feature = "rt-multi-thread")] - #[cfg_attr(docsrs, doc(cfg(feature = "rt-multi-thread")))] - pub fn new_multi_thread() -> Builder { - Builder::new(Kind::MultiThread) + cfg_not_wasi! { + /// Returns a new builder with the multi thread scheduler selected. + /// + /// Configuration methods can be chained on the return value. + #[cfg(feature = "rt-multi-thread")] + #[cfg_attr(docsrs, doc(cfg(feature = "rt-multi-thread")))] + pub fn new_multi_thread() -> Builder { + // The number `61` is fairly arbitrary. I believe this value was copied from golang. + Builder::new(Kind::MultiThread, 61) + } } /// Returns a new runtime builder initialized with default configuration /// values. /// /// Configuration methods can be chained on the return value. - pub(crate) fn new(kind: Kind) -> Builder { + pub(crate) fn new(kind: Kind, event_interval: u32) -> Builder { Builder { kind, // I/O defaults to "off" enable_io: false, + nevents: 1024, // Time defaults to "off" enable_time: false, @@ -121,6 +250,7 @@ impl Builder { // The clock starts not-paused start_paused: false, + // Read from environment variable first in multi-threaded mode. // Default to lazy auto-detection (one thread per CPU core) worker_threads: None, @@ -135,8 +265,26 @@ impl Builder { // No worker thread callbacks after_start: None, before_stop: None, + before_park: None, + after_unpark: None, keep_alive: None, + + // Defaults for these values depend on the scheduler kind, so we get them + // as parameters. + global_queue_interval: None, + event_interval, + + seed_generator: RngSeedGenerator::new(RngSeed::new()), + + #[cfg(tokio_unstable)] + unhandled_panic: UnhandledPanic::Ignore, + + metrics_poll_count_histogram_enable: false, + + metrics_poll_count_histogram: Default::default(), + + disable_lifo_slot: false, } } @@ -157,7 +305,11 @@ impl Builder { /// .unwrap(); /// ``` pub fn enable_all(&mut self) -> &mut Self { - #[cfg(any(feature = "net", feature = "process", all(unix, feature = "signal")))] + #[cfg(any( + feature = "net", + all(unix, feature = "process"), + all(unix, feature = "signal") + ))] self.enable_io(); #[cfg(feature = "time")] self.enable_time(); @@ -170,15 +322,13 @@ impl Builder { /// This can be any number above 0 though it is advised to keep this value /// on the smaller side. /// + /// This will override the value read from environment variable `TOKIO_WORKER_THREADS`. + /// /// # Default /// /// The default value is the number of cores available to the system. /// - /// # Panic - /// - /// When using the `current_thread` runtime this method will panic, since - /// those variants do not allow setting worker thread counts. - /// + /// When using the `current_thread` runtime this method has no effect. /// /// # Examples /// @@ -211,9 +361,10 @@ impl Builder { /// rt.block_on(async move {}); /// ``` /// - /// # Panic + /// # Panics /// /// This will panic if `val` is not larger than `0`. + #[track_caller] pub fn worker_threads(&mut self, val: usize) -> &mut Self { assert!(val > 0, "Worker threads cannot be set to 0"); self.worker_threads = Some(val); @@ -229,7 +380,7 @@ impl Builder { /// /// The default value is 512. /// - /// # Panic + /// # Panics /// /// This will panic if `val` is not larger than `0`. /// @@ -241,6 +392,7 @@ impl Builder { /// [`spawn_blocking`]: fn@crate::task::spawn_blocking /// [`worker_threads`]: Self::worker_threads /// [`thread_keep_alive`]: Self::thread_keep_alive + #[track_caller] #[cfg_attr(docsrs, doc(alias = "max_threads"))] pub fn max_blocking_threads(&mut self, val: usize) -> &mut Self { assert!(val > 0, "Max blocking threads cannot be set to 0"); @@ -278,7 +430,6 @@ impl Builder { /// ``` /// # use tokio::runtime; /// # use std::sync::atomic::{AtomicUsize, Ordering}; - /// /// # pub fn main() { /// let rt = runtime::Builder::new_multi_thread() /// .thread_name_fn(|| { @@ -330,7 +481,6 @@ impl Builder { /// /// ``` /// # use tokio::runtime; - /// /// # pub fn main() { /// let runtime = runtime::Builder::new_multi_thread() /// .on_thread_start(|| { @@ -356,7 +506,6 @@ impl Builder { /// /// ``` /// # use tokio::runtime; - /// /// # pub fn main() { /// let runtime = runtime::Builder::new_multi_thread() /// .on_thread_stop(|| { @@ -374,6 +523,119 @@ impl Builder { self } + /// Executes function `f` just before a thread is parked (goes idle). + /// `f` is called within the Tokio context, so functions like [`tokio::spawn`](crate::spawn) + /// can be called, and may result in this thread being unparked immediately. + /// + /// This can be used to start work only when the executor is idle, or for bookkeeping + /// and monitoring purposes. + /// + /// Note: There can only be one park callback for a runtime; calling this function + /// more than once replaces the last callback defined, rather than adding to it. + /// + /// # Examples + /// + /// ## Multithreaded executor + /// ``` + /// # use std::sync::Arc; + /// # use std::sync::atomic::{AtomicBool, Ordering}; + /// # use tokio::runtime; + /// # use tokio::sync::Barrier; + /// # pub fn main() { + /// let once = AtomicBool::new(true); + /// let barrier = Arc::new(Barrier::new(2)); + /// + /// let runtime = runtime::Builder::new_multi_thread() + /// .worker_threads(1) + /// .on_thread_park({ + /// let barrier = barrier.clone(); + /// move || { + /// let barrier = barrier.clone(); + /// if once.swap(false, Ordering::Relaxed) { + /// tokio::spawn(async move { barrier.wait().await; }); + /// } + /// } + /// }) + /// .build() + /// .unwrap(); + /// + /// runtime.block_on(async { + /// barrier.wait().await; + /// }) + /// # } + /// ``` + /// ## Current thread executor + /// ``` + /// # use std::sync::Arc; + /// # use std::sync::atomic::{AtomicBool, Ordering}; + /// # use tokio::runtime; + /// # use tokio::sync::Barrier; + /// # pub fn main() { + /// let once = AtomicBool::new(true); + /// let barrier = Arc::new(Barrier::new(2)); + /// + /// let runtime = runtime::Builder::new_current_thread() + /// .on_thread_park({ + /// let barrier = barrier.clone(); + /// move || { + /// let barrier = barrier.clone(); + /// if once.swap(false, Ordering::Relaxed) { + /// tokio::spawn(async move { barrier.wait().await; }); + /// } + /// } + /// }) + /// .build() + /// .unwrap(); + /// + /// runtime.block_on(async { + /// barrier.wait().await; + /// }) + /// # } + /// ``` + #[cfg(not(loom))] + pub fn on_thread_park<F>(&mut self, f: F) -> &mut Self + where + F: Fn() + Send + Sync + 'static, + { + self.before_park = Some(std::sync::Arc::new(f)); + self + } + + /// Executes function `f` just after a thread unparks (starts executing tasks). + /// + /// This is intended for bookkeeping and monitoring use cases; note that work + /// in this callback will increase latencies when the application has allowed one or + /// more runtime threads to go idle. + /// + /// Note: There can only be one unpark callback for a runtime; calling this function + /// more than once replaces the last callback defined, rather than adding to it. + /// + /// # Examples + /// + /// ``` + /// # use tokio::runtime; + /// # pub fn main() { + /// let runtime = runtime::Builder::new_multi_thread() + /// .on_thread_unpark(|| { + /// println!("thread unparking"); + /// }) + /// .build(); + /// + /// runtime.unwrap().block_on(async { + /// tokio::task::yield_now().await; + /// println!("Hello from Tokio!"); + /// }) + /// # } + /// ``` + #[cfg(not(loom))] + pub fn on_thread_unpark<F>(&mut self, f: F) -> &mut Self + where + F: Fn() + Send + Sync + 'static, + { + self.after_unpark = Some(std::sync::Arc::new(f)); + self + } + /// Creates the configured `Runtime`. /// /// The returned `Runtime` instance is ready to spawn tasks. @@ -391,8 +653,8 @@ impl Builder { /// ``` pub fn build(&mut self) -> io::Result<Runtime> { match &self.kind { - Kind::CurrentThread => self.build_basic_runtime(), - #[cfg(feature = "rt-multi-thread")] + Kind::CurrentThread => self.build_current_thread_runtime(), + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] Kind::MultiThread => self.build_threaded_runtime(), } } @@ -401,12 +663,13 @@ impl Builder { driver::Cfg { enable_pause_time: match self.kind { Kind::CurrentThread => true, - #[cfg(feature = "rt-multi-thread")] + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] Kind::MultiThread => false, }, enable_io: self.enable_io, enable_time: self.enable_time, start_paused: self.start_paused, + nevents: self.nevents, } } @@ -420,7 +683,6 @@ impl Builder { /// ``` /// # use tokio::runtime; /// # use std::time::Duration; - /// /// # pub fn main() { /// let rt = runtime::Builder::new_multi_thread() /// .thread_keep_alive(Duration::from_millis(100)) @@ -432,34 +694,385 @@ impl Builder { self } - fn build_basic_runtime(&mut self) -> io::Result<Runtime> { - use crate::runtime::{BasicScheduler, Kind}; + /// Sets the number of scheduler ticks after which the scheduler will poll the global + /// task queue. + /// + /// A scheduler "tick" roughly corresponds to one `poll` invocation on a task. + /// + /// By default the global queue interval is: + /// + /// * `31` for the current-thread scheduler. + /// * `61` for the multithreaded scheduler. + /// + /// Schedulers have a local queue of already-claimed tasks, and a global queue of incoming + /// tasks. Setting the interval to a smaller value increases the fairness of the scheduler, + /// at the cost of more synchronization overhead. That can be beneficial for prioritizing + /// getting started on new work, especially if tasks frequently yield rather than complete + /// or await on further I/O. Conversely, a higher value prioritizes existing work, and + /// is a good choice when most tasks quickly complete polling. + /// + /// # Examples + /// + /// ``` + /// # use tokio::runtime; + /// # pub fn main() { + /// let rt = runtime::Builder::new_multi_thread() + /// .global_queue_interval(31) + /// .build(); + /// # } + /// ``` + pub fn global_queue_interval(&mut self, val: u32) -> &mut Self { + self.global_queue_interval = Some(val); + self + } - let (driver, resources) = driver::Driver::new(self.get_cfg())?; + /// Sets the number of scheduler ticks after which the scheduler will poll for + /// external events (timers, I/O, and so on). + /// + /// A scheduler "tick" roughly corresponds to one `poll` invocation on a task. + /// + /// By default, the event interval is `61` for all scheduler types. + /// + /// Setting the event interval determines the effective "priority" of delivering + /// these external events (which may wake up additional tasks), compared to + /// executing tasks that are currently ready to run. A smaller value is useful + /// when tasks frequently spend a long time in polling, or frequently yield, + /// which can result in overly long delays picking up I/O events. Conversely, + /// picking up new events requires extra synchronization and syscall overhead, + /// so if tasks generally complete their polling quickly, a higher event interval + /// will minimize that overhead while still keeping the scheduler responsive to + /// events. + /// + /// # Examples + /// + /// ``` + /// # use tokio::runtime; + /// # pub fn main() { + /// let rt = runtime::Builder::new_multi_thread() + /// .event_interval(31) + /// .build(); + /// # } + /// ``` + pub fn event_interval(&mut self, val: u32) -> &mut Self { + self.event_interval = val; + self + } - // And now put a single-threaded scheduler on top of the timer. When - // there are no futures ready to do something, it'll let the timer or - // the reactor to generate some new stimuli for the futures to continue - // in their life. - let scheduler = BasicScheduler::new(driver); - let spawner = Spawner::Basic(scheduler.spawner().clone()); + cfg_unstable! { + /// Configure how the runtime responds to an unhandled panic on a + /// spawned task. + /// + /// By default, an unhandled panic (i.e. a panic not caught by + /// [`std::panic::catch_unwind`]) has no impact on the runtime's + /// execution. The panic is error value is forwarded to the task's + /// [`JoinHandle`] and all other spawned tasks continue running. + /// + /// The `unhandled_panic` option enables configuring this behavior. + /// + /// * `UnhandledPanic::Ignore` is the default behavior. Panics on + /// spawned tasks have no impact on the runtime's execution. + /// * `UnhandledPanic::ShutdownRuntime` will force the runtime to + /// shutdown immediately when a spawned task panics even if that + /// task's `JoinHandle` has not been dropped. All other spawned tasks + /// will immediately terminate and further calls to + /// [`Runtime::block_on`] will panic. + /// + /// # Unstable + /// + /// This option is currently unstable and its implementation is + /// incomplete. The API may change or be removed in the future. See + /// tokio-rs/tokio#4516 for more details. + /// + /// # Examples + /// + /// The following demonstrates a runtime configured to shutdown on + /// panic. The first spawned task panics and results in the runtime + /// shutting down. The second spawned task never has a chance to + /// execute. The call to `block_on` will panic due to the runtime being + /// forcibly shutdown. + /// + /// ```should_panic + /// use tokio::runtime::{self, UnhandledPanic}; + /// + /// # pub fn main() { + /// let rt = runtime::Builder::new_current_thread() + /// .unhandled_panic(UnhandledPanic::ShutdownRuntime) + /// .build() + /// .unwrap(); + /// + /// rt.spawn(async { panic!("boom"); }); + /// rt.spawn(async { + /// // This task never completes. + /// }); + /// + /// rt.block_on(async { + /// // Do some work + /// # loop { tokio::task::yield_now().await; } + /// }) + /// # } + /// ``` + /// + /// [`JoinHandle`]: struct@crate::task::JoinHandle + pub fn unhandled_panic(&mut self, behavior: UnhandledPanic) -> &mut Self { + self.unhandled_panic = behavior; + self + } + + /// Disables the LIFO task scheduler heuristic. + /// + /// The multi-threaded scheduler includes a heuristic for optimizing + /// message-passing patterns. This heuristic results in the **last** + /// scheduled task being polled first. + /// + /// To implement this heuristic, each worker thread has a slot which + /// holds the task that should be polled next. However, this slot cannot + /// be stolen by other worker threads, which can result in lower total + /// throughput when tasks tend to have longer poll times. + /// + /// This configuration option will disable this heuristic resulting in + /// all scheduled tasks being pushed into the worker-local queue, which + /// is stealable. + /// + /// Consider trying this option when the task "scheduled" time is high + /// but the runtime is underutilized. Use tokio-rs/tokio-metrics to + /// collect this data. + /// + /// # Unstable + /// + /// This configuration option is considered a workaround for the LIFO + /// slot not being stealable. When the slot becomes stealable, we will + /// revisit whether or not this option is necessary. See + /// tokio-rs/tokio#4941. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime; + /// + /// let rt = runtime::Builder::new_multi_thread() + /// .disable_lifo_slot() + /// .build() + /// .unwrap(); + /// ``` + pub fn disable_lifo_slot(&mut self) -> &mut Self { + self.disable_lifo_slot = true; + self + } + + /// Specifies the random number generation seed to use within all + /// threads associated with the runtime being built. + /// + /// This option is intended to make certain parts of the runtime + /// deterministic (e.g. the [`tokio::select!`] macro). In the case of + /// [`tokio::select!`] it will ensure that the order that branches are + /// polled is deterministic. + /// + /// In addition to the code specifying `rng_seed` and interacting with + /// the runtime, the internals of Tokio and the Rust compiler may affect + /// the sequences of random numbers. In order to ensure repeatable + /// results, the version of Tokio, the versions of all other + /// dependencies that interact with Tokio, and the Rust compiler version + /// should also all remain constant. + /// + /// # Examples + /// + /// ``` + /// # use tokio::runtime::{self, RngSeed}; + /// # pub fn main() { + /// let seed = RngSeed::from_bytes(b"place your seed here"); + /// let rt = runtime::Builder::new_current_thread() + /// .rng_seed(seed) + /// .build(); + /// # } + /// ``` + /// + /// [`tokio::select!`]: crate::select + pub fn rng_seed(&mut self, seed: RngSeed) -> &mut Self { + self.seed_generator = RngSeedGenerator::new(seed); + self + } + } + + cfg_metrics! { + /// Enables tracking the distribution of task poll times. + /// + /// Task poll times are not instrumented by default as doing so requires + /// calling [`Instant::now()`] twice per task poll, which could add + /// measurable overhead. Use the [`Handle::metrics()`] to access the + /// metrics data. + /// + /// The histogram uses fixed bucket sizes. In other words, the histogram + /// buckets are not dynamic based on input values. Use the + /// `metrics_poll_count_histogram_` builder methods to configure the + /// histogram details. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime; + /// + /// let rt = runtime::Builder::new_multi_thread() + /// .enable_metrics_poll_count_histogram() + /// .build() + /// .unwrap(); + /// # // Test default values here + /// # fn us(n: u64) -> std::time::Duration { std::time::Duration::from_micros(n) } + /// # let m = rt.handle().metrics(); + /// # assert_eq!(m.poll_count_histogram_num_buckets(), 10); + /// # assert_eq!(m.poll_count_histogram_bucket_range(0), us(0)..us(100)); + /// # assert_eq!(m.poll_count_histogram_bucket_range(1), us(100)..us(200)); + /// ``` + /// + /// [`Handle::metrics()`]: crate::runtime::Handle::metrics + /// [`Instant::now()`]: std::time::Instant::now + pub fn enable_metrics_poll_count_histogram(&mut self) -> &mut Self { + self.metrics_poll_count_histogram_enable = true; + self + } + + /// Sets the histogram scale for tracking the distribution of task poll + /// times. + /// + /// Tracking the distribution of task poll times can be done using a + /// linear or log scale. When using linear scale, each histogram bucket + /// will represent the same range of poll times. When using log scale, + /// each histogram bucket will cover a range twice as big as the + /// previous bucket. + /// + /// **Default:** linear scale. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::{self, HistogramScale}; + /// + /// let rt = runtime::Builder::new_multi_thread() + /// .enable_metrics_poll_count_histogram() + /// .metrics_poll_count_histogram_scale(HistogramScale::Log) + /// .build() + /// .unwrap(); + /// ``` + pub fn metrics_poll_count_histogram_scale(&mut self, histogram_scale: crate::runtime::HistogramScale) -> &mut Self { + self.metrics_poll_count_histogram.scale = histogram_scale; + self + } + + /// Sets the histogram resolution for tracking the distribution of task + /// poll times. + /// + /// The resolution is the histogram's first bucket's range. When using a + /// linear histogram scale, each bucket will cover the same range. When + /// using a log scale, each bucket will cover a range twice as big as + /// the previous bucket. In the log case, the resolution represents the + /// smallest bucket range. + /// + /// Note that, when using log scale, the resolution is rounded up to the + /// nearest power of 2 in nanoseconds. + /// + /// **Default:** 100 microseconds. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime; + /// use std::time::Duration; + /// + /// let rt = runtime::Builder::new_multi_thread() + /// .enable_metrics_poll_count_histogram() + /// .metrics_poll_count_histogram_resolution(Duration::from_micros(100)) + /// .build() + /// .unwrap(); + /// ``` + pub fn metrics_poll_count_histogram_resolution(&mut self, resolution: Duration) -> &mut Self { + assert!(resolution > Duration::from_secs(0)); + // Sanity check the argument and also make the cast below safe. + assert!(resolution <= Duration::from_secs(1)); + + let resolution = resolution.as_nanos() as u64; + self.metrics_poll_count_histogram.resolution = resolution; + self + } + + /// Sets the number of buckets for the histogram tracking the + /// distribution of task poll times. + /// + /// The last bucket tracks all greater values that fall out of other + /// ranges. So, configuring the histogram using a linear scale, + /// resolution of 50ms, and 10 buckets, the 10th bucket will track task + /// polls that take more than 450ms to complete. + /// + /// **Default:** 10 + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime; + /// + /// let rt = runtime::Builder::new_multi_thread() + /// .enable_metrics_poll_count_histogram() + /// .metrics_poll_count_histogram_buckets(15) + /// .build() + /// .unwrap(); + /// ``` + pub fn metrics_poll_count_histogram_buckets(&mut self, buckets: usize) -> &mut Self { + self.metrics_poll_count_histogram.num_buckets = buckets; + self + } + } + + fn build_current_thread_runtime(&mut self) -> io::Result<Runtime> { + use crate::runtime::scheduler::{self, CurrentThread}; + use crate::runtime::{runtime::Scheduler, Config}; + + let (driver, driver_handle) = driver::Driver::new(self.get_cfg())?; // Blocking pool let blocking_pool = blocking::create_blocking_pool(self, self.max_blocking_threads); let blocking_spawner = blocking_pool.spawner().clone(); - Ok(Runtime { - kind: Kind::CurrentThread(scheduler), - handle: Handle { - spawner, - io_handle: resources.io_handle, - time_handle: resources.time_handle, - signal_handle: resources.signal_handle, - clock: resources.clock, - blocking_spawner, + // Generate a rng seed for this runtime. + let seed_generator_1 = self.seed_generator.next_generator(); + let seed_generator_2 = self.seed_generator.next_generator(); + + // And now put a single-threaded scheduler on top of the timer. When + // there are no futures ready to do something, it'll let the timer or + // the reactor to generate some new stimuli for the futures to continue + // in their life. + let (scheduler, handle) = CurrentThread::new( + driver, + driver_handle, + blocking_spawner, + seed_generator_2, + Config { + before_park: self.before_park.clone(), + after_unpark: self.after_unpark.clone(), + global_queue_interval: self.global_queue_interval, + event_interval: self.event_interval, + #[cfg(tokio_unstable)] + unhandled_panic: self.unhandled_panic.clone(), + disable_lifo_slot: self.disable_lifo_slot, + seed_generator: seed_generator_1, + metrics_poll_count_histogram: self.metrics_poll_count_histogram_builder(), }, + ); + + let handle = Handle { + inner: scheduler::Handle::CurrentThread(handle), + }; + + Ok(Runtime::from_parts( + Scheduler::CurrentThread(scheduler), + handle, blocking_pool, - }) + )) + } + + fn metrics_poll_count_histogram_builder(&self) -> Option<HistogramBuilder> { + if self.metrics_poll_count_histogram_enable { + Some(self.metrics_poll_count_histogram.clone()) + } else { + None + } } } @@ -484,6 +1097,25 @@ cfg_io_driver! { self.enable_io = true; self } + + /// Enables the I/O driver and configures the max number of events to be + /// processed per tick. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime; + /// + /// let rt = runtime::Builder::new_current_thread() + /// .enable_io() + /// .max_io_events_per_tick(1024) + /// .build() + /// .unwrap(); + /// ``` + pub fn max_io_events_per_tick(&mut self, capacity: usize) -> &mut Self { + self.nevents = capacity; + self + } } } @@ -539,39 +1171,48 @@ cfg_rt_multi_thread! { impl Builder { fn build_threaded_runtime(&mut self) -> io::Result<Runtime> { use crate::loom::sys::num_cpus; - use crate::runtime::{Kind, ThreadPool}; - use crate::runtime::park::Parker; + use crate::runtime::{Config, runtime::Scheduler}; + use crate::runtime::scheduler::{self, MultiThread}; let core_threads = self.worker_threads.unwrap_or_else(num_cpus); - let (driver, resources) = driver::Driver::new(self.get_cfg())?; - - let (scheduler, launch) = ThreadPool::new(core_threads, Parker::new(driver)); - let spawner = Spawner::ThreadPool(scheduler.spawner().clone()); + let (driver, driver_handle) = driver::Driver::new(self.get_cfg())?; // Create the blocking pool - let blocking_pool = blocking::create_blocking_pool(self, self.max_blocking_threads + core_threads); + let blocking_pool = + blocking::create_blocking_pool(self, self.max_blocking_threads + core_threads); let blocking_spawner = blocking_pool.spawner().clone(); - // Create the runtime handle - let handle = Handle { - spawner, - io_handle: resources.io_handle, - time_handle: resources.time_handle, - signal_handle: resources.signal_handle, - clock: resources.clock, + // Generate a rng seed for this runtime. + let seed_generator_1 = self.seed_generator.next_generator(); + let seed_generator_2 = self.seed_generator.next_generator(); + + let (scheduler, handle, launch) = MultiThread::new( + core_threads, + driver, + driver_handle, blocking_spawner, - }; + seed_generator_2, + Config { + before_park: self.before_park.clone(), + after_unpark: self.after_unpark.clone(), + global_queue_interval: self.global_queue_interval, + event_interval: self.event_interval, + #[cfg(tokio_unstable)] + unhandled_panic: self.unhandled_panic.clone(), + disable_lifo_slot: self.disable_lifo_slot, + seed_generator: seed_generator_1, + metrics_poll_count_histogram: self.metrics_poll_count_histogram_builder(), + }, + ); + + let handle = Handle { inner: scheduler::Handle::MultiThread(handle) }; // Spawn the thread pool workers - let _enter = crate::runtime::context::enter(handle.clone()); + let _enter = handle.enter(); launch.launch(); - Ok(Runtime { - kind: Kind::ThreadPool(scheduler), - handle, - blocking_pool, - }) + Ok(Runtime::from_parts(Scheduler::MultiThread(scheduler), handle, blocking_pool)) } } } @@ -587,7 +1228,9 @@ impl fmt::Debug for Builder { ) .field("thread_stack_size", &self.thread_stack_size) .field("after_start", &self.after_start.as_ref().map(|_| "...")) - .field("before_stop", &self.after_start.as_ref().map(|_| "...")) + .field("before_stop", &self.before_stop.as_ref().map(|_| "...")) + .field("before_park", &self.before_park.as_ref().map(|_| "...")) + .field("after_unpark", &self.after_unpark.as_ref().map(|_| "...")) .finish() } } diff --git a/vendor/tokio/src/runtime/config.rs b/vendor/tokio/src/runtime/config.rs new file mode 100644 index 000000000..c42e4fe5a --- /dev/null +++ b/vendor/tokio/src/runtime/config.rs @@ -0,0 +1,37 @@ +#![cfg_attr(any(not(feature = "full"), tokio_wasm), allow(dead_code))] +use crate::runtime::Callback; +use crate::util::RngSeedGenerator; + +pub(crate) struct Config { + /// How many ticks before pulling a task from the global/remote queue? + pub(crate) global_queue_interval: Option<u32>, + + /// How many ticks before yielding to the driver for timer and I/O events? + pub(crate) event_interval: u32, + + /// Callback for a worker parking itself + pub(crate) before_park: Option<Callback>, + + /// Callback for a worker unparking itself + pub(crate) after_unpark: Option<Callback>, + + /// The multi-threaded scheduler includes a per-worker LIFO slot used to + /// store the last scheduled task. This can improve certain usage patterns, + /// especially message passing between tasks. However, this LIFO slot is not + /// currently stealable. + /// + /// Eventually, the LIFO slot **will** become stealable, however as a + /// stop-gap, this unstable option lets users disable the LIFO task. + pub(crate) disable_lifo_slot: bool, + + /// Random number generator seed to configure runtimes to act in a + /// deterministic way. + pub(crate) seed_generator: RngSeedGenerator, + + /// How to build poll time histograms + pub(crate) metrics_poll_count_histogram: Option<crate::runtime::HistogramBuilder>, + + #[cfg(tokio_unstable)] + /// How to respond to unhandled task panics. + pub(crate) unhandled_panic: crate::runtime::UnhandledPanic, +} diff --git a/vendor/tokio/src/runtime/context.rs b/vendor/tokio/src/runtime/context.rs index a727ed497..5943e9aa9 100644 --- a/vendor/tokio/src/runtime/context.rs +++ b/vendor/tokio/src/runtime/context.rs @@ -1,73 +1,191 @@ -//! Thread local runtime context -use crate::runtime::Handle; +use crate::loom::thread::AccessError; +use crate::runtime::coop; -use std::cell::RefCell; +use std::cell::Cell; -thread_local! { - static CONTEXT: RefCell<Option<Handle>> = RefCell::new(None) -} +#[cfg(any(feature = "rt", feature = "macros"))] +use crate::util::rand::FastRand; -pub(crate) fn current() -> Option<Handle> { - CONTEXT.with(|ctx| ctx.borrow().clone()) -} +cfg_rt! { + mod blocking; + pub(crate) use blocking::{disallow_block_in_place, try_enter_blocking_region, BlockingRegionGuard}; -cfg_io_driver! { - pub(crate) fn io_handle() -> crate::runtime::driver::IoHandle { - CONTEXT.with(|ctx| { - let ctx = ctx.borrow(); - ctx.as_ref().expect(crate::util::error::CONTEXT_MISSING_ERROR).io_handle.clone() - }) + mod current; + pub(crate) use current::{with_current, try_set_current, SetCurrentGuard}; + + mod runtime; + pub(crate) use runtime::{EnterRuntime, enter_runtime}; + + mod scoped; + use scoped::Scoped; + + use crate::runtime::{scheduler, task::Id}; + + use std::task::Waker; + + cfg_taskdump! { + use crate::runtime::task::trace; } } -cfg_signal_internal! { - #[cfg(unix)] - pub(crate) fn signal_handle() -> crate::runtime::driver::SignalHandle { - CONTEXT.with(|ctx| { - let ctx = ctx.borrow(); - ctx.as_ref().expect(crate::util::error::CONTEXT_MISSING_ERROR).signal_handle.clone() - }) - } +cfg_rt_multi_thread! { + mod runtime_mt; + pub(crate) use runtime_mt::{current_enter_context, exit_runtime}; } -cfg_time! { - pub(crate) fn time_handle() -> crate::runtime::driver::TimeHandle { - CONTEXT.with(|ctx| { - let ctx = ctx.borrow(); - ctx.as_ref().expect(crate::util::error::CONTEXT_MISSING_ERROR).time_handle.clone() - }) - } +struct Context { + /// Uniquely identifies the current thread + #[cfg(feature = "rt")] + thread_id: Cell<Option<ThreadId>>, - cfg_test_util! { - pub(crate) fn clock() -> Option<crate::runtime::driver::Clock> { - CONTEXT.with(|ctx| (*ctx.borrow()).as_ref().map(|ctx| ctx.clock.clone())) - } - } + /// Handle to the runtime scheduler running on the current thread. + #[cfg(feature = "rt")] + current: current::HandleCell, + + /// Handle to the scheduler's internal "context" + #[cfg(feature = "rt")] + scheduler: Scoped<scheduler::Context>, + + #[cfg(feature = "rt")] + current_task_id: Cell<Option<Id>>, + + /// Tracks if the current thread is currently driving a runtime. + /// Note, that if this is set to "entered", the current scheduler + /// handle may not reference the runtime currently executing. This + /// is because other runtime handles may be set to current from + /// within a runtime. + #[cfg(feature = "rt")] + runtime: Cell<EnterRuntime>, + + #[cfg(any(feature = "rt", feature = "macros"))] + rng: Cell<Option<FastRand>>, + + /// Tracks the amount of "work" a task may still do before yielding back to + /// the sheduler + budget: Cell<coop::Budget>, + + #[cfg(all( + tokio_unstable, + tokio_taskdump, + feature = "rt", + target_os = "linux", + any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64") + ))] + trace: trace::Context, } -cfg_rt! { - pub(crate) fn spawn_handle() -> Option<crate::runtime::Spawner> { - CONTEXT.with(|ctx| (*ctx.borrow()).as_ref().map(|ctx| ctx.spawner.clone())) +tokio_thread_local! { + static CONTEXT: Context = const { + Context { + #[cfg(feature = "rt")] + thread_id: Cell::new(None), + + /// Tracks the current runtime handle to use when spawning, + /// accessing drivers, etc... + #[cfg(feature = "rt")] + current: current::HandleCell::new(), + + /// Tracks the current scheduler internal context + #[cfg(feature = "rt")] + scheduler: Scoped::new(), + + #[cfg(feature = "rt")] + current_task_id: Cell::new(None), + + /// Tracks if the current thread is currently driving a runtime. + /// Note, that if this is set to "entered", the current scheduler + /// handle may not reference the runtime currently executing. This + /// is because other runtime handles may be set to current from + /// within a runtime. + #[cfg(feature = "rt")] + runtime: Cell::new(EnterRuntime::NotEntered), + + #[cfg(any(feature = "rt", feature = "macros"))] + rng: Cell::new(None), + + budget: Cell::new(coop::Budget::unconstrained()), + + #[cfg(all( + tokio_unstable, + tokio_taskdump, + feature = "rt", + target_os = "linux", + any( + target_arch = "aarch64", + target_arch = "x86", + target_arch = "x86_64" + ) + ))] + trace: trace::Context::new(), + } } } -/// Set this [`Handle`] as the current active [`Handle`]. -/// -/// [`Handle`]: Handle -pub(crate) fn enter(new: Handle) -> EnterGuard { +#[cfg(any(feature = "macros", all(feature = "sync", feature = "rt")))] +pub(crate) fn thread_rng_n(n: u32) -> u32 { CONTEXT.with(|ctx| { - let old = ctx.borrow_mut().replace(new); - EnterGuard(old) + let mut rng = ctx.rng.get().unwrap_or_else(FastRand::new); + let ret = rng.fastrand_n(n); + ctx.rng.set(Some(rng)); + ret }) } -#[derive(Debug)] -pub(crate) struct EnterGuard(Option<Handle>); +pub(super) fn budget<R>(f: impl FnOnce(&Cell<coop::Budget>) -> R) -> Result<R, AccessError> { + CONTEXT.try_with(|ctx| f(&ctx.budget)) +} + +cfg_rt! { + use crate::runtime::ThreadId; + + pub(crate) fn thread_id() -> Result<ThreadId, AccessError> { + CONTEXT.try_with(|ctx| { + match ctx.thread_id.get() { + Some(id) => id, + None => { + let id = ThreadId::next(); + ctx.thread_id.set(Some(id)); + id + } + } + }) + } -impl Drop for EnterGuard { - fn drop(&mut self) { - CONTEXT.with(|ctx| { - *ctx.borrow_mut() = self.0.take(); + pub(crate) fn set_current_task_id(id: Option<Id>) -> Option<Id> { + CONTEXT.try_with(|ctx| ctx.current_task_id.replace(id)).unwrap_or(None) + } + + pub(crate) fn current_task_id() -> Option<Id> { + CONTEXT.try_with(|ctx| ctx.current_task_id.get()).unwrap_or(None) + } + + #[track_caller] + pub(crate) fn defer(waker: &Waker) { + with_scheduler(|maybe_scheduler| { + if let Some(scheduler) = maybe_scheduler { + scheduler.defer(waker); + } else { + // Called from outside of the runtime, immediately wake the + // task. + waker.wake_by_ref(); + } }); } + + pub(super) fn set_scheduler<R>(v: &scheduler::Context, f: impl FnOnce() -> R) -> R { + CONTEXT.with(|c| c.scheduler.set(v, f)) + } + + #[track_caller] + pub(super) fn with_scheduler<R>(f: impl FnOnce(Option<&scheduler::Context>) -> R) -> R { + CONTEXT.with(|c| c.scheduler.with(f)) + } + + cfg_taskdump! { + /// SAFETY: Callers of this function must ensure that trace frames always + /// form a valid linked list. + pub(crate) unsafe fn with_trace<R>(f: impl FnOnce(&trace::Context) -> R) -> Option<R> { + CONTEXT.try_with(|c| f(&c.trace)).ok() + } + } } diff --git a/vendor/tokio/src/runtime/context/blocking.rs b/vendor/tokio/src/runtime/context/blocking.rs new file mode 100644 index 000000000..8ae4f570e --- /dev/null +++ b/vendor/tokio/src/runtime/context/blocking.rs @@ -0,0 +1,121 @@ +use super::{EnterRuntime, CONTEXT}; + +use crate::loom::thread::AccessError; +use crate::util::markers::NotSendOrSync; + +use std::marker::PhantomData; +use std::time::Duration; + +/// Guard tracking that a caller has entered a blocking region. +#[must_use] +pub(crate) struct BlockingRegionGuard { + _p: PhantomData<NotSendOrSync>, +} + +pub(crate) struct DisallowBlockInPlaceGuard(bool); + +pub(crate) fn try_enter_blocking_region() -> Option<BlockingRegionGuard> { + CONTEXT + .try_with(|c| { + if c.runtime.get().is_entered() { + None + } else { + Some(BlockingRegionGuard::new()) + } + // If accessing the thread-local fails, the thread is terminating + // and thread-locals are being destroyed. Because we don't know if + // we are currently in a runtime or not, we default to being + // permissive. + }) + .unwrap_or_else(|_| Some(BlockingRegionGuard::new())) +} + +/// Disallows blocking in the current runtime context until the guard is dropped. +pub(crate) fn disallow_block_in_place() -> DisallowBlockInPlaceGuard { + let reset = CONTEXT.with(|c| { + if let EnterRuntime::Entered { + allow_block_in_place: true, + } = c.runtime.get() + { + c.runtime.set(EnterRuntime::Entered { + allow_block_in_place: false, + }); + true + } else { + false + } + }); + + DisallowBlockInPlaceGuard(reset) +} + +impl BlockingRegionGuard { + pub(super) fn new() -> BlockingRegionGuard { + BlockingRegionGuard { _p: PhantomData } + } + + /// Blocks the thread on the specified future, returning the value with + /// which that future completes. + pub(crate) fn block_on<F>(&mut self, f: F) -> Result<F::Output, AccessError> + where + F: std::future::Future, + { + use crate::runtime::park::CachedParkThread; + + let mut park = CachedParkThread::new(); + park.block_on(f) + } + + /// Blocks the thread on the specified future for **at most** `timeout` + /// + /// If the future completes before `timeout`, the result is returned. If + /// `timeout` elapses, then `Err` is returned. + pub(crate) fn block_on_timeout<F>(&mut self, f: F, timeout: Duration) -> Result<F::Output, ()> + where + F: std::future::Future, + { + use crate::runtime::park::CachedParkThread; + use std::task::Context; + use std::task::Poll::Ready; + use std::time::Instant; + + let mut park = CachedParkThread::new(); + let waker = park.waker().map_err(|_| ())?; + let mut cx = Context::from_waker(&waker); + + pin!(f); + let when = Instant::now() + timeout; + + loop { + if let Ready(v) = crate::runtime::coop::budget(|| f.as_mut().poll(&mut cx)) { + return Ok(v); + } + + let now = Instant::now(); + + if now >= when { + return Err(()); + } + + park.park_timeout(when - now); + } + } +} + +impl Drop for DisallowBlockInPlaceGuard { + fn drop(&mut self) { + if self.0 { + // XXX: Do we want some kind of assertion here, or is "best effort" okay? + CONTEXT.with(|c| { + if let EnterRuntime::Entered { + allow_block_in_place: false, + } = c.runtime.get() + { + c.runtime.set(EnterRuntime::Entered { + allow_block_in_place: true, + }); + } + }) + } + } +} diff --git a/vendor/tokio/src/runtime/context/current.rs b/vendor/tokio/src/runtime/context/current.rs new file mode 100644 index 000000000..c3dc5c899 --- /dev/null +++ b/vendor/tokio/src/runtime/context/current.rs @@ -0,0 +1,99 @@ +use super::{Context, CONTEXT}; + +use crate::runtime::{scheduler, TryCurrentError}; +use crate::util::markers::SyncNotSend; + +use std::cell::{Cell, RefCell}; +use std::marker::PhantomData; + +#[derive(Debug)] +#[must_use] +pub(crate) struct SetCurrentGuard { + // The previous handle + prev: Option<scheduler::Handle>, + + // The depth for this guard + depth: usize, + + // Don't let the type move across threads. + _p: PhantomData<SyncNotSend>, +} + +pub(super) struct HandleCell { + /// Current handle + handle: RefCell<Option<scheduler::Handle>>, + + /// Tracks the number of nested calls to `try_set_current`. + depth: Cell<usize>, +} + +/// Sets this [`Handle`] as the current active [`Handle`]. +/// +/// [`Handle`]: crate::runtime::scheduler::Handle +pub(crate) fn try_set_current(handle: &scheduler::Handle) -> Option<SetCurrentGuard> { + CONTEXT.try_with(|ctx| ctx.set_current(handle)).ok() +} + +pub(crate) fn with_current<F, R>(f: F) -> Result<R, TryCurrentError> +where + F: FnOnce(&scheduler::Handle) -> R, +{ + match CONTEXT.try_with(|ctx| ctx.current.handle.borrow().as_ref().map(f)) { + Ok(Some(ret)) => Ok(ret), + Ok(None) => Err(TryCurrentError::new_no_context()), + Err(_access_error) => Err(TryCurrentError::new_thread_local_destroyed()), + } +} + +impl Context { + pub(super) fn set_current(&self, handle: &scheduler::Handle) -> SetCurrentGuard { + let old_handle = self.current.handle.borrow_mut().replace(handle.clone()); + let depth = self.current.depth.get(); + + if depth == usize::MAX { + panic!("reached max `enter` depth"); + } + + let depth = depth + 1; + self.current.depth.set(depth); + + SetCurrentGuard { + prev: old_handle, + depth, + _p: PhantomData, + } + } +} + +impl HandleCell { + pub(super) const fn new() -> HandleCell { + HandleCell { + handle: RefCell::new(None), + depth: Cell::new(0), + } + } +} + +impl Drop for SetCurrentGuard { + fn drop(&mut self) { + CONTEXT.with(|ctx| { + let depth = ctx.current.depth.get(); + + if depth != self.depth { + if !std::thread::panicking() { + panic!( + "`EnterGuard` values dropped out of order. Guards returned by \ + `tokio::runtime::Handle::enter()` must be dropped in the reverse \ + order as they were acquired." + ); + } else { + // Just return... this will leave handles in a wonky state though... + return; + } + } + + *ctx.current.handle.borrow_mut() = self.prev.take(); + ctx.current.depth.set(depth - 1); + }); + } +} diff --git a/vendor/tokio/src/runtime/context/runtime.rs b/vendor/tokio/src/runtime/context/runtime.rs new file mode 100644 index 000000000..f2e29899a --- /dev/null +++ b/vendor/tokio/src/runtime/context/runtime.rs @@ -0,0 +1,99 @@ +use super::{BlockingRegionGuard, SetCurrentGuard, CONTEXT}; + +use crate::runtime::scheduler; +use crate::util::rand::{FastRand, RngSeed}; + +use std::fmt; + +#[derive(Debug, Clone, Copy)] +#[must_use] +pub(crate) enum EnterRuntime { + /// Currently in a runtime context. + #[cfg_attr(not(feature = "rt"), allow(dead_code))] + Entered { allow_block_in_place: bool }, + + /// Not in a runtime context **or** a blocking region. + NotEntered, +} + +/// Guard tracking that a caller has entered a runtime context. +#[must_use] +pub(crate) struct EnterRuntimeGuard { + /// Tracks that the current thread has entered a blocking function call. + pub(crate) blocking: BlockingRegionGuard, + + #[allow(dead_code)] // Only tracking the guard. + pub(crate) handle: SetCurrentGuard, + + // Tracks the previous random number generator seed + old_seed: RngSeed, +} + +/// Marks the current thread as being within the dynamic extent of an +/// executor. +#[track_caller] +pub(crate) fn enter_runtime<F, R>(handle: &scheduler::Handle, allow_block_in_place: bool, f: F) -> R +where + F: FnOnce(&mut BlockingRegionGuard) -> R, +{ + let maybe_guard = CONTEXT.with(|c| { + if c.runtime.get().is_entered() { + None + } else { + // Set the entered flag + c.runtime.set(EnterRuntime::Entered { + allow_block_in_place, + }); + + // Generate a new seed + let rng_seed = handle.seed_generator().next_seed(); + + // Swap the RNG seed + let mut rng = c.rng.get().unwrap_or_else(FastRand::new); + let old_seed = rng.replace_seed(rng_seed); + c.rng.set(Some(rng)); + + Some(EnterRuntimeGuard { + blocking: BlockingRegionGuard::new(), + handle: c.set_current(handle), + old_seed, + }) + } + }); + + if let Some(mut guard) = maybe_guard { + return f(&mut guard.blocking); + } + + panic!( + "Cannot start a runtime from within a runtime. This happens \ + because a function (like `block_on`) attempted to block the \ + current thread while the thread is being used to drive \ + asynchronous tasks." + ); +} + +impl fmt::Debug for EnterRuntimeGuard { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Enter").finish() + } +} + +impl Drop for EnterRuntimeGuard { + fn drop(&mut self) { + CONTEXT.with(|c| { + assert!(c.runtime.get().is_entered()); + c.runtime.set(EnterRuntime::NotEntered); + // Replace the previous RNG seed + let mut rng = c.rng.get().unwrap_or_else(FastRand::new); + rng.replace_seed(self.old_seed.clone()); + c.rng.set(Some(rng)); + }); + } +} + +impl EnterRuntime { + pub(crate) fn is_entered(self) -> bool { + matches!(self, EnterRuntime::Entered { .. }) + } +} diff --git a/vendor/tokio/src/runtime/context/runtime_mt.rs b/vendor/tokio/src/runtime/context/runtime_mt.rs new file mode 100644 index 000000000..728caeae9 --- /dev/null +++ b/vendor/tokio/src/runtime/context/runtime_mt.rs @@ -0,0 +1,36 @@ +use super::{EnterRuntime, CONTEXT}; + +/// Returns true if in a runtime context. +pub(crate) fn current_enter_context() -> EnterRuntime { + CONTEXT.with(|c| c.runtime.get()) +} + +/// Forces the current "entered" state to be cleared while the closure +/// is executed. +pub(crate) fn exit_runtime<F: FnOnce() -> R, R>(f: F) -> R { + // Reset in case the closure panics + struct Reset(EnterRuntime); + + impl Drop for Reset { + fn drop(&mut self) { + CONTEXT.with(|c| { + assert!( + !c.runtime.get().is_entered(), + "closure claimed permanent executor" + ); + c.runtime.set(self.0); + }); + } + } + + let was = CONTEXT.with(|c| { + let e = c.runtime.get(); + assert!(e.is_entered(), "asked to exit when not entered"); + c.runtime.set(EnterRuntime::NotEntered); + e + }); + + let _reset = Reset(was); + // dropping _reset after f() will reset ENTERED + f() +} diff --git a/vendor/tokio/src/runtime/context/scoped.rs b/vendor/tokio/src/runtime/context/scoped.rs new file mode 100644 index 000000000..7b202a16c --- /dev/null +++ b/vendor/tokio/src/runtime/context/scoped.rs @@ -0,0 +1,56 @@ +use std::cell::Cell; +use std::ptr; + +/// Scoped thread-local storage +pub(super) struct Scoped<T> { + pub(super) inner: Cell<*const T>, +} + +impl<T> Scoped<T> { + pub(super) const fn new() -> Scoped<T> { + Scoped { + inner: Cell::new(ptr::null()), + } + } + + /// Inserts a value into the scoped cell for the duration of the closure + pub(super) fn set<F, R>(&self, t: &T, f: F) -> R + where + F: FnOnce() -> R, + { + struct Reset<'a, T> { + cell: &'a Cell<*const T>, + prev: *const T, + } + + impl<T> Drop for Reset<'_, T> { + fn drop(&mut self) { + self.cell.set(self.prev); + } + } + + let prev = self.inner.get(); + self.inner.set(t as *const _); + + let _reset = Reset { + cell: &self.inner, + prev, + }; + + f() + } + + /// Gets the value out of the scoped cell; + pub(super) fn with<F, R>(&self, f: F) -> R + where + F: FnOnce(Option<&T>) -> R, + { + let val = self.inner.get(); + + if val.is_null() { + f(None) + } else { + unsafe { f(Some(&*val)) } + } + } +} diff --git a/vendor/tokio/src/coop.rs b/vendor/tokio/src/runtime/coop.rs index 16d93fb75..2dba24615 100644 --- a/vendor/tokio/src/coop.rs +++ b/vendor/tokio/src/runtime/coop.rs @@ -29,17 +29,18 @@ // other futures. By doing this, you avoid double-counting each iteration of // the outer future against the cooperating budget. -use std::cell::Cell; - -thread_local! { - static CURRENT: Cell<Budget> = Cell::new(Budget::unconstrained()); -} +use crate::runtime::context; /// Opaque type tracking the amount of "work" a task may still do before /// yielding back to the scheduler. #[derive(Debug, Copy, Clone)] pub(crate) struct Budget(Option<u8>); +pub(crate) struct BudgetDecrement { + success: bool, + hit_zero: bool, +} + impl Budget { /// Budget assigned to a task on each poll. /// @@ -56,27 +57,23 @@ impl Budget { } /// Returns an unconstrained budget. Operations will not be limited. - const fn unconstrained() -> Budget { + pub(super) const fn unconstrained() -> Budget { Budget(None) } -} -cfg_rt_multi_thread! { - impl Budget { - fn has_remaining(self) -> bool { - self.0.map(|budget| budget > 0).unwrap_or(true) - } + fn has_remaining(self) -> bool { + self.0.map(|budget| budget > 0).unwrap_or(true) } } -/// Run the given closure with a cooperative task budget. When the function +/// Runs the given closure with a cooperative task budget. When the function /// returns, the budget is reset to the value prior to calling the function. #[inline(always)] pub(crate) fn budget<R>(f: impl FnOnce() -> R) -> R { with_budget(Budget::initial(), f) } -/// Run the given closure with an unconstrained task budget. When the function returns, the budget +/// Runs the given closure with an unconstrained task budget. When the function returns, the budget /// is reset to the value prior to calling the function. #[inline(always)] pub(crate) fn with_unconstrained<R>(f: impl FnOnce() -> R) -> R { @@ -85,54 +82,60 @@ pub(crate) fn with_unconstrained<R>(f: impl FnOnce() -> R) -> R { #[inline(always)] fn with_budget<R>(budget: Budget, f: impl FnOnce() -> R) -> R { - struct ResetGuard<'a> { - cell: &'a Cell<Budget>, + struct ResetGuard { prev: Budget, } - impl<'a> Drop for ResetGuard<'a> { + impl Drop for ResetGuard { fn drop(&mut self) { - self.cell.set(self.prev); + let _ = context::budget(|cell| { + cell.set(self.prev); + }); } } - CURRENT.with(move |cell| { + #[allow(unused_variables)] + let maybe_guard = context::budget(|cell| { let prev = cell.get(); - cell.set(budget); - let _guard = ResetGuard { cell, prev }; + ResetGuard { prev } + }); + + // The function is called regardless even if the budget is not successfully + // set due to the thread-local being destroyed. + f() +} - f() - }) +#[inline(always)] +pub(crate) fn has_budget_remaining() -> bool { + // If the current budget cannot be accessed due to the thread-local being + // shutdown, then we assume there is budget remaining. + context::budget(|cell| cell.get().has_remaining()).unwrap_or(true) } cfg_rt_multi_thread! { - /// Set the current task's budget + /// Sets the current task's budget. pub(crate) fn set(budget: Budget) { - CURRENT.with(|cell| cell.set(budget)) - } - - #[inline(always)] - pub(crate) fn has_budget_remaining() -> bool { - CURRENT.with(|cell| cell.get().has_remaining()) + let _ = context::budget(|cell| cell.set(budget)); } } cfg_rt! { - /// Forcibly remove the budgeting constraints early. + /// Forcibly removes the budgeting constraints early. /// /// Returns the remaining budget pub(crate) fn stop() -> Budget { - CURRENT.with(|cell| { + context::budget(|cell| { let prev = cell.get(); cell.set(Budget::unconstrained()); prev - }) + }).unwrap_or(Budget::unconstrained()) } } cfg_coop! { + use std::cell::Cell; use std::task::{Context, Poll}; #[must_use] @@ -150,7 +153,7 @@ cfg_coop! { // They are both represented as the remembered budget being unconstrained. let budget = self.0.get(); if !budget.is_unconstrained() { - CURRENT.with(|cell| { + let _ = context::budget(|cell| { cell.set(budget); }); } @@ -171,33 +174,65 @@ cfg_coop! { /// that progress was made. #[inline] pub(crate) fn poll_proceed(cx: &mut Context<'_>) -> Poll<RestoreOnPending> { - CURRENT.with(|cell| { + context::budget(|cell| { let mut budget = cell.get(); - if budget.decrement() { + let decrement = budget.decrement(); + + if decrement.success { let restore = RestoreOnPending(Cell::new(cell.get())); cell.set(budget); + + // avoid double counting + if decrement.hit_zero { + inc_budget_forced_yield_count(); + } + Poll::Ready(restore) } else { cx.waker().wake_by_ref(); Poll::Pending } - }) + }).unwrap_or(Poll::Ready(RestoreOnPending(Cell::new(Budget::unconstrained())))) + } + + cfg_rt! { + cfg_metrics! { + #[inline(always)] + fn inc_budget_forced_yield_count() { + let _ = context::with_current(|handle| { + handle.scheduler_metrics().inc_budget_forced_yield_count(); + }); + } + } + + cfg_not_metrics! { + #[inline(always)] + fn inc_budget_forced_yield_count() {} + } + } + + cfg_not_rt! { + #[inline(always)] + fn inc_budget_forced_yield_count() {} } impl Budget { - /// Decrement the budget. Returns `true` if successful. Decrementing fails + /// Decrements the budget. Returns `true` if successful. Decrementing fails /// when there is not enough remaining budget. - fn decrement(&mut self) -> bool { + fn decrement(&mut self) -> BudgetDecrement { if let Some(num) = &mut self.0 { if *num > 0 { *num -= 1; - true + + let hit_zero = *num == 0; + + BudgetDecrement { success: true, hit_zero } } else { - false + BudgetDecrement { success: false, hit_zero: false } } } else { - true + BudgetDecrement { success: true, hit_zero: false } } } @@ -211,12 +246,15 @@ cfg_coop! { mod test { use super::*; + #[cfg(tokio_wasm_not_wasi)] + use wasm_bindgen_test::wasm_bindgen_test as test; + fn get() -> Budget { - CURRENT.with(|cell| cell.get()) + context::budget(|cell| cell.get()).unwrap_or(Budget::unconstrained()) } #[test] - fn bugeting() { + fn budgeting() { use futures::future::poll_fn; use tokio_test::*; diff --git a/vendor/tokio/src/runtime/driver.rs b/vendor/tokio/src/runtime/driver.rs index 7e459779b..572fdefb0 100644 --- a/vendor/tokio/src/runtime/driver.rs +++ b/vendor/tokio/src/runtime/driver.rs @@ -1,63 +1,233 @@ //! Abstracts out the entire chain of runtime sub-drivers into common types. -use crate::park::thread::ParkThread; -use crate::park::Park; + +// Eventually, this file will see significant refactoring / cleanup. For now, we +// don't need to worry much about dead code with certain feature permutations. +#![cfg_attr(not(feature = "full"), allow(dead_code))] + +use crate::runtime::park::{ParkThread, UnparkThread}; use std::io; use std::time::Duration; +#[derive(Debug)] +pub(crate) struct Driver { + inner: TimeDriver, +} + +#[derive(Debug)] +pub(crate) struct Handle { + /// IO driver handle + pub(crate) io: IoHandle, + + /// Signal driver handle + #[cfg_attr(any(not(unix), loom), allow(dead_code))] + pub(crate) signal: SignalHandle, + + /// Time driver handle + pub(crate) time: TimeHandle, + + /// Source of `Instant::now()` + #[cfg_attr(not(all(feature = "time", feature = "test-util")), allow(dead_code))] + pub(crate) clock: Clock, +} + +pub(crate) struct Cfg { + pub(crate) enable_io: bool, + pub(crate) enable_time: bool, + pub(crate) enable_pause_time: bool, + pub(crate) start_paused: bool, + pub(crate) nevents: usize, +} + +impl Driver { + pub(crate) fn new(cfg: Cfg) -> io::Result<(Self, Handle)> { + let (io_stack, io_handle, signal_handle) = create_io_stack(cfg.enable_io, cfg.nevents)?; + + let clock = create_clock(cfg.enable_pause_time, cfg.start_paused); + + let (time_driver, time_handle) = create_time_driver(cfg.enable_time, io_stack, &clock); + + Ok(( + Self { inner: time_driver }, + Handle { + io: io_handle, + signal: signal_handle, + time: time_handle, + clock, + }, + )) + } + + pub(crate) fn park(&mut self, handle: &Handle) { + self.inner.park(handle) + } + + pub(crate) fn park_timeout(&mut self, handle: &Handle, duration: Duration) { + self.inner.park_timeout(handle, duration) + } + + pub(crate) fn shutdown(&mut self, handle: &Handle) { + self.inner.shutdown(handle) + } +} + +impl Handle { + pub(crate) fn unpark(&self) { + #[cfg(feature = "time")] + if let Some(handle) = &self.time { + handle.unpark(); + } + + self.io.unpark(); + } + + cfg_io_driver! { + #[track_caller] + pub(crate) fn io(&self) -> &crate::runtime::io::Handle { + self.io + .as_ref() + .expect("A Tokio 1.x context was found, but IO is disabled. Call `enable_io` on the runtime builder to enable IO.") + } + } + + cfg_signal_internal_and_unix! { + #[track_caller] + pub(crate) fn signal(&self) -> &crate::runtime::signal::Handle { + self.signal + .as_ref() + .expect("there is no signal driver running, must be called from the context of Tokio runtime") + } + } + + cfg_time! { + /// Returns a reference to the time driver handle. + /// + /// Panics if no time driver is present. + #[track_caller] + pub(crate) fn time(&self) -> &crate::runtime::time::Handle { + self.time + .as_ref() + .expect("A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers.") + } + + pub(crate) fn clock(&self) -> &Clock { + &self.clock + } + } +} + // ===== io driver ===== cfg_io_driver! { - type IoDriver = crate::io::driver::Driver; - type IoStack = crate::park::either::Either<ProcessDriver, ParkThread>; - pub(crate) type IoHandle = Option<crate::io::driver::Handle>; + pub(crate) type IoDriver = crate::runtime::io::Driver; - fn create_io_stack(enabled: bool) -> io::Result<(IoStack, IoHandle, SignalHandle)> { - use crate::park::either::Either; + #[derive(Debug)] + pub(crate) enum IoStack { + Enabled(ProcessDriver), + Disabled(ParkThread), + } + + #[derive(Debug)] + pub(crate) enum IoHandle { + Enabled(crate::runtime::io::Handle), + Disabled(UnparkThread), + } + fn create_io_stack(enabled: bool, nevents: usize) -> io::Result<(IoStack, IoHandle, SignalHandle)> { #[cfg(loom)] assert!(!enabled); let ret = if enabled { - let io_driver = crate::io::driver::Driver::new()?; - let io_handle = io_driver.handle(); + let (io_driver, io_handle) = crate::runtime::io::Driver::new(nevents)?; - let (signal_driver, signal_handle) = create_signal_driver(io_driver)?; + let (signal_driver, signal_handle) = create_signal_driver(io_driver, &io_handle)?; let process_driver = create_process_driver(signal_driver); - (Either::A(process_driver), Some(io_handle), signal_handle) + (IoStack::Enabled(process_driver), IoHandle::Enabled(io_handle), signal_handle) } else { - (Either::B(ParkThread::new()), Default::default(), Default::default()) + let park_thread = ParkThread::new(); + let unpark_thread = park_thread.unpark(); + (IoStack::Disabled(park_thread), IoHandle::Disabled(unpark_thread), Default::default()) }; Ok(ret) } + + impl IoStack { + pub(crate) fn park(&mut self, handle: &Handle) { + match self { + IoStack::Enabled(v) => v.park(handle), + IoStack::Disabled(v) => v.park(), + } + } + + pub(crate) fn park_timeout(&mut self, handle: &Handle, duration: Duration) { + match self { + IoStack::Enabled(v) => v.park_timeout(handle, duration), + IoStack::Disabled(v) => v.park_timeout(duration), + } + } + + pub(crate) fn shutdown(&mut self, handle: &Handle) { + match self { + IoStack::Enabled(v) => v.shutdown(handle), + IoStack::Disabled(v) => v.shutdown(), + } + } + } + + impl IoHandle { + pub(crate) fn unpark(&self) { + match self { + IoHandle::Enabled(handle) => handle.unpark(), + IoHandle::Disabled(handle) => handle.unpark(), + } + } + + pub(crate) fn as_ref(&self) -> Option<&crate::runtime::io::Handle> { + match self { + IoHandle::Enabled(v) => Some(v), + IoHandle::Disabled(..) => None, + } + } + } } cfg_not_io_driver! { - pub(crate) type IoHandle = (); - type IoStack = ParkThread; + pub(crate) type IoHandle = UnparkThread; + + #[derive(Debug)] + pub(crate) struct IoStack(ParkThread); - fn create_io_stack(_enabled: bool) -> io::Result<(IoStack, IoHandle, SignalHandle)> { - Ok((ParkThread::new(), Default::default(), Default::default())) + fn create_io_stack(_enabled: bool, _nevents: usize) -> io::Result<(IoStack, IoHandle, SignalHandle)> { + let park_thread = ParkThread::new(); + let unpark_thread = park_thread.unpark(); + Ok((IoStack(park_thread), unpark_thread, Default::default())) } -} -// ===== signal driver ===== + impl IoStack { + pub(crate) fn park(&mut self, _handle: &Handle) { + self.0.park(); + } + + pub(crate) fn park_timeout(&mut self, _handle: &Handle, duration: Duration) { + self.0.park_timeout(duration); + } -macro_rules! cfg_signal_internal_and_unix { - ($($item:item)*) => { - #[cfg(unix)] - cfg_signal_internal! { $($item)* } + pub(crate) fn shutdown(&mut self, _handle: &Handle) { + self.0.shutdown(); + } } } +// ===== signal driver ===== + cfg_signal_internal_and_unix! { - type SignalDriver = crate::signal::unix::driver::Driver; - pub(crate) type SignalHandle = Option<crate::signal::unix::driver::Handle>; + type SignalDriver = crate::runtime::signal::Driver; + pub(crate) type SignalHandle = Option<crate::runtime::signal::Handle>; - fn create_signal_driver(io_driver: IoDriver) -> io::Result<(SignalDriver, SignalHandle)> { - let driver = crate::signal::unix::driver::Driver::new(io_driver)?; + fn create_signal_driver(io_driver: IoDriver, io_handle: &crate::runtime::io::Handle) -> io::Result<(SignalDriver, SignalHandle)> { + let driver = crate::runtime::signal::Driver::new(io_driver, io_handle)?; let handle = driver.handle(); Ok((driver, Some(handle))) } @@ -69,7 +239,7 @@ cfg_not_signal_internal! { cfg_io_driver! { type SignalDriver = IoDriver; - fn create_signal_driver(io_driver: IoDriver) -> io::Result<(SignalDriver, SignalHandle)> { + fn create_signal_driver(io_driver: IoDriver, _io_handle: &crate::runtime::io::Handle) -> io::Result<(SignalDriver, SignalHandle)> { Ok((io_driver, ())) } } @@ -78,10 +248,10 @@ cfg_not_signal_internal! { // ===== process driver ===== cfg_process_driver! { - type ProcessDriver = crate::process::unix::driver::Driver; + type ProcessDriver = crate::runtime::process::Driver; fn create_process_driver(signal_driver: SignalDriver) -> ProcessDriver { - crate::process::unix::driver::Driver::new(signal_driver) + ProcessDriver::new(signal_driver) } } @@ -98,10 +268,16 @@ cfg_not_process_driver! { // ===== time driver ===== cfg_time! { - type TimeDriver = crate::park::either::Either<crate::time::driver::Driver<IoStack>, IoStack>; + #[derive(Debug)] + pub(crate) enum TimeDriver { + Enabled { + driver: crate::runtime::time::Driver, + }, + Disabled(IoStack), + } pub(crate) type Clock = crate::time::Clock; - pub(crate) type TimeHandle = Option<crate::time::driver::Handle>; + pub(crate) type TimeHandle = Option<crate::runtime::time::Handle>; fn create_clock(enable_pausing: bool, start_paused: bool) -> Clock { crate::time::Clock::new(enable_pausing, start_paused) @@ -110,17 +286,37 @@ cfg_time! { fn create_time_driver( enable: bool, io_stack: IoStack, - clock: Clock, + clock: &Clock, ) -> (TimeDriver, TimeHandle) { - use crate::park::either::Either; - if enable { - let driver = crate::time::driver::Driver::new(io_stack, clock); - let handle = driver.handle(); + let (driver, handle) = crate::runtime::time::Driver::new(io_stack, clock); - (Either::A(driver), Some(handle)) + (TimeDriver::Enabled { driver }, Some(handle)) } else { - (Either::B(io_stack), None) + (TimeDriver::Disabled(io_stack), None) + } + } + + impl TimeDriver { + pub(crate) fn park(&mut self, handle: &Handle) { + match self { + TimeDriver::Enabled { driver, .. } => driver.park(handle), + TimeDriver::Disabled(v) => v.park(handle), + } + } + + pub(crate) fn park_timeout(&mut self, handle: &Handle, duration: Duration) { + match self { + TimeDriver::Enabled { driver } => driver.park_timeout(handle, duration), + TimeDriver::Disabled(v) => v.park_timeout(handle, duration), + } + } + + pub(crate) fn shutdown(&mut self, handle: &Handle) { + match self { + TimeDriver::Enabled { driver } => driver.shutdown(handle), + TimeDriver::Disabled(v) => v.shutdown(handle), + } } } } @@ -138,71 +334,8 @@ cfg_not_time! { fn create_time_driver( _enable: bool, io_stack: IoStack, - _clock: Clock, + _clock: &Clock, ) -> (TimeDriver, TimeHandle) { (io_stack, ()) } } - -// ===== runtime driver ===== - -#[derive(Debug)] -pub(crate) struct Driver { - inner: TimeDriver, -} - -pub(crate) struct Resources { - pub(crate) io_handle: IoHandle, - pub(crate) signal_handle: SignalHandle, - pub(crate) time_handle: TimeHandle, - pub(crate) clock: Clock, -} - -pub(crate) struct Cfg { - pub(crate) enable_io: bool, - pub(crate) enable_time: bool, - pub(crate) enable_pause_time: bool, - pub(crate) start_paused: bool, -} - -impl Driver { - pub(crate) fn new(cfg: Cfg) -> io::Result<(Self, Resources)> { - let (io_stack, io_handle, signal_handle) = create_io_stack(cfg.enable_io)?; - - let clock = create_clock(cfg.enable_pause_time, cfg.start_paused); - - let (time_driver, time_handle) = - create_time_driver(cfg.enable_time, io_stack, clock.clone()); - - Ok(( - Self { inner: time_driver }, - Resources { - io_handle, - signal_handle, - time_handle, - clock, - }, - )) - } -} - -impl Park for Driver { - type Unpark = <TimeDriver as Park>::Unpark; - type Error = <TimeDriver as Park>::Error; - - fn unpark(&self) -> Self::Unpark { - self.inner.unpark() - } - - fn park(&mut self) -> Result<(), Self::Error> { - self.inner.park() - } - - fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> { - self.inner.park_timeout(duration) - } - - fn shutdown(&mut self) { - self.inner.shutdown() - } -} diff --git a/vendor/tokio/src/runtime/dump.rs b/vendor/tokio/src/runtime/dump.rs new file mode 100644 index 000000000..994b7f9c0 --- /dev/null +++ b/vendor/tokio/src/runtime/dump.rs @@ -0,0 +1,76 @@ +//! Snapshots of runtime state. +//! +//! See [Handle::dump][crate::runtime::Handle::dump]. + +use std::fmt; + +/// A snapshot of a runtime's state. +/// +/// See [Handle::dump][crate::runtime::Handle::dump]. +#[derive(Debug)] +pub struct Dump { + tasks: Tasks, +} + +/// Snapshots of tasks. +/// +/// See [Handle::dump][crate::runtime::Handle::dump]. +#[derive(Debug)] +pub struct Tasks { + tasks: Vec<Task>, +} + +/// A snapshot of a task. +/// +/// See [Handle::dump][crate::runtime::Handle::dump]. +#[derive(Debug)] +pub struct Task { + trace: Trace, +} + +/// An execution trace of a task's last poll. +/// +/// See [Handle::dump][crate::runtime::Handle::dump]. +#[derive(Debug)] +pub struct Trace { + inner: super::task::trace::Trace, +} + +impl Dump { + pub(crate) fn new(tasks: Vec<Task>) -> Self { + Self { + tasks: Tasks { tasks }, + } + } + + /// Tasks in this snapshot. + pub fn tasks(&self) -> &Tasks { + &self.tasks + } +} + +impl Tasks { + /// Iterate over tasks. + pub fn iter(&self) -> impl Iterator<Item = &Task> { + self.tasks.iter() + } +} + +impl Task { + pub(crate) fn new(trace: super::task::trace::Trace) -> Self { + Self { + trace: Trace { inner: trace }, + } + } + + /// A trace of this task's state. + pub fn trace(&self) -> &Trace { + &self.trace + } +} + +impl fmt::Display for Trace { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.inner.fmt(f) + } +} diff --git a/vendor/tokio/src/runtime/enter.rs b/vendor/tokio/src/runtime/enter.rs deleted file mode 100644 index e91408f38..000000000 --- a/vendor/tokio/src/runtime/enter.rs +++ /dev/null @@ -1,205 +0,0 @@ -use std::cell::{Cell, RefCell}; -use std::fmt; -use std::marker::PhantomData; - -#[derive(Debug, Clone, Copy)] -pub(crate) enum EnterContext { - #[cfg_attr(not(feature = "rt"), allow(dead_code))] - Entered { - allow_blocking: bool, - }, - NotEntered, -} - -impl EnterContext { - pub(crate) fn is_entered(self) -> bool { - matches!(self, EnterContext::Entered { .. }) - } -} - -thread_local!(static ENTERED: Cell<EnterContext> = Cell::new(EnterContext::NotEntered)); - -/// Represents an executor context. -pub(crate) struct Enter { - _p: PhantomData<RefCell<()>>, -} - -cfg_rt! { - use crate::park::thread::ParkError; - - use std::time::Duration; - - /// Marks the current thread as being within the dynamic extent of an - /// executor. - pub(crate) fn enter(allow_blocking: bool) -> Enter { - if let Some(enter) = try_enter(allow_blocking) { - return enter; - } - - panic!( - "Cannot start a runtime from within a runtime. This happens \ - because a function (like `block_on`) attempted to block the \ - current thread while the thread is being used to drive \ - asynchronous tasks." - ); - } - - /// Tries to enter a runtime context, returns `None` if already in a runtime - /// context. - pub(crate) fn try_enter(allow_blocking: bool) -> Option<Enter> { - ENTERED.with(|c| { - if c.get().is_entered() { - None - } else { - c.set(EnterContext::Entered { allow_blocking }); - Some(Enter { _p: PhantomData }) - } - }) - } -} - -// Forces the current "entered" state to be cleared while the closure -// is executed. -// -// # Warning -// -// This is hidden for a reason. Do not use without fully understanding -// executors. Misusing can easily cause your program to deadlock. -cfg_rt_multi_thread! { - pub(crate) fn exit<F: FnOnce() -> R, R>(f: F) -> R { - // Reset in case the closure panics - struct Reset(EnterContext); - impl Drop for Reset { - fn drop(&mut self) { - ENTERED.with(|c| { - assert!(!c.get().is_entered(), "closure claimed permanent executor"); - c.set(self.0); - }); - } - } - - let was = ENTERED.with(|c| { - let e = c.get(); - assert!(e.is_entered(), "asked to exit when not entered"); - c.set(EnterContext::NotEntered); - e - }); - - let _reset = Reset(was); - // dropping _reset after f() will reset ENTERED - f() - } -} - -cfg_rt! { - /// Disallow blocking in the current runtime context until the guard is dropped. - pub(crate) fn disallow_blocking() -> DisallowBlockingGuard { - let reset = ENTERED.with(|c| { - if let EnterContext::Entered { - allow_blocking: true, - } = c.get() - { - c.set(EnterContext::Entered { - allow_blocking: false, - }); - true - } else { - false - } - }); - DisallowBlockingGuard(reset) - } - - pub(crate) struct DisallowBlockingGuard(bool); - impl Drop for DisallowBlockingGuard { - fn drop(&mut self) { - if self.0 { - // XXX: Do we want some kind of assertion here, or is "best effort" okay? - ENTERED.with(|c| { - if let EnterContext::Entered { - allow_blocking: false, - } = c.get() - { - c.set(EnterContext::Entered { - allow_blocking: true, - }); - } - }) - } - } - } -} - -cfg_rt_multi_thread! { - /// Returns true if in a runtime context. - pub(crate) fn context() -> EnterContext { - ENTERED.with(|c| c.get()) - } -} - -cfg_rt! { - impl Enter { - /// Blocks the thread on the specified future, returning the value with - /// which that future completes. - pub(crate) fn block_on<F>(&mut self, f: F) -> Result<F::Output, ParkError> - where - F: std::future::Future, - { - use crate::park::thread::CachedParkThread; - - let mut park = CachedParkThread::new(); - park.block_on(f) - } - - /// Blocks the thread on the specified future for **at most** `timeout` - /// - /// If the future completes before `timeout`, the result is returned. If - /// `timeout` elapses, then `Err` is returned. - pub(crate) fn block_on_timeout<F>(&mut self, f: F, timeout: Duration) -> Result<F::Output, ParkError> - where - F: std::future::Future, - { - use crate::park::Park; - use crate::park::thread::CachedParkThread; - use std::task::Context; - use std::task::Poll::Ready; - use std::time::Instant; - - let mut park = CachedParkThread::new(); - let waker = park.get_unpark()?.into_waker(); - let mut cx = Context::from_waker(&waker); - - pin!(f); - let when = Instant::now() + timeout; - - loop { - if let Ready(v) = crate::coop::budget(|| f.as_mut().poll(&mut cx)) { - return Ok(v); - } - - let now = Instant::now(); - - if now >= when { - return Err(()); - } - - park.park_timeout(when - now)?; - } - } - } -} - -impl fmt::Debug for Enter { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("Enter").finish() - } -} - -impl Drop for Enter { - fn drop(&mut self) { - ENTERED.with(|c| { - assert!(c.get().is_entered()); - c.set(EnterContext::NotEntered); - }); - } -} diff --git a/vendor/tokio/src/runtime/handle.rs b/vendor/tokio/src/runtime/handle.rs index 7dff91448..be4743d47 100644 --- a/vendor/tokio/src/runtime/handle.rs +++ b/vendor/tokio/src/runtime/handle.rs @@ -1,10 +1,4 @@ -use crate::runtime::blocking::task::BlockingTask; -use crate::runtime::task::{self, JoinHandle}; -use crate::runtime::{blocking, context, driver, Spawner}; -use crate::util::error::CONTEXT_MISSING_ERROR; - -use std::future::Future; -use std::{error, fmt}; +use crate::runtime::{context, scheduler, RuntimeFlavor}; /// Handle to the runtime. /// @@ -13,24 +7,18 @@ use std::{error, fmt}; /// /// [`Runtime::handle`]: crate::runtime::Runtime::handle() #[derive(Debug, Clone)] +// When the `rt` feature is *not* enabled, this type is still defined, but not +// included in the public API. pub struct Handle { - pub(super) spawner: Spawner, - - /// Handles to the I/O drivers - pub(super) io_handle: driver::IoHandle, - - /// Handles to the signal drivers - pub(super) signal_handle: driver::SignalHandle, - - /// Handles to the time drivers - pub(super) time_handle: driver::TimeHandle, + pub(crate) inner: scheduler::Handle, +} - /// Source of `Instant::now()` - pub(super) clock: driver::Clock, +use crate::runtime::task::JoinHandle; +use crate::util::error::{CONTEXT_MISSING_ERROR, THREAD_LOCAL_DESTROYED_ERROR}; - /// Blocking pool spawner - pub(super) blocking_spawner: blocking::Spawner, -} +use std::future::Future; +use std::marker::PhantomData; +use std::{error, fmt}; /// Runtime context guard. /// @@ -41,32 +29,72 @@ pub struct Handle { #[derive(Debug)] #[must_use = "Creating and dropping a guard does nothing"] pub struct EnterGuard<'a> { - handle: &'a Handle, - guard: context::EnterGuard, + _guard: context::SetCurrentGuard, + _handle_lifetime: PhantomData<&'a Handle>, } impl Handle { - /// Enter the runtime context. This allows you to construct types that must - /// have an executor available on creation such as [`Sleep`] or [`TcpStream`]. - /// It will also allow you to call methods such as [`tokio::spawn`]. + /// Enters the runtime context. This allows you to construct types that must + /// have an executor available on creation such as [`Sleep`] or + /// [`TcpStream`]. It will also allow you to call methods such as + /// [`tokio::spawn`] and [`Handle::current`] without panicking. + /// + /// # Panics + /// + /// When calling `Handle::enter` multiple times, the returned guards + /// **must** be dropped in the reverse order that they were acquired. + /// Failure to do so will result in a panic and possible memory leaks. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Runtime; + /// + /// let rt = Runtime::new().unwrap(); + /// + /// let _guard = rt.enter(); + /// tokio::spawn(async { + /// println!("Hello world!"); + /// }); + /// ``` + /// + /// Do **not** do the following, this shows a scenario that will result in a + /// panic and possible memory leak. + /// + /// ```should_panic + /// use tokio::runtime::Runtime; + /// + /// let rt1 = Runtime::new().unwrap(); + /// let rt2 = Runtime::new().unwrap(); + /// + /// let enter1 = rt1.enter(); + /// let enter2 = rt2.enter(); + /// + /// drop(enter1); + /// drop(enter2); + /// ``` /// /// [`Sleep`]: struct@crate::time::Sleep /// [`TcpStream`]: struct@crate::net::TcpStream /// [`tokio::spawn`]: fn@crate::spawn pub fn enter(&self) -> EnterGuard<'_> { EnterGuard { - handle: self, - guard: context::enter(self.clone()), + _guard: match context::try_set_current(&self.inner) { + Some(guard) => guard, + None => panic!("{}", crate::util::error::THREAD_LOCAL_DESTROYED_ERROR), + }, + _handle_lifetime: PhantomData, } } - /// Returns a `Handle` view over the currently running `Runtime` + /// Returns a `Handle` view over the currently running `Runtime`. /// - /// # Panic + /// # Panics /// /// This will panic if called outside the context of a Tokio runtime. That means that you must - /// call this on one of the threads **being run by the runtime**. Calling this from within a - /// thread created by `std::thread::spawn` (for example) will cause a panic. + /// call this on one of the threads **being run by the runtime**, or from a thread with an active + /// `EnterGuard`. Calling this from within a thread created by `std::thread::spawn` (for example) + /// will cause a panic unless that thread has an active `EnterGuard`. /// /// # Examples /// @@ -90,16 +118,24 @@ impl Handle { /// # let handle = /// thread::spawn(move || { /// // Notice that the handle is created outside of this thread and then moved in - /// handle.spawn(async { /* ... */ }) - /// // This next line would cause a panic - /// // let handle2 = Handle::current(); + /// handle.spawn(async { /* ... */ }); + /// // This next line would cause a panic because we haven't entered the runtime + /// // and created an EnterGuard + /// // let handle2 = Handle::current(); // panic + /// // So we create a guard here with Handle::enter(); + /// let _guard = handle.enter(); + /// // Now we can call Handle::current(); + /// let handle2 = Handle::current(); /// }); /// # handle.join().unwrap(); /// # }); /// # } /// ``` + #[track_caller] pub fn current() -> Self { - context::current().expect(CONTEXT_MISSING_ERROR) + Handle { + inner: scheduler::Handle::current(), + } } /// Returns a Handle view over the currently running Runtime @@ -108,15 +144,21 @@ impl Handle { /// /// Contrary to `current`, this never panics pub fn try_current() -> Result<Self, TryCurrentError> { - context::current().ok_or(TryCurrentError(())) + context::with_current(|inner| Handle { + inner: inner.clone(), + }) } - /// Spawn a future onto the Tokio runtime. + /// Spawns a future onto the Tokio runtime. /// /// This spawns the given future onto the runtime's executor, usually a /// thread pool. The thread pool is then responsible for polling the future /// until it completes. /// + /// The provided future will start running in the background immediately + /// when `spawn` is called, even if you don't await the returned + /// `JoinHandle`. + /// /// See [module level][mod] documentation for more details. /// /// [mod]: index.html @@ -138,18 +180,16 @@ impl Handle { /// }); /// # } /// ``` - #[cfg_attr(tokio_track_caller, track_caller)] + #[track_caller] pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output> where F: Future + Send + 'static, F::Output: Send + 'static, { - #[cfg(all(tokio_unstable, feature = "tracing"))] - let future = crate::util::trace::task(future, "task", None); - self.spawner.spawn(future) + self.spawn_named(future, None) } - /// Run the provided function on an executor dedicated to blocking + /// Runs the provided function on an executor dedicated to blocking /// operations. /// /// # Examples @@ -168,57 +208,16 @@ impl Handle { /// println!("now running on a worker thread"); /// }); /// # } - #[cfg_attr(tokio_track_caller, track_caller)] + #[track_caller] pub fn spawn_blocking<F, R>(&self, func: F) -> JoinHandle<R> where F: FnOnce() -> R + Send + 'static, R: Send + 'static, { - self.spawn_blocking_inner(func, None) + self.inner.blocking_spawner().spawn_blocking(self, func) } - #[cfg_attr(tokio_track_caller, track_caller)] - pub(crate) fn spawn_blocking_inner<F, R>(&self, func: F, name: Option<&str>) -> JoinHandle<R> - where - F: FnOnce() -> R + Send + 'static, - R: Send + 'static, - { - let fut = BlockingTask::new(func); - - #[cfg(all(tokio_unstable, feature = "tracing"))] - let fut = { - use tracing::Instrument; - #[cfg(tokio_track_caller)] - let location = std::panic::Location::caller(); - #[cfg(tokio_track_caller)] - let span = tracing::trace_span!( - target: "tokio::task", - "task", - kind = %"blocking", - function = %std::any::type_name::<F>(), - task.name = %name.unwrap_or_default(), - spawn.location = %format_args!("{}:{}:{}", location.file(), location.line(), location.column()), - ); - #[cfg(not(tokio_track_caller))] - let span = tracing::trace_span!( - target: "tokio::task", - "task", - kind = %"blocking", - task.name = %name.unwrap_or_default(), - function = %std::any::type_name::<F>(), - ); - fut.instrument(span) - }; - - #[cfg(not(all(tokio_unstable, feature = "tracing")))] - let _ = name; - - let (task, handle) = task::joinable(fut); - let _ = self.blocking_spawner.spawn(task, &self); - handle - } - - /// Run a future to completion on this `Handle`'s associated `Runtime`. + /// Runs a future to completion on this `Handle`'s associated `Runtime`. /// /// This runs the given future on the current thread, blocking until it is /// complete, and yielding its resolved result. Any tasks or timers which @@ -288,36 +287,300 @@ impl Handle { /// [`tokio::fs`]: crate::fs /// [`tokio::net`]: crate::net /// [`tokio::time`]: crate::time + #[track_caller] pub fn block_on<F: Future>(&self, future: F) -> F::Output { - // Enter the **runtime** context. This configures spawning, the current I/O driver, ... - let _rt_enter = self.enter(); + #[cfg(all( + tokio_unstable, + tokio_taskdump, + feature = "rt", + target_os = "linux", + any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64") + ))] + let future = super::task::trace::Trace::root(future); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + let future = + crate::util::trace::task(future, "block_on", None, super::task::Id::next().as_u64()); + + // Enter the runtime context. This sets the current driver handles and + // prevents blocking an existing runtime. + context::enter_runtime(&self.inner, true, |blocking| { + blocking.block_on(future).expect("failed to park thread") + }) + } - // Enter a **blocking** context. This prevents blocking from a runtime. - let mut blocking_enter = crate::runtime::enter(true); + #[track_caller] + pub(crate) fn spawn_named<F>(&self, future: F, _name: Option<&str>) -> JoinHandle<F::Output> + where + F: Future + Send + 'static, + F::Output: Send + 'static, + { + let id = crate::runtime::task::Id::next(); + #[cfg(all( + tokio_unstable, + tokio_taskdump, + feature = "rt", + target_os = "linux", + any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64") + ))] + let future = super::task::trace::Trace::root(future); + #[cfg(all(tokio_unstable, feature = "tracing"))] + let future = crate::util::trace::task(future, "task", _name, id.as_u64()); + self.inner.spawn(future, id) + } - // Block on the future - blocking_enter - .block_on(future) - .expect("failed to park thread") + /// Returns the flavor of the current `Runtime`. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::{Handle, RuntimeFlavor}; + /// + /// #[tokio::main(flavor = "current_thread")] + /// async fn main() { + /// assert_eq!(RuntimeFlavor::CurrentThread, Handle::current().runtime_flavor()); + /// } + /// ``` + /// + /// ``` + /// use tokio::runtime::{Handle, RuntimeFlavor}; + /// + /// #[tokio::main(flavor = "multi_thread", worker_threads = 4)] + /// async fn main() { + /// assert_eq!(RuntimeFlavor::MultiThread, Handle::current().runtime_flavor()); + /// } + /// ``` + pub fn runtime_flavor(&self) -> RuntimeFlavor { + match self.inner { + scheduler::Handle::CurrentThread(_) => RuntimeFlavor::CurrentThread, + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + scheduler::Handle::MultiThread(_) => RuntimeFlavor::MultiThread, + } } +} - pub(crate) fn shutdown(mut self) { - self.spawner.shutdown(); +cfg_metrics! { + use crate::runtime::RuntimeMetrics; + + impl Handle { + /// Returns a view that lets you get information about how the runtime + /// is performing. + pub fn metrics(&self) -> RuntimeMetrics { + RuntimeMetrics::new(self.clone()) + } + } +} + +cfg_taskdump! { + impl Handle { + /// Captures a snapshot of the runtime's state. + /// + /// This functionality is experimental, and comes with a number of + /// requirements and limitations. + /// + /// # Examples + /// + /// This can be used to get call traces of each task in the runtime. + /// Calls to `Handle::dump` should usually be enclosed in a + /// [timeout][crate::time::timeout], so that dumping does not escalate a + /// single blocked runtime thread into an entirely blocked runtime. + /// + /// ``` + /// # use tokio::runtime::Runtime; + /// # fn dox() { + /// # let rt = Runtime::new().unwrap(); + /// # rt.spawn(async { + /// use tokio::runtime::Handle; + /// use tokio::time::{timeout, Duration}; + /// + /// // Inside an async block or function. + /// let handle = Handle::current(); + /// if let Ok(dump) = timeout(Duration::from_secs(2), handle.dump()).await { + /// for (i, task) in dump.tasks().iter().enumerate() { + /// let trace = task.trace(); + /// println!("TASK {i}:"); + /// println!("{trace}\n"); + /// } + /// } + /// # }); + /// # } + /// ``` + /// + /// This produces highly detailed traces of tasks; e.g.: + /// + /// ```plain + /// TASK 0: + /// ╼ dump::main::{{closure}}::a::{{closure}} at /tokio/examples/dump.rs:18:20 + /// └╼ dump::main::{{closure}}::b::{{closure}} at /tokio/examples/dump.rs:23:20 + /// └╼ dump::main::{{closure}}::c::{{closure}} at /tokio/examples/dump.rs:28:24 + /// └╼ tokio::sync::barrier::Barrier::wait::{{closure}} at /tokio/tokio/src/sync/barrier.rs:129:10 + /// └╼ <tokio::util::trace::InstrumentedAsyncOp<F> as core::future::future::Future>::poll at /tokio/tokio/src/util/trace.rs:77:46 + /// └╼ tokio::sync::barrier::Barrier::wait_internal::{{closure}} at /tokio/tokio/src/sync/barrier.rs:183:36 + /// └╼ tokio::sync::watch::Receiver<T>::changed::{{closure}} at /tokio/tokio/src/sync/watch.rs:604:55 + /// └╼ tokio::sync::watch::changed_impl::{{closure}} at /tokio/tokio/src/sync/watch.rs:755:18 + /// └╼ <tokio::sync::notify::Notified as core::future::future::Future>::poll at /tokio/tokio/src/sync/notify.rs:1103:9 + /// └╼ tokio::sync::notify::Notified::poll_notified at /tokio/tokio/src/sync/notify.rs:996:32 + /// ``` + /// + /// # Requirements + /// + /// ## Debug Info Must Be Available + /// + /// To produce task traces, the application must **not** be compiled + /// with split debuginfo. On Linux, including debuginfo within the + /// application binary is the (correct) default. You can further ensure + /// this behavior with the following directive in your `Cargo.toml`: + /// + /// ```toml + /// [profile.*] + /// split-debuginfo = "off" + /// ``` + /// + /// ## Unstable Features + /// + /// This functionality is **unstable**, and requires both the + /// `tokio_unstable` and `tokio_taskdump` cfg flags to be set. + /// + /// You can do this by setting the `RUSTFLAGS` environment variable + /// before invoking `cargo`; e.g.: + /// ```bash + /// RUSTFLAGS="--cfg tokio_unstable --cfg tokio_taskdump" cargo run --example dump + /// ``` + /// + /// Or by [configuring][cargo-config] `rustflags` in + /// `.cargo/config.toml`: + /// ```text + /// [build] + /// rustflags = ["--cfg tokio_unstable", "--cfg tokio_taskdump"] + /// ``` + /// + /// [cargo-config]: + /// https://doc.rust-lang.org/cargo/reference/config.html + /// + /// ## Platform Requirements + /// + /// Task dumps are supported on Linux atop aarch64, x86 and x86_64. + /// + /// ## Current Thread Runtime Requirements + /// + /// On the `current_thread` runtime, task dumps may only be requested + /// from *within* the context of the runtime being dumped. Do not, for + /// example, await `Handle::dump()` on a different runtime. + /// + /// # Limitations + /// + /// ## Performance + /// + /// Although enabling the `tokio_taskdump` feature imposes virtually no + /// additional runtime overhead, actually calling `Handle::dump` is + /// expensive. The runtime must synchronize and pause its workers, then + /// re-poll every task in a special tracing mode. Avoid requesting dumps + /// often. + /// + /// ## Local Executors + /// + /// Tasks managed by local executors (e.g., `FuturesUnordered` and + /// [`LocalSet`][crate::task::LocalSet]) may not appear in task dumps. + /// + /// ## Non-Termination When Workers Are Blocked + /// + /// The future produced by `Handle::dump` may never produce `Ready` if + /// another runtime worker is blocked for more than 250ms. This may + /// occur if a dump is requested during shutdown, or if another runtime + /// worker is infinite looping or synchronously deadlocked. For these + /// reasons, task dumping should usually be paired with an explicit + /// [timeout][crate::time::timeout]. + pub async fn dump(&self) -> crate::runtime::Dump { + match &self.inner { + scheduler::Handle::CurrentThread(handle) => handle.dump(), + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + scheduler::Handle::MultiThread(handle) => { + // perform the trace in a separate thread so that the + // trace itself does not appear in the taskdump. + let handle = handle.clone(); + spawn_thread(async { + let handle = handle; + handle.dump().await + }).await + }, + } + } + } + + cfg_rt_multi_thread! { + /// Spawn a new thread and asynchronously await on its result. + async fn spawn_thread<F>(f: F) -> <F as Future>::Output + where + F: Future + Send + 'static, + <F as Future>::Output: Send + 'static + { + let (tx, rx) = crate::sync::oneshot::channel(); + crate::loom::thread::spawn(|| { + let rt = crate::runtime::Builder::new_current_thread().build().unwrap(); + rt.block_on(async { + let _ = tx.send(f.await); + }); + }); + rx.await.unwrap() + } } } /// Error returned by `try_current` when no Runtime has been started -pub struct TryCurrentError(()); +#[derive(Debug)] +pub struct TryCurrentError { + kind: TryCurrentErrorKind, +} + +impl TryCurrentError { + pub(crate) fn new_no_context() -> Self { + Self { + kind: TryCurrentErrorKind::NoContext, + } + } -impl fmt::Debug for TryCurrentError { + pub(crate) fn new_thread_local_destroyed() -> Self { + Self { + kind: TryCurrentErrorKind::ThreadLocalDestroyed, + } + } + + /// Returns true if the call failed because there is currently no runtime in + /// the Tokio context. + pub fn is_missing_context(&self) -> bool { + matches!(self.kind, TryCurrentErrorKind::NoContext) + } + + /// Returns true if the call failed because the Tokio context thread-local + /// had been destroyed. This can usually only happen if in the destructor of + /// other thread-locals. + pub fn is_thread_local_destroyed(&self) -> bool { + matches!(self.kind, TryCurrentErrorKind::ThreadLocalDestroyed) + } +} + +enum TryCurrentErrorKind { + NoContext, + ThreadLocalDestroyed, +} + +impl fmt::Debug for TryCurrentErrorKind { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("TryCurrentError").finish() + use TryCurrentErrorKind::*; + match self { + NoContext => f.write_str("NoContext"), + ThreadLocalDestroyed => f.write_str("ThreadLocalDestroyed"), + } } } impl fmt::Display for TryCurrentError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str(CONTEXT_MISSING_ERROR) + use TryCurrentErrorKind::*; + match self.kind { + NoContext => f.write_str(CONTEXT_MISSING_ERROR), + ThreadLocalDestroyed => f.write_str(THREAD_LOCAL_DESTROYED_ERROR), + } } } diff --git a/vendor/tokio/src/runtime/io/metrics.rs b/vendor/tokio/src/runtime/io/metrics.rs new file mode 100644 index 000000000..ec341efe6 --- /dev/null +++ b/vendor/tokio/src/runtime/io/metrics.rs @@ -0,0 +1,24 @@ +//! This file contains mocks of the metrics types used in the I/O driver. +//! +//! The reason these mocks don't live in `src/runtime/mock.rs` is because +//! these need to be available in the case when `net` is enabled but +//! `rt` is not. + +cfg_not_rt_and_metrics_and_net! { + #[derive(Default)] + pub(crate) struct IoDriverMetrics {} + + impl IoDriverMetrics { + pub(crate) fn incr_fd_count(&self) {} + pub(crate) fn dec_fd_count(&self) {} + pub(crate) fn incr_ready_count_by(&self, _amt: u64) {} + } +} + +cfg_net! { + cfg_rt! { + cfg_metrics! { + pub(crate) use crate::runtime::IoDriverMetrics; + } + } +} diff --git a/vendor/tokio/src/runtime/io/mod.rs b/vendor/tokio/src/runtime/io/mod.rs new file mode 100644 index 000000000..2dd426f11 --- /dev/null +++ b/vendor/tokio/src/runtime/io/mod.rs @@ -0,0 +1,356 @@ +#![cfg_attr(not(all(feature = "rt", feature = "net")), allow(dead_code))] + +mod registration; +pub(crate) use registration::Registration; + +mod scheduled_io; +use scheduled_io::ScheduledIo; + +mod metrics; + +use crate::io::interest::Interest; +use crate::io::ready::Ready; +use crate::runtime::driver; +use crate::util::slab::{self, Slab}; +use crate::{loom::sync::RwLock, util::bit}; + +use metrics::IoDriverMetrics; + +use std::fmt; +use std::io; +use std::time::Duration; + +/// I/O driver, backed by Mio. +pub(crate) struct Driver { + /// Tracks the number of times `turn` is called. It is safe for this to wrap + /// as it is mostly used to determine when to call `compact()`. + tick: u8, + + /// True when an event with the signal token is received + signal_ready: bool, + + /// Reuse the `mio::Events` value across calls to poll. + events: mio::Events, + + /// Primary slab handle containing the state for each resource registered + /// with this driver. + resources: Slab<ScheduledIo>, + + /// The system event queue. + poll: mio::Poll, +} + +/// A reference to an I/O driver. +pub(crate) struct Handle { + /// Registers I/O resources. + registry: mio::Registry, + + /// Allocates `ScheduledIo` handles when creating new resources. + io_dispatch: RwLock<IoDispatcher>, + + /// Used to wake up the reactor from a call to `turn`. + /// Not supported on Wasi due to lack of threading support. + #[cfg(not(tokio_wasi))] + waker: mio::Waker, + + pub(crate) metrics: IoDriverMetrics, +} + +#[derive(Debug)] +pub(crate) struct ReadyEvent { + tick: u8, + pub(crate) ready: Ready, + is_shutdown: bool, +} + +cfg_net_unix!( + impl ReadyEvent { + pub(crate) fn with_ready(&self, ready: Ready) -> Self { + Self { + ready, + tick: self.tick, + is_shutdown: self.is_shutdown, + } + } + } +); + +struct IoDispatcher { + allocator: slab::Allocator<ScheduledIo>, + is_shutdown: bool, +} + +#[derive(Debug, Eq, PartialEq, Clone, Copy)] +enum Direction { + Read, + Write, +} + +enum Tick { + Set(u8), + Clear(u8), +} + +// TODO: Don't use a fake token. Instead, reserve a slot entry for the wakeup +// token. +const TOKEN_WAKEUP: mio::Token = mio::Token(1 << 31); +const TOKEN_SIGNAL: mio::Token = mio::Token(1 + (1 << 31)); + +const ADDRESS: bit::Pack = bit::Pack::least_significant(24); + +// Packs the generation value in the `readiness` field. +// +// The generation prevents a race condition where a slab slot is reused for a +// new socket while the I/O driver is about to apply a readiness event. The +// generation value is checked when setting new readiness. If the generation do +// not match, then the readiness event is discarded. +const GENERATION: bit::Pack = ADDRESS.then(7); + +fn _assert_kinds() { + fn _assert<T: Send + Sync>() {} + + _assert::<Handle>(); +} + +// ===== impl Driver ===== + +impl Driver { + /// Creates a new event loop, returning any error that happened during the + /// creation. + pub(crate) fn new(nevents: usize) -> io::Result<(Driver, Handle)> { + let poll = mio::Poll::new()?; + #[cfg(not(tokio_wasi))] + let waker = mio::Waker::new(poll.registry(), TOKEN_WAKEUP)?; + let registry = poll.registry().try_clone()?; + + let slab = Slab::new(); + let allocator = slab.allocator(); + + let driver = Driver { + tick: 0, + signal_ready: false, + events: mio::Events::with_capacity(nevents), + poll, + resources: slab, + }; + + let handle = Handle { + registry, + io_dispatch: RwLock::new(IoDispatcher::new(allocator)), + #[cfg(not(tokio_wasi))] + waker, + metrics: IoDriverMetrics::default(), + }; + + Ok((driver, handle)) + } + + pub(crate) fn park(&mut self, rt_handle: &driver::Handle) { + let handle = rt_handle.io(); + self.turn(handle, None); + } + + pub(crate) fn park_timeout(&mut self, rt_handle: &driver::Handle, duration: Duration) { + let handle = rt_handle.io(); + self.turn(handle, Some(duration)); + } + + pub(crate) fn shutdown(&mut self, rt_handle: &driver::Handle) { + let handle = rt_handle.io(); + + if handle.shutdown() { + self.resources.for_each(|io| { + // If a task is waiting on the I/O resource, notify it that the + // runtime is being shutdown. And shutdown will clear all wakers. + io.shutdown(); + }); + } + } + + fn turn(&mut self, handle: &Handle, max_wait: Option<Duration>) { + // How often to call `compact()` on the resource slab + const COMPACT_INTERVAL: u8 = 255; + + self.tick = self.tick.wrapping_add(1); + + if self.tick == COMPACT_INTERVAL { + self.resources.compact() + } + + let events = &mut self.events; + + // Block waiting for an event to happen, peeling out how many events + // happened. + match self.poll.poll(events, max_wait) { + Ok(_) => {} + Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {} + #[cfg(tokio_wasi)] + Err(e) if e.kind() == io::ErrorKind::InvalidInput => { + // In case of wasm32_wasi this error happens, when trying to poll without subscriptions + // just return from the park, as there would be nothing, which wakes us up. + } + Err(e) => panic!("unexpected error when polling the I/O driver: {:?}", e), + } + + // Process all the events that came in, dispatching appropriately + let mut ready_count = 0; + for event in events.iter() { + let token = event.token(); + + if token == TOKEN_WAKEUP { + // Nothing to do, the event is used to unblock the I/O driver + } else if token == TOKEN_SIGNAL { + self.signal_ready = true; + } else { + Self::dispatch( + &mut self.resources, + self.tick, + token, + Ready::from_mio(event), + ); + ready_count += 1; + } + } + + handle.metrics.incr_ready_count_by(ready_count); + } + + fn dispatch(resources: &mut Slab<ScheduledIo>, tick: u8, token: mio::Token, ready: Ready) { + let addr = slab::Address::from_usize(ADDRESS.unpack(token.0)); + + let io = match resources.get(addr) { + Some(io) => io, + None => return, + }; + + let res = io.set_readiness(Some(token.0), Tick::Set(tick), |curr| curr | ready); + + if res.is_err() { + // token no longer valid! + return; + } + + io.wake(ready); + } +} + +impl fmt::Debug for Driver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "Driver") + } +} + +impl Handle { + /// Forces a reactor blocked in a call to `turn` to wakeup, or otherwise + /// makes the next call to `turn` return immediately. + /// + /// This method is intended to be used in situations where a notification + /// needs to otherwise be sent to the main reactor. If the reactor is + /// currently blocked inside of `turn` then it will wake up and soon return + /// after this method has been called. If the reactor is not currently + /// blocked in `turn`, then the next call to `turn` will not block and + /// return immediately. + pub(crate) fn unpark(&self) { + #[cfg(not(tokio_wasi))] + self.waker.wake().expect("failed to wake I/O driver"); + } + + /// Registers an I/O resource with the reactor for a given `mio::Ready` state. + /// + /// The registration token is returned. + pub(super) fn add_source( + &self, + source: &mut impl mio::event::Source, + interest: Interest, + ) -> io::Result<slab::Ref<ScheduledIo>> { + let (address, shared) = self.allocate()?; + + let token = GENERATION.pack(shared.generation(), ADDRESS.pack(address.as_usize(), 0)); + + self.registry + .register(source, mio::Token(token), interest.to_mio())?; + + self.metrics.incr_fd_count(); + + Ok(shared) + } + + /// Deregisters an I/O resource from the reactor. + pub(super) fn deregister_source(&self, source: &mut impl mio::event::Source) -> io::Result<()> { + self.registry.deregister(source)?; + + self.metrics.dec_fd_count(); + + Ok(()) + } + + /// shutdown the dispatcher. + fn shutdown(&self) -> bool { + let mut io = self.io_dispatch.write().unwrap(); + if io.is_shutdown { + return false; + } + io.is_shutdown = true; + true + } + + fn allocate(&self) -> io::Result<(slab::Address, slab::Ref<ScheduledIo>)> { + let io = self.io_dispatch.read().unwrap(); + if io.is_shutdown { + return Err(io::Error::new( + io::ErrorKind::Other, + crate::util::error::RUNTIME_SHUTTING_DOWN_ERROR, + )); + } + io.allocator.allocate().ok_or_else(|| { + io::Error::new( + io::ErrorKind::Other, + "reactor at max registered I/O resources", + ) + }) + } +} + +impl fmt::Debug for Handle { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "Handle") + } +} + +// ===== impl IoDispatcher ===== + +impl IoDispatcher { + fn new(allocator: slab::Allocator<ScheduledIo>) -> Self { + Self { + allocator, + is_shutdown: false, + } + } +} + +impl Direction { + pub(super) fn mask(self) -> Ready { + match self { + Direction::Read => Ready::READABLE | Ready::READ_CLOSED, + Direction::Write => Ready::WRITABLE | Ready::WRITE_CLOSED, + } + } +} + +// Signal handling +cfg_signal_internal_and_unix! { + impl Handle { + pub(crate) fn register_signal_receiver(&self, receiver: &mut mio::net::UnixStream) -> io::Result<()> { + self.registry.register(receiver, TOKEN_SIGNAL, mio::Interest::READABLE)?; + Ok(()) + } + } + + impl Driver { + pub(crate) fn consume_signal_ready(&mut self) -> bool { + let ret = self.signal_ready; + self.signal_ready = false; + ret + } + } +} diff --git a/vendor/tokio/src/io/driver/registration.rs b/vendor/tokio/src/runtime/io/registration.rs index 7350be634..341fa0539 100644 --- a/vendor/tokio/src/io/driver/registration.rs +++ b/vendor/tokio/src/runtime/io/registration.rs @@ -1,6 +1,8 @@ #![cfg_attr(not(feature = "net"), allow(dead_code))] -use crate::io::driver::{Direction, Handle, Interest, ReadyEvent, ScheduledIo}; +use crate::io::interest::Interest; +use crate::runtime::io::{Direction, Handle, ReadyEvent, ScheduledIo}; +use crate::runtime::scheduler; use crate::util::slab; use mio::event::Source; @@ -42,8 +44,8 @@ cfg_io_driver! { /// [`poll_write_ready`]: method@Self::poll_write_ready` #[derive(Debug)] pub(crate) struct Registration { - /// Handle to the associated driver. - handle: Handle, + /// Handle to the associated runtime. + handle: scheduler::Handle, /// Reference to state stored by the driver. shared: slab::Ref<ScheduledIo>, @@ -56,10 +58,8 @@ unsafe impl Sync for Registration {} // ===== impl Registration ===== impl Registration { - /// Registers the I/O resource with the default reactor, for a specific - /// `Interest`. `new_with_interest` should be used over `new` when you need - /// control over the readiness state, such as when a file descriptor only - /// allows reads. This does not add `hup` or `error` so if you are + /// Registers the I/O resource with the reactor for the provided handle, for + /// a specific `Interest`. This does not add `hup` or `error` so if you are /// interested in those states, you will need to add them to the readiness /// state passed to this function. /// @@ -67,19 +67,13 @@ impl Registration { /// /// - `Ok` if the registration happened successfully /// - `Err` if an error was encountered during registration + #[track_caller] pub(crate) fn new_with_interest_and_handle( io: &mut impl Source, interest: Interest, - handle: Handle, + handle: scheduler::Handle, ) -> io::Result<Registration> { - let shared = if let Some(inner) = handle.inner() { - inner.add_source(io, interest)? - } else { - return Err(io::Error::new( - io::ErrorKind::Other, - "failed to find event loop", - )); - }; + let shared = handle.driver().io().add_source(io, interest)?; Ok(Registration { handle, shared }) } @@ -101,11 +95,7 @@ impl Registration { /// /// `Err` is returned if an error is encountered. pub(crate) fn deregister(&mut self, io: &mut impl Source) -> io::Result<()> { - let inner = match self.handle.inner() { - Some(inner) => inner, - None => return Err(io::Error::new(io::ErrorKind::Other, "reactor gone")), - }; - inner.deregister_source(io) + self.handle().deregister_source(io) } pub(crate) fn clear_readiness(&self, event: ReadyEvent) { @@ -126,6 +116,7 @@ impl Registration { // Uses the poll path, requiring the caller to ensure mutual exclusion for // correctness. Only the last task to call this function is notified. + #[cfg(not(tokio_wasi))] pub(crate) fn poll_read_io<R>( &self, cx: &mut Context<'_>, @@ -153,11 +144,12 @@ impl Registration { cx: &mut Context<'_>, direction: Direction, ) -> Poll<io::Result<ReadyEvent>> { + ready!(crate::trace::trace_leaf(cx)); // Keep track of task budget - let coop = ready!(crate::coop::poll_proceed(cx)); + let coop = ready!(crate::runtime::coop::poll_proceed(cx)); let ev = ready!(self.shared.poll_readiness(cx, direction)); - if self.handle.inner().is_none() { + if ev.is_shutdown { return Poll::Ready(Err(gone())); } @@ -206,6 +198,10 @@ impl Registration { res => res, } } + + fn handle(&self) -> &Handle { + self.handle.driver().io() + } } impl Drop for Registration { @@ -222,28 +218,22 @@ impl Drop for Registration { } fn gone() -> io::Error { - io::Error::new(io::ErrorKind::Other, "IO driver has terminated") + io::Error::new( + io::ErrorKind::Other, + crate::util::error::RUNTIME_SHUTTING_DOWN_ERROR, + ) } cfg_io_readiness! { impl Registration { pub(crate) async fn readiness(&self, interest: Interest) -> io::Result<ReadyEvent> { - use std::future::Future; - use std::pin::Pin; - - let fut = self.shared.readiness(interest); - pin!(fut); + let ev = self.shared.readiness(interest).await; - crate::future::poll_fn(|cx| { - if self.handle.inner().is_none() { - return Poll::Ready(Err(io::Error::new( - io::ErrorKind::Other, - crate::util::error::RUNTIME_SHUTTING_DOWN_ERROR - ))); - } + if ev.is_shutdown { + return Err(gone()) + } - Pin::new(&mut fut).poll(cx).map(Ok) - }).await + Ok(ev) } pub(crate) async fn async_io<R>(&self, interest: Interest, mut f: impl FnMut() -> io::Result<R>) -> io::Result<R> { diff --git a/vendor/tokio/src/io/driver/scheduled_io.rs b/vendor/tokio/src/runtime/io/scheduled_io.rs index 517801079..197a4e0e2 100644 --- a/vendor/tokio/src/io/driver/scheduled_io.rs +++ b/vendor/tokio/src/runtime/io/scheduled_io.rs @@ -1,8 +1,11 @@ -use super::{Interest, Ready, ReadyEvent, Tick}; +use super::{ReadyEvent, Tick}; +use crate::io::interest::Interest; +use crate::io::ready::Ready; use crate::loom::sync::atomic::AtomicUsize; use crate::loom::sync::Mutex; use crate::util::bit; use crate::util::slab::Entry; +use crate::util::WakeList; use std::sync::atomic::Ordering::{AcqRel, Acquire, Release}; use std::task::{Context, Poll, Waker}; @@ -35,17 +38,14 @@ cfg_io_readiness! { #[derive(Debug, Default)] struct Waiters { #[cfg(feature = "net")] - /// List of all current waiters + /// List of all current waiters. list: WaitList, - /// Waker used for AsyncRead + /// Waker used for AsyncRead. reader: Option<Waker>, - /// Waker used for AsyncWrite + /// Waker used for AsyncWrite. writer: Option<Waker>, - - /// True if this ScheduledIo has been killed due to IO driver shutdown - is_shutdown: bool, } cfg_io_readiness! { @@ -53,19 +53,27 @@ cfg_io_readiness! { struct Waiter { pointers: linked_list::Pointers<Waiter>, - /// The waker for this task + /// The waker for this task. waker: Option<Waker>, - /// The interest this waiter is waiting on + /// The interest this waiter is waiting on. interest: Interest, is_ready: bool, - /// Should never be `!Unpin` + /// Should never be `!Unpin`. _p: PhantomPinned, } - /// Future returned by `readiness()` + generate_addr_of_methods! { + impl<> Waiter { + unsafe fn addr_of_pointers(self: NonNull<Self>) -> NonNull<linked_list::Pointers<Waiter>> { + &self.pointers + } + } + } + + /// Future returned by `readiness()`. struct Readiness<'a> { scheduled_io: &'a ScheduledIo, @@ -84,7 +92,7 @@ cfg_io_readiness! { // The `ScheduledIo::readiness` (`AtomicUsize`) is packed full of goodness. // -// | reserved | generation | driver tick | readiness | +// | shutdown | generation | driver tick | readiness | // |----------+------------+--------------+-----------| // | 1 bit | 7 bits + 8 bits + 16 bits | @@ -94,6 +102,8 @@ const TICK: bit::Pack = READINESS.then(8); const GENERATION: bit::Pack = TICK.then(7); +const SHUTDOWN: bit::Pack = GENERATION.then(1); + #[test] fn test_generations_assert_same() { assert_eq!(super::GENERATION, GENERATION); @@ -127,9 +137,11 @@ impl ScheduledIo { } /// Invoked when the IO driver is shut down; forces this ScheduledIo into a - /// permanently ready state. + /// permanently shutdown state. pub(super) fn shutdown(&self) { - self.wake0(Ready::ALL, true) + let mask = SHUTDOWN.pack(1, 0); + self.readiness.fetch_or(mask, AcqRel); + self.wake(Ready::ALL); } /// Sets the readiness on this `ScheduledIo` by invoking the given closure on @@ -208,32 +220,21 @@ impl ScheduledIo { /// than 32 wakers to notify, if the stack array fills up, the lock is /// released, the array is cleared, and the iteration continues. pub(super) fn wake(&self, ready: Ready) { - self.wake0(ready, false); - } - - fn wake0(&self, ready: Ready, shutdown: bool) { - const NUM_WAKERS: usize = 32; - - let mut wakers: [Option<Waker>; NUM_WAKERS] = Default::default(); - let mut curr = 0; + let mut wakers = WakeList::new(); let mut waiters = self.waiters.lock(); - waiters.is_shutdown |= shutdown; - // check for AsyncRead slot if ready.is_readable() { if let Some(waker) = waiters.reader.take() { - wakers[curr] = Some(waker); - curr += 1; + wakers.push(waker); } } // check for AsyncWrite slot if ready.is_writable() { if let Some(waker) = waiters.writer.take() { - wakers[curr] = Some(waker); - curr += 1; + wakers.push(waker); } } @@ -241,15 +242,14 @@ impl ScheduledIo { 'outer: loop { let mut iter = waiters.list.drain_filter(|w| ready.satisfies(w.interest)); - while curr < NUM_WAKERS { + while wakers.can_push() { match iter.next() { Some(waiter) => { let waiter = unsafe { &mut *waiter.as_ptr() }; if let Some(waker) = waiter.waker.take() { waiter.is_ready = true; - wakers[curr] = Some(waker); - curr += 1; + wakers.push(waker); } } None => { @@ -260,11 +260,7 @@ impl ScheduledIo { drop(waiters); - for waker in wakers.iter_mut().take(curr) { - waker.take().unwrap().wake(); - } - - curr = 0; + wakers.wake_all(); // Acquire the lock again. waiters = self.waiters.lock(); @@ -273,9 +269,7 @@ impl ScheduledIo { // Release the lock before notifying drop(waiters); - for waker in wakers.iter_mut().take(curr) { - waker.take().unwrap().wake(); - } + wakers.wake_all(); } pub(super) fn ready_event(&self, interest: Interest) -> ReadyEvent { @@ -284,10 +278,11 @@ impl ScheduledIo { ReadyEvent { tick: TICK.unpack(curr) as u8, ready: interest.mask() & Ready::from_usize(READINESS.unpack(curr)), + is_shutdown: SHUTDOWN.unpack(curr) != 0, } } - /// Poll version of checking readiness for a certain direction. + /// Polls for readiness events in a given direction. /// /// These are to support `AsyncRead` and `AsyncWrite` polling methods, /// which cannot use the `async fn` version. This uses reserved reader @@ -300,8 +295,9 @@ impl ScheduledIo { let curr = self.readiness.load(Acquire); let ready = direction.mask() & Ready::from_usize(READINESS.unpack(curr)); + let is_shutdown = SHUTDOWN.unpack(curr) != 0; - if ready.is_empty() { + if ready.is_empty() && !is_shutdown { // Update the task info let mut waiters = self.waiters.lock(); let slot = match direction { @@ -326,10 +322,12 @@ impl ScheduledIo { // taking the waiters lock let curr = self.readiness.load(Acquire); let ready = direction.mask() & Ready::from_usize(READINESS.unpack(curr)); - if waiters.is_shutdown { + let is_shutdown = SHUTDOWN.unpack(curr) != 0; + if is_shutdown { Poll::Ready(ReadyEvent { tick: TICK.unpack(curr) as u8, ready: direction.mask(), + is_shutdown, }) } else if ready.is_empty() { Poll::Pending @@ -337,12 +335,14 @@ impl ScheduledIo { Poll::Ready(ReadyEvent { tick: TICK.unpack(curr) as u8, ready, + is_shutdown, }) } } else { Poll::Ready(ReadyEvent { tick: TICK.unpack(curr) as u8, ready, + is_shutdown, }) } } @@ -374,7 +374,7 @@ unsafe impl Sync for ScheduledIo {} cfg_io_readiness! { impl ScheduledIo { - /// An async version of `poll_readiness` which uses a linked list of wakers + /// An async version of `poll_readiness` which uses a linked list of wakers. pub(crate) async fn readiness(&self, interest: Interest) -> ReadyEvent { self.readiness_fut(interest).await } @@ -410,8 +410,8 @@ cfg_io_readiness! { ptr } - unsafe fn pointers(mut target: NonNull<Waiter>) -> NonNull<linked_list::Pointers<Waiter>> { - NonNull::from(&mut target.as_mut().pointers) + unsafe fn pointers(target: NonNull<Waiter>) -> NonNull<linked_list::Pointers<Waiter>> { + Waiter::addr_of_pointers(target) } } @@ -434,16 +434,17 @@ cfg_io_readiness! { // Optimistically check existing readiness let curr = scheduled_io.readiness.load(SeqCst); let ready = Ready::from_usize(READINESS.unpack(curr)); + let is_shutdown = SHUTDOWN.unpack(curr) != 0; // Safety: `waiter.interest` never changes let interest = unsafe { (*waiter.get()).interest }; let ready = ready.intersection(interest); - if !ready.is_empty() { + if !ready.is_empty() || is_shutdown { // Currently ready! let tick = TICK.unpack(curr) as u8; *state = State::Done; - return Poll::Ready(ReadyEvent { tick, ready }); + return Poll::Ready(ReadyEvent { tick, ready, is_shutdown }); } // Wasn't ready, take the lock (and check again while locked). @@ -451,18 +452,19 @@ cfg_io_readiness! { let curr = scheduled_io.readiness.load(SeqCst); let mut ready = Ready::from_usize(READINESS.unpack(curr)); + let is_shutdown = SHUTDOWN.unpack(curr) != 0; - if waiters.is_shutdown { + if is_shutdown { ready = Ready::ALL; } let ready = ready.intersection(interest); - if !ready.is_empty() { + if !ready.is_empty() || is_shutdown { // Currently ready! let tick = TICK.unpack(curr) as u8; *state = State::Done; - return Poll::Ready(ReadyEvent { tick, ready }); + return Poll::Ready(ReadyEvent { tick, ready, is_shutdown }); } // Not ready even after locked, insert into list... @@ -511,14 +513,26 @@ cfg_io_readiness! { drop(waiters); } State::Done => { - let tick = TICK.unpack(scheduled_io.readiness.load(Acquire)) as u8; - // Safety: State::Done means it is no longer shared let w = unsafe { &mut *waiter.get() }; + let curr = scheduled_io.readiness.load(Acquire); + let is_shutdown = SHUTDOWN.unpack(curr) != 0; + + // The returned tick might be newer than the event + // which notified our waker. This is ok because the future + // still didn't return `Poll::Ready`. + let tick = TICK.unpack(curr) as u8; + + // The readiness state could have been cleared in the meantime, + // but we allow the returned ready set to be empty. + let curr_ready = Ready::from_usize(READINESS.unpack(curr)); + let ready = curr_ready.intersection(w.interest); + return Poll::Ready(ReadyEvent { tick, - ready: Ready::from_interest(w.interest), + ready, + is_shutdown, }); } } diff --git a/vendor/tokio/src/runtime/metrics/batch.rs b/vendor/tokio/src/runtime/metrics/batch.rs new file mode 100644 index 000000000..1bb4e261f --- /dev/null +++ b/vendor/tokio/src/runtime/metrics/batch.rs @@ -0,0 +1,162 @@ +use crate::runtime::metrics::{HistogramBatch, WorkerMetrics}; + +use std::sync::atomic::Ordering::Relaxed; +use std::time::{Duration, Instant}; + +pub(crate) struct MetricsBatch { + /// Number of times the worker parked. + park_count: u64, + + /// Number of times the worker woke w/o doing work. + noop_count: u64, + + /// Number of tasks stolen. + steal_count: u64, + + /// Number of times tasks where stolen. + steal_operations: u64, + + /// Number of tasks that were polled by the worker. + poll_count: u64, + + /// Number of tasks polled when the worker entered park. This is used to + /// track the noop count. + poll_count_on_last_park: u64, + + /// Number of tasks that were scheduled locally on this worker. + local_schedule_count: u64, + + /// Number of tasks moved to the global queue to make space in the local + /// queue + overflow_count: u64, + + /// The total busy duration in nanoseconds. + busy_duration_total: u64, + + /// Instant at which work last resumed (continued after park). + processing_scheduled_tasks_started_at: Instant, + + /// If `Some`, tracks poll times in nanoseconds + poll_timer: Option<PollTimer>, +} + +struct PollTimer { + /// Histogram of poll counts within each band. + poll_counts: HistogramBatch, + + /// Instant when the most recent task started polling. + poll_started_at: Instant, +} + +impl MetricsBatch { + pub(crate) fn new(worker_metrics: &WorkerMetrics) -> MetricsBatch { + let now = Instant::now(); + + MetricsBatch { + park_count: 0, + noop_count: 0, + steal_count: 0, + steal_operations: 0, + poll_count: 0, + poll_count_on_last_park: 0, + local_schedule_count: 0, + overflow_count: 0, + busy_duration_total: 0, + processing_scheduled_tasks_started_at: now, + poll_timer: worker_metrics + .poll_count_histogram + .as_ref() + .map(|worker_poll_counts| PollTimer { + poll_counts: HistogramBatch::from_histogram(worker_poll_counts), + poll_started_at: now, + }), + } + } + + pub(crate) fn submit(&mut self, worker: &WorkerMetrics) { + worker.park_count.store(self.park_count, Relaxed); + worker.noop_count.store(self.noop_count, Relaxed); + worker.steal_count.store(self.steal_count, Relaxed); + worker + .steal_operations + .store(self.steal_operations, Relaxed); + worker.poll_count.store(self.poll_count, Relaxed); + + worker + .busy_duration_total + .store(self.busy_duration_total, Relaxed); + + worker + .local_schedule_count + .store(self.local_schedule_count, Relaxed); + worker.overflow_count.store(self.overflow_count, Relaxed); + + if let Some(poll_timer) = &self.poll_timer { + let dst = worker.poll_count_histogram.as_ref().unwrap(); + poll_timer.poll_counts.submit(dst); + } + } + + /// The worker is about to park. + pub(crate) fn about_to_park(&mut self) { + self.park_count += 1; + + if self.poll_count_on_last_park == self.poll_count { + self.noop_count += 1; + } else { + self.poll_count_on_last_park = self.poll_count; + } + } + + /// Start processing a batch of tasks + pub(crate) fn start_processing_scheduled_tasks(&mut self) { + self.processing_scheduled_tasks_started_at = Instant::now(); + } + + /// Stop processing a batch of tasks + pub(crate) fn end_processing_scheduled_tasks(&mut self) { + let busy_duration = self.processing_scheduled_tasks_started_at.elapsed(); + self.busy_duration_total += duration_as_u64(busy_duration); + } + + /// Start polling an individual task + pub(crate) fn start_poll(&mut self) { + self.poll_count += 1; + + if let Some(poll_timer) = &mut self.poll_timer { + poll_timer.poll_started_at = Instant::now(); + } + } + + /// Stop polling an individual task + pub(crate) fn end_poll(&mut self) { + if let Some(poll_timer) = &mut self.poll_timer { + let elapsed = duration_as_u64(poll_timer.poll_started_at.elapsed()); + poll_timer.poll_counts.measure(elapsed, 1); + } + } + + pub(crate) fn inc_local_schedule_count(&mut self) { + self.local_schedule_count += 1; + } +} + +cfg_rt_multi_thread! { + impl MetricsBatch { + pub(crate) fn incr_steal_count(&mut self, by: u16) { + self.steal_count += by as u64; + } + + pub(crate) fn incr_steal_operations(&mut self) { + self.steal_operations += 1; + } + + pub(crate) fn incr_overflow_count(&mut self) { + self.overflow_count += 1; + } + } +} + +fn duration_as_u64(dur: Duration) -> u64 { + u64::try_from(dur.as_nanos()).unwrap_or(u64::MAX) +} diff --git a/vendor/tokio/src/runtime/metrics/histogram.rs b/vendor/tokio/src/runtime/metrics/histogram.rs new file mode 100644 index 000000000..976f54fe8 --- /dev/null +++ b/vendor/tokio/src/runtime/metrics/histogram.rs @@ -0,0 +1,502 @@ +use crate::loom::sync::atomic::{AtomicU64, Ordering::Relaxed}; + +use std::cmp; +use std::ops::Range; + +#[derive(Debug)] +pub(crate) struct Histogram { + /// The histogram buckets + buckets: Box<[AtomicU64]>, + + /// Bucket scale, linear or log + scale: HistogramScale, + + /// Minimum resolution + resolution: u64, +} + +#[derive(Debug, Clone)] +pub(crate) struct HistogramBuilder { + /// Histogram scale + pub(crate) scale: HistogramScale, + + /// Must be a power of 2 + pub(crate) resolution: u64, + + /// Number of buckets + pub(crate) num_buckets: usize, +} + +#[derive(Debug)] +pub(crate) struct HistogramBatch { + buckets: Box<[u64]>, + scale: HistogramScale, + resolution: u64, +} + +cfg_unstable! { + /// Whether the histogram used to aggregate a metric uses a linear or + /// logarithmic scale. + #[derive(Debug, Copy, Clone, Eq, PartialEq)] + #[non_exhaustive] + pub enum HistogramScale { + /// Linear bucket scale + Linear, + + /// Logarithmic bucket scale + Log, + } +} + +impl Histogram { + pub(crate) fn num_buckets(&self) -> usize { + self.buckets.len() + } + + pub(crate) fn get(&self, bucket: usize) -> u64 { + self.buckets[bucket].load(Relaxed) + } + + pub(crate) fn bucket_range(&self, bucket: usize) -> Range<u64> { + match self.scale { + HistogramScale::Log => Range { + start: if bucket == 0 { + 0 + } else { + self.resolution << (bucket - 1) + }, + end: if bucket == self.buckets.len() - 1 { + u64::MAX + } else { + self.resolution << bucket + }, + }, + HistogramScale::Linear => Range { + start: self.resolution * bucket as u64, + end: if bucket == self.buckets.len() - 1 { + u64::MAX + } else { + self.resolution * (bucket as u64 + 1) + }, + }, + } + } +} + +impl HistogramBatch { + pub(crate) fn from_histogram(histogram: &Histogram) -> HistogramBatch { + let buckets = vec![0; histogram.buckets.len()].into_boxed_slice(); + + HistogramBatch { + buckets, + scale: histogram.scale, + resolution: histogram.resolution, + } + } + + pub(crate) fn measure(&mut self, value: u64, count: u64) { + self.buckets[self.value_to_bucket(value)] += count; + } + + pub(crate) fn submit(&self, histogram: &Histogram) { + debug_assert_eq!(self.scale, histogram.scale); + debug_assert_eq!(self.resolution, histogram.resolution); + debug_assert_eq!(self.buckets.len(), histogram.buckets.len()); + + for i in 0..self.buckets.len() { + histogram.buckets[i].store(self.buckets[i], Relaxed); + } + } + + fn value_to_bucket(&self, value: u64) -> usize { + match self.scale { + HistogramScale::Linear => { + let max = self.buckets.len() - 1; + cmp::min(value / self.resolution, max as u64) as usize + } + HistogramScale::Log => { + let max = self.buckets.len() - 1; + + if value < self.resolution { + 0 + } else { + let significant_digits = 64 - value.leading_zeros(); + let bucket_digits = 64 - (self.resolution - 1).leading_zeros(); + cmp::min(significant_digits as usize - bucket_digits as usize, max) + } + } + } + } +} + +impl HistogramBuilder { + pub(crate) fn new() -> HistogramBuilder { + HistogramBuilder { + scale: HistogramScale::Linear, + // Resolution is in nanoseconds. + resolution: 100_000, + num_buckets: 10, + } + } + + pub(crate) fn build(&self) -> Histogram { + let mut resolution = self.resolution; + + assert!(resolution > 0); + + if matches!(self.scale, HistogramScale::Log) { + resolution = resolution.next_power_of_two(); + } + + Histogram { + buckets: (0..self.num_buckets) + .map(|_| AtomicU64::new(0)) + .collect::<Vec<_>>() + .into_boxed_slice(), + resolution, + scale: self.scale, + } + } +} + +impl Default for HistogramBuilder { + fn default() -> HistogramBuilder { + HistogramBuilder::new() + } +} + +#[cfg(test)] +mod test { + use super::*; + + macro_rules! assert_bucket_eq { + ($h:expr, $bucket:expr, $val:expr) => {{ + assert_eq!($h.buckets[$bucket], $val); + }}; + } + + #[test] + fn log_scale_resolution_1() { + let h = HistogramBuilder { + scale: HistogramScale::Log, + resolution: 1, + num_buckets: 10, + } + .build(); + + assert_eq!(h.bucket_range(0), 0..1); + assert_eq!(h.bucket_range(1), 1..2); + assert_eq!(h.bucket_range(2), 2..4); + assert_eq!(h.bucket_range(3), 4..8); + assert_eq!(h.bucket_range(9), 256..u64::MAX); + + let mut b = HistogramBatch::from_histogram(&h); + + b.measure(0, 1); + assert_bucket_eq!(b, 0, 1); + assert_bucket_eq!(b, 1, 0); + + b.measure(1, 1); + assert_bucket_eq!(b, 0, 1); + assert_bucket_eq!(b, 1, 1); + assert_bucket_eq!(b, 2, 0); + + b.measure(2, 1); + assert_bucket_eq!(b, 0, 1); + assert_bucket_eq!(b, 1, 1); + assert_bucket_eq!(b, 2, 1); + + b.measure(3, 1); + assert_bucket_eq!(b, 0, 1); + assert_bucket_eq!(b, 1, 1); + assert_bucket_eq!(b, 2, 2); + + b.measure(4, 1); + assert_bucket_eq!(b, 0, 1); + assert_bucket_eq!(b, 1, 1); + assert_bucket_eq!(b, 2, 2); + assert_bucket_eq!(b, 3, 1); + + b.measure(100, 1); + assert_bucket_eq!(b, 7, 1); + + b.measure(128, 1); + assert_bucket_eq!(b, 8, 1); + + b.measure(4096, 1); + assert_bucket_eq!(b, 9, 1); + } + + #[test] + fn log_scale_resolution_2() { + let h = HistogramBuilder { + scale: HistogramScale::Log, + resolution: 2, + num_buckets: 10, + } + .build(); + + assert_eq!(h.bucket_range(0), 0..2); + assert_eq!(h.bucket_range(1), 2..4); + assert_eq!(h.bucket_range(2), 4..8); + assert_eq!(h.bucket_range(3), 8..16); + assert_eq!(h.bucket_range(9), 512..u64::MAX); + + let mut b = HistogramBatch::from_histogram(&h); + + b.measure(0, 1); + assert_bucket_eq!(b, 0, 1); + assert_bucket_eq!(b, 1, 0); + + b.measure(1, 1); + assert_bucket_eq!(b, 0, 2); + assert_bucket_eq!(b, 1, 0); + + b.measure(2, 1); + assert_bucket_eq!(b, 0, 2); + assert_bucket_eq!(b, 1, 1); + assert_bucket_eq!(b, 2, 0); + + b.measure(3, 1); + assert_bucket_eq!(b, 0, 2); + assert_bucket_eq!(b, 1, 2); + assert_bucket_eq!(b, 2, 0); + + b.measure(4, 1); + assert_bucket_eq!(b, 0, 2); + assert_bucket_eq!(b, 1, 2); + assert_bucket_eq!(b, 2, 1); + + b.measure(5, 1); + assert_bucket_eq!(b, 0, 2); + assert_bucket_eq!(b, 1, 2); + assert_bucket_eq!(b, 2, 2); + + b.measure(6, 1); + assert_bucket_eq!(b, 0, 2); + assert_bucket_eq!(b, 1, 2); + assert_bucket_eq!(b, 2, 3); + + b.measure(7, 1); + assert_bucket_eq!(b, 0, 2); + assert_bucket_eq!(b, 1, 2); + assert_bucket_eq!(b, 2, 4); + + b.measure(8, 1); + assert_bucket_eq!(b, 0, 2); + assert_bucket_eq!(b, 1, 2); + assert_bucket_eq!(b, 2, 4); + assert_bucket_eq!(b, 3, 1); + + b.measure(100, 1); + assert_bucket_eq!(b, 6, 1); + + b.measure(128, 1); + assert_bucket_eq!(b, 7, 1); + + b.measure(4096, 1); + assert_bucket_eq!(b, 9, 1); + + for bucket in h.buckets.iter() { + assert_eq!(bucket.load(Relaxed), 0); + } + + b.submit(&h); + + for i in 0..h.buckets.len() { + assert_eq!(h.buckets[i].load(Relaxed), b.buckets[i]); + } + + b.submit(&h); + + for i in 0..h.buckets.len() { + assert_eq!(h.buckets[i].load(Relaxed), b.buckets[i]); + } + } + + #[test] + fn linear_scale_resolution_1() { + let h = HistogramBuilder { + scale: HistogramScale::Linear, + resolution: 1, + num_buckets: 10, + } + .build(); + + assert_eq!(h.bucket_range(0), 0..1); + assert_eq!(h.bucket_range(1), 1..2); + assert_eq!(h.bucket_range(2), 2..3); + assert_eq!(h.bucket_range(3), 3..4); + assert_eq!(h.bucket_range(9), 9..u64::MAX); + + let mut b = HistogramBatch::from_histogram(&h); + + b.measure(0, 1); + assert_bucket_eq!(b, 0, 1); + assert_bucket_eq!(b, 1, 0); + + b.measure(1, 1); + assert_bucket_eq!(b, 0, 1); + assert_bucket_eq!(b, 1, 1); + assert_bucket_eq!(b, 2, 0); + + b.measure(2, 1); + assert_bucket_eq!(b, 0, 1); + assert_bucket_eq!(b, 1, 1); + assert_bucket_eq!(b, 2, 1); + assert_bucket_eq!(b, 3, 0); + + b.measure(3, 1); + assert_bucket_eq!(b, 0, 1); + assert_bucket_eq!(b, 1, 1); + assert_bucket_eq!(b, 2, 1); + assert_bucket_eq!(b, 3, 1); + + b.measure(5, 1); + assert_bucket_eq!(b, 5, 1); + + b.measure(4096, 1); + assert_bucket_eq!(b, 9, 1); + + for bucket in h.buckets.iter() { + assert_eq!(bucket.load(Relaxed), 0); + } + + b.submit(&h); + + for i in 0..h.buckets.len() { + assert_eq!(h.buckets[i].load(Relaxed), b.buckets[i]); + } + + b.submit(&h); + + for i in 0..h.buckets.len() { + assert_eq!(h.buckets[i].load(Relaxed), b.buckets[i]); + } + } + + #[test] + fn linear_scale_resolution_100() { + let h = HistogramBuilder { + scale: HistogramScale::Linear, + resolution: 100, + num_buckets: 10, + } + .build(); + + assert_eq!(h.bucket_range(0), 0..100); + assert_eq!(h.bucket_range(1), 100..200); + assert_eq!(h.bucket_range(2), 200..300); + assert_eq!(h.bucket_range(3), 300..400); + assert_eq!(h.bucket_range(9), 900..u64::MAX); + + let mut b = HistogramBatch::from_histogram(&h); + + b.measure(0, 1); + assert_bucket_eq!(b, 0, 1); + assert_bucket_eq!(b, 1, 0); + + b.measure(50, 1); + assert_bucket_eq!(b, 0, 2); + assert_bucket_eq!(b, 1, 0); + + b.measure(100, 1); + assert_bucket_eq!(b, 0, 2); + assert_bucket_eq!(b, 1, 1); + assert_bucket_eq!(b, 2, 0); + + b.measure(101, 1); + assert_bucket_eq!(b, 0, 2); + assert_bucket_eq!(b, 1, 2); + assert_bucket_eq!(b, 2, 0); + + b.measure(200, 1); + assert_bucket_eq!(b, 0, 2); + assert_bucket_eq!(b, 1, 2); + assert_bucket_eq!(b, 2, 1); + + b.measure(299, 1); + assert_bucket_eq!(b, 0, 2); + assert_bucket_eq!(b, 1, 2); + assert_bucket_eq!(b, 2, 2); + + b.measure(222, 1); + assert_bucket_eq!(b, 0, 2); + assert_bucket_eq!(b, 1, 2); + assert_bucket_eq!(b, 2, 3); + + b.measure(300, 1); + assert_bucket_eq!(b, 0, 2); + assert_bucket_eq!(b, 1, 2); + assert_bucket_eq!(b, 2, 3); + assert_bucket_eq!(b, 3, 1); + + b.measure(888, 1); + assert_bucket_eq!(b, 8, 1); + + b.measure(4096, 1); + assert_bucket_eq!(b, 9, 1); + + for bucket in h.buckets.iter() { + assert_eq!(bucket.load(Relaxed), 0); + } + + b.submit(&h); + + for i in 0..h.buckets.len() { + assert_eq!(h.buckets[i].load(Relaxed), b.buckets[i]); + } + + b.submit(&h); + + for i in 0..h.buckets.len() { + assert_eq!(h.buckets[i].load(Relaxed), b.buckets[i]); + } + } + + #[test] + fn inc_by_more_than_one() { + let h = HistogramBuilder { + scale: HistogramScale::Linear, + resolution: 100, + num_buckets: 10, + } + .build(); + + let mut b = HistogramBatch::from_histogram(&h); + + b.measure(0, 3); + assert_bucket_eq!(b, 0, 3); + assert_bucket_eq!(b, 1, 0); + + b.measure(50, 5); + assert_bucket_eq!(b, 0, 8); + assert_bucket_eq!(b, 1, 0); + + b.measure(100, 2); + assert_bucket_eq!(b, 0, 8); + assert_bucket_eq!(b, 1, 2); + assert_bucket_eq!(b, 2, 0); + + b.measure(101, 19); + assert_bucket_eq!(b, 0, 8); + assert_bucket_eq!(b, 1, 21); + assert_bucket_eq!(b, 2, 0); + + for bucket in h.buckets.iter() { + assert_eq!(bucket.load(Relaxed), 0); + } + + b.submit(&h); + + for i in 0..h.buckets.len() { + assert_eq!(h.buckets[i].load(Relaxed), b.buckets[i]); + } + + b.submit(&h); + + for i in 0..h.buckets.len() { + assert_eq!(h.buckets[i].load(Relaxed), b.buckets[i]); + } + } +} diff --git a/vendor/tokio/src/runtime/metrics/io.rs b/vendor/tokio/src/runtime/metrics/io.rs new file mode 100644 index 000000000..06efdd42d --- /dev/null +++ b/vendor/tokio/src/runtime/metrics/io.rs @@ -0,0 +1,24 @@ +#![cfg_attr(not(feature = "net"), allow(dead_code))] + +use crate::loom::sync::atomic::{AtomicU64, Ordering::Relaxed}; + +#[derive(Default)] +pub(crate) struct IoDriverMetrics { + pub(super) fd_registered_count: AtomicU64, + pub(super) fd_deregistered_count: AtomicU64, + pub(super) ready_count: AtomicU64, +} + +impl IoDriverMetrics { + pub(crate) fn incr_fd_count(&self) { + self.fd_registered_count.fetch_add(1, Relaxed); + } + + pub(crate) fn dec_fd_count(&self) { + self.fd_deregistered_count.fetch_add(1, Relaxed); + } + + pub(crate) fn incr_ready_count_by(&self, amt: u64) { + self.ready_count.fetch_add(amt, Relaxed); + } +} diff --git a/vendor/tokio/src/runtime/metrics/mock.rs b/vendor/tokio/src/runtime/metrics/mock.rs new file mode 100644 index 000000000..8f8345c08 --- /dev/null +++ b/vendor/tokio/src/runtime/metrics/mock.rs @@ -0,0 +1,55 @@ +//! This file contains mocks of the types in src/runtime/metrics + +pub(crate) struct SchedulerMetrics {} + +pub(crate) struct WorkerMetrics {} + +pub(crate) struct MetricsBatch {} + +#[derive(Clone, Default)] +pub(crate) struct HistogramBuilder {} + +impl SchedulerMetrics { + pub(crate) fn new() -> Self { + Self {} + } + + /// Increment the number of tasks scheduled externally + pub(crate) fn inc_remote_schedule_count(&self) {} +} + +impl WorkerMetrics { + pub(crate) fn new() -> Self { + Self {} + } + + pub(crate) fn from_config(config: &crate::runtime::Config) -> Self { + // Prevent the dead-code warning from being triggered + let _ = &config.metrics_poll_count_histogram; + Self::new() + } + + pub(crate) fn set_queue_depth(&self, _len: usize) {} +} + +impl MetricsBatch { + pub(crate) fn new(_: &WorkerMetrics) -> Self { + Self {} + } + + pub(crate) fn submit(&mut self, _to: &WorkerMetrics) {} + pub(crate) fn about_to_park(&mut self) {} + pub(crate) fn inc_local_schedule_count(&mut self) {} + pub(crate) fn start_processing_scheduled_tasks(&mut self) {} + pub(crate) fn end_processing_scheduled_tasks(&mut self) {} + pub(crate) fn start_poll(&mut self) {} + pub(crate) fn end_poll(&mut self) {} +} + +cfg_rt_multi_thread! { + impl MetricsBatch { + pub(crate) fn incr_steal_count(&mut self, _by: u16) {} + pub(crate) fn incr_steal_operations(&mut self) {} + pub(crate) fn incr_overflow_count(&mut self) {} + } +} diff --git a/vendor/tokio/src/runtime/metrics/mod.rs b/vendor/tokio/src/runtime/metrics/mod.rs new file mode 100644 index 000000000..88be4a521 --- /dev/null +++ b/vendor/tokio/src/runtime/metrics/mod.rs @@ -0,0 +1,40 @@ +//! This module contains information need to view information about how the +//! runtime is performing. +//! +//! **Note**: This is an [unstable API][unstable]. The public API of types in +//! this module may break in 1.x releases. See [the documentation on unstable +//! features][unstable] for details. +//! +//! [unstable]: crate#unstable-features +#![allow(clippy::module_inception)] + +cfg_metrics! { + mod batch; + pub(crate) use batch::MetricsBatch; + + mod histogram; + pub(crate) use histogram::{Histogram, HistogramBatch, HistogramBuilder}; + #[allow(unreachable_pub)] // rust-lang/rust#57411 + pub use histogram::HistogramScale; + + mod runtime; + #[allow(unreachable_pub)] // rust-lang/rust#57411 + pub use runtime::RuntimeMetrics; + + mod scheduler; + pub(crate) use scheduler::SchedulerMetrics; + + mod worker; + pub(crate) use worker::WorkerMetrics; + + cfg_net! { + mod io; + pub(crate) use io::IoDriverMetrics; + } +} + +cfg_not_metrics! { + mod mock; + + pub(crate) use mock::{SchedulerMetrics, WorkerMetrics, MetricsBatch, HistogramBuilder}; +} diff --git a/vendor/tokio/src/runtime/metrics/runtime.rs b/vendor/tokio/src/runtime/metrics/runtime.rs new file mode 100644 index 000000000..1f990a1f8 --- /dev/null +++ b/vendor/tokio/src/runtime/metrics/runtime.rs @@ -0,0 +1,883 @@ +use crate::runtime::Handle; + +use std::ops::Range; +use std::sync::atomic::Ordering::Relaxed; +use std::time::Duration; + +/// Handle to the runtime's metrics. +/// +/// This handle is internally reference-counted and can be freely cloned. A +/// `RuntimeMetrics` handle is obtained using the [`Runtime::metrics`] method. +/// +/// [`Runtime::metrics`]: crate::runtime::Runtime::metrics() +#[derive(Clone, Debug)] +pub struct RuntimeMetrics { + handle: Handle, +} + +impl RuntimeMetrics { + pub(crate) fn new(handle: Handle) -> RuntimeMetrics { + RuntimeMetrics { handle } + } + + /// Returns the number of worker threads used by the runtime. + /// + /// The number of workers is set by configuring `worker_threads` on + /// `runtime::Builder`. When using the `current_thread` runtime, the return + /// value is always `1`. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Handle; + /// + /// #[tokio::main] + /// async fn main() { + /// let metrics = Handle::current().metrics(); + /// + /// let n = metrics.num_workers(); + /// println!("Runtime is using {} workers", n); + /// } + /// ``` + pub fn num_workers(&self) -> usize { + self.handle.inner.num_workers() + } + + /// Returns the number of additional threads spawned by the runtime. + /// + /// The number of workers is set by configuring `max_blocking_threads` on + /// `runtime::Builder`. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Handle; + /// + /// #[tokio::main] + /// async fn main() { + /// let _ = tokio::task::spawn_blocking(move || { + /// // Stand-in for compute-heavy work or using synchronous APIs + /// 1 + 1 + /// }).await; + /// let metrics = Handle::current().metrics(); + /// + /// let n = metrics.num_blocking_threads(); + /// println!("Runtime has created {} threads", n); + /// } + /// ``` + pub fn num_blocking_threads(&self) -> usize { + self.handle.inner.num_blocking_threads() + } + + /// Returns the number of active tasks in the runtime. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Handle; + /// + /// #[tokio::main] + /// async fn main() { + /// let metrics = Handle::current().metrics(); + /// + /// let n = metrics.active_tasks_count(); + /// println!("Runtime has {} active tasks", n); + /// } + /// ``` + pub fn active_tasks_count(&self) -> usize { + self.handle.inner.active_tasks_count() + } + + /// Returns the number of idle threads, which have spawned by the runtime + /// for `spawn_blocking` calls. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Handle; + /// + /// #[tokio::main] + /// async fn main() { + /// let _ = tokio::task::spawn_blocking(move || { + /// // Stand-in for compute-heavy work or using synchronous APIs + /// 1 + 1 + /// }).await; + /// let metrics = Handle::current().metrics(); + /// + /// let n = metrics.num_idle_blocking_threads(); + /// println!("Runtime has {} idle blocking thread pool threads", n); + /// } + /// ``` + pub fn num_idle_blocking_threads(&self) -> usize { + self.handle.inner.num_idle_blocking_threads() + } + + /// Returns the number of tasks scheduled from **outside** of the runtime. + /// + /// The remote schedule count starts at zero when the runtime is created and + /// increases by one each time a task is woken from **outside** of the + /// runtime. This usually means that a task is spawned or notified from a + /// non-runtime thread and must be queued using the Runtime's injection + /// queue, which tends to be slower. + /// + /// The counter is monotonically increasing. It is never decremented or + /// reset to zero. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Handle; + /// + /// #[tokio::main] + /// async fn main() { + /// let metrics = Handle::current().metrics(); + /// + /// let n = metrics.remote_schedule_count(); + /// println!("{} tasks were scheduled from outside the runtime", n); + /// } + /// ``` + pub fn remote_schedule_count(&self) -> u64 { + self.handle + .inner + .scheduler_metrics() + .remote_schedule_count + .load(Relaxed) + } + + /// Returns the number of times that tasks have been forced to yield back to the scheduler + /// after exhausting their task budgets. + /// + /// This count starts at zero when the runtime is created and increases by one each time a task yields due to exhausting its budget. + /// + /// The counter is monotonically increasing. It is never decremented or + /// reset to zero. + pub fn budget_forced_yield_count(&self) -> u64 { + self.handle + .inner + .scheduler_metrics() + .budget_forced_yield_count + .load(Relaxed) + } + + /// Returns the total number of times the given worker thread has parked. + /// + /// The worker park count starts at zero when the runtime is created and + /// increases by one each time the worker parks the thread waiting for new + /// inbound events to process. This usually means the worker has processed + /// all pending work and is currently idle. + /// + /// The counter is monotonically increasing. It is never decremented or + /// reset to zero. + /// + /// # Arguments + /// + /// `worker` is the index of the worker being queried. The given value must + /// be between 0 and `num_workers()`. The index uniquely identifies a single + /// worker and will continue to identify the worker throughout the lifetime + /// of the runtime instance. + /// + /// # Panics + /// + /// The method panics when `worker` represents an invalid worker, i.e. is + /// greater than or equal to `num_workers()`. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Handle; + /// + /// #[tokio::main] + /// async fn main() { + /// let metrics = Handle::current().metrics(); + /// + /// let n = metrics.worker_park_count(0); + /// println!("worker 0 parked {} times", n); + /// } + /// ``` + pub fn worker_park_count(&self, worker: usize) -> u64 { + self.handle + .inner + .worker_metrics(worker) + .park_count + .load(Relaxed) + } + + /// Returns the number of times the given worker thread unparked but + /// performed no work before parking again. + /// + /// The worker no-op count starts at zero when the runtime is created and + /// increases by one each time the worker unparks the thread but finds no + /// new work and goes back to sleep. This indicates a false-positive wake up. + /// + /// The counter is monotonically increasing. It is never decremented or + /// reset to zero. + /// + /// # Arguments + /// + /// `worker` is the index of the worker being queried. The given value must + /// be between 0 and `num_workers()`. The index uniquely identifies a single + /// worker and will continue to identify the worker throughout the lifetime + /// of the runtime instance. + /// + /// # Panics + /// + /// The method panics when `worker` represents an invalid worker, i.e. is + /// greater than or equal to `num_workers()`. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Handle; + /// + /// #[tokio::main] + /// async fn main() { + /// let metrics = Handle::current().metrics(); + /// + /// let n = metrics.worker_noop_count(0); + /// println!("worker 0 had {} no-op unparks", n); + /// } + /// ``` + pub fn worker_noop_count(&self, worker: usize) -> u64 { + self.handle + .inner + .worker_metrics(worker) + .noop_count + .load(Relaxed) + } + + /// Returns the number of tasks the given worker thread stole from + /// another worker thread. + /// + /// This metric only applies to the **multi-threaded** runtime and will + /// always return `0` when using the current thread runtime. + /// + /// The worker steal count starts at zero when the runtime is created and + /// increases by `N` each time the worker has processed its scheduled queue + /// and successfully steals `N` more pending tasks from another worker. + /// + /// The counter is monotonically increasing. It is never decremented or + /// reset to zero. + /// + /// # Arguments + /// + /// `worker` is the index of the worker being queried. The given value must + /// be between 0 and `num_workers()`. The index uniquely identifies a single + /// worker and will continue to identify the worker throughout the lifetime + /// of the runtime instance. + /// + /// # Panics + /// + /// The method panics when `worker` represents an invalid worker, i.e. is + /// greater than or equal to `num_workers()`. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Handle; + /// + /// #[tokio::main] + /// async fn main() { + /// let metrics = Handle::current().metrics(); + /// + /// let n = metrics.worker_steal_count(0); + /// println!("worker 0 has stolen {} tasks", n); + /// } + /// ``` + pub fn worker_steal_count(&self, worker: usize) -> u64 { + self.handle + .inner + .worker_metrics(worker) + .steal_count + .load(Relaxed) + } + + /// Returns the number of times the given worker thread stole tasks from + /// another worker thread. + /// + /// This metric only applies to the **multi-threaded** runtime and will + /// always return `0` when using the current thread runtime. + /// + /// The worker steal count starts at zero when the runtime is created and + /// increases by one each time the worker has processed its scheduled queue + /// and successfully steals more pending tasks from another worker. + /// + /// The counter is monotonically increasing. It is never decremented or + /// reset to zero. + /// + /// # Arguments + /// + /// `worker` is the index of the worker being queried. The given value must + /// be between 0 and `num_workers()`. The index uniquely identifies a single + /// worker and will continue to identify the worker throughout the lifetime + /// of the runtime instance. + /// + /// # Panics + /// + /// The method panics when `worker` represents an invalid worker, i.e. is + /// greater than or equal to `num_workers()`. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Handle; + /// + /// #[tokio::main] + /// async fn main() { + /// let metrics = Handle::current().metrics(); + /// + /// let n = metrics.worker_steal_operations(0); + /// println!("worker 0 has stolen tasks {} times", n); + /// } + /// ``` + pub fn worker_steal_operations(&self, worker: usize) -> u64 { + self.handle + .inner + .worker_metrics(worker) + .steal_operations + .load(Relaxed) + } + + /// Returns the number of tasks the given worker thread has polled. + /// + /// The worker poll count starts at zero when the runtime is created and + /// increases by one each time the worker polls a scheduled task. + /// + /// The counter is monotonically increasing. It is never decremented or + /// reset to zero. + /// + /// # Arguments + /// + /// `worker` is the index of the worker being queried. The given value must + /// be between 0 and `num_workers()`. The index uniquely identifies a single + /// worker and will continue to identify the worker throughout the lifetime + /// of the runtime instance. + /// + /// # Panics + /// + /// The method panics when `worker` represents an invalid worker, i.e. is + /// greater than or equal to `num_workers()`. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Handle; + /// + /// #[tokio::main] + /// async fn main() { + /// let metrics = Handle::current().metrics(); + /// + /// let n = metrics.worker_poll_count(0); + /// println!("worker 0 has polled {} tasks", n); + /// } + /// ``` + pub fn worker_poll_count(&self, worker: usize) -> u64 { + self.handle + .inner + .worker_metrics(worker) + .poll_count + .load(Relaxed) + } + + /// Returns the amount of time the given worker thread has been busy. + /// + /// The worker busy duration starts at zero when the runtime is created and + /// increases whenever the worker is spending time processing work. Using + /// this value can indicate the load of the given worker. If a lot of time + /// is spent busy, then the worker is under load and will check for inbound + /// events less often. + /// + /// The timer is monotonically increasing. It is never decremented or reset + /// to zero. + /// + /// # Arguments + /// + /// `worker` is the index of the worker being queried. The given value must + /// be between 0 and `num_workers()`. The index uniquely identifies a single + /// worker and will continue to identify the worker throughout the lifetime + /// of the runtime instance. + /// + /// # Panics + /// + /// The method panics when `worker` represents an invalid worker, i.e. is + /// greater than or equal to `num_workers()`. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Handle; + /// + /// #[tokio::main] + /// async fn main() { + /// let metrics = Handle::current().metrics(); + /// + /// let n = metrics.worker_total_busy_duration(0); + /// println!("worker 0 was busy for a total of {:?}", n); + /// } + /// ``` + pub fn worker_total_busy_duration(&self, worker: usize) -> Duration { + let nanos = self + .handle + .inner + .worker_metrics(worker) + .busy_duration_total + .load(Relaxed); + Duration::from_nanos(nanos) + } + + /// Returns the number of tasks scheduled from **within** the runtime on the + /// given worker's local queue. + /// + /// The local schedule count starts at zero when the runtime is created and + /// increases by one each time a task is woken from **inside** of the + /// runtime on the given worker. This usually means that a task is spawned + /// or notified from within a runtime thread and will be queued on the + /// worker-local queue. + /// + /// The counter is monotonically increasing. It is never decremented or + /// reset to zero. + /// + /// # Arguments + /// + /// `worker` is the index of the worker being queried. The given value must + /// be between 0 and `num_workers()`. The index uniquely identifies a single + /// worker and will continue to identify the worker throughout the lifetime + /// of the runtime instance. + /// + /// # Panics + /// + /// The method panics when `worker` represents an invalid worker, i.e. is + /// greater than or equal to `num_workers()`. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Handle; + /// + /// #[tokio::main] + /// async fn main() { + /// let metrics = Handle::current().metrics(); + /// + /// let n = metrics.worker_local_schedule_count(0); + /// println!("{} tasks were scheduled on the worker's local queue", n); + /// } + /// ``` + pub fn worker_local_schedule_count(&self, worker: usize) -> u64 { + self.handle + .inner + .worker_metrics(worker) + .local_schedule_count + .load(Relaxed) + } + + /// Returns the number of times the given worker thread saturated its local + /// queue. + /// + /// This metric only applies to the **multi-threaded** scheduler. + /// + /// The worker steal count starts at zero when the runtime is created and + /// increases by one each time the worker attempts to schedule a task + /// locally, but its local queue is full. When this happens, half of the + /// local queue is moved to the injection queue. + /// + /// The counter is monotonically increasing. It is never decremented or + /// reset to zero. + /// + /// # Arguments + /// + /// `worker` is the index of the worker being queried. The given value must + /// be between 0 and `num_workers()`. The index uniquely identifies a single + /// worker and will continue to identify the worker throughout the lifetime + /// of the runtime instance. + /// + /// # Panics + /// + /// The method panics when `worker` represents an invalid worker, i.e. is + /// greater than or equal to `num_workers()`. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Handle; + /// + /// #[tokio::main] + /// async fn main() { + /// let metrics = Handle::current().metrics(); + /// + /// let n = metrics.worker_overflow_count(0); + /// println!("worker 0 has overflowed its queue {} times", n); + /// } + /// ``` + pub fn worker_overflow_count(&self, worker: usize) -> u64 { + self.handle + .inner + .worker_metrics(worker) + .overflow_count + .load(Relaxed) + } + + /// Returns the number of tasks currently scheduled in the runtime's + /// injection queue. + /// + /// Tasks that are spawned or notified from a non-runtime thread are + /// scheduled using the runtime's injection queue. This metric returns the + /// **current** number of tasks pending in the injection queue. As such, the + /// returned value may increase or decrease as new tasks are scheduled and + /// processed. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Handle; + /// + /// #[tokio::main] + /// async fn main() { + /// let metrics = Handle::current().metrics(); + /// + /// let n = metrics.injection_queue_depth(); + /// println!("{} tasks currently pending in the runtime's injection queue", n); + /// } + /// ``` + pub fn injection_queue_depth(&self) -> usize { + self.handle.inner.injection_queue_depth() + } + + /// Returns the number of tasks currently scheduled in the given worker's + /// local queue. + /// + /// Tasks that are spawned or notified from within a runtime thread are + /// scheduled using that worker's local queue. This metric returns the + /// **current** number of tasks pending in the worker's local queue. As + /// such, the returned value may increase or decrease as new tasks are + /// scheduled and processed. + /// + /// # Arguments + /// + /// `worker` is the index of the worker being queried. The given value must + /// be between 0 and `num_workers()`. The index uniquely identifies a single + /// worker and will continue to identify the worker throughout the lifetime + /// of the runtime instance. + /// + /// # Panics + /// + /// The method panics when `worker` represents an invalid worker, i.e. is + /// greater than or equal to `num_workers()`. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Handle; + /// + /// #[tokio::main] + /// async fn main() { + /// let metrics = Handle::current().metrics(); + /// + /// let n = metrics.worker_local_queue_depth(0); + /// println!("{} tasks currently pending in worker 0's local queue", n); + /// } + /// ``` + pub fn worker_local_queue_depth(&self, worker: usize) -> usize { + self.handle.inner.worker_local_queue_depth(worker) + } + + /// Returns `true` if the runtime is tracking the distribution of task poll + /// times. + /// + /// Task poll times are not instrumented by default as doing so requires + /// calling [`Instant::now()`] twice per task poll. The feature is enabled + /// by calling [`enable_metrics_poll_count_histogram()`] when building the + /// runtime. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::{self, Handle}; + /// + /// fn main() { + /// runtime::Builder::new_current_thread() + /// .enable_metrics_poll_count_histogram() + /// .build() + /// .unwrap() + /// .block_on(async { + /// let metrics = Handle::current().metrics(); + /// let enabled = metrics.poll_count_histogram_enabled(); + /// + /// println!("Tracking task poll time distribution: {:?}", enabled); + /// }); + /// } + /// ``` + /// + /// [`enable_metrics_poll_count_histogram()`]: crate::runtime::Builder::enable_metrics_poll_count_histogram + /// [`Instant::now()`]: std::time::Instant::now + pub fn poll_count_histogram_enabled(&self) -> bool { + self.handle + .inner + .worker_metrics(0) + .poll_count_histogram + .is_some() + } + + /// Returns the number of histogram buckets tracking the distribution of + /// task poll times. + /// + /// This value is configured by calling + /// [`metrics_poll_count_histogram_buckets()`] when building the runtime. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::{self, Handle}; + /// + /// fn main() { + /// runtime::Builder::new_current_thread() + /// .enable_metrics_poll_count_histogram() + /// .build() + /// .unwrap() + /// .block_on(async { + /// let metrics = Handle::current().metrics(); + /// let buckets = metrics.poll_count_histogram_num_buckets(); + /// + /// println!("Histogram buckets: {:?}", buckets); + /// }); + /// } + /// ``` + /// + /// [`metrics_poll_count_histogram_buckets()`]: + /// crate::runtime::Builder::metrics_poll_count_histogram_buckets + pub fn poll_count_histogram_num_buckets(&self) -> usize { + self.handle + .inner + .worker_metrics(0) + .poll_count_histogram + .as_ref() + .map(|histogram| histogram.num_buckets()) + .unwrap_or_default() + } + + /// Returns the range of task poll times tracked by the given bucket. + /// + /// This value is configured by calling + /// [`metrics_poll_count_histogram_resolution()`] when building the runtime. + /// + /// # Panics + /// + /// The method panics if `bucket` represents an invalid bucket index, i.e. + /// is greater than or equal to `poll_count_histogram_num_buckets()`. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::{self, Handle}; + /// + /// fn main() { + /// runtime::Builder::new_current_thread() + /// .enable_metrics_poll_count_histogram() + /// .build() + /// .unwrap() + /// .block_on(async { + /// let metrics = Handle::current().metrics(); + /// let buckets = metrics.poll_count_histogram_num_buckets(); + /// + /// for i in 0..buckets { + /// let range = metrics.poll_count_histogram_bucket_range(i); + /// println!("Histogram bucket {} range: {:?}", i, range); + /// } + /// }); + /// } + /// ``` + /// + /// [`metrics_poll_count_histogram_resolution()`]: + /// crate::runtime::Builder::metrics_poll_count_histogram_resolution + #[track_caller] + pub fn poll_count_histogram_bucket_range(&self, bucket: usize) -> Range<Duration> { + self.handle + .inner + .worker_metrics(0) + .poll_count_histogram + .as_ref() + .map(|histogram| { + let range = histogram.bucket_range(bucket); + std::ops::Range { + start: Duration::from_nanos(range.start), + end: Duration::from_nanos(range.end), + } + }) + .unwrap_or_default() + } + + /// Returns the number of times the given worker polled tasks with a poll + /// duration within the given bucket's range. + /// + /// Each worker maintains its own histogram and the counts for each bucket + /// starts at zero when the runtime is created. Each time the worker polls a + /// task, it tracks the duration the task poll time took and increments the + /// associated bucket by 1. + /// + /// Each bucket is a monotonically increasing counter. It is never + /// decremented or reset to zero. + /// + /// # Arguments + /// + /// `worker` is the index of the worker being queried. The given value must + /// be between 0 and `num_workers()`. The index uniquely identifies a single + /// worker and will continue to identify the worker throughout the lifetime + /// of the runtime instance. + /// + /// `bucket` is the index of the bucket being queried. The bucket is scoped + /// to the worker. The range represented by the bucket can be queried by + /// calling [`poll_count_histogram_bucket_range()`]. Each worker maintains + /// identical bucket ranges. + /// + /// # Panics + /// + /// The method panics when `worker` represents an invalid worker, i.e. is + /// greater than or equal to `num_workers()` or if `bucket` represents an + /// invalid bucket. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::{self, Handle}; + /// + /// fn main() { + /// runtime::Builder::new_current_thread() + /// .enable_metrics_poll_count_histogram() + /// .build() + /// .unwrap() + /// .block_on(async { + /// let metrics = Handle::current().metrics(); + /// let buckets = metrics.poll_count_histogram_num_buckets(); + /// + /// for worker in 0..metrics.num_workers() { + /// for i in 0..buckets { + /// let count = metrics.poll_count_histogram_bucket_count(worker, i); + /// println!("Poll count {}", count); + /// } + /// } + /// }); + /// } + /// ``` + /// + /// [`poll_count_histogram_bucket_range()`]: crate::runtime::RuntimeMetrics::poll_count_histogram_bucket_range + #[track_caller] + pub fn poll_count_histogram_bucket_count(&self, worker: usize, bucket: usize) -> u64 { + self.handle + .inner + .worker_metrics(worker) + .poll_count_histogram + .as_ref() + .map(|histogram| histogram.get(bucket)) + .unwrap_or_default() + } + + /// Returns the number of tasks currently scheduled in the blocking + /// thread pool, spawned using `spawn_blocking`. + /// + /// This metric returns the **current** number of tasks pending in + /// blocking thread pool. As such, the returned value may increase + /// or decrease as new tasks are scheduled and processed. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Handle; + /// + /// #[tokio::main] + /// async fn main() { + /// let metrics = Handle::current().metrics(); + /// + /// let n = metrics.blocking_queue_depth(); + /// println!("{} tasks currently pending in the blocking thread pool", n); + /// } + /// ``` + pub fn blocking_queue_depth(&self) -> usize { + self.handle.inner.blocking_queue_depth() + } +} + +cfg_net! { + impl RuntimeMetrics { + /// Returns the number of file descriptors that have been registered with the + /// runtime's I/O driver. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Handle; + /// + /// #[tokio::main] + /// async fn main() { + /// let metrics = Handle::current().metrics(); + /// + /// let registered_fds = metrics.io_driver_fd_registered_count(); + /// println!("{} fds have been registered with the runtime's I/O driver.", registered_fds); + /// + /// let deregistered_fds = metrics.io_driver_fd_deregistered_count(); + /// + /// let current_fd_count = registered_fds - deregistered_fds; + /// println!("{} fds are currently registered by the runtime's I/O driver.", current_fd_count); + /// } + /// ``` + pub fn io_driver_fd_registered_count(&self) -> u64 { + self.with_io_driver_metrics(|m| { + m.fd_registered_count.load(Relaxed) + }) + } + + /// Returns the number of file descriptors that have been deregistered by the + /// runtime's I/O driver. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Handle; + /// + /// #[tokio::main] + /// async fn main() { + /// let metrics = Handle::current().metrics(); + /// + /// let n = metrics.io_driver_fd_deregistered_count(); + /// println!("{} fds have been deregistered by the runtime's I/O driver.", n); + /// } + /// ``` + pub fn io_driver_fd_deregistered_count(&self) -> u64 { + self.with_io_driver_metrics(|m| { + m.fd_deregistered_count.load(Relaxed) + }) + } + + /// Returns the number of ready events processed by the runtime's + /// I/O driver. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Handle; + /// + /// #[tokio::main] + /// async fn main() { + /// let metrics = Handle::current().metrics(); + /// + /// let n = metrics.io_driver_ready_count(); + /// println!("{} ready events processed by the runtime's I/O driver.", n); + /// } + /// ``` + pub fn io_driver_ready_count(&self) -> u64 { + self.with_io_driver_metrics(|m| m.ready_count.load(Relaxed)) + } + + fn with_io_driver_metrics<F>(&self, f: F) -> u64 + where + F: Fn(&super::IoDriverMetrics) -> u64, + { + // TODO: Investigate if this should return 0, most of our metrics always increase + // thus this breaks that guarantee. + self.handle + .inner + .driver() + .io + .as_ref() + .map(|h| f(&h.metrics)) + .unwrap_or(0) + } + } +} diff --git a/vendor/tokio/src/runtime/metrics/scheduler.rs b/vendor/tokio/src/runtime/metrics/scheduler.rs new file mode 100644 index 000000000..d9f8edfaa --- /dev/null +++ b/vendor/tokio/src/runtime/metrics/scheduler.rs @@ -0,0 +1,34 @@ +use crate::loom::sync::atomic::{AtomicU64, Ordering::Relaxed}; + +/// Retrieves metrics from the Tokio runtime. +/// +/// **Note**: This is an [unstable API][unstable]. The public API of this type +/// may break in 1.x releases. See [the documentation on unstable +/// features][unstable] for details. +/// +/// [unstable]: crate#unstable-features +#[derive(Debug)] +pub(crate) struct SchedulerMetrics { + /// Number of tasks that are scheduled from outside the runtime. + pub(super) remote_schedule_count: AtomicU64, + pub(super) budget_forced_yield_count: AtomicU64, +} + +impl SchedulerMetrics { + pub(crate) fn new() -> SchedulerMetrics { + SchedulerMetrics { + remote_schedule_count: AtomicU64::new(0), + budget_forced_yield_count: AtomicU64::new(0), + } + } + + /// Increment the number of tasks scheduled externally + pub(crate) fn inc_remote_schedule_count(&self) { + self.remote_schedule_count.fetch_add(1, Relaxed); + } + + /// Increment the number of tasks forced to yield due to budget exhaustion + pub(crate) fn inc_budget_forced_yield_count(&self) { + self.budget_forced_yield_count.fetch_add(1, Relaxed); + } +} diff --git a/vendor/tokio/src/runtime/metrics/worker.rs b/vendor/tokio/src/runtime/metrics/worker.rs new file mode 100644 index 000000000..e0f23e6a0 --- /dev/null +++ b/vendor/tokio/src/runtime/metrics/worker.rs @@ -0,0 +1,80 @@ +use crate::loom::sync::atomic::Ordering::Relaxed; +use crate::loom::sync::atomic::{AtomicU64, AtomicUsize}; +use crate::runtime::metrics::Histogram; +use crate::runtime::Config; + +/// Retrieve runtime worker metrics. +/// +/// **Note**: This is an [unstable API][unstable]. The public API of this type +/// may break in 1.x releases. See [the documentation on unstable +/// features][unstable] for details. +/// +/// [unstable]: crate#unstable-features +#[derive(Debug)] +#[repr(align(128))] +pub(crate) struct WorkerMetrics { + /// Number of times the worker parked. + pub(crate) park_count: AtomicU64, + + /// Number of times the worker woke then parked again without doing work. + pub(crate) noop_count: AtomicU64, + + /// Number of tasks the worker stole. + pub(crate) steal_count: AtomicU64, + + /// Number of times the worker stole + pub(crate) steal_operations: AtomicU64, + + /// Number of tasks the worker polled. + pub(crate) poll_count: AtomicU64, + + /// Amount of time the worker spent doing work vs. parking. + pub(crate) busy_duration_total: AtomicU64, + + /// Number of tasks scheduled for execution on the worker's local queue. + pub(crate) local_schedule_count: AtomicU64, + + /// Number of tasks moved from the local queue to the global queue to free space. + pub(crate) overflow_count: AtomicU64, + + /// Number of tasks currently in the local queue. Used only by the + /// current-thread scheduler. + pub(crate) queue_depth: AtomicUsize, + + /// If `Some`, tracks the the number of polls by duration range. + pub(super) poll_count_histogram: Option<Histogram>, +} + +impl WorkerMetrics { + pub(crate) fn from_config(config: &Config) -> WorkerMetrics { + let mut worker_metrics = WorkerMetrics::new(); + worker_metrics.poll_count_histogram = config + .metrics_poll_count_histogram + .as_ref() + .map(|histogram_builder| histogram_builder.build()); + worker_metrics + } + + pub(crate) fn new() -> WorkerMetrics { + WorkerMetrics { + park_count: AtomicU64::new(0), + noop_count: AtomicU64::new(0), + steal_count: AtomicU64::new(0), + steal_operations: AtomicU64::new(0), + poll_count: AtomicU64::new(0), + overflow_count: AtomicU64::new(0), + busy_duration_total: AtomicU64::new(0), + local_schedule_count: AtomicU64::new(0), + queue_depth: AtomicUsize::new(0), + poll_count_histogram: None, + } + } + + pub(crate) fn queue_depth(&self) -> usize { + self.queue_depth.load(Relaxed) + } + + pub(crate) fn set_queue_depth(&self, len: usize) { + self.queue_depth.store(len, Relaxed); + } +} diff --git a/vendor/tokio/src/runtime/mod.rs b/vendor/tokio/src/runtime/mod.rs index 52532ec6f..cb198f51f 100644 --- a/vendor/tokio/src/runtime/mod.rs +++ b/vendor/tokio/src/runtime/mod.rs @@ -137,7 +137,7 @@ //! use tokio::runtime; //! //! # fn main() -> Result<(), Box<dyn std::error::Error>> { -//! let basic_rt = runtime::Builder::new_current_thread() +//! let rt = runtime::Builder::new_current_thread() //! .build()?; //! # Ok(()) } //! ``` @@ -156,9 +156,11 @@ //! multi-thread scheduler spawns threads to schedule tasks and for `spawn_blocking` //! calls. //! -//! While the `Runtime` is active, threads may shutdown after periods of being -//! idle. Once `Runtime` is dropped, all runtime threads are forcibly shutdown. -//! Any tasks that have not yet completed will be dropped. +//! While the `Runtime` is active, threads may shut down after periods of being +//! idle. Once `Runtime` is dropped, all runtime threads have usually been +//! terminated, but in the presence of unstoppable spawned work are not +//! guaranteed to have been terminated. See the +//! [struct level documentation](Runtime#shutdown) for more details. //! //! [tasks]: crate::task //! [`Runtime`]: Runtime @@ -166,7 +168,6 @@ //! [`tokio::main`]: ../attr.main.html //! [runtime builder]: crate::runtime::Builder //! [`Runtime::new`]: crate::runtime::Runtime::new -//! [`Builder::basic_scheduler`]: crate::runtime::Builder::basic_scheduler //! [`Builder::threaded_scheduler`]: crate::runtime::Builder::threaded_scheduler //! [`Builder::enable_io`]: crate::runtime::Builder::enable_io //! [`Builder::enable_time`]: crate::runtime::Builder::enable_time @@ -174,390 +175,91 @@ // At the top due to macros #[cfg(test)] +#[cfg(not(tokio_wasm))] #[macro_use] mod tests; -pub(crate) mod enter; +pub(crate) mod context; -pub(crate) mod task; +pub(crate) mod coop; -cfg_rt! { - mod basic_scheduler; - use basic_scheduler::BasicScheduler; +pub(crate) mod park; - mod blocking; - use blocking::BlockingPool; - pub(crate) use blocking::spawn_blocking; +mod driver; - mod builder; - pub use self::builder::Builder; - - pub(crate) mod context; - pub(crate) mod driver; - - use self::enter::enter; - - mod handle; - pub use handle::{EnterGuard, Handle}; +pub(crate) mod scheduler; - mod spawner; - use self::spawner::Spawner; +cfg_io_driver_impl! { + pub(crate) mod io; } -cfg_rt_multi_thread! { - mod park; - use park::Parker; +cfg_process_driver! { + mod process; } -cfg_rt_multi_thread! { - mod queue; +cfg_time! { + pub(crate) mod time; +} - pub(crate) mod thread_pool; - use self::thread_pool::ThreadPool; +cfg_signal_internal_and_unix! { + pub(crate) mod signal; } cfg_rt! { - use crate::task::JoinHandle; - - use std::future::Future; - use std::time::Duration; + pub(crate) mod task; - /// The Tokio runtime. - /// - /// The runtime provides an I/O driver, task scheduler, [timer], and - /// blocking pool, necessary for running asynchronous tasks. - /// - /// Instances of `Runtime` can be created using [`new`], or [`Builder`]. - /// However, most users will use the `#[tokio::main]` annotation on their - /// entry point instead. - /// - /// See [module level][mod] documentation for more details. - /// - /// # Shutdown - /// - /// Shutting down the runtime is done by dropping the value. The current - /// thread will block until the shut down operation has completed. - /// - /// * Drain any scheduled work queues. - /// * Drop any futures that have not yet completed. - /// * Drop the reactor. - /// - /// Once the reactor has dropped, any outstanding I/O resources bound to - /// that reactor will no longer function. Calling any method on them will - /// result in an error. - /// - /// # Sharing - /// - /// The Tokio runtime implements `Sync` and `Send` to allow you to wrap it - /// in a `Arc`. Most fn take `&self` to allow you to call them concurrently - /// across multiple threads. - /// - /// Calls to `shutdown` and `shutdown_timeout` require exclusive ownership of - /// the runtime type and this can be achieved via `Arc::try_unwrap` when only - /// one strong count reference is left over. - /// - /// [timer]: crate::time - /// [mod]: index.html - /// [`new`]: method@Self::new - /// [`Builder`]: struct@Builder - #[derive(Debug)] - pub struct Runtime { - /// Task executor - kind: Kind, + mod config; + use config::Config; - /// Handle to runtime, also contains driver handles - handle: Handle, + mod blocking; + #[cfg_attr(tokio_wasi, allow(unused_imports))] + pub(crate) use blocking::spawn_blocking; - /// Blocking pool handle, used to signal shutdown - blocking_pool: BlockingPool, + cfg_trace! { + pub(crate) use blocking::Mandatory; } - /// The runtime executor is either a thread-pool or a current-thread executor. - #[derive(Debug)] - enum Kind { - /// Execute all tasks on the current-thread. - CurrentThread(BasicScheduler<driver::Driver>), - - /// Execute tasks across multiple threads. - #[cfg(feature = "rt-multi-thread")] - ThreadPool(ThreadPool), + cfg_fs! { + pub(crate) use blocking::spawn_mandatory_blocking; } - /// After thread starts / before thread stops - type Callback = std::sync::Arc<dyn Fn() + Send + Sync>; - - impl Runtime { - /// Create a new runtime instance with default configuration values. - /// - /// This results in the multi threaded scheduler, I/O driver, and time driver being - /// initialized. - /// - /// Most applications will not need to call this function directly. Instead, - /// they will use the [`#[tokio::main]` attribute][main]. When a more complex - /// configuration is necessary, the [runtime builder] may be used. - /// - /// See [module level][mod] documentation for more details. - /// - /// # Examples - /// - /// Creating a new `Runtime` with default configuration values. - /// - /// ``` - /// use tokio::runtime::Runtime; - /// - /// let rt = Runtime::new() - /// .unwrap(); - /// - /// // Use the runtime... - /// ``` - /// - /// [mod]: index.html - /// [main]: ../attr.main.html - /// [threaded scheduler]: index.html#threaded-scheduler - /// [basic scheduler]: index.html#basic-scheduler - /// [runtime builder]: crate::runtime::Builder - #[cfg(feature = "rt-multi-thread")] - #[cfg_attr(docsrs, doc(cfg(feature = "rt-multi-thread")))] - pub fn new() -> std::io::Result<Runtime> { - Builder::new_multi_thread().enable_all().build() - } + mod builder; + pub use self::builder::Builder; + cfg_unstable! { + pub use self::builder::UnhandledPanic; + pub use crate::util::rand::RngSeed; + } - /// Return a handle to the runtime's spawner. - /// - /// The returned handle can be used to spawn tasks that run on this runtime, and can - /// be cloned to allow moving the `Handle` to other threads. - /// - /// # Examples - /// - /// ``` - /// use tokio::runtime::Runtime; - /// - /// let rt = Runtime::new() - /// .unwrap(); - /// - /// let handle = rt.handle(); - /// - /// // Use the handle... - /// ``` - pub fn handle(&self) -> &Handle { - &self.handle - } + cfg_taskdump! { + pub mod dump; + pub use dump::Dump; + } - /// Spawn a future onto the Tokio runtime. - /// - /// This spawns the given future onto the runtime's executor, usually a - /// thread pool. The thread pool is then responsible for polling the future - /// until it completes. - /// - /// See [module level][mod] documentation for more details. - /// - /// [mod]: index.html - /// - /// # Examples - /// - /// ``` - /// use tokio::runtime::Runtime; - /// - /// # fn dox() { - /// // Create the runtime - /// let rt = Runtime::new().unwrap(); - /// - /// // Spawn a future onto the runtime - /// rt.spawn(async { - /// println!("now running on a worker thread"); - /// }); - /// # } - /// ``` - #[cfg_attr(tokio_track_caller, track_caller)] - pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output> - where - F: Future + Send + 'static, - F::Output: Send + 'static, - { - self.handle.spawn(future) - } + mod handle; + pub use handle::{EnterGuard, Handle, TryCurrentError}; - /// Run the provided function on an executor dedicated to blocking operations. - /// - /// # Examples - /// - /// ``` - /// use tokio::runtime::Runtime; - /// - /// # fn dox() { - /// // Create the runtime - /// let rt = Runtime::new().unwrap(); - /// - /// // Spawn a blocking function onto the runtime - /// rt.spawn_blocking(|| { - /// println!("now running on a worker thread"); - /// }); - /// # } - #[cfg_attr(tokio_track_caller, track_caller)] - pub fn spawn_blocking<F, R>(&self, func: F) -> JoinHandle<R> - where - F: FnOnce() -> R + Send + 'static, - R: Send + 'static, - { - self.handle.spawn_blocking(func) - } + mod runtime; + pub use runtime::{Runtime, RuntimeFlavor}; - /// Run a future to completion on the Tokio runtime. This is the - /// runtime's entry point. - /// - /// This runs the given future on the current thread, blocking until it is - /// complete, and yielding its resolved result. Any tasks or timers - /// which the future spawns internally will be executed on the runtime. - /// - /// # Multi thread scheduler - /// - /// When the multi thread scheduler is used this will allow futures - /// to run within the io driver and timer context of the overall runtime. - /// - /// # Current thread scheduler - /// - /// When the current thread scheduler is enabled `block_on` - /// can be called concurrently from multiple threads. The first call - /// will take ownership of the io and timer drivers. This means - /// other threads which do not own the drivers will hook into that one. - /// When the first `block_on` completes, other threads will be able to - /// "steal" the driver to allow continued execution of their futures. - /// - /// # Panics - /// - /// This function panics if the provided future panics, or if called within an - /// asynchronous execution context. - /// - /// # Examples - /// - /// ```no_run - /// use tokio::runtime::Runtime; - /// - /// // Create the runtime - /// let rt = Runtime::new().unwrap(); - /// - /// // Execute the future, blocking the current thread until completion - /// rt.block_on(async { - /// println!("hello"); - /// }); - /// ``` - /// - /// [handle]: fn@Handle::block_on - pub fn block_on<F: Future>(&self, future: F) -> F::Output { - let _enter = self.enter(); + mod thread_id; + pub(crate) use thread_id::ThreadId; - match &self.kind { - Kind::CurrentThread(exec) => exec.block_on(future), - #[cfg(feature = "rt-multi-thread")] - Kind::ThreadPool(exec) => exec.block_on(future), - } - } + cfg_metrics! { + mod metrics; + pub use metrics::{RuntimeMetrics, HistogramScale}; - /// Enter the runtime context. - /// - /// This allows you to construct types that must have an executor - /// available on creation such as [`Sleep`] or [`TcpStream`]. It will - /// also allow you to call methods such as [`tokio::spawn`]. - /// - /// [`Sleep`]: struct@crate::time::Sleep - /// [`TcpStream`]: struct@crate::net::TcpStream - /// [`tokio::spawn`]: fn@crate::spawn - /// - /// # Example - /// - /// ``` - /// use tokio::runtime::Runtime; - /// - /// fn function_that_spawns(msg: String) { - /// // Had we not used `rt.enter` below, this would panic. - /// tokio::spawn(async move { - /// println!("{}", msg); - /// }); - /// } - /// - /// fn main() { - /// let rt = Runtime::new().unwrap(); - /// - /// let s = "Hello World!".to_string(); - /// - /// // By entering the context, we tie `tokio::spawn` to this executor. - /// let _guard = rt.enter(); - /// function_that_spawns(s); - /// } - /// ``` - pub fn enter(&self) -> EnterGuard<'_> { - self.handle.enter() - } + pub(crate) use metrics::{MetricsBatch, SchedulerMetrics, WorkerMetrics, HistogramBuilder}; - /// Shutdown the runtime, waiting for at most `duration` for all spawned - /// task to shutdown. - /// - /// Usually, dropping a `Runtime` handle is sufficient as tasks are able to - /// shutdown in a timely fashion. However, dropping a `Runtime` will wait - /// indefinitely for all tasks to terminate, and there are cases where a long - /// blocking task has been spawned, which can block dropping `Runtime`. - /// - /// In this case, calling `shutdown_timeout` with an explicit wait timeout - /// can work. The `shutdown_timeout` will signal all tasks to shutdown and - /// will wait for at most `duration` for all spawned tasks to terminate. If - /// `timeout` elapses before all tasks are dropped, the function returns and - /// outstanding tasks are potentially leaked. - /// - /// # Examples - /// - /// ``` - /// use tokio::runtime::Runtime; - /// use tokio::task; - /// - /// use std::thread; - /// use std::time::Duration; - /// - /// fn main() { - /// let runtime = Runtime::new().unwrap(); - /// - /// runtime.block_on(async move { - /// task::spawn_blocking(move || { - /// thread::sleep(Duration::from_secs(10_000)); - /// }); - /// }); - /// - /// runtime.shutdown_timeout(Duration::from_millis(100)); - /// } - /// ``` - pub fn shutdown_timeout(mut self, duration: Duration) { - // Wakeup and shutdown all the worker threads - self.handle.shutdown(); - self.blocking_pool.shutdown(Some(duration)); + cfg_net! { + pub(crate) use metrics::IoDriverMetrics; } + } - /// Shutdown the runtime, without waiting for any spawned tasks to shutdown. - /// - /// This can be useful if you want to drop a runtime from within another runtime. - /// Normally, dropping a runtime will block indefinitely for spawned blocking tasks - /// to complete, which would normally not be permitted within an asynchronous context. - /// By calling `shutdown_background()`, you can drop the runtime from such a context. - /// - /// Note however, that because we do not wait for any blocking tasks to complete, this - /// may result in a resource leak (in that any blocking tasks are still running until they - /// return. - /// - /// This function is equivalent to calling `shutdown_timeout(Duration::of_nanos(0))`. - /// - /// ``` - /// use tokio::runtime::Runtime; - /// - /// fn main() { - /// let runtime = Runtime::new().unwrap(); - /// - /// runtime.block_on(async move { - /// let inner_runtime = Runtime::new().unwrap(); - /// // ... - /// inner_runtime.shutdown_background(); - /// }); - /// } - /// ``` - pub fn shutdown_background(self) { - self.shutdown_timeout(Duration::from_nanos(0)) - } + cfg_not_metrics! { + pub(crate) mod metrics; + pub(crate) use metrics::{SchedulerMetrics, WorkerMetrics, MetricsBatch, HistogramBuilder}; } + + /// After thread starts / before thread stops + type Callback = std::sync::Arc<dyn Fn() + Send + Sync>; } diff --git a/vendor/tokio/src/runtime/park.rs b/vendor/tokio/src/runtime/park.rs index 033b9f20b..2392846ab 100644 --- a/vendor/tokio/src/runtime/park.rs +++ b/vendor/tokio/src/runtime/park.rs @@ -1,153 +1,102 @@ -//! Parks the runtime. -//! -//! A combination of the various resource driver park handles. +#![cfg_attr(not(feature = "full"), allow(dead_code))] use crate::loom::sync::atomic::AtomicUsize; use crate::loom::sync::{Arc, Condvar, Mutex}; -use crate::loom::thread; -use crate::park::{Park, Unpark}; -use crate::runtime::driver::Driver; -use crate::util::TryLock; use std::sync::atomic::Ordering::SeqCst; use std::time::Duration; -pub(crate) struct Parker { +#[derive(Debug)] +pub(crate) struct ParkThread { inner: Arc<Inner>, } -pub(crate) struct Unparker { +/// Unblocks a thread that was blocked by `ParkThread`. +#[derive(Clone, Debug)] +pub(crate) struct UnparkThread { inner: Arc<Inner>, } +#[derive(Debug)] struct Inner { - /// Avoids entering the park if possible state: AtomicUsize, - - /// Used to coordinate access to the driver / condvar mutex: Mutex<()>, - - /// Condvar to block on if the driver is unavailable. condvar: Condvar, - - /// Resource (I/O, time, ...) driver - shared: Arc<Shared>, } const EMPTY: usize = 0; -const PARKED_CONDVAR: usize = 1; -const PARKED_DRIVER: usize = 2; -const NOTIFIED: usize = 3; +const PARKED: usize = 1; +const NOTIFIED: usize = 2; -/// Shared across multiple Parker handles -struct Shared { - /// Shared driver. Only one thread at a time can use this - driver: TryLock<Driver>, +tokio_thread_local! { + static CURRENT_PARKER: ParkThread = ParkThread::new(); +} - /// Unpark handle - handle: <Driver as Park>::Unpark, +// Bit of a hack, but it is only for loom +#[cfg(loom)] +tokio_thread_local! { + static CURRENT_THREAD_PARK_COUNT: AtomicUsize = AtomicUsize::new(0); } -impl Parker { - pub(crate) fn new(driver: Driver) -> Parker { - let handle = driver.unpark(); +// ==== impl ParkThread ==== - Parker { +impl ParkThread { + pub(crate) fn new() -> Self { + Self { inner: Arc::new(Inner { state: AtomicUsize::new(EMPTY), mutex: Mutex::new(()), condvar: Condvar::new(), - shared: Arc::new(Shared { - driver: TryLock::new(driver), - handle, - }), }), } } -} -impl Clone for Parker { - fn clone(&self) -> Parker { - Parker { - inner: Arc::new(Inner { - state: AtomicUsize::new(EMPTY), - mutex: Mutex::new(()), - condvar: Condvar::new(), - shared: self.inner.shared.clone(), - }), - } + pub(crate) fn unpark(&self) -> UnparkThread { + let inner = self.inner.clone(); + UnparkThread { inner } } -} - -impl Park for Parker { - type Unpark = Unparker; - type Error = (); - fn unpark(&self) -> Unparker { - Unparker { - inner: self.inner.clone(), - } - } - - fn park(&mut self) -> Result<(), Self::Error> { + pub(crate) fn park(&mut self) { + #[cfg(loom)] + CURRENT_THREAD_PARK_COUNT.with(|count| count.fetch_add(1, SeqCst)); self.inner.park(); - Ok(()) } - fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> { - // Only parking with zero is supported... - assert_eq!(duration, Duration::from_millis(0)); + pub(crate) fn park_timeout(&mut self, duration: Duration) { + #[cfg(loom)] + CURRENT_THREAD_PARK_COUNT.with(|count| count.fetch_add(1, SeqCst)); - if let Some(mut driver) = self.inner.shared.driver.try_lock() { - driver.park_timeout(duration).map_err(|_| ()) - } else { - Ok(()) - } + // Wasm doesn't have threads, so just sleep. + #[cfg(not(tokio_wasm))] + self.inner.park_timeout(duration); + #[cfg(tokio_wasm)] + std::thread::sleep(duration); } - fn shutdown(&mut self) { + pub(crate) fn shutdown(&mut self) { self.inner.shutdown(); } } -impl Unpark for Unparker { - fn unpark(&self) { - self.inner.unpark(); - } -} +// ==== impl Inner ==== impl Inner { /// Parks the current thread for at most `dur`. fn park(&self) { - for _ in 0..3 { - // If we were previously notified then we consume this notification and - // return quickly. - if self - .state - .compare_exchange(NOTIFIED, EMPTY, SeqCst, SeqCst) - .is_ok() - { - return; - } - - thread::yield_now(); - } - - if let Some(mut driver) = self.shared.driver.try_lock() { - self.park_driver(&mut driver); - } else { - self.park_condvar(); + // If we were previously notified then we consume this notification and + // return quickly. + if self + .state + .compare_exchange(NOTIFIED, EMPTY, SeqCst, SeqCst) + .is_ok() + { + return; } - } - fn park_condvar(&self) { // Otherwise we need to coordinate going to sleep let mut m = self.mutex.lock(); - match self - .state - .compare_exchange(EMPTY, PARKED_CONDVAR, SeqCst, SeqCst) - { + match self.state.compare_exchange(EMPTY, PARKED, SeqCst, SeqCst) { Ok(_) => {} Err(NOTIFIED) => { // We must read here, even though we know it will be `NOTIFIED`. @@ -180,33 +129,44 @@ impl Inner { } } - fn park_driver(&self, driver: &mut Driver) { - match self + fn park_timeout(&self, dur: Duration) { + // Like `park` above we have a fast path for an already-notified thread, + // and afterwards we start coordinating for a sleep. Return quickly. + if self .state - .compare_exchange(EMPTY, PARKED_DRIVER, SeqCst, SeqCst) + .compare_exchange(NOTIFIED, EMPTY, SeqCst, SeqCst) + .is_ok() { + return; + } + + if dur == Duration::from_millis(0) { + return; + } + + let m = self.mutex.lock(); + + match self.state.compare_exchange(EMPTY, PARKED, SeqCst, SeqCst) { Ok(_) => {} Err(NOTIFIED) => { - // We must read here, even though we know it will be `NOTIFIED`. - // This is because `unpark` may have been called again since we read - // `NOTIFIED` in the `compare_exchange` above. We must perform an - // acquire operation that synchronizes with that `unpark` to observe - // any writes it made before the call to unpark. To do that we must - // read from the write it made to `state`. + // We must read again here, see `park`. let old = self.state.swap(EMPTY, SeqCst); debug_assert_eq!(old, NOTIFIED, "park state changed unexpectedly"); return; } - Err(actual) => panic!("inconsistent park state; actual = {}", actual), + Err(actual) => panic!("inconsistent park_timeout state; actual = {}", actual), } - // TODO: don't unwrap - driver.park().unwrap(); + // Wait with a timeout, and if we spuriously wake up or otherwise wake up + // from a notification, we just want to unconditionally set the state back to + // empty, either consuming a notification or un-flagging ourselves as + // parked. + let (_m, _result) = self.condvar.wait_timeout(m, dur).unwrap(); match self.state.swap(EMPTY, SeqCst) { - NOTIFIED => {} // got a notification, hurray! - PARKED_DRIVER => {} // no notification, alas + NOTIFIED => {} // got a notification, hurray! + PARKED => {} // no notification, alas n => panic!("inconsistent park_timeout state: {}", n), } } @@ -218,15 +178,12 @@ impl Inner { // is already `NOTIFIED`. That is why this must be a swap rather than a // compare-and-swap that returns if it reads `NOTIFIED` on failure. match self.state.swap(NOTIFIED, SeqCst) { - EMPTY => {} // no one was waiting - NOTIFIED => {} // already unparked - PARKED_CONDVAR => self.unpark_condvar(), - PARKED_DRIVER => self.unpark_driver(), - actual => panic!("inconsistent state in unpark; actual = {}", actual), + EMPTY => return, // no one was waiting + NOTIFIED => return, // already unparked + PARKED => {} // gotta go wake someone up + _ => panic!("inconsistent state in unpark"), } - } - fn unpark_condvar(&self) { // There is a period between when the parked thread sets `state` to // `PARKED` (or last checked `state` in the case of a spurious wake // up) and when it actually waits on `cvar`. If we were to notify @@ -243,15 +200,149 @@ impl Inner { self.condvar.notify_one() } - fn unpark_driver(&self) { - self.shared.handle.unpark(); + fn shutdown(&self) { + self.condvar.notify_all(); } +} - fn shutdown(&self) { - if let Some(mut driver) = self.shared.driver.try_lock() { - driver.shutdown(); +impl Default for ParkThread { + fn default() -> Self { + Self::new() + } +} + +// ===== impl UnparkThread ===== + +impl UnparkThread { + pub(crate) fn unpark(&self) { + self.inner.unpark(); + } +} + +use crate::loom::thread::AccessError; +use std::future::Future; +use std::marker::PhantomData; +use std::mem; +use std::rc::Rc; +use std::task::{RawWaker, RawWakerVTable, Waker}; + +/// Blocks the current thread using a condition variable. +#[derive(Debug)] +pub(crate) struct CachedParkThread { + _anchor: PhantomData<Rc<()>>, +} + +impl CachedParkThread { + /// Creates a new `ParkThread` handle for the current thread. + /// + /// This type cannot be moved to other threads, so it should be created on + /// the thread that the caller intends to park. + pub(crate) fn new() -> CachedParkThread { + CachedParkThread { + _anchor: PhantomData, } + } - self.condvar.notify_all(); + pub(crate) fn waker(&self) -> Result<Waker, AccessError> { + self.unpark().map(|unpark| unpark.into_waker()) } + + fn unpark(&self) -> Result<UnparkThread, AccessError> { + self.with_current(|park_thread| park_thread.unpark()) + } + + pub(crate) fn park(&mut self) { + self.with_current(|park_thread| park_thread.inner.park()) + .unwrap(); + } + + pub(crate) fn park_timeout(&mut self, duration: Duration) { + self.with_current(|park_thread| park_thread.inner.park_timeout(duration)) + .unwrap(); + } + + /// Gets a reference to the `ParkThread` handle for this thread. + fn with_current<F, R>(&self, f: F) -> Result<R, AccessError> + where + F: FnOnce(&ParkThread) -> R, + { + CURRENT_PARKER.try_with(|inner| f(inner)) + } + + pub(crate) fn block_on<F: Future>(&mut self, f: F) -> Result<F::Output, AccessError> { + use std::task::Context; + use std::task::Poll::Ready; + + // `get_unpark()` should not return a Result + let waker = self.waker()?; + let mut cx = Context::from_waker(&waker); + + pin!(f); + + loop { + if let Ready(v) = crate::runtime::coop::budget(|| f.as_mut().poll(&mut cx)) { + return Ok(v); + } + + self.park(); + } + } +} + +impl UnparkThread { + pub(crate) fn into_waker(self) -> Waker { + unsafe { + let raw = unparker_to_raw_waker(self.inner); + Waker::from_raw(raw) + } + } +} + +impl Inner { + #[allow(clippy::wrong_self_convention)] + fn into_raw(this: Arc<Inner>) -> *const () { + Arc::into_raw(this) as *const () + } + + unsafe fn from_raw(ptr: *const ()) -> Arc<Inner> { + Arc::from_raw(ptr as *const Inner) + } +} + +unsafe fn unparker_to_raw_waker(unparker: Arc<Inner>) -> RawWaker { + RawWaker::new( + Inner::into_raw(unparker), + &RawWakerVTable::new(clone, wake, wake_by_ref, drop_waker), + ) +} + +unsafe fn clone(raw: *const ()) -> RawWaker { + let unparker = Inner::from_raw(raw); + + // Increment the ref count + mem::forget(unparker.clone()); + + unparker_to_raw_waker(unparker) +} + +unsafe fn drop_waker(raw: *const ()) { + let _ = Inner::from_raw(raw); +} + +unsafe fn wake(raw: *const ()) { + let unparker = Inner::from_raw(raw); + unparker.unpark(); +} + +unsafe fn wake_by_ref(raw: *const ()) { + let unparker = Inner::from_raw(raw); + unparker.unpark(); + + // We don't actually own a reference to the unparker + mem::forget(unparker); +} + +#[cfg(loom)] +pub(crate) fn current_thread_park_count() -> usize { + CURRENT_THREAD_PARK_COUNT.with(|count| count.load(SeqCst)) } diff --git a/vendor/tokio/src/process/unix/driver.rs b/vendor/tokio/src/runtime/process.rs index 43b2efa0a..df339b0e7 100644 --- a/vendor/tokio/src/process/unix/driver.rs +++ b/vendor/tokio/src/runtime/process.rs @@ -1,12 +1,11 @@ #![cfg_attr(not(feature = "rt"), allow(dead_code))] -//! Process driver +//! Process driver. -use crate::park::Park; use crate::process::unix::GlobalOrphanQueue; -use crate::signal::unix::driver::{Driver as SignalDriver, Handle as SignalHandle}; +use crate::runtime::driver; +use crate::runtime::signal::{Driver as SignalDriver, Handle as SignalHandle}; -use std::io; use std::time::Duration; /// Responsible for cleaning up orphaned child processes on Unix platforms. @@ -28,31 +27,18 @@ impl Driver { signal_handle, } } -} - -// ===== impl Park for Driver ===== - -impl Park for Driver { - type Unpark = <SignalDriver as Park>::Unpark; - type Error = io::Error; - - fn unpark(&self) -> Self::Unpark { - self.park.unpark() - } - fn park(&mut self) -> Result<(), Self::Error> { - self.park.park()?; + pub(crate) fn park(&mut self, handle: &driver::Handle) { + self.park.park(handle); GlobalOrphanQueue::reap_orphans(&self.signal_handle); - Ok(()) } - fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> { - self.park.park_timeout(duration)?; + pub(crate) fn park_timeout(&mut self, handle: &driver::Handle, duration: Duration) { + self.park.park_timeout(handle, duration); GlobalOrphanQueue::reap_orphans(&self.signal_handle); - Ok(()) } - fn shutdown(&mut self) { - self.park.shutdown() + pub(crate) fn shutdown(&mut self, handle: &driver::Handle) { + self.park.shutdown(handle) } } diff --git a/vendor/tokio/src/runtime/runtime.rs b/vendor/tokio/src/runtime/runtime.rs new file mode 100644 index 000000000..3f3499975 --- /dev/null +++ b/vendor/tokio/src/runtime/runtime.rs @@ -0,0 +1,445 @@ +use crate::runtime::blocking::BlockingPool; +use crate::runtime::scheduler::CurrentThread; +use crate::runtime::{context, EnterGuard, Handle}; +use crate::task::JoinHandle; + +use std::future::Future; +use std::time::Duration; + +cfg_rt_multi_thread! { + use crate::runtime::Builder; + use crate::runtime::scheduler::MultiThread; +} + +/// The Tokio runtime. +/// +/// The runtime provides an I/O driver, task scheduler, [timer], and +/// blocking pool, necessary for running asynchronous tasks. +/// +/// Instances of `Runtime` can be created using [`new`], or [`Builder`]. +/// However, most users will use the `#[tokio::main]` annotation on their +/// entry point instead. +/// +/// See [module level][mod] documentation for more details. +/// +/// # Shutdown +/// +/// Shutting down the runtime is done by dropping the value, or calling +/// [`Runtime::shutdown_background`] or [`Runtime::shutdown_timeout`]. +/// +/// Tasks spawned through [`Runtime::spawn`] keep running until they yield. +/// Then they are dropped. They are not *guaranteed* to run to completion, but +/// *might* do so if they do not yield until completion. +/// +/// Blocking functions spawned through [`Runtime::spawn_blocking`] keep running +/// until they return. +/// +/// The thread initiating the shutdown blocks until all spawned work has been +/// stopped. This can take an indefinite amount of time. The `Drop` +/// implementation waits forever for this. +/// +/// `shutdown_background` and `shutdown_timeout` can be used if waiting forever +/// is undesired. When the timeout is reached, spawned work that did not stop +/// in time and threads running it are leaked. The work continues to run until +/// one of the stopping conditions is fulfilled, but the thread initiating the +/// shutdown is unblocked. +/// +/// Once the runtime has been dropped, any outstanding I/O resources bound to +/// it will no longer function. Calling any method on them will result in an +/// error. +/// +/// # Sharing +/// +/// The Tokio runtime implements `Sync` and `Send` to allow you to wrap it +/// in a `Arc`. Most fn take `&self` to allow you to call them concurrently +/// across multiple threads. +/// +/// Calls to `shutdown` and `shutdown_timeout` require exclusive ownership of +/// the runtime type and this can be achieved via `Arc::try_unwrap` when only +/// one strong count reference is left over. +/// +/// [timer]: crate::time +/// [mod]: index.html +/// [`new`]: method@Self::new +/// [`Builder`]: struct@Builder +#[derive(Debug)] +pub struct Runtime { + /// Task scheduler + scheduler: Scheduler, + + /// Handle to runtime, also contains driver handles + handle: Handle, + + /// Blocking pool handle, used to signal shutdown + blocking_pool: BlockingPool, +} + +/// The flavor of a `Runtime`. +/// +/// This is the return type for [`Handle::runtime_flavor`](crate::runtime::Handle::runtime_flavor()). +#[derive(Debug, PartialEq, Eq)] +#[non_exhaustive] +pub enum RuntimeFlavor { + /// The flavor that executes all tasks on the current thread. + CurrentThread, + /// The flavor that executes tasks across multiple threads. + MultiThread, +} + +/// The runtime scheduler is either a multi-thread or a current-thread executor. +#[derive(Debug)] +pub(super) enum Scheduler { + /// Execute all tasks on the current-thread. + CurrentThread(CurrentThread), + + /// Execute tasks across multiple threads. + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + MultiThread(MultiThread), +} + +impl Runtime { + pub(super) fn from_parts( + scheduler: Scheduler, + handle: Handle, + blocking_pool: BlockingPool, + ) -> Runtime { + Runtime { + scheduler, + handle, + blocking_pool, + } + } + + cfg_not_wasi! { + /// Creates a new runtime instance with default configuration values. + /// + /// This results in the multi threaded scheduler, I/O driver, and time driver being + /// initialized. + /// + /// Most applications will not need to call this function directly. Instead, + /// they will use the [`#[tokio::main]` attribute][main]. When a more complex + /// configuration is necessary, the [runtime builder] may be used. + /// + /// See [module level][mod] documentation for more details. + /// + /// # Examples + /// + /// Creating a new `Runtime` with default configuration values. + /// + /// ``` + /// use tokio::runtime::Runtime; + /// + /// let rt = Runtime::new() + /// .unwrap(); + /// + /// // Use the runtime... + /// ``` + /// + /// [mod]: index.html + /// [main]: ../attr.main.html + /// [threaded scheduler]: index.html#threaded-scheduler + /// [runtime builder]: crate::runtime::Builder + #[cfg(feature = "rt-multi-thread")] + #[cfg_attr(docsrs, doc(cfg(feature = "rt-multi-thread")))] + pub fn new() -> std::io::Result<Runtime> { + Builder::new_multi_thread().enable_all().build() + } + } + + /// Returns a handle to the runtime's spawner. + /// + /// The returned handle can be used to spawn tasks that run on this runtime, and can + /// be cloned to allow moving the `Handle` to other threads. + /// + /// Calling [`Handle::block_on`] on a handle to a `current_thread` runtime is error-prone. + /// Refer to the documentation of [`Handle::block_on`] for more. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Runtime; + /// + /// let rt = Runtime::new() + /// .unwrap(); + /// + /// let handle = rt.handle(); + /// + /// // Use the handle... + /// ``` + pub fn handle(&self) -> &Handle { + &self.handle + } + + /// Spawns a future onto the Tokio runtime. + /// + /// This spawns the given future onto the runtime's executor, usually a + /// thread pool. The thread pool is then responsible for polling the future + /// until it completes. + /// + /// The provided future will start running in the background immediately + /// when `spawn` is called, even if you don't await the returned + /// `JoinHandle`. + /// + /// See [module level][mod] documentation for more details. + /// + /// [mod]: index.html + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Runtime; + /// + /// # fn dox() { + /// // Create the runtime + /// let rt = Runtime::new().unwrap(); + /// + /// // Spawn a future onto the runtime + /// rt.spawn(async { + /// println!("now running on a worker thread"); + /// }); + /// # } + /// ``` + #[track_caller] + pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output> + where + F: Future + Send + 'static, + F::Output: Send + 'static, + { + self.handle.spawn(future) + } + + /// Runs the provided function on an executor dedicated to blocking operations. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Runtime; + /// + /// # fn dox() { + /// // Create the runtime + /// let rt = Runtime::new().unwrap(); + /// + /// // Spawn a blocking function onto the runtime + /// rt.spawn_blocking(|| { + /// println!("now running on a worker thread"); + /// }); + /// # } + #[track_caller] + pub fn spawn_blocking<F, R>(&self, func: F) -> JoinHandle<R> + where + F: FnOnce() -> R + Send + 'static, + R: Send + 'static, + { + self.handle.spawn_blocking(func) + } + + /// Runs a future to completion on the Tokio runtime. This is the + /// runtime's entry point. + /// + /// This runs the given future on the current thread, blocking until it is + /// complete, and yielding its resolved result. Any tasks or timers + /// which the future spawns internally will be executed on the runtime. + /// + /// # Non-worker future + /// + /// Note that the future required by this function does not run as a + /// worker. The expectation is that other tasks are spawned by the future here. + /// Awaiting on other futures from the future provided here will not + /// perform as fast as those spawned as workers. + /// + /// # Multi thread scheduler + /// + /// When the multi thread scheduler is used this will allow futures + /// to run within the io driver and timer context of the overall runtime. + /// + /// Any spawned tasks will continue running after `block_on` returns. + /// + /// # Current thread scheduler + /// + /// When the current thread scheduler is enabled `block_on` + /// can be called concurrently from multiple threads. The first call + /// will take ownership of the io and timer drivers. This means + /// other threads which do not own the drivers will hook into that one. + /// When the first `block_on` completes, other threads will be able to + /// "steal" the driver to allow continued execution of their futures. + /// + /// Any spawned tasks will be suspended after `block_on` returns. Calling + /// `block_on` again will resume previously spawned tasks. + /// + /// # Panics + /// + /// This function panics if the provided future panics, or if called within an + /// asynchronous execution context. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::runtime::Runtime; + /// + /// // Create the runtime + /// let rt = Runtime::new().unwrap(); + /// + /// // Execute the future, blocking the current thread until completion + /// rt.block_on(async { + /// println!("hello"); + /// }); + /// ``` + /// + /// [handle]: fn@Handle::block_on + #[track_caller] + pub fn block_on<F: Future>(&self, future: F) -> F::Output { + #[cfg(all( + tokio_unstable, + tokio_taskdump, + feature = "rt", + target_os = "linux", + any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64") + ))] + let future = super::task::trace::Trace::root(future); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + let future = crate::util::trace::task( + future, + "block_on", + None, + crate::runtime::task::Id::next().as_u64(), + ); + + let _enter = self.enter(); + + match &self.scheduler { + Scheduler::CurrentThread(exec) => exec.block_on(&self.handle.inner, future), + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + Scheduler::MultiThread(exec) => exec.block_on(&self.handle.inner, future), + } + } + + /// Enters the runtime context. + /// + /// This allows you to construct types that must have an executor + /// available on creation such as [`Sleep`] or [`TcpStream`]. It will + /// also allow you to call methods such as [`tokio::spawn`]. + /// + /// [`Sleep`]: struct@crate::time::Sleep + /// [`TcpStream`]: struct@crate::net::TcpStream + /// [`tokio::spawn`]: fn@crate::spawn + /// + /// # Example + /// + /// ``` + /// use tokio::runtime::Runtime; + /// + /// fn function_that_spawns(msg: String) { + /// // Had we not used `rt.enter` below, this would panic. + /// tokio::spawn(async move { + /// println!("{}", msg); + /// }); + /// } + /// + /// fn main() { + /// let rt = Runtime::new().unwrap(); + /// + /// let s = "Hello World!".to_string(); + /// + /// // By entering the context, we tie `tokio::spawn` to this executor. + /// let _guard = rt.enter(); + /// function_that_spawns(s); + /// } + /// ``` + pub fn enter(&self) -> EnterGuard<'_> { + self.handle.enter() + } + + /// Shuts down the runtime, waiting for at most `duration` for all spawned + /// work to stop. + /// + /// See the [struct level documentation](Runtime#shutdown) for more details. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Runtime; + /// use tokio::task; + /// + /// use std::thread; + /// use std::time::Duration; + /// + /// fn main() { + /// let runtime = Runtime::new().unwrap(); + /// + /// runtime.block_on(async move { + /// task::spawn_blocking(move || { + /// thread::sleep(Duration::from_secs(10_000)); + /// }); + /// }); + /// + /// runtime.shutdown_timeout(Duration::from_millis(100)); + /// } + /// ``` + pub fn shutdown_timeout(mut self, duration: Duration) { + // Wakeup and shutdown all the worker threads + self.handle.inner.shutdown(); + self.blocking_pool.shutdown(Some(duration)); + } + + /// Shuts down the runtime, without waiting for any spawned work to stop. + /// + /// This can be useful if you want to drop a runtime from within another runtime. + /// Normally, dropping a runtime will block indefinitely for spawned blocking tasks + /// to complete, which would normally not be permitted within an asynchronous context. + /// By calling `shutdown_background()`, you can drop the runtime from such a context. + /// + /// Note however, that because we do not wait for any blocking tasks to complete, this + /// may result in a resource leak (in that any blocking tasks are still running until they + /// return. + /// + /// See the [struct level documentation](Runtime#shutdown) for more details. + /// + /// This function is equivalent to calling `shutdown_timeout(Duration::from_nanos(0))`. + /// + /// ``` + /// use tokio::runtime::Runtime; + /// + /// fn main() { + /// let runtime = Runtime::new().unwrap(); + /// + /// runtime.block_on(async move { + /// let inner_runtime = Runtime::new().unwrap(); + /// // ... + /// inner_runtime.shutdown_background(); + /// }); + /// } + /// ``` + pub fn shutdown_background(self) { + self.shutdown_timeout(Duration::from_nanos(0)) + } +} + +#[allow(clippy::single_match)] // there are comments in the error branch, so we don't want if-let +impl Drop for Runtime { + fn drop(&mut self) { + match &mut self.scheduler { + Scheduler::CurrentThread(current_thread) => { + // This ensures that tasks spawned on the current-thread + // runtime are dropped inside the runtime's context. + let _guard = context::try_set_current(&self.handle.inner); + current_thread.shutdown(&self.handle.inner); + } + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + Scheduler::MultiThread(multi_thread) => { + // The threaded scheduler drops its tasks on its worker threads, which is + // already in the runtime's context. + multi_thread.shutdown(&self.handle.inner); + } + } + } +} + +cfg_metrics! { + impl Runtime { + /// TODO + pub fn metrics(&self) -> crate::runtime::RuntimeMetrics { + self.handle.metrics() + } + } +} diff --git a/vendor/tokio/src/runtime/scheduler/current_thread.rs b/vendor/tokio/src/runtime/scheduler/current_thread.rs new file mode 100644 index 000000000..ac4a8d6fa --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/current_thread.rs @@ -0,0 +1,750 @@ +use crate::future::poll_fn; +use crate::loom::sync::atomic::AtomicBool; +use crate::loom::sync::Arc; +use crate::runtime::driver::{self, Driver}; +use crate::runtime::scheduler::{self, Defer, Inject}; +use crate::runtime::task::{self, JoinHandle, OwnedTasks, Schedule, Task}; +use crate::runtime::{blocking, context, Config, MetricsBatch, SchedulerMetrics, WorkerMetrics}; +use crate::sync::notify::Notify; +use crate::util::atomic_cell::AtomicCell; +use crate::util::{waker_ref, RngSeedGenerator, Wake, WakerRef}; + +use std::cell::RefCell; +use std::collections::VecDeque; +use std::fmt; +use std::future::Future; +use std::sync::atomic::Ordering::{AcqRel, Release}; +use std::task::Poll::{Pending, Ready}; +use std::task::Waker; +use std::time::Duration; + +/// Executes tasks on the current thread +pub(crate) struct CurrentThread { + /// Core scheduler data is acquired by a thread entering `block_on`. + core: AtomicCell<Core>, + + /// Notifier for waking up other threads to steal the + /// driver. + notify: Notify, +} + +/// Handle to the current thread scheduler +pub(crate) struct Handle { + /// Scheduler state shared across threads + shared: Shared, + + /// Resource driver handles + pub(crate) driver: driver::Handle, + + /// Blocking pool spawner + pub(crate) blocking_spawner: blocking::Spawner, + + /// Current random number generator seed + pub(crate) seed_generator: RngSeedGenerator, +} + +/// Data required for executing the scheduler. The struct is passed around to +/// a function that will perform the scheduling work and acts as a capability token. +struct Core { + /// Scheduler run queue + tasks: VecDeque<Notified>, + + /// Current tick + tick: u32, + + /// Runtime driver + /// + /// The driver is removed before starting to park the thread + driver: Option<Driver>, + + /// Metrics batch + metrics: MetricsBatch, + + /// How often to check the global queue + global_queue_interval: u32, + + /// True if a task panicked without being handled and the runtime is + /// configured to shutdown on unhandled panic. + unhandled_panic: bool, +} + +/// Scheduler state shared between threads. +struct Shared { + /// Remote run queue + inject: Inject<Arc<Handle>>, + + /// Collection of all active tasks spawned onto this executor. + owned: OwnedTasks<Arc<Handle>>, + + /// Indicates whether the blocked on thread was woken. + woken: AtomicBool, + + /// Scheduler configuration options + config: Config, + + /// Keeps track of various runtime metrics. + scheduler_metrics: SchedulerMetrics, + + /// This scheduler only has one worker. + worker_metrics: WorkerMetrics, +} + +/// Thread-local context. +/// +/// pub(crate) to store in `runtime::context`. +pub(crate) struct Context { + /// Scheduler handle + handle: Arc<Handle>, + + /// Scheduler core, enabling the holder of `Context` to execute the + /// scheduler. + core: RefCell<Option<Box<Core>>>, + + /// Deferred tasks, usually ones that called `task::yield_now()`. + pub(crate) defer: Defer, +} + +type Notified = task::Notified<Arc<Handle>>; + +/// Initial queue capacity. +const INITIAL_CAPACITY: usize = 64; + +/// Used if none is specified. This is a temporary constant and will be removed +/// as we unify tuning logic between the multi-thread and current-thread +/// schedulers. +const DEFAULT_GLOBAL_QUEUE_INTERVAL: u32 = 31; + +impl CurrentThread { + pub(crate) fn new( + driver: Driver, + driver_handle: driver::Handle, + blocking_spawner: blocking::Spawner, + seed_generator: RngSeedGenerator, + config: Config, + ) -> (CurrentThread, Arc<Handle>) { + let worker_metrics = WorkerMetrics::from_config(&config); + + // Get the configured global queue interval, or use the default. + let global_queue_interval = config + .global_queue_interval + .unwrap_or(DEFAULT_GLOBAL_QUEUE_INTERVAL); + + let handle = Arc::new(Handle { + shared: Shared { + inject: Inject::new(), + owned: OwnedTasks::new(), + woken: AtomicBool::new(false), + config, + scheduler_metrics: SchedulerMetrics::new(), + worker_metrics, + }, + driver: driver_handle, + blocking_spawner, + seed_generator, + }); + + let core = AtomicCell::new(Some(Box::new(Core { + tasks: VecDeque::with_capacity(INITIAL_CAPACITY), + tick: 0, + driver: Some(driver), + metrics: MetricsBatch::new(&handle.shared.worker_metrics), + global_queue_interval, + unhandled_panic: false, + }))); + + let scheduler = CurrentThread { + core, + notify: Notify::new(), + }; + + (scheduler, handle) + } + + #[track_caller] + pub(crate) fn block_on<F: Future>(&self, handle: &scheduler::Handle, future: F) -> F::Output { + pin!(future); + + crate::runtime::context::enter_runtime(handle, false, |blocking| { + let handle = handle.as_current_thread(); + + // Attempt to steal the scheduler core and block_on the future if we can + // there, otherwise, lets select on a notification that the core is + // available or the future is complete. + loop { + if let Some(core) = self.take_core(handle) { + return core.block_on(future); + } else { + let notified = self.notify.notified(); + pin!(notified); + + if let Some(out) = blocking + .block_on(poll_fn(|cx| { + if notified.as_mut().poll(cx).is_ready() { + return Ready(None); + } + + if let Ready(out) = future.as_mut().poll(cx) { + return Ready(Some(out)); + } + + Pending + })) + .expect("Failed to `Enter::block_on`") + { + return out; + } + } + } + }) + } + + fn take_core(&self, handle: &Arc<Handle>) -> Option<CoreGuard<'_>> { + let core = self.core.take()?; + + Some(CoreGuard { + context: scheduler::Context::CurrentThread(Context { + handle: handle.clone(), + core: RefCell::new(Some(core)), + defer: Defer::new(), + }), + scheduler: self, + }) + } + + pub(crate) fn shutdown(&mut self, handle: &scheduler::Handle) { + let handle = handle.as_current_thread(); + + // Avoid a double panic if we are currently panicking and + // the lock may be poisoned. + + let core = match self.take_core(handle) { + Some(core) => core, + None if std::thread::panicking() => return, + None => panic!("Oh no! We never placed the Core back, this is a bug!"), + }; + + // Check that the thread-local is not being destroyed + let tls_available = context::with_current(|_| ()).is_ok(); + + if tls_available { + core.enter(|core, _context| { + let core = shutdown2(core, handle); + (core, ()) + }); + } else { + // Shutdown without setting the context. `tokio::spawn` calls will + // fail, but those will fail either way because the thread-local is + // not available anymore. + let context = core.context.expect_current_thread(); + let core = context.core.borrow_mut().take().unwrap(); + + let core = shutdown2(core, handle); + *context.core.borrow_mut() = Some(core); + } + } +} + +fn shutdown2(mut core: Box<Core>, handle: &Handle) -> Box<Core> { + // Drain the OwnedTasks collection. This call also closes the + // collection, ensuring that no tasks are ever pushed after this + // call returns. + handle.shared.owned.close_and_shutdown_all(); + + // Drain local queue + // We already shut down every task, so we just need to drop the task. + while let Some(task) = core.next_local_task(handle) { + drop(task); + } + + // Close the injection queue + handle.shared.inject.close(); + + // Drain remote queue + while let Some(task) = handle.shared.inject.pop() { + drop(task); + } + + assert!(handle.shared.owned.is_empty()); + + // Submit metrics + core.submit_metrics(handle); + + // Shutdown the resource drivers + if let Some(driver) = core.driver.as_mut() { + driver.shutdown(&handle.driver); + } + + core +} + +impl fmt::Debug for CurrentThread { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt.debug_struct("CurrentThread").finish() + } +} + +// ===== impl Core ===== + +impl Core { + /// Get and increment the current tick + fn tick(&mut self) { + self.tick = self.tick.wrapping_add(1); + } + + fn next_task(&mut self, handle: &Handle) -> Option<Notified> { + if self.tick % self.global_queue_interval == 0 { + handle + .next_remote_task() + .or_else(|| self.next_local_task(handle)) + } else { + self.next_local_task(handle) + .or_else(|| handle.next_remote_task()) + } + } + + fn next_local_task(&mut self, handle: &Handle) -> Option<Notified> { + let ret = self.tasks.pop_front(); + handle + .shared + .worker_metrics + .set_queue_depth(self.tasks.len()); + ret + } + + fn push_task(&mut self, handle: &Handle, task: Notified) { + self.tasks.push_back(task); + self.metrics.inc_local_schedule_count(); + handle + .shared + .worker_metrics + .set_queue_depth(self.tasks.len()); + } + + fn submit_metrics(&mut self, handle: &Handle) { + self.metrics.submit(&handle.shared.worker_metrics); + } +} + +#[cfg(tokio_taskdump)] +fn wake_deferred_tasks_and_free(context: &Context) { + let wakers = context.defer.take_deferred(); + for waker in wakers { + waker.wake(); + } +} + +// ===== impl Context ===== + +impl Context { + /// Execute the closure with the given scheduler core stored in the + /// thread-local context. + fn run_task<R>(&self, mut core: Box<Core>, f: impl FnOnce() -> R) -> (Box<Core>, R) { + core.metrics.start_poll(); + let mut ret = self.enter(core, || crate::runtime::coop::budget(f)); + ret.0.metrics.end_poll(); + ret + } + + /// Blocks the current thread until an event is received by the driver, + /// including I/O events, timer events, ... + fn park(&self, mut core: Box<Core>, handle: &Handle) -> Box<Core> { + let mut driver = core.driver.take().expect("driver missing"); + + if let Some(f) = &handle.shared.config.before_park { + // Incorrect lint, the closures are actually different types so `f` + // cannot be passed as an argument to `enter`. + #[allow(clippy::redundant_closure)] + let (c, _) = self.enter(core, || f()); + core = c; + } + + // This check will fail if `before_park` spawns a task for us to run + // instead of parking the thread + if core.tasks.is_empty() { + // Park until the thread is signaled + core.metrics.about_to_park(); + core.submit_metrics(handle); + + let (c, _) = self.enter(core, || { + driver.park(&handle.driver); + self.defer.wake(); + }); + + core = c; + } + + if let Some(f) = &handle.shared.config.after_unpark { + // Incorrect lint, the closures are actually different types so `f` + // cannot be passed as an argument to `enter`. + #[allow(clippy::redundant_closure)] + let (c, _) = self.enter(core, || f()); + core = c; + } + + core.driver = Some(driver); + core + } + + /// Checks the driver for new events without blocking the thread. + fn park_yield(&self, mut core: Box<Core>, handle: &Handle) -> Box<Core> { + let mut driver = core.driver.take().expect("driver missing"); + + core.submit_metrics(handle); + + let (mut core, _) = self.enter(core, || { + driver.park_timeout(&handle.driver, Duration::from_millis(0)); + self.defer.wake(); + }); + + core.driver = Some(driver); + core + } + + fn enter<R>(&self, core: Box<Core>, f: impl FnOnce() -> R) -> (Box<Core>, R) { + // Store the scheduler core in the thread-local context + // + // A drop-guard is employed at a higher level. + *self.core.borrow_mut() = Some(core); + + // Execute the closure while tracking the execution budget + let ret = f(); + + // Take the scheduler core back + let core = self.core.borrow_mut().take().expect("core missing"); + (core, ret) + } + + pub(crate) fn defer(&self, waker: &Waker) { + self.defer.defer(waker); + } +} + +// ===== impl Handle ===== + +impl Handle { + /// Spawns a future onto the `CurrentThread` scheduler + pub(crate) fn spawn<F>( + me: &Arc<Self>, + future: F, + id: crate::runtime::task::Id, + ) -> JoinHandle<F::Output> + where + F: crate::future::Future + Send + 'static, + F::Output: Send + 'static, + { + let (handle, notified) = me.shared.owned.bind(future, me.clone(), id); + + if let Some(notified) = notified { + me.schedule(notified); + } + + handle + } + + /// Capture a snapshot of this runtime's state. + #[cfg(all( + tokio_unstable, + tokio_taskdump, + target_os = "linux", + any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64") + ))] + pub(crate) fn dump(&self) -> crate::runtime::Dump { + use crate::runtime::dump; + use task::trace::trace_current_thread; + + let mut traces = vec![]; + + // todo: how to make this work outside of a runtime context? + context::with_scheduler(|maybe_context| { + // drain the local queue + let context = if let Some(context) = maybe_context { + context.expect_current_thread() + } else { + return; + }; + let mut maybe_core = context.core.borrow_mut(); + let core = if let Some(core) = maybe_core.as_mut() { + core + } else { + return; + }; + let local = &mut core.tasks; + + if self.shared.inject.is_closed() { + return; + } + + traces = trace_current_thread(&self.shared.owned, local, &self.shared.inject) + .into_iter() + .map(dump::Task::new) + .collect(); + + // Avoid double borrow panic + drop(maybe_core); + + // Taking a taskdump could wakes every task, but we probably don't want + // the `yield_now` vector to be that large under normal circumstances. + // Therefore, we free its allocation. + wake_deferred_tasks_and_free(context); + }); + + dump::Dump::new(traces) + } + + fn next_remote_task(&self) -> Option<Notified> { + self.shared.inject.pop() + } + + fn waker_ref(me: &Arc<Self>) -> WakerRef<'_> { + // Set woken to true when enter block_on, ensure outer future + // be polled for the first time when enter loop + me.shared.woken.store(true, Release); + waker_ref(me) + } + + // reset woken to false and return original value + pub(crate) fn reset_woken(&self) -> bool { + self.shared.woken.swap(false, AcqRel) + } +} + +cfg_metrics! { + impl Handle { + pub(crate) fn scheduler_metrics(&self) -> &SchedulerMetrics { + &self.shared.scheduler_metrics + } + + pub(crate) fn injection_queue_depth(&self) -> usize { + self.shared.inject.len() + } + + pub(crate) fn worker_metrics(&self, worker: usize) -> &WorkerMetrics { + assert_eq!(0, worker); + &self.shared.worker_metrics + } + + pub(crate) fn num_blocking_threads(&self) -> usize { + self.blocking_spawner.num_threads() + } + + pub(crate) fn num_idle_blocking_threads(&self) -> usize { + self.blocking_spawner.num_idle_threads() + } + + pub(crate) fn blocking_queue_depth(&self) -> usize { + self.blocking_spawner.queue_depth() + } + + pub(crate) fn active_tasks_count(&self) -> usize { + self.shared.owned.active_tasks_count() + } + } +} + +impl fmt::Debug for Handle { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt.debug_struct("current_thread::Handle { ... }").finish() + } +} + +// ===== impl Shared ===== + +impl Schedule for Arc<Handle> { + fn release(&self, task: &Task<Self>) -> Option<Task<Self>> { + self.shared.owned.remove(task) + } + + fn schedule(&self, task: task::Notified<Self>) { + use scheduler::Context::CurrentThread; + + context::with_scheduler(|maybe_cx| match maybe_cx { + Some(CurrentThread(cx)) if Arc::ptr_eq(self, &cx.handle) => { + let mut core = cx.core.borrow_mut(); + + // If `None`, the runtime is shutting down, so there is no need + // to schedule the task. + if let Some(core) = core.as_mut() { + core.push_task(self, task); + } + } + _ => { + // Track that a task was scheduled from **outside** of the runtime. + self.shared.scheduler_metrics.inc_remote_schedule_count(); + + // Schedule the task + self.shared.inject.push(task); + self.driver.unpark(); + } + }); + } + + cfg_unstable! { + fn unhandled_panic(&self) { + use crate::runtime::UnhandledPanic; + + match self.shared.config.unhandled_panic { + UnhandledPanic::Ignore => { + // Do nothing + } + UnhandledPanic::ShutdownRuntime => { + use scheduler::Context::CurrentThread; + + // This hook is only called from within the runtime, so + // `context::with_scheduler` should match with `&self`, i.e. + // there is no opportunity for a nested scheduler to be + // called. + context::with_scheduler(|maybe_cx| match maybe_cx { + Some(CurrentThread(cx)) if Arc::ptr_eq(self, &cx.handle) => { + let mut core = cx.core.borrow_mut(); + + // If `None`, the runtime is shutting down, so there is no need to signal shutdown + if let Some(core) = core.as_mut() { + core.unhandled_panic = true; + self.shared.owned.close_and_shutdown_all(); + } + } + _ => unreachable!("runtime core not set in CURRENT thread-local"), + }) + } + } + } + } +} + +impl Wake for Handle { + fn wake(arc_self: Arc<Self>) { + Wake::wake_by_ref(&arc_self) + } + + /// Wake by reference + fn wake_by_ref(arc_self: &Arc<Self>) { + arc_self.shared.woken.store(true, Release); + arc_self.driver.unpark(); + } +} + +// ===== CoreGuard ===== + +/// Used to ensure we always place the `Core` value back into its slot in +/// `CurrentThread`, even if the future panics. +struct CoreGuard<'a> { + context: scheduler::Context, + scheduler: &'a CurrentThread, +} + +impl CoreGuard<'_> { + #[track_caller] + fn block_on<F: Future>(self, future: F) -> F::Output { + let ret = self.enter(|mut core, context| { + let waker = Handle::waker_ref(&context.handle); + let mut cx = std::task::Context::from_waker(&waker); + + pin!(future); + + core.metrics.start_processing_scheduled_tasks(); + + 'outer: loop { + let handle = &context.handle; + + if handle.reset_woken() { + let (c, res) = context.enter(core, || { + crate::runtime::coop::budget(|| future.as_mut().poll(&mut cx)) + }); + + core = c; + + if let Ready(v) = res { + return (core, Some(v)); + } + } + + for _ in 0..handle.shared.config.event_interval { + // Make sure we didn't hit an unhandled_panic + if core.unhandled_panic { + return (core, None); + } + + core.tick(); + + let entry = core.next_task(handle); + + let task = match entry { + Some(entry) => entry, + None => { + core.metrics.end_processing_scheduled_tasks(); + + core = if !context.defer.is_empty() { + context.park_yield(core, handle) + } else { + context.park(core, handle) + }; + + core.metrics.start_processing_scheduled_tasks(); + + // Try polling the `block_on` future next + continue 'outer; + } + }; + + let task = context.handle.shared.owned.assert_owner(task); + + let (c, _) = context.run_task(core, || { + task.run(); + }); + + core = c; + } + + core.metrics.end_processing_scheduled_tasks(); + + // Yield to the driver, this drives the timer and pulls any + // pending I/O events. + core = context.park_yield(core, handle); + + core.metrics.start_processing_scheduled_tasks(); + } + }); + + match ret { + Some(ret) => ret, + None => { + // `block_on` panicked. + panic!("a spawned task panicked and the runtime is configured to shut down on unhandled panic"); + } + } + } + + /// Enters the scheduler context. This sets the queue and other necessary + /// scheduler state in the thread-local. + fn enter<F, R>(self, f: F) -> R + where + F: FnOnce(Box<Core>, &Context) -> (Box<Core>, R), + { + let context = self.context.expect_current_thread(); + + // Remove `core` from `context` to pass into the closure. + let core = context.core.borrow_mut().take().expect("core missing"); + + // Call the closure and place `core` back + let (core, ret) = context::set_scheduler(&self.context, || f(core, context)); + + *context.core.borrow_mut() = Some(core); + + ret + } +} + +impl Drop for CoreGuard<'_> { + fn drop(&mut self) { + let context = self.context.expect_current_thread(); + + if let Some(core) = context.core.borrow_mut().take() { + // Replace old scheduler back into the state to allow + // other threads to pick it up and drive it. + self.scheduler.core.set(core); + + // Wake up other possible threads that could steal the driver. + self.scheduler.notify.notify_one() + } + } +} diff --git a/vendor/tokio/src/runtime/scheduler/defer.rs b/vendor/tokio/src/runtime/scheduler/defer.rs new file mode 100644 index 000000000..a4be8ef2e --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/defer.rs @@ -0,0 +1,43 @@ +use std::cell::RefCell; +use std::task::Waker; + +pub(crate) struct Defer { + deferred: RefCell<Vec<Waker>>, +} + +impl Defer { + pub(crate) fn new() -> Defer { + Defer { + deferred: Default::default(), + } + } + + pub(crate) fn defer(&self, waker: &Waker) { + let mut deferred = self.deferred.borrow_mut(); + + // If the same task adds itself a bunch of times, then only add it once. + if let Some(last) = deferred.last() { + if last.will_wake(waker) { + return; + } + } + + deferred.push(waker.clone()); + } + + pub(crate) fn is_empty(&self) -> bool { + self.deferred.borrow().is_empty() + } + + pub(crate) fn wake(&self) { + while let Some(waker) = self.deferred.borrow_mut().pop() { + waker.wake(); + } + } + + #[cfg(tokio_taskdump)] + pub(crate) fn take_deferred(&self) -> Vec<Waker> { + let mut deferred = self.deferred.borrow_mut(); + std::mem::take(&mut *deferred) + } +} diff --git a/vendor/tokio/src/runtime/scheduler/inject.rs b/vendor/tokio/src/runtime/scheduler/inject.rs new file mode 100644 index 000000000..39976fcd7 --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/inject.rs @@ -0,0 +1,72 @@ +//! Inject queue used to send wakeups to a work-stealing scheduler + +use crate::loom::sync::Mutex; +use crate::runtime::task; + +mod pop; +pub(crate) use pop::Pop; + +mod shared; +pub(crate) use shared::Shared; + +mod synced; +pub(crate) use synced::Synced; + +cfg_rt_multi_thread! { + mod rt_multi_thread; +} + +cfg_metrics! { + mod metrics; +} + +/// Growable, MPMC queue used to inject new tasks into the scheduler and as an +/// overflow queue when the local, fixed-size, array queue overflows. +pub(crate) struct Inject<T: 'static> { + shared: Shared<T>, + synced: Mutex<Synced>, +} + +impl<T: 'static> Inject<T> { + pub(crate) fn new() -> Inject<T> { + let (shared, synced) = Shared::new(); + + Inject { + shared, + synced: Mutex::new(synced), + } + } + + // Kind of annoying to have to include the cfg here + #[cfg(tokio_taskdump)] + pub(crate) fn is_closed(&self) -> bool { + let synced = self.synced.lock(); + self.shared.is_closed(&synced) + } + + /// Closes the injection queue, returns `true` if the queue is open when the + /// transition is made. + pub(crate) fn close(&self) -> bool { + let mut synced = self.synced.lock(); + self.shared.close(&mut synced) + } + + /// Pushes a value into the queue. + /// + /// This does nothing if the queue is closed. + pub(crate) fn push(&self, task: task::Notified<T>) { + let mut synced = self.synced.lock(); + // safety: passing correct `Synced` + unsafe { self.shared.push(&mut synced, task) } + } + + pub(crate) fn pop(&self) -> Option<task::Notified<T>> { + if self.shared.is_empty() { + return None; + } + + let mut synced = self.synced.lock(); + // safety: passing correct `Synced` + unsafe { self.shared.pop(&mut synced) } + } +} diff --git a/vendor/tokio/src/runtime/scheduler/inject/metrics.rs b/vendor/tokio/src/runtime/scheduler/inject/metrics.rs new file mode 100644 index 000000000..76f045fdb --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/inject/metrics.rs @@ -0,0 +1,7 @@ +use super::Inject; + +impl<T: 'static> Inject<T> { + pub(crate) fn len(&self) -> usize { + self.shared.len() + } +} diff --git a/vendor/tokio/src/runtime/scheduler/inject/pop.rs b/vendor/tokio/src/runtime/scheduler/inject/pop.rs new file mode 100644 index 000000000..4e6d5d3be --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/inject/pop.rs @@ -0,0 +1,55 @@ +use super::Synced; + +use crate::runtime::task; + +use std::marker::PhantomData; + +pub(crate) struct Pop<'a, T: 'static> { + len: usize, + synced: &'a mut Synced, + _p: PhantomData<T>, +} + +impl<'a, T: 'static> Pop<'a, T> { + pub(super) fn new(len: usize, synced: &'a mut Synced) -> Pop<'a, T> { + Pop { + len, + synced, + _p: PhantomData, + } + } +} + +impl<'a, T: 'static> Iterator for Pop<'a, T> { + type Item = task::Notified<T>; + + fn next(&mut self) -> Option<Self::Item> { + if self.len == 0 { + return None; + } + + let ret = self.synced.pop(); + + // Should be `Some` when `len > 0` + debug_assert!(ret.is_some()); + + self.len -= 1; + ret + } + + fn size_hint(&self) -> (usize, Option<usize>) { + (self.len, Some(self.len)) + } +} + +impl<'a, T: 'static> ExactSizeIterator for Pop<'a, T> { + fn len(&self) -> usize { + self.len + } +} + +impl<'a, T: 'static> Drop for Pop<'a, T> { + fn drop(&mut self) { + for _ in self.by_ref() {} + } +} diff --git a/vendor/tokio/src/runtime/scheduler/inject/rt_multi_thread.rs b/vendor/tokio/src/runtime/scheduler/inject/rt_multi_thread.rs new file mode 100644 index 000000000..07d1063c5 --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/inject/rt_multi_thread.rs @@ -0,0 +1,98 @@ +use super::{Shared, Synced}; + +use crate::runtime::scheduler::Lock; +use crate::runtime::task; + +use std::sync::atomic::Ordering::Release; + +impl<'a> Lock<Synced> for &'a mut Synced { + type Handle = &'a mut Synced; + + fn lock(self) -> Self::Handle { + self + } +} + +impl AsMut<Synced> for Synced { + fn as_mut(&mut self) -> &mut Synced { + self + } +} + +impl<T: 'static> Shared<T> { + /// Pushes several values into the queue. + /// + /// # Safety + /// + /// Must be called with the same `Synced` instance returned by `Inject::new` + #[inline] + pub(crate) unsafe fn push_batch<L, I>(&self, shared: L, mut iter: I) + where + L: Lock<Synced>, + I: Iterator<Item = task::Notified<T>>, + { + let first = match iter.next() { + Some(first) => first.into_raw(), + None => return, + }; + + // Link up all the tasks. + let mut prev = first; + let mut counter = 1; + + // We are going to be called with an `std::iter::Chain`, and that + // iterator overrides `for_each` to something that is easier for the + // compiler to optimize than a loop. + iter.for_each(|next| { + let next = next.into_raw(); + + // safety: Holding the Notified for a task guarantees exclusive + // access to the `queue_next` field. + unsafe { prev.set_queue_next(Some(next)) }; + prev = next; + counter += 1; + }); + + // Now that the tasks are linked together, insert them into the + // linked list. + self.push_batch_inner(shared, first, prev, counter); + } + + /// Inserts several tasks that have been linked together into the queue. + /// + /// The provided head and tail may be be the same task. In this case, a + /// single task is inserted. + #[inline] + unsafe fn push_batch_inner<L>( + &self, + shared: L, + batch_head: task::RawTask, + batch_tail: task::RawTask, + num: usize, + ) where + L: Lock<Synced>, + { + debug_assert!(unsafe { batch_tail.get_queue_next().is_none() }); + + let mut synced = shared.lock(); + let synced = synced.as_mut(); + + if let Some(tail) = synced.tail { + unsafe { + tail.set_queue_next(Some(batch_head)); + } + } else { + synced.head = Some(batch_head); + } + + synced.tail = Some(batch_tail); + + // Increment the count. + // + // safety: All updates to the len atomic are guarded by the mutex. As + // such, a non-atomic load followed by a store is safe. + let len = self.len.unsync_load(); + + self.len.store(len + num, Release); + } +} diff --git a/vendor/tokio/src/runtime/scheduler/inject/shared.rs b/vendor/tokio/src/runtime/scheduler/inject/shared.rs new file mode 100644 index 000000000..7fdd2839d --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/inject/shared.rs @@ -0,0 +1,119 @@ +use super::{Pop, Synced}; + +use crate::loom::sync::atomic::AtomicUsize; +use crate::runtime::task; + +use std::marker::PhantomData; +use std::sync::atomic::Ordering::{Acquire, Release}; + +pub(crate) struct Shared<T: 'static> { + /// Number of pending tasks in the queue. This helps prevent unnecessary + /// locking in the hot path. + pub(super) len: AtomicUsize, + + _p: PhantomData<T>, +} + +unsafe impl<T> Send for Shared<T> {} +unsafe impl<T> Sync for Shared<T> {} + +impl<T: 'static> Shared<T> { + pub(crate) fn new() -> (Shared<T>, Synced) { + let inject = Shared { + len: AtomicUsize::new(0), + _p: PhantomData, + }; + + let synced = Synced { + is_closed: false, + head: None, + tail: None, + }; + + (inject, synced) + } + + pub(crate) fn is_empty(&self) -> bool { + self.len() == 0 + } + + // Kind of annoying to have to include the cfg here + #[cfg(any(tokio_taskdump, all(feature = "rt-multi-thread", not(tokio_wasi))))] + pub(crate) fn is_closed(&self, synced: &Synced) -> bool { + synced.is_closed + } + + /// Closes the injection queue, returns `true` if the queue is open when the + /// transition is made. + pub(crate) fn close(&self, synced: &mut Synced) -> bool { + if synced.is_closed { + return false; + } + + synced.is_closed = true; + true + } + + pub(crate) fn len(&self) -> usize { + self.len.load(Acquire) + } + + /// Pushes a value into the queue. + /// + /// This does nothing if the queue is closed. + /// + /// # Safety + /// + /// Must be called with the same `Synced` instance returned by `Inject::new` + pub(crate) unsafe fn push(&self, synced: &mut Synced, task: task::Notified<T>) { + if synced.is_closed { + return; + } + + // safety: only mutated with the lock held + let len = self.len.unsync_load(); + let task = task.into_raw(); + + // The next pointer should already be null + debug_assert!(unsafe { task.get_queue_next().is_none() }); + + if let Some(tail) = synced.tail { + // safety: Holding the Notified for a task guarantees exclusive + // access to the `queue_next` field. + unsafe { tail.set_queue_next(Some(task)) }; + } else { + synced.head = Some(task); + } + + synced.tail = Some(task); + self.len.store(len + 1, Release); + } + + /// Pop a value from the queue. + /// + /// # Safety + /// + /// Must be called with the same `Synced` instance returned by `Inject::new` + pub(crate) unsafe fn pop(&self, synced: &mut Synced) -> Option<task::Notified<T>> { + self.pop_n(synced, 1).next() + } + + /// Pop `n` values from the queue + /// + /// # Safety + /// + /// Must be called with the same `Synced` instance returned by `Inject::new` + pub(crate) unsafe fn pop_n<'a>(&'a self, synced: &'a mut Synced, n: usize) -> Pop<'a, T> { + use std::cmp; + + // safety: All updates to the len atomic are guarded by the mutex. As + // such, a non-atomic load followed by a store is safe. + let len = self.len.unsync_load(); + let n = cmp::min(n, len); + + // Decrement the count. + self.len.store(len - n, Release); + + Pop::new(n, synced) + } +} diff --git a/vendor/tokio/src/runtime/scheduler/inject/synced.rs b/vendor/tokio/src/runtime/scheduler/inject/synced.rs new file mode 100644 index 000000000..6847f68e5 --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/inject/synced.rs @@ -0,0 +1,32 @@ +use crate::runtime::task; + +pub(crate) struct Synced { + /// True if the queue is closed. + pub(super) is_closed: bool, + + /// Linked-list head. + pub(super) head: Option<task::RawTask>, + + /// Linked-list tail. + pub(super) tail: Option<task::RawTask>, +} + +unsafe impl Send for Synced {} +unsafe impl Sync for Synced {} + +impl Synced { + pub(super) fn pop<T: 'static>(&mut self) -> Option<task::Notified<T>> { + let task = self.head?; + + self.head = unsafe { task.get_queue_next() }; + + if self.head.is_none() { + self.tail = None; + } + + unsafe { task.set_queue_next(None) }; + + // safety: a `Notified` is pushed into the queue and now it is popped! + Some(unsafe { task::Notified::from_raw(task) }) + } +} diff --git a/vendor/tokio/src/runtime/scheduler/lock.rs b/vendor/tokio/src/runtime/scheduler/lock.rs new file mode 100644 index 000000000..0901c2b37 --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/lock.rs @@ -0,0 +1,6 @@ +/// A lock (mutex) yielding generic data. +pub(crate) trait Lock<T> { + type Handle: AsMut<T>; + + fn lock(self) -> Self::Handle; +} diff --git a/vendor/tokio/src/runtime/scheduler/mod.rs b/vendor/tokio/src/runtime/scheduler/mod.rs new file mode 100644 index 000000000..3e3151711 --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/mod.rs @@ -0,0 +1,249 @@ +cfg_rt! { + pub(crate) mod current_thread; + pub(crate) use current_thread::CurrentThread; + + mod defer; + use defer::Defer; + + pub(crate) mod inject; + pub(crate) use inject::Inject; +} + +cfg_rt_multi_thread! { + mod lock; + use lock::Lock; + + pub(crate) mod multi_thread; + pub(crate) use multi_thread::MultiThread; +} + +use crate::runtime::driver; + +#[derive(Debug, Clone)] +pub(crate) enum Handle { + #[cfg(feature = "rt")] + CurrentThread(Arc<current_thread::Handle>), + + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + MultiThread(Arc<multi_thread::Handle>), + + // TODO: This is to avoid triggering "dead code" warnings many other places + // in the codebase. Remove this during a later cleanup + #[cfg(not(feature = "rt"))] + #[allow(dead_code)] + Disabled, +} + +#[cfg(feature = "rt")] +pub(super) enum Context { + CurrentThread(current_thread::Context), + + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + MultiThread(multi_thread::Context), +} + +impl Handle { + #[cfg_attr(not(feature = "full"), allow(dead_code))] + pub(crate) fn driver(&self) -> &driver::Handle { + match *self { + #[cfg(feature = "rt")] + Handle::CurrentThread(ref h) => &h.driver, + + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + Handle::MultiThread(ref h) => &h.driver, + + #[cfg(not(feature = "rt"))] + Handle::Disabled => unreachable!(), + } + } +} + +cfg_rt! { + use crate::future::Future; + use crate::loom::sync::Arc; + use crate::runtime::{blocking, task::Id}; + use crate::runtime::context; + use crate::task::JoinHandle; + use crate::util::RngSeedGenerator; + use std::task::Waker; + + impl Handle { + #[track_caller] + pub(crate) fn current() -> Handle { + match context::with_current(Clone::clone) { + Ok(handle) => handle, + Err(e) => panic!("{}", e), + } + } + + pub(crate) fn blocking_spawner(&self) -> &blocking::Spawner { + match self { + Handle::CurrentThread(h) => &h.blocking_spawner, + + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + Handle::MultiThread(h) => &h.blocking_spawner, + } + } + + pub(crate) fn spawn<F>(&self, future: F, id: Id) -> JoinHandle<F::Output> + where + F: Future + Send + 'static, + F::Output: Send + 'static, + { + match self { + Handle::CurrentThread(h) => current_thread::Handle::spawn(h, future, id), + + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + Handle::MultiThread(h) => multi_thread::Handle::spawn(h, future, id), + } + } + + pub(crate) fn shutdown(&self) { + match *self { + Handle::CurrentThread(_) => {}, + + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + Handle::MultiThread(ref h) => h.shutdown(), + } + } + + pub(crate) fn seed_generator(&self) -> &RngSeedGenerator { + match self { + Handle::CurrentThread(h) => &h.seed_generator, + + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + Handle::MultiThread(h) => &h.seed_generator, + } + } + + pub(crate) fn as_current_thread(&self) -> &Arc<current_thread::Handle> { + match self { + Handle::CurrentThread(handle) => handle, + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + _ => panic!("not a CurrentThread handle"), + } + } + } + + cfg_metrics! { + use crate::runtime::{SchedulerMetrics, WorkerMetrics}; + + impl Handle { + pub(crate) fn num_workers(&self) -> usize { + match self { + Handle::CurrentThread(_) => 1, + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + Handle::MultiThread(handle) => handle.num_workers(), + } + } + + pub(crate) fn num_blocking_threads(&self) -> usize { + match self { + Handle::CurrentThread(handle) => handle.num_blocking_threads(), + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + Handle::MultiThread(handle) => handle.num_blocking_threads(), + } + } + + pub(crate) fn num_idle_blocking_threads(&self) -> usize { + match self { + Handle::CurrentThread(handle) => handle.num_idle_blocking_threads(), + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + Handle::MultiThread(handle) => handle.num_idle_blocking_threads(), + } + } + + pub(crate) fn active_tasks_count(&self) -> usize { + match self { + Handle::CurrentThread(handle) => handle.active_tasks_count(), + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + Handle::MultiThread(handle) => handle.active_tasks_count(), + } + } + + pub(crate) fn scheduler_metrics(&self) -> &SchedulerMetrics { + match self { + Handle::CurrentThread(handle) => handle.scheduler_metrics(), + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + Handle::MultiThread(handle) => handle.scheduler_metrics(), + } + } + + pub(crate) fn worker_metrics(&self, worker: usize) -> &WorkerMetrics { + match self { + Handle::CurrentThread(handle) => handle.worker_metrics(worker), + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + Handle::MultiThread(handle) => handle.worker_metrics(worker), + } + } + + pub(crate) fn injection_queue_depth(&self) -> usize { + match self { + Handle::CurrentThread(handle) => handle.injection_queue_depth(), + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + Handle::MultiThread(handle) => handle.injection_queue_depth(), + } + } + + pub(crate) fn worker_local_queue_depth(&self, worker: usize) -> usize { + match self { + Handle::CurrentThread(handle) => handle.worker_metrics(worker).queue_depth(), + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + Handle::MultiThread(handle) => handle.worker_local_queue_depth(worker), + } + } + + pub(crate) fn blocking_queue_depth(&self) -> usize { + match self { + Handle::CurrentThread(handle) => handle.blocking_queue_depth(), + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + Handle::MultiThread(handle) => handle.blocking_queue_depth(), + } + } + } + } + + impl Context { + #[track_caller] + pub(crate) fn expect_current_thread(&self) -> ¤t_thread::Context { + match self { + Context::CurrentThread(context) => context, + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + _ => panic!("expected `CurrentThread::Context`") + } + } + + pub(crate) fn defer(&self, waker: &Waker) { + match self { + Context::CurrentThread(context) => context.defer(waker), + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + Context::MultiThread(context) => context.defer(waker), + } + } + + cfg_rt_multi_thread! { + #[track_caller] + pub(crate) fn expect_multi_thread(&self) -> &multi_thread::Context { + match self { + Context::MultiThread(context) => context, + _ => panic!("expected `MultiThread::Context`") + } + } + } + } +} + +cfg_not_rt! { + #[cfg(any( + feature = "net", + all(unix, feature = "process"), + all(unix, feature = "signal"), + feature = "time", + ))] + impl Handle { + #[track_caller] + pub(crate) fn current() -> Handle { + panic!("{}", crate::util::error::CONTEXT_MISSING_ERROR) + } + } +} diff --git a/vendor/tokio/src/runtime/scheduler/multi_thread/counters.rs b/vendor/tokio/src/runtime/scheduler/multi_thread/counters.rs new file mode 100644 index 000000000..50bcc1198 --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/multi_thread/counters.rs @@ -0,0 +1,62 @@ +#[cfg(tokio_internal_mt_counters)] +mod imp { + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering::Relaxed; + + static NUM_MAINTENANCE: AtomicUsize = AtomicUsize::new(0); + static NUM_NOTIFY_LOCAL: AtomicUsize = AtomicUsize::new(0); + static NUM_UNPARKS_LOCAL: AtomicUsize = AtomicUsize::new(0); + static NUM_LIFO_SCHEDULES: AtomicUsize = AtomicUsize::new(0); + static NUM_LIFO_CAPPED: AtomicUsize = AtomicUsize::new(0); + + impl Drop for super::Counters { + fn drop(&mut self) { + let notifies_local = NUM_NOTIFY_LOCAL.load(Relaxed); + let unparks_local = NUM_UNPARKS_LOCAL.load(Relaxed); + let maintenance = NUM_MAINTENANCE.load(Relaxed); + let lifo_scheds = NUM_LIFO_SCHEDULES.load(Relaxed); + let lifo_capped = NUM_LIFO_CAPPED.load(Relaxed); + + println!("---"); + println!("notifies (local): {}", notifies_local); + println!(" unparks (local): {}", unparks_local); + println!(" maintenance: {}", maintenance); + println!(" LIFO schedules: {}", lifo_scheds); + println!(" LIFO capped: {}", lifo_capped); + } + } + + pub(crate) fn inc_num_inc_notify_local() { + NUM_NOTIFY_LOCAL.fetch_add(1, Relaxed); + } + + pub(crate) fn inc_num_unparks_local() { + NUM_UNPARKS_LOCAL.fetch_add(1, Relaxed); + } + + pub(crate) fn inc_num_maintenance() { + NUM_MAINTENANCE.fetch_add(1, Relaxed); + } + + pub(crate) fn inc_lifo_schedules() { + NUM_LIFO_SCHEDULES.fetch_add(1, Relaxed); + } + + pub(crate) fn inc_lifo_capped() { + NUM_LIFO_CAPPED.fetch_add(1, Relaxed); + } +} + +#[cfg(not(tokio_internal_mt_counters))] +mod imp { + pub(crate) fn inc_num_inc_notify_local() {} + pub(crate) fn inc_num_unparks_local() {} + pub(crate) fn inc_num_maintenance() {} + pub(crate) fn inc_lifo_schedules() {} + pub(crate) fn inc_lifo_capped() {} +} + +#[derive(Debug)] +pub(crate) struct Counters; + +pub(super) use imp::*; diff --git a/vendor/tokio/src/runtime/scheduler/multi_thread/handle.rs b/vendor/tokio/src/runtime/scheduler/multi_thread/handle.rs new file mode 100644 index 000000000..98e476585 --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/multi_thread/handle.rs @@ -0,0 +1,68 @@ +use crate::future::Future; +use crate::loom::sync::Arc; +use crate::runtime::scheduler::multi_thread::worker; +use crate::runtime::{ + blocking, driver, + task::{self, JoinHandle}, +}; +use crate::util::RngSeedGenerator; + +use std::fmt; + +cfg_metrics! { + mod metrics; +} + +cfg_taskdump! { + mod taskdump; +} + +/// Handle to the multi thread scheduler +pub(crate) struct Handle { + /// Task spawner + pub(super) shared: worker::Shared, + + /// Resource driver handles + pub(crate) driver: driver::Handle, + + /// Blocking pool spawner + pub(crate) blocking_spawner: blocking::Spawner, + + /// Current random number generator seed + pub(crate) seed_generator: RngSeedGenerator, +} + +impl Handle { + /// Spawns a future onto the thread pool + pub(crate) fn spawn<F>(me: &Arc<Self>, future: F, id: task::Id) -> JoinHandle<F::Output> + where + F: crate::future::Future + Send + 'static, + F::Output: Send + 'static, + { + Self::bind_new_task(me, future, id) + } + + pub(crate) fn shutdown(&self) { + self.close(); + } + + pub(super) fn bind_new_task<T>(me: &Arc<Self>, future: T, id: task::Id) -> JoinHandle<T::Output> + where + T: Future + Send + 'static, + T::Output: Send + 'static, + { + let (handle, notified) = me.shared.owned.bind(future, me.clone(), id); + + if let Some(notified) = notified { + me.schedule_task(notified, false); + } + + handle + } +} + +impl fmt::Debug for Handle { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt.debug_struct("multi_thread::Handle { ... }").finish() + } +} diff --git a/vendor/tokio/src/runtime/scheduler/multi_thread/handle/metrics.rs b/vendor/tokio/src/runtime/scheduler/multi_thread/handle/metrics.rs new file mode 100644 index 000000000..838694fc8 --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/multi_thread/handle/metrics.rs @@ -0,0 +1,41 @@ +use super::Handle; + +use crate::runtime::{SchedulerMetrics, WorkerMetrics}; + +impl Handle { + pub(crate) fn num_workers(&self) -> usize { + self.shared.worker_metrics.len() + } + + pub(crate) fn num_blocking_threads(&self) -> usize { + self.blocking_spawner.num_threads() + } + + pub(crate) fn num_idle_blocking_threads(&self) -> usize { + self.blocking_spawner.num_idle_threads() + } + + pub(crate) fn active_tasks_count(&self) -> usize { + self.shared.owned.active_tasks_count() + } + + pub(crate) fn scheduler_metrics(&self) -> &SchedulerMetrics { + &self.shared.scheduler_metrics + } + + pub(crate) fn worker_metrics(&self, worker: usize) -> &WorkerMetrics { + &self.shared.worker_metrics[worker] + } + + pub(crate) fn injection_queue_depth(&self) -> usize { + self.shared.injection_queue_depth() + } + + pub(crate) fn worker_local_queue_depth(&self, worker: usize) -> usize { + self.shared.worker_local_queue_depth(worker) + } + + pub(crate) fn blocking_queue_depth(&self) -> usize { + self.blocking_spawner.queue_depth() + } +} diff --git a/vendor/tokio/src/runtime/scheduler/multi_thread/handle/taskdump.rs b/vendor/tokio/src/runtime/scheduler/multi_thread/handle/taskdump.rs new file mode 100644 index 000000000..477d857d8 --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/multi_thread/handle/taskdump.rs @@ -0,0 +1,26 @@ +use super::Handle; + +use crate::runtime::Dump; + +impl Handle { + pub(crate) async fn dump(&self) -> Dump { + let trace_status = &self.shared.trace_status; + + // If a dump is in progress, block. + trace_status.start_trace_request(&self).await; + + let result = loop { + if let Some(result) = trace_status.take_result() { + break result; + } else { + self.notify_all(); + trace_status.result_ready.notified().await; + } + }; + + // Allow other queued dumps to proceed. + trace_status.end_trace_request(&self).await; + + result + } +} diff --git a/vendor/tokio/src/runtime/thread_pool/idle.rs b/vendor/tokio/src/runtime/scheduler/multi_thread/idle.rs index 2cac30ee8..834bc2b66 100644 --- a/vendor/tokio/src/runtime/thread_pool/idle.rs +++ b/vendor/tokio/src/runtime/scheduler/multi_thread/idle.rs @@ -1,7 +1,7 @@ //! Coordinates idling workers use crate::loom::sync::atomic::AtomicUsize; -use crate::loom::sync::Mutex; +use crate::runtime::scheduler::multi_thread::Shared; use std::fmt; use std::sync::atomic::Ordering::{self, SeqCst}; @@ -13,13 +13,16 @@ pub(super) struct Idle { /// Used as a fast-path to avoid acquiring the lock when needed. state: AtomicUsize, - /// Sleeping workers - sleepers: Mutex<Vec<usize>>, - /// Total number of workers. num_workers: usize, } +/// Data synchronized by the scheduler mutex +pub(super) struct Synced { + /// Sleeping workers + sleepers: Vec<usize>, +} + const UNPARK_SHIFT: usize = 16; const UNPARK_MASK: usize = !SEARCH_MASK; const SEARCH_MASK: usize = (1 << UNPARK_SHIFT) - 1; @@ -28,19 +31,24 @@ const SEARCH_MASK: usize = (1 << UNPARK_SHIFT) - 1; struct State(usize); impl Idle { - pub(super) fn new(num_workers: usize) -> Idle { + pub(super) fn new(num_workers: usize) -> (Idle, Synced) { let init = State::new(num_workers); - Idle { + let idle = Idle { state: AtomicUsize::new(init.into()), - sleepers: Mutex::new(Vec::with_capacity(num_workers)), num_workers, - } + }; + + let synced = Synced { + sleepers: Vec::with_capacity(num_workers), + }; + + (idle, synced) } /// If there are no workers actively searching, returns the index of a /// worker currently sleeping. - pub(super) fn worker_to_notify(&self) -> Option<usize> { + pub(super) fn worker_to_notify(&self, shared: &Shared) -> Option<usize> { // If at least one worker is spinning, work being notified will // eventually be found. A searching thread will find **some** work and // notify another worker, eventually leading to our work being found. @@ -55,7 +63,7 @@ impl Idle { } // Acquire the lock - let mut sleepers = self.sleepers.lock(); + let mut lock = shared.synced.lock(); // Check again, now that the lock is acquired if !self.notify_should_wakeup() { @@ -64,10 +72,10 @@ impl Idle { // A worker should be woken up, atomically increment the number of // searching workers as well as the number of unparked workers. - State::unpark_one(&self.state); + State::unpark_one(&self.state, 1); // Get the worker to unpark - let ret = sleepers.pop(); + let ret = lock.idle.sleepers.pop(); debug_assert!(ret.is_some()); ret @@ -75,15 +83,20 @@ impl Idle { /// Returns `true` if the worker needs to do a final check for submitted /// work. - pub(super) fn transition_worker_to_parked(&self, worker: usize, is_searching: bool) -> bool { + pub(super) fn transition_worker_to_parked( + &self, + shared: &Shared, + worker: usize, + is_searching: bool, + ) -> bool { // Acquire the lock - let mut sleepers = self.sleepers.lock(); + let mut lock = shared.synced.lock(); // Decrement the number of unparked threads let ret = State::dec_num_unparked(&self.state, is_searching); // Track the sleeping worker - sleepers.push(worker); + lock.idle.sleepers.push(worker); ret } @@ -111,25 +124,30 @@ impl Idle { /// Unpark a specific worker. This happens if tasks are submitted from /// within the worker's park routine. - pub(super) fn unpark_worker_by_id(&self, worker_id: usize) { - let mut sleepers = self.sleepers.lock(); + /// + /// Returns `true` if the worker was parked before calling the method. + pub(super) fn unpark_worker_by_id(&self, shared: &Shared, worker_id: usize) -> bool { + let mut lock = shared.synced.lock(); + let sleepers = &mut lock.idle.sleepers; for index in 0..sleepers.len() { if sleepers[index] == worker_id { sleepers.swap_remove(index); // Update the state accordingly while the lock is held. - State::unpark_one(&self.state); + State::unpark_one(&self.state, 0); - return; + return true; } } + + false } - /// Returns `true` if `worker_id` is contained in the sleep set - pub(super) fn is_parked(&self, worker_id: usize) -> bool { - let sleepers = self.sleepers.lock(); - sleepers.contains(&worker_id) + /// Returns `true` if `worker_id` is contained in the sleep set. + pub(super) fn is_parked(&self, shared: &Shared, worker_id: usize) -> bool { + let lock = shared.synced.lock(); + lock.idle.sleepers.contains(&worker_id) } fn notify_should_wakeup(&self) -> bool { @@ -151,8 +169,8 @@ impl State { State(cell.load(ordering)) } - fn unpark_one(cell: &AtomicUsize) { - cell.fetch_add(1 | (1 << UNPARK_SHIFT), SeqCst); + fn unpark_one(cell: &AtomicUsize, num_searching: usize) { + cell.fetch_add(num_searching | (1 << UNPARK_SHIFT), SeqCst); } fn inc_num_searching(cell: &AtomicUsize, ordering: Ordering) { diff --git a/vendor/tokio/src/runtime/scheduler/multi_thread/mod.rs b/vendor/tokio/src/runtime/scheduler/multi_thread/mod.rs new file mode 100644 index 000000000..d85a0ae0a --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/multi_thread/mod.rs @@ -0,0 +1,103 @@ +//! Multi-threaded runtime + +mod counters; +use counters::Counters; + +mod handle; +pub(crate) use handle::Handle; + +mod overflow; +pub(crate) use overflow::Overflow; + +mod idle; +use self::idle::Idle; + +mod stats; +pub(crate) use stats::Stats; + +mod park; +pub(crate) use park::{Parker, Unparker}; + +pub(crate) mod queue; + +mod worker; +pub(crate) use worker::{Context, Launch, Shared}; + +cfg_taskdump! { + mod trace; + use trace::TraceStatus; + + pub(crate) use worker::Synced; +} + +cfg_not_taskdump! { + mod trace_mock; + use trace_mock::TraceStatus; +} + +pub(crate) use worker::block_in_place; + +use crate::loom::sync::Arc; +use crate::runtime::{ + blocking, + driver::{self, Driver}, + scheduler, Config, +}; +use crate::util::RngSeedGenerator; + +use std::fmt; +use std::future::Future; + +/// Work-stealing based thread pool for executing futures. +pub(crate) struct MultiThread; + +// ===== impl MultiThread ===== + +impl MultiThread { + pub(crate) fn new( + size: usize, + driver: Driver, + driver_handle: driver::Handle, + blocking_spawner: blocking::Spawner, + seed_generator: RngSeedGenerator, + config: Config, + ) -> (MultiThread, Arc<Handle>, Launch) { + let parker = Parker::new(driver); + let (handle, launch) = worker::create( + size, + parker, + driver_handle, + blocking_spawner, + seed_generator, + config, + ); + + (MultiThread, handle, launch) + } + + /// Blocks the current thread waiting for the future to complete. + /// + /// The future will execute on the current thread, but all spawned tasks + /// will be executed on the thread pool. + pub(crate) fn block_on<F>(&self, handle: &scheduler::Handle, future: F) -> F::Output + where + F: Future, + { + crate::runtime::context::enter_runtime(handle, true, |blocking| { + blocking.block_on(future).expect("failed to park thread") + }) + } + + pub(crate) fn shutdown(&mut self, handle: &scheduler::Handle) { + match handle { + scheduler::Handle::MultiThread(handle) => handle.shutdown(), + _ => panic!("expected MultiThread scheduler"), + } + } +} + +impl fmt::Debug for MultiThread { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt.debug_struct("MultiThread").finish() + } +} diff --git a/vendor/tokio/src/runtime/scheduler/multi_thread/overflow.rs b/vendor/tokio/src/runtime/scheduler/multi_thread/overflow.rs new file mode 100644 index 000000000..ab664811c --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/multi_thread/overflow.rs @@ -0,0 +1,26 @@ +use crate::runtime::task; + +#[cfg(test)] +use std::cell::RefCell; + +pub(crate) trait Overflow<T: 'static> { + fn push(&self, task: task::Notified<T>); + + fn push_batch<I>(&self, iter: I) + where + I: Iterator<Item = task::Notified<T>>; +} + +#[cfg(test)] +impl<T: 'static> Overflow<T> for RefCell<Vec<task::Notified<T>>> { + fn push(&self, task: task::Notified<T>) { + self.borrow_mut().push(task); + } + + fn push_batch<I>(&self, iter: I) + where + I: Iterator<Item = task::Notified<T>>, + { + self.borrow_mut().extend(iter); + } +} diff --git a/vendor/tokio/src/runtime/scheduler/multi_thread/park.rs b/vendor/tokio/src/runtime/scheduler/multi_thread/park.rs new file mode 100644 index 000000000..0a00ea004 --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/multi_thread/park.rs @@ -0,0 +1,232 @@ +//! Parks the runtime. +//! +//! A combination of the various resource driver park handles. + +use crate::loom::sync::atomic::AtomicUsize; +use crate::loom::sync::{Arc, Condvar, Mutex}; +use crate::runtime::driver::{self, Driver}; +use crate::util::TryLock; + +use std::sync::atomic::Ordering::SeqCst; +use std::time::Duration; + +pub(crate) struct Parker { + inner: Arc<Inner>, +} + +pub(crate) struct Unparker { + inner: Arc<Inner>, +} + +struct Inner { + /// Avoids entering the park if possible + state: AtomicUsize, + + /// Used to coordinate access to the driver / condvar + mutex: Mutex<()>, + + /// Condvar to block on if the driver is unavailable. + condvar: Condvar, + + /// Resource (I/O, time, ...) driver + shared: Arc<Shared>, +} + +const EMPTY: usize = 0; +const PARKED_CONDVAR: usize = 1; +const PARKED_DRIVER: usize = 2; +const NOTIFIED: usize = 3; + +/// Shared across multiple Parker handles +struct Shared { + /// Shared driver. Only one thread at a time can use this + driver: TryLock<Driver>, +} + +impl Parker { + pub(crate) fn new(driver: Driver) -> Parker { + Parker { + inner: Arc::new(Inner { + state: AtomicUsize::new(EMPTY), + mutex: Mutex::new(()), + condvar: Condvar::new(), + shared: Arc::new(Shared { + driver: TryLock::new(driver), + }), + }), + } + } + + pub(crate) fn unpark(&self) -> Unparker { + Unparker { + inner: self.inner.clone(), + } + } + + pub(crate) fn park(&mut self, handle: &driver::Handle) { + self.inner.park(handle); + } + + pub(crate) fn park_timeout(&mut self, handle: &driver::Handle, duration: Duration) { + // Only parking with zero is supported... + assert_eq!(duration, Duration::from_millis(0)); + + if let Some(mut driver) = self.inner.shared.driver.try_lock() { + driver.park_timeout(handle, duration) + } + } + + pub(crate) fn shutdown(&mut self, handle: &driver::Handle) { + self.inner.shutdown(handle); + } +} + +impl Clone for Parker { + fn clone(&self) -> Parker { + Parker { + inner: Arc::new(Inner { + state: AtomicUsize::new(EMPTY), + mutex: Mutex::new(()), + condvar: Condvar::new(), + shared: self.inner.shared.clone(), + }), + } + } +} + +impl Unparker { + pub(crate) fn unpark(&self, driver: &driver::Handle) { + self.inner.unpark(driver); + } +} + +impl Inner { + /// Parks the current thread for at most `dur`. + fn park(&self, handle: &driver::Handle) { + // If we were previously notified then we consume this notification and + // return quickly. + if self + .state + .compare_exchange(NOTIFIED, EMPTY, SeqCst, SeqCst) + .is_ok() + { + return; + } + + if let Some(mut driver) = self.shared.driver.try_lock() { + self.park_driver(&mut driver, handle); + } else { + self.park_condvar(); + } + } + + fn park_condvar(&self) { + // Otherwise we need to coordinate going to sleep + let mut m = self.mutex.lock(); + + match self + .state + .compare_exchange(EMPTY, PARKED_CONDVAR, SeqCst, SeqCst) + { + Ok(_) => {} + Err(NOTIFIED) => { + // We must read here, even though we know it will be `NOTIFIED`. + // This is because `unpark` may have been called again since we read + // `NOTIFIED` in the `compare_exchange` above. We must perform an + // acquire operation that synchronizes with that `unpark` to observe + // any writes it made before the call to unpark. To do that we must + // read from the write it made to `state`. + let old = self.state.swap(EMPTY, SeqCst); + debug_assert_eq!(old, NOTIFIED, "park state changed unexpectedly"); + + return; + } + Err(actual) => panic!("inconsistent park state; actual = {}", actual), + } + + loop { + m = self.condvar.wait(m).unwrap(); + + if self + .state + .compare_exchange(NOTIFIED, EMPTY, SeqCst, SeqCst) + .is_ok() + { + // got a notification + return; + } + + // spurious wakeup, go back to sleep + } + } + + fn park_driver(&self, driver: &mut Driver, handle: &driver::Handle) { + match self + .state + .compare_exchange(EMPTY, PARKED_DRIVER, SeqCst, SeqCst) + { + Ok(_) => {} + Err(NOTIFIED) => { + // We must read here, even though we know it will be `NOTIFIED`. + // This is because `unpark` may have been called again since we read + // `NOTIFIED` in the `compare_exchange` above. We must perform an + // acquire operation that synchronizes with that `unpark` to observe + // any writes it made before the call to unpark. To do that we must + // read from the write it made to `state`. + let old = self.state.swap(EMPTY, SeqCst); + debug_assert_eq!(old, NOTIFIED, "park state changed unexpectedly"); + + return; + } + Err(actual) => panic!("inconsistent park state; actual = {}", actual), + } + + driver.park(handle); + + match self.state.swap(EMPTY, SeqCst) { + NOTIFIED => {} // got a notification, hurray! + PARKED_DRIVER => {} // no notification, alas + n => panic!("inconsistent park_timeout state: {}", n), + } + } + + fn unpark(&self, driver: &driver::Handle) { + // To ensure the unparked thread will observe any writes we made before + // this call, we must perform a release operation that `park` can + // synchronize with. To do that we must write `NOTIFIED` even if `state` + // is already `NOTIFIED`. That is why this must be a swap rather than a + // compare-and-swap that returns if it reads `NOTIFIED` on failure. + match self.state.swap(NOTIFIED, SeqCst) { + EMPTY => {} // no one was waiting + NOTIFIED => {} // already unparked + PARKED_CONDVAR => self.unpark_condvar(), + PARKED_DRIVER => driver.unpark(), + actual => panic!("inconsistent state in unpark; actual = {}", actual), + } + } + + fn unpark_condvar(&self) { + // There is a period between when the parked thread sets `state` to + // `PARKED` (or last checked `state` in the case of a spurious wake + // up) and when it actually waits on `cvar`. If we were to notify + // during this period it would be ignored and then when the parked + // thread went to sleep it would never wake up. Fortunately, it has + // `lock` locked at this stage so we can acquire `lock` to wait until + // it is ready to receive the notification. + // + // Releasing `lock` before the call to `notify_one` means that when the + // parked thread wakes it doesn't get woken only to have to wait for us + // to release `lock`. + drop(self.mutex.lock()); + + self.condvar.notify_one() + } + + fn shutdown(&self, handle: &driver::Handle) { + if let Some(mut driver) = self.shared.driver.try_lock() { + driver.shutdown(handle); + } + + self.condvar.notify_all(); + } +} diff --git a/vendor/tokio/src/runtime/queue.rs b/vendor/tokio/src/runtime/scheduler/multi_thread/queue.rs index 818ef7bac..dd66fa2dd 100644 --- a/vendor/tokio/src/runtime/queue.rs +++ b/vendor/tokio/src/runtime/scheduler/multi_thread/queue.rs @@ -1,72 +1,63 @@ //! Run-queue structures to support a work-stealing scheduler use crate::loom::cell::UnsafeCell; -use crate::loom::sync::atomic::{AtomicU16, AtomicU32, AtomicUsize}; -use crate::loom::sync::{Arc, Mutex}; +use crate::loom::sync::Arc; +use crate::runtime::scheduler::multi_thread::{Overflow, Stats}; use crate::runtime::task; -use std::marker::PhantomData; -use std::mem::MaybeUninit; -use std::ptr::{self, NonNull}; +use std::mem::{self, MaybeUninit}; +use std::ptr; use std::sync::atomic::Ordering::{AcqRel, Acquire, Relaxed, Release}; +// Use wider integers when possible to increase ABA resilience. +// +// See issue #5041: <https://github.com/tokio-rs/tokio/issues/5041>. +cfg_has_atomic_u64! { + type UnsignedShort = u32; + type UnsignedLong = u64; + type AtomicUnsignedShort = crate::loom::sync::atomic::AtomicU32; + type AtomicUnsignedLong = crate::loom::sync::atomic::AtomicU64; +} +cfg_not_has_atomic_u64! { + type UnsignedShort = u16; + type UnsignedLong = u32; + type AtomicUnsignedShort = crate::loom::sync::atomic::AtomicU16; + type AtomicUnsignedLong = crate::loom::sync::atomic::AtomicU32; +} + /// Producer handle. May only be used from a single thread. -pub(super) struct Local<T: 'static> { +pub(crate) struct Local<T: 'static> { inner: Arc<Inner<T>>, } /// Consumer handle. May be used from many threads. -pub(super) struct Steal<T: 'static>(Arc<Inner<T>>); - -/// Growable, MPMC queue used to inject new tasks into the scheduler and as an -/// overflow queue when the local, fixed-size, array queue overflows. -pub(super) struct Inject<T: 'static> { - /// Pointers to the head and tail of the queue - pointers: Mutex<Pointers>, - - /// Number of pending tasks in the queue. This helps prevent unnecessary - /// locking in the hot path. - len: AtomicUsize, +pub(crate) struct Steal<T: 'static>(Arc<Inner<T>>); - _p: PhantomData<T>, -} - -pub(super) struct Inner<T: 'static> { +pub(crate) struct Inner<T: 'static> { /// Concurrently updated by many threads. /// - /// Contains two `u16` values. The LSB byte is the "real" head of the queue. - /// The `u16` in the MSB is set by a stealer in process of stealing values. - /// It represents the first value being stolen in the batch. `u16` is used - /// in order to distinguish between `head == tail` and `head == tail - - /// capacity`. + /// Contains two `UnsignedShort` values. The LSB byte is the "real" head of + /// the queue. The `UnsignedShort` in the MSB is set by a stealer in process + /// of stealing values. It represents the first value being stolen in the + /// batch. The `UnsignedShort` indices are intentionally wider than strictly + /// required for buffer indexing in order to provide ABA mitigation and make + /// it possible to distinguish between full and empty buffers. /// - /// When both `u16` values are the same, there is no active stealer. + /// When both `UnsignedShort` values are the same, there is no active + /// stealer. /// /// Tracking an in-progress stealer prevents a wrapping scenario. - head: AtomicU32, + head: AtomicUnsignedLong, /// Only updated by producer thread but read by many threads. - tail: AtomicU16, + tail: AtomicUnsignedShort, /// Elements - buffer: Box<[UnsafeCell<MaybeUninit<task::Notified<T>>>]>, -} - -struct Pointers { - /// True if the queue is closed - is_closed: bool, - - /// Linked-list head - head: Option<NonNull<task::Header>>, - - /// Linked-list tail - tail: Option<NonNull<task::Header>>, + buffer: Box<[UnsafeCell<MaybeUninit<task::Notified<T>>>; LOCAL_QUEUE_CAPACITY]>, } unsafe impl<T> Send for Inner<T> {} unsafe impl<T> Sync for Inner<T> {} -unsafe impl<T> Send for Inject<T> {} -unsafe impl<T> Sync for Inject<T> {} #[cfg(not(loom))] const LOCAL_QUEUE_CAPACITY: usize = 256; @@ -79,8 +70,19 @@ const LOCAL_QUEUE_CAPACITY: usize = 4; const MASK: usize = LOCAL_QUEUE_CAPACITY - 1; +// Constructing the fixed size array directly is very awkward. The only way to +// do it is to repeat `UnsafeCell::new(MaybeUninit::uninit())` 256 times, as +// the contents are not Copy. The trick with defining a const doesn't work for +// generic types. +fn make_fixed_size<T>(buffer: Box<[T]>) -> Box<[T; LOCAL_QUEUE_CAPACITY]> { + assert_eq!(buffer.len(), LOCAL_QUEUE_CAPACITY); + + // safety: We check that the length is correct. + unsafe { Box::from_raw(Box::into_raw(buffer).cast()) } +} + /// Create a new local run-queue -pub(super) fn local<T: 'static>() -> (Steal<T>, Local<T>) { +pub(crate) fn local<T: 'static>() -> (Steal<T>, Local<T>) { let mut buffer = Vec::with_capacity(LOCAL_QUEUE_CAPACITY); for _ in 0..LOCAL_QUEUE_CAPACITY { @@ -88,9 +90,9 @@ pub(super) fn local<T: 'static>() -> (Steal<T>, Local<T>) { } let inner = Arc::new(Inner { - head: AtomicU32::new(0), - tail: AtomicU16::new(0), - buffer: buffer.into(), + head: AtomicUnsignedLong::new(0), + tail: AtomicUnsignedShort::new(0), + buffer: make_fixed_size(buffer.into_boxed_slice()), }); let local = Local { @@ -103,16 +105,89 @@ pub(super) fn local<T: 'static>() -> (Steal<T>, Local<T>) { } impl<T> Local<T> { - /// Returns true if the queue has entries that can be stealed. - pub(super) fn is_stealable(&self) -> bool { + /// Returns the number of entries in the queue + pub(crate) fn len(&self) -> usize { + self.inner.len() as usize + } + + /// How many tasks can be pushed into the queue + pub(crate) fn remaining_slots(&self) -> usize { + self.inner.remaining_slots() + } + + pub(crate) fn max_capacity(&self) -> usize { + LOCAL_QUEUE_CAPACITY + } + + /// Returns false if there are any entries in the queue + /// + /// Separate to is_stealable so that refactors of is_stealable to "protect" + /// some tasks from stealing won't affect this + pub(crate) fn has_tasks(&self) -> bool { !self.inner.is_empty() } - /// Pushes a task to the back of the local queue, skipping the LIFO slot. - pub(super) fn push_back(&mut self, mut task: task::Notified<T>, inject: &Inject<T>) - where - T: crate::runtime::task::Schedule, - { + /// Pushes a batch of tasks to the back of the queue. All tasks must fit in + /// the local queue. + /// + /// # Panics + /// + /// The method panics if there is not enough capacity to fit in the queue. + pub(crate) fn push_back(&mut self, tasks: impl ExactSizeIterator<Item = task::Notified<T>>) { + let len = tasks.len(); + assert!(len <= LOCAL_QUEUE_CAPACITY); + + if len == 0 { + // Nothing to do + return; + } + + let head = self.inner.head.load(Acquire); + let (steal, _) = unpack(head); + + // safety: this is the **only** thread that updates this cell. + let mut tail = unsafe { self.inner.tail.unsync_load() }; + + if tail.wrapping_sub(steal) <= (LOCAL_QUEUE_CAPACITY - len) as UnsignedShort { + // Yes, this if condition is structured a bit weird (first block + // does nothing, second returns an error). It is this way to match + // `push_back_or_overflow`. + } else { + panic!() + } + + for task in tasks { + let idx = tail as usize & MASK; + + self.inner.buffer[idx].with_mut(|ptr| { + // Write the task to the slot + // + // Safety: There is only one producer and the above `if` + // condition ensures we don't touch a cell if there is a + // value, thus no consumer. + unsafe { + ptr::write((*ptr).as_mut_ptr(), task); + } + }); + + tail = tail.wrapping_add(1); + } + + self.inner.tail.store(tail, Release); + } + + /// Pushes a task to the back of the local queue, if there is not enough + /// capacity in the queue, this triggers the overflow operation. + /// + /// When the queue overflows, half of the curent contents of the queue is + /// moved to the given Injection queue. This frees up capacity for more + /// tasks to be pushed into the local queue. + pub(crate) fn push_back_or_overflow<O: Overflow<T>>( + &mut self, + mut task: task::Notified<T>, + overflow: &O, + stats: &mut Stats, + ) { let tail = loop { let head = self.inner.head.load(Acquire); let (steal, real) = unpack(head); @@ -120,23 +195,18 @@ impl<T> Local<T> { // safety: this is the **only** thread that updates this cell. let tail = unsafe { self.inner.tail.unsync_load() }; - if tail.wrapping_sub(steal) < LOCAL_QUEUE_CAPACITY as u16 { + if tail.wrapping_sub(steal) < LOCAL_QUEUE_CAPACITY as UnsignedShort { // There is capacity for the task break tail; } else if steal != real { // Concurrently stealing, this will free up capacity, so only - // push the new task onto the inject queue - // - // If the task fails to be pushed on the injection queue, there - // is nothing to be done at this point as the task cannot be a - // newly spawned task. Shutting down this task is handled by the - // worker shutdown process. - let _ = inject.push(task); + // push the task onto the inject queue + overflow.push(task); return; } else { // Push the current task and half of the queue into the // inject queue. - match self.push_overflow(task, real, tail, inject) { + match self.push_overflow(task, real, tail, overflow, stats) { Ok(_) => return, // Lost the race, try again Err(v) => { @@ -146,6 +216,11 @@ impl<T> Local<T> { } }; + self.push_back_finish(task, tail); + } + + // Second half of `push_back` + fn push_back_finish(&self, task: task::Notified<T>, tail: UnsignedShort) { // Map the position to a slot index. let idx = tail as usize & MASK; @@ -172,16 +247,20 @@ impl<T> Local<T> { /// workers "missed" some of the tasks during a steal, they will get /// another opportunity. #[inline(never)] - fn push_overflow( + fn push_overflow<O: Overflow<T>>( &mut self, task: task::Notified<T>, - head: u16, - tail: u16, - inject: &Inject<T>, + head: UnsignedShort, + tail: UnsignedShort, + overflow: &O, + stats: &mut Stats, ) -> Result<(), task::Notified<T>> { - const BATCH_LEN: usize = LOCAL_QUEUE_CAPACITY / 2 + 1; + /// How many elements are we taking from the local queue. + /// + /// This is one less than the number of tasks pushed to the inject + /// queue as we are also inserting the `task` argument. + const NUM_TASKS_TAKEN: UnsignedShort = (LOCAL_QUEUE_CAPACITY / 2) as UnsignedShort; - let n = (LOCAL_QUEUE_CAPACITY / 2) as u16; assert_eq!( tail.wrapping_sub(head) as usize, LOCAL_QUEUE_CAPACITY, @@ -207,7 +286,10 @@ impl<T> Local<T> { .head .compare_exchange( prev, - pack(head.wrapping_add(n), head.wrapping_add(n)), + pack( + head.wrapping_add(NUM_TASKS_TAKEN), + head.wrapping_add(NUM_TASKS_TAKEN), + ), Release, Relaxed, ) @@ -219,47 +301,50 @@ impl<T> Local<T> { return Err(task); } - // link the tasks - for i in 0..n { - let j = i + 1; - - let i_idx = i.wrapping_add(head) as usize & MASK; - let j_idx = j.wrapping_add(head) as usize & MASK; - - // Get the next pointer - let next = if j == n { - // The last task in the local queue being moved - task.header().into() - } else { - // safety: The above CAS prevents a stealer from accessing these - // tasks and we are the only producer. - self.inner.buffer[j_idx].with(|ptr| unsafe { - let value = (*ptr).as_ptr(); - (*value).header().into() - }) - }; - - // safety: the above CAS prevents a stealer from accessing these - // tasks and we are the only producer. - self.inner.buffer[i_idx].with_mut(|ptr| unsafe { - let ptr = (*ptr).as_ptr(); - (*ptr).header().set_next(Some(next)) - }); + /// An iterator that takes elements out of the run queue. + struct BatchTaskIter<'a, T: 'static> { + buffer: &'a [UnsafeCell<MaybeUninit<task::Notified<T>>>; LOCAL_QUEUE_CAPACITY], + head: UnsignedLong, + i: UnsignedLong, + } + impl<'a, T: 'static> Iterator for BatchTaskIter<'a, T> { + type Item = task::Notified<T>; + + #[inline] + fn next(&mut self) -> Option<task::Notified<T>> { + if self.i == UnsignedLong::from(NUM_TASKS_TAKEN) { + None + } else { + let i_idx = self.i.wrapping_add(self.head) as usize & MASK; + let slot = &self.buffer[i_idx]; + + // safety: Our CAS from before has assumed exclusive ownership + // of the task pointers in this range. + let task = slot.with(|ptr| unsafe { ptr::read((*ptr).as_ptr()) }); + + self.i += 1; + Some(task) + } + } } - // safety: the above CAS prevents a stealer from accessing these tasks - // and we are the only producer. - let head = self.inner.buffer[head as usize & MASK] - .with(|ptr| unsafe { ptr::read((*ptr).as_ptr()) }); + // safety: The CAS above ensures that no consumer will look at these + // values again, and we are the only producer. + let batch_iter = BatchTaskIter { + buffer: &self.inner.buffer, + head: head as UnsignedLong, + i: 0, + }; + overflow.push_batch(batch_iter.chain(std::iter::once(task))); - // Push the tasks onto the inject queue - inject.push_batch(head, task, BATCH_LEN); + // Add 1 to factor in the task currently being scheduled. + stats.incr_overflow_count(); Ok(()) } /// Pops a task from the local queue. - pub(super) fn pop(&mut self) -> Option<task::Notified<T>> { + pub(crate) fn pop(&mut self) -> Option<task::Notified<T>> { let mut head = self.inner.head.load(Acquire); let idx = loop { @@ -301,12 +386,16 @@ impl<T> Local<T> { } impl<T> Steal<T> { - pub(super) fn is_empty(&self) -> bool { + pub(crate) fn is_empty(&self) -> bool { self.0.is_empty() } /// Steals half the tasks from self and place them into `dst`. - pub(super) fn steal_into(&self, dst: &mut Local<T>) -> Option<task::Notified<T>> { + pub(crate) fn steal_into( + &self, + dst: &mut Local<T>, + dst_stats: &mut Stats, + ) -> Option<task::Notified<T>> { // Safety: the caller is the only thread that mutates `dst.tail` and // holds a mutable reference. let dst_tail = unsafe { dst.inner.tail.unsync_load() }; @@ -316,7 +405,7 @@ impl<T> Steal<T> { // from `dst` there may not be enough capacity to steal. let (steal, _) = unpack(dst.inner.head.load(Acquire)); - if dst_tail.wrapping_sub(steal) > LOCAL_QUEUE_CAPACITY as u16 / 2 { + if dst_tail.wrapping_sub(steal) > LOCAL_QUEUE_CAPACITY as UnsignedShort / 2 { // we *could* try to steal less here, but for simplicity, we're just // going to abort. return None; @@ -331,6 +420,9 @@ impl<T> Steal<T> { return None; } + dst_stats.incr_steal_count(n as u16); + dst_stats.incr_steal_operations(); + // We are returning a task here n -= 1; @@ -354,7 +446,7 @@ impl<T> Steal<T> { // Steal tasks from `self`, placing them into `dst`. Returns the number of // tasks that were stolen. - fn steal_into2(&self, dst: &mut Local<T>, dst_tail: u16) -> u16 { + fn steal_into2(&self, dst: &mut Local<T>, dst_tail: UnsignedShort) -> UnsignedShort { let mut prev_packed = self.0.head.load(Acquire); let mut next_packed; @@ -396,7 +488,11 @@ impl<T> Steal<T> { } }; - assert!(n <= LOCAL_QUEUE_CAPACITY as u16 / 2, "actual = {}", n); + assert!( + n <= LOCAL_QUEUE_CAPACITY as UnsignedShort / 2, + "actual = {}", + n + ); let (first, _) = unpack(next_packed); @@ -450,6 +546,14 @@ impl<T> Steal<T> { } } +cfg_metrics! { + impl<T> Steal<T> { + pub(crate) fn len(&self) -> usize { + self.0.len() as _ + } + } +} + impl<T> Clone for Steal<T> { fn clone(&self) -> Steal<T> { Steal(self.0.clone()) @@ -465,179 +569,37 @@ impl<T> Drop for Local<T> { } impl<T> Inner<T> { - fn is_empty(&self) -> bool { - let (_, head) = unpack(self.head.load(Acquire)); + fn remaining_slots(&self) -> usize { + let (steal, _) = unpack(self.head.load(Acquire)); let tail = self.tail.load(Acquire); - head == tail + LOCAL_QUEUE_CAPACITY - (tail.wrapping_sub(steal) as usize) } -} -impl<T: 'static> Inject<T> { - pub(super) fn new() -> Inject<T> { - Inject { - pointers: Mutex::new(Pointers { - is_closed: false, - head: None, - tail: None, - }), - len: AtomicUsize::new(0), - _p: PhantomData, - } - } - - pub(super) fn is_empty(&self) -> bool { - self.len() == 0 - } - - /// Close the injection queue, returns `true` if the queue is open when the - /// transition is made. - pub(super) fn close(&self) -> bool { - let mut p = self.pointers.lock(); - - if p.is_closed { - return false; - } - - p.is_closed = true; - true - } - - pub(super) fn is_closed(&self) -> bool { - self.pointers.lock().is_closed - } - - pub(super) fn len(&self) -> usize { - self.len.load(Acquire) - } - - /// Pushes a value into the queue. - /// - /// Returns `Err(task)` if pushing fails due to the queue being shutdown. - /// The caller is expected to call `shutdown()` on the task **if and only - /// if** it is a newly spawned task. - pub(super) fn push(&self, task: task::Notified<T>) -> Result<(), task::Notified<T>> - where - T: crate::runtime::task::Schedule, - { - // Acquire queue lock - let mut p = self.pointers.lock(); - - if p.is_closed { - return Err(task); - } - - // safety: only mutated with the lock held - let len = unsafe { self.len.unsync_load() }; - let task = task.into_raw(); - - // The next pointer should already be null - debug_assert!(get_next(task).is_none()); - - if let Some(tail) = p.tail { - set_next(tail, Some(task)); - } else { - p.head = Some(task); - } - - p.tail = Some(task); - - self.len.store(len + 1, Release); - Ok(()) - } - - pub(super) fn push_batch( - &self, - batch_head: task::Notified<T>, - batch_tail: task::Notified<T>, - num: usize, - ) { - let batch_head = batch_head.into_raw(); - let batch_tail = batch_tail.into_raw(); - - debug_assert!(get_next(batch_tail).is_none()); - - let mut p = self.pointers.lock(); - - if let Some(tail) = p.tail { - set_next(tail, Some(batch_head)); - } else { - p.head = Some(batch_head); - } - - p.tail = Some(batch_tail); - - // Increment the count. - // - // safety: All updates to the len atomic are guarded by the mutex. As - // such, a non-atomic load followed by a store is safe. - let len = unsafe { self.len.unsync_load() }; - - self.len.store(len + num, Release); - } - - pub(super) fn pop(&self) -> Option<task::Notified<T>> { - // Fast path, if len == 0, then there are no values - if self.is_empty() { - return None; - } - - let mut p = self.pointers.lock(); - - // It is possible to hit null here if another thread popped the last - // task between us checking `len` and acquiring the lock. - let task = p.head?; - - p.head = get_next(task); - - if p.head.is_none() { - p.tail = None; - } - - set_next(task, None); - - // Decrement the count. - // - // safety: All updates to the len atomic are guarded by the mutex. As - // such, a non-atomic load followed by a store is safe. - self.len - .store(unsafe { self.len.unsync_load() } - 1, Release); - - // safety: a `Notified` is pushed into the queue and now it is popped! - Some(unsafe { task::Notified::from_raw(task) }) - } -} + fn len(&self) -> UnsignedShort { + let (_, head) = unpack(self.head.load(Acquire)); + let tail = self.tail.load(Acquire); -impl<T: 'static> Drop for Inject<T> { - fn drop(&mut self) { - if !std::thread::panicking() { - assert!(self.pop().is_none(), "queue not empty"); - } + tail.wrapping_sub(head) } -} -fn get_next(header: NonNull<task::Header>) -> Option<NonNull<task::Header>> { - unsafe { header.as_ref().queue_next.with(|ptr| *ptr) } -} - -fn set_next(header: NonNull<task::Header>, val: Option<NonNull<task::Header>>) { - unsafe { - header.as_ref().set_next(val); + fn is_empty(&self) -> bool { + self.len() == 0 } } /// Split the head value into the real head and the index a stealer is working /// on. -fn unpack(n: u32) -> (u16, u16) { - let real = n & u16::MAX as u32; - let steal = n >> 16; +fn unpack(n: UnsignedLong) -> (UnsignedShort, UnsignedShort) { + let real = n & UnsignedShort::MAX as UnsignedLong; + let steal = n >> (mem::size_of::<UnsignedShort>() * 8); - (steal as u16, real as u16) + (steal as UnsignedShort, real as UnsignedShort) } /// Join the two head values -fn pack(steal: u16, real: u16) -> u32 { - (real as u32) | ((steal as u32) << 16) +fn pack(steal: UnsignedShort, real: UnsignedShort) -> UnsignedLong { + (real as UnsignedLong) | ((steal as UnsignedLong) << (mem::size_of::<UnsignedShort>() * 8)) } #[test] diff --git a/vendor/tokio/src/runtime/scheduler/multi_thread/stats.rs b/vendor/tokio/src/runtime/scheduler/multi_thread/stats.rs new file mode 100644 index 000000000..f01daaa1b --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/multi_thread/stats.rs @@ -0,0 +1,140 @@ +use crate::runtime::{Config, MetricsBatch, WorkerMetrics}; + +use std::cmp; +use std::time::{Duration, Instant}; + +/// Per-worker statistics. This is used for both tuning the scheduler and +/// reporting runtime-level metrics/stats. +pub(crate) struct Stats { + /// The metrics batch used to report runtime-level metrics/stats to the + /// user. + batch: MetricsBatch, + + /// Instant at which work last resumed (continued after park). + /// + /// This duplicates the value stored in `MetricsBatch`. We will unify + /// `Stats` and `MetricsBatch` when we stabilize metrics. + processing_scheduled_tasks_started_at: Instant, + + /// Number of tasks polled in the batch of scheduled tasks + tasks_polled_in_batch: usize, + + /// Exponentially-weighted moving average of time spent polling scheduled a + /// task. + /// + /// Tracked in nanoseconds, stored as a f64 since that is what we use with + /// the EWMA calculations + task_poll_time_ewma: f64, +} + +/// How to weigh each individual poll time, value is plucked from thin air. +const TASK_POLL_TIME_EWMA_ALPHA: f64 = 0.1; + +/// Ideally, we wouldn't go above this, value is plucked from thin air. +const TARGET_GLOBAL_QUEUE_INTERVAL: f64 = Duration::from_micros(200).as_nanos() as f64; + +/// Max value for the global queue interval. This is 2x the previous default +const MAX_TASKS_POLLED_PER_GLOBAL_QUEUE_INTERVAL: u32 = 127; + +/// This is the previous default +const TARGET_TASKS_POLLED_PER_GLOBAL_QUEUE_INTERVAL: u32 = 61; + +impl Stats { + pub(crate) fn new(worker_metrics: &WorkerMetrics) -> Stats { + // Seed the value with what we hope to see. + let task_poll_time_ewma = + TARGET_GLOBAL_QUEUE_INTERVAL / TARGET_TASKS_POLLED_PER_GLOBAL_QUEUE_INTERVAL as f64; + + Stats { + batch: MetricsBatch::new(worker_metrics), + processing_scheduled_tasks_started_at: Instant::now(), + tasks_polled_in_batch: 0, + task_poll_time_ewma, + } + } + + pub(crate) fn tuned_global_queue_interval(&self, config: &Config) -> u32 { + // If an interval is explicitly set, don't tune. + if let Some(configured) = config.global_queue_interval { + return configured; + } + + // As of Rust 1.45, casts from f64 -> u32 are saturating, which is fine here. + let tasks_per_interval = (TARGET_GLOBAL_QUEUE_INTERVAL / self.task_poll_time_ewma) as u32; + + cmp::max( + // We don't want to return less than 2 as that would result in the + // global queue always getting checked first. + 2, + cmp::min( + MAX_TASKS_POLLED_PER_GLOBAL_QUEUE_INTERVAL, + tasks_per_interval, + ), + ) + } + + pub(crate) fn submit(&mut self, to: &WorkerMetrics) { + self.batch.submit(to); + } + + pub(crate) fn about_to_park(&mut self) { + self.batch.about_to_park(); + } + + pub(crate) fn inc_local_schedule_count(&mut self) { + self.batch.inc_local_schedule_count(); + } + + pub(crate) fn start_processing_scheduled_tasks(&mut self) { + self.batch.start_processing_scheduled_tasks(); + + self.processing_scheduled_tasks_started_at = Instant::now(); + self.tasks_polled_in_batch = 0; + } + + pub(crate) fn end_processing_scheduled_tasks(&mut self) { + self.batch.end_processing_scheduled_tasks(); + + // Update the EWMA task poll time + if self.tasks_polled_in_batch > 0 { + let now = Instant::now(); + + // If we "overflow" this conversion, we have bigger problems than + // slightly off stats. + let elapsed = (now - self.processing_scheduled_tasks_started_at).as_nanos() as f64; + let num_polls = self.tasks_polled_in_batch as f64; + + // Calculate the mean poll duration for a single task in the batch + let mean_poll_duration = elapsed / num_polls; + + // Compute the alpha weighted by the number of tasks polled this batch. + let weighted_alpha = 1.0 - (1.0 - TASK_POLL_TIME_EWMA_ALPHA).powf(num_polls); + + // Now compute the new weighted average task poll time. + self.task_poll_time_ewma = weighted_alpha * mean_poll_duration + + (1.0 - weighted_alpha) * self.task_poll_time_ewma; + } + } + + pub(crate) fn start_poll(&mut self) { + self.batch.start_poll(); + + self.tasks_polled_in_batch += 1; + } + + pub(crate) fn end_poll(&mut self) { + self.batch.end_poll(); + } + + pub(crate) fn incr_steal_count(&mut self, by: u16) { + self.batch.incr_steal_count(by); + } + + pub(crate) fn incr_steal_operations(&mut self) { + self.batch.incr_steal_operations(); + } + + pub(crate) fn incr_overflow_count(&mut self) { + self.batch.incr_overflow_count(); + } +} diff --git a/vendor/tokio/src/runtime/scheduler/multi_thread/trace.rs b/vendor/tokio/src/runtime/scheduler/multi_thread/trace.rs new file mode 100644 index 000000000..7b4aeb5c1 --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/multi_thread/trace.rs @@ -0,0 +1,61 @@ +use crate::loom::sync::atomic::{AtomicBool, Ordering}; +use crate::loom::sync::{Barrier, Mutex}; +use crate::runtime::dump::Dump; +use crate::runtime::scheduler::multi_thread::Handle; +use crate::sync::notify::Notify; + +/// Tracing status of the worker. +pub(super) struct TraceStatus { + pub(super) trace_requested: AtomicBool, + pub(super) trace_start: Barrier, + pub(super) trace_end: Barrier, + pub(super) result_ready: Notify, + pub(super) trace_result: Mutex<Option<Dump>>, +} + +impl TraceStatus { + pub(super) fn new(remotes_len: usize) -> Self { + Self { + trace_requested: AtomicBool::new(false), + trace_start: Barrier::new(remotes_len), + trace_end: Barrier::new(remotes_len), + result_ready: Notify::new(), + trace_result: Mutex::new(None), + } + } + + pub(super) fn trace_requested(&self) -> bool { + self.trace_requested.load(Ordering::Relaxed) + } + + pub(super) async fn start_trace_request(&self, handle: &Handle) { + while self + .trace_requested + .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed) + .is_err() + { + handle.notify_all(); + crate::task::yield_now().await; + } + } + + pub(super) fn stash_result(&self, dump: Dump) { + let _ = self.trace_result.lock().insert(dump); + self.result_ready.notify_one(); + } + + pub(super) fn take_result(&self) -> Option<Dump> { + self.trace_result.lock().take() + } + + pub(super) async fn end_trace_request(&self, handle: &Handle) { + while self + .trace_requested + .compare_exchange(true, false, Ordering::Acquire, Ordering::Relaxed) + .is_err() + { + handle.notify_all(); + crate::task::yield_now().await; + } + } +} diff --git a/vendor/tokio/src/runtime/scheduler/multi_thread/trace_mock.rs b/vendor/tokio/src/runtime/scheduler/multi_thread/trace_mock.rs new file mode 100644 index 000000000..2c17a4e38 --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/multi_thread/trace_mock.rs @@ -0,0 +1,11 @@ +pub(super) struct TraceStatus {} + +impl TraceStatus { + pub(super) fn new(_: usize) -> Self { + Self {} + } + + pub(super) fn trace_requested(&self) -> bool { + false + } +} diff --git a/vendor/tokio/src/runtime/scheduler/multi_thread/worker.rs b/vendor/tokio/src/runtime/scheduler/multi_thread/worker.rs new file mode 100644 index 000000000..6ae114633 --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/multi_thread/worker.rs @@ -0,0 +1,1216 @@ +//! A scheduler is initialized with a fixed number of workers. Each worker is +//! driven by a thread. Each worker has a "core" which contains data such as the +//! run queue and other state. When `block_in_place` is called, the worker's +//! "core" is handed off to a new thread allowing the scheduler to continue to +//! make progress while the originating thread blocks. +//! +//! # Shutdown +//! +//! Shutting down the runtime involves the following steps: +//! +//! 1. The Shared::close method is called. This closes the inject queue and +//! OwnedTasks instance and wakes up all worker threads. +//! +//! 2. Each worker thread observes the close signal next time it runs +//! Core::maintenance by checking whether the inject queue is closed. +//! The Core::is_shutdown flag is set to true. +//! +//! 3. The worker thread calls `pre_shutdown` in parallel. Here, the worker +//! will keep removing tasks from OwnedTasks until it is empty. No new +//! tasks can be pushed to the OwnedTasks during or after this step as it +//! was closed in step 1. +//! +//! 5. The workers call Shared::shutdown to enter the single-threaded phase of +//! shutdown. These calls will push their core to Shared::shutdown_cores, +//! and the last thread to push its core will finish the shutdown procedure. +//! +//! 6. The local run queue of each core is emptied, then the inject queue is +//! emptied. +//! +//! At this point, shutdown has completed. It is not possible for any of the +//! collections to contain any tasks at this point, as each collection was +//! closed first, then emptied afterwards. +//! +//! ## Spawns during shutdown +//! +//! When spawning tasks during shutdown, there are two cases: +//! +//! * The spawner observes the OwnedTasks being open, and the inject queue is +//! closed. +//! * The spawner observes the OwnedTasks being closed and doesn't check the +//! inject queue. +//! +//! The first case can only happen if the OwnedTasks::bind call happens before +//! or during step 1 of shutdown. In this case, the runtime will clean up the +//! task in step 3 of shutdown. +//! +//! In the latter case, the task was not spawned and the task is immediately +//! cancelled by the spawner. +//! +//! The correctness of shutdown requires both the inject queue and OwnedTasks +//! collection to have a closed bit. With a close bit on only the inject queue, +//! spawning could run in to a situation where a task is successfully bound long +//! after the runtime has shut down. With a close bit on only the OwnedTasks, +//! the first spawning situation could result in the notification being pushed +//! to the inject queue after step 6 of shutdown, which would leave a task in +//! the inject queue indefinitely. This would be a ref-count cycle and a memory +//! leak. + +use crate::loom::sync::{Arc, Mutex}; +use crate::runtime; +use crate::runtime::context; +use crate::runtime::scheduler::multi_thread::{ + idle, queue, Counters, Handle, Idle, Overflow, Parker, Stats, TraceStatus, Unparker, +}; +use crate::runtime::scheduler::{inject, Defer, Lock}; +use crate::runtime::task::OwnedTasks; +use crate::runtime::{ + blocking, coop, driver, scheduler, task, Config, SchedulerMetrics, WorkerMetrics, +}; +use crate::util::atomic_cell::AtomicCell; +use crate::util::rand::{FastRand, RngSeedGenerator}; + +use std::cell::RefCell; +use std::task::Waker; +use std::time::Duration; + +cfg_metrics! { + mod metrics; +} + +cfg_taskdump! { + mod taskdump; +} + +cfg_not_taskdump! { + mod taskdump_mock; +} + +/// A scheduler worker +pub(super) struct Worker { + /// Reference to scheduler's handle + handle: Arc<Handle>, + + /// Index holding this worker's remote state + index: usize, + + /// Used to hand-off a worker's core to another thread. + core: AtomicCell<Core>, +} + +/// Core data +struct Core { + /// Used to schedule bookkeeping tasks every so often. + tick: u32, + + /// When a task is scheduled from a worker, it is stored in this slot. The + /// worker will check this slot for a task **before** checking the run + /// queue. This effectively results in the **last** scheduled task to be run + /// next (LIFO). This is an optimization for improving locality which + /// benefits message passing patterns and helps to reduce latency. + lifo_slot: Option<Notified>, + + /// When `true`, locally scheduled tasks go to the LIFO slot. When `false`, + /// they go to the back of the `run_queue`. + lifo_enabled: bool, + + /// The worker-local run queue. + run_queue: queue::Local<Arc<Handle>>, + + /// True if the worker is currently searching for more work. Searching + /// involves attempting to steal from other workers. + is_searching: bool, + + /// True if the scheduler is being shutdown + is_shutdown: bool, + + /// True if the scheduler is being traced + is_traced: bool, + + /// Parker + /// + /// Stored in an `Option` as the parker is added / removed to make the + /// borrow checker happy. + park: Option<Parker>, + + /// Per-worker runtime stats + stats: Stats, + + /// How often to check the global queue + global_queue_interval: u32, + + /// Fast random number generator. + rand: FastRand, +} + +/// State shared across all workers +pub(crate) struct Shared { + /// Per-worker remote state. All other workers have access to this and is + /// how they communicate between each other. + remotes: Box<[Remote]>, + + /// Global task queue used for: + /// 1. Submit work to the scheduler while **not** currently on a worker thread. + /// 2. Submit work to the scheduler when a worker run queue is saturated + pub(super) inject: inject::Shared<Arc<Handle>>, + + /// Coordinates idle workers + idle: Idle, + + /// Collection of all active tasks spawned onto this executor. + pub(super) owned: OwnedTasks<Arc<Handle>>, + + /// Data synchronized by the scheduler mutex + pub(super) synced: Mutex<Synced>, + + /// Cores that have observed the shutdown signal + /// + /// The core is **not** placed back in the worker to avoid it from being + /// stolen by a thread that was spawned as part of `block_in_place`. + #[allow(clippy::vec_box)] // we're moving an already-boxed value + shutdown_cores: Mutex<Vec<Box<Core>>>, + + /// The number of cores that have observed the trace signal. + pub(super) trace_status: TraceStatus, + + /// Scheduler configuration options + config: Config, + + /// Collects metrics from the runtime. + pub(super) scheduler_metrics: SchedulerMetrics, + + pub(super) worker_metrics: Box<[WorkerMetrics]>, + + /// Only held to trigger some code on drop. This is used to get internal + /// runtime metrics that can be useful when doing performance + /// investigations. This does nothing (empty struct, no drop impl) unless + /// the `tokio_internal_mt_counters` cfg flag is set. + _counters: Counters, +} + +/// Data synchronized by the scheduler mutex +pub(crate) struct Synced { + /// Synchronized state for `Idle`. + pub(super) idle: idle::Synced, + + /// Synchronized state for `Inject`. + pub(crate) inject: inject::Synced, +} + +/// Used to communicate with a worker from other threads. +struct Remote { + /// Steals tasks from this worker. + pub(super) steal: queue::Steal<Arc<Handle>>, + + /// Unparks the associated worker thread + unpark: Unparker, +} + +/// Thread-local context +pub(crate) struct Context { + /// Worker + worker: Arc<Worker>, + + /// Core data + core: RefCell<Option<Box<Core>>>, + + /// Tasks to wake after resource drivers are polled. This is mostly to + /// handle yielded tasks. + pub(crate) defer: Defer, +} + +/// Starts the workers +pub(crate) struct Launch(Vec<Arc<Worker>>); + +/// Running a task may consume the core. If the core is still available when +/// running the task completes, it is returned. Otherwise, the worker will need +/// to stop processing. +type RunResult = Result<Box<Core>, ()>; + +/// A task handle +type Task = task::Task<Arc<Handle>>; + +/// A notified task handle +type Notified = task::Notified<Arc<Handle>>; + +/// Value picked out of thin-air. Running the LIFO slot a handful of times +/// seemms sufficient to benefit from locality. More than 3 times probably is +/// overweighing. The value can be tuned in the future with data that shows +/// improvements. +const MAX_LIFO_POLLS_PER_TICK: usize = 3; + +pub(super) fn create( + size: usize, + park: Parker, + driver_handle: driver::Handle, + blocking_spawner: blocking::Spawner, + seed_generator: RngSeedGenerator, + config: Config, +) -> (Arc<Handle>, Launch) { + let mut cores = Vec::with_capacity(size); + let mut remotes = Vec::with_capacity(size); + let mut worker_metrics = Vec::with_capacity(size); + + // Create the local queues + for _ in 0..size { + let (steal, run_queue) = queue::local(); + + let park = park.clone(); + let unpark = park.unpark(); + let metrics = WorkerMetrics::from_config(&config); + let stats = Stats::new(&metrics); + + cores.push(Box::new(Core { + tick: 0, + lifo_slot: None, + lifo_enabled: !config.disable_lifo_slot, + run_queue, + is_searching: false, + is_shutdown: false, + is_traced: false, + park: Some(park), + global_queue_interval: stats.tuned_global_queue_interval(&config), + stats, + rand: FastRand::from_seed(config.seed_generator.next_seed()), + })); + + remotes.push(Remote { steal, unpark }); + worker_metrics.push(metrics); + } + + let (idle, idle_synced) = Idle::new(size); + let (inject, inject_synced) = inject::Shared::new(); + + let remotes_len = remotes.len(); + let handle = Arc::new(Handle { + shared: Shared { + remotes: remotes.into_boxed_slice(), + inject, + idle, + owned: OwnedTasks::new(), + synced: Mutex::new(Synced { + idle: idle_synced, + inject: inject_synced, + }), + shutdown_cores: Mutex::new(vec![]), + trace_status: TraceStatus::new(remotes_len), + config, + scheduler_metrics: SchedulerMetrics::new(), + worker_metrics: worker_metrics.into_boxed_slice(), + _counters: Counters, + }, + driver: driver_handle, + blocking_spawner, + seed_generator, + }); + + let mut launch = Launch(vec![]); + + for (index, core) in cores.drain(..).enumerate() { + launch.0.push(Arc::new(Worker { + handle: handle.clone(), + index, + core: AtomicCell::new(Some(core)), + })); + } + + (handle, launch) +} + +#[track_caller] +pub(crate) fn block_in_place<F, R>(f: F) -> R +where + F: FnOnce() -> R, +{ + // Try to steal the worker core back + struct Reset { + take_core: bool, + budget: coop::Budget, + } + + impl Drop for Reset { + fn drop(&mut self) { + with_current(|maybe_cx| { + if let Some(cx) = maybe_cx { + if self.take_core { + let core = cx.worker.core.take(); + let mut cx_core = cx.core.borrow_mut(); + assert!(cx_core.is_none()); + *cx_core = core; + } + + // Reset the task budget as we are re-entering the + // runtime. + coop::set(self.budget); + } + }); + } + } + + let mut had_entered = false; + let mut take_core = false; + + let setup_result = with_current(|maybe_cx| { + match ( + crate::runtime::context::current_enter_context(), + maybe_cx.is_some(), + ) { + (context::EnterRuntime::Entered { .. }, true) => { + // We are on a thread pool runtime thread, so we just need to + // set up blocking. + had_entered = true; + } + ( + context::EnterRuntime::Entered { + allow_block_in_place, + }, + false, + ) => { + // We are on an executor, but _not_ on the thread pool. That is + // _only_ okay if we are in a thread pool runtime's block_on + // method: + if allow_block_in_place { + had_entered = true; + return Ok(()); + } else { + // This probably means we are on the current_thread runtime or in a + // LocalSet, where it is _not_ okay to block. + return Err( + "can call blocking only when running on the multi-threaded runtime", + ); + } + } + (context::EnterRuntime::NotEntered, true) => { + // This is a nested call to block_in_place (we already exited). + // All the necessary setup has already been done. + return Ok(()); + } + (context::EnterRuntime::NotEntered, false) => { + // We are outside of the tokio runtime, so blocking is fine. + // We can also skip all of the thread pool blocking setup steps. + return Ok(()); + } + } + + let cx = maybe_cx.expect("no .is_some() == false cases above should lead here"); + + // Get the worker core. If none is set, then blocking is fine! + let core = match cx.core.borrow_mut().take() { + Some(core) => core, + None => return Ok(()), + }; + + // We are taking the core from the context and sending it to another + // thread. + take_core = true; + + // The parker should be set here + assert!(core.park.is_some()); + + // In order to block, the core must be sent to another thread for + // execution. + // + // First, move the core back into the worker's shared core slot. + cx.worker.core.set(core); + + // Next, clone the worker handle and send it to a new thread for + // processing. + // + // Once the blocking task is done executing, we will attempt to + // steal the core back. + let worker = cx.worker.clone(); + runtime::spawn_blocking(move || run(worker)); + Ok(()) + }); + + if let Err(panic_message) = setup_result { + panic!("{}", panic_message); + } + + if had_entered { + // Unset the current task's budget. Blocking sections are not + // constrained by task budgets. + let _reset = Reset { + take_core, + budget: coop::stop(), + }; + + crate::runtime::context::exit_runtime(f) + } else { + f() + } +} + +impl Launch { + pub(crate) fn launch(mut self) { + for worker in self.0.drain(..) { + runtime::spawn_blocking(move || run(worker)); + } + } +} + +fn run(worker: Arc<Worker>) { + struct AbortOnPanic; + + impl Drop for AbortOnPanic { + fn drop(&mut self) { + if std::thread::panicking() { + eprintln!("worker thread panicking; aborting process"); + std::process::abort(); + } + } + } + + // Catching panics on worker threads in tests is quite tricky. Instead, when + // debug assertions are enabled, we just abort the process. + #[cfg(debug_assertions)] + let _abort_on_panic = AbortOnPanic; + + // Acquire a core. If this fails, then another thread is running this + // worker and there is nothing further to do. + let core = match worker.core.take() { + Some(core) => core, + None => return, + }; + + let handle = scheduler::Handle::MultiThread(worker.handle.clone()); + + crate::runtime::context::enter_runtime(&handle, true, |_| { + // Set the worker context. + let cx = scheduler::Context::MultiThread(Context { + worker, + core: RefCell::new(None), + defer: Defer::new(), + }); + + context::set_scheduler(&cx, || { + let cx = cx.expect_multi_thread(); + + // This should always be an error. It only returns a `Result` to support + // using `?` to short circuit. + assert!(cx.run(core).is_err()); + + // Check if there are any deferred tasks to notify. This can happen when + // the worker core is lost due to `block_in_place()` being called from + // within the task. + cx.defer.wake(); + }); + }); +} + +impl Context { + fn run(&self, mut core: Box<Core>) -> RunResult { + // Reset `lifo_enabled` here in case the core was previously stolen from + // a task that had the LIFO slot disabled. + self.reset_lifo_enabled(&mut core); + + // Start as "processing" tasks as polling tasks from the local queue + // will be one of the first things we do. + core.stats.start_processing_scheduled_tasks(); + + while !core.is_shutdown { + self.assert_lifo_enabled_is_correct(&core); + + if core.is_traced { + core = self.worker.handle.trace_core(core); + } + + // Increment the tick + core.tick(); + + // Run maintenance, if needed + core = self.maintenance(core); + + // First, check work available to the current worker. + if let Some(task) = core.next_task(&self.worker) { + core = self.run_task(task, core)?; + continue; + } + + // We consumed all work in the queues and will start searching for work. + core.stats.end_processing_scheduled_tasks(); + + // There is no more **local** work to process, try to steal work + // from other workers. + if let Some(task) = core.steal_work(&self.worker) { + // Found work, switch back to processing + core.stats.start_processing_scheduled_tasks(); + core = self.run_task(task, core)?; + } else { + // Wait for work + core = if !self.defer.is_empty() { + self.park_timeout(core, Some(Duration::from_millis(0))) + } else { + self.park(core) + }; + } + } + + core.pre_shutdown(&self.worker); + + // Signal shutdown + self.worker.handle.shutdown_core(core); + Err(()) + } + + fn run_task(&self, task: Notified, mut core: Box<Core>) -> RunResult { + let task = self.worker.handle.shared.owned.assert_owner(task); + + // Make sure the worker is not in the **searching** state. This enables + // another idle worker to try to steal work. + core.transition_from_searching(&self.worker); + + self.assert_lifo_enabled_is_correct(&core); + + // Measure the poll start time. Note that we may end up polling other + // tasks under this measurement. In this case, the tasks came from the + // LIFO slot and are considered part of the current task for scheduling + // purposes. These tasks inherent the "parent"'s limits. + core.stats.start_poll(); + + // Make the core available to the runtime context + *self.core.borrow_mut() = Some(core); + + // Run the task + coop::budget(|| { + task.run(); + let mut lifo_polls = 0; + + // As long as there is budget remaining and a task exists in the + // `lifo_slot`, then keep running. + loop { + // Check if we still have the core. If not, the core was stolen + // by another worker. + let mut core = match self.core.borrow_mut().take() { + Some(core) => core, + None => { + // In this case, we cannot call `reset_lifo_enabled()` + // because the core was stolen. The stealer will handle + // that at the top of `Context::run` + return Err(()); + } + }; + + // Check for a task in the LIFO slot + let task = match core.lifo_slot.take() { + Some(task) => task, + None => { + self.reset_lifo_enabled(&mut core); + core.stats.end_poll(); + return Ok(core); + } + }; + + if !coop::has_budget_remaining() { + core.stats.end_poll(); + + // Not enough budget left to run the LIFO task, push it to + // the back of the queue and return. + core.run_queue.push_back_or_overflow( + task, + &*self.worker.handle, + &mut core.stats, + ); + // If we hit this point, the LIFO slot should be enabled. + // There is no need to reset it. + debug_assert!(core.lifo_enabled); + return Ok(core); + } + + // Track that we are about to run a task from the LIFO slot. + lifo_polls += 1; + super::counters::inc_lifo_schedules(); + + // Disable the LIFO slot if we reach our limit + // + // In ping-ping style workloads where task A notifies task B, + // which notifies task A again, continuously prioritizing the + // LIFO slot can cause starvation as these two tasks will + // repeatedly schedule the other. To mitigate this, we limit the + // number of times the LIFO slot is prioritized. + if lifo_polls >= MAX_LIFO_POLLS_PER_TICK { + core.lifo_enabled = false; + super::counters::inc_lifo_capped(); + } + + // Run the LIFO task, then loop + *self.core.borrow_mut() = Some(core); + let task = self.worker.handle.shared.owned.assert_owner(task); + task.run(); + } + }) + } + + fn reset_lifo_enabled(&self, core: &mut Core) { + core.lifo_enabled = !self.worker.handle.shared.config.disable_lifo_slot; + } + + fn assert_lifo_enabled_is_correct(&self, core: &Core) { + debug_assert_eq!( + core.lifo_enabled, + !self.worker.handle.shared.config.disable_lifo_slot + ); + } + + fn maintenance(&self, mut core: Box<Core>) -> Box<Core> { + if core.tick % self.worker.handle.shared.config.event_interval == 0 { + super::counters::inc_num_maintenance(); + + core.stats.end_processing_scheduled_tasks(); + + // Call `park` with a 0 timeout. This enables the I/O driver, timer, ... + // to run without actually putting the thread to sleep. + core = self.park_timeout(core, Some(Duration::from_millis(0))); + + // Run regularly scheduled maintenance + core.maintenance(&self.worker); + + core.stats.start_processing_scheduled_tasks(); + } + + core + } + + /// Parks the worker thread while waiting for tasks to execute. + /// + /// This function checks if indeed there's no more work left to be done before parking. + /// Also important to notice that, before parking, the worker thread will try to take + /// ownership of the Driver (IO/Time) and dispatch any events that might have fired. + /// Whenever a worker thread executes the Driver loop, all waken tasks are scheduled + /// in its own local queue until the queue saturates (ntasks > LOCAL_QUEUE_CAPACITY). + /// When the local queue is saturated, the overflow tasks are added to the injection queue + /// from where other workers can pick them up. + /// Also, we rely on the workstealing algorithm to spread the tasks amongst workers + /// after all the IOs get dispatched + fn park(&self, mut core: Box<Core>) -> Box<Core> { + if let Some(f) = &self.worker.handle.shared.config.before_park { + f(); + } + + if core.transition_to_parked(&self.worker) { + while !core.is_shutdown && !core.is_traced { + core.stats.about_to_park(); + core = self.park_timeout(core, None); + + // Run regularly scheduled maintenance + core.maintenance(&self.worker); + + if core.transition_from_parked(&self.worker) { + break; + } + } + } + + if let Some(f) = &self.worker.handle.shared.config.after_unpark { + f(); + } + core + } + + fn park_timeout(&self, mut core: Box<Core>, duration: Option<Duration>) -> Box<Core> { + self.assert_lifo_enabled_is_correct(&core); + + // Take the parker out of core + let mut park = core.park.take().expect("park missing"); + + // Store `core` in context + *self.core.borrow_mut() = Some(core); + + // Park thread + if let Some(timeout) = duration { + park.park_timeout(&self.worker.handle.driver, timeout); + } else { + park.park(&self.worker.handle.driver); + } + + self.defer.wake(); + + // Remove `core` from context + core = self.core.borrow_mut().take().expect("core missing"); + + // Place `park` back in `core` + core.park = Some(park); + + if core.should_notify_others() { + self.worker.handle.notify_parked_local(); + } + + core + } + + pub(crate) fn defer(&self, waker: &Waker) { + self.defer.defer(waker); + } +} + +impl Core { + /// Increment the tick + fn tick(&mut self) { + self.tick = self.tick.wrapping_add(1); + } + + /// Return the next notified task available to this worker. + fn next_task(&mut self, worker: &Worker) -> Option<Notified> { + if self.tick % self.global_queue_interval == 0 { + // Update the global queue interval, if needed + self.tune_global_queue_interval(worker); + + worker + .handle + .next_remote_task() + .or_else(|| self.next_local_task()) + } else { + let maybe_task = self.next_local_task(); + + if maybe_task.is_some() { + return maybe_task; + } + + if worker.inject().is_empty() { + return None; + } + + // Other threads can only **remove** tasks from the current worker's + // `run_queue`. So, we can be confident that by the time we call + // `run_queue.push_back` below, there will be *at least* `cap` + // available slots in the queue. + let cap = usize::min( + self.run_queue.remaining_slots(), + self.run_queue.max_capacity() / 2, + ); + + // The worker is currently idle, pull a batch of work from the + // injection queue. We don't want to pull *all* the work so other + // workers can also get some. + let n = usize::min( + worker.inject().len() / worker.handle.shared.remotes.len() + 1, + cap, + ); + + let mut synced = worker.handle.shared.synced.lock(); + // safety: passing in the correct `inject::Synced`. + let mut tasks = unsafe { worker.inject().pop_n(&mut synced.inject, n) }; + + // Pop the first task to return immedietly + let ret = tasks.next(); + + // Push the rest of the on the run queue + self.run_queue.push_back(tasks); + + ret + } + } + + fn next_local_task(&mut self) -> Option<Notified> { + self.lifo_slot.take().or_else(|| self.run_queue.pop()) + } + + /// Function responsible for stealing tasks from another worker + /// + /// Note: Only if less than half the workers are searching for tasks to steal + /// a new worker will actually try to steal. The idea is to make sure not all + /// workers will be trying to steal at the same time. + fn steal_work(&mut self, worker: &Worker) -> Option<Notified> { + if !self.transition_to_searching(worker) { + return None; + } + + let num = worker.handle.shared.remotes.len(); + // Start from a random worker + let start = self.rand.fastrand_n(num as u32) as usize; + + for i in 0..num { + let i = (start + i) % num; + + // Don't steal from ourself! We know we don't have work. + if i == worker.index { + continue; + } + + let target = &worker.handle.shared.remotes[i]; + if let Some(task) = target + .steal + .steal_into(&mut self.run_queue, &mut self.stats) + { + return Some(task); + } + } + + // Fallback on checking the global queue + worker.handle.next_remote_task() + } + + fn transition_to_searching(&mut self, worker: &Worker) -> bool { + if !self.is_searching { + self.is_searching = worker.handle.shared.idle.transition_worker_to_searching(); + } + + self.is_searching + } + + fn transition_from_searching(&mut self, worker: &Worker) { + if !self.is_searching { + return; + } + + self.is_searching = false; + worker.handle.transition_worker_from_searching(); + } + + fn has_tasks(&self) -> bool { + self.lifo_slot.is_some() || self.run_queue.has_tasks() + } + + fn should_notify_others(&self) -> bool { + // If there are tasks available to steal, but this worker is not + // looking for tasks to steal, notify another worker. + if self.is_searching { + return false; + } + self.lifo_slot.is_some() as usize + self.run_queue.len() > 1 + } + + /// Prepares the worker state for parking. + /// + /// Returns true if the transition happened, false if there is work to do first. + fn transition_to_parked(&mut self, worker: &Worker) -> bool { + // Workers should not park if they have work to do + if self.has_tasks() || self.is_traced { + return false; + } + + // When the final worker transitions **out** of searching to parked, it + // must check all the queues one last time in case work materialized + // between the last work scan and transitioning out of searching. + let is_last_searcher = worker.handle.shared.idle.transition_worker_to_parked( + &worker.handle.shared, + worker.index, + self.is_searching, + ); + + // The worker is no longer searching. Setting this is the local cache + // only. + self.is_searching = false; + + if is_last_searcher { + worker.handle.notify_if_work_pending(); + } + + true + } + + /// Returns `true` if the transition happened. + fn transition_from_parked(&mut self, worker: &Worker) -> bool { + // If a task is in the lifo slot/run queue, then we must unpark regardless of + // being notified + if self.has_tasks() { + // When a worker wakes, it should only transition to the "searching" + // state when the wake originates from another worker *or* a new task + // is pushed. We do *not* want the worker to transition to "searching" + // when it wakes when the I/O driver receives new events. + self.is_searching = !worker + .handle + .shared + .idle + .unpark_worker_by_id(&worker.handle.shared, worker.index); + return true; + } + + if worker + .handle + .shared + .idle + .is_parked(&worker.handle.shared, worker.index) + { + return false; + } + + // When unparked, the worker is in the searching state. + self.is_searching = true; + true + } + + /// Runs maintenance work such as checking the pool's state. + fn maintenance(&mut self, worker: &Worker) { + self.stats + .submit(&worker.handle.shared.worker_metrics[worker.index]); + + if !self.is_shutdown { + // Check if the scheduler has been shutdown + let synced = worker.handle.shared.synced.lock(); + self.is_shutdown = worker.inject().is_closed(&synced.inject); + } + + if !self.is_traced { + // Check if the worker should be tracing. + self.is_traced = worker.handle.shared.trace_status.trace_requested(); + } + } + + /// Signals all tasks to shut down, and waits for them to complete. Must run + /// before we enter the single-threaded phase of shutdown processing. + fn pre_shutdown(&mut self, worker: &Worker) { + // Signal to all tasks to shut down. + worker.handle.shared.owned.close_and_shutdown_all(); + + self.stats + .submit(&worker.handle.shared.worker_metrics[worker.index]); + } + + /// Shuts down the core. + fn shutdown(&mut self, handle: &Handle) { + // Take the core + let mut park = self.park.take().expect("park missing"); + + // Drain the queue + while self.next_local_task().is_some() {} + + park.shutdown(&handle.driver); + } + + fn tune_global_queue_interval(&mut self, worker: &Worker) { + let next = self + .stats + .tuned_global_queue_interval(&worker.handle.shared.config); + + debug_assert!(next > 1); + + // Smooth out jitter + if abs_diff(self.global_queue_interval, next) > 2 { + self.global_queue_interval = next; + } + } +} + +impl Worker { + /// Returns a reference to the scheduler's injection queue. + fn inject(&self) -> &inject::Shared<Arc<Handle>> { + &self.handle.shared.inject + } +} + +// TODO: Move `Handle` impls into handle.rs +impl task::Schedule for Arc<Handle> { + fn release(&self, task: &Task) -> Option<Task> { + self.shared.owned.remove(task) + } + + fn schedule(&self, task: Notified) { + self.schedule_task(task, false); + } + + fn yield_now(&self, task: Notified) { + self.schedule_task(task, true); + } +} + +impl Handle { + pub(super) fn schedule_task(&self, task: Notified, is_yield: bool) { + with_current(|maybe_cx| { + if let Some(cx) = maybe_cx { + // Make sure the task is part of the **current** scheduler. + if self.ptr_eq(&cx.worker.handle) { + // And the current thread still holds a core + if let Some(core) = cx.core.borrow_mut().as_mut() { + self.schedule_local(core, task, is_yield); + return; + } + } + } + + // Otherwise, use the inject queue. + self.push_remote_task(task); + self.notify_parked_remote(); + }) + } + + fn schedule_local(&self, core: &mut Core, task: Notified, is_yield: bool) { + core.stats.inc_local_schedule_count(); + + // Spawning from the worker thread. If scheduling a "yield" then the + // task must always be pushed to the back of the queue, enabling other + // tasks to be executed. If **not** a yield, then there is more + // flexibility and the task may go to the front of the queue. + let should_notify = if is_yield || !core.lifo_enabled { + core.run_queue + .push_back_or_overflow(task, self, &mut core.stats); + true + } else { + // Push to the LIFO slot + let prev = core.lifo_slot.take(); + let ret = prev.is_some(); + + if let Some(prev) = prev { + core.run_queue + .push_back_or_overflow(prev, self, &mut core.stats); + } + + core.lifo_slot = Some(task); + + ret + }; + + // Only notify if not currently parked. If `park` is `None`, then the + // scheduling is from a resource driver. As notifications often come in + // batches, the notification is delayed until the park is complete. + if should_notify && core.park.is_some() { + self.notify_parked_local(); + } + } + + fn next_remote_task(&self) -> Option<Notified> { + if self.shared.inject.is_empty() { + return None; + } + + let mut synced = self.shared.synced.lock(); + // safety: passing in correct `idle::Synced` + unsafe { self.shared.inject.pop(&mut synced.inject) } + } + + fn push_remote_task(&self, task: Notified) { + self.shared.scheduler_metrics.inc_remote_schedule_count(); + + let mut synced = self.shared.synced.lock(); + // safety: passing in correct `idle::Synced` + unsafe { + self.shared.inject.push(&mut synced.inject, task); + } + } + + pub(super) fn close(&self) { + if self + .shared + .inject + .close(&mut self.shared.synced.lock().inject) + { + self.notify_all(); + } + } + + fn notify_parked_local(&self) { + super::counters::inc_num_inc_notify_local(); + + if let Some(index) = self.shared.idle.worker_to_notify(&self.shared) { + super::counters::inc_num_unparks_local(); + self.shared.remotes[index].unpark.unpark(&self.driver); + } + } + + fn notify_parked_remote(&self) { + if let Some(index) = self.shared.idle.worker_to_notify(&self.shared) { + self.shared.remotes[index].unpark.unpark(&self.driver); + } + } + + pub(super) fn notify_all(&self) { + for remote in &self.shared.remotes[..] { + remote.unpark.unpark(&self.driver); + } + } + + fn notify_if_work_pending(&self) { + for remote in &self.shared.remotes[..] { + if !remote.steal.is_empty() { + self.notify_parked_local(); + return; + } + } + + if !self.shared.inject.is_empty() { + self.notify_parked_local(); + } + } + + fn transition_worker_from_searching(&self) { + if self.shared.idle.transition_worker_from_searching() { + // We are the final searching worker. Because work was found, we + // need to notify another worker. + self.notify_parked_local(); + } + } + + /// Signals that a worker has observed the shutdown signal and has replaced + /// its core back into its handle. + /// + /// If all workers have reached this point, the final cleanup is performed. + fn shutdown_core(&self, core: Box<Core>) { + let mut cores = self.shared.shutdown_cores.lock(); + cores.push(core); + + if cores.len() != self.shared.remotes.len() { + return; + } + + debug_assert!(self.shared.owned.is_empty()); + + for mut core in cores.drain(..) { + core.shutdown(self); + } + + // Drain the injection queue + // + // We already shut down every task, so we can simply drop the tasks. + while let Some(task) = self.next_remote_task() { + drop(task); + } + } + + fn ptr_eq(&self, other: &Handle) -> bool { + std::ptr::eq(self, other) + } +} + +impl Overflow<Arc<Handle>> for Handle { + fn push(&self, task: task::Notified<Arc<Handle>>) { + self.push_remote_task(task); + } + + fn push_batch<I>(&self, iter: I) + where + I: Iterator<Item = task::Notified<Arc<Handle>>>, + { + unsafe { + self.shared.inject.push_batch(self, iter); + } + } +} + +pub(crate) struct InjectGuard<'a> { + lock: crate::loom::sync::MutexGuard<'a, Synced>, +} + +impl<'a> AsMut<inject::Synced> for InjectGuard<'a> { + fn as_mut(&mut self) -> &mut inject::Synced { + &mut self.lock.inject + } +} + +impl<'a> Lock<inject::Synced> for &'a Handle { + type Handle = InjectGuard<'a>; + + fn lock(self) -> Self::Handle { + InjectGuard { + lock: self.shared.synced.lock(), + } + } +} + +#[track_caller] +fn with_current<R>(f: impl FnOnce(Option<&Context>) -> R) -> R { + use scheduler::Context::MultiThread; + + context::with_scheduler(|ctx| match ctx { + Some(MultiThread(ctx)) => f(Some(ctx)), + _ => f(None), + }) +} + +// `u32::abs_diff` is not available on Tokio's MSRV. +fn abs_diff(a: u32, b: u32) -> u32 { + if a > b { + a - b + } else { + b - a + } +} diff --git a/vendor/tokio/src/runtime/scheduler/multi_thread/worker/metrics.rs b/vendor/tokio/src/runtime/scheduler/multi_thread/worker/metrics.rs new file mode 100644 index 000000000..a9a5ab3ed --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/multi_thread/worker/metrics.rs @@ -0,0 +1,11 @@ +use super::Shared; + +impl Shared { + pub(crate) fn injection_queue_depth(&self) -> usize { + self.inject.len() + } + + pub(crate) fn worker_local_queue_depth(&self, worker: usize) -> usize { + self.remotes[worker].steal.len() + } +} diff --git a/vendor/tokio/src/runtime/scheduler/multi_thread/worker/taskdump.rs b/vendor/tokio/src/runtime/scheduler/multi_thread/worker/taskdump.rs new file mode 100644 index 000000000..d310d9f6d --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/multi_thread/worker/taskdump.rs @@ -0,0 +1,79 @@ +use super::{Core, Handle, Shared}; + +use crate::loom::sync::Arc; +use crate::runtime::scheduler::multi_thread::Stats; +use crate::runtime::task::trace::trace_multi_thread; +use crate::runtime::{dump, WorkerMetrics}; + +use std::time::Duration; + +impl Handle { + pub(super) fn trace_core(&self, mut core: Box<Core>) -> Box<Core> { + core.is_traced = false; + + if core.is_shutdown { + return core; + } + + // wait for other workers, or timeout without tracing + let timeout = Duration::from_millis(250); // a _very_ generous timeout + let barrier = + if let Some(barrier) = self.shared.trace_status.trace_start.wait_timeout(timeout) { + barrier + } else { + // don't attempt to trace + return core; + }; + + if !barrier.is_leader() { + // wait for leader to finish tracing + self.shared.trace_status.trace_end.wait(); + return core; + } + + // trace + + let owned = &self.shared.owned; + let mut local = self.shared.steal_all(); + let synced = &self.shared.synced; + let injection = &self.shared.inject; + + // safety: `trace_multi_thread` is invoked with the same `synced` that `injection` + // was created with. + let traces = unsafe { trace_multi_thread(owned, &mut local, synced, injection) } + .into_iter() + .map(dump::Task::new) + .collect(); + + let result = dump::Dump::new(traces); + + // stash the result + self.shared.trace_status.stash_result(result); + + // allow other workers to proceed + self.shared.trace_status.trace_end.wait(); + + core + } +} + +impl Shared { + /// Steal all tasks from remotes into a single local queue. + pub(super) fn steal_all(&self) -> super::queue::Local<Arc<Handle>> { + let (_steal, mut local) = super::queue::local(); + + let worker_metrics = WorkerMetrics::new(); + let mut stats = Stats::new(&worker_metrics); + + for remote in self.remotes.iter() { + let steal = &remote.steal; + while !steal.is_empty() { + if let Some(task) = steal.steal_into(&mut local, &mut stats) { + local.push_back([task].into_iter()); + } + } + } + + local + } +} diff --git a/vendor/tokio/src/runtime/scheduler/multi_thread/worker/taskdump_mock.rs b/vendor/tokio/src/runtime/scheduler/multi_thread/worker/taskdump_mock.rs new file mode 100644 index 000000000..24c5600ce --- /dev/null +++ b/vendor/tokio/src/runtime/scheduler/multi_thread/worker/taskdump_mock.rs @@ -0,0 +1,7 @@ +use super::{Core, Handle}; + +impl Handle { + pub(super) fn trace_core(&self, core: Box<Core>) -> Box<Core> { + core + } +} diff --git a/vendor/tokio/src/runtime/signal/mod.rs b/vendor/tokio/src/runtime/signal/mod.rs new file mode 100644 index 000000000..24f2f4c6c --- /dev/null +++ b/vendor/tokio/src/runtime/signal/mod.rs @@ -0,0 +1,142 @@ +#![cfg_attr(not(feature = "rt"), allow(dead_code))] + +//! Signal driver + +use crate::runtime::{driver, io}; +use crate::signal::registry::globals; + +use mio::net::UnixStream; +use std::io::{self as std_io, Read}; +use std::sync::{Arc, Weak}; +use std::time::Duration; + +/// Responsible for registering wakeups when an OS signal is received, and +/// subsequently dispatching notifications to any signal listeners as appropriate. +/// +/// Note: this driver relies on having an enabled IO driver in order to listen to +/// pipe write wakeups. +#[derive(Debug)] +pub(crate) struct Driver { + /// Thread parker. The `Driver` park implementation delegates to this. + io: io::Driver, + + /// A pipe for receiving wake events from the signal handler + receiver: UnixStream, + + /// Shared state. The driver keeps a strong ref and the handle keeps a weak + /// ref. The weak ref is used to check if the driver is still active before + /// trying to register a signal handler. + inner: Arc<()>, +} + +#[derive(Debug, Default)] +pub(crate) struct Handle { + /// Paired w/ the `Arc` above and is used to check if the driver is still + /// around before attempting to register a signal handler. + inner: Weak<()>, +} + +// ===== impl Driver ===== + +impl Driver { + /// Creates a new signal `Driver` instance that delegates wakeups to `park`. + pub(crate) fn new(io: io::Driver, io_handle: &io::Handle) -> std_io::Result<Self> { + use std::mem::ManuallyDrop; + use std::os::unix::io::{AsRawFd, FromRawFd}; + + // NB: We give each driver a "fresh" receiver file descriptor to avoid + // the issues described in alexcrichton/tokio-process#42. + // + // In the past we would reuse the actual receiver file descriptor and + // swallow any errors around double registration of the same descriptor. + // I'm not sure if the second (failed) registration simply doesn't end + // up receiving wake up notifications, or there could be some race + // condition when consuming readiness events, but having distinct + // descriptors appears to mitigate this. + // + // Unfortunately we cannot just use a single global UnixStream instance + // either, since we can't assume they will always be registered with the + // exact same reactor. + // + // Mio 0.7 removed `try_clone()` as an API due to unexpected behavior + // with registering dups with the same reactor. In this case, duping is + // safe as each dup is registered with separate reactors **and** we + // only expect at least one dup to receive the notification. + + // Manually drop as we don't actually own this instance of UnixStream. + let receiver_fd = globals().receiver.as_raw_fd(); + + // safety: there is nothing unsafe about this, but the `from_raw_fd` fn is marked as unsafe. + let original = + ManuallyDrop::new(unsafe { std::os::unix::net::UnixStream::from_raw_fd(receiver_fd) }); + let mut receiver = UnixStream::from_std(original.try_clone()?); + + io_handle.register_signal_receiver(&mut receiver)?; + + Ok(Self { + io, + receiver, + inner: Arc::new(()), + }) + } + + /// Returns a handle to this event loop which can be sent across threads + /// and can be used as a proxy to the event loop itself. + pub(crate) fn handle(&self) -> Handle { + Handle { + inner: Arc::downgrade(&self.inner), + } + } + + pub(crate) fn park(&mut self, handle: &driver::Handle) { + self.io.park(handle); + self.process(); + } + + pub(crate) fn park_timeout(&mut self, handle: &driver::Handle, duration: Duration) { + self.io.park_timeout(handle, duration); + self.process(); + } + + pub(crate) fn shutdown(&mut self, handle: &driver::Handle) { + self.io.shutdown(handle) + } + + fn process(&mut self) { + // If the signal pipe has not received a readiness event, then there is + // nothing else to do. + if !self.io.consume_signal_ready() { + return; + } + + // Drain the pipe completely so we can receive a new readiness event + // if another signal has come in. + let mut buf = [0; 128]; + loop { + match self.receiver.read(&mut buf) { + Ok(0) => panic!("EOF on self-pipe"), + Ok(_) => continue, // Keep reading + Err(e) if e.kind() == std_io::ErrorKind::WouldBlock => break, + Err(e) => panic!("Bad read on self-pipe: {}", e), + } + } + + // Broadcast any signals which were received + globals().broadcast(); + } +} + +// ===== impl Handle ===== + +impl Handle { + pub(crate) fn check_inner(&self) -> std_io::Result<()> { + if self.inner.strong_count() > 0 { + Ok(()) + } else { + Err(std_io::Error::new( + std_io::ErrorKind::Other, + "signal driver gone", + )) + } + } +} diff --git a/vendor/tokio/src/runtime/spawner.rs b/vendor/tokio/src/runtime/spawner.rs deleted file mode 100644 index fbcde2cfa..000000000 --- a/vendor/tokio/src/runtime/spawner.rs +++ /dev/null @@ -1,45 +0,0 @@ -cfg_rt! { - use crate::future::Future; - use crate::runtime::basic_scheduler; - use crate::task::JoinHandle; -} - -cfg_rt_multi_thread! { - use crate::runtime::thread_pool; -} - -#[derive(Debug, Clone)] -pub(crate) enum Spawner { - #[cfg(feature = "rt")] - Basic(basic_scheduler::Spawner), - #[cfg(feature = "rt-multi-thread")] - ThreadPool(thread_pool::Spawner), -} - -impl Spawner { - pub(crate) fn shutdown(&mut self) { - #[cfg(feature = "rt-multi-thread")] - { - if let Spawner::ThreadPool(spawner) = self { - spawner.shutdown(); - } - } - } -} - -cfg_rt! { - impl Spawner { - pub(crate) fn spawn<F>(&self, future: F) -> JoinHandle<F::Output> - where - F: Future + Send + 'static, - F::Output: Send + 'static, - { - match self { - #[cfg(feature = "rt")] - Spawner::Basic(spawner) => spawner.spawn(future), - #[cfg(feature = "rt-multi-thread")] - Spawner::ThreadPool(spawner) => spawner.spawn(future), - } - } - } -} diff --git a/vendor/tokio/src/runtime/task/abort.rs b/vendor/tokio/src/runtime/task/abort.rs new file mode 100644 index 000000000..6edca1004 --- /dev/null +++ b/vendor/tokio/src/runtime/task/abort.rs @@ -0,0 +1,87 @@ +use crate::runtime::task::{Header, RawTask}; +use std::fmt; +use std::panic::{RefUnwindSafe, UnwindSafe}; + +/// An owned permission to abort a spawned task, without awaiting its completion. +/// +/// Unlike a [`JoinHandle`], an `AbortHandle` does *not* represent the +/// permission to await the task's completion, only to terminate it. +/// +/// The task may be aborted by calling the [`AbortHandle::abort`] method. +/// Dropping an `AbortHandle` releases the permission to terminate the task +/// --- it does *not* abort the task. +/// +/// [`JoinHandle`]: crate::task::JoinHandle +#[cfg_attr(docsrs, doc(cfg(feature = "rt")))] +pub struct AbortHandle { + raw: RawTask, +} + +impl AbortHandle { + pub(super) fn new(raw: RawTask) -> Self { + Self { raw } + } + + /// Abort the task associated with the handle. + /// + /// Awaiting a cancelled task might complete as usual if the task was + /// already completed at the time it was cancelled, but most likely it + /// will fail with a [cancelled] `JoinError`. + /// + /// If the task was already cancelled, such as by [`JoinHandle::abort`], + /// this method will do nothing. + /// + /// [cancelled]: method@super::error::JoinError::is_cancelled + /// [`JoinHandle::abort`]: method@super::JoinHandle::abort + pub fn abort(&self) { + self.raw.remote_abort(); + } + + /// Checks if the task associated with this `AbortHandle` has finished. + /// + /// Please note that this method can return `false` even if `abort` has been + /// called on the task. This is because the cancellation process may take + /// some time, and this method does not return `true` until it has + /// completed. + pub fn is_finished(&self) -> bool { + let state = self.raw.state().load(); + state.is_complete() + } + + /// Returns a [task ID] that uniquely identifies this task relative to other + /// currently spawned tasks. + /// + /// **Note**: This is an [unstable API][unstable]. The public API of this type + /// may break in 1.x releases. See [the documentation on unstable + /// features][unstable] for details. + /// + /// [task ID]: crate::task::Id + /// [unstable]: crate#unstable-features + #[cfg(tokio_unstable)] + #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))] + pub fn id(&self) -> super::Id { + // Safety: The header pointer is valid. + unsafe { Header::get_id(self.raw.header_ptr()) } + } +} + +unsafe impl Send for AbortHandle {} +unsafe impl Sync for AbortHandle {} + +impl UnwindSafe for AbortHandle {} +impl RefUnwindSafe for AbortHandle {} + +impl fmt::Debug for AbortHandle { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + // Safety: The header pointer is valid. + let id_ptr = unsafe { Header::get_id_ptr(self.raw.header_ptr()) }; + let id = unsafe { id_ptr.as_ref() }; + fmt.debug_struct("AbortHandle").field("id", id).finish() + } +} + +impl Drop for AbortHandle { + fn drop(&mut self) { + self.raw.drop_abort_handle(); + } +} diff --git a/vendor/tokio/src/runtime/task/core.rs b/vendor/tokio/src/runtime/task/core.rs index 428c921fe..110933e58 100644 --- a/vendor/tokio/src/runtime/task/core.rs +++ b/vendor/tokio/src/runtime/task/core.rs @@ -11,9 +11,10 @@ use crate::future::Future; use crate::loom::cell::UnsafeCell; +use crate::runtime::context; use crate::runtime::task::raw::{self, Vtable}; use crate::runtime::task::state::State; -use crate::runtime::task::{Notified, Schedule, Task}; +use crate::runtime::task::{Id, Schedule}; use crate::util::linked_list; use std::pin::Pin; @@ -24,6 +25,97 @@ use std::task::{Context, Poll, Waker}; /// /// It is critical for `Header` to be the first field as the task structure will /// be referenced by both *mut Cell and *mut Header. +/// +/// Any changes to the layout of this struct _must_ also be reflected in the +/// const fns in raw.rs. +/// +// # This struct should be cache padded to avoid false sharing. The cache padding rules are copied +// from crossbeam-utils/src/cache_padded.rs +// +// Starting from Intel's Sandy Bridge, spatial prefetcher is now pulling pairs of 64-byte cache +// lines at a time, so we have to align to 128 bytes rather than 64. +// +// Sources: +// - https://www.intel.com/content/dam/www/public/us/en/documents/manuals/64-ia-32-architectures-optimization-manual.pdf +// - https://github.com/facebook/folly/blob/1b5288e6eea6df074758f877c849b6e73bbb9fbb/folly/lang/Align.h#L107 +// +// ARM's big.LITTLE architecture has asymmetric cores and "big" cores have 128-byte cache line size. +// +// Sources: +// - https://www.mono-project.com/news/2016/09/12/arm64-icache/ +// +// powerpc64 has 128-byte cache line size. +// +// Sources: +// - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_ppc64x.go#L9 +#[cfg_attr( + any( + target_arch = "x86_64", + target_arch = "aarch64", + target_arch = "powerpc64", + ), + repr(align(128)) +)] +// arm, mips, mips64, riscv64, sparc, and hexagon have 32-byte cache line size. +// +// Sources: +// - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_arm.go#L7 +// - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_mips.go#L7 +// - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_mipsle.go#L7 +// - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_mips64x.go#L9 +// - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_riscv64.go#L7 +// - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/sparc/include/asm/cache.h#L17 +// - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/hexagon/include/asm/cache.h#L12 +// +// riscv32 is assumed not to exceed the cache line size of riscv64. +#[cfg_attr( + any( + target_arch = "arm", + target_arch = "mips", + target_arch = "mips64", + target_arch = "riscv32", + target_arch = "riscv64", + target_arch = "sparc", + target_arch = "hexagon", + ), + repr(align(32)) +)] +// m68k has 16-byte cache line size. +// +// Sources: +// - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/m68k/include/asm/cache.h#L9 +#[cfg_attr(target_arch = "m68k", repr(align(16)))] +// s390x has 256-byte cache line size. +// +// Sources: +// - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_s390x.go#L7 +// - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/s390/include/asm/cache.h#L13 +#[cfg_attr(target_arch = "s390x", repr(align(256)))] +// x86, wasm, and sparc64 have 64-byte cache line size. +// +// Sources: +// - https://github.com/golang/go/blob/dda2991c2ea0c5914714469c4defc2562a907230/src/internal/cpu/cpu_x86.go#L9 +// - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_wasm.go#L7 +// - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/sparc/include/asm/cache.h#L19 +// +// All others are assumed to have 64-byte cache line size. +#[cfg_attr( + not(any( + target_arch = "x86_64", + target_arch = "aarch64", + target_arch = "powerpc64", + target_arch = "arm", + target_arch = "mips", + target_arch = "mips64", + target_arch = "riscv32", + target_arch = "riscv64", + target_arch = "sparc", + target_arch = "hexagon", + target_arch = "m68k", + target_arch = "s390x", + )), + repr(align(64)) +)] #[repr(C)] pub(super) struct Cell<T: Future, S> { /// Hot task state data @@ -36,10 +128,6 @@ pub(super) struct Cell<T: Future, S> { pub(super) trailer: Trailer, } -pub(super) struct Scheduler<S> { - scheduler: UnsafeCell<Option<S>>, -} - pub(super) struct CoreStage<T: Future> { stage: UnsafeCell<Stage<T>>, } @@ -47,45 +135,71 @@ pub(super) struct CoreStage<T: Future> { /// The core of the task. /// /// Holds the future or output, depending on the stage of execution. +/// +/// Any changes to the layout of this struct _must_ also be reflected in the +/// const fns in raw.rs. +#[repr(C)] pub(super) struct Core<T: Future, S> { - /// Scheduler used to drive this future - pub(super) scheduler: Scheduler<S>, + /// Scheduler used to drive this future. + pub(super) scheduler: S, + + /// The task's ID, used for populating `JoinError`s. + pub(super) task_id: Id, - /// Either the future or the output + /// Either the future or the output. pub(super) stage: CoreStage<T>, } /// Crate public as this is also needed by the pool. #[repr(C)] pub(crate) struct Header { - /// Task state + /// Task state. pub(super) state: State, - pub(crate) owned: UnsafeCell<linked_list::Pointers<Header>>, - - /// Pointer to next task, used with the injection queue - pub(crate) queue_next: UnsafeCell<Option<NonNull<Header>>>, - - /// Pointer to the next task in the transfer stack - pub(super) stack_next: UnsafeCell<Option<NonNull<Header>>>, + /// Pointer to next task, used with the injection queue. + pub(super) queue_next: UnsafeCell<Option<NonNull<Header>>>, /// Table of function pointers for executing actions on the task. pub(super) vtable: &'static Vtable, + /// This integer contains the id of the OwnedTasks or LocalOwnedTasks that + /// this task is stored in. If the task is not in any list, should be the + /// id of the list that it was previously in, or zero if it has never been + /// in any list. + /// + /// Once a task has been bound to a list, it can never be bound to another + /// list, even if removed from the first list. + /// + /// The id is not unset when removed from a list because we want to be able + /// to read the id without synchronization, even if it is concurrently being + /// removed from the list. + pub(super) owner_id: UnsafeCell<u64>, + /// The tracing ID for this instrumented task. #[cfg(all(tokio_unstable, feature = "tracing"))] - pub(super) id: Option<tracing::Id>, + pub(super) tracing_id: Option<tracing::Id>, } unsafe impl Send for Header {} unsafe impl Sync for Header {} -/// Cold data is stored after the future. +/// Cold data is stored after the future. Data is considered cold if it is only +/// used during creation or shutdown of the task. pub(super) struct Trailer { + /// Pointers for the linked list in the `OwnedTasks` that owns this task. + pub(super) owned: linked_list::Pointers<Header>, /// Consumer task waiting on completion of this task. pub(super) waker: UnsafeCell<Option<Waker>>, } +generate_addr_of_methods! { + impl<> Trailer { + pub(super) unsafe fn addr_of_owned(self: NonNull<Self>) -> NonNull<linked_list::Pointers<Header>> { + &self.owned + } + } +} + /// Either the future or the output. pub(super) enum Stage<T: Future> { Running(T), @@ -96,117 +210,48 @@ pub(super) enum Stage<T: Future> { impl<T: Future, S: Schedule> Cell<T, S> { /// Allocates a new task cell, containing the header, trailer, and core /// structures. - pub(super) fn new(future: T, state: State) -> Box<Cell<T, S>> { + pub(super) fn new(future: T, scheduler: S, state: State, task_id: Id) -> Box<Cell<T, S>> { #[cfg(all(tokio_unstable, feature = "tracing"))] - let id = future.id(); - Box::new(Cell { + let tracing_id = future.id(); + let result = Box::new(Cell { header: Header { state, - owned: UnsafeCell::new(linked_list::Pointers::new()), queue_next: UnsafeCell::new(None), - stack_next: UnsafeCell::new(None), vtable: raw::vtable::<T, S>(), + owner_id: UnsafeCell::new(0), #[cfg(all(tokio_unstable, feature = "tracing"))] - id, + tracing_id, }, core: Core { - scheduler: Scheduler { - scheduler: UnsafeCell::new(None), - }, + scheduler, stage: CoreStage { stage: UnsafeCell::new(Stage::Running(future)), }, + task_id, }, trailer: Trailer { waker: UnsafeCell::new(None), + owned: linked_list::Pointers::new(), }, - }) - } -} - -impl<S: Schedule> Scheduler<S> { - pub(super) fn with_mut<R>(&self, f: impl FnOnce(*mut Option<S>) -> R) -> R { - self.scheduler.with_mut(f) - } - - /// Bind a scheduler to the task. - /// - /// This only happens on the first poll and must be preceded by a call to - /// `is_bound` to determine if binding is appropriate or not. - /// - /// # Safety - /// - /// Binding must not be done concurrently since it will mutate the task - /// core through a shared reference. - pub(super) fn bind_scheduler(&self, task: Task<S>) { - // This function may be called concurrently, but the __first__ time it - // is called, the caller has unique access to this field. All subsequent - // concurrent calls will be via the `Waker`, which will "happens after" - // the first poll. - // - // In other words, it is always safe to read the field and it is safe to - // write to the field when it is `None`. - debug_assert!(!self.is_bound()); - - // Bind the task to the scheduler - let scheduler = S::bind(task); - - // Safety: As `scheduler` is not set, this is the first poll - self.scheduler.with_mut(|ptr| unsafe { - *ptr = Some(scheduler); }); - } - /// Returns true if the task is bound to a scheduler. - pub(super) fn is_bound(&self) -> bool { - // Safety: never called concurrently w/ a mutation. - self.scheduler.with(|ptr| unsafe { (*ptr).is_some() }) - } + #[cfg(debug_assertions)] + { + let trailer_addr = (&result.trailer) as *const Trailer as usize; + let trailer_ptr = unsafe { Header::get_trailer(NonNull::from(&result.header)) }; + assert_eq!(trailer_addr, trailer_ptr.as_ptr() as usize); - /// Schedule the future for execution - pub(super) fn schedule(&self, task: Notified<S>) { - self.scheduler.with(|ptr| { - // Safety: Can only be called after initial `poll`, which is the - // only time the field is mutated. - match unsafe { &*ptr } { - Some(scheduler) => scheduler.schedule(task), - None => panic!("no scheduler set"), - } - }); - } + let scheduler_addr = (&result.core.scheduler) as *const S as usize; + let scheduler_ptr = + unsafe { Header::get_scheduler::<S>(NonNull::from(&result.header)) }; + assert_eq!(scheduler_addr, scheduler_ptr.as_ptr() as usize); - /// Schedule the future for execution in the near future, yielding the - /// thread to other tasks. - pub(super) fn yield_now(&self, task: Notified<S>) { - self.scheduler.with(|ptr| { - // Safety: Can only be called after initial `poll`, which is the - // only time the field is mutated. - match unsafe { &*ptr } { - Some(scheduler) => scheduler.yield_now(task), - None => panic!("no scheduler set"), - } - }); - } + let id_addr = (&result.core.task_id) as *const Id as usize; + let id_ptr = unsafe { Header::get_id_ptr(NonNull::from(&result.header)) }; + assert_eq!(id_addr, id_ptr.as_ptr() as usize); + } - /// Release the task - /// - /// If the `Scheduler` implementation is able to, it returns the `Task` - /// handle immediately. The caller of this function will batch a ref-dec - /// with a state change. - pub(super) fn release(&self, task: Task<S>) -> Option<Task<S>> { - use std::mem::ManuallyDrop; - - let task = ManuallyDrop::new(task); - - self.scheduler.with(|ptr| { - // Safety: Can only be called after initial `poll`, which is the - // only time the field is mutated. - match unsafe { &*ptr } { - Some(scheduler) => scheduler.release(&*task), - // Task was never polled - None => None, - } - }) + result } } @@ -214,8 +259,30 @@ impl<T: Future> CoreStage<T> { pub(super) fn with_mut<R>(&self, f: impl FnOnce(*mut Stage<T>) -> R) -> R { self.stage.with_mut(f) } +} + +/// Set and clear the task id in the context when the future is executed or +/// dropped, or when the output produced by the future is dropped. +pub(crate) struct TaskIdGuard { + parent_task_id: Option<Id>, +} + +impl TaskIdGuard { + fn enter(id: Id) -> Self { + TaskIdGuard { + parent_task_id: context::set_current_task_id(Some(id)), + } + } +} - /// Poll the future +impl Drop for TaskIdGuard { + fn drop(&mut self) { + context::set_current_task_id(self.parent_task_id); + } +} + +impl<T: Future, S: Schedule> Core<T, S> { + /// Polls the future. /// /// # Safety /// @@ -230,7 +297,7 @@ impl<T: Future> CoreStage<T> { /// heap. pub(super) fn poll(&self, mut cx: Context<'_>) -> Poll<T::Output> { let res = { - self.stage.with_mut(|ptr| { + self.stage.stage.with_mut(|ptr| { // Safety: The caller ensures mutual exclusion to the field. let future = match unsafe { &mut *ptr } { Stage::Running(future) => future, @@ -240,6 +307,7 @@ impl<T: Future> CoreStage<T> { // Safety: The caller ensures the future is pinned. let future = unsafe { Pin::new_unchecked(future) }; + let _guard = TaskIdGuard::enter(self.task_id); future.poll(&mut cx) }) }; @@ -251,7 +319,7 @@ impl<T: Future> CoreStage<T> { res } - /// Drop the future + /// Drops the future. /// /// # Safety /// @@ -263,7 +331,7 @@ impl<T: Future> CoreStage<T> { } } - /// Store the task output + /// Stores the task output. /// /// # Safety /// @@ -275,7 +343,7 @@ impl<T: Future> CoreStage<T> { } } - /// Take the task output + /// Takes the task output. /// /// # Safety /// @@ -283,7 +351,7 @@ impl<T: Future> CoreStage<T> { pub(super) fn take_output(&self) -> super::Result<T::Output> { use std::mem; - self.stage.with_mut(|ptr| { + self.stage.stage.with_mut(|ptr| { // Safety:: the caller ensures mutual exclusion to the field. match mem::replace(unsafe { &mut *ptr }, Stage::Consumed) { Stage::Finished(output) => output, @@ -293,38 +361,99 @@ impl<T: Future> CoreStage<T> { } unsafe fn set_stage(&self, stage: Stage<T>) { - self.stage.with_mut(|ptr| *ptr = stage) + let _guard = TaskIdGuard::enter(self.task_id); + self.stage.stage.with_mut(|ptr| *ptr = stage) } } -cfg_rt_multi_thread! { - impl Header { - pub(crate) fn shutdown(&self) { - use crate::runtime::task::RawTask; +impl Header { + pub(super) unsafe fn set_next(&self, next: Option<NonNull<Header>>) { + self.queue_next.with_mut(|ptr| *ptr = next); + } + + // safety: The caller must guarantee exclusive access to this field, and + // must ensure that the id is either 0 or the id of the OwnedTasks + // containing this task. + pub(super) unsafe fn set_owner_id(&self, owner: u64) { + self.owner_id.with_mut(|ptr| *ptr = owner); + } - let task = unsafe { RawTask::from_raw(self.into()) }; - task.shutdown(); - } + pub(super) fn get_owner_id(&self) -> u64 { + // safety: If there are concurrent writes, then that write has violated + // the safety requirements on `set_owner_id`. + unsafe { self.owner_id.with(|ptr| *ptr) } + } - pub(crate) unsafe fn set_next(&self, next: Option<NonNull<Header>>) { - self.queue_next.with_mut(|ptr| *ptr = next); - } + /// Gets a pointer to the `Trailer` of the task containing this `Header`. + /// + /// # Safety + /// + /// The provided raw pointer must point at the header of a task. + pub(super) unsafe fn get_trailer(me: NonNull<Header>) -> NonNull<Trailer> { + let offset = me.as_ref().vtable.trailer_offset; + let trailer = me.as_ptr().cast::<u8>().add(offset).cast::<Trailer>(); + NonNull::new_unchecked(trailer) + } + + /// Gets a pointer to the scheduler of the task containing this `Header`. + /// + /// # Safety + /// + /// The provided raw pointer must point at the header of a task. + /// + /// The generic type S must be set to the correct scheduler type for this + /// task. + pub(super) unsafe fn get_scheduler<S>(me: NonNull<Header>) -> NonNull<S> { + let offset = me.as_ref().vtable.scheduler_offset; + let scheduler = me.as_ptr().cast::<u8>().add(offset).cast::<S>(); + NonNull::new_unchecked(scheduler) + } + + /// Gets a pointer to the id of the task containing this `Header`. + /// + /// # Safety + /// + /// The provided raw pointer must point at the header of a task. + pub(super) unsafe fn get_id_ptr(me: NonNull<Header>) -> NonNull<Id> { + let offset = me.as_ref().vtable.id_offset; + let id = me.as_ptr().cast::<u8>().add(offset).cast::<Id>(); + NonNull::new_unchecked(id) + } + + /// Gets the id of the task containing this `Header`. + /// + /// # Safety + /// + /// The provided raw pointer must point at the header of a task. + pub(super) unsafe fn get_id(me: NonNull<Header>) -> Id { + let ptr = Header::get_id_ptr(me).as_ptr(); + *ptr + } + + /// Gets the tracing id of the task containing this `Header`. + /// + /// # Safety + /// + /// The provided raw pointer must point at the header of a task. + #[cfg(all(tokio_unstable, feature = "tracing"))] + pub(super) unsafe fn get_tracing_id(me: &NonNull<Header>) -> Option<&tracing::Id> { + me.as_ref().tracing_id.as_ref() } } impl Trailer { - pub(crate) unsafe fn set_waker(&self, waker: Option<Waker>) { + pub(super) unsafe fn set_waker(&self, waker: Option<Waker>) { self.waker.with_mut(|ptr| { *ptr = waker; }); } - pub(crate) unsafe fn will_wake(&self, waker: &Waker) -> bool { + pub(super) unsafe fn will_wake(&self, waker: &Waker) -> bool { self.waker .with(|ptr| (*ptr).as_ref().unwrap().will_wake(waker)) } - pub(crate) fn wake_join(&self) { + pub(super) fn wake_join(&self) { self.waker.with(|ptr| match unsafe { &*ptr } { Some(waker) => waker.wake_by_ref(), None => panic!("waker missing"), diff --git a/vendor/tokio/src/runtime/task/error.rs b/vendor/tokio/src/runtime/task/error.rs index 177fe65e9..f7ead77b7 100644 --- a/vendor/tokio/src/runtime/task/error.rs +++ b/vendor/tokio/src/runtime/task/error.rs @@ -1,39 +1,43 @@ use std::any::Any; use std::fmt; use std::io; -use std::sync::Mutex; +use super::Id; +use crate::util::SyncWrapper; cfg_rt! { /// Task failed to execute to completion. pub struct JoinError { repr: Repr, + id: Id, } } enum Repr { Cancelled, - Panic(Mutex<Box<dyn Any + Send + 'static>>), + Panic(SyncWrapper<Box<dyn Any + Send + 'static>>), } impl JoinError { - pub(crate) fn cancelled() -> JoinError { + pub(crate) fn cancelled(id: Id) -> JoinError { JoinError { repr: Repr::Cancelled, + id, } } - pub(crate) fn panic(err: Box<dyn Any + Send + 'static>) -> JoinError { + pub(crate) fn panic(id: Id, err: Box<dyn Any + Send + 'static>) -> JoinError { JoinError { - repr: Repr::Panic(Mutex::new(err)), + repr: Repr::Panic(SyncWrapper::new(err)), + id, } } - /// Returns true if the error was caused by the task being cancelled + /// Returns true if the error was caused by the task being cancelled. pub fn is_cancelled(&self) -> bool { matches!(&self.repr, Repr::Cancelled) } - /// Returns true if the error was caused by the task panicking + /// Returns true if the error was caused by the task panicking. /// /// # Examples /// @@ -78,6 +82,7 @@ impl JoinError { /// } /// } /// ``` + #[track_caller] pub fn into_panic(self) -> Box<dyn Any + Send + 'static> { self.try_into_panic() .expect("`JoinError` reason is not a panic.") @@ -106,17 +111,32 @@ impl JoinError { /// ``` pub fn try_into_panic(self) -> Result<Box<dyn Any + Send + 'static>, JoinError> { match self.repr { - Repr::Panic(p) => Ok(p.into_inner().expect("Extracting panic from mutex")), + Repr::Panic(p) => Ok(p.into_inner()), _ => Err(self), } } + + /// Returns a [task ID] that identifies the task which errored relative to + /// other currently spawned tasks. + /// + /// **Note**: This is an [unstable API][unstable]. The public API of this type + /// may break in 1.x releases. See [the documentation on unstable + /// features][unstable] for details. + /// + /// [task ID]: crate::task::Id + /// [unstable]: crate#unstable-features + #[cfg(tokio_unstable)] + #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))] + pub fn id(&self) -> Id { + self.id + } } impl fmt::Display for JoinError { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { match &self.repr { - Repr::Cancelled => write!(fmt, "cancelled"), - Repr::Panic(_) => write!(fmt, "panic"), + Repr::Cancelled => write!(fmt, "task {} was cancelled", self.id), + Repr::Panic(_) => write!(fmt, "task {} panicked", self.id), } } } @@ -124,8 +144,8 @@ impl fmt::Display for JoinError { impl fmt::Debug for JoinError { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { match &self.repr { - Repr::Cancelled => write!(fmt, "JoinError::Cancelled"), - Repr::Panic(_) => write!(fmt, "JoinError::Panic(...)"), + Repr::Cancelled => write!(fmt, "JoinError::Cancelled({:?})", self.id), + Repr::Panic(_) => write!(fmt, "JoinError::Panic({:?}, ...)", self.id), } } } diff --git a/vendor/tokio/src/runtime/task/harness.rs b/vendor/tokio/src/runtime/task/harness.rs index 8cd649dc7..8e3c3d14f 100644 --- a/vendor/tokio/src/runtime/task/harness.rs +++ b/vendor/tokio/src/runtime/task/harness.rs @@ -1,15 +1,16 @@ use crate::future::Future; -use crate::runtime::task::core::{Cell, Core, CoreStage, Header, Scheduler, Trailer}; -use crate::runtime::task::state::Snapshot; +use crate::runtime::task::core::{Cell, Core, Header, Trailer}; +use crate::runtime::task::state::{Snapshot, State}; use crate::runtime::task::waker::waker_ref; -use crate::runtime::task::{JoinError, Notified, Schedule, Task}; +use crate::runtime::task::{JoinError, Notified, RawTask, Schedule, Task}; use std::mem; +use std::mem::ManuallyDrop; use std::panic; use std::ptr::NonNull; use std::task::{Context, Poll, Waker}; -/// Typed raw task handle +/// Typed raw task handle. pub(super) struct Harness<T: Future, S: 'static> { cell: NonNull<Cell<T, S>>, } @@ -25,8 +26,16 @@ where } } + fn header_ptr(&self) -> NonNull<Header> { + self.cell.cast() + } + fn header(&self) -> &Header { - unsafe { &self.cell.as_ref().header } + unsafe { &*self.header_ptr().as_ptr() } + } + + fn state(&self) -> &State { + &self.header().state } fn trailer(&self) -> &Trailer { @@ -36,13 +45,91 @@ where fn core(&self) -> &Core<T, S> { unsafe { &self.cell.as_ref().core } } +} + +/// Task operations that can be implemented without being generic over the +/// scheduler or task. Only one version of these methods should exist in the +/// final binary. +impl RawTask { + pub(super) fn drop_reference(self) { + if self.state().ref_dec() { + self.dealloc(); + } + } + + /// This call consumes a ref-count and notifies the task. This will create a + /// new Notified and submit it if necessary. + /// + /// The caller does not need to hold a ref-count besides the one that was + /// passed to this call. + pub(super) fn wake_by_val(&self) { + use super::state::TransitionToNotifiedByVal; + + match self.state().transition_to_notified_by_val() { + TransitionToNotifiedByVal::Submit => { + // The caller has given us a ref-count, and the transition has + // created a new ref-count, so we now hold two. We turn the new + // ref-count Notified and pass it to the call to `schedule`. + // + // The old ref-count is retained for now to ensure that the task + // is not dropped during the call to `schedule` if the call + // drops the task it was given. + self.schedule(); + + // Now that we have completed the call to schedule, we can + // release our ref-count. + self.drop_reference(); + } + TransitionToNotifiedByVal::Dealloc => { + self.dealloc(); + } + TransitionToNotifiedByVal::DoNothing => {} + } + } + + /// This call notifies the task. It will not consume any ref-counts, but the + /// caller should hold a ref-count. This will create a new Notified and + /// submit it if necessary. + pub(super) fn wake_by_ref(&self) { + use super::state::TransitionToNotifiedByRef; + + match self.state().transition_to_notified_by_ref() { + TransitionToNotifiedByRef::Submit => { + // The transition above incremented the ref-count for a new task + // and the caller also holds a ref-count. The caller's ref-count + // ensures that the task is not destroyed even if the new task + // is dropped before `schedule` returns. + self.schedule(); + } + TransitionToNotifiedByRef::DoNothing => {} + } + } - fn scheduler_view(&self) -> SchedulerView<'_, S> { - SchedulerView { - header: self.header(), - scheduler: &self.core().scheduler, + /// Remotely aborts the task. + /// + /// The caller should hold a ref-count, but we do not consume it. + /// + /// This is similar to `shutdown` except that it asks the runtime to perform + /// the shutdown. This is necessary to avoid the shutdown happening in the + /// wrong thread for non-Send tasks. + pub(super) fn remote_abort(&self) { + if self.state().transition_to_notified_and_cancel() { + // The transition has created a new ref-count, which we turn into + // a Notified and pass to the task. + // + // Since the caller holds a ref-count, the task cannot be destroyed + // before the call to `schedule` returns even if the call drops the + // `Notified` internally. + self.schedule(); } } + + /// Try to set the waker notified when the task is complete. Returns true if + /// the task has already completed. If this call returns false, then the + /// waker will not be notified. + pub(super) fn try_set_join_waker(&self, waker: &Waker) -> bool { + can_read_output(self.header(), self.trailer(), waker) + } } impl<T, S> Harness<T, S> @@ -50,43 +137,109 @@ where T: Future, S: Schedule, { - /// Polls the inner future. + pub(super) fn drop_reference(self) { + if self.state().ref_dec() { + self.dealloc(); + } + } + + /// Polls the inner future. A ref-count is consumed. /// /// All necessary state checks and transitions are performed. - /// /// Panics raised while polling the future are handled. pub(super) fn poll(self) { + // We pass our ref-count to `poll_inner`. match self.poll_inner() { PollFuture::Notified => { - // Signal yield - self.core().scheduler.yield_now(Notified(self.to_task())); - // The ref-count was incremented as part of - // `transition_to_idle`. + // The `poll_inner` call has given us two ref-counts back. + // We give one of them to a new task and call `yield_now`. + self.core() + .scheduler + .yield_now(Notified(self.get_new_task())); + + // The remaining ref-count is now dropped. We kept the extra + // ref-count until now to ensure that even if the `yield_now` + // call drops the provided task, the task isn't deallocated + // before after `yield_now` returns. self.drop_reference(); } - PollFuture::DropReference => { - self.drop_reference(); + PollFuture::Complete => { + self.complete(); } - PollFuture::Complete(out, is_join_interested) => { - self.complete(out, is_join_interested); + PollFuture::Dealloc => { + self.dealloc(); } - PollFuture::None => (), + PollFuture::Done => (), } } - fn poll_inner(&self) -> PollFuture<T::Output> { - let snapshot = match self.scheduler_view().transition_to_running() { - TransitionToRunning::Ok(snapshot) => snapshot, - TransitionToRunning::DropReference => return PollFuture::DropReference, - }; + /// Polls the task and cancel it if necessary. This takes ownership of a + /// ref-count. + /// + /// If the return value is Notified, the caller is given ownership of two + /// ref-counts. + /// + /// If the return value is Complete, the caller is given ownership of a + /// single ref-count, which should be passed on to `complete`. + /// + /// If the return value is Dealloc, then this call consumed the last + /// ref-count and the caller should call `dealloc`. + /// + /// Otherwise the ref-count is consumed and the caller should not access + /// `self` again. + fn poll_inner(&self) -> PollFuture { + use super::state::{TransitionToIdle, TransitionToRunning}; + + match self.state().transition_to_running() { + TransitionToRunning::Success => { + let header_ptr = self.header_ptr(); + let waker_ref = waker_ref::<T, S>(&header_ptr); + let cx = Context::from_waker(&waker_ref); + let res = poll_future(self.core(), cx); + + if res == Poll::Ready(()) { + // The future completed. Move on to complete the task. + return PollFuture::Complete; + } - // The transition to `Running` done above ensures that a lock on the - // future has been obtained. This also ensures the `*mut T` pointer - // contains the future (as opposed to the output) and is initialized. + match self.state().transition_to_idle() { + TransitionToIdle::Ok => PollFuture::Done, + TransitionToIdle::OkNotified => PollFuture::Notified, + TransitionToIdle::OkDealloc => PollFuture::Dealloc, + TransitionToIdle::Cancelled => { + // The transition to idle failed because the task was + // cancelled during the poll. + cancel_task(self.core()); + PollFuture::Complete + } + } + } + TransitionToRunning::Cancelled => { + cancel_task(self.core()); + PollFuture::Complete + } + TransitionToRunning::Failed => PollFuture::Done, + TransitionToRunning::Dealloc => PollFuture::Dealloc, + } + } - let waker_ref = waker_ref::<T, S>(self.header()); - let cx = Context::from_waker(&*waker_ref); - poll_future(self.header(), &self.core().stage, snapshot, cx) + /// Forcibly shuts down the task. + /// + /// Attempt to transition to `Running` in order to forcibly shutdown the + /// task. If the task is currently running or in a state of completion, then + /// there is nothing further to do. When the task completes running, it will + /// notice the `CANCELLED` bit and finalize the task. + pub(super) fn shutdown(self) { + if !self.state().transition_to_shutdown() { + // The task is concurrently running. No further work needed. + self.drop_reference(); + return; + } + + // By transitioning the lifecycle to `Running`, we have permission to + // drop the future. + cancel_task(self.core()); + self.complete(); } pub(super) fn dealloc(self) { @@ -95,8 +248,20 @@ where // Check causality self.core().stage.with_mut(drop); - self.core().scheduler.with_mut(drop); + // Safety: The caller of this method just transitioned our ref-count to + // zero, so it is our responsibility to release the allocation. + // + // We don't hold any references into the allocation at this point, but + // it is possible for another thread to still hold a `&State` into the + // allocation if that other thread has decremented its last ref-count, + // but has not yet returned from the relevant method on `State`. + // + // However, the `State` type consists of just an `AtomicUsize`, and an + // `AtomicUsize` wraps the entirety of its contents in an `UnsafeCell`. + // As explained in the documentation for `UnsafeCell`, such references + // are allowed to be dangling after their last use, even if the + // reference has not yet gone out of scope. unsafe { drop(Box::from_raw(self.cell.as_ptr())); } @@ -107,227 +272,92 @@ where /// Read the task output into `dst`. pub(super) fn try_read_output(self, dst: &mut Poll<super::Result<T::Output>>, waker: &Waker) { if can_read_output(self.header(), self.trailer(), waker) { - *dst = Poll::Ready(self.core().stage.take_output()); + *dst = Poll::Ready(self.core().take_output()); } } pub(super) fn drop_join_handle_slow(self) { - let mut maybe_panic = None; - // Try to unset `JOIN_INTEREST`. This must be done as a first step in // case the task concurrently completed. - if self.header().state.unset_join_interested().is_err() { + if self.state().unset_join_interested().is_err() { // It is our responsibility to drop the output. This is critical as // the task output may not be `Send` and as such must remain with // the scheduler or `JoinHandle`. i.e. if the output remains in the // task structure until the task is deallocated, it may be dropped // by a Waker on any arbitrary thread. - let panic = panic::catch_unwind(panic::AssertUnwindSafe(|| { - self.core().stage.drop_future_or_output(); + // + // Panics are delivered to the user via the `JoinHandle`. Given that + // they are dropping the `JoinHandle`, we assume they are not + // interested in the panic and swallow it. + let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| { + self.core().drop_future_or_output(); })); - if let Err(panic) = panic { - maybe_panic = Some(panic); - } } // Drop the `JoinHandle` reference, possibly deallocating the task self.drop_reference(); - - if let Some(panic) = maybe_panic { - panic::resume_unwind(panic); - } - } - - // ===== waker behavior ===== - - pub(super) fn wake_by_val(self) { - self.wake_by_ref(); - self.drop_reference(); - } - - pub(super) fn wake_by_ref(&self) { - if self.header().state.transition_to_notified() { - self.core().scheduler.schedule(Notified(self.to_task())); - } - } - - pub(super) fn drop_reference(self) { - if self.header().state.ref_dec() { - self.dealloc(); - } } - #[cfg(all(tokio_unstable, feature = "tracing"))] - pub(super) fn id(&self) -> Option<&tracing::Id> { - self.header().id.as_ref() - } - - /// Forcibly shutdown the task - /// - /// Attempt to transition to `Running` in order to forcibly shutdown the - /// task. If the task is currently running or in a state of completion, then - /// there is nothing further to do. When the task completes running, it will - /// notice the `CANCELLED` bit and finalize the task. - pub(super) fn shutdown(self) { - if !self.header().state.transition_to_shutdown() { - // The task is concurrently running. No further work needed. - return; - } - - // By transitioning the lifecycle to `Running`, we have permission to - // drop the future. - let err = cancel_task(&self.core().stage); - self.complete(Err(err), true) - } + // ====== internal ====== - /// Remotely abort the task - /// - /// This is similar to `shutdown` except that it asks the runtime to perform - /// the shutdown. This is necessary to avoid the shutdown happening in the - /// wrong thread for non-Send tasks. - pub(super) fn remote_abort(self) { - if self.header().state.transition_to_notified_and_cancel() { - self.core().scheduler.schedule(Notified(self.to_task())); - } - } + /// Completes the task. This method assumes that the state is RUNNING. + fn complete(self) { + // The future has completed and its output has been written to the task + // stage. We transition from running to complete. - // ====== internal ====== + let snapshot = self.state().transition_to_complete(); - fn complete(self, output: super::Result<T::Output>, is_join_interested: bool) { - // We catch panics here because dropping the output may panic. - // - // Dropping the output can also happen in the first branch inside - // transition_to_complete. + // We catch panics here in case dropping the future or waking the + // JoinHandle panics. let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| { - if is_join_interested { - // Store the output. The future has already been dropped - // - // Safety: Mutual exclusion is obtained by having transitioned the task - // state -> Running - let stage = &self.core().stage; - stage.store_output(output); - - // Transition to `Complete`, notifying the `JoinHandle` if necessary. - transition_to_complete(self.header(), stage, &self.trailer()); - } else { - drop(output); + if !snapshot.is_join_interested() { + // The `JoinHandle` is not interested in the output of + // this task. It is our responsibility to drop the + // output. + self.core().drop_future_or_output(); + } else if snapshot.is_join_waker_set() { + // Notify the waker. Reading the waker field is safe per rule 4 + // in task/mod.rs, since the JOIN_WAKER bit is set and the call + // to transition_to_complete() above set the COMPLETE bit. + self.trailer().wake_join(); } })); // The task has completed execution and will no longer be scheduled. - // - // Attempts to batch a ref-dec with the state transition below. + let num_release = self.release(); - if self - .scheduler_view() - .transition_to_terminal(is_join_interested) - { - self.dealloc() + if self.state().transition_to_terminal(num_release) { + self.dealloc(); } } - fn to_task(&self) -> Task<S> { - self.scheduler_view().to_task() - } -} + /// Releases the task from the scheduler. Returns the number of ref-counts + /// that should be decremented. + fn release(&self) -> usize { + // We don't actually increment the ref-count here, but the new task is + // never destroyed, so that's ok. + let me = ManuallyDrop::new(self.get_new_task()); -enum TransitionToRunning { - Ok(Snapshot), - DropReference, -} - -struct SchedulerView<'a, S> { - header: &'a Header, - scheduler: &'a Scheduler<S>, -} - -impl<'a, S> SchedulerView<'a, S> -where - S: Schedule, -{ - fn to_task(&self) -> Task<S> { - // SAFETY The header is from the same struct containing the scheduler `S` so the cast is safe - unsafe { Task::from_raw(self.header.into()) } - } - - /// Returns true if the task should be deallocated. - fn transition_to_terminal(&self, is_join_interested: bool) -> bool { - let ref_dec = if self.scheduler.is_bound() { - if let Some(task) = self.scheduler.release(self.to_task()) { - mem::forget(task); - true - } else { - false - } + if let Some(task) = self.core().scheduler.release(&me) { + mem::forget(task); + 2 } else { - false - }; - - // This might deallocate - let snapshot = self - .header - .state - .transition_to_terminal(!is_join_interested, ref_dec); - - snapshot.ref_count() == 0 - } - - fn transition_to_running(&self) -> TransitionToRunning { - // If this is the first time the task is polled, the task will be bound - // to the scheduler, in which case the task ref count must be - // incremented. - let is_not_bound = !self.scheduler.is_bound(); - - // Transition the task to the running state. - // - // A failure to transition here indicates the task has been cancelled - // while in the run queue pending execution. - let snapshot = match self.header.state.transition_to_running(is_not_bound) { - Ok(snapshot) => snapshot, - Err(_) => { - // The task was shutdown while in the run queue. At this point, - // we just hold a ref counted reference. Since we do not have access to it here - // return `DropReference` so the caller drops it. - return TransitionToRunning::DropReference; - } - }; - - if is_not_bound { - // Ensure the task is bound to a scheduler instance. Since this is - // the first time polling the task, a scheduler instance is pulled - // from the local context and assigned to the task. - // - // The scheduler maintains ownership of the task and responds to - // `wake` calls. - // - // The task reference count has been incremented. - // - // Safety: Since we have unique access to the task so that we can - // safely call `bind_scheduler`. - self.scheduler.bind_scheduler(self.to_task()); + 1 } - TransitionToRunning::Ok(snapshot) } -} -/// Transitions the task's lifecycle to `Complete`. Notifies the -/// `JoinHandle` if it still has interest in the completion. -fn transition_to_complete<T>(header: &Header, stage: &CoreStage<T>, trailer: &Trailer) -where - T: Future, -{ - // Transition the task's lifecycle to `Complete` and get a snapshot of - // the task's sate. - let snapshot = header.state.transition_to_complete(); - - if !snapshot.is_join_interested() { - // The `JoinHandle` is not interested in the output of this task. It - // is our responsibility to drop the output. - stage.drop_future_or_output(); - } else if snapshot.has_join_waker() { - // Notify the join handle. The previous transition obtains the - // lock on the waker cell. - trailer.wake_join(); + /// Creates a new task that holds its own ref-count. + /// + /// # Safety + /// + /// Any use of `self` after this call must ensure that a ref-count to the + /// task holds the task alive until after the use of `self`. Passing the + /// returned Task to any method on `self` is unsound if dropping the Task + /// could drop `self` before the call on `self` returned. + fn get_new_task(&self) -> Task<S> { + // safety: The header is at the beginning of the cell, so this cast is + // safe. + unsafe { Task::from_raw(self.cell.cast()) } } } @@ -338,36 +368,30 @@ fn can_read_output(header: &Header, trailer: &Trailer, waker: &Waker) -> bool { debug_assert!(snapshot.is_join_interested()); if !snapshot.is_complete() { - // The waker must be stored in the task struct. - let res = if snapshot.has_join_waker() { - // There already is a waker stored in the struct. If it matches - // the provided waker, then there is no further work to do. - // Otherwise, the waker must be swapped. - let will_wake = unsafe { - // Safety: when `JOIN_INTEREST` is set, only `JOIN_HANDLE` - // may mutate the `waker` field. - trailer.will_wake(waker) - }; - - if will_wake { - // The task is not complete **and** the waker is up to date, - // there is nothing further that needs to be done. + // If the task is not complete, try storing the provided waker in the + // task's waker field. + + let res = if snapshot.is_join_waker_set() { + // If JOIN_WAKER is set, then JoinHandle has previously stored a + // waker in the waker field per step (iii) of rule 5 in task/mod.rs. + + // Optimization: if the stored waker and the provided waker wake the + // same task, then return without touching the waker field. (Reading + // the waker field below is safe per rule 3 in task/mod.rs.) + if unsafe { trailer.will_wake(waker) } { return false; } - // Unset the `JOIN_WAKER` to gain mutable access to the `waker` - // field then update the field with the new join worker. - // - // This requires two atomic operations, unsetting the bit and - // then resetting it. If the task transitions to complete - // concurrently to either one of those operations, then setting - // the join waker fails and we proceed to reading the task - // output. + // Otherwise swap the stored waker with the provided waker by + // following the rule 5 in task/mod.rs. header .state .unset_waker() .and_then(|snapshot| set_join_waker(header, trailer, waker.clone(), snapshot)) } else { + // If JOIN_WAKER is unset, then JoinHandle has mutable access to the + // waker field per rule 2 in task/mod.rs; therefore, skip step (i) + // of rule 5 and try to store the provided waker in the waker field. set_join_waker(header, trailer, waker.clone(), snapshot) }; @@ -388,7 +412,7 @@ fn set_join_waker( snapshot: Snapshot, ) -> Result<Snapshot, Snapshot> { assert!(snapshot.is_join_interested()); - assert!(!snapshot.has_join_waker()); + assert!(!snapshot.is_join_waker_set()); // Safety: Only the `JoinHandle` may set the `waker` field. When // `JOIN_INTEREST` is **not** set, nothing else will touch the field. @@ -409,73 +433,69 @@ fn set_join_waker( res } -enum PollFuture<T> { - Complete(Result<T, JoinError>, bool), - DropReference, +enum PollFuture { + Complete, Notified, - None, + Done, + Dealloc, } -fn cancel_task<T: Future>(stage: &CoreStage<T>) -> JoinError { +/// Cancels the task and store the appropriate error in the stage field. +fn cancel_task<T: Future, S: Schedule>(core: &Core<T, S>) { // Drop the future from a panic guard. let res = panic::catch_unwind(panic::AssertUnwindSafe(|| { - stage.drop_future_or_output(); + core.drop_future_or_output(); })); - if let Err(err) = res { - // Dropping the future panicked, complete the join - // handle with the panic to avoid dropping the panic - // on the ground. - JoinError::panic(err) - } else { - JoinError::cancelled() + match res { + Ok(()) => { + core.store_output(Err(JoinError::cancelled(core.task_id))); + } + Err(panic) => { + core.store_output(Err(JoinError::panic(core.task_id, panic))); + } } } -fn poll_future<T: Future>( - header: &Header, - core: &CoreStage<T>, - snapshot: Snapshot, - cx: Context<'_>, -) -> PollFuture<T::Output> { - if snapshot.is_cancelled() { - PollFuture::Complete(Err(cancel_task(core)), snapshot.is_join_interested()) - } else { - let res = panic::catch_unwind(panic::AssertUnwindSafe(|| { - struct Guard<'a, T: Future> { - core: &'a CoreStage<T>, - } - - impl<T: Future> Drop for Guard<'_, T> { - fn drop(&mut self) { - self.core.drop_future_or_output(); - } +/// Polls the future. If the future completes, the output is written to the +/// stage field. +fn poll_future<T: Future, S: Schedule>(core: &Core<T, S>, cx: Context<'_>) -> Poll<()> { + // Poll the future. + let output = panic::catch_unwind(panic::AssertUnwindSafe(|| { + struct Guard<'a, T: Future, S: Schedule> { + core: &'a Core<T, S>, + } + impl<'a, T: Future, S: Schedule> Drop for Guard<'a, T, S> { + fn drop(&mut self) { + // If the future panics on poll, we drop it inside the panic + // guard. + self.core.drop_future_or_output(); } + } + let guard = Guard { core }; + let res = guard.core.poll(cx); + mem::forget(guard); + res + })); - let guard = Guard { core }; - - let res = guard.core.poll(cx); + // Prepare output for being placed in the core stage. + let output = match output { + Ok(Poll::Pending) => return Poll::Pending, + Ok(Poll::Ready(output)) => Ok(output), + Err(panic) => { + core.scheduler.unhandled_panic(); + Err(JoinError::panic(core.task_id, panic)) + } + }; - // prevent the guard from dropping the future - mem::forget(guard); + // Catch and ignore panics if the future panics on drop. + let res = panic::catch_unwind(panic::AssertUnwindSafe(|| { + core.store_output(output); + })); - res - })); - match res { - Ok(Poll::Pending) => match header.state.transition_to_idle() { - Ok(snapshot) => { - if snapshot.is_notified() { - PollFuture::Notified - } else { - PollFuture::None - } - } - Err(_) => PollFuture::Complete(Err(cancel_task(core)), true), - }, - Ok(Poll::Ready(ok)) => PollFuture::Complete(Ok(ok), snapshot.is_join_interested()), - Err(err) => { - PollFuture::Complete(Err(JoinError::panic(err)), snapshot.is_join_interested()) - } - } + if res.is_err() { + core.scheduler.unhandled_panic(); } + + Poll::Ready(()) } diff --git a/vendor/tokio/src/runtime/task/id.rs b/vendor/tokio/src/runtime/task/id.rs new file mode 100644 index 000000000..2b0d95c02 --- /dev/null +++ b/vendor/tokio/src/runtime/task/id.rs @@ -0,0 +1,87 @@ +use crate::runtime::context; + +use std::fmt; + +/// An opaque ID that uniquely identifies a task relative to all other currently +/// running tasks. +/// +/// # Notes +/// +/// - Task IDs are unique relative to other *currently running* tasks. When a +/// task completes, the same ID may be used for another task. +/// - Task IDs are *not* sequential, and do not indicate the order in which +/// tasks are spawned, what runtime a task is spawned on, or any other data. +/// - The task ID of the currently running task can be obtained from inside the +/// task via the [`task::try_id()`](crate::task::try_id()) and +/// [`task::id()`](crate::task::id()) functions and from outside the task via +/// the [`JoinHandle::id()`](crate::task::JoinHandle::id()) function. +/// +/// **Note**: This is an [unstable API][unstable]. The public API of this type +/// may break in 1.x releases. See [the documentation on unstable +/// features][unstable] for details. +/// +/// [unstable]: crate#unstable-features +#[cfg_attr(docsrs, doc(cfg(all(feature = "rt", tokio_unstable))))] +#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))] +#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)] +pub struct Id(u64); + +/// Returns the [`Id`] of the currently running task. +/// +/// # Panics +/// +/// This function panics if called from outside a task. Please note that calls +/// to `block_on` do not have task IDs, so the method will panic if called from +/// within a call to `block_on`. For a version of this function that doesn't +/// panic, see [`task::try_id()`](crate::runtime::task::try_id()). +/// +/// **Note**: This is an [unstable API][unstable]. The public API of this type +/// may break in 1.x releases. See [the documentation on unstable +/// features][unstable] for details. +/// +/// [task ID]: crate::task::Id +/// [unstable]: crate#unstable-features +#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))] +#[track_caller] +pub fn id() -> Id { + context::current_task_id().expect("Can't get a task id when not inside a task") +} + +/// Returns the [`Id`] of the currently running task, or `None` if called outside +/// of a task. +/// +/// This function is similar to [`task::id()`](crate::runtime::task::id()), except +/// that it returns `None` rather than panicking if called outside of a task +/// context. +/// +/// **Note**: This is an [unstable API][unstable]. The public API of this type +/// may break in 1.x releases. See [the documentation on unstable +/// features][unstable] for details. +/// +/// [task ID]: crate::task::Id +/// [unstable]: crate#unstable-features +#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))] +#[track_caller] +pub fn try_id() -> Option<Id> { + context::current_task_id() +} + +impl fmt::Display for Id { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(f) + } +} + +impl Id { + pub(crate) fn next() -> Self { + use crate::loom::sync::atomic::{Ordering::Relaxed, StaticAtomicU64}; + + static NEXT_ID: StaticAtomicU64 = StaticAtomicU64::new(1); + + Self(NEXT_ID.fetch_add(1, Relaxed)) + } + + pub(crate) fn as_u64(&self) -> u64 { + self.0 + } +} diff --git a/vendor/tokio/src/runtime/task/join.rs b/vendor/tokio/src/runtime/task/join.rs index 2fe40a721..ee3925884 100644 --- a/vendor/tokio/src/runtime/task/join.rs +++ b/vendor/tokio/src/runtime/task/join.rs @@ -1,16 +1,19 @@ -use crate::runtime::task::RawTask; +use crate::runtime::task::{Header, RawTask}; use std::fmt; use std::future::Future; use std::marker::PhantomData; +use std::panic::{RefUnwindSafe, UnwindSafe}; use std::pin::Pin; -use std::task::{Context, Poll}; +use std::task::{Context, Poll, Waker}; cfg_rt! { /// An owned permission to join on a task (await its termination). /// - /// This can be thought of as the equivalent of [`std::thread::JoinHandle`] for - /// a task rather than a thread. + /// This can be thought of as the equivalent of [`std::thread::JoinHandle`] + /// for a Tokio task rather than a thread. Note that the background task + /// associated with this `JoinHandle` started running immediately when you + /// called spawn, even if you have not yet awaited the `JoinHandle`. /// /// A `JoinHandle` *detaches* the associated task when it is dropped, which /// means that there is no longer any handle to the task, and no way to `join` @@ -19,6 +22,15 @@ cfg_rt! { /// This `struct` is created by the [`task::spawn`] and [`task::spawn_blocking`] /// functions. /// + /// # Cancel safety + /// + /// The `&mut JoinHandle<T>` type is cancel safe. If it is used as the event + /// in a `tokio::select!` statement and some other branch completes first, + /// then it is guaranteed that the output of the task is not lost. + /// + /// If a `JoinHandle` is dropped, then the task continues running in the + /// background and its return value is lost. + /// /// # Examples /// /// Creation from [`task::spawn`]: @@ -142,7 +154,7 @@ cfg_rt! { /// [`std::thread::JoinHandle`]: std::thread::JoinHandle /// [`JoinError`]: crate::task::JoinError pub struct JoinHandle<T> { - raw: Option<RawTask>, + raw: RawTask, _p: PhantomData<T>, } } @@ -150,10 +162,13 @@ cfg_rt! { unsafe impl<T: Send> Send for JoinHandle<T> {} unsafe impl<T: Send> Sync for JoinHandle<T> {} +impl<T> UnwindSafe for JoinHandle<T> {} +impl<T> RefUnwindSafe for JoinHandle<T> {} + impl<T> JoinHandle<T> { pub(super) fn new(raw: RawTask) -> JoinHandle<T> { JoinHandle { - raw: Some(raw), + raw, _p: PhantomData, } } @@ -162,39 +177,134 @@ impl<T> JoinHandle<T> { /// /// Awaiting a cancelled task might complete as usual if the task was /// already completed at the time it was cancelled, but most likely it - /// will complete with a `Err(JoinError::Cancelled)`. + /// will fail with a [cancelled] `JoinError`. /// /// ```rust /// use tokio::time; /// - /// #[tokio::main] - /// async fn main() { - /// let mut handles = Vec::new(); - /// - /// handles.push(tokio::spawn(async { - /// time::sleep(time::Duration::from_secs(10)).await; - /// true - /// })); - /// - /// handles.push(tokio::spawn(async { - /// time::sleep(time::Duration::from_secs(10)).await; - /// false - /// })); - /// - /// for handle in &handles { - /// handle.abort(); - /// } - /// - /// for handle in handles { - /// assert!(handle.await.unwrap_err().is_cancelled()); - /// } + /// # #[tokio::main(flavor = "current_thread", start_paused = true)] + /// # async fn main() { + /// let mut handles = Vec::new(); + /// + /// handles.push(tokio::spawn(async { + /// time::sleep(time::Duration::from_secs(10)).await; + /// true + /// })); + /// + /// handles.push(tokio::spawn(async { + /// time::sleep(time::Duration::from_secs(10)).await; + /// false + /// })); + /// + /// for handle in &handles { + /// handle.abort(); + /// } + /// + /// for handle in handles { + /// assert!(handle.await.unwrap_err().is_cancelled()); /// } + /// # } /// ``` + /// [cancelled]: method@super::error::JoinError::is_cancelled pub fn abort(&self) { - if let Some(raw) = self.raw { - raw.remote_abort(); + self.raw.remote_abort(); + } + + /// Checks if the task associated with this `JoinHandle` has finished. + /// + /// Please note that this method can return `false` even if [`abort`] has been + /// called on the task. This is because the cancellation process may take + /// some time, and this method does not return `true` until it has + /// completed. + /// + /// ```rust + /// use tokio::time; + /// + /// # #[tokio::main(flavor = "current_thread", start_paused = true)] + /// # async fn main() { + /// let handle1 = tokio::spawn(async { + /// // do some stuff here + /// }); + /// let handle2 = tokio::spawn(async { + /// // do some other stuff here + /// time::sleep(time::Duration::from_secs(10)).await; + /// }); + /// // Wait for the task to finish + /// handle2.abort(); + /// time::sleep(time::Duration::from_secs(1)).await; + /// assert!(handle1.is_finished()); + /// assert!(handle2.is_finished()); + /// # } + /// ``` + /// [`abort`]: method@JoinHandle::abort + pub fn is_finished(&self) -> bool { + let state = self.raw.header().state.load(); + state.is_complete() + } + + /// Set the waker that is notified when the task completes. + pub(crate) fn set_join_waker(&mut self, waker: &Waker) { + if self.raw.try_set_join_waker(waker) { + // In this case the task has already completed. We wake the waker immediately. + waker.wake_by_ref(); } } + + /// Returns a new `AbortHandle` that can be used to remotely abort this task. + /// + /// Awaiting a task cancelled by the `AbortHandle` might complete as usual if the task was + /// already completed at the time it was cancelled, but most likely it + /// will fail with a [cancelled] `JoinError`. + /// + /// ```rust + /// use tokio::{time, task}; + /// + /// # #[tokio::main(flavor = "current_thread", start_paused = true)] + /// # async fn main() { + /// let mut handles = Vec::new(); + /// + /// handles.push(tokio::spawn(async { + /// time::sleep(time::Duration::from_secs(10)).await; + /// true + /// })); + /// + /// handles.push(tokio::spawn(async { + /// time::sleep(time::Duration::from_secs(10)).await; + /// false + /// })); + /// + /// let abort_handles: Vec<task::AbortHandle> = handles.iter().map(|h| h.abort_handle()).collect(); + /// + /// for handle in abort_handles { + /// handle.abort(); + /// } + /// + /// for handle in handles { + /// assert!(handle.await.unwrap_err().is_cancelled()); + /// } + /// # } + /// ``` + /// [cancelled]: method@super::error::JoinError::is_cancelled + pub fn abort_handle(&self) -> super::AbortHandle { + self.raw.ref_inc(); + super::AbortHandle::new(self.raw) + } + + /// Returns a [task ID] that uniquely identifies this task relative to other + /// currently spawned tasks. + /// + /// **Note**: This is an [unstable API][unstable]. The public API of this type + /// may break in 1.x releases. See [the documentation on unstable + /// features][unstable] for details. + /// + /// [task ID]: crate::task::Id + /// [unstable]: crate#unstable-features + #[cfg(tokio_unstable)] + #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))] + pub fn id(&self) -> super::Id { + // Safety: The header pointer is valid. + unsafe { Header::get_id(self.raw.header_ptr()) } + } } impl<T> Unpin for JoinHandle<T> {} @@ -203,17 +313,11 @@ impl<T> Future for JoinHandle<T> { type Output = super::Result<T>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + ready!(crate::trace::trace_leaf(cx)); let mut ret = Poll::Pending; // Keep track of task budget - let coop = ready!(crate::coop::poll_proceed(cx)); - - // Raw should always be set. If it is not, this is due to polling after - // completion - let raw = self - .raw - .as_ref() - .expect("polling after `JoinHandle` already completed"); + let coop = ready!(crate::runtime::coop::poll_proceed(cx)); // Try to read the task output. If the task is not yet complete, the // waker is stored and is notified once the task does complete. @@ -227,7 +331,8 @@ impl<T> Future for JoinHandle<T> { // // The type of `T` must match the task's output type. unsafe { - raw.try_read_output(&mut ret as *mut _ as *mut (), cx.waker()); + self.raw + .try_read_output(&mut ret as *mut _ as *mut (), cx.waker()); } if ret.is_ready() { @@ -240,13 +345,11 @@ impl<T> Future for JoinHandle<T> { impl<T> Drop for JoinHandle<T> { fn drop(&mut self) { - if let Some(raw) = self.raw.take() { - if raw.header().state.drop_join_handle_fast().is_ok() { - return; - } - - raw.drop_join_handle_slow(); + if self.raw.state().drop_join_handle_fast().is_ok() { + return; } + + self.raw.drop_join_handle_slow(); } } @@ -255,6 +358,9 @@ where T: fmt::Debug, { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - fmt.debug_struct("JoinHandle").finish() + // Safety: The header pointer is valid. + let id_ptr = unsafe { Header::get_id_ptr(self.raw.header_ptr()) }; + let id = unsafe { id_ptr.as_ref() }; + fmt.debug_struct("JoinHandle").field("id", id).finish() } } diff --git a/vendor/tokio/src/runtime/task/list.rs b/vendor/tokio/src/runtime/task/list.rs new file mode 100644 index 000000000..fb7dbdc1d --- /dev/null +++ b/vendor/tokio/src/runtime/task/list.rs @@ -0,0 +1,319 @@ +//! This module has containers for storing the tasks spawned on a scheduler. The +//! `OwnedTasks` container is thread-safe but can only store tasks that +//! implement Send. The `LocalOwnedTasks` container is not thread safe, but can +//! store non-Send tasks. +//! +//! The collections can be closed to prevent adding new tasks during shutdown of +//! the scheduler with the collection. + +use crate::future::Future; +use crate::loom::cell::UnsafeCell; +use crate::loom::sync::Mutex; +use crate::runtime::task::{JoinHandle, LocalNotified, Notified, Schedule, Task}; +use crate::util::linked_list::{CountedLinkedList, Link, LinkedList}; + +use std::marker::PhantomData; + +// The id from the module below is used to verify whether a given task is stored +// in this OwnedTasks, or some other task. The counter starts at one so we can +// use zero for tasks not owned by any list. +// +// The safety checks in this file can technically be violated if the counter is +// overflown, but the checks are not supposed to ever fail unless there is a +// bug in Tokio, so we accept that certain bugs would not be caught if the two +// mixed up runtimes happen to have the same id. + +cfg_has_atomic_u64! { + use std::sync::atomic::{AtomicU64, Ordering}; + + static NEXT_OWNED_TASKS_ID: AtomicU64 = AtomicU64::new(1); + + fn get_next_id() -> u64 { + loop { + let id = NEXT_OWNED_TASKS_ID.fetch_add(1, Ordering::Relaxed); + if id != 0 { + return id; + } + } + } +} + +cfg_not_has_atomic_u64! { + use std::sync::atomic::{AtomicU32, Ordering}; + + static NEXT_OWNED_TASKS_ID: AtomicU32 = AtomicU32::new(1); + + fn get_next_id() -> u64 { + loop { + let id = NEXT_OWNED_TASKS_ID.fetch_add(1, Ordering::Relaxed); + if id != 0 { + return u64::from(id); + } + } + } +} + +pub(crate) struct OwnedTasks<S: 'static> { + inner: Mutex<CountedOwnedTasksInner<S>>, + id: u64, +} +struct CountedOwnedTasksInner<S: 'static> { + list: CountedLinkedList<Task<S>, <Task<S> as Link>::Target>, + closed: bool, +} +pub(crate) struct LocalOwnedTasks<S: 'static> { + inner: UnsafeCell<OwnedTasksInner<S>>, + id: u64, + _not_send_or_sync: PhantomData<*const ()>, +} +struct OwnedTasksInner<S: 'static> { + list: LinkedList<Task<S>, <Task<S> as Link>::Target>, + closed: bool, +} + +impl<S: 'static> OwnedTasks<S> { + pub(crate) fn new() -> Self { + Self { + inner: Mutex::new(CountedOwnedTasksInner { + list: CountedLinkedList::new(), + closed: false, + }), + id: get_next_id(), + } + } + + /// Binds the provided task to this OwnedTasks instance. This fails if the + /// OwnedTasks has been closed. + pub(crate) fn bind<T>( + &self, + task: T, + scheduler: S, + id: super::Id, + ) -> (JoinHandle<T::Output>, Option<Notified<S>>) + where + S: Schedule, + T: Future + Send + 'static, + T::Output: Send + 'static, + { + let (task, notified, join) = super::new_task(task, scheduler, id); + + unsafe { + // safety: We just created the task, so we have exclusive access + // to the field. + task.header().set_owner_id(self.id); + } + + let mut lock = self.inner.lock(); + if lock.closed { + drop(lock); + drop(notified); + task.shutdown(); + (join, None) + } else { + lock.list.push_front(task); + (join, Some(notified)) + } + } + + /// Asserts that the given task is owned by this OwnedTasks and convert it to + /// a LocalNotified, giving the thread permission to poll this task. + #[inline] + pub(crate) fn assert_owner(&self, task: Notified<S>) -> LocalNotified<S> { + assert_eq!(task.header().get_owner_id(), self.id); + + // safety: All tasks bound to this OwnedTasks are Send, so it is safe + // to poll it on this thread no matter what thread we are on. + LocalNotified { + task: task.0, + _not_send: PhantomData, + } + } + + /// Shuts down all tasks in the collection. This call also closes the + /// collection, preventing new items from being added. + pub(crate) fn close_and_shutdown_all(&self) + where + S: Schedule, + { + // The first iteration of the loop was unrolled so it can set the + // closed bool. + let first_task = { + let mut lock = self.inner.lock(); + lock.closed = true; + lock.list.pop_back() + }; + match first_task { + Some(task) => task.shutdown(), + None => return, + } + + loop { + let task = match self.inner.lock().list.pop_back() { + Some(task) => task, + None => return, + }; + + task.shutdown(); + } + } + + pub(crate) fn active_tasks_count(&self) -> usize { + self.inner.lock().list.count() + } + + pub(crate) fn remove(&self, task: &Task<S>) -> Option<Task<S>> { + let task_id = task.header().get_owner_id(); + if task_id == 0 { + // The task is unowned. + return None; + } + + assert_eq!(task_id, self.id); + + // safety: We just checked that the provided task is not in some other + // linked list. + unsafe { self.inner.lock().list.remove(task.header_ptr()) } + } + + pub(crate) fn is_empty(&self) -> bool { + self.inner.lock().list.is_empty() + } +} + +cfg_taskdump! { + impl<S: 'static> OwnedTasks<S> { + /// Locks the tasks, and calls `f` on an iterator over them. + pub(crate) fn for_each<F>(&self, f: F) + where + F: FnMut(&Task<S>) + { + self.inner.lock().list.for_each(f) + } + } +} + +impl<S: 'static> LocalOwnedTasks<S> { + pub(crate) fn new() -> Self { + Self { + inner: UnsafeCell::new(OwnedTasksInner { + list: LinkedList::new(), + closed: false, + }), + id: get_next_id(), + _not_send_or_sync: PhantomData, + } + } + + pub(crate) fn bind<T>( + &self, + task: T, + scheduler: S, + id: super::Id, + ) -> (JoinHandle<T::Output>, Option<Notified<S>>) + where + S: Schedule, + T: Future + 'static, + T::Output: 'static, + { + let (task, notified, join) = super::new_task(task, scheduler, id); + + unsafe { + // safety: We just created the task, so we have exclusive access + // to the field. + task.header().set_owner_id(self.id); + } + + if self.is_closed() { + drop(notified); + task.shutdown(); + (join, None) + } else { + self.with_inner(|inner| { + inner.list.push_front(task); + }); + (join, Some(notified)) + } + } + + /// Shuts down all tasks in the collection. This call also closes the + /// collection, preventing new items from being added. + pub(crate) fn close_and_shutdown_all(&self) + where + S: Schedule, + { + self.with_inner(|inner| inner.closed = true); + + while let Some(task) = self.with_inner(|inner| inner.list.pop_back()) { + task.shutdown(); + } + } + + pub(crate) fn remove(&self, task: &Task<S>) -> Option<Task<S>> { + let task_id = task.header().get_owner_id(); + if task_id == 0 { + // The task is unowned. + return None; + } + + assert_eq!(task_id, self.id); + + self.with_inner(|inner| + // safety: We just checked that the provided task is not in some + // other linked list. + unsafe { inner.list.remove(task.header_ptr()) }) + } + + /// Asserts that the given task is owned by this LocalOwnedTasks and convert + /// it to a LocalNotified, giving the thread permission to poll this task. + #[inline] + pub(crate) fn assert_owner(&self, task: Notified<S>) -> LocalNotified<S> { + assert_eq!(task.header().get_owner_id(), self.id); + + // safety: The task was bound to this LocalOwnedTasks, and the + // LocalOwnedTasks is not Send or Sync, so we are on the right thread + // for polling this task. + LocalNotified { + task: task.0, + _not_send: PhantomData, + } + } + + #[inline] + fn with_inner<F, T>(&self, f: F) -> T + where + F: FnOnce(&mut OwnedTasksInner<S>) -> T, + { + // safety: This type is not Sync, so concurrent calls of this method + // can't happen. Furthermore, all uses of this method in this file make + // sure that they don't call `with_inner` recursively. + self.inner.with_mut(|ptr| unsafe { f(&mut *ptr) }) + } + + pub(crate) fn is_closed(&self) -> bool { + self.with_inner(|inner| inner.closed) + } + + pub(crate) fn is_empty(&self) -> bool { + self.with_inner(|inner| inner.list.is_empty()) + } +} + +#[cfg(all(test))] +mod tests { + use super::*; + + // This test may run in parallel with other tests, so we only test that ids + // come in increasing order. + #[test] + fn test_id_not_broken() { + let mut last_id = get_next_id(); + assert_ne!(last_id, 0); + + for _ in 0..1000 { + let next_id = get_next_id(); + assert_ne!(next_id, 0); + assert!(last_id < next_id); + last_id = next_id; + } + } +} diff --git a/vendor/tokio/src/runtime/task/mod.rs b/vendor/tokio/src/runtime/task/mod.rs index 58b8c2a15..932552fb9 100644 --- a/vendor/tokio/src/runtime/task/mod.rs +++ b/vendor/tokio/src/runtime/task/mod.rs @@ -1,29 +1,209 @@ +//! The task module. +//! +//! The task module contains the code that manages spawned tasks and provides a +//! safe API for the rest of the runtime to use. Each task in a runtime is +//! stored in an OwnedTasks or LocalOwnedTasks object. +//! +//! # Task reference types +//! +//! A task is usually referenced by multiple handles, and there are several +//! types of handles. +//! +//! * OwnedTask - tasks stored in an OwnedTasks or LocalOwnedTasks are of this +//! reference type. +//! +//! * JoinHandle - each task has a JoinHandle that allows access to the output +//! of the task. +//! +//! * Waker - every waker for a task has this reference type. There can be any +//! number of waker references. +//! +//! * Notified - tracks whether the task is notified. +//! +//! * Unowned - this task reference type is used for tasks not stored in any +//! runtime. Mainly used for blocking tasks, but also in tests. +//! +//! The task uses a reference count to keep track of how many active references +//! exist. The Unowned reference type takes up two ref-counts. All other +//! reference types take up a single ref-count. +//! +//! Besides the waker type, each task has at most one of each reference type. +//! +//! # State +//! +//! The task stores its state in an atomic usize with various bitfields for the +//! necessary information. The state has the following bitfields: +//! +//! * RUNNING - Tracks whether the task is currently being polled or cancelled. +//! This bit functions as a lock around the task. +//! +//! * COMPLETE - Is one once the future has fully completed and has been +//! dropped. Never unset once set. Never set together with RUNNING. +//! +//! * NOTIFIED - Tracks whether a Notified object currently exists. +//! +//! * CANCELLED - Is set to one for tasks that should be cancelled as soon as +//! possible. May take any value for completed tasks. +//! +//! * JOIN_INTEREST - Is set to one if there exists a JoinHandle. +//! +//! * JOIN_WAKER - Acts as an access control bit for the join handle waker. The +//! protocol for its usage is described below. +//! +//! The rest of the bits are used for the ref-count. +//! +//! # Fields in the task +//! +//! The task has various fields. This section describes how and when it is safe +//! to access a field. +//! +//! * The state field is accessed with atomic instructions. +//! +//! * The OwnedTask reference has exclusive access to the `owned` field. +//! +//! * The Notified reference has exclusive access to the `queue_next` field. +//! +//! * The `owner_id` field can be set as part of construction of the task, but +//! is otherwise immutable and anyone can access the field immutably without +//! synchronization. +//! +//! * If COMPLETE is one, then the JoinHandle has exclusive access to the +//! stage field. If COMPLETE is zero, then the RUNNING bitfield functions as +//! a lock for the stage field, and it can be accessed only by the thread +//! that set RUNNING to one. +//! +//! * The waker field may be concurrently accessed by different threads: in one +//! thread the runtime may complete a task and *read* the waker field to +//! invoke the waker, and in another thread the task's JoinHandle may be +//! polled, and if the task hasn't yet completed, the JoinHandle may *write* +//! a waker to the waker field. The JOIN_WAKER bit ensures safe access by +//! multiple threads to the waker field using the following rules: +//! +//! 1. JOIN_WAKER is initialized to zero. +//! +//! 2. If JOIN_WAKER is zero, then the JoinHandle has exclusive (mutable) +//! access to the waker field. +//! +//! 3. If JOIN_WAKER is one, then the JoinHandle has shared (read-only) +//! access to the waker field. +//! +//! 4. If JOIN_WAKER is one and COMPLETE is one, then the runtime has shared +//! (read-only) access to the waker field. +//! +//! 5. If the JoinHandle needs to write to the waker field, then the +//! JoinHandle needs to (i) successfully set JOIN_WAKER to zero if it is +//! not already zero to gain exclusive access to the waker field per rule +//! 2, (ii) write a waker, and (iii) successfully set JOIN_WAKER to one. +//! +//! 6. The JoinHandle can change JOIN_WAKER only if COMPLETE is zero (i.e. +//! the task hasn't yet completed). +//! +//! Rule 6 implies that the steps (i) or (iii) of rule 5 may fail due to a +//! race. If step (i) fails, then the attempt to write a waker is aborted. If +//! step (iii) fails because COMPLETE is set to one by another thread after +//! step (i), then the waker field is cleared. Once COMPLETE is one (i.e. +//! task has completed), the JoinHandle will not modify JOIN_WAKER. After the +//! runtime sets COMPLETE to one, it invokes the waker if there is one. +//! +//! All other fields are immutable and can be accessed immutably without +//! synchronization by anyone. +//! +//! # Safety +//! +//! This section goes through various situations and explains why the API is +//! safe in that situation. +//! +//! ## Polling or dropping the future +//! +//! Any mutable access to the future happens after obtaining a lock by modifying +//! the RUNNING field, so exclusive access is ensured. +//! +//! When the task completes, exclusive access to the output is transferred to +//! the JoinHandle. If the JoinHandle is already dropped when the transition to +//! complete happens, the thread performing that transition retains exclusive +//! access to the output and should immediately drop it. +//! +//! ## Non-Send futures +//! +//! If a future is not Send, then it is bound to a LocalOwnedTasks. The future +//! will only ever be polled or dropped given a LocalNotified or inside a call +//! to LocalOwnedTasks::shutdown_all. In either case, it is guaranteed that the +//! future is on the right thread. +//! +//! If the task is never removed from the LocalOwnedTasks, then it is leaked, so +//! there is no risk that the task is dropped on some other thread when the last +//! ref-count drops. +//! +//! ## Non-Send output +//! +//! When a task completes, the output is placed in the stage of the task. Then, +//! a transition that sets COMPLETE to true is performed, and the value of +//! JOIN_INTEREST when this transition happens is read. +//! +//! If JOIN_INTEREST is zero when the transition to COMPLETE happens, then the +//! output is immediately dropped. +//! +//! If JOIN_INTEREST is one when the transition to COMPLETE happens, then the +//! JoinHandle is responsible for cleaning up the output. If the output is not +//! Send, then this happens: +//! +//! 1. The output is created on the thread that the future was polled on. Since +//! only non-Send futures can have non-Send output, the future was polled on +//! the thread that the future was spawned from. +//! 2. Since `JoinHandle<Output>` is not Send if Output is not Send, the +//! JoinHandle is also on the thread that the future was spawned from. +//! 3. Thus, the JoinHandle will not move the output across threads when it +//! takes or drops the output. +//! +//! ## Recursive poll/shutdown +//! +//! Calling poll from inside a shutdown call or vice-versa is not prevented by +//! the API exposed by the task module, so this has to be safe. In either case, +//! the lock in the RUNNING bitfield makes the inner call return immediately. If +//! the inner call is a `shutdown` call, then the CANCELLED bit is set, and the +//! poll call will notice it when the poll finishes, and the task is cancelled +//! at that point. + +// Some task infrastructure is here to support `JoinSet`, which is currently +// unstable. This should be removed once `JoinSet` is stabilized. +#![cfg_attr(not(tokio_unstable), allow(dead_code))] + mod core; use self::core::Cell; -pub(crate) use self::core::Header; +use self::core::Header; mod error; -#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 pub use self::error::JoinError; mod harness; use self::harness::Harness; +mod id; +#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))] +pub use id::{id, try_id, Id}; + +#[cfg(feature = "rt")] +mod abort; mod join; -#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 + +#[cfg(feature = "rt")] +pub use self::abort::AbortHandle; + pub use self::join::JoinHandle; +mod list; +pub(crate) use self::list::{LocalOwnedTasks, OwnedTasks}; + mod raw; -use self::raw::RawTask; +pub(crate) use self::raw::RawTask; mod state; use self::state::State; mod waker; -cfg_rt_multi_thread! { - mod stack; - pub(crate) use self::stack::TransferStack; +cfg_taskdump! { + pub(crate) mod trace; } use crate::future::Future; @@ -33,7 +213,7 @@ use std::marker::PhantomData; use std::ptr::NonNull; use std::{fmt, mem}; -/// An owned handle to the task, tracked by ref count +/// An owned handle to the task, tracked by ref count. #[repr(transparent)] pub(crate) struct Task<S: 'static> { raw: RawTask, @@ -43,30 +223,43 @@ pub(crate) struct Task<S: 'static> { unsafe impl<S> Send for Task<S> {} unsafe impl<S> Sync for Task<S> {} -/// A task was notified +/// A task was notified. #[repr(transparent)] pub(crate) struct Notified<S: 'static>(Task<S>); +// safety: This type cannot be used to touch the task without first verifying +// that the value is on a thread where it is safe to poll the task. unsafe impl<S: Schedule> Send for Notified<S> {} unsafe impl<S: Schedule> Sync for Notified<S> {} -/// Task result sent back +/// A non-Send variant of Notified with the invariant that it is on a thread +/// where it is safe to poll it. +#[repr(transparent)] +pub(crate) struct LocalNotified<S: 'static> { + task: Task<S>, + _not_send: PhantomData<*const ()>, +} + +/// A task that is not owned by any OwnedTasks. Used for blocking tasks. +/// This type holds two ref-counts. +pub(crate) struct UnownedTask<S: 'static> { + raw: RawTask, + _p: PhantomData<S>, +} + +// safety: This type can only be created given a Send task. +unsafe impl<S> Send for UnownedTask<S> {} +unsafe impl<S> Sync for UnownedTask<S> {} + +/// Task result sent back. pub(crate) type Result<T> = std::result::Result<T, JoinError>; pub(crate) trait Schedule: Sync + Sized + 'static { - /// Bind a task to the executor. - /// - /// Guaranteed to be called from the thread that called `poll` on the task. - /// The returned `Schedule` instance is associated with the task and is used - /// as `&self` in the other methods on this trait. - fn bind(task: Task<Self>) -> Self; - /// The task has completed work and is ready to be released. The scheduler - /// is free to drop it whenever. + /// should release it immediately and return it. The task module will batch + /// the ref-dec with setting other options. /// - /// If the scheduler can immediately release the task, it should return - /// it as part of the function. This enables the task module to batch - /// the ref-dec with other options. + /// If the scheduler has already released the task, then None is returned. fn release(&self, task: &Task<Self>) -> Option<Task<Self>>; /// Schedule the task @@ -77,104 +270,177 @@ pub(crate) trait Schedule: Sync + Sized + 'static { fn yield_now(&self, task: Notified<Self>) { self.schedule(task); } + + /// Polling the task resulted in a panic. Should the runtime shutdown? + fn unhandled_panic(&self) { + // By default, do nothing. This maintains the 1.0 behavior. + } } cfg_rt! { - /// Create a new task with an associated join handle - pub(crate) fn joinable<T, S>(task: T) -> (Notified<S>, JoinHandle<T::Output>) + /// This is the constructor for a new task. Three references to the task are + /// created. The first task reference is usually put into an OwnedTasks + /// immediately. The Notified is sent to the scheduler as an ordinary + /// notification. + fn new_task<T, S>( + task: T, + scheduler: S, + id: Id, + ) -> (Task<S>, Notified<S>, JoinHandle<T::Output>) where - T: Future + Send + 'static, S: Schedule, + T: Future + 'static, + T::Output: 'static, { - let raw = RawTask::new::<_, S>(task); - + let raw = RawTask::new::<T, S>(task, scheduler, id); let task = Task { raw, _p: PhantomData, }; - + let notified = Notified(Task { + raw, + _p: PhantomData, + }); let join = JoinHandle::new(raw); - (Notified(task), join) + (task, notified, join) } -} -cfg_rt! { - /// Create a new `!Send` task with an associated join handle - pub(crate) unsafe fn joinable_local<T, S>(task: T) -> (Notified<S>, JoinHandle<T::Output>) + /// Creates a new task with an associated join handle. This method is used + /// only when the task is not going to be stored in an `OwnedTasks` list. + /// + /// Currently only blocking tasks use this method. + pub(crate) fn unowned<T, S>(task: T, scheduler: S, id: Id) -> (UnownedTask<S>, JoinHandle<T::Output>) where - T: Future + 'static, S: Schedule, + T: Send + Future + 'static, + T::Output: Send + 'static, { - let raw = RawTask::new::<_, S>(task); + let (task, notified, join) = new_task(task, scheduler, id); - let task = Task { - raw, + // This transfers the ref-count of task and notified into an UnownedTask. + // This is valid because an UnownedTask holds two ref-counts. + let unowned = UnownedTask { + raw: task.raw, _p: PhantomData, }; + std::mem::forget(task); + std::mem::forget(notified); - let join = JoinHandle::new(raw); - - (Notified(task), join) + (unowned, join) } } impl<S: 'static> Task<S> { - pub(crate) unsafe fn from_raw(ptr: NonNull<Header>) -> Task<S> { + unsafe fn new(raw: RawTask) -> Task<S> { Task { - raw: RawTask::from_raw(ptr), + raw, _p: PhantomData, } } - pub(crate) fn header(&self) -> &Header { + unsafe fn from_raw(ptr: NonNull<Header>) -> Task<S> { + Task::new(RawTask::from_raw(ptr)) + } + + #[cfg(all( + tokio_unstable, + tokio_taskdump, + feature = "rt", + target_os = "linux", + any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64") + ))] + pub(super) fn as_raw(&self) -> RawTask { + self.raw + } + + fn header(&self) -> &Header { self.raw.header() } -} -cfg_rt_multi_thread! { - impl<S: 'static> Notified<S> { - pub(crate) unsafe fn from_raw(ptr: NonNull<Header>) -> Notified<S> { - Notified(Task::from_raw(ptr)) - } + fn header_ptr(&self) -> NonNull<Header> { + self.raw.header_ptr() + } +} - pub(crate) fn header(&self) -> &Header { - self.0.header() - } +impl<S: 'static> Notified<S> { + fn header(&self) -> &Header { + self.0.header() } +} - impl<S: 'static> Task<S> { - pub(crate) fn into_raw(self) -> NonNull<Header> { - let ret = self.header().into(); - mem::forget(self); - ret - } +impl<S: 'static> Notified<S> { + pub(crate) unsafe fn from_raw(ptr: RawTask) -> Notified<S> { + Notified(Task::new(ptr)) } +} - impl<S: 'static> Notified<S> { - pub(crate) fn into_raw(self) -> NonNull<Header> { - self.0.into_raw() - } +impl<S: 'static> Notified<S> { + pub(crate) fn into_raw(self) -> RawTask { + let raw = self.0.raw; + mem::forget(self); + raw } } impl<S: Schedule> Task<S> { - /// Pre-emptively cancel the task as part of the shutdown process. - pub(crate) fn shutdown(&self) { - self.raw.shutdown(); + /// Preemptively cancels the task as part of the shutdown process. + pub(crate) fn shutdown(self) { + let raw = self.raw; + mem::forget(self); + raw.shutdown(); } } -impl<S: Schedule> Notified<S> { - /// Run the task +impl<S: Schedule> LocalNotified<S> { + /// Runs the task. + pub(crate) fn run(self) { + let raw = self.task.raw; + mem::forget(self); + raw.poll(); + } +} + +impl<S: Schedule> UnownedTask<S> { + // Used in test of the inject queue. + #[cfg(test)] + #[cfg_attr(tokio_wasm, allow(dead_code))] + pub(super) fn into_notified(self) -> Notified<S> { + Notified(self.into_task()) + } + + fn into_task(self) -> Task<S> { + // Convert into a task. + let task = Task { + raw: self.raw, + _p: PhantomData, + }; + mem::forget(self); + + // Drop a ref-count since an UnownedTask holds two. + task.header().state.ref_dec(); + + task + } + pub(crate) fn run(self) { - self.0.raw.poll(); + let raw = self.raw; mem::forget(self); + + // Transfer one ref-count to a Task object. + let task = Task::<S> { + raw, + _p: PhantomData, + }; + + // Use the other ref-count to poll the task. + raw.poll(); + // Decrement our extra ref-count + drop(task); } - /// Pre-emptively cancel the task as part of the shutdown process. pub(crate) fn shutdown(self) { - self.0.shutdown(); + self.into_task().shutdown() } } @@ -188,6 +454,16 @@ impl<S: 'static> Drop for Task<S> { } } +impl<S: 'static> Drop for UnownedTask<S> { + fn drop(&mut self) { + // Decrement the ref count + if self.raw.header().state.ref_dec_twice() { + // Deallocate if this is the final ref count + self.raw.dealloc(); + } + } +} + impl<S> fmt::Debug for Task<S> { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { write!(fmt, "Task({:p})", self.header()) @@ -202,13 +478,13 @@ impl<S> fmt::Debug for Notified<S> { /// # Safety /// -/// Tasks are pinned +/// Tasks are pinned. unsafe impl<S> linked_list::Link for Task<S> { type Handle = Task<S>; type Target = Header; fn as_raw(handle: &Task<S>) -> NonNull<Header> { - handle.header().into() + handle.raw.header_ptr() } unsafe fn from_raw(ptr: NonNull<Header>) -> Task<S> { @@ -216,7 +492,6 @@ unsafe impl<S> linked_list::Link for Task<S> { } unsafe fn pointers(target: NonNull<Header>) -> NonNull<linked_list::Pointers<Header>> { - // Not super great as it avoids some of looms checking... - NonNull::from(target.as_ref().owned.with_mut(|ptr| &mut *ptr)) + self::core::Trailer::addr_of_owned(Header::get_trailer(target)) } } diff --git a/vendor/tokio/src/runtime/task/raw.rs b/vendor/tokio/src/runtime/task/raw.rs index 56d65d5a6..807885928 100644 --- a/vendor/tokio/src/runtime/task/raw.rs +++ b/vendor/tokio/src/runtime/task/raw.rs @@ -1,53 +1,167 @@ use crate::future::Future; -use crate::runtime::task::{Cell, Harness, Header, Schedule, State}; +use crate::runtime::task::core::{Core, Trailer}; +use crate::runtime::task::{Cell, Harness, Header, Id, Schedule, State}; use std::ptr::NonNull; use std::task::{Poll, Waker}; /// Raw task handle -pub(super) struct RawTask { +pub(crate) struct RawTask { ptr: NonNull<Header>, } pub(super) struct Vtable { - /// Poll the future + /// Polls the future. pub(super) poll: unsafe fn(NonNull<Header>), - /// Deallocate the memory + /// Schedules the task for execution on the runtime. + pub(super) schedule: unsafe fn(NonNull<Header>), + + /// Deallocates the memory. pub(super) dealloc: unsafe fn(NonNull<Header>), - /// Read the task output, if complete + /// Reads the task output, if complete. pub(super) try_read_output: unsafe fn(NonNull<Header>, *mut (), &Waker), - /// The join handle has been dropped + /// The join handle has been dropped. pub(super) drop_join_handle_slow: unsafe fn(NonNull<Header>), - /// The task is remotely aborted - pub(super) remote_abort: unsafe fn(NonNull<Header>), + /// An abort handle has been dropped. + pub(super) drop_abort_handle: unsafe fn(NonNull<Header>), - /// Scheduler is being shutdown + /// Scheduler is being shutdown. pub(super) shutdown: unsafe fn(NonNull<Header>), + + /// The number of bytes that the `trailer` field is offset from the header. + pub(super) trailer_offset: usize, + + /// The number of bytes that the `scheduler` field is offset from the header. + pub(super) scheduler_offset: usize, + + /// The number of bytes that the `id` field is offset from the header. + pub(super) id_offset: usize, } /// Get the vtable for the requested `T` and `S` generics. pub(super) fn vtable<T: Future, S: Schedule>() -> &'static Vtable { &Vtable { poll: poll::<T, S>, + schedule: schedule::<S>, dealloc: dealloc::<T, S>, try_read_output: try_read_output::<T, S>, drop_join_handle_slow: drop_join_handle_slow::<T, S>, - remote_abort: remote_abort::<T, S>, + drop_abort_handle: drop_abort_handle::<T, S>, shutdown: shutdown::<T, S>, + trailer_offset: OffsetHelper::<T, S>::TRAILER_OFFSET, + scheduler_offset: OffsetHelper::<T, S>::SCHEDULER_OFFSET, + id_offset: OffsetHelper::<T, S>::ID_OFFSET, } } +/// Calling `get_trailer_offset` directly in vtable doesn't work because it +/// prevents the vtable from being promoted to a static reference. +/// +/// See this thread for more info: +/// <https://users.rust-lang.org/t/custom-vtables-with-integers/78508> +struct OffsetHelper<T, S>(T, S); +impl<T: Future, S: Schedule> OffsetHelper<T, S> { + // Pass `size_of`/`align_of` as arguments rather than calling them directly + // inside `get_trailer_offset` because trait bounds on generic parameters + // of const fn are unstable on our MSRV. + const TRAILER_OFFSET: usize = get_trailer_offset( + std::mem::size_of::<Header>(), + std::mem::size_of::<Core<T, S>>(), + std::mem::align_of::<Core<T, S>>(), + std::mem::align_of::<Trailer>(), + ); + + // The `scheduler` is the first field of `Core`, so it has the same + // offset as `Core`. + const SCHEDULER_OFFSET: usize = get_core_offset( + std::mem::size_of::<Header>(), + std::mem::align_of::<Core<T, S>>(), + ); + + const ID_OFFSET: usize = get_id_offset( + std::mem::size_of::<Header>(), + std::mem::align_of::<Core<T, S>>(), + std::mem::size_of::<S>(), + std::mem::align_of::<Id>(), + ); +} + +/// Compute the offset of the `Trailer` field in `Cell<T, S>` using the +/// `#[repr(C)]` algorithm. +/// +/// Pseudo-code for the `#[repr(C)]` algorithm can be found here: +/// <https://doc.rust-lang.org/reference/type-layout.html#reprc-structs> +const fn get_trailer_offset( + header_size: usize, + core_size: usize, + core_align: usize, + trailer_align: usize, +) -> usize { + let mut offset = header_size; + + let core_misalign = offset % core_align; + if core_misalign > 0 { + offset += core_align - core_misalign; + } + offset += core_size; + + let trailer_misalign = offset % trailer_align; + if trailer_misalign > 0 { + offset += trailer_align - trailer_misalign; + } + + offset +} + +/// Compute the offset of the `Core<T, S>` field in `Cell<T, S>` using the +/// `#[repr(C)]` algorithm. +/// +/// Pseudo-code for the `#[repr(C)]` algorithm can be found here: +/// <https://doc.rust-lang.org/reference/type-layout.html#reprc-structs> +const fn get_core_offset(header_size: usize, core_align: usize) -> usize { + let mut offset = header_size; + + let core_misalign = offset % core_align; + if core_misalign > 0 { + offset += core_align - core_misalign; + } + + offset +} + +/// Compute the offset of the `Id` field in `Cell<T, S>` using the +/// `#[repr(C)]` algorithm. +/// +/// Pseudo-code for the `#[repr(C)]` algorithm can be found here: +/// <https://doc.rust-lang.org/reference/type-layout.html#reprc-structs> +const fn get_id_offset( + header_size: usize, + core_align: usize, + scheduler_size: usize, + id_align: usize, +) -> usize { + let mut offset = get_core_offset(header_size, core_align); + offset += scheduler_size; + + let id_misalign = offset % id_align; + if id_misalign > 0 { + offset += id_align - id_misalign; + } + + offset +} + impl RawTask { - pub(super) fn new<T, S>(task: T) -> RawTask + pub(super) fn new<T, S>(task: T, scheduler: S, id: Id) -> RawTask where T: Future, S: Schedule, { - let ptr = Box::into_raw(Cell::<_, S>::new(task, State::new())); + let ptr = Box::into_raw(Cell::<_, S>::new(task, scheduler, State::new(), id)); let ptr = unsafe { NonNull::new_unchecked(ptr as *mut Header) }; RawTask { ptr } @@ -57,19 +171,40 @@ impl RawTask { RawTask { ptr } } - /// Returns a reference to the task's meta structure. - /// - /// Safe as `Header` is `Sync`. + pub(super) fn header_ptr(&self) -> NonNull<Header> { + self.ptr + } + + pub(super) fn trailer_ptr(&self) -> NonNull<Trailer> { + unsafe { Header::get_trailer(self.ptr) } + } + + /// Returns a reference to the task's header. pub(super) fn header(&self) -> &Header { unsafe { self.ptr.as_ref() } } + /// Returns a reference to the task's trailer. + pub(super) fn trailer(&self) -> &Trailer { + unsafe { &*self.trailer_ptr().as_ptr() } + } + + /// Returns a reference to the task's state. + pub(super) fn state(&self) -> &State { + &self.header().state + } + /// Safety: mutual exclusion is required to call this function. - pub(super) fn poll(self) { + pub(crate) fn poll(self) { let vtable = self.header().vtable; unsafe { (vtable.poll)(self.ptr) } } + pub(super) fn schedule(self) { + let vtable = self.header().vtable; + unsafe { (vtable.schedule)(self.ptr) } + } + pub(super) fn dealloc(self) { let vtable = self.header().vtable; unsafe { @@ -89,14 +224,42 @@ impl RawTask { unsafe { (vtable.drop_join_handle_slow)(self.ptr) } } + pub(super) fn drop_abort_handle(self) { + let vtable = self.header().vtable; + unsafe { (vtable.drop_abort_handle)(self.ptr) } + } + pub(super) fn shutdown(self) { let vtable = self.header().vtable; unsafe { (vtable.shutdown)(self.ptr) } } - pub(super) fn remote_abort(self) { - let vtable = self.header().vtable; - unsafe { (vtable.remote_abort)(self.ptr) } + /// Increment the task's reference count. + /// + /// Currently, this is used only when creating an `AbortHandle`. + pub(super) fn ref_inc(self) { + self.header().state.ref_inc(); + } + + /// Get the queue-next pointer + /// + /// This is for usage by the injection queue + /// + /// Safety: make sure only one queue uses this and access is synchronized. + pub(crate) unsafe fn get_queue_next(self) -> Option<RawTask> { + self.header() + .queue_next + .with(|ptr| *ptr) + .map(|p| RawTask::from_raw(p)) + } + + /// Sets the queue-next pointer + /// + /// This is for usage by the injection queue + /// + /// Safety: make sure only one queue uses this and access is synchronized. + pub(crate) unsafe fn set_queue_next(self, val: Option<RawTask>) { + self.header().set_next(val.map(|task| task.ptr)); } } @@ -113,6 +276,15 @@ unsafe fn poll<T: Future, S: Schedule>(ptr: NonNull<Header>) { harness.poll(); } +unsafe fn schedule<S: Schedule>(ptr: NonNull<Header>) { + use crate::runtime::task::{Notified, Task}; + + let scheduler = Header::get_scheduler::<S>(ptr); + scheduler + .as_ref() + .schedule(Notified(Task::from_raw(ptr.cast()))); +} + unsafe fn dealloc<T: Future, S: Schedule>(ptr: NonNull<Header>) { let harness = Harness::<T, S>::from_raw(ptr); harness.dealloc(); @@ -134,9 +306,9 @@ unsafe fn drop_join_handle_slow<T: Future, S: Schedule>(ptr: NonNull<Header>) { harness.drop_join_handle_slow() } -unsafe fn remote_abort<T: Future, S: Schedule>(ptr: NonNull<Header>) { +unsafe fn drop_abort_handle<T: Future, S: Schedule>(ptr: NonNull<Header>) { let harness = Harness::<T, S>::from_raw(ptr); - harness.remote_abort() + harness.drop_reference(); } unsafe fn shutdown<T: Future, S: Schedule>(ptr: NonNull<Header>) { diff --git a/vendor/tokio/src/runtime/task/stack.rs b/vendor/tokio/src/runtime/task/stack.rs deleted file mode 100644 index 9dd8d3f43..000000000 --- a/vendor/tokio/src/runtime/task/stack.rs +++ /dev/null @@ -1,83 +0,0 @@ -use crate::loom::sync::atomic::AtomicPtr; -use crate::runtime::task::{Header, Task}; - -use std::marker::PhantomData; -use std::ptr::{self, NonNull}; -use std::sync::atomic::Ordering::{Acquire, Relaxed, Release}; - -/// Concurrent stack of tasks, used to pass ownership of a task from one worker -/// to another. -pub(crate) struct TransferStack<T: 'static> { - head: AtomicPtr<Header>, - _p: PhantomData<T>, -} - -impl<T: 'static> TransferStack<T> { - pub(crate) fn new() -> TransferStack<T> { - TransferStack { - head: AtomicPtr::new(ptr::null_mut()), - _p: PhantomData, - } - } - - pub(crate) fn push(&self, task: Task<T>) { - let task = task.into_raw(); - - // We don't care about any memory associated w/ setting the `head` - // field, just the current value. - // - // The compare-exchange creates a release sequence. - let mut curr = self.head.load(Relaxed); - - loop { - unsafe { - task.as_ref() - .stack_next - .with_mut(|ptr| *ptr = NonNull::new(curr)) - }; - - let res = self - .head - .compare_exchange(curr, task.as_ptr() as *mut _, Release, Relaxed); - - match res { - Ok(_) => return, - Err(actual) => { - curr = actual; - } - } - } - } - - pub(crate) fn drain(&self) -> impl Iterator<Item = Task<T>> { - struct Iter<T: 'static>(Option<NonNull<Header>>, PhantomData<T>); - - impl<T: 'static> Iterator for Iter<T> { - type Item = Task<T>; - - fn next(&mut self) -> Option<Task<T>> { - let task = self.0?; - - // Move the cursor forward - self.0 = unsafe { task.as_ref().stack_next.with(|ptr| *ptr) }; - - // Return the task - unsafe { Some(Task::from_raw(task)) } - } - } - - impl<T: 'static> Drop for Iter<T> { - fn drop(&mut self) { - use std::process; - - if self.0.is_some() { - // we have bugs - process::abort(); - } - } - } - - let ptr = self.head.swap(ptr::null_mut(), Acquire); - Iter(NonNull::new(ptr), PhantomData) - } -} diff --git a/vendor/tokio/src/runtime/task/state.rs b/vendor/tokio/src/runtime/task/state.rs index 603772162..12f544918 100644 --- a/vendor/tokio/src/runtime/task/state.rs +++ b/vendor/tokio/src/runtime/task/state.rs @@ -8,7 +8,7 @@ pub(super) struct State { val: AtomicUsize, } -/// Current state value +/// Current state value. #[derive(Copy, Clone)] pub(super) struct Snapshot(usize); @@ -19,20 +19,20 @@ const RUNNING: usize = 0b0001; /// The task is complete. /// -/// Once this bit is set, it is never unset +/// Once this bit is set, it is never unset. const COMPLETE: usize = 0b0010; -/// Extracts the task's lifecycle value from the state +/// Extracts the task's lifecycle value from the state. const LIFECYCLE_MASK: usize = 0b11; /// Flag tracking if the task has been pushed into a run queue. const NOTIFIED: usize = 0b100; -/// The join handle is still around +/// The join handle is still around. #[allow(clippy::unusual_byte_groupings)] // https://github.com/rust-lang/rust-clippy/issues/6556 const JOIN_INTEREST: usize = 0b1_000; -/// A join handle waker has been set +/// A join handle waker has been set. #[allow(clippy::unusual_byte_groupings)] // https://github.com/rust-lang/rust-clippy/issues/6556 const JOIN_WAKER: usize = 0b10_000; @@ -40,36 +40,66 @@ const JOIN_WAKER: usize = 0b10_000; #[allow(clippy::unusual_byte_groupings)] // https://github.com/rust-lang/rust-clippy/issues/6556 const CANCELLED: usize = 0b100_000; -/// All bits +/// All bits. const STATE_MASK: usize = LIFECYCLE_MASK | NOTIFIED | JOIN_INTEREST | JOIN_WAKER | CANCELLED; /// Bits used by the ref count portion of the state. const REF_COUNT_MASK: usize = !STATE_MASK; -/// Number of positions to shift the ref count +/// Number of positions to shift the ref count. const REF_COUNT_SHIFT: usize = REF_COUNT_MASK.count_zeros() as usize; -/// One ref count +/// One ref count. const REF_ONE: usize = 1 << REF_COUNT_SHIFT; -/// State a task is initialized with +/// State a task is initialized with. /// -/// A task is initialized with two references: one for the scheduler and one for -/// the `JoinHandle`. As the task starts with a `JoinHandle`, `JOIN_INTEREST` is -/// set. A new task is immediately pushed into the run queue for execution and -/// starts with the `NOTIFIED` flag set. -const INITIAL_STATE: usize = (REF_ONE * 2) | JOIN_INTEREST | NOTIFIED; +/// A task is initialized with three references: +/// +/// * A reference that will be stored in an OwnedTasks or LocalOwnedTasks. +/// * A reference that will be sent to the scheduler as an ordinary notification. +/// * A reference for the JoinHandle. +/// +/// As the task starts with a `JoinHandle`, `JOIN_INTEREST` is set. +/// As the task starts with a `Notified`, `NOTIFIED` is set. +const INITIAL_STATE: usize = (REF_ONE * 3) | JOIN_INTEREST | NOTIFIED; + +#[must_use] +pub(super) enum TransitionToRunning { + Success, + Cancelled, + Failed, + Dealloc, +} + +#[must_use] +pub(super) enum TransitionToIdle { + Ok, + OkNotified, + OkDealloc, + Cancelled, +} + +#[must_use] +pub(super) enum TransitionToNotifiedByVal { + DoNothing, + Submit, + Dealloc, +} + +#[must_use] +pub(crate) enum TransitionToNotifiedByRef { + DoNothing, + Submit, +} /// All transitions are performed via RMW operations. This establishes an /// unambiguous modification order. impl State { - /// Return a task's initial state + /// Returns a task's initial state. pub(super) fn new() -> State { - // A task is initialized with three references: one for the scheduler, - // one for the `JoinHandle`, one for the task handle made available in - // release. As the task starts with a `JoinHandle`, `JOIN_INTEREST` is - // set. A new task is immediately pushed into the run queue for - // execution and starts with the `NOTIFIED` flag set. + // The raw task returned by this method has a ref-count of three. See + // the comment on INITIAL_STATE for more. State { val: AtomicUsize::new(INITIAL_STATE), } @@ -80,57 +110,72 @@ impl State { Snapshot(self.val.load(Acquire)) } - /// Attempt to transition the lifecycle to `Running`. - /// - /// If `ref_inc` is set, the reference count is also incremented. - /// - /// The `NOTIFIED` bit is always unset. - pub(super) fn transition_to_running(&self, ref_inc: bool) -> UpdateResult { - self.fetch_update(|curr| { - assert!(curr.is_notified()); - - let mut next = curr; + /// Attempts to transition the lifecycle to `Running`. This sets the + /// notified bit to false so notifications during the poll can be detected. + pub(super) fn transition_to_running(&self) -> TransitionToRunning { + self.fetch_update_action(|mut next| { + let action; + assert!(next.is_notified()); if !next.is_idle() { - return None; - } - - if ref_inc { - next.ref_inc(); + // This happens if the task is either currently running or if it + // has already completed, e.g. if it was cancelled during + // shutdown. Consume the ref-count and return. + next.ref_dec(); + if next.ref_count() == 0 { + action = TransitionToRunning::Dealloc; + } else { + action = TransitionToRunning::Failed; + } + } else { + // We are able to lock the RUNNING bit. + next.set_running(); + next.unset_notified(); + + if next.is_cancelled() { + action = TransitionToRunning::Cancelled; + } else { + action = TransitionToRunning::Success; + } } - - next.set_running(); - next.unset_notified(); - Some(next) + (action, Some(next)) }) } /// Transitions the task from `Running` -> `Idle`. /// - /// Returns `Ok` if the transition to `Idle` is successful, `Err` otherwise. - /// In both cases, a snapshot of the state from **after** the transition is - /// returned. - /// + /// Returns `true` if the transition to `Idle` is successful, `false` otherwise. /// The transition to `Idle` fails if the task has been flagged to be /// cancelled. - pub(super) fn transition_to_idle(&self) -> UpdateResult { - self.fetch_update(|curr| { + pub(super) fn transition_to_idle(&self) -> TransitionToIdle { + self.fetch_update_action(|curr| { assert!(curr.is_running()); if curr.is_cancelled() { - return None; + return (TransitionToIdle::Cancelled, None); } let mut next = curr; + let action; next.unset_running(); - if next.is_notified() { - // The caller needs to schedule the task. To do this, it needs a - // waker. The waker requires a ref count. + if !next.is_notified() { + // Polling the future consumes the ref-count of the Notified. + next.ref_dec(); + if next.ref_count() == 0 { + action = TransitionToIdle::OkDealloc; + } else { + action = TransitionToIdle::Ok; + } + } else { + // The caller will schedule a new notification, so we create a + // new ref-count for the notification. Our own ref-count is kept + // for now, and the caller will drop it shortly. next.ref_inc(); + action = TransitionToIdle::OkNotified; } - Some(next) + (action, Some(next)) }) } @@ -145,51 +190,139 @@ impl State { Snapshot(prev.0 ^ DELTA) } - /// Transition from `Complete` -> `Terminal`, decrementing the reference - /// count by 1. + /// Transitions from `Complete` -> `Terminal`, decrementing the reference + /// count the specified number of times. /// - /// When `ref_dec` is set, an additional ref count decrement is performed. - /// This is used to batch atomic ops when possible. - pub(super) fn transition_to_terminal(&self, complete: bool, ref_dec: bool) -> Snapshot { - self.fetch_update(|mut snapshot| { - if complete { - snapshot.set_complete(); - } else { - assert!(snapshot.is_complete()); - } + /// Returns true if the task should be deallocated. + pub(super) fn transition_to_terminal(&self, count: usize) -> bool { + let prev = Snapshot(self.val.fetch_sub(count * REF_ONE, AcqRel)); + assert!( + prev.ref_count() >= count, + "current: {}, sub: {}", + prev.ref_count(), + count + ); + prev.ref_count() == count + } + + /// Transitions the state to `NOTIFIED`. + /// + /// If no task needs to be submitted, a ref-count is consumed. + /// + /// If a task needs to be submitted, the ref-count is incremented for the + /// new Notified. + pub(super) fn transition_to_notified_by_val(&self) -> TransitionToNotifiedByVal { + self.fetch_update_action(|mut snapshot| { + let action; + + if snapshot.is_running() { + // If the task is running, we mark it as notified, but we should + // not submit anything as the thread currently running the + // future is responsible for that. + snapshot.set_notified(); + snapshot.ref_dec(); - // Decrement the primary handle - snapshot.ref_dec(); + // The thread that set the running bit also holds a ref-count. + assert!(snapshot.ref_count() > 0); - if ref_dec { - // Decrement a second time + action = TransitionToNotifiedByVal::DoNothing; + } else if snapshot.is_complete() || snapshot.is_notified() { + // We do not need to submit any notifications, but we have to + // decrement the ref-count. snapshot.ref_dec(); + + if snapshot.ref_count() == 0 { + action = TransitionToNotifiedByVal::Dealloc; + } else { + action = TransitionToNotifiedByVal::DoNothing; + } + } else { + // We create a new notified that we can submit. The caller + // retains ownership of the ref-count they passed in. + snapshot.set_notified(); + snapshot.ref_inc(); + action = TransitionToNotifiedByVal::Submit; } - Some(snapshot) + (action, Some(snapshot)) }) - .unwrap() } /// Transitions the state to `NOTIFIED`. - /// - /// Returns `true` if the task needs to be submitted to the pool for - /// execution - pub(super) fn transition_to_notified(&self) -> bool { - let prev = Snapshot(self.val.fetch_or(NOTIFIED, AcqRel)); - prev.will_need_queueing() + pub(super) fn transition_to_notified_by_ref(&self) -> TransitionToNotifiedByRef { + self.fetch_update_action(|mut snapshot| { + if snapshot.is_complete() || snapshot.is_notified() { + // There is nothing to do in this case. + (TransitionToNotifiedByRef::DoNothing, None) + } else if snapshot.is_running() { + // If the task is running, we mark it as notified, but we should + // not submit as the thread currently running the future is + // responsible for that. + snapshot.set_notified(); + (TransitionToNotifiedByRef::DoNothing, Some(snapshot)) + } else { + // The task is idle and not notified. We should submit a + // notification. + snapshot.set_notified(); + snapshot.ref_inc(); + (TransitionToNotifiedByRef::Submit, Some(snapshot)) + } + }) } - /// Set the cancelled bit and transition the state to `NOTIFIED`. + /// Transitions the state to `NOTIFIED`, unconditionally increasing the ref count. + #[cfg(all( + tokio_unstable, + tokio_taskdump, + feature = "rt", + target_os = "linux", + any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64") + ))] + pub(super) fn transition_to_notified_for_tracing(&self) { + self.fetch_update_action(|mut snapshot| { + snapshot.set_notified(); + snapshot.ref_inc(); + ((), Some(snapshot)) + }); + } + + /// Sets the cancelled bit and transitions the state to `NOTIFIED` if idle. /// /// Returns `true` if the task needs to be submitted to the pool for - /// execution + /// execution. pub(super) fn transition_to_notified_and_cancel(&self) -> bool { - let prev = Snapshot(self.val.fetch_or(NOTIFIED | CANCELLED, AcqRel)); - prev.will_need_queueing() + self.fetch_update_action(|mut snapshot| { + if snapshot.is_cancelled() || snapshot.is_complete() { + // Aborts to completed or cancelled tasks are no-ops. + (false, None) + } else if snapshot.is_running() { + // If the task is running, we mark it as cancelled. The thread + // running the task will notice the cancelled bit when it + // stops polling and it will kill the task. + // + // The set_notified() call is not strictly necessary but it will + // in some cases let a wake_by_ref call return without having + // to perform a compare_exchange. + snapshot.set_notified(); + snapshot.set_cancelled(); + (false, Some(snapshot)) + } else { + // The task is idle. We set the cancelled and notified bits and + // submit a notification if the notified bit was not already + // set. + snapshot.set_cancelled(); + if !snapshot.is_notified() { + snapshot.set_notified(); + snapshot.ref_inc(); + (true, Some(snapshot)) + } else { + (false, Some(snapshot)) + } + } + }) } - /// Set the `CANCELLED` bit and attempt to transition to `Running`. + /// Sets the `CANCELLED` bit and attempts to transition to `Running`. /// /// Returns `true` if the transition to `Running` succeeded. pub(super) fn transition_to_shutdown(&self) -> bool { @@ -200,17 +333,11 @@ impl State { if snapshot.is_idle() { snapshot.set_running(); - - if snapshot.is_notified() { - // If the task is idle and notified, this indicates the task is - // in the run queue and is considered owned by the scheduler. - // The shutdown operation claims ownership of the task, which - // means we need to assign an additional ref-count to the task - // in the queue. - snapshot.ref_inc(); - } } + // If the task was not idle, the thread currently running the task + // will notice the cancelled bit and cancel it once the poll + // completes. snapshot.set_cancelled(); Some(snapshot) }); @@ -219,7 +346,7 @@ impl State { } /// Optimistically tries to swap the state assuming the join handle is - /// __immediately__ dropped on spawn + /// __immediately__ dropped on spawn. pub(super) fn drop_join_handle_fast(&self) -> Result<(), ()> { use std::sync::atomic::Ordering::Relaxed; @@ -241,7 +368,7 @@ impl State { .map_err(|_| ()) } - /// Try to unset the JOIN_INTEREST flag. + /// Tries to unset the JOIN_INTEREST flag. /// /// Returns `Ok` if the operation happens before the task transitions to a /// completed state, `Err` otherwise. @@ -260,14 +387,14 @@ impl State { }) } - /// Set the `JOIN_WAKER` bit. + /// Sets the `JOIN_WAKER` bit. /// /// Returns `Ok` if the bit is set, `Err` otherwise. This operation fails if /// the task has completed. pub(super) fn set_join_waker(&self) -> UpdateResult { self.fetch_update(|curr| { assert!(curr.is_join_interested()); - assert!(!curr.has_join_waker()); + assert!(!curr.is_join_waker_set()); if curr.is_complete() { return None; @@ -287,7 +414,7 @@ impl State { pub(super) fn unset_waker(&self) -> UpdateResult { self.fetch_update(|curr| { assert!(curr.is_join_interested()); - assert!(curr.has_join_waker()); + assert!(curr.is_join_waker_set()); if curr.is_complete() { return None; @@ -326,9 +453,39 @@ impl State { /// Returns `true` if the task should be released. pub(super) fn ref_dec(&self) -> bool { let prev = Snapshot(self.val.fetch_sub(REF_ONE, AcqRel)); + assert!(prev.ref_count() >= 1); prev.ref_count() == 1 } + /// Returns `true` if the task should be released. + pub(super) fn ref_dec_twice(&self) -> bool { + let prev = Snapshot(self.val.fetch_sub(2 * REF_ONE, AcqRel)); + assert!(prev.ref_count() >= 2); + prev.ref_count() == 2 + } + + fn fetch_update_action<F, T>(&self, mut f: F) -> T + where + F: FnMut(Snapshot) -> (T, Option<Snapshot>), + { + let mut curr = self.load(); + + loop { + let (output, next) = f(curr); + let next = match next { + Some(next) => next, + None => return output, + }; + + let res = self.val.compare_exchange(curr.0, next.0, AcqRel, Acquire); + + match res { + Ok(_) => return output, + Err(actual) => curr = Snapshot(actual), + } + } + } + fn fetch_update<F>(&self, mut f: F) -> Result<Snapshot, Snapshot> where F: FnMut(Snapshot) -> Option<Snapshot>, @@ -368,6 +525,10 @@ impl Snapshot { self.0 &= !NOTIFIED } + fn set_notified(&mut self) { + self.0 |= NOTIFIED + } + pub(super) fn is_running(self) -> bool { self.0 & RUNNING == RUNNING } @@ -388,10 +549,6 @@ impl Snapshot { self.0 |= CANCELLED; } - fn set_complete(&mut self) { - self.0 |= COMPLETE; - } - /// Returns `true` if the task's future has completed execution. pub(super) fn is_complete(self) -> bool { self.0 & COMPLETE == COMPLETE @@ -405,7 +562,7 @@ impl Snapshot { self.0 &= !JOIN_INTEREST } - pub(super) fn has_join_waker(self) -> bool { + pub(super) fn is_join_waker_set(self) -> bool { self.0 & JOIN_WAKER == JOIN_WAKER } @@ -430,10 +587,6 @@ impl Snapshot { assert!(self.ref_count() > 0); self.0 -= REF_ONE } - - fn will_need_queueing(self) -> bool { - !self.is_notified() && self.is_idle() - } } impl fmt::Debug for State { @@ -451,7 +604,7 @@ impl fmt::Debug for Snapshot { .field("is_notified", &self.is_notified()) .field("is_cancelled", &self.is_cancelled()) .field("is_join_interested", &self.is_join_interested()) - .field("has_join_waker", &self.has_join_waker()) + .field("is_join_waker_set", &self.is_join_waker_set()) .field("ref_count", &self.ref_count()) .finish() } diff --git a/vendor/tokio/src/runtime/task/trace/mod.rs b/vendor/tokio/src/runtime/task/trace/mod.rs new file mode 100644 index 000000000..543b7eee9 --- /dev/null +++ b/vendor/tokio/src/runtime/task/trace/mod.rs @@ -0,0 +1,330 @@ +use crate::loom::sync::Arc; +use crate::runtime::context; +use crate::runtime::scheduler::{self, current_thread, Inject}; + +use backtrace::BacktraceFrame; +use std::cell::Cell; +use std::collections::VecDeque; +use std::ffi::c_void; +use std::fmt; +use std::future::Future; +use std::pin::Pin; +use std::ptr::{self, NonNull}; +use std::task::{self, Poll}; + +mod symbol; +mod tree; + +use symbol::Symbol; +use tree::Tree; + +use super::{Notified, OwnedTasks}; + +type Backtrace = Vec<BacktraceFrame>; +type SymbolTrace = Vec<Symbol>; + +/// The ambiant backtracing context. +pub(crate) struct Context { + /// The address of [`Trace::root`] establishes an upper unwinding bound on + /// the backtraces in `Trace`. + active_frame: Cell<Option<NonNull<Frame>>>, + /// The place to stash backtraces. + collector: Cell<Option<Trace>>, +} + +/// A [`Frame`] in an intrusive, doubly-linked tree of [`Frame`]s. +struct Frame { + /// The location associated with this frame. + inner_addr: *const c_void, + + /// The parent frame, if any. + parent: Option<NonNull<Frame>>, +} + +/// An tree execution trace. +/// +/// Traces are captured with [`Trace::capture`], rooted with [`Trace::root`] +/// and leaved with [`trace_leaf`]. +#[derive(Clone, Debug)] +pub(crate) struct Trace { + // The linear backtraces that comprise this trace. These linear traces can + // be re-knitted into a tree. + backtraces: Vec<Backtrace>, +} + +pin_project_lite::pin_project! { + #[derive(Debug, Clone)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub(crate) struct Root<T> { + #[pin] + future: T, + } +} + +const FAIL_NO_THREAD_LOCAL: &str = "The Tokio thread-local has been destroyed \ + as part of shutting down the current \ + thread, so collecting a taskdump is not \ + possible."; + +impl Context { + pub(crate) const fn new() -> Self { + Context { + active_frame: Cell::new(None), + collector: Cell::new(None), + } + } + + /// SAFETY: Callers of this function must ensure that trace frames always + /// form a valid linked list. + unsafe fn try_with_current<F, R>(f: F) -> Option<R> + where + F: FnOnce(&Self) -> R, + { + crate::runtime::context::with_trace(f) + } + + unsafe fn with_current_frame<F, R>(f: F) -> R + where + F: FnOnce(&Cell<Option<NonNull<Frame>>>) -> R, + { + Self::try_with_current(|context| f(&context.active_frame)).expect(FAIL_NO_THREAD_LOCAL) + } + + fn with_current_collector<F, R>(f: F) -> R + where + F: FnOnce(&Cell<Option<Trace>>) -> R, + { + // SAFETY: This call can only access the collector field, so it cannot + // break the trace frame linked list. + unsafe { + Self::try_with_current(|context| f(&context.collector)).expect(FAIL_NO_THREAD_LOCAL) + } + } +} + +impl Trace { + /// Invokes `f`, returning both its result and the collection of backtraces + /// captured at each sub-invocation of [`trace_leaf`]. + #[inline(never)] + pub(crate) fn capture<F, R>(f: F) -> (R, Trace) + where + F: FnOnce() -> R, + { + let collector = Trace { backtraces: vec![] }; + + let previous = Context::with_current_collector(|current| current.replace(Some(collector))); + + let result = f(); + + let collector = + Context::with_current_collector(|current| current.replace(previous)).unwrap(); + + (result, collector) + } + + /// The root of a trace. + #[inline(never)] + pub(crate) fn root<F>(future: F) -> Root<F> { + Root { future } + } +} + +/// If this is a sub-invocation of [`Trace::capture`], capture a backtrace. +/// +/// The captured backtrace will be returned by [`Trace::capture`]. +/// +/// Invoking this function does nothing when it is not a sub-invocation +/// [`Trace::capture`]. +// This function is marked `#[inline(never)]` to ensure that it gets a distinct `Frame` in the +// backtrace, below which frames should not be included in the backtrace (since they reflect the +// internal implementation details of this crate). +#[inline(never)] +pub(crate) fn trace_leaf(cx: &mut task::Context<'_>) -> Poll<()> { + // Safety: We don't manipulate the current context's active frame. + let did_trace = unsafe { + Context::try_with_current(|context_cell| { + if let Some(mut collector) = context_cell.collector.take() { + let mut frames = vec![]; + let mut above_leaf = false; + + if let Some(active_frame) = context_cell.active_frame.get() { + let active_frame = active_frame.as_ref(); + + backtrace::trace(|frame| { + let below_root = !ptr::eq(frame.symbol_address(), active_frame.inner_addr); + + // only capture frames above `Trace::leaf` and below + // `Trace::root`. + if above_leaf && below_root { + frames.push(frame.to_owned().into()); + } + + if ptr::eq(frame.symbol_address(), trace_leaf as *const _) { + above_leaf = true; + } + + // only continue unwinding if we're below `Trace::root` + below_root + }); + } + collector.backtraces.push(frames); + context_cell.collector.set(Some(collector)); + true + } else { + false + } + }) + .unwrap_or(false) + }; + + if did_trace { + // Use the same logic that `yield_now` uses to send out wakeups after + // the task yields. + context::with_scheduler(|scheduler| { + if let Some(scheduler) = scheduler { + match scheduler { + scheduler::Context::CurrentThread(s) => s.defer.defer(cx.waker()), + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + scheduler::Context::MultiThread(s) => s.defer.defer(cx.waker()), + } + } + }); + + Poll::Pending + } else { + Poll::Ready(()) + } +} + +impl fmt::Display for Trace { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + Tree::from_trace(self.clone()).fmt(f) + } +} + +fn defer<F: FnOnce() -> R, R>(f: F) -> impl Drop { + use std::mem::ManuallyDrop; + + struct Defer<F: FnOnce() -> R, R>(ManuallyDrop<F>); + + impl<F: FnOnce() -> R, R> Drop for Defer<F, R> { + #[inline(always)] + fn drop(&mut self) { + unsafe { + ManuallyDrop::take(&mut self.0)(); + } + } + } + + Defer(ManuallyDrop::new(f)) +} + +impl<T: Future> Future for Root<T> { + type Output = T::Output; + + #[inline(never)] + fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> { + // SAFETY: The context's current frame is restored to its original state + // before `frame` is dropped. + unsafe { + let mut frame = Frame { + inner_addr: Self::poll as *const c_void, + parent: None, + }; + + Context::with_current_frame(|current| { + frame.parent = current.take(); + current.set(Some(NonNull::from(&frame))); + }); + + let _restore = defer(|| { + Context::with_current_frame(|current| { + current.set(frame.parent); + }); + }); + + let this = self.project(); + this.future.poll(cx) + } + } +} + +/// Trace and poll all tasks of the current_thread runtime. +pub(in crate::runtime) fn trace_current_thread( + owned: &OwnedTasks<Arc<current_thread::Handle>>, + local: &mut VecDeque<Notified<Arc<current_thread::Handle>>>, + injection: &Inject<Arc<current_thread::Handle>>, +) -> Vec<Trace> { + // clear the local and injection queues + local.clear(); + + while let Some(task) = injection.pop() { + drop(task); + } + + // notify each task + let mut tasks = vec![]; + owned.for_each(|task| { + // set the notified bit + task.as_raw().state().transition_to_notified_for_tracing(); + // store the raw tasks into a vec + tasks.push(task.as_raw()); + }); + + tasks + .into_iter() + .map(|task| { + let ((), trace) = Trace::capture(|| task.poll()); + trace + }) + .collect() +} + +cfg_rt_multi_thread! { + use crate::loom::sync::Mutex; + use crate::runtime::scheduler::multi_thread; + use crate::runtime::scheduler::multi_thread::Synced; + use crate::runtime::scheduler::inject::Shared; + + /// Trace and poll all tasks of the current_thread runtime. + /// + /// ## Safety + /// + /// Must be called with the same `synced` that `injection` was created with. + pub(in crate::runtime) unsafe fn trace_multi_thread( + owned: &OwnedTasks<Arc<multi_thread::Handle>>, + local: &mut multi_thread::queue::Local<Arc<multi_thread::Handle>>, + synced: &Mutex<Synced>, + injection: &Shared<Arc<multi_thread::Handle>>, + ) -> Vec<Trace> { + // clear the local queue + while let Some(notified) = local.pop() { + drop(notified); + } + + // clear the injection queue + let mut synced = synced.lock(); + while let Some(notified) = injection.pop(&mut synced.inject) { + drop(notified); + } + + drop(synced); + + // notify each task + let mut traces = vec![]; + owned.for_each(|task| { + // set the notified bit + task.as_raw().state().transition_to_notified_for_tracing(); + + // trace the task + let ((), trace) = Trace::capture(|| task.as_raw().poll()); + traces.push(trace); + + // reschedule the task + let _ = task.as_raw().state().transition_to_notified_by_ref(); + task.as_raw().schedule(); + }); + + traces + } +} diff --git a/vendor/tokio/src/runtime/task/trace/symbol.rs b/vendor/tokio/src/runtime/task/trace/symbol.rs new file mode 100644 index 000000000..49d7ba37f --- /dev/null +++ b/vendor/tokio/src/runtime/task/trace/symbol.rs @@ -0,0 +1,92 @@ +use backtrace::BacktraceSymbol; +use std::fmt; +use std::hash::{Hash, Hasher}; +use std::ptr; + +/// A symbol in a backtrace. +/// +/// This wrapper type serves two purposes. The first is that it provides a +/// representation of a symbol that can be inserted into hashmaps and hashsets; +/// the [`backtrace`] crate does not define [`Hash`], [`PartialEq`], or [`Eq`] +/// on [`BacktraceSymbol`], and recommends that users define their own wrapper +/// which implements these traits. +/// +/// Second, this wrapper includes a `parent_hash` field that uniquely +/// identifies this symbol's position in its trace. Otherwise, e.g., our code +/// would not be able to distinguish between recursive calls of a function at +/// different depths. +#[derive(Clone)] +pub(super) struct Symbol { + pub(super) symbol: BacktraceSymbol, + pub(super) parent_hash: u64, +} + +impl Hash for Symbol { + fn hash<H>(&self, state: &mut H) + where + H: Hasher, + { + if let Some(name) = self.symbol.name() { + name.as_bytes().hash(state); + } + + if let Some(addr) = self.symbol.addr() { + ptr::hash(addr, state); + } + + self.symbol.filename().hash(state); + self.symbol.lineno().hash(state); + self.symbol.colno().hash(state); + self.parent_hash.hash(state); + } +} + +impl PartialEq for Symbol { + fn eq(&self, other: &Self) -> bool { + (self.parent_hash == other.parent_hash) + && match (self.symbol.name(), other.symbol.name()) { + (None, None) => true, + (Some(lhs_name), Some(rhs_name)) => lhs_name.as_bytes() == rhs_name.as_bytes(), + _ => false, + } + && match (self.symbol.addr(), other.symbol.addr()) { + (None, None) => true, + (Some(lhs_addr), Some(rhs_addr)) => ptr::eq(lhs_addr, rhs_addr), + _ => false, + } + && (self.symbol.filename() == other.symbol.filename()) + && (self.symbol.lineno() == other.symbol.lineno()) + && (self.symbol.colno() == other.symbol.colno()) + } +} + +impl Eq for Symbol {} + +impl fmt::Display for Symbol { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + if let Some(name) = self.symbol.name() { + let name = name.to_string(); + let name = if let Some((name, _)) = name.rsplit_once("::") { + name + } else { + &name + }; + fmt::Display::fmt(&name, f)?; + } + + if let Some(filename) = self.symbol.filename() { + f.write_str(" at ")?; + filename.to_string_lossy().fmt(f)?; + if let Some(lineno) = self.symbol.lineno() { + f.write_str(":")?; + fmt::Display::fmt(&lineno, f)?; + if let Some(colno) = self.symbol.colno() { + f.write_str(":")?; + fmt::Display::fmt(&colno, f)?; + } + } + } + + Ok(()) + } +} diff --git a/vendor/tokio/src/runtime/task/trace/tree.rs b/vendor/tokio/src/runtime/task/trace/tree.rs new file mode 100644 index 000000000..7e6f8efec --- /dev/null +++ b/vendor/tokio/src/runtime/task/trace/tree.rs @@ -0,0 +1,126 @@ +use std::collections::{hash_map::DefaultHasher, HashMap, HashSet}; +use std::fmt; +use std::hash::{Hash, Hasher}; + +use super::{Backtrace, Symbol, SymbolTrace, Trace}; + +/// An adjacency list representation of an execution tree. +/// +/// This tree provides a convenient intermediate representation for formatting +/// [`Trace`] as a tree. +pub(super) struct Tree { + /// The roots of the trees. + /// + /// There should only be one root, but the code is robust to multiple roots. + roots: HashSet<Symbol>, + + /// The adjacency list of symbols in the execution tree(s). + edges: HashMap<Symbol, HashSet<Symbol>>, +} + +impl Tree { + /// Constructs a [`Tree`] from [`Trace`] + pub(super) fn from_trace(trace: Trace) -> Self { + let mut roots: HashSet<Symbol> = HashSet::default(); + let mut edges: HashMap<Symbol, HashSet<Symbol>> = HashMap::default(); + + for trace in trace.backtraces { + let trace = to_symboltrace(trace); + + if let Some(first) = trace.first() { + roots.insert(first.to_owned()); + } + + let mut trace = trace.into_iter().peekable(); + while let Some(frame) = trace.next() { + let subframes = edges.entry(frame).or_default(); + if let Some(subframe) = trace.peek() { + subframes.insert(subframe.clone()); + } + } + } + + Tree { roots, edges } + } + + /// Produces the sub-symbols of a given symbol. + fn consequences(&self, frame: &Symbol) -> Option<impl ExactSizeIterator<Item = &Symbol>> { + Some(self.edges.get(frame)?.iter()) + } + + /// Format this [`Tree`] as a textual tree. + fn display<W: fmt::Write>( + &self, + f: &mut W, + root: &Symbol, + is_last: bool, + prefix: &str, + ) -> fmt::Result { + let root_fmt = format!("{}", root); + + let current; + let next; + + if is_last { + current = format!("{prefix}└╼\u{a0}{root_fmt}"); + next = format!("{}\u{a0}\u{a0}\u{a0}", prefix); + } else { + current = format!("{prefix}├╼\u{a0}{root_fmt}"); + next = format!("{}│\u{a0}\u{a0}", prefix); + } + + write!(f, "{}", { + let mut current = current.chars(); + current.next().unwrap(); + current.next().unwrap(); + ¤t.as_str() + })?; + + if let Some(consequences) = self.consequences(root) { + let len = consequences.len(); + for (i, consequence) in consequences.enumerate() { + let is_last = i == len - 1; + writeln!(f)?; + self.display(f, consequence, is_last, &next)?; + } + } + + Ok(()) + } +} + +impl fmt::Display for Tree { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + for root in &self.roots { + self.display(f, root, true, " ")?; + } + Ok(()) + } +} + +/// Resolve a sequence of [`backtrace::BacktraceFrame`]s into a sequence of +/// [`Symbol`]s. +fn to_symboltrace(backtrace: Backtrace) -> SymbolTrace { + // Resolve the backtrace frames to symbols. + let backtrace: Backtrace = { + let mut backtrace = backtrace::Backtrace::from(backtrace); + backtrace.resolve(); + backtrace.into() + }; + + // Accumulate the symbols in descending order into `symboltrace`. + let mut symboltrace: SymbolTrace = vec![]; + let mut state = DefaultHasher::new(); + for frame in backtrace.into_iter().rev() { + for symbol in frame.symbols().iter().rev() { + let symbol = Symbol { + symbol: symbol.clone(), + parent_hash: state.finish(), + }; + symbol.hash(&mut state); + symboltrace.push(symbol); + } + } + + symboltrace +} diff --git a/vendor/tokio/src/runtime/task/waker.rs b/vendor/tokio/src/runtime/task/waker.rs index b7313b4c5..b5f5ace9e 100644 --- a/vendor/tokio/src/runtime/task/waker.rs +++ b/vendor/tokio/src/runtime/task/waker.rs @@ -1,6 +1,5 @@ use crate::future::Future; -use crate::runtime::task::harness::Harness; -use crate::runtime::task::{Header, Schedule}; +use crate::runtime::task::{Header, RawTask, Schedule}; use std::marker::PhantomData; use std::mem::ManuallyDrop; @@ -13,9 +12,9 @@ pub(super) struct WakerRef<'a, S: 'static> { _p: PhantomData<(&'a Header, S)>, } -/// Returns a `WakerRef` which avoids having to pre-emptively increase the +/// Returns a `WakerRef` which avoids having to preemptively increase the /// refcount if there is no need to do so. -pub(super) fn waker_ref<T, S>(header: &Header) -> WakerRef<'_, S> +pub(super) fn waker_ref<T, S>(header: &NonNull<Header>) -> WakerRef<'_, S> where T: Future, S: Schedule, @@ -28,7 +27,7 @@ where // point and not an *owned* waker, we must ensure that `drop` is never // called on this waker instance. This is done by wrapping it with // `ManuallyDrop` and then never calling drop. - let waker = unsafe { ManuallyDrop::new(Waker::from_raw(raw_waker::<T, S>(header))) }; + let waker = unsafe { ManuallyDrop::new(Waker::from_raw(raw_waker(*header))) }; WakerRef { waker, @@ -46,8 +45,8 @@ impl<S> ops::Deref for WakerRef<'_, S> { cfg_trace! { macro_rules! trace { - ($harness:expr, $op:expr) => { - if let Some(id) = $harness.id() { + ($header:expr, $op:expr) => { + if let Some(id) = Header::get_tracing_id(&$header) { tracing::trace!( target: "tokio::task::waker", op = $op, @@ -60,71 +59,46 @@ cfg_trace! { cfg_not_trace! { macro_rules! trace { - ($harness:expr, $op:expr) => { + ($header:expr, $op:expr) => { // noop - let _ = &$harness; + let _ = &$header; } } } -unsafe fn clone_waker<T, S>(ptr: *const ()) -> RawWaker -where - T: Future, - S: Schedule, -{ - let header = ptr as *const Header; - let ptr = NonNull::new_unchecked(ptr as *mut Header); - let harness = Harness::<T, S>::from_raw(ptr); - trace!(harness, "waker.clone"); - (*header).state.ref_inc(); - raw_waker::<T, S>(header) +unsafe fn clone_waker(ptr: *const ()) -> RawWaker { + let header = NonNull::new_unchecked(ptr as *mut Header); + trace!(header, "waker.clone"); + header.as_ref().state.ref_inc(); + raw_waker(header) } -unsafe fn drop_waker<T, S>(ptr: *const ()) -where - T: Future, - S: Schedule, -{ +unsafe fn drop_waker(ptr: *const ()) { let ptr = NonNull::new_unchecked(ptr as *mut Header); - let harness = Harness::<T, S>::from_raw(ptr); - trace!(harness, "waker.drop"); - harness.drop_reference(); + trace!(ptr, "waker.drop"); + let raw = RawTask::from_raw(ptr); + raw.drop_reference(); } -unsafe fn wake_by_val<T, S>(ptr: *const ()) -where - T: Future, - S: Schedule, -{ +unsafe fn wake_by_val(ptr: *const ()) { let ptr = NonNull::new_unchecked(ptr as *mut Header); - let harness = Harness::<T, S>::from_raw(ptr); - trace!(harness, "waker.wake"); - harness.wake_by_val(); + trace!(ptr, "waker.wake"); + let raw = RawTask::from_raw(ptr); + raw.wake_by_val(); } // Wake without consuming the waker -unsafe fn wake_by_ref<T, S>(ptr: *const ()) -where - T: Future, - S: Schedule, -{ +unsafe fn wake_by_ref(ptr: *const ()) { let ptr = NonNull::new_unchecked(ptr as *mut Header); - let harness = Harness::<T, S>::from_raw(ptr); - trace!(harness, "waker.wake_by_ref"); - harness.wake_by_ref(); + trace!(ptr, "waker.wake_by_ref"); + let raw = RawTask::from_raw(ptr); + raw.wake_by_ref(); } -fn raw_waker<T, S>(header: *const Header) -> RawWaker -where - T: Future, - S: Schedule, -{ - let ptr = header as *const (); - let vtable = &RawWakerVTable::new( - clone_waker::<T, S>, - wake_by_val::<T, S>, - wake_by_ref::<T, S>, - drop_waker::<T, S>, - ); - RawWaker::new(ptr, vtable) +static WAKER_VTABLE: RawWakerVTable = + RawWakerVTable::new(clone_waker, wake_by_val, wake_by_ref, drop_waker); + +fn raw_waker(header: NonNull<Header>) -> RawWaker { + let ptr = header.as_ptr() as *const (); + RawWaker::new(ptr, &WAKER_VTABLE) } diff --git a/vendor/tokio/src/runtime/tests/inject.rs b/vendor/tokio/src/runtime/tests/inject.rs new file mode 100644 index 000000000..ccead5e02 --- /dev/null +++ b/vendor/tokio/src/runtime/tests/inject.rs @@ -0,0 +1,54 @@ +use crate::runtime::scheduler::inject; + +#[test] +fn push_and_pop() { + const N: usize = 2; + + let (inject, mut synced) = inject::Shared::new(); + + for i in 0..N { + assert_eq!(inject.len(), i); + let (task, _) = super::unowned(async {}); + unsafe { inject.push(&mut synced, task) }; + } + + for i in 0..N { + assert_eq!(inject.len(), N - i); + assert!(unsafe { inject.pop(&mut synced) }.is_some()); + } + + println!("--------------"); + + assert!(unsafe { inject.pop(&mut synced) }.is_none()); +} + +#[test] +fn push_batch_and_pop() { + let (inject, mut inject_synced) = inject::Shared::new(); + + unsafe { + inject.push_batch( + &mut inject_synced, + (0..10).map(|_| super::unowned(async {}).0), + ); + + assert_eq!(5, inject.pop_n(&mut inject_synced, 5).count()); + assert_eq!(5, inject.pop_n(&mut inject_synced, 5).count()); + assert_eq!(0, inject.pop_n(&mut inject_synced, 5).count()); + } +} + +#[test] +fn pop_n_drains_on_drop() { + let (inject, mut inject_synced) = inject::Shared::new(); + + unsafe { + inject.push_batch( + &mut inject_synced, + (0..10).map(|_| super::unowned(async {}).0), + ); + let _ = inject.pop_n(&mut inject_synced, 10); + + assert_eq!(inject.len(), 0); + } +} diff --git a/vendor/tokio/src/runtime/tests/loom_basic_scheduler.rs b/vendor/tokio/src/runtime/tests/loom_basic_scheduler.rs deleted file mode 100644 index e6221d3b1..000000000 --- a/vendor/tokio/src/runtime/tests/loom_basic_scheduler.rs +++ /dev/null @@ -1,82 +0,0 @@ -use crate::loom::sync::atomic::AtomicUsize; -use crate::loom::sync::Arc; -use crate::loom::thread; -use crate::runtime::{Builder, Runtime}; -use crate::sync::oneshot::{self, Receiver}; -use crate::task; -use std::future::Future; -use std::pin::Pin; -use std::sync::atomic::Ordering::{Acquire, Release}; -use std::task::{Context, Poll}; - -fn assert_at_most_num_polls(rt: Arc<Runtime>, at_most_polls: usize) { - let (tx, rx) = oneshot::channel(); - let num_polls = Arc::new(AtomicUsize::new(0)); - rt.spawn(async move { - for _ in 0..12 { - task::yield_now().await; - } - tx.send(()).unwrap(); - }); - - rt.block_on(async { - BlockedFuture { - rx, - num_polls: num_polls.clone(), - } - .await; - }); - - let polls = num_polls.load(Acquire); - assert!(polls <= at_most_polls); -} - -#[test] -fn block_on_num_polls() { - loom::model(|| { - // we expect at most 3 number of polls because there are - // three points at which we poll the future. At any of these - // points it can be ready: - // - // - when we fail to steal the parker and we block on a - // notification that it is available. - // - // - when we steal the parker and we schedule the future - // - // - when the future is woken up and we have ran the max - // number of tasks for the current tick or there are no - // more tasks to run. - // - let at_most = 3; - - let rt1 = Arc::new(Builder::new_current_thread().build().unwrap()); - let rt2 = rt1.clone(); - let rt3 = rt1.clone(); - - let th1 = thread::spawn(move || assert_at_most_num_polls(rt1, at_most)); - let th2 = thread::spawn(move || assert_at_most_num_polls(rt2, at_most)); - let th3 = thread::spawn(move || assert_at_most_num_polls(rt3, at_most)); - - th1.join().unwrap(); - th2.join().unwrap(); - th3.join().unwrap(); - }); -} - -struct BlockedFuture { - rx: Receiver<()>, - num_polls: Arc<AtomicUsize>, -} - -impl Future for BlockedFuture { - type Output = (); - - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { - self.num_polls.fetch_add(1, Release); - - match Pin::new(&mut self.rx).poll(cx) { - Poll::Pending => Poll::Pending, - _ => Poll::Ready(()), - } - } -} diff --git a/vendor/tokio/src/runtime/tests/loom_blocking.rs b/vendor/tokio/src/runtime/tests/loom_blocking.rs index 8fb54c565..5c4aeae39 100644 --- a/vendor/tokio/src/runtime/tests/loom_blocking.rs +++ b/vendor/tokio/src/runtime/tests/loom_blocking.rs @@ -23,6 +23,77 @@ fn blocking_shutdown() { }); } +#[test] +fn spawn_mandatory_blocking_should_always_run() { + use crate::runtime::tests::loom_oneshot; + loom::model(|| { + let rt = runtime::Builder::new_current_thread().build().unwrap(); + + let (tx, rx) = loom_oneshot::channel(); + let _enter = rt.enter(); + runtime::spawn_blocking(|| {}); + runtime::spawn_mandatory_blocking(move || { + let _ = tx.send(()); + }) + .unwrap(); + + drop(rt); + + // This call will deadlock if `spawn_mandatory_blocking` doesn't run. + let () = rx.recv(); + }); +} + +#[test] +fn spawn_mandatory_blocking_should_run_even_when_shutting_down_from_other_thread() { + use crate::runtime::tests::loom_oneshot; + loom::model(|| { + let rt = runtime::Builder::new_current_thread().build().unwrap(); + let handle = rt.handle().clone(); + + // Drop the runtime in a different thread + { + loom::thread::spawn(move || { + drop(rt); + }); + } + + let _enter = handle.enter(); + let (tx, rx) = loom_oneshot::channel(); + let handle = runtime::spawn_mandatory_blocking(move || { + let _ = tx.send(()); + }); + + // handle.is_some() means that `spawn_mandatory_blocking` + // promised us to run the blocking task + if handle.is_some() { + // This call will deadlock if `spawn_mandatory_blocking` doesn't run. + let () = rx.recv(); + } + }); +} + +#[test] +fn spawn_blocking_when_paused() { + use std::time::Duration; + loom::model(|| { + let rt = crate::runtime::Builder::new_current_thread() + .enable_time() + .start_paused(true) + .build() + .unwrap(); + let handle = rt.handle(); + let _enter = handle.enter(); + let a = crate::task::spawn_blocking(|| {}); + let b = crate::task::spawn_blocking(|| {}); + rt.block_on(crate::time::timeout(Duration::from_millis(1), async move { + a.await.expect("blocking task should finish"); + b.await.expect("blocking task should finish"); + })) + .expect("timeout should not trigger"); + }); +} + fn mk_runtime(num_threads: usize) -> Runtime { runtime::Builder::new_multi_thread() .worker_threads(num_threads) diff --git a/vendor/tokio/src/runtime/tests/loom_current_thread_scheduler.rs b/vendor/tokio/src/runtime/tests/loom_current_thread_scheduler.rs new file mode 100644 index 000000000..a772603f7 --- /dev/null +++ b/vendor/tokio/src/runtime/tests/loom_current_thread_scheduler.rs @@ -0,0 +1,142 @@ +use crate::loom::sync::atomic::AtomicUsize; +use crate::loom::sync::Arc; +use crate::loom::thread; +use crate::runtime::{Builder, Runtime}; +use crate::sync::oneshot::{self, Receiver}; +use crate::task; +use std::future::Future; +use std::pin::Pin; +use std::sync::atomic::Ordering::{Acquire, Release}; +use std::task::{Context, Poll}; + +fn assert_at_most_num_polls(rt: Arc<Runtime>, at_most_polls: usize) { + let (tx, rx) = oneshot::channel(); + let num_polls = Arc::new(AtomicUsize::new(0)); + rt.spawn(async move { + for _ in 0..12 { + task::yield_now().await; + } + tx.send(()).unwrap(); + }); + + rt.block_on(async { + BlockedFuture { + rx, + num_polls: num_polls.clone(), + } + .await; + }); + + let polls = num_polls.load(Acquire); + assert!(polls <= at_most_polls); +} + +#[test] +fn block_on_num_polls() { + loom::model(|| { + // we expect at most 4 number of polls because there are three points at + // which we poll the future and an opportunity for a false-positive.. At + // any of these points it can be ready: + // + // - when we fail to steal the parker and we block on a notification + // that it is available. + // + // - when we steal the parker and we schedule the future + // + // - when the future is woken up and we have ran the max number of tasks + // for the current tick or there are no more tasks to run. + // + // - a thread is notified that the parker is available but a third + // thread acquires it before the notified thread can. + // + let at_most = 4; + + let rt1 = Arc::new(Builder::new_current_thread().build().unwrap()); + let rt2 = rt1.clone(); + let rt3 = rt1.clone(); + + let th1 = thread::spawn(move || assert_at_most_num_polls(rt1, at_most)); + let th2 = thread::spawn(move || assert_at_most_num_polls(rt2, at_most)); + let th3 = thread::spawn(move || assert_at_most_num_polls(rt3, at_most)); + + th1.join().unwrap(); + th2.join().unwrap(); + th3.join().unwrap(); + }); +} + +#[test] +fn assert_no_unnecessary_polls() { + loom::model(|| { + // // After we poll outer future, woken should reset to false + let rt = Builder::new_current_thread().build().unwrap(); + let (tx, rx) = oneshot::channel(); + let pending_cnt = Arc::new(AtomicUsize::new(0)); + + rt.spawn(async move { + for _ in 0..24 { + task::yield_now().await; + } + tx.send(()).unwrap(); + }); + + let pending_cnt_clone = pending_cnt.clone(); + rt.block_on(async move { + // use task::yield_now() to ensure woken set to true + // ResetFuture will be polled at most once + // Here comes two cases + // 1. recv no message from channel, ResetFuture will be polled + // but get Pending and we record ResetFuture.pending_cnt ++. + // Then when message arrive, ResetFuture returns Ready. So we + // expect ResetFuture.pending_cnt = 1 + // 2. recv message from channel, ResetFuture returns Ready immediately. + // We expect ResetFuture.pending_cnt = 0 + task::yield_now().await; + ResetFuture { + rx, + pending_cnt: pending_cnt_clone, + } + .await; + }); + + let pending_cnt = pending_cnt.load(Acquire); + assert!(pending_cnt <= 1); + }); +} + +struct BlockedFuture { + rx: Receiver<()>, + num_polls: Arc<AtomicUsize>, +} + +impl Future for BlockedFuture { + type Output = (); + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + self.num_polls.fetch_add(1, Release); + + match Pin::new(&mut self.rx).poll(cx) { + Poll::Pending => Poll::Pending, + _ => Poll::Ready(()), + } + } +} + +struct ResetFuture { + rx: Receiver<()>, + pending_cnt: Arc<AtomicUsize>, +} + +impl Future for ResetFuture { + type Output = (); + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + match Pin::new(&mut self.rx).poll(cx) { + Poll::Pending => { + self.pending_cnt.fetch_add(1, Release); + Poll::Pending + } + _ => Poll::Ready(()), + } + } +} diff --git a/vendor/tokio/src/runtime/tests/loom_join_set.rs b/vendor/tokio/src/runtime/tests/loom_join_set.rs new file mode 100644 index 000000000..bd343876a --- /dev/null +++ b/vendor/tokio/src/runtime/tests/loom_join_set.rs @@ -0,0 +1,82 @@ +use crate::runtime::Builder; +use crate::task::JoinSet; + +#[test] +fn test_join_set() { + loom::model(|| { + let rt = Builder::new_multi_thread() + .worker_threads(1) + .build() + .unwrap(); + let mut set = JoinSet::new(); + + rt.block_on(async { + assert_eq!(set.len(), 0); + set.spawn(async { () }); + assert_eq!(set.len(), 1); + set.spawn(async { () }); + assert_eq!(set.len(), 2); + let () = set.join_next().await.unwrap().unwrap(); + assert_eq!(set.len(), 1); + set.spawn(async { () }); + assert_eq!(set.len(), 2); + let () = set.join_next().await.unwrap().unwrap(); + assert_eq!(set.len(), 1); + let () = set.join_next().await.unwrap().unwrap(); + assert_eq!(set.len(), 0); + set.spawn(async { () }); + assert_eq!(set.len(), 1); + }); + + drop(set); + drop(rt); + }); +} + +#[test] +fn abort_all_during_completion() { + use std::sync::{ + atomic::{AtomicBool, Ordering::SeqCst}, + Arc, + }; + + // These booleans assert that at least one execution had the task complete first, and that at + // least one execution had the task be cancelled before it completed. + let complete_happened = Arc::new(AtomicBool::new(false)); + let cancel_happened = Arc::new(AtomicBool::new(false)); + + { + let complete_happened = complete_happened.clone(); + let cancel_happened = cancel_happened.clone(); + loom::model(move || { + let rt = Builder::new_multi_thread() + .worker_threads(1) + .build() + .unwrap(); + + let mut set = JoinSet::new(); + + rt.block_on(async { + set.spawn(async { () }); + set.abort_all(); + + match set.join_next().await { + Some(Ok(())) => complete_happened.store(true, SeqCst), + Some(Err(err)) if err.is_cancelled() => cancel_happened.store(true, SeqCst), + Some(Err(err)) => panic!("fail: {}", err), + None => { + unreachable!("Aborting the task does not remove it from the JoinSet.") + } + } + + assert!(matches!(set.join_next().await, None)); + }); + + drop(set); + drop(rt); + }); + } + + assert!(complete_happened.load(SeqCst)); + assert!(cancel_happened.load(SeqCst)); +} diff --git a/vendor/tokio/src/runtime/tests/loom_pool.rs b/vendor/tokio/src/runtime/tests/loom_pool.rs index 06ad6412f..fb42e1eb4 100644 --- a/vendor/tokio/src/runtime/tests/loom_pool.rs +++ b/vendor/tokio/src/runtime/tests/loom_pool.rs @@ -11,7 +11,7 @@ use crate::{spawn, task}; use tokio_test::assert_ok; use loom::sync::atomic::{AtomicBool, AtomicUsize}; -use loom::sync::{Arc, Mutex}; +use loom::sync::Arc; use pin_project_lite::pin_project; use std::future::Future; @@ -19,6 +19,57 @@ use std::pin::Pin; use std::sync::atomic::Ordering::{Relaxed, SeqCst}; use std::task::{Context, Poll}; +mod atomic_take { + use loom::sync::atomic::AtomicBool; + use std::mem::MaybeUninit; + use std::sync::atomic::Ordering::SeqCst; + + pub(super) struct AtomicTake<T> { + inner: MaybeUninit<T>, + taken: AtomicBool, + } + + impl<T> AtomicTake<T> { + pub(super) fn new(value: T) -> Self { + Self { + inner: MaybeUninit::new(value), + taken: AtomicBool::new(false), + } + } + + pub(super) fn take(&self) -> Option<T> { + // safety: Only one thread will see the boolean change from false + // to true, so that thread is able to take the value. + match self.taken.fetch_or(true, SeqCst) { + false => unsafe { Some(std::ptr::read(self.inner.as_ptr())) }, + true => None, + } + } + } + + impl<T> Drop for AtomicTake<T> { + fn drop(&mut self) { + drop(self.take()); + } + } +} + +#[derive(Clone)] +struct AtomicOneshot<T> { + value: std::sync::Arc<atomic_take::AtomicTake<oneshot::Sender<T>>>, +} +impl<T> AtomicOneshot<T> { + fn new(sender: oneshot::Sender<T>) -> Self { + Self { + value: std::sync::Arc::new(atomic_take::AtomicTake::new(sender)), + } + } + + fn assert_send(&self, value: T) { + self.value.take().unwrap().send(value); + } +} + /// Tests are divided into groups to make the runs faster on CI. mod group_a { use super::*; @@ -52,7 +103,7 @@ mod group_a { let c1 = Arc::new(AtomicUsize::new(0)); let (tx, rx) = oneshot::channel(); - let tx1 = Arc::new(Mutex::new(Some(tx))); + let tx1 = AtomicOneshot::new(tx); // Spawn a task let c2 = c1.clone(); @@ -60,7 +111,7 @@ mod group_a { pool.spawn(track(async move { spawn(track(async move { if 1 == c1.fetch_add(1, Relaxed) { - tx1.lock().unwrap().take().unwrap().send(()); + tx1.assert_send(()); } })); })); @@ -69,7 +120,7 @@ mod group_a { pool.spawn(track(async move { spawn(track(async move { if 1 == c2.fetch_add(1, Relaxed) { - tx2.lock().unwrap().take().unwrap().send(()); + tx2.assert_send(()); } })); })); @@ -119,7 +170,7 @@ mod group_b { let (block_tx, block_rx) = oneshot::channel(); let (done_tx, done_rx) = oneshot::channel(); - let done_tx = Arc::new(Mutex::new(Some(done_tx))); + let done_tx = AtomicOneshot::new(done_tx); pool.spawn(track(async move { crate::task::block_in_place(move || { @@ -136,7 +187,7 @@ mod group_b { pool.spawn(track(async move { if NUM == cnt.fetch_add(1, Relaxed) + 1 { - done_tx.lock().unwrap().take().unwrap().send(()); + done_tx.assert_send(()); } })); } @@ -159,23 +210,6 @@ mod group_b { } #[test] - fn pool_shutdown() { - loom::model(|| { - let pool = mk_pool(2); - - pool.spawn(track(async move { - gated2(true).await; - })); - - pool.spawn(track(async move { - gated2(false).await; - })); - - drop(pool); - }); - } - - #[test] fn join_output() { loom::model(|| { let rt = mk_pool(1); @@ -223,10 +257,6 @@ mod group_b { }); }); } -} - -mod group_c { - use super::*; #[test] fn shutdown_with_notification() { @@ -255,6 +285,27 @@ mod group_c { } } +mod group_c { + use super::*; + + #[test] + fn pool_shutdown() { + loom::model(|| { + let pool = mk_pool(2); + + pool.spawn(track(async move { + gated2(true).await; + })); + + pool.spawn(track(async move { + gated2(false).await; + })); + + drop(pool); + }); + } +} + mod group_d { use super::*; @@ -266,27 +317,25 @@ mod group_d { let c1 = Arc::new(AtomicUsize::new(0)); let (done_tx, done_rx) = oneshot::channel(); - let done_tx1 = Arc::new(Mutex::new(Some(done_tx))); + let done_tx1 = AtomicOneshot::new(done_tx); + let done_tx2 = done_tx1.clone(); // Spawn a task let c2 = c1.clone(); - let done_tx2 = done_tx1.clone(); pool.spawn(track(async move { - gated().await; - gated().await; + multi_gated().await; if 1 == c1.fetch_add(1, Relaxed) { - done_tx1.lock().unwrap().take().unwrap().send(()); + done_tx1.assert_send(()); } })); // Spawn a second task pool.spawn(track(async move { - gated().await; - gated().await; + multi_gated().await; if 1 == c2.fetch_add(1, Relaxed) { - done_tx2.lock().unwrap().take().unwrap().send(()); + done_tx2.assert_send(()); } })); @@ -298,14 +347,12 @@ mod group_d { fn mk_pool(num_threads: usize) -> Runtime { runtime::Builder::new_multi_thread() .worker_threads(num_threads) + // Set the intervals to avoid tuning logic + .event_interval(2) .build() .unwrap() } -fn gated() -> impl Future<Output = &'static str> { - gated2(false) -} - fn gated2(thread: bool) -> impl Future<Output = &'static str> { use loom::thread; use std::sync::Arc; @@ -343,6 +390,38 @@ fn gated2(thread: bool) -> impl Future<Output = &'static str> { }) } +async fn multi_gated() { + struct Gate { + waker: loom::future::AtomicWaker, + count: AtomicUsize, + } + + let gate = Arc::new(Gate { + waker: loom::future::AtomicWaker::new(), + count: AtomicUsize::new(0), + }); + + { + let gate = gate.clone(); + spawn(track(async move { + for i in 1..3 { + gate.count.store(i, SeqCst); + gate.waker.wake(); + } + })); + } + + poll_fn(move |cx| { + if gate.count.load(SeqCst) < 2 { + gate.waker.register_by_ref(cx.waker()); + Poll::Pending + } else { + Poll::Ready(()) + } + }) + .await; +} + fn track<T: Future>(f: T) -> Track<T> { Track { inner: f, diff --git a/vendor/tokio/src/runtime/tests/loom_queue.rs b/vendor/tokio/src/runtime/tests/loom_queue.rs index 34da7fd66..b60e039b9 100644 --- a/vendor/tokio/src/runtime/tests/loom_queue.rs +++ b/vendor/tokio/src/runtime/tests/loom_queue.rs @@ -1,20 +1,27 @@ -use crate::runtime::queue; -use crate::runtime::task::{self, Schedule, Task}; +use crate::runtime::scheduler::multi_thread::{queue, Stats}; +use crate::runtime::tests::NoopSchedule; use loom::thread; +use std::cell::RefCell; + +fn new_stats() -> Stats { + Stats::new(&crate::runtime::WorkerMetrics::new()) +} #[test] fn basic() { loom::model(|| { let (steal, mut local) = queue::local(); - let inject = queue::Inject::new(); + let inject = RefCell::new(vec![]); + let mut stats = new_stats(); let th = thread::spawn(move || { + let mut stats = new_stats(); let (_, mut local) = queue::local(); let mut n = 0; for _ in 0..3 { - if steal.steal_into(&mut local).is_some() { + if steal.steal_into(&mut local, &mut stats).is_some() { n += 1; } @@ -30,8 +37,8 @@ fn basic() { for _ in 0..2 { for _ in 0..2 { - let (task, _) = super::joinable::<_, Runtime>(async {}); - local.push_back(task, &inject); + let (task, _) = super::unowned(async {}); + local.push_back_or_overflow(task, &inject, &mut stats); } if local.pop().is_some() { @@ -39,17 +46,15 @@ fn basic() { } // Push another task - let (task, _) = super::joinable::<_, Runtime>(async {}); - local.push_back(task, &inject); + let (task, _) = super::unowned(async {}); + local.push_back_or_overflow(task, &inject, &mut stats); while local.pop().is_some() { n += 1; } } - while inject.pop().is_some() { - n += 1; - } + n += inject.borrow_mut().drain(..).count(); n += th.join().unwrap(); @@ -61,13 +66,15 @@ fn basic() { fn steal_overflow() { loom::model(|| { let (steal, mut local) = queue::local(); - let inject = queue::Inject::new(); + let inject = RefCell::new(vec![]); + let mut stats = new_stats(); let th = thread::spawn(move || { + let mut stats = new_stats(); let (_, mut local) = queue::local(); let mut n = 0; - if steal.steal_into(&mut local).is_some() { + if steal.steal_into(&mut local, &mut stats).is_some() { n += 1; } @@ -81,16 +88,16 @@ fn steal_overflow() { let mut n = 0; // push a task, pop a task - let (task, _) = super::joinable::<_, Runtime>(async {}); - local.push_back(task, &inject); + let (task, _) = super::unowned(async {}); + local.push_back_or_overflow(task, &inject, &mut stats); if local.pop().is_some() { n += 1; } for _ in 0..6 { - let (task, _) = super::joinable::<_, Runtime>(async {}); - local.push_back(task, &inject); + let (task, _) = super::unowned(async {}); + local.push_back_or_overflow(task, &inject, &mut stats); } n += th.join().unwrap(); @@ -99,9 +106,7 @@ fn steal_overflow() { n += 1; } - while inject.pop().is_some() { - n += 1; - } + n += inject.borrow_mut().drain(..).count(); assert_eq!(7, n); }); @@ -111,10 +116,11 @@ fn steal_overflow() { fn multi_stealer() { const NUM_TASKS: usize = 5; - fn steal_tasks(steal: queue::Steal<Runtime>) -> usize { + fn steal_tasks(steal: queue::Steal<NoopSchedule>) -> usize { + let mut stats = new_stats(); let (_, mut local) = queue::local(); - if steal.steal_into(&mut local).is_none() { + if steal.steal_into(&mut local, &mut stats).is_none() { return 0; } @@ -129,12 +135,13 @@ fn multi_stealer() { loom::model(|| { let (steal, mut local) = queue::local(); - let inject = queue::Inject::new(); + let inject = RefCell::new(vec![]); + let mut stats = new_stats(); // Push work for _ in 0..NUM_TASKS { - let (task, _) = super::joinable::<_, Runtime>(async {}); - local.push_back(task, &inject); + let (task, _) = super::unowned(async {}); + local.push_back_or_overflow(task, &inject, &mut stats); } let th1 = { @@ -150,9 +157,7 @@ fn multi_stealer() { n += 1; } - while inject.pop().is_some() { - n += 1; - } + n += inject.borrow_mut().drain(..).count(); n += th1.join().unwrap(); n += th2.join().unwrap(); @@ -164,23 +169,25 @@ fn multi_stealer() { #[test] fn chained_steal() { loom::model(|| { + let mut stats = new_stats(); let (s1, mut l1) = queue::local(); let (s2, mut l2) = queue::local(); - let inject = queue::Inject::new(); + let inject = RefCell::new(vec![]); // Load up some tasks for _ in 0..4 { - let (task, _) = super::joinable::<_, Runtime>(async {}); - l1.push_back(task, &inject); + let (task, _) = super::unowned(async {}); + l1.push_back_or_overflow(task, &inject, &mut stats); - let (task, _) = super::joinable::<_, Runtime>(async {}); - l2.push_back(task, &inject); + let (task, _) = super::unowned(async {}); + l2.push_back_or_overflow(task, &inject, &mut stats); } // Spawn a task to steal from **our** queue let th = thread::spawn(move || { + let mut stats = new_stats(); let (_, mut local) = queue::local(); - s1.steal_into(&mut local); + s1.steal_into(&mut local, &mut stats); while local.pop().is_some() {} }); @@ -188,29 +195,11 @@ fn chained_steal() { // Drain our tasks, then attempt to steal while l1.pop().is_some() {} - s2.steal_into(&mut l1); + s2.steal_into(&mut l1, &mut stats); th.join().unwrap(); while l1.pop().is_some() {} while l2.pop().is_some() {} - while inject.pop().is_some() {} }); } - -struct Runtime; - -impl Schedule for Runtime { - fn bind(task: Task<Self>) -> Runtime { - std::mem::forget(task); - Runtime - } - - fn release(&self, _task: &Task<Self>) -> Option<Task<Self>> { - None - } - - fn schedule(&self, _task: task::Notified<Self>) { - unreachable!(); - } -} diff --git a/vendor/tokio/src/runtime/tests/loom_yield.rs b/vendor/tokio/src/runtime/tests/loom_yield.rs new file mode 100644 index 000000000..ba506e5a4 --- /dev/null +++ b/vendor/tokio/src/runtime/tests/loom_yield.rs @@ -0,0 +1,37 @@ +use crate::runtime::park; +use crate::runtime::tests::loom_oneshot as oneshot; +use crate::runtime::{self, Runtime}; + +#[test] +fn yield_calls_park_before_scheduling_again() { + // Don't need to check all permutations + let mut loom = loom::model::Builder::default(); + loom.max_permutations = Some(1); + loom.check(|| { + let rt = mk_runtime(2); + let (tx, rx) = oneshot::channel::<()>(); + + rt.spawn(async { + let tid = loom::thread::current().id(); + let park_count = park::current_thread_park_count(); + + crate::task::yield_now().await; + + if tid == loom::thread::current().id() { + let new_park_count = park::current_thread_park_count(); + assert_eq!(park_count + 1, new_park_count); + } + + tx.send(()); + }); + + rx.recv(); + }); +} + +fn mk_runtime(num_threads: usize) -> Runtime { + runtime::Builder::new_multi_thread() + .worker_threads(num_threads) + .build() + .unwrap() +} diff --git a/vendor/tokio/src/runtime/tests/mod.rs b/vendor/tokio/src/runtime/tests/mod.rs index 3f2cc9825..b12a76e26 100644 --- a/vendor/tokio/src/runtime/tests/mod.rs +++ b/vendor/tokio/src/runtime/tests/mod.rs @@ -1,35 +1,73 @@ -#[cfg(not(all(tokio_unstable, feature = "tracing")))] -use crate::runtime::task::joinable; +// Enable dead_code / unreachable_pub here. It has been disabled in lib.rs for +// other code when running loom tests. +#![cfg_attr(loom, warn(dead_code, unreachable_pub))] -#[cfg(all(tokio_unstable, feature = "tracing"))] -use self::joinable_wrapper::joinable; +use self::noop_scheduler::NoopSchedule; +use self::unowned_wrapper::unowned; -#[cfg(all(tokio_unstable, feature = "tracing"))] -mod joinable_wrapper { - use crate::runtime::task::{JoinHandle, Notified, Schedule}; - use tracing::Instrument; +mod noop_scheduler { + use crate::runtime::task::{self, Task}; - pub(crate) fn joinable<T, S>(task: T) -> (Notified<S>, JoinHandle<T::Output>) + /// `task::Schedule` implementation that does nothing, for testing. + pub(crate) struct NoopSchedule; + + impl task::Schedule for NoopSchedule { + fn release(&self, _task: &Task<Self>) -> Option<Task<Self>> { + None + } + + fn schedule(&self, _task: task::Notified<Self>) { + unreachable!(); + } + } +} + +mod unowned_wrapper { + use crate::runtime::task::{Id, JoinHandle, Notified}; + use crate::runtime::tests::NoopSchedule; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + pub(crate) fn unowned<T>(task: T) -> (Notified<NoopSchedule>, JoinHandle<T::Output>) where T: std::future::Future + Send + 'static, - S: Schedule, + T::Output: Send + 'static, { + use tracing::Instrument; let span = tracing::trace_span!("test_span"); - crate::runtime::task::joinable(task.instrument(span)) + let task = task.instrument(span); + let (task, handle) = crate::runtime::task::unowned(task, NoopSchedule, Id::next()); + (task.into_notified(), handle) + } + + #[cfg(not(all(tokio_unstable, feature = "tracing")))] + pub(crate) fn unowned<T>(task: T) -> (Notified<NoopSchedule>, JoinHandle<T::Output>) + where + T: std::future::Future + Send + 'static, + T::Output: Send + 'static, + { + let (task, handle) = crate::runtime::task::unowned(task, NoopSchedule, Id::next()); + (task.into_notified(), handle) } } cfg_loom! { - mod loom_basic_scheduler; - mod loom_local; mod loom_blocking; + mod loom_current_thread_scheduler; + mod loom_local; mod loom_oneshot; mod loom_pool; mod loom_queue; mod loom_shutdown_join; + mod loom_join_set; + mod loom_yield; + + // Make sure debug assertions are enabled + #[cfg(not(debug_assertions))] + compiler_error!("these tests require debug assertions to be enabled"); } cfg_not_loom! { + mod inject; mod queue; #[cfg(not(miri))] diff --git a/vendor/tokio/src/runtime/tests/queue.rs b/vendor/tokio/src/runtime/tests/queue.rs index b2962f154..5df92b7a2 100644 --- a/vendor/tokio/src/runtime/tests/queue.rs +++ b/vendor/tokio/src/runtime/tests/queue.rs @@ -1,39 +1,107 @@ -use crate::runtime::queue; +use crate::runtime::scheduler::multi_thread::{queue, Stats}; use crate::runtime::task::{self, Schedule, Task}; +use std::cell::RefCell; use std::thread; use std::time::Duration; +#[allow(unused)] +macro_rules! assert_metrics { + ($stats:ident, $field:ident == $v:expr) => {{ + use crate::runtime::WorkerMetrics; + use std::sync::atomic::Ordering::Relaxed; + + let worker = WorkerMetrics::new(); + $stats.submit(&worker); + + let expect = $v; + let actual = worker.$field.load(Relaxed); + + assert!(actual == expect, "expect = {}; actual = {}", expect, actual) + }}; +} + +fn new_stats() -> Stats { + use crate::runtime::WorkerMetrics; + Stats::new(&WorkerMetrics::new()) +} + #[test] -fn fits_256() { +fn fits_256_one_at_a_time() { let (_, mut local) = queue::local(); - let inject = queue::Inject::new(); + let inject = RefCell::new(vec![]); + let mut stats = new_stats(); for _ in 0..256 { - let (task, _) = super::joinable::<_, Runtime>(async {}); - local.push_back(task, &inject); + let (task, _) = super::unowned(async {}); + local.push_back_or_overflow(task, &inject, &mut stats); } - assert!(inject.pop().is_none()); + cfg_metrics! { + assert_metrics!(stats, overflow_count == 0); + } + + assert!(inject.borrow_mut().pop().is_none()); while local.pop().is_some() {} } #[test] +fn fits_256_all_at_once() { + let (_, mut local) = queue::local(); + + let mut tasks = (0..256) + .map(|_| super::unowned(async {}).0) + .collect::<Vec<_>>(); + local.push_back(tasks.drain(..)); + + let mut i = 0; + while local.pop().is_some() { + i += 1; + } + + assert_eq!(i, 256); +} + +#[test] +fn fits_256_all_in_chunks() { + let (_, mut local) = queue::local(); + + let mut tasks = (0..256) + .map(|_| super::unowned(async {}).0) + .collect::<Vec<_>>(); + + local.push_back(tasks.drain(..10)); + local.push_back(tasks.drain(..100)); + local.push_back(tasks.drain(..46)); + local.push_back(tasks.drain(..100)); + + let mut i = 0; + while local.pop().is_some() { + i += 1; + } + + assert_eq!(i, 256); +} + +#[test] fn overflow() { let (_, mut local) = queue::local(); - let inject = queue::Inject::new(); + let inject = RefCell::new(vec![]); + let mut stats = new_stats(); for _ in 0..257 { - let (task, _) = super::joinable::<_, Runtime>(async {}); - local.push_back(task, &inject); + let (task, _) = super::unowned(async {}); + local.push_back_or_overflow(task, &inject, &mut stats); + } + + cfg_metrics! { + assert_metrics!(stats, overflow_count == 1); } let mut n = 0; - while inject.pop().is_some() { - n += 1; - } + n += inject.borrow_mut().drain(..).count(); while local.pop().is_some() { n += 1; @@ -44,16 +112,22 @@ fn overflow() { #[test] fn steal_batch() { + let mut stats = new_stats(); + let (steal1, mut local1) = queue::local(); let (_, mut local2) = queue::local(); - let inject = queue::Inject::new(); + let inject = RefCell::new(vec![]); for _ in 0..4 { - let (task, _) = super::joinable::<_, Runtime>(async {}); - local1.push_back(task, &inject); + let (task, _) = super::unowned(async {}); + local1.push_back_or_overflow(task, &inject, &mut stats); } - assert!(steal1.steal_into(&mut local2).is_some()); + assert!(steal1.steal_into(&mut local2, &mut stats).is_some()); + + cfg_metrics! { + assert_metrics!(stats, steal_count == 2); + } for _ in 0..1 { assert!(local2.pop().is_some()); @@ -68,24 +142,35 @@ fn steal_batch() { assert!(local1.pop().is_none()); } +const fn normal_or_miri(normal: usize, miri: usize) -> usize { + if cfg!(miri) { + miri + } else { + normal + } +} + #[test] fn stress1() { - const NUM_ITER: usize = 1; - const NUM_STEAL: usize = 1_000; - const NUM_LOCAL: usize = 1_000; - const NUM_PUSH: usize = 500; - const NUM_POP: usize = 250; + const NUM_ITER: usize = 5; + const NUM_STEAL: usize = normal_or_miri(1_000, 10); + const NUM_LOCAL: usize = normal_or_miri(1_000, 10); + const NUM_PUSH: usize = normal_or_miri(500, 10); + const NUM_POP: usize = normal_or_miri(250, 10); + + let mut stats = new_stats(); for _ in 0..NUM_ITER { let (steal, mut local) = queue::local(); - let inject = queue::Inject::new(); + let inject = RefCell::new(vec![]); let th = thread::spawn(move || { + let mut stats = new_stats(); let (_, mut local) = queue::local(); let mut n = 0; for _ in 0..NUM_STEAL { - if steal.steal_into(&mut local).is_some() { + if steal.steal_into(&mut local, &mut stats).is_some() { n += 1; } @@ -96,6 +181,10 @@ fn stress1() { thread::yield_now(); } + cfg_metrics! { + assert_metrics!(stats, steal_count == n as _); + } + n }); @@ -103,8 +192,8 @@ fn stress1() { for _ in 0..NUM_LOCAL { for _ in 0..NUM_PUSH { - let (task, _) = super::joinable::<_, Runtime>(async {}); - local.push_back(task, &inject); + let (task, _) = super::unowned(async {}); + local.push_back_or_overflow(task, &inject, &mut stats); } for _ in 0..NUM_POP { @@ -116,9 +205,7 @@ fn stress1() { } } - while inject.pop().is_some() { - n += 1; - } + n += inject.borrow_mut().drain(..).count(); n += th.join().unwrap(); @@ -129,19 +216,22 @@ fn stress1() { #[test] fn stress2() { const NUM_ITER: usize = 1; - const NUM_TASKS: usize = 1_000_000; - const NUM_STEAL: usize = 1_000; + const NUM_TASKS: usize = normal_or_miri(1_000_000, 50); + const NUM_STEAL: usize = normal_or_miri(1_000, 10); + + let mut stats = new_stats(); for _ in 0..NUM_ITER { let (steal, mut local) = queue::local(); - let inject = queue::Inject::new(); + let inject = RefCell::new(vec![]); let th = thread::spawn(move || { + let mut stats = new_stats(); let (_, mut local) = queue::local(); let mut n = 0; for _ in 0..NUM_STEAL { - if steal.steal_into(&mut local).is_some() { + if steal.steal_into(&mut local, &mut stats).is_some() { n += 1; } @@ -158,16 +248,14 @@ fn stress2() { let mut num_pop = 0; for i in 0..NUM_TASKS { - let (task, _) = super::joinable::<_, Runtime>(async {}); - local.push_back(task, &inject); + let (task, _) = super::unowned(async {}); + local.push_back_or_overflow(task, &inject, &mut stats); if i % 128 == 0 && local.pop().is_some() { num_pop += 1; } - while inject.pop().is_some() { - num_pop += 1; - } + num_pop += inject.borrow_mut().drain(..).count(); } num_pop += th.join().unwrap(); @@ -176,9 +264,7 @@ fn stress2() { num_pop += 1; } - while inject.pop().is_some() { - num_pop += 1; - } + num_pop += inject.borrow_mut().drain(..).count(); assert_eq!(num_pop, NUM_TASKS); } @@ -187,11 +273,6 @@ fn stress2() { struct Runtime; impl Schedule for Runtime { - fn bind(task: Task<Self>) -> Runtime { - std::mem::forget(task); - Runtime - } - fn release(&self, _task: &Task<Self>) -> Option<Task<Self>> { None } diff --git a/vendor/tokio/src/runtime/tests/task.rs b/vendor/tokio/src/runtime/tests/task.rs index 7c2012523..a79c0f50d 100644 --- a/vendor/tokio/src/runtime/tests/task.rs +++ b/vendor/tokio/src/runtime/tests/task.rs @@ -1,44 +1,235 @@ -use crate::runtime::task::{self, Schedule, Task}; -use crate::util::linked_list::{Link, LinkedList}; +use crate::runtime::task::{self, unowned, Id, JoinHandle, OwnedTasks, Schedule, Task}; +use crate::runtime::tests::NoopSchedule; use crate::util::TryLock; use std::collections::VecDeque; +use std::future::Future; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; +struct AssertDropHandle { + is_dropped: Arc<AtomicBool>, +} +impl AssertDropHandle { + #[track_caller] + fn assert_dropped(&self) { + assert!(self.is_dropped.load(Ordering::SeqCst)); + } + + #[track_caller] + fn assert_not_dropped(&self) { + assert!(!self.is_dropped.load(Ordering::SeqCst)); + } +} + +struct AssertDrop { + is_dropped: Arc<AtomicBool>, +} +impl AssertDrop { + fn new() -> (Self, AssertDropHandle) { + let shared = Arc::new(AtomicBool::new(false)); + ( + AssertDrop { + is_dropped: shared.clone(), + }, + AssertDropHandle { + is_dropped: shared.clone(), + }, + ) + } +} +impl Drop for AssertDrop { + fn drop(&mut self) { + self.is_dropped.store(true, Ordering::SeqCst); + } +} + +// A Notified does not shut down on drop, but it is dropped once the ref-count +// hits zero. +#[test] +fn create_drop1() { + let (ad, handle) = AssertDrop::new(); + let (notified, join) = unowned( + async { + drop(ad); + unreachable!() + }, + NoopSchedule, + Id::next(), + ); + drop(notified); + handle.assert_not_dropped(); + drop(join); + handle.assert_dropped(); +} + #[test] -fn create_drop() { - let _ = super::joinable::<_, Runtime>(async { unreachable!() }); +fn create_drop2() { + let (ad, handle) = AssertDrop::new(); + let (notified, join) = unowned( + async { + drop(ad); + unreachable!() + }, + NoopSchedule, + Id::next(), + ); + drop(join); + handle.assert_not_dropped(); + drop(notified); + handle.assert_dropped(); +} + +#[test] +fn drop_abort_handle1() { + let (ad, handle) = AssertDrop::new(); + let (notified, join) = unowned( + async { + drop(ad); + unreachable!() + }, + NoopSchedule, + Id::next(), + ); + let abort = join.abort_handle(); + drop(join); + handle.assert_not_dropped(); + drop(notified); + handle.assert_not_dropped(); + drop(abort); + handle.assert_dropped(); +} + +#[test] +fn drop_abort_handle2() { + let (ad, handle) = AssertDrop::new(); + let (notified, join) = unowned( + async { + drop(ad); + unreachable!() + }, + NoopSchedule, + Id::next(), + ); + let abort = join.abort_handle(); + drop(notified); + handle.assert_not_dropped(); + drop(abort); + handle.assert_not_dropped(); + drop(join); + handle.assert_dropped(); +} + +// Shutting down through Notified works +#[test] +fn create_shutdown1() { + let (ad, handle) = AssertDrop::new(); + let (notified, join) = unowned( + async { + drop(ad); + unreachable!() + }, + NoopSchedule, + Id::next(), + ); + drop(join); + handle.assert_not_dropped(); + notified.shutdown(); + handle.assert_dropped(); +} + +#[test] +fn create_shutdown2() { + let (ad, handle) = AssertDrop::new(); + let (notified, join) = unowned( + async { + drop(ad); + unreachable!() + }, + NoopSchedule, + Id::next(), + ); + handle.assert_not_dropped(); + notified.shutdown(); + handle.assert_dropped(); + drop(join); +} + +#[test] +fn unowned_poll() { + let (task, _) = unowned(async {}, NoopSchedule, Id::next()); + task.run(); } #[test] fn schedule() { with(|rt| { - let (task, _) = super::joinable(async { + rt.spawn(async { crate::task::yield_now().await; }); - rt.schedule(task); - assert_eq!(2, rt.tick()); + rt.shutdown(); }) } #[test] fn shutdown() { with(|rt| { - let (task, _) = super::joinable(async { + rt.spawn(async { loop { crate::task::yield_now().await; } }); - rt.schedule(task); rt.tick_max(1); rt.shutdown(); }) } +#[test] +fn shutdown_immediately() { + with(|rt| { + rt.spawn(async { + loop { + crate::task::yield_now().await; + } + }); + + rt.shutdown(); + }) +} + +#[test] +fn spawn_during_shutdown() { + static DID_SPAWN: AtomicBool = AtomicBool::new(false); + + struct SpawnOnDrop(Runtime); + impl Drop for SpawnOnDrop { + fn drop(&mut self) { + DID_SPAWN.store(true, Ordering::SeqCst); + self.0.spawn(async {}); + } + } + + with(|rt| { + let rt2 = rt.clone(); + rt.spawn(async move { + let _spawn_on_drop = SpawnOnDrop(rt2); + + loop { + crate::task::yield_now().await; + } + }); + + rt.tick_max(1); + rt.shutdown(); + }); + + assert!(DID_SPAWN.load(Ordering::SeqCst)); +} + fn with(f: impl FnOnce(Runtime)) { struct Reset; @@ -51,10 +242,9 @@ fn with(f: impl FnOnce(Runtime)) { let _reset = Reset; let rt = Runtime(Arc::new(Inner { - released: task::TransferStack::new(), + owned: OwnedTasks::new(), core: TryLock::new(Core { queue: VecDeque::new(), - tasks: LinkedList::new(), }), })); @@ -66,18 +256,31 @@ fn with(f: impl FnOnce(Runtime)) { struct Runtime(Arc<Inner>); struct Inner { - released: task::TransferStack<Runtime>, core: TryLock<Core>, + owned: OwnedTasks<Runtime>, } struct Core { queue: VecDeque<task::Notified<Runtime>>, - tasks: LinkedList<Task<Runtime>, <Task<Runtime> as Link>::Target>, } static CURRENT: TryLock<Option<Runtime>> = TryLock::new(None); impl Runtime { + fn spawn<T>(&self, future: T) -> JoinHandle<T::Output> + where + T: 'static + Send + Future, + T::Output: 'static + Send, + { + let (handle, notified) = self.0.owned.bind(future, self.clone(), Id::next()); + + if let Some(notified) = notified { + self.schedule(notified); + } + + handle + } + fn tick(&self) -> usize { self.tick_max(usize::MAX) } @@ -88,11 +291,10 @@ impl Runtime { while !self.is_empty() && n < max { let task = self.next_task(); n += 1; + let task = self.0.owned.assert_owner(task); task.run(); } - self.0.maintenance(); - n } @@ -107,50 +309,21 @@ impl Runtime { fn shutdown(&self) { let mut core = self.0.core.try_lock().unwrap(); - for task in core.tasks.iter() { - task.shutdown(); - } + self.0.owned.close_and_shutdown_all(); while let Some(task) = core.queue.pop_back() { - task.shutdown(); + drop(task); } drop(core); - while !self.0.core.try_lock().unwrap().tasks.is_empty() { - self.0.maintenance(); - } - } -} - -impl Inner { - fn maintenance(&self) { - use std::mem::ManuallyDrop; - - for task in self.released.drain() { - let task = ManuallyDrop::new(task); - - // safety: see worker.rs - unsafe { - let ptr = task.header().into(); - self.core.try_lock().unwrap().tasks.remove(ptr); - } - } + assert!(self.0.owned.is_empty()); } } impl Schedule for Runtime { - fn bind(task: Task<Self>) -> Runtime { - let rt = CURRENT.try_lock().unwrap().as_ref().unwrap().clone(); - rt.0.core.try_lock().unwrap().tasks.push_front(task); - rt - } - fn release(&self, task: &Task<Self>) -> Option<Task<Self>> { - // safety: copying worker.rs - let task = unsafe { Task::from_raw(task.header().into()) }; - self.0.released.push(task); - None + self.0.owned.remove(task) } fn schedule(&self, task: task::Notified<Self>) { diff --git a/vendor/tokio/src/runtime/tests/task_combinations.rs b/vendor/tokio/src/runtime/tests/task_combinations.rs index 76ce2330c..73a20d976 100644 --- a/vendor/tokio/src/runtime/tests/task_combinations.rs +++ b/vendor/tokio/src/runtime/tests/task_combinations.rs @@ -1,8 +1,10 @@ +use std::fmt; use std::future::Future; use std::panic; use std::pin::Pin; use std::task::{Context, Poll}; +use crate::runtime::task::AbortHandle; use crate::runtime::Builder; use crate::sync::oneshot; use crate::task::JoinHandle; @@ -56,6 +58,12 @@ enum CombiAbort { AbortedAfterConsumeOutput = 4, } +#[derive(Copy, Clone, Debug, PartialEq)] +enum CombiAbortSource { + JoinHandle, + AbortHandle, +} + #[test] fn test_combinations() { let mut rt = &[ @@ -90,6 +98,13 @@ fn test_combinations() { CombiAbort::AbortedAfterFinish, CombiAbort::AbortedAfterConsumeOutput, ]; + let ah = [ + None, + Some(CombiJoinHandle::DropImmediately), + Some(CombiJoinHandle::DropFirstPoll), + Some(CombiJoinHandle::DropAfterNoConsume), + Some(CombiJoinHandle::DropAfterConsume), + ]; for rt in rt.iter().copied() { for ls in ls.iter().copied() { @@ -98,7 +113,34 @@ fn test_combinations() { for ji in ji.iter().copied() { for jh in jh.iter().copied() { for abort in abort.iter().copied() { - test_combination(rt, ls, task, output, ji, jh, abort); + // abort via join handle --- abort handles + // may be dropped at any point + for ah in ah.iter().copied() { + test_combination( + rt, + ls, + task, + output, + ji, + jh, + ah, + abort, + CombiAbortSource::JoinHandle, + ); + } + // if aborting via AbortHandle, it will + // never be dropped. + test_combination( + rt, + ls, + task, + output, + ji, + jh, + None, + abort, + CombiAbortSource::AbortHandle, + ); } } } @@ -108,6 +150,9 @@ fn test_combinations() { } } +fn is_debug<T: fmt::Debug>(_: &T) {} + +#[allow(clippy::too_many_arguments)] fn test_combination( rt: CombiRuntime, ls: CombiLocalSet, @@ -115,12 +160,24 @@ fn test_combination( output: CombiOutput, ji: CombiJoinInterest, jh: CombiJoinHandle, + ah: Option<CombiJoinHandle>, abort: CombiAbort, + abort_src: CombiAbortSource, ) { - if (jh as usize) < (abort as usize) { - // drop before abort not possible - return; + match (abort_src, ah) { + (CombiAbortSource::JoinHandle, _) if (jh as usize) < (abort as usize) => { + // join handle dropped prior to abort + return; + } + (CombiAbortSource::AbortHandle, Some(_)) => { + // abort handle dropped, we can't abort through the + // abort handle + return; + } + + _ => {} } + if (task == CombiTask::PanicOnDrop) && (output == CombiOutput::PanicOnDrop) { // this causes double panic return; @@ -130,7 +187,15 @@ fn test_combination( return; } - println!("Runtime {:?}, LocalSet {:?}, Task {:?}, Output {:?}, JoinInterest {:?}, JoinHandle {:?}, Abort {:?}", rt, ls, task, output, ji, jh, abort); + is_debug(&rt); + is_debug(&ls); + is_debug(&task); + is_debug(&output); + is_debug(&ji); + is_debug(&jh); + is_debug(&ah); + is_debug(&abort); + is_debug(&abort_src); // A runtime optionally with a LocalSet struct Rt { @@ -282,8 +347,24 @@ fn test_combination( ); } + // If we are either aborting the task via an abort handle, or dropping via + // an abort handle, do that now. + let mut abort_handle = if ah.is_some() || abort_src == CombiAbortSource::AbortHandle { + handle.as_ref().map(JoinHandle::abort_handle) + } else { + None + }; + + let do_abort = |abort_handle: &mut Option<AbortHandle>, + join_handle: Option<&mut JoinHandle<_>>| { + match abort_src { + CombiAbortSource::AbortHandle => abort_handle.take().unwrap().abort(), + CombiAbortSource::JoinHandle => join_handle.unwrap().abort(), + } + }; + if abort == CombiAbort::AbortedImmediately { - handle.as_mut().unwrap().abort(); + do_abort(&mut abort_handle, handle.as_mut()); aborted = true; } if jh == CombiJoinHandle::DropImmediately { @@ -301,12 +382,15 @@ fn test_combination( } if abort == CombiAbort::AbortedFirstPoll { - handle.as_mut().unwrap().abort(); + do_abort(&mut abort_handle, handle.as_mut()); aborted = true; } if jh == CombiJoinHandle::DropFirstPoll { drop(handle.take().unwrap()); } + if ah == Some(CombiJoinHandle::DropFirstPoll) { + drop(abort_handle.take().unwrap()); + } // Signal the future that it can return now let _ = on_complete.send(()); @@ -318,23 +402,42 @@ fn test_combination( if abort == CombiAbort::AbortedAfterFinish { // Don't set aborted to true here as the task already finished - handle.as_mut().unwrap().abort(); + do_abort(&mut abort_handle, handle.as_mut()); } if jh == CombiJoinHandle::DropAfterNoConsume { - // The runtime will usually have dropped every ref-count at this point, - // in which case dropping the JoinHandle drops the output. - // - // (But it might race and still hold a ref-count) - let panic = panic::catch_unwind(panic::AssertUnwindSafe(|| { + if ah == Some(CombiJoinHandle::DropAfterNoConsume) { drop(handle.take().unwrap()); - })); - if panic.is_err() { - assert!( - (output == CombiOutput::PanicOnDrop) - && (!matches!(task, CombiTask::PanicOnRun | CombiTask::PanicOnRunAndDrop)) - && !aborted, - "Dropping JoinHandle shouldn't panic here" - ); + // The runtime will usually have dropped every ref-count at this point, + // in which case dropping the AbortHandle drops the output. + // + // (But it might race and still hold a ref-count) + let panic = panic::catch_unwind(panic::AssertUnwindSafe(|| { + drop(abort_handle.take().unwrap()); + })); + if panic.is_err() { + assert!( + (output == CombiOutput::PanicOnDrop) + && (!matches!(task, CombiTask::PanicOnRun | CombiTask::PanicOnRunAndDrop)) + && !aborted, + "Dropping AbortHandle shouldn't panic here" + ); + } + } else { + // The runtime will usually have dropped every ref-count at this point, + // in which case dropping the JoinHandle drops the output. + // + // (But it might race and still hold a ref-count) + let panic = panic::catch_unwind(panic::AssertUnwindSafe(|| { + drop(handle.take().unwrap()); + })); + if panic.is_err() { + assert!( + (output == CombiOutput::PanicOnDrop) + && (!matches!(task, CombiTask::PanicOnRun | CombiTask::PanicOnRunAndDrop)) + && !aborted, + "Dropping JoinHandle shouldn't panic here" + ); + } } } @@ -362,11 +465,15 @@ fn test_combination( _ => unreachable!(), } - let handle = handle.take().unwrap(); + let mut handle = handle.take().unwrap(); if abort == CombiAbort::AbortedAfterConsumeOutput { - handle.abort(); + do_abort(&mut abort_handle, Some(&mut handle)); } drop(handle); + + if ah == Some(CombiJoinHandle::DropAfterConsume) { + drop(abort_handle.take()); + } } // The output should have been dropped now. Check whether the output diff --git a/vendor/tokio/src/runtime/thread_id.rs b/vendor/tokio/src/runtime/thread_id.rs new file mode 100644 index 000000000..ef3928979 --- /dev/null +++ b/vendor/tokio/src/runtime/thread_id.rs @@ -0,0 +1,31 @@ +use std::num::NonZeroU64; + +#[derive(Eq, PartialEq, Clone, Copy, Hash, Debug)] +pub(crate) struct ThreadId(NonZeroU64); + +impl ThreadId { + pub(crate) fn next() -> Self { + use crate::loom::sync::atomic::{Ordering::Relaxed, StaticAtomicU64}; + + static NEXT_ID: StaticAtomicU64 = StaticAtomicU64::new(0); + + let mut last = NEXT_ID.load(Relaxed); + loop { + let id = match last.checked_add(1) { + Some(id) => id, + None => exhausted(), + }; + + match NEXT_ID.compare_exchange_weak(last, id, Relaxed, Relaxed) { + Ok(_) => return ThreadId(NonZeroU64::new(id).unwrap()), + Err(id) => last = id, + } + } + } +} + +#[cold] +#[allow(dead_code)] +fn exhausted() -> ! { + panic!("failed to generate unique thread ID: bitspace exhausted") +} diff --git a/vendor/tokio/src/runtime/thread_pool/mod.rs b/vendor/tokio/src/runtime/thread_pool/mod.rs deleted file mode 100644 index 96312d346..000000000 --- a/vendor/tokio/src/runtime/thread_pool/mod.rs +++ /dev/null @@ -1,116 +0,0 @@ -//! Threadpool - -mod atomic_cell; -use atomic_cell::AtomicCell; - -mod idle; -use self::idle::Idle; - -mod worker; -pub(crate) use worker::Launch; - -pub(crate) use worker::block_in_place; - -use crate::loom::sync::Arc; -use crate::runtime::task::{self, JoinHandle}; -use crate::runtime::Parker; - -use std::fmt; -use std::future::Future; - -/// Work-stealing based thread pool for executing futures. -pub(crate) struct ThreadPool { - spawner: Spawner, -} - -/// Submit futures to the associated thread pool for execution. -/// -/// A `Spawner` instance is a handle to a single thread pool that allows the owner -/// of the handle to spawn futures onto the thread pool. -/// -/// The `Spawner` handle is *only* used for spawning new futures. It does not -/// impact the lifecycle of the thread pool in any way. The thread pool may -/// shutdown while there are outstanding `Spawner` instances. -/// -/// `Spawner` instances are obtained by calling [`ThreadPool::spawner`]. -/// -/// [`ThreadPool::spawner`]: method@ThreadPool::spawner -#[derive(Clone)] -pub(crate) struct Spawner { - shared: Arc<worker::Shared>, -} - -// ===== impl ThreadPool ===== - -impl ThreadPool { - pub(crate) fn new(size: usize, parker: Parker) -> (ThreadPool, Launch) { - let (shared, launch) = worker::create(size, parker); - let spawner = Spawner { shared }; - let thread_pool = ThreadPool { spawner }; - - (thread_pool, launch) - } - - /// Returns reference to `Spawner`. - /// - /// The `Spawner` handle can be cloned and enables spawning tasks from other - /// threads. - pub(crate) fn spawner(&self) -> &Spawner { - &self.spawner - } - - /// Blocks the current thread waiting for the future to complete. - /// - /// The future will execute on the current thread, but all spawned tasks - /// will be executed on the thread pool. - pub(crate) fn block_on<F>(&self, future: F) -> F::Output - where - F: Future, - { - let mut enter = crate::runtime::enter(true); - enter.block_on(future).expect("failed to park thread") - } -} - -impl fmt::Debug for ThreadPool { - fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - fmt.debug_struct("ThreadPool").finish() - } -} - -impl Drop for ThreadPool { - fn drop(&mut self) { - self.spawner.shutdown(); - } -} - -// ==== impl Spawner ===== - -impl Spawner { - /// Spawns a future onto the thread pool - pub(crate) fn spawn<F>(&self, future: F) -> JoinHandle<F::Output> - where - F: crate::future::Future + Send + 'static, - F::Output: Send + 'static, - { - let (task, handle) = task::joinable(future); - - if let Err(task) = self.shared.schedule(task, false) { - // The newly spawned task could not be scheduled because the runtime - // is shutting down. The task must be explicitly shutdown at this point. - task.shutdown(); - } - - handle - } - - pub(crate) fn shutdown(&mut self) { - self.shared.close(); - } -} - -impl fmt::Debug for Spawner { - fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - fmt.debug_struct("Spawner").finish() - } -} diff --git a/vendor/tokio/src/runtime/thread_pool/worker.rs b/vendor/tokio/src/runtime/thread_pool/worker.rs deleted file mode 100644 index 70cbddbd0..000000000 --- a/vendor/tokio/src/runtime/thread_pool/worker.rs +++ /dev/null @@ -1,841 +0,0 @@ -//! A scheduler is initialized with a fixed number of workers. Each worker is -//! driven by a thread. Each worker has a "core" which contains data such as the -//! run queue and other state. When `block_in_place` is called, the worker's -//! "core" is handed off to a new thread allowing the scheduler to continue to -//! make progress while the originating thread blocks. - -use crate::coop; -use crate::loom::rand::seed; -use crate::loom::sync::{Arc, Mutex}; -use crate::park::{Park, Unpark}; -use crate::runtime; -use crate::runtime::enter::EnterContext; -use crate::runtime::park::{Parker, Unparker}; -use crate::runtime::thread_pool::{AtomicCell, Idle}; -use crate::runtime::{queue, task}; -use crate::util::linked_list::{Link, LinkedList}; -use crate::util::FastRand; - -use std::cell::RefCell; -use std::time::Duration; - -/// A scheduler worker -pub(super) struct Worker { - /// Reference to shared state - shared: Arc<Shared>, - - /// Index holding this worker's remote state - index: usize, - - /// Used to hand-off a worker's core to another thread. - core: AtomicCell<Core>, -} - -/// Core data -struct Core { - /// Used to schedule bookkeeping tasks every so often. - tick: u8, - - /// When a task is scheduled from a worker, it is stored in this slot. The - /// worker will check this slot for a task **before** checking the run - /// queue. This effectively results in the **last** scheduled task to be run - /// next (LIFO). This is an optimization for message passing patterns and - /// helps to reduce latency. - lifo_slot: Option<Notified>, - - /// The worker-local run queue. - run_queue: queue::Local<Arc<Worker>>, - - /// True if the worker is currently searching for more work. Searching - /// involves attempting to steal from other workers. - is_searching: bool, - - /// True if the scheduler is being shutdown - is_shutdown: bool, - - /// Tasks owned by the core - tasks: LinkedList<Task, <Task as Link>::Target>, - - /// Parker - /// - /// Stored in an `Option` as the parker is added / removed to make the - /// borrow checker happy. - park: Option<Parker>, - - /// Fast random number generator. - rand: FastRand, -} - -/// State shared across all workers -pub(super) struct Shared { - /// Per-worker remote state. All other workers have access to this and is - /// how they communicate between each other. - remotes: Box<[Remote]>, - - /// Submit work to the scheduler while **not** currently on a worker thread. - inject: queue::Inject<Arc<Worker>>, - - /// Coordinates idle workers - idle: Idle, - - /// Cores that have observed the shutdown signal - /// - /// The core is **not** placed back in the worker to avoid it from being - /// stolen by a thread that was spawned as part of `block_in_place`. - #[allow(clippy::vec_box)] // we're moving an already-boxed value - shutdown_cores: Mutex<Vec<Box<Core>>>, -} - -/// Used to communicate with a worker from other threads. -struct Remote { - /// Steal tasks from this worker. - steal: queue::Steal<Arc<Worker>>, - - /// Transfers tasks to be released. Any worker pushes tasks, only the owning - /// worker pops. - pending_drop: task::TransferStack<Arc<Worker>>, - - /// Unparks the associated worker thread - unpark: Unparker, -} - -/// Thread-local context -struct Context { - /// Worker - worker: Arc<Worker>, - - /// Core data - core: RefCell<Option<Box<Core>>>, -} - -/// Starts the workers -pub(crate) struct Launch(Vec<Arc<Worker>>); - -/// Running a task may consume the core. If the core is still available when -/// running the task completes, it is returned. Otherwise, the worker will need -/// to stop processing. -type RunResult = Result<Box<Core>, ()>; - -/// A task handle -type Task = task::Task<Arc<Worker>>; - -/// A notified task handle -type Notified = task::Notified<Arc<Worker>>; - -// Tracks thread-local state -scoped_thread_local!(static CURRENT: Context); - -pub(super) fn create(size: usize, park: Parker) -> (Arc<Shared>, Launch) { - let mut cores = vec![]; - let mut remotes = vec![]; - - // Create the local queues - for _ in 0..size { - let (steal, run_queue) = queue::local(); - - let park = park.clone(); - let unpark = park.unpark(); - - cores.push(Box::new(Core { - tick: 0, - lifo_slot: None, - run_queue, - is_searching: false, - is_shutdown: false, - tasks: LinkedList::new(), - park: Some(park), - rand: FastRand::new(seed()), - })); - - remotes.push(Remote { - steal, - pending_drop: task::TransferStack::new(), - unpark, - }); - } - - let shared = Arc::new(Shared { - remotes: remotes.into_boxed_slice(), - inject: queue::Inject::new(), - idle: Idle::new(size), - shutdown_cores: Mutex::new(vec![]), - }); - - let mut launch = Launch(vec![]); - - for (index, core) in cores.drain(..).enumerate() { - launch.0.push(Arc::new(Worker { - shared: shared.clone(), - index, - core: AtomicCell::new(Some(core)), - })); - } - - (shared, launch) -} - -pub(crate) fn block_in_place<F, R>(f: F) -> R -where - F: FnOnce() -> R, -{ - // Try to steal the worker core back - struct Reset(coop::Budget); - - impl Drop for Reset { - fn drop(&mut self) { - CURRENT.with(|maybe_cx| { - if let Some(cx) = maybe_cx { - let core = cx.worker.core.take(); - let mut cx_core = cx.core.borrow_mut(); - assert!(cx_core.is_none()); - *cx_core = core; - - // Reset the task budget as we are re-entering the - // runtime. - coop::set(self.0); - } - }); - } - } - - let mut had_entered = false; - - CURRENT.with(|maybe_cx| { - match (crate::runtime::enter::context(), maybe_cx.is_some()) { - (EnterContext::Entered { .. }, true) => { - // We are on a thread pool runtime thread, so we just need to set up blocking. - had_entered = true; - } - (EnterContext::Entered { allow_blocking }, false) => { - // We are on an executor, but _not_ on the thread pool. - // That is _only_ okay if we are in a thread pool runtime's block_on method: - if allow_blocking { - had_entered = true; - return; - } else { - // This probably means we are on the basic_scheduler or in a LocalSet, - // where it is _not_ okay to block. - panic!("can call blocking only when running on the multi-threaded runtime"); - } - } - (EnterContext::NotEntered, true) => { - // This is a nested call to block_in_place (we already exited). - // All the necessary setup has already been done. - return; - } - (EnterContext::NotEntered, false) => { - // We are outside of the tokio runtime, so blocking is fine. - // We can also skip all of the thread pool blocking setup steps. - return; - } - } - - let cx = maybe_cx.expect("no .is_some() == false cases above should lead here"); - - // Get the worker core. If none is set, then blocking is fine! - let core = match cx.core.borrow_mut().take() { - Some(core) => core, - None => return, - }; - - // The parker should be set here - assert!(core.park.is_some()); - - // In order to block, the core must be sent to another thread for - // execution. - // - // First, move the core back into the worker's shared core slot. - cx.worker.core.set(core); - - // Next, clone the worker handle and send it to a new thread for - // processing. - // - // Once the blocking task is done executing, we will attempt to - // steal the core back. - let worker = cx.worker.clone(); - runtime::spawn_blocking(move || run(worker)); - }); - - if had_entered { - // Unset the current task's budget. Blocking sections are not - // constrained by task budgets. - let _reset = Reset(coop::stop()); - - crate::runtime::enter::exit(f) - } else { - f() - } -} - -/// After how many ticks is the global queue polled. This helps to ensure -/// fairness. -/// -/// The number is fairly arbitrary. I believe this value was copied from golang. -const GLOBAL_POLL_INTERVAL: u8 = 61; - -impl Launch { - pub(crate) fn launch(mut self) { - for worker in self.0.drain(..) { - runtime::spawn_blocking(move || run(worker)); - } - } -} - -fn run(worker: Arc<Worker>) { - // Acquire a core. If this fails, then another thread is running this - // worker and there is nothing further to do. - let core = match worker.core.take() { - Some(core) => core, - None => return, - }; - - // Set the worker context. - let cx = Context { - worker, - core: RefCell::new(None), - }; - - let _enter = crate::runtime::enter(true); - - CURRENT.set(&cx, || { - // This should always be an error. It only returns a `Result` to support - // using `?` to short circuit. - assert!(cx.run(core).is_err()); - }); -} - -impl Context { - fn run(&self, mut core: Box<Core>) -> RunResult { - while !core.is_shutdown { - // Increment the tick - core.tick(); - - // Run maintenance, if needed - core = self.maintenance(core); - - // First, check work available to the current worker. - if let Some(task) = core.next_task(&self.worker) { - core = self.run_task(task, core)?; - continue; - } - - // There is no more **local** work to process, try to steal work - // from other workers. - if let Some(task) = core.steal_work(&self.worker) { - core = self.run_task(task, core)?; - } else { - // Wait for work - core = self.park(core); - } - } - - core.pre_shutdown(&self.worker); - - // Signal shutdown - self.worker.shared.shutdown(core); - Err(()) - } - - fn run_task(&self, task: Notified, mut core: Box<Core>) -> RunResult { - // Make sure the worker is not in the **searching** state. This enables - // another idle worker to try to steal work. - core.transition_from_searching(&self.worker); - - // Make the core available to the runtime context - *self.core.borrow_mut() = Some(core); - - // Run the task - coop::budget(|| { - task.run(); - - // As long as there is budget remaining and a task exists in the - // `lifo_slot`, then keep running. - loop { - // Check if we still have the core. If not, the core was stolen - // by another worker. - let mut core = match self.core.borrow_mut().take() { - Some(core) => core, - None => return Err(()), - }; - - // Check for a task in the LIFO slot - let task = match core.lifo_slot.take() { - Some(task) => task, - None => return Ok(core), - }; - - if coop::has_budget_remaining() { - // Run the LIFO task, then loop - *self.core.borrow_mut() = Some(core); - task.run(); - } else { - // Not enough budget left to run the LIFO task, push it to - // the back of the queue and return. - core.run_queue.push_back(task, self.worker.inject()); - return Ok(core); - } - } - }) - } - - fn maintenance(&self, mut core: Box<Core>) -> Box<Core> { - if core.tick % GLOBAL_POLL_INTERVAL == 0 { - // Call `park` with a 0 timeout. This enables the I/O driver, timer, ... - // to run without actually putting the thread to sleep. - core = self.park_timeout(core, Some(Duration::from_millis(0))); - - // Run regularly scheduled maintenance - core.maintenance(&self.worker); - } - - core - } - - fn park(&self, mut core: Box<Core>) -> Box<Core> { - core.transition_to_parked(&self.worker); - - while !core.is_shutdown { - core = self.park_timeout(core, None); - - // Run regularly scheduled maintenance - core.maintenance(&self.worker); - - if core.transition_from_parked(&self.worker) { - return core; - } - } - - core - } - - fn park_timeout(&self, mut core: Box<Core>, duration: Option<Duration>) -> Box<Core> { - // Take the parker out of core - let mut park = core.park.take().expect("park missing"); - - // Store `core` in context - *self.core.borrow_mut() = Some(core); - - // Park thread - if let Some(timeout) = duration { - park.park_timeout(timeout).expect("park failed"); - } else { - park.park().expect("park failed"); - } - - // Remove `core` from context - core = self.core.borrow_mut().take().expect("core missing"); - - // Place `park` back in `core` - core.park = Some(park); - - // If there are tasks available to steal, notify a worker - if core.run_queue.is_stealable() { - self.worker.shared.notify_parked(); - } - - core - } -} - -impl Core { - /// Increment the tick - fn tick(&mut self) { - self.tick = self.tick.wrapping_add(1); - } - - /// Return the next notified task available to this worker. - fn next_task(&mut self, worker: &Worker) -> Option<Notified> { - if self.tick % GLOBAL_POLL_INTERVAL == 0 { - worker.inject().pop().or_else(|| self.next_local_task()) - } else { - self.next_local_task().or_else(|| worker.inject().pop()) - } - } - - fn next_local_task(&mut self) -> Option<Notified> { - self.lifo_slot.take().or_else(|| self.run_queue.pop()) - } - - fn steal_work(&mut self, worker: &Worker) -> Option<Notified> { - if !self.transition_to_searching(worker) { - return None; - } - - let num = worker.shared.remotes.len(); - // Start from a random worker - let start = self.rand.fastrand_n(num as u32) as usize; - - for i in 0..num { - let i = (start + i) % num; - - // Don't steal from ourself! We know we don't have work. - if i == worker.index { - continue; - } - - let target = &worker.shared.remotes[i]; - if let Some(task) = target.steal.steal_into(&mut self.run_queue) { - return Some(task); - } - } - - // Fallback on checking the global queue - worker.shared.inject.pop() - } - - fn transition_to_searching(&mut self, worker: &Worker) -> bool { - if !self.is_searching { - self.is_searching = worker.shared.idle.transition_worker_to_searching(); - } - - self.is_searching - } - - fn transition_from_searching(&mut self, worker: &Worker) { - if !self.is_searching { - return; - } - - self.is_searching = false; - worker.shared.transition_worker_from_searching(); - } - - /// Prepare the worker state for parking - fn transition_to_parked(&mut self, worker: &Worker) { - // When the final worker transitions **out** of searching to parked, it - // must check all the queues one last time in case work materialized - // between the last work scan and transitioning out of searching. - let is_last_searcher = worker - .shared - .idle - .transition_worker_to_parked(worker.index, self.is_searching); - - // The worker is no longer searching. Setting this is the local cache - // only. - self.is_searching = false; - - if is_last_searcher { - worker.shared.notify_if_work_pending(); - } - } - - /// Returns `true` if the transition happened. - fn transition_from_parked(&mut self, worker: &Worker) -> bool { - // If a task is in the lifo slot, then we must unpark regardless of - // being notified - if self.lifo_slot.is_some() { - worker.shared.idle.unpark_worker_by_id(worker.index); - self.is_searching = true; - return true; - } - - if worker.shared.idle.is_parked(worker.index) { - return false; - } - - // When unparked, the worker is in the searching state. - self.is_searching = true; - true - } - - /// Runs maintenance work such as free pending tasks and check the pool's - /// state. - fn maintenance(&mut self, worker: &Worker) { - self.drain_pending_drop(worker); - - if !self.is_shutdown { - // Check if the scheduler has been shutdown - self.is_shutdown = worker.inject().is_closed(); - } - } - - // Signals all tasks to shut down, and waits for them to complete. Must run - // before we enter the single-threaded phase of shutdown processing. - fn pre_shutdown(&mut self, worker: &Worker) { - // Signal to all tasks to shut down. - for header in self.tasks.iter() { - header.shutdown(); - } - - loop { - self.drain_pending_drop(worker); - - if self.tasks.is_empty() { - break; - } - - // Wait until signalled - let park = self.park.as_mut().expect("park missing"); - park.park().expect("park failed"); - } - } - - // Shutdown the core - fn shutdown(&mut self) { - assert!(self.tasks.is_empty()); - - // Take the core - let mut park = self.park.take().expect("park missing"); - - // Drain the queue - while self.next_local_task().is_some() {} - - park.shutdown(); - } - - fn drain_pending_drop(&mut self, worker: &Worker) { - use std::mem::ManuallyDrop; - - for task in worker.remote().pending_drop.drain() { - let task = ManuallyDrop::new(task); - - // safety: tasks are only pushed into the `pending_drop` stacks that - // are associated with the list they are inserted into. When a task - // is pushed into `pending_drop`, the ref-inc is skipped, so we must - // not ref-dec here. - // - // See `bind` and `release` implementations. - unsafe { - self.tasks.remove(task.header().into()); - } - } - } -} - -impl Worker { - /// Returns a reference to the scheduler's injection queue - fn inject(&self) -> &queue::Inject<Arc<Worker>> { - &self.shared.inject - } - - /// Return a reference to this worker's remote data - fn remote(&self) -> &Remote { - &self.shared.remotes[self.index] - } - - fn eq(&self, other: &Worker) -> bool { - self.shared.ptr_eq(&other.shared) && self.index == other.index - } -} - -impl task::Schedule for Arc<Worker> { - fn bind(task: Task) -> Arc<Worker> { - CURRENT.with(|maybe_cx| { - let cx = maybe_cx.expect("scheduler context missing"); - - // Track the task - cx.core - .borrow_mut() - .as_mut() - .expect("scheduler core missing") - .tasks - .push_front(task); - - // Return a clone of the worker - cx.worker.clone() - }) - } - - fn release(&self, task: &Task) -> Option<Task> { - use std::ptr::NonNull; - - enum Immediate { - // Task has been synchronously removed from the Core owned by the - // current thread - Removed(Option<Task>), - // Task is owned by another thread, so we need to notify it to clean - // up the task later. - MaybeRemote, - } - - let immediate = CURRENT.with(|maybe_cx| { - let cx = match maybe_cx { - Some(cx) => cx, - None => return Immediate::MaybeRemote, - }; - - if !self.eq(&cx.worker) { - // Task owned by another core, so we need to notify it. - return Immediate::MaybeRemote; - } - - let mut maybe_core = cx.core.borrow_mut(); - - if let Some(core) = &mut *maybe_core { - // Directly remove the task - // - // safety: the task is inserted in the list in `bind`. - unsafe { - let ptr = NonNull::from(task.header()); - return Immediate::Removed(core.tasks.remove(ptr)); - } - } - - Immediate::MaybeRemote - }); - - // Checks if we were called from within a worker, allowing for immediate - // removal of a scheduled task. Else we have to go through the slower - // process below where we remotely mark a task as dropped. - match immediate { - Immediate::Removed(task) => return task, - Immediate::MaybeRemote => (), - }; - - // Track the task to be released by the worker that owns it - // - // Safety: We get a new handle without incrementing the ref-count. - // A ref-count is held by the "owned" linked list and it is only - // ever removed from that list as part of the release process: this - // method or popping the task from `pending_drop`. Thus, we can rely - // on the ref-count held by the linked-list to keep the memory - // alive. - // - // When the task is removed from the stack, it is forgotten instead - // of dropped. - let task = unsafe { Task::from_raw(task.header().into()) }; - - self.remote().pending_drop.push(task); - - // The worker core has been handed off to another thread. In the - // event that the scheduler is currently shutting down, the thread - // that owns the task may be waiting on the release to complete - // shutdown. - if self.inject().is_closed() { - self.remote().unpark.unpark(); - } - - None - } - - fn schedule(&self, task: Notified) { - // Because this is not a newly spawned task, if scheduling fails due to - // the runtime shutting down, there is no special work that must happen - // here. - let _ = self.shared.schedule(task, false); - } - - fn yield_now(&self, task: Notified) { - // Because this is not a newly spawned task, if scheduling fails due to - // the runtime shutting down, there is no special work that must happen - // here. - let _ = self.shared.schedule(task, true); - } -} - -impl Shared { - pub(super) fn schedule(&self, task: Notified, is_yield: bool) -> Result<(), Notified> { - CURRENT.with(|maybe_cx| { - if let Some(cx) = maybe_cx { - // Make sure the task is part of the **current** scheduler. - if self.ptr_eq(&cx.worker.shared) { - // And the current thread still holds a core - if let Some(core) = cx.core.borrow_mut().as_mut() { - self.schedule_local(core, task, is_yield); - return Ok(()); - } - } - } - - // Otherwise, use the inject queue - self.inject.push(task)?; - self.notify_parked(); - Ok(()) - }) - } - - fn schedule_local(&self, core: &mut Core, task: Notified, is_yield: bool) { - // Spawning from the worker thread. If scheduling a "yield" then the - // task must always be pushed to the back of the queue, enabling other - // tasks to be executed. If **not** a yield, then there is more - // flexibility and the task may go to the front of the queue. - let should_notify = if is_yield { - core.run_queue.push_back(task, &self.inject); - true - } else { - // Push to the LIFO slot - let prev = core.lifo_slot.take(); - let ret = prev.is_some(); - - if let Some(prev) = prev { - core.run_queue.push_back(prev, &self.inject); - } - - core.lifo_slot = Some(task); - - ret - }; - - // Only notify if not currently parked. If `park` is `None`, then the - // scheduling is from a resource driver. As notifications often come in - // batches, the notification is delayed until the park is complete. - if should_notify && core.park.is_some() { - self.notify_parked(); - } - } - - pub(super) fn close(&self) { - if self.inject.close() { - self.notify_all(); - } - } - - fn notify_parked(&self) { - if let Some(index) = self.idle.worker_to_notify() { - self.remotes[index].unpark.unpark(); - } - } - - fn notify_all(&self) { - for remote in &self.remotes[..] { - remote.unpark.unpark(); - } - } - - fn notify_if_work_pending(&self) { - for remote in &self.remotes[..] { - if !remote.steal.is_empty() { - self.notify_parked(); - return; - } - } - - if !self.inject.is_empty() { - self.notify_parked(); - } - } - - fn transition_worker_from_searching(&self) { - if self.idle.transition_worker_from_searching() { - // We are the final searching worker. Because work was found, we - // need to notify another worker. - self.notify_parked(); - } - } - - /// Signals that a worker has observed the shutdown signal and has replaced - /// its core back into its handle. - /// - /// If all workers have reached this point, the final cleanup is performed. - fn shutdown(&self, core: Box<Core>) { - let mut cores = self.shutdown_cores.lock(); - cores.push(core); - - if cores.len() != self.remotes.len() { - return; - } - - for mut core in cores.drain(..) { - core.shutdown(); - } - - // Drain the injection queue - while let Some(task) = self.inject.pop() { - task.shutdown(); - } - } - - fn ptr_eq(&self, other: &Shared) -> bool { - std::ptr::eq(self, other) - } -} diff --git a/vendor/tokio/src/time/driver/entry.rs b/vendor/tokio/src/runtime/time/entry.rs index 168e0b995..798d3c11e 100644 --- a/vendor/tokio/src/time/driver/entry.rs +++ b/vendor/tokio/src/runtime/time/entry.rs @@ -5,9 +5,9 @@ //! //! # Ground rules //! -//! The heart of the timer implementation here is the `TimerShared` structure, -//! shared between the `TimerEntry` and the driver. Generally, we permit access -//! to `TimerShared` ONLY via either 1) a mutable reference to `TimerEntry` or +//! The heart of the timer implementation here is the [`TimerShared`] structure, +//! shared between the [`TimerEntry`] and the driver. Generally, we permit access +//! to [`TimerShared`] ONLY via either 1) a mutable reference to [`TimerEntry`] or //! 2) a held driver lock. //! //! It follows from this that any changes made while holding BOTH 1 and 2 will @@ -36,7 +36,7 @@ //! point than it was originally scheduled for. //! //! This is accomplished by lazily rescheduling timers. That is, we update the -//! state field field with the true expiration of the timer from the holder of +//! state field with the true expiration of the timer from the holder of //! the [`TimerEntry`]. When the driver services timers (ie, whenever it's //! walking lists of timers), it checks this "true when" value, and reschedules //! based on it. @@ -49,19 +49,20 @@ //! There is of course a race condition between timer reset and timer //! expiration. If the driver fails to observe the updated expiration time, it //! could trigger expiration of the timer too early. However, because -//! `mark_pending` performs a compare-and-swap, it will identify this race and +//! [`mark_pending`][mark_pending] performs a compare-and-swap, it will identify this race and //! refuse to mark the timer as pending. +//! +//! [mark_pending]: TimerHandle::mark_pending use crate::loom::cell::UnsafeCell; use crate::loom::sync::atomic::AtomicU64; use crate::loom::sync::atomic::Ordering; +use crate::runtime::scheduler; use crate::sync::AtomicWaker; use crate::time::Instant; use crate::util::linked_list; -use super::Handle; - use std::cell::UnsafeCell as StdUnsafeCell; use std::task::{Context, Poll, Waker}; use std::{marker::PhantomPinned, pin::Pin, ptr::NonNull}; @@ -71,6 +72,10 @@ type TimerResult = Result<(), crate::time::error::Error>; const STATE_DEREGISTERED: u64 = u64::MAX; const STATE_PENDING_FIRE: u64 = STATE_DEREGISTERED - 1; const STATE_MIN_VALUE: u64 = STATE_PENDING_FIRE; +/// The largest safe integer to use for ticks. +/// +/// This value should be updated if any other signal values are added above. +pub(super) const MAX_SAFE_MILLIS_DURATION: u64 = u64::MAX - 2; /// This structure holds the current shared state of the timer - its scheduled /// time (if registered), or otherwise the result of the timer completing, as @@ -93,7 +98,7 @@ pub(super) struct StateCell { /// without holding the driver lock is undefined behavior. result: UnsafeCell<TimerResult>, /// The currently-registered waker - waker: CachePadded<AtomicWaker>, + waker: AtomicWaker, } impl Default for StateCell { @@ -113,7 +118,7 @@ impl StateCell { Self { state: AtomicU64::new(STATE_DEREGISTERED), result: UnsafeCell::new(Ok(())), - waker: CachePadded(AtomicWaker::new()), + waker: AtomicWaker::new(), } } @@ -125,7 +130,7 @@ impl StateCell { fn when(&self) -> Option<u64> { let cur_state = self.state.load(Ordering::Relaxed); - if cur_state == u64::MAX { + if cur_state == STATE_DEREGISTERED { None } else { Some(cur_state) @@ -138,7 +143,7 @@ impl StateCell { // We must register first. This ensures that either `fire` will // observe the new waker, or we will observe a racing fire to have set // the state, or both. - self.waker.0.register_by_ref(waker); + self.waker.register_by_ref(waker); self.read_state() } @@ -171,7 +176,12 @@ impl StateCell { let mut cur_state = self.state.load(Ordering::Relaxed); loop { - debug_assert!(cur_state < STATE_MIN_VALUE); + // improve the error message for things like + // https://github.com/tokio-rs/tokio/issues/3675 + assert!( + cur_state < STATE_MIN_VALUE, + "mark_pending called when the timer entry is in an invalid state" + ); if cur_state > not_after { break Err(cur_state); @@ -221,7 +231,7 @@ impl StateCell { self.state.store(STATE_DEREGISTERED, Ordering::Release); - self.waker.0.take_waker() + self.waker.take_waker() } /// Marks the timer as registered (poll will return None) and sets the @@ -281,10 +291,10 @@ impl StateCell { /// timer. As this participates in intrusive data structures, it must be pinned /// before polling. #[derive(Debug)] -pub(super) struct TimerEntry { - /// Arc reference to the driver. We can only free the driver after +pub(crate) struct TimerEntry { + /// Arc reference to the runtime handle. We can only free the driver after /// deregistering everything from their respective timer wheels. - driver: Handle, + driver: scheduler::Handle, /// Shared inner structure; this is part of an intrusive linked list, and /// therefore other references can exist to it while mutable references to /// Entry exist. @@ -292,9 +302,11 @@ pub(super) struct TimerEntry { /// This is manipulated only under the inner mutex. TODO: Can we use loom /// cells for this? inner: StdUnsafeCell<TimerShared>, - /// Initial deadline for the timer. This is used to register on the first + /// Deadline for the timer. This is used to register on the first /// poll, as we can't register prior to being pinned. - initial_deadline: Option<Instant>, + deadline: Instant, + /// Whether the deadline has been registered. + registered: bool, /// Ensure the type is !Unpin _m: std::marker::PhantomPinned, } @@ -323,32 +335,65 @@ pub(super) type EntryList = crate::util::linked_list::LinkedList<TimerShared, Ti /// frontend (`Entry`) and driver backend. /// /// Note that this structure is located inside the `TimerEntry` structure. -#[derive(Debug)] pub(crate) struct TimerShared { + /// A link within the doubly-linked list of timers on a particular level and + /// slot. Valid only if state is equal to Registered. + /// + /// Only accessed under the entry lock. + pointers: linked_list::Pointers<TimerShared>, + + /// The expiration time for which this entry is currently registered. + /// Generally owned by the driver, but is accessed by the entry when not + /// registered. + cached_when: AtomicU64, + + /// The true expiration time. Set by the timer future, read by the driver. + true_when: AtomicU64, + /// Current state. This records whether the timer entry is currently under /// the ownership of the driver, and if not, its current state (not /// complete, fired, error, etc). state: StateCell, - /// Data manipulated by the driver thread itself, only. - driver_state: CachePadded<TimerSharedPadded>, - _p: PhantomPinned, } +unsafe impl Send for TimerShared {} +unsafe impl Sync for TimerShared {} + +impl std::fmt::Debug for TimerShared { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TimerShared") + .field("when", &self.true_when.load(Ordering::Relaxed)) + .field("cached_when", &self.cached_when.load(Ordering::Relaxed)) + .field("state", &self.state) + .finish() + } +} + +generate_addr_of_methods! { + impl<> TimerShared { + unsafe fn addr_of_pointers(self: NonNull<Self>) -> NonNull<linked_list::Pointers<TimerShared>> { + &self.pointers + } + } +} + impl TimerShared { pub(super) fn new() -> Self { Self { + cached_when: AtomicU64::new(0), + true_when: AtomicU64::new(0), + pointers: linked_list::Pointers::new(), state: StateCell::default(), - driver_state: CachePadded(TimerSharedPadded::new()), _p: PhantomPinned, } } - /// Gets the cached time-of-expiration value + /// Gets the cached time-of-expiration value. pub(super) fn cached_when(&self) -> u64 { // Cached-when is only accessed under the driver lock, so we can use relaxed - self.driver_state.0.cached_when.load(Ordering::Relaxed) + self.cached_when.load(Ordering::Relaxed) } /// Gets the true time-of-expiration value, and copies it into the cached @@ -359,10 +404,7 @@ impl TimerShared { pub(super) unsafe fn sync_when(&self) -> u64 { let true_when = self.true_when(); - self.driver_state - .0 - .cached_when - .store(true_when, Ordering::Relaxed); + self.cached_when.store(true_when, Ordering::Relaxed); true_when } @@ -372,10 +414,7 @@ impl TimerShared { /// SAFETY: Must be called with the driver lock held, and when this entry is /// not in any timer wheel lists. unsafe fn set_cached_when(&self, when: u64) { - self.driver_state - .0 - .cached_when - .store(when, Ordering::Relaxed); + self.cached_when.store(when, Ordering::Relaxed); } /// Returns the true time-of-expiration value, with relaxed memory ordering. @@ -390,7 +429,7 @@ impl TimerShared { /// in the timer wheel. pub(super) unsafe fn set_expiration(&self, t: u64) { self.state.set_expiration(t); - self.driver_state.0.cached_when.store(t, Ordering::Relaxed); + self.cached_when.store(t, Ordering::Relaxed); } /// Sets the true time-of-expiration only if it is after the current. @@ -414,48 +453,6 @@ impl TimerShared { } } -/// Additional shared state between the driver and the timer which is cache -/// padded. This contains the information that the driver thread accesses most -/// frequently to minimize contention. In particular, we move it away from the -/// waker, as the waker is updated on every poll. -struct TimerSharedPadded { - /// The expiration time for which this entry is currently registered. - /// Generally owned by the driver, but is accessed by the entry when not - /// registered. - cached_when: AtomicU64, - - /// The true expiration time. Set by the timer future, read by the driver. - true_when: AtomicU64, - - /// A link within the doubly-linked list of timers on a particular level and - /// slot. Valid only if state is equal to Registered. - /// - /// Only accessed under the entry lock. - pointers: StdUnsafeCell<linked_list::Pointers<TimerShared>>, -} - -impl std::fmt::Debug for TimerSharedPadded { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("TimerSharedPadded") - .field("when", &self.true_when.load(Ordering::Relaxed)) - .field("cached_when", &self.cached_when.load(Ordering::Relaxed)) - .finish() - } -} - -impl TimerSharedPadded { - fn new() -> Self { - Self { - cached_when: AtomicU64::new(0), - true_when: AtomicU64::new(0), - pointers: StdUnsafeCell::new(linked_list::Pointers::new()), - } - } -} - -unsafe impl Send for TimerShared {} -unsafe impl Sync for TimerShared {} - unsafe impl linked_list::Link for TimerShared { type Handle = TimerHandle; @@ -472,20 +469,25 @@ unsafe impl linked_list::Link for TimerShared { unsafe fn pointers( target: NonNull<Self::Target>, ) -> NonNull<linked_list::Pointers<Self::Target>> { - unsafe { NonNull::new(target.as_ref().driver_state.0.pointers.get()).unwrap() } + TimerShared::addr_of_pointers(target) } } // ===== impl Entry ===== impl TimerEntry { - pub(crate) fn new(handle: &Handle, deadline: Instant) -> Self { + #[track_caller] + pub(crate) fn new(handle: &scheduler::Handle, deadline: Instant) -> Self { + // Panic if the time driver is not enabled + let _ = handle.driver().time(); + let driver = handle.clone(); Self { driver, inner: StdUnsafeCell::new(TimerShared::new()), - initial_deadline: Some(deadline), + deadline, + registered: false, _m: std::marker::PhantomPinned, } } @@ -494,8 +496,12 @@ impl TimerEntry { unsafe { &*self.inner.get() } } + pub(crate) fn deadline(&self) -> Instant { + self.deadline + } + pub(crate) fn is_elapsed(&self) -> bool { - !self.inner().state.might_be_registered() && self.initial_deadline.is_none() + !self.inner().state.might_be_registered() && self.registered } /// Cancels and deregisters the timer. This operation is irreversible. @@ -522,20 +528,24 @@ impl TimerEntry { // driver did so far and happens-before everything the driver does in // the future. While we have the lock held, we also go ahead and // deregister the entry if necessary. - unsafe { self.driver.clear_entry(NonNull::from(self.inner())) }; + unsafe { self.driver().clear_entry(NonNull::from(self.inner())) }; } - pub(crate) fn reset(mut self: Pin<&mut Self>, new_time: Instant) { - unsafe { self.as_mut().get_unchecked_mut() }.initial_deadline = None; + pub(crate) fn reset(mut self: Pin<&mut Self>, new_time: Instant, reregister: bool) { + unsafe { self.as_mut().get_unchecked_mut() }.deadline = new_time; + unsafe { self.as_mut().get_unchecked_mut() }.registered = reregister; - let tick = self.driver.time_source().deadline_to_tick(new_time); + let tick = self.driver().time_source().deadline_to_tick(new_time); if self.inner().extend_expiration(tick).is_ok() { return; } - unsafe { - self.driver.reregister(tick, self.inner().into()); + if reregister { + unsafe { + self.driver() + .reregister(&self.driver.driver().io, tick, self.inner().into()); + } } } @@ -543,18 +553,28 @@ impl TimerEntry { mut self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll<Result<(), super::Error>> { - if self.driver.is_shutdown() { + if self.driver().is_shutdown() { panic!("{}", crate::util::error::RUNTIME_SHUTTING_DOWN_ERROR); } - if let Some(deadline) = self.initial_deadline { - self.as_mut().reset(deadline); + if !self.registered { + let deadline = self.deadline; + self.as_mut().reset(deadline, true); } let this = unsafe { self.get_unchecked_mut() }; this.inner().state.poll(cx.waker()) } + + pub(crate) fn driver(&self) -> &super::Handle { + self.driver.driver().time() + } + + #[cfg(all(tokio_unstable, feature = "tracing"))] + pub(crate) fn clock(&self) -> &super::Clock { + self.driver.driver().clock() + } } impl TimerHandle { @@ -622,8 +642,3 @@ impl Drop for TimerEntry { unsafe { Pin::new_unchecked(self) }.as_mut().cancel() } } - -#[cfg_attr(target_arch = "x86_64", repr(align(128)))] -#[cfg_attr(not(target_arch = "x86_64"), repr(align(64)))] -#[derive(Debug, Default)] -struct CachePadded<T>(T); diff --git a/vendor/tokio/src/runtime/time/handle.rs b/vendor/tokio/src/runtime/time/handle.rs new file mode 100644 index 000000000..fce791d99 --- /dev/null +++ b/vendor/tokio/src/runtime/time/handle.rs @@ -0,0 +1,62 @@ +use crate::runtime::time::TimeSource; +use std::fmt; + +/// Handle to time driver instance. +pub(crate) struct Handle { + pub(super) time_source: TimeSource, + pub(super) inner: super::Inner, +} + +impl Handle { + /// Returns the time source associated with this handle. + pub(crate) fn time_source(&self) -> &TimeSource { + &self.time_source + } + + /// Checks whether the driver has been shutdown. + pub(super) fn is_shutdown(&self) -> bool { + self.inner.is_shutdown() + } + + /// Track that the driver is being unparked + pub(crate) fn unpark(&self) { + #[cfg(feature = "test-util")] + self.inner + .did_wake + .store(true, std::sync::atomic::Ordering::SeqCst); + } +} + +cfg_not_rt! { + impl Handle { + /// Tries to get a handle to the current timer. + /// + /// # Panics + /// + /// This function panics if there is no current timer set. + /// + /// It can be triggered when [`Builder::enable_time`] or + /// [`Builder::enable_all`] are not included in the builder. + /// + /// It can also panic whenever a timer is created outside of a + /// Tokio runtime. That is why `rt.block_on(sleep(...))` will panic, + /// since the function is executed outside of the runtime. + /// Whereas `rt.block_on(async {sleep(...).await})` doesn't panic. + /// And this is because wrapping the function on an async makes it lazy, + /// and so gets executed inside the runtime successfully without + /// panicking. + /// + /// [`Builder::enable_time`]: crate::runtime::Builder::enable_time + /// [`Builder::enable_all`]: crate::runtime::Builder::enable_all + #[track_caller] + pub(crate) fn current() -> Self { + panic!("{}", crate::util::error::CONTEXT_MISSING_ERROR) + } + } +} + +impl fmt::Debug for Handle { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "Handle") + } +} diff --git a/vendor/tokio/src/time/driver/mod.rs b/vendor/tokio/src/runtime/time/mod.rs index 37d2231c3..423ad79ab 100644 --- a/vendor/tokio/src/time/driver/mod.rs +++ b/vendor/tokio/src/runtime/time/mod.rs @@ -4,25 +4,26 @@ #![allow(unused_unsafe)] #![cfg_attr(not(feature = "rt"), allow(dead_code))] -//! Time driver +//! Time driver. mod entry; -pub(self) use self::entry::{EntryList, TimerEntry, TimerHandle, TimerShared}; +pub(crate) use entry::TimerEntry; +use entry::{EntryList, TimerHandle, TimerShared, MAX_SAFE_MILLIS_DURATION}; mod handle; pub(crate) use self::handle::Handle; -mod wheel; +mod source; +pub(crate) use source::TimeSource; -pub(super) mod sleep; +mod wheel; use crate::loom::sync::atomic::{AtomicBool, Ordering}; -use crate::loom::sync::{Arc, Mutex}; -use crate::park::{Park, Unpark}; +use crate::loom::sync::Mutex; +use crate::runtime::driver::{self, IoHandle, IoStack}; use crate::time::error::Error; -use crate::time::{Clock, Duration, Instant}; +use crate::time::{Clock, Duration}; -use std::convert::TryInto; use std::fmt; use std::{num::NonZeroU64, ptr::NonNull, task::Waker}; @@ -82,15 +83,18 @@ use std::{num::NonZeroU64, ptr::NonNull, task::Waker}; /// [timeout]: crate::time::Timeout /// [interval]: crate::time::Interval #[derive(Debug)] -pub(crate) struct Driver<P: Park + 'static> { - /// Timing backend in use - time_source: ClockTime, +pub(crate) struct Driver { + /// Parker to delegate to. + park: IoStack, +} - /// Shared state - handle: Handle, +/// Timer state shared between `Driver`, `Handle`, and `Registration`. +struct Inner { + // The state is split like this so `Handle` can access `is_shutdown` without locking the mutex + pub(super) state: Mutex<InnerState>, - /// Parker to delegate to - park: P, + /// True if the driver is being shutdown. + pub(super) is_shutdown: AtomicBool, // When `true`, a call to `park_timeout` should immediately return and time // should not advance. One reason for this to be `true` is if the task @@ -99,113 +103,80 @@ pub(crate) struct Driver<P: Park + 'static> { // While it may look racy, it only has any effect when the clock is paused // and pausing the clock is restricted to a single-threaded runtime. #[cfg(feature = "test-util")] - did_wake: Arc<AtomicBool>, -} - -/// A structure which handles conversion from Instants to u64 timestamps. -#[derive(Debug, Clone)] -pub(self) struct ClockTime { - clock: super::clock::Clock, - start_time: Instant, -} - -impl ClockTime { - pub(self) fn new(clock: Clock) -> Self { - Self { - start_time: clock.now(), - clock, - } - } - - pub(self) fn deadline_to_tick(&self, t: Instant) -> u64 { - // Round up to the end of a ms - self.instant_to_tick(t + Duration::from_nanos(999_999)) - } - - pub(self) fn instant_to_tick(&self, t: Instant) -> u64 { - // round up - let dur: Duration = t - .checked_duration_since(self.start_time) - .unwrap_or_else(|| Duration::from_secs(0)); - let ms = dur.as_millis(); - - ms.try_into().expect("Duration too far into the future") - } - - pub(self) fn tick_to_duration(&self, t: u64) -> Duration { - Duration::from_millis(t) - } - - pub(self) fn now(&self) -> u64 { - self.instant_to_tick(self.clock.now()) - } -} - -/// Timer state shared between `Driver`, `Handle`, and `Registration`. -struct Inner { - // The state is split like this so `Handle` can access `is_shutdown` without locking the mutex - pub(super) state: Mutex<InnerState>, - - /// True if the driver is being shutdown - pub(super) is_shutdown: AtomicBool, + did_wake: AtomicBool, } /// Time state shared which must be protected by a `Mutex` struct InnerState { - /// Timing backend in use - time_source: ClockTime, - /// The last published timer `elapsed` value. elapsed: u64, - /// The earliest time at which we promise to wake up without unparking + /// The earliest time at which we promise to wake up without unparking. next_wake: Option<NonZeroU64>, - /// Timer wheel + /// Timer wheel. wheel: wheel::Wheel, - - /// Unparker that can be used to wake the time driver - unpark: Box<dyn Unpark>, } // ===== impl Driver ===== -impl<P> Driver<P> -where - P: Park + 'static, -{ +impl Driver { /// Creates a new `Driver` instance that uses `park` to block the current /// thread and `time_source` to get the current time and convert to ticks. /// /// Specifying the source of time is useful when testing. - pub(crate) fn new(park: P, clock: Clock) -> Driver<P> { - let time_source = ClockTime::new(clock); + pub(crate) fn new(park: IoStack, clock: &Clock) -> (Driver, Handle) { + let time_source = TimeSource::new(clock); - let inner = Inner::new(time_source.clone(), Box::new(park.unpark())); - - Driver { + let handle = Handle { time_source, - handle: Handle::new(Arc::new(inner)), - park, - #[cfg(feature = "test-util")] - did_wake: Arc::new(AtomicBool::new(false)), - } + inner: Inner { + state: Mutex::new(InnerState { + elapsed: 0, + next_wake: None, + wheel: wheel::Wheel::new(), + }), + is_shutdown: AtomicBool::new(false), + + #[cfg(feature = "test-util")] + did_wake: AtomicBool::new(false), + }, + }; + + let driver = Driver { park }; + + (driver, handle) } - /// Returns a handle to the timer. - /// - /// The `Handle` is how `Sleep` instances are created. The `Sleep` instances - /// can either be created directly or the `Handle` instance can be passed to - /// `with_default`, setting the timer as the default timer for the execution - /// context. - pub(crate) fn handle(&self) -> Handle { - self.handle.clone() + pub(crate) fn park(&mut self, handle: &driver::Handle) { + self.park_internal(handle, None) } - fn park_internal(&mut self, limit: Option<Duration>) -> Result<(), P::Error> { - let mut lock = self.handle.get().state.lock(); + pub(crate) fn park_timeout(&mut self, handle: &driver::Handle, duration: Duration) { + self.park_internal(handle, Some(duration)) + } + + pub(crate) fn shutdown(&mut self, rt_handle: &driver::Handle) { + let handle = rt_handle.time(); + + if handle.is_shutdown() { + return; + } + + handle.inner.is_shutdown.store(true, Ordering::SeqCst); + + // Advance time forward to the end of time. - assert!(!self.handle.is_shutdown()); + handle.process_at_time(u64::MAX); + + self.park.shutdown(rt_handle); + } + + fn park_internal(&mut self, rt_handle: &driver::Handle, limit: Option<Duration>) { + let handle = rt_handle.time(); + let mut lock = handle.inner.state.lock(); + + assert!(!handle.is_shutdown()); let next_wake = lock.wheel.next_expiration_time(); lock.next_wake = @@ -215,86 +186,91 @@ where match next_wake { Some(when) => { - let now = self.time_source.now(); + let now = handle.time_source.now(rt_handle.clock()); // Note that we effectively round up to 1ms here - this avoids // very short-duration microsecond-resolution sleeps that the OS // might treat as zero-length. - let mut duration = self.time_source.tick_to_duration(when.saturating_sub(now)); + let mut duration = handle + .time_source + .tick_to_duration(when.saturating_sub(now)); if duration > Duration::from_millis(0) { if let Some(limit) = limit { duration = std::cmp::min(limit, duration); } - self.park_timeout(duration)?; + self.park_thread_timeout(rt_handle, duration); } else { - self.park.park_timeout(Duration::from_secs(0))?; + self.park.park_timeout(rt_handle, Duration::from_secs(0)); } } None => { if let Some(duration) = limit { - self.park_timeout(duration)?; + self.park_thread_timeout(rt_handle, duration); } else { - self.park.park()?; + self.park.park(rt_handle); } } } // Process pending timers after waking up - self.handle.process(); - - Ok(()) + handle.process(rt_handle.clock()); } cfg_test_util! { - fn park_timeout(&mut self, duration: Duration) -> Result<(), P::Error> { - let clock = &self.time_source.clock; + fn park_thread_timeout(&mut self, rt_handle: &driver::Handle, duration: Duration) { + let handle = rt_handle.time(); + let clock = rt_handle.clock(); - if clock.is_paused() { - self.park.park_timeout(Duration::from_secs(0))?; + if clock.can_auto_advance() { + self.park.park_timeout(rt_handle, Duration::from_secs(0)); // If the time driver was woken, then the park completed // before the "duration" elapsed (usually caused by a // yield in `Runtime::block_on`). In this case, we don't // advance the clock. - if !self.did_wake() { + if !handle.did_wake() { // Simulate advancing time - clock.advance(duration); + if let Err(msg) = clock.advance(duration) { + panic!("{}", msg); + } } } else { - self.park.park_timeout(duration)?; + self.park.park_timeout(rt_handle, duration); } - - Ok(()) - } - - fn did_wake(&self) -> bool { - self.did_wake.swap(false, Ordering::SeqCst) } } cfg_not_test_util! { - fn park_timeout(&mut self, duration: Duration) -> Result<(), P::Error> { - self.park.park_timeout(duration) + fn park_thread_timeout(&mut self, rt_handle: &driver::Handle, duration: Duration) { + self.park.park_timeout(rt_handle, duration); } } } impl Handle { /// Runs timer related logic, and returns the next wakeup time - pub(self) fn process(&self) { - let now = self.time_source().now(); + pub(self) fn process(&self, clock: &Clock) { + let now = self.time_source().now(clock); self.process_at_time(now) } - pub(self) fn process_at_time(&self, now: u64) { + pub(self) fn process_at_time(&self, mut now: u64) { let mut waker_list: [Option<Waker>; 32] = Default::default(); let mut waker_idx = 0; - let mut lock = self.get().lock(); + let mut lock = self.inner.lock(); - assert!(now >= lock.elapsed); + if now < lock.elapsed { + // Time went backwards! This normally shouldn't happen as the Rust language + // guarantees that an Instant is monotonic, but can happen when running + // Linux in a VM on a Windows host due to std incorrectly trusting the + // hardware clock to be monotonic. + // + // See <https://github.com/tokio-rs/tokio/issues/3619> for more information. + now = lock.elapsed; + } while let Some(entry) = lock.wheel.poll(now) { debug_assert!(unsafe { entry.is_pending() }); @@ -315,7 +291,7 @@ impl Handle { waker_idx = 0; - lock = self.get().lock(); + lock = self.inner.lock(); } } } @@ -346,7 +322,7 @@ impl Handle { /// `add_entry` must not be called concurrently. pub(self) unsafe fn clear_entry(&self, entry: NonNull<TimerShared>) { unsafe { - let mut lock = self.get().lock(); + let mut lock = self.inner.lock(); if entry.as_ref().might_be_registered() { lock.wheel.remove(entry); @@ -362,9 +338,14 @@ impl Handle { /// driver. No other threads are allowed to concurrently manipulate the /// timer at all (the current thread should hold an exclusive reference to /// the `TimerEntry`) - pub(self) unsafe fn reregister(&self, new_tick: u64, entry: NonNull<TimerShared>) { + pub(self) unsafe fn reregister( + &self, + unpark: &IoHandle, + new_tick: u64, + entry: NonNull<TimerShared>, + ) { let waker = unsafe { - let mut lock = self.get().lock(); + let mut lock = self.inner.lock(); // We may have raced with a firing/deregistration, so check before // deregistering. @@ -390,12 +371,12 @@ impl Handle { .map(|next_wake| when < next_wake.get()) .unwrap_or(true) { - lock.unpark.unpark(); + unpark.unpark(); } None } - Err((entry, super::error::InsertError::Elapsed)) => unsafe { + Err((entry, crate::time::error::InsertError::Elapsed)) => unsafe { entry.fire(Ok(())) }, } @@ -411,94 +392,17 @@ impl Handle { waker.wake(); } } -} - -impl<P> Park for Driver<P> -where - P: Park + 'static, -{ - type Unpark = TimerUnpark<P>; - type Error = P::Error; - - fn unpark(&self) -> Self::Unpark { - TimerUnpark::new(self) - } - - fn park(&mut self) -> Result<(), Self::Error> { - self.park_internal(None) - } - - fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> { - self.park_internal(Some(duration)) - } - - fn shutdown(&mut self) { - if self.handle.is_shutdown() { - return; - } - - self.handle.get().is_shutdown.store(true, Ordering::SeqCst); - - // Advance time forward to the end of time. - - self.handle.process_at_time(u64::MAX); - - self.park.shutdown(); - } -} - -impl<P> Drop for Driver<P> -where - P: Park + 'static, -{ - fn drop(&mut self) { - self.shutdown(); - } -} -pub(crate) struct TimerUnpark<P: Park + 'static> { - inner: P::Unpark, - - #[cfg(feature = "test-util")] - did_wake: Arc<AtomicBool>, -} - -impl<P: Park + 'static> TimerUnpark<P> { - fn new(driver: &Driver<P>) -> TimerUnpark<P> { - TimerUnpark { - inner: driver.park.unpark(), - - #[cfg(feature = "test-util")] - did_wake: driver.did_wake.clone(), + cfg_test_util! { + fn did_wake(&self) -> bool { + self.inner.did_wake.swap(false, Ordering::SeqCst) } } } -impl<P: Park + 'static> Unpark for TimerUnpark<P> { - fn unpark(&self) { - #[cfg(feature = "test-util")] - self.did_wake.store(true, Ordering::SeqCst); - - self.inner.unpark(); - } -} - // ===== impl Inner ===== impl Inner { - pub(self) fn new(time_source: ClockTime, unpark: Box<dyn Unpark>) -> Self { - Inner { - state: Mutex::new(InnerState { - time_source, - elapsed: 0, - next_wake: None, - unpark, - wheel: wheel::Wheel::new(), - }), - is_shutdown: AtomicBool::new(false), - } - } - /// Locks the driver's inner structure pub(super) fn lock(&self) -> crate::loom::sync::MutexGuard<'_, InnerState> { self.state.lock() diff --git a/vendor/tokio/src/runtime/time/source.rs b/vendor/tokio/src/runtime/time/source.rs new file mode 100644 index 000000000..4647bc412 --- /dev/null +++ b/vendor/tokio/src/runtime/time/source.rs @@ -0,0 +1,39 @@ +use super::MAX_SAFE_MILLIS_DURATION; +use crate::time::{Clock, Duration, Instant}; + +/// A structure which handles conversion from Instants to u64 timestamps. +#[derive(Debug)] +pub(crate) struct TimeSource { + start_time: Instant, +} + +impl TimeSource { + pub(crate) fn new(clock: &Clock) -> Self { + Self { + start_time: clock.now(), + } + } + + pub(crate) fn deadline_to_tick(&self, t: Instant) -> u64 { + // Round up to the end of a ms + self.instant_to_tick(t + Duration::from_nanos(999_999)) + } + + pub(crate) fn instant_to_tick(&self, t: Instant) -> u64 { + // round up + let dur: Duration = t + .checked_duration_since(self.start_time) + .unwrap_or_else(|| Duration::from_secs(0)); + let ms = dur.as_millis(); + + ms.try_into().unwrap_or(MAX_SAFE_MILLIS_DURATION) + } + + pub(crate) fn tick_to_duration(&self, t: u64) -> Duration { + Duration::from_millis(t) + } + + pub(crate) fn now(&self, clock: &Clock) -> u64 { + self.instant_to_tick(clock.now()) + } +} diff --git a/vendor/tokio/src/time/driver/tests/mod.rs b/vendor/tokio/src/runtime/time/tests/mod.rs index 7c5cf1fd0..155d99a34 100644 --- a/vendor/tokio/src/time/driver/tests/mod.rs +++ b/vendor/tokio/src/runtime/time/tests/mod.rs @@ -1,33 +1,27 @@ +#![cfg(not(tokio_wasi))] + use std::{task::Context, time::Duration}; #[cfg(not(loom))] use futures::task::noop_waker_ref; +use crate::loom::sync::atomic::{AtomicBool, Ordering}; use crate::loom::sync::Arc; use crate::loom::thread; -use crate::{ - loom::sync::atomic::{AtomicBool, Ordering}, - park::Unpark, -}; - -use super::{Handle, TimerEntry}; -struct MockUnpark {} -impl Unpark for MockUnpark { - fn unpark(&self) {} -} -impl MockUnpark { - fn mock() -> Box<dyn Unpark> { - Box::new(Self {}) - } -} +use super::TimerEntry; fn block_on<T>(f: impl std::future::Future<Output = T>) -> T { #[cfg(loom)] return loom::future::block_on(f); #[cfg(not(loom))] - return futures::executor::block_on(f); + { + let rt = crate::runtime::Builder::new_current_thread() + .build() + .unwrap(); + rt.block_on(f) + } } fn model(f: impl Fn() + Send + Sync + 'static) { @@ -38,18 +32,26 @@ fn model(f: impl Fn() + Send + Sync + 'static) { f(); } +fn rt(start_paused: bool) -> crate::runtime::Runtime { + crate::runtime::Builder::new_current_thread() + .enable_time() + .start_paused(start_paused) + .build() + .unwrap() +} + #[test] fn single_timer() { model(|| { - let clock = crate::time::clock::Clock::new(true, false); - let time_source = super::ClockTime::new(clock.clone()); - - let inner = super::Inner::new(time_source.clone(), MockUnpark::mock()); - let handle = Handle::new(Arc::new(inner)); + let rt = rt(false); + let handle = rt.handle(); let handle_ = handle.clone(); let jh = thread::spawn(move || { - let entry = TimerEntry::new(&handle_, clock.now() + Duration::from_secs(1)); + let entry = TimerEntry::new( + &handle_.inner, + handle_.inner.driver().clock().now() + Duration::from_secs(1), + ); pin!(entry); block_on(futures::future::poll_fn(|cx| { @@ -60,10 +62,13 @@ fn single_timer() { thread::yield_now(); + let time = handle.inner.driver().time(); + let clock = handle.inner.driver().clock(); + // This may or may not return Some (depending on how it races with the // thread). If it does return None, however, the timer should complete // synchronously. - handle.process_at_time(time_source.now() + 2_000_000_000); + time.process_at_time(time.time_source().now(clock) + 2_000_000_000); jh.join().unwrap(); }) @@ -72,15 +77,15 @@ fn single_timer() { #[test] fn drop_timer() { model(|| { - let clock = crate::time::clock::Clock::new(true, false); - let time_source = super::ClockTime::new(clock.clone()); - - let inner = super::Inner::new(time_source.clone(), MockUnpark::mock()); - let handle = Handle::new(Arc::new(inner)); + let rt = rt(false); + let handle = rt.handle(); let handle_ = handle.clone(); let jh = thread::spawn(move || { - let entry = TimerEntry::new(&handle_, clock.now() + Duration::from_secs(1)); + let entry = TimerEntry::new( + &handle_.inner, + handle_.inner.driver().clock().now() + Duration::from_secs(1), + ); pin!(entry); let _ = entry @@ -93,8 +98,11 @@ fn drop_timer() { thread::yield_now(); + let time = handle.inner.driver().time(); + let clock = handle.inner.driver().clock(); + // advance 2s in the future. - handle.process_at_time(time_source.now() + 2_000_000_000); + time.process_at_time(time.time_source().now(clock) + 2_000_000_000); jh.join().unwrap(); }) @@ -103,15 +111,15 @@ fn drop_timer() { #[test] fn change_waker() { model(|| { - let clock = crate::time::clock::Clock::new(true, false); - let time_source = super::ClockTime::new(clock.clone()); - - let inner = super::Inner::new(time_source.clone(), MockUnpark::mock()); - let handle = Handle::new(Arc::new(inner)); + let rt = rt(false); + let handle = rt.handle(); let handle_ = handle.clone(); let jh = thread::spawn(move || { - let entry = TimerEntry::new(&handle_, clock.now() + Duration::from_secs(1)); + let entry = TimerEntry::new( + &handle_.inner, + handle_.inner.driver().clock().now() + Duration::from_secs(1), + ); pin!(entry); let _ = entry @@ -126,8 +134,11 @@ fn change_waker() { thread::yield_now(); + let time = handle.inner.driver().time(); + let clock = handle.inner.driver().clock(); + // advance 2s - handle.process_at_time(time_source.now() + 2_000_000_000); + time.process_at_time(time.time_source().now(clock) + 2_000_000_000); jh.join().unwrap(); }) @@ -138,25 +149,22 @@ fn reset_future() { model(|| { let finished_early = Arc::new(AtomicBool::new(false)); - let clock = crate::time::clock::Clock::new(true, false); - let time_source = super::ClockTime::new(clock.clone()); - - let inner = super::Inner::new(time_source.clone(), MockUnpark::mock()); - let handle = Handle::new(Arc::new(inner)); + let rt = rt(false); + let handle = rt.handle(); let handle_ = handle.clone(); let finished_early_ = finished_early.clone(); - let start = clock.now(); + let start = handle.inner.driver().clock().now(); let jh = thread::spawn(move || { - let entry = TimerEntry::new(&handle_, start + Duration::from_secs(1)); + let entry = TimerEntry::new(&handle_.inner, start + Duration::from_secs(1)); pin!(entry); let _ = entry .as_mut() .poll_elapsed(&mut Context::from_waker(futures::task::noop_waker_ref())); - entry.as_mut().reset(start + Duration::from_secs(2)); + entry.as_mut().reset(start + Duration::from_secs(2), true); // shouldn't complete before 2s block_on(futures::future::poll_fn(|cx| { @@ -169,12 +177,22 @@ fn reset_future() { thread::yield_now(); + let handle = handle.inner.driver().time(); + // This may or may not return a wakeup time. - handle.process_at_time(time_source.instant_to_tick(start + Duration::from_millis(1500))); + handle.process_at_time( + handle + .time_source() + .instant_to_tick(start + Duration::from_millis(1500)), + ); assert!(!finished_early.load(Ordering::Relaxed)); - handle.process_at_time(time_source.instant_to_tick(start + Duration::from_millis(2500))); + handle.process_at_time( + handle + .time_source() + .instant_to_tick(start + Duration::from_millis(2500)), + ); jh.join().unwrap(); @@ -182,23 +200,27 @@ fn reset_future() { }) } +#[cfg(not(loom))] +fn normal_or_miri<T>(normal: T, miri: T) -> T { + if cfg!(miri) { + miri + } else { + normal + } +} + #[test] #[cfg(not(loom))] fn poll_process_levels() { - let clock = crate::time::clock::Clock::new(true, false); - clock.pause(); - - let time_source = super::ClockTime::new(clock.clone()); - - let inner = super::Inner::new(time_source, MockUnpark::mock()); - let handle = Handle::new(Arc::new(inner)); + let rt = rt(true); + let handle = rt.handle(); let mut entries = vec![]; - for i in 0..1024 { + for i in 0..normal_or_miri(1024, 64) { let mut entry = Box::pin(TimerEntry::new( - &handle, - clock.now() + Duration::from_millis(i), + &handle.inner, + handle.inner.driver().clock().now() + Duration::from_millis(i), )); let _ = entry @@ -208,8 +230,9 @@ fn poll_process_levels() { entries.push(entry); } - for t in 1..1024 { - handle.process_at_time(t as u64); + for t in 1..normal_or_miri(1024, 64) { + handle.inner.driver().time().process_at_time(t as u64); + for (deadline, future) in entries.iter_mut().enumerate() { let mut context = Context::from_waker(noop_waker_ref()); if deadline <= t { @@ -226,62 +249,19 @@ fn poll_process_levels() { fn poll_process_levels_targeted() { let mut context = Context::from_waker(noop_waker_ref()); - let clock = crate::time::clock::Clock::new(true, false); - clock.pause(); - - let time_source = super::ClockTime::new(clock.clone()); + let rt = rt(true); + let handle = rt.handle(); - let inner = super::Inner::new(time_source, MockUnpark::mock()); - let handle = Handle::new(Arc::new(inner)); - - let e1 = TimerEntry::new(&handle, clock.now() + Duration::from_millis(193)); + let e1 = TimerEntry::new( + &handle.inner, + handle.inner.driver().clock().now() + Duration::from_millis(193), + ); pin!(e1); + let handle = handle.inner.driver().time(); + handle.process_at_time(62); assert!(e1.as_mut().poll_elapsed(&mut context).is_pending()); handle.process_at_time(192); handle.process_at_time(192); } - -/* -#[test] -fn balanced_incr_and_decr() { - const OPS: usize = 5; - - fn incr(inner: Arc<Inner>) { - for _ in 0..OPS { - inner.increment().expect("increment should not have failed"); - thread::yield_now(); - } - } - - fn decr(inner: Arc<Inner>) { - let mut ops_performed = 0; - while ops_performed < OPS { - if inner.num(Ordering::Relaxed) > 0 { - ops_performed += 1; - inner.decrement(); - } - thread::yield_now(); - } - } - - loom::model(|| { - let unpark = Box::new(MockUnpark); - let instant = Instant::now(); - - let inner = Arc::new(Inner::new(instant, unpark)); - - let incr_inner = inner.clone(); - let decr_inner = inner.clone(); - - let incr_hndle = thread::spawn(move || incr(incr_inner)); - let decr_hndle = thread::spawn(move || decr(decr_inner)); - - incr_hndle.join().expect("should never fail"); - decr_hndle.join().expect("should never fail"); - - assert_eq!(inner.num(Ordering::SeqCst), 0); - }) -} -*/ diff --git a/vendor/tokio/src/time/driver/wheel/level.rs b/vendor/tokio/src/runtime/time/wheel/level.rs index 81d6b58c7..7e48ff5c5 100644 --- a/vendor/tokio/src/time/driver/wheel/level.rs +++ b/vendor/tokio/src/runtime/time/wheel/level.rs @@ -1,6 +1,4 @@ -use crate::time::driver::TimerHandle; - -use crate::time::driver::{EntryList, TimerShared}; +use crate::runtime::time::{EntryList, TimerHandle, TimerShared}; use std::{fmt, ptr::NonNull}; @@ -53,7 +51,7 @@ impl Level { // However, that is only supported for arrays of size // 32 or fewer. So in our case we have to explicitly // invoke the constructor for each array element. - let ctor = || EntryList::default(); + let ctor = EntryList::default; Level { level, @@ -143,8 +141,9 @@ impl Level { let level_range = level_range(self.level); let slot_range = slot_range(self.level); - // TODO: This can probably be simplified w/ power of 2 math - let level_start = now - (now % level_range); + // Compute the start date of the current level by masking the low bits + // of `now` (`level_range` is a power of 2). + let level_start = now & !(level_range - 1); let mut deadline = level_start + slot as u64 * slot_range; if deadline <= now { @@ -250,7 +249,7 @@ fn level_range(level: usize) -> u64 { LEVEL_MULT as u64 * slot_range(level) } -/// Convert a duration (milliseconds) and a level to a slot position +/// Converts a duration (milliseconds) and a level to a slot position. fn slot_for(duration: u64, level: usize) -> usize { ((duration >> (level * 6)) % LEVEL_MULT as u64) as usize } diff --git a/vendor/tokio/src/time/driver/wheel/mod.rs b/vendor/tokio/src/runtime/time/wheel/mod.rs index 5a40f6db8..bf13b7b24 100644 --- a/vendor/tokio/src/time/driver/wheel/mod.rs +++ b/vendor/tokio/src/runtime/time/wheel/mod.rs @@ -1,4 +1,4 @@ -use crate::time::driver::{TimerHandle, TimerShared}; +use crate::runtime::time::{TimerHandle, TimerShared}; use crate::time::error::InsertError; mod level; @@ -46,11 +46,11 @@ pub(crate) struct Wheel { /// precision of 1 millisecond. const NUM_LEVELS: usize = 6; -/// The maximum duration of a `Sleep` +/// The maximum duration of a `Sleep`. pub(super) const MAX_DURATION: u64 = (1 << (6 * NUM_LEVELS)) - 1; impl Wheel { - /// Create a new timing wheel + /// Creates a new timing wheel. pub(crate) fn new() -> Wheel { let levels = (0..NUM_LEVELS).map(Level::new).collect(); @@ -61,13 +61,13 @@ impl Wheel { } } - /// Return the number of milliseconds that have elapsed since the timing + /// Returns the number of milliseconds that have elapsed since the timing /// wheel's creation. pub(crate) fn elapsed(&self) -> u64 { self.elapsed } - /// Insert an entry into the timing wheel. + /// Inserts an entry into the timing wheel. /// /// # Arguments /// @@ -115,7 +115,7 @@ impl Wheel { Ok(when) } - /// Remove `item` from the timing wheel. + /// Removes `item` from the timing wheel. pub(crate) unsafe fn remove(&mut self, item: NonNull<TimerShared>) { unsafe { let when = item.as_ref().cached_when(); @@ -136,7 +136,7 @@ impl Wheel { } } - /// Instant at which to poll + /// Instant at which to poll. pub(crate) fn poll_at(&self) -> Option<u64> { self.next_expiration().map(|expiration| expiration.deadline) } @@ -148,23 +148,13 @@ impl Wheel { return Some(handle); } - // under what circumstances is poll.expiration Some vs. None? - let expiration = self.next_expiration().and_then(|expiration| { - if expiration.deadline > now { - None - } else { - Some(expiration) - } - }); - - match expiration { - Some(ref expiration) if expiration.deadline > now => return None, - Some(ref expiration) => { + match self.next_expiration() { + Some(ref expiration) if expiration.deadline <= now => { self.process_expiration(expiration); self.set_elapsed(expiration.deadline); } - None => { + _ => { // in this case the poll did not indicate an expiration // _and_ we were not able to find a next expiration in // the current list of timers. advance to the poll's diff --git a/vendor/tokio/src/signal/ctrl_c.rs b/vendor/tokio/src/signal/ctrl_c.rs index 1eeeb85aa..b26ab7ead 100644 --- a/vendor/tokio/src/signal/ctrl_c.rs +++ b/vendor/tokio/src/signal/ctrl_c.rs @@ -47,6 +47,15 @@ use std::io; /// println!("received ctrl-c event"); /// } /// ``` +/// +/// Listen in the background: +/// +/// ```rust,no_run +/// tokio::spawn(async move { +/// tokio::signal::ctrl_c().await.unwrap(); +/// // Your handler here +/// }); +/// ``` pub async fn ctrl_c() -> io::Result<()> { os_impl::ctrl_c()?.recv().await; Ok(()) diff --git a/vendor/tokio/src/signal/mod.rs b/vendor/tokio/src/signal/mod.rs index fe572f041..3aacc60ef 100644 --- a/vendor/tokio/src/signal/mod.rs +++ b/vendor/tokio/src/signal/mod.rs @@ -1,4 +1,4 @@ -//! Asynchronous signal handling for Tokio +//! Asynchronous signal handling for Tokio. //! //! Note that signal handling is in general a very tricky topic and should be //! used with great care. This crate attempts to implement 'best practice' for @@ -48,7 +48,7 @@ use std::task::{Context, Poll}; mod ctrl_c; pub use ctrl_c::ctrl_c; -mod registry; +pub(crate) mod registry; mod os { #[cfg(unix)] diff --git a/vendor/tokio/src/signal/registry.rs b/vendor/tokio/src/signal/registry.rs index 8b89108a6..48e98c832 100644 --- a/vendor/tokio/src/signal/registry.rs +++ b/vendor/tokio/src/signal/registry.rs @@ -1,12 +1,10 @@ #![allow(clippy::unit_arg)] use crate::signal::os::{OsExtraData, OsStorage}; - use crate::sync::watch; +use crate::util::once_cell::OnceCell; -use once_cell::sync::Lazy; use std::ops; -use std::pin::Pin; use std::sync::atomic::{AtomicBool, Ordering}; pub(crate) type EventId = usize; @@ -152,19 +150,25 @@ impl Globals { } } -pub(crate) fn globals() -> Pin<&'static Globals> +fn globals_init() -> Globals where OsExtraData: 'static + Send + Sync + Init, OsStorage: 'static + Send + Sync + Init, { - static GLOBALS: Lazy<Pin<Box<Globals>>> = Lazy::new(|| { - Box::pin(Globals { - extra: OsExtraData::init(), - registry: Registry::new(OsStorage::init()), - }) - }); - - GLOBALS.as_ref() + Globals { + extra: OsExtraData::init(), + registry: Registry::new(OsStorage::init()), + } +} + +pub(crate) fn globals() -> &'static Globals +where + OsExtraData: 'static + Send + Sync + Init, + OsStorage: 'static + Send + Sync + Init, +{ + static GLOBALS: OnceCell<Globals> = OnceCell::new(); + + GLOBALS.get(globals_init) } #[cfg(all(test, not(loom)))] @@ -202,7 +206,12 @@ mod tests { registry.broadcast(); // Yield so the previous broadcast can get received - crate::time::sleep(std::time::Duration::from_millis(10)).await; + // + // This yields many times since the block_on task is only polled every 61 + // ticks. + for _ in 0..100 { + crate::task::yield_now().await; + } // Send subsequent signal registry.record_event(0); @@ -232,7 +241,7 @@ mod tests { #[test] fn record_invalid_event_does_nothing() { let registry = Registry::new(vec![EventInfo::default()]); - registry.record_event(42); + registry.record_event(1302); } #[test] @@ -240,17 +249,17 @@ mod tests { let registry = Registry::new(vec![EventInfo::default(), EventInfo::default()]); registry.record_event(0); - assert_eq!(false, registry.broadcast()); + assert!(!registry.broadcast()); let first = registry.register_listener(0); let second = registry.register_listener(1); registry.record_event(0); - assert_eq!(true, registry.broadcast()); + assert!(registry.broadcast()); drop(first); registry.record_event(0); - assert_eq!(false, registry.broadcast()); + assert!(!registry.broadcast()); drop(second); } diff --git a/vendor/tokio/src/signal/reusable_box.rs b/vendor/tokio/src/signal/reusable_box.rs index 426ecb06f..796fa210b 100644 --- a/vendor/tokio/src/signal/reusable_box.rs +++ b/vendor/tokio/src/signal/reusable_box.rs @@ -30,7 +30,7 @@ impl<T> ReusableBoxFuture<T> { Self { boxed } } - /// Replace the future currently stored in this box. + /// Replaces the future currently stored in this box. /// /// This reallocates if and only if the layout of the provided future is /// different from the layout of the currently stored future. @@ -43,7 +43,7 @@ impl<T> ReusableBoxFuture<T> { } } - /// Replace the future currently stored in this box. + /// Replaces the future currently stored in this box. /// /// This function never reallocates, but returns an error if the provided /// future has a different size or alignment from the currently stored @@ -70,7 +70,7 @@ impl<T> ReusableBoxFuture<T> { } } - /// Set the current future. + /// Sets the current future. /// /// # Safety /// @@ -103,14 +103,14 @@ impl<T> ReusableBoxFuture<T> { } } - /// Get a pinned reference to the underlying future. + /// Gets a pinned reference to the underlying future. pub(crate) fn get_pin(&mut self) -> Pin<&mut (dyn Future<Output = T> + Send)> { // SAFETY: The user of this box cannot move the box, and we do not move it // either. unsafe { Pin::new_unchecked(self.boxed.as_mut()) } } - /// Poll the future stored inside this box. + /// Polls the future stored inside this box. pub(crate) fn poll(&mut self, cx: &mut Context<'_>) -> Poll<T> { self.get_pin().poll(cx) } @@ -119,7 +119,7 @@ impl<T> ReusableBoxFuture<T> { impl<T> Future for ReusableBoxFuture<T> { type Output = T; - /// Poll the future stored inside this box. + /// Polls the future stored inside this box. fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> { Pin::into_inner(self).get_pin().poll(cx) } diff --git a/vendor/tokio/src/signal/unix.rs b/vendor/tokio/src/signal/unix.rs index f96b2f4c2..ae5c13085 100644 --- a/vendor/tokio/src/signal/unix.rs +++ b/vendor/tokio/src/signal/unix.rs @@ -4,30 +4,33 @@ //! `Signal` type for receiving notifications of signals. #![cfg(unix)] +#![cfg_attr(docsrs, doc(cfg(all(unix, feature = "signal"))))] +use crate::runtime::scheduler; +use crate::runtime::signal::Handle; use crate::signal::registry::{globals, EventId, EventInfo, Globals, Init, Storage}; use crate::signal::RxFuture; use crate::sync::watch; use mio::net::UnixStream; use std::io::{self, Error, ErrorKind, Write}; -use std::pin::Pin; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Once; use std::task::{Context, Poll}; -pub(crate) mod driver; -use self::driver::Handle; - pub(crate) type OsStorage = Vec<SignalInfo>; -// Number of different unix signals -// (FreeBSD has 33) -const SIGNUM: usize = 33; - impl Init for OsStorage { fn init() -> Self { - (0..SIGNUM).map(|_| SignalInfo::default()).collect() + // There are reliable signals ranging from 1 to 33 available on every Unix platform. + #[cfg(not(target_os = "linux"))] + let possible = 0..=33; + + // On Linux, there are additional real-time signals available. + #[cfg(target_os = "linux")] + let possible = 0..=libc::SIGRTMAX(); + + possible.map(|_| SignalInfo::default()).collect() } } @@ -47,7 +50,7 @@ impl Storage for OsStorage { #[derive(Debug)] pub(crate) struct OsExtraData { sender: UnixStream, - receiver: UnixStream, + pub(crate) receiver: UnixStream, } impl Init for OsExtraData { @@ -59,7 +62,7 @@ impl Init for OsExtraData { } /// Represents the specific kind of signal to listen for. -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)] pub struct SignalKind(libc::c_int); impl SignalKind { @@ -79,15 +82,26 @@ impl SignalKind { // unlikely to change to other types, but technically libc can change this // in the future minor version. // See https://github.com/tokio-rs/tokio/issues/3767 for more. - pub fn from_raw(signum: std::os::raw::c_int) -> Self { + pub const fn from_raw(signum: std::os::raw::c_int) -> Self { Self(signum as libc::c_int) } + /// Get the signal's numeric value. + /// + /// ```rust + /// # use tokio::signal::unix::SignalKind; + /// let kind = SignalKind::interrupt(); + /// assert_eq!(kind.as_raw_value(), libc::SIGINT); + /// ``` + pub const fn as_raw_value(&self) -> std::os::raw::c_int { + self.0 + } + /// Represents the SIGALRM signal. /// /// On Unix systems this signal is sent when a real-time timer has expired. /// By default, the process is terminated by this signal. - pub fn alarm() -> Self { + pub const fn alarm() -> Self { Self(libc::SIGALRM) } @@ -95,7 +109,7 @@ impl SignalKind { /// /// On Unix systems this signal is sent when the status of a child process /// has changed. By default, this signal is ignored. - pub fn child() -> Self { + pub const fn child() -> Self { Self(libc::SIGCHLD) } @@ -103,7 +117,7 @@ impl SignalKind { /// /// On Unix systems this signal is sent when the terminal is disconnected. /// By default, the process is terminated by this signal. - pub fn hangup() -> Self { + pub const fn hangup() -> Self { Self(libc::SIGHUP) } @@ -118,7 +132,7 @@ impl SignalKind { target_os = "netbsd", target_os = "openbsd" ))] - pub fn info() -> Self { + pub const fn info() -> Self { Self(libc::SIGINFO) } @@ -126,7 +140,7 @@ impl SignalKind { /// /// On Unix systems this signal is sent to interrupt a program. /// By default, the process is terminated by this signal. - pub fn interrupt() -> Self { + pub const fn interrupt() -> Self { Self(libc::SIGINT) } @@ -134,7 +148,7 @@ impl SignalKind { /// /// On Unix systems this signal is sent when I/O operations are possible /// on some file descriptor. By default, this signal is ignored. - pub fn io() -> Self { + pub const fn io() -> Self { Self(libc::SIGIO) } @@ -143,7 +157,7 @@ impl SignalKind { /// On Unix systems this signal is sent when the process attempts to write /// to a pipe which has no reader. By default, the process is terminated by /// this signal. - pub fn pipe() -> Self { + pub const fn pipe() -> Self { Self(libc::SIGPIPE) } @@ -152,7 +166,7 @@ impl SignalKind { /// On Unix systems this signal is sent to issue a shutdown of the /// process, after which the OS will dump the process core. /// By default, the process is terminated by this signal. - pub fn quit() -> Self { + pub const fn quit() -> Self { Self(libc::SIGQUIT) } @@ -160,7 +174,7 @@ impl SignalKind { /// /// On Unix systems this signal is sent to issue a shutdown of the /// process. By default, the process is terminated by this signal. - pub fn terminate() -> Self { + pub const fn terminate() -> Self { Self(libc::SIGTERM) } @@ -168,7 +182,7 @@ impl SignalKind { /// /// On Unix systems this is a user defined signal. /// By default, the process is terminated by this signal. - pub fn user_defined1() -> Self { + pub const fn user_defined1() -> Self { Self(libc::SIGUSR1) } @@ -176,7 +190,7 @@ impl SignalKind { /// /// On Unix systems this is a user defined signal. /// By default, the process is terminated by this signal. - pub fn user_defined2() -> Self { + pub const fn user_defined2() -> Self { Self(libc::SIGUSR2) } @@ -184,11 +198,23 @@ impl SignalKind { /// /// On Unix systems this signal is sent when the terminal window is resized. /// By default, this signal is ignored. - pub fn window_change() -> Self { + pub const fn window_change() -> Self { Self(libc::SIGWINCH) } } +impl From<std::os::raw::c_int> for SignalKind { + fn from(signum: std::os::raw::c_int) -> Self { + Self::from_raw(signum as libc::c_int) + } +} + +impl From<SignalKind> for std::os::raw::c_int { + fn from(kind: SignalKind) -> Self { + kind.as_raw_value() + } +} + pub(crate) struct SignalInfo { event_info: EventInfo, init: Once, @@ -213,7 +239,7 @@ impl Default for SignalInfo { /// 2. Wake up the driver by writing a byte to a pipe /// /// Those two operations should both be async-signal safe. -fn action(globals: Pin<&'static Globals>, signal: libc::c_int) { +fn action(globals: &'static Globals, signal: libc::c_int) { globals.record_event(signal as EventId); // Send a wakeup, ignore any errors (anything reasonably possible is @@ -266,7 +292,11 @@ fn signal_enable(signal: SignalKind, handle: &Handle) -> io::Result<()> { } } -/// A stream of events for receiving a particular type of OS signal. +/// An listener for receiving a particular type of OS signal. +/// +/// The listener can be turned into a `Stream` using [`SignalStream`]. +/// +/// [`SignalStream`]: https://docs.rs/tokio-stream/latest/tokio_stream/wrappers/struct.SignalStream.html /// /// In general signal handling on Unix is a pretty tricky topic, and this /// structure is no exception! There are some important limitations to keep in @@ -281,7 +311,7 @@ fn signal_enable(signal: SignalKind, handle: &Handle) -> io::Result<()> { /// Once `poll` has been called, however, a further signal is guaranteed to /// be yielded as an item. /// -/// Put another way, any element pulled off the returned stream corresponds to +/// Put another way, any element pulled off the returned listener corresponds to /// *at least one* signal, but possibly more. /// /// * Signal handling in general is relatively inefficient. Although some @@ -319,11 +349,11 @@ fn signal_enable(signal: SignalKind, handle: &Handle) -> io::Result<()> { /// #[tokio::main] /// async fn main() -> Result<(), Box<dyn std::error::Error>> { /// // An infinite stream of hangup signals. -/// let mut stream = signal(SignalKind::hangup())?; +/// let mut sig = signal(SignalKind::hangup())?; /// /// // Print whenever a HUP signal is received /// loop { -/// stream.recv().await; +/// sig.recv().await; /// println!("got signal HUP"); /// } /// } @@ -334,7 +364,7 @@ pub struct Signal { inner: RxFuture, } -/// Creates a new stream which will receive notifications when the current +/// Creates a new listener which will receive notifications when the current /// process receives the specified signal `kind`. /// /// This function will create a new stream which binds to the default reactor. @@ -356,8 +386,15 @@ pub struct Signal { /// * If the previous initialization of this specific signal failed. /// * If the signal is one of /// [`signal_hook::FORBIDDEN`](fn@signal_hook_registry::register#panics) +/// +/// # Panics +/// +/// This function panics if there is no current reactor set, or if the `rt` +/// feature flag is not enabled. +#[track_caller] pub fn signal(kind: SignalKind) -> io::Result<Signal> { - let rx = signal_with_handle(kind, &Handle::current())?; + let handle = scheduler::Handle::current(); + let rx = signal_with_handle(kind, handle.driver().signal())?; Ok(Signal { inner: RxFuture::new(rx), @@ -379,6 +416,12 @@ impl Signal { /// /// `None` is returned if no more events can be received by this stream. /// + /// # Cancel safety + /// + /// This method is cancel safe. If you use it as the event in a + /// [`tokio::select!`](crate::select) statement and some other branch + /// completes first, then it is guaranteed that no signal is lost. + /// /// # Examples /// /// Wait for SIGHUP @@ -473,4 +516,15 @@ mod tests { ) .unwrap_err(); } + + #[test] + fn from_c_int() { + assert_eq!(SignalKind::from(2), SignalKind::interrupt()); + } + + #[test] + fn into_c_int() { + let value: std::os::raw::c_int = SignalKind::interrupt().into(); + assert_eq!(value, libc::SIGINT as _); + } } diff --git a/vendor/tokio/src/signal/unix/driver.rs b/vendor/tokio/src/signal/unix/driver.rs deleted file mode 100644 index 5fe7c354c..000000000 --- a/vendor/tokio/src/signal/unix/driver.rs +++ /dev/null @@ -1,207 +0,0 @@ -#![cfg_attr(not(feature = "rt"), allow(dead_code))] - -//! Signal driver - -use crate::io::driver::{Driver as IoDriver, Interest}; -use crate::io::PollEvented; -use crate::park::Park; -use crate::signal::registry::globals; - -use mio::net::UnixStream; -use std::io::{self, Read}; -use std::ptr; -use std::sync::{Arc, Weak}; -use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker}; -use std::time::Duration; - -/// Responsible for registering wakeups when an OS signal is received, and -/// subsequently dispatching notifications to any signal listeners as appropriate. -/// -/// Note: this driver relies on having an enabled IO driver in order to listen to -/// pipe write wakeups. -#[derive(Debug)] -pub(crate) struct Driver { - /// Thread parker. The `Driver` park implementation delegates to this. - park: IoDriver, - - /// A pipe for receiving wake events from the signal handler - receiver: PollEvented<UnixStream>, - - /// Shared state - inner: Arc<Inner>, -} - -#[derive(Clone, Debug, Default)] -pub(crate) struct Handle { - inner: Weak<Inner>, -} - -#[derive(Debug)] -pub(super) struct Inner(()); - -// ===== impl Driver ===== - -impl Driver { - /// Creates a new signal `Driver` instance that delegates wakeups to `park`. - pub(crate) fn new(park: IoDriver) -> io::Result<Self> { - use std::mem::ManuallyDrop; - use std::os::unix::io::{AsRawFd, FromRawFd}; - - // NB: We give each driver a "fresh" receiver file descriptor to avoid - // the issues described in alexcrichton/tokio-process#42. - // - // In the past we would reuse the actual receiver file descriptor and - // swallow any errors around double registration of the same descriptor. - // I'm not sure if the second (failed) registration simply doesn't end - // up receiving wake up notifications, or there could be some race - // condition when consuming readiness events, but having distinct - // descriptors for distinct PollEvented instances appears to mitigate - // this. - // - // Unfortunately we cannot just use a single global PollEvented instance - // either, since we can't compare Handles or assume they will always - // point to the exact same reactor. - // - // Mio 0.7 removed `try_clone()` as an API due to unexpected behavior - // with registering dups with the same reactor. In this case, duping is - // safe as each dup is registered with separate reactors **and** we - // only expect at least one dup to receive the notification. - - // Manually drop as we don't actually own this instance of UnixStream. - let receiver_fd = globals().receiver.as_raw_fd(); - - // safety: there is nothing unsafe about this, but the `from_raw_fd` fn is marked as unsafe. - let original = - ManuallyDrop::new(unsafe { std::os::unix::net::UnixStream::from_raw_fd(receiver_fd) }); - let receiver = UnixStream::from_std(original.try_clone()?); - let receiver = PollEvented::new_with_interest_and_handle( - receiver, - Interest::READABLE | Interest::WRITABLE, - park.handle(), - )?; - - Ok(Self { - park, - receiver, - inner: Arc::new(Inner(())), - }) - } - - /// Returns a handle to this event loop which can be sent across threads - /// and can be used as a proxy to the event loop itself. - pub(crate) fn handle(&self) -> Handle { - Handle { - inner: Arc::downgrade(&self.inner), - } - } - - fn process(&self) { - // Check if the pipe is ready to read and therefore has "woken" us up - // - // To do so, we will `poll_read_ready` with a noop waker, since we don't - // need to actually be notified when read ready... - let waker = unsafe { Waker::from_raw(RawWaker::new(ptr::null(), &NOOP_WAKER_VTABLE)) }; - let mut cx = Context::from_waker(&waker); - - let ev = match self.receiver.registration().poll_read_ready(&mut cx) { - Poll::Ready(Ok(ev)) => ev, - Poll::Ready(Err(e)) => panic!("reactor gone: {}", e), - Poll::Pending => return, // No wake has arrived, bail - }; - - // Drain the pipe completely so we can receive a new readiness event - // if another signal has come in. - let mut buf = [0; 128]; - loop { - match (&*self.receiver).read(&mut buf) { - Ok(0) => panic!("EOF on self-pipe"), - Ok(_) => continue, // Keep reading - Err(e) if e.kind() == io::ErrorKind::WouldBlock => break, - Err(e) => panic!("Bad read on self-pipe: {}", e), - } - } - - self.receiver.registration().clear_readiness(ev); - - // Broadcast any signals which were received - globals().broadcast(); - } -} - -const NOOP_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new(noop_clone, noop, noop, noop); - -unsafe fn noop_clone(_data: *const ()) -> RawWaker { - RawWaker::new(ptr::null(), &NOOP_WAKER_VTABLE) -} - -unsafe fn noop(_data: *const ()) {} - -// ===== impl Park for Driver ===== - -impl Park for Driver { - type Unpark = <IoDriver as Park>::Unpark; - type Error = io::Error; - - fn unpark(&self) -> Self::Unpark { - self.park.unpark() - } - - fn park(&mut self) -> Result<(), Self::Error> { - self.park.park()?; - self.process(); - Ok(()) - } - - fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> { - self.park.park_timeout(duration)?; - self.process(); - Ok(()) - } - - fn shutdown(&mut self) { - self.park.shutdown() - } -} - -// ===== impl Handle ===== - -impl Handle { - pub(super) fn check_inner(&self) -> io::Result<()> { - if self.inner.strong_count() > 0 { - Ok(()) - } else { - Err(io::Error::new(io::ErrorKind::Other, "signal driver gone")) - } - } -} - -cfg_rt! { - impl Handle { - /// Returns a handle to the current driver - /// - /// # Panics - /// - /// This function panics if there is no current signal driver set. - pub(super) fn current() -> Self { - crate::runtime::context::signal_handle().expect( - "there is no signal driver running, must be called from the context of Tokio runtime", - ) - } - } -} - -cfg_not_rt! { - impl Handle { - /// Returns a handle to the current driver - /// - /// # Panics - /// - /// This function panics if there is no current signal driver set. - pub(super) fn current() -> Self { - panic!( - "there is no signal driver running, must be called from the context of Tokio runtime or with\ - `rt` enabled.", - ) - } - } -} diff --git a/vendor/tokio/src/signal/windows.rs b/vendor/tokio/src/signal/windows.rs index c231d6268..2f70f98b1 100644 --- a/vendor/tokio/src/signal/windows.rs +++ b/vendor/tokio/src/signal/windows.rs @@ -1,133 +1,28 @@ //! Windows-specific types for signal handling. //! -//! This module is only defined on Windows and allows receiving "ctrl-c" -//! and "ctrl-break" notifications. These events are listened for via the -//! `SetConsoleCtrlHandler` function which receives events of the type -//! `CTRL_C_EVENT` and `CTRL_BREAK_EVENT`. +//! This module is only defined on Windows and allows receiving "ctrl-c", +//! "ctrl-break", "ctrl-logoff", "ctrl-shutdown", and "ctrl-close" +//! notifications. These events are listened for via the `SetConsoleCtrlHandler` +//! function which receives the corresponding windows_sys event type. -#![cfg(windows)] +#![cfg(any(windows, docsrs))] +#![cfg_attr(docsrs, doc(cfg(all(windows, feature = "signal"))))] -use crate::signal::registry::{globals, EventId, EventInfo, Init, Storage}; use crate::signal::RxFuture; - -use std::convert::TryFrom; use std::io; -use std::sync::Once; use std::task::{Context, Poll}; -use winapi::shared::minwindef::{BOOL, DWORD, FALSE, TRUE}; -use winapi::um::consoleapi::SetConsoleCtrlHandler; -use winapi::um::wincon::{CTRL_BREAK_EVENT, CTRL_C_EVENT}; - -#[derive(Debug)] -pub(crate) struct OsStorage { - ctrl_c: EventInfo, - ctrl_break: EventInfo, -} -impl Init for OsStorage { - fn init() -> Self { - Self { - ctrl_c: EventInfo::default(), - ctrl_break: EventInfo::default(), - } - } -} - -impl Storage for OsStorage { - fn event_info(&self, id: EventId) -> Option<&EventInfo> { - match DWORD::try_from(id) { - Ok(CTRL_C_EVENT) => Some(&self.ctrl_c), - Ok(CTRL_BREAK_EVENT) => Some(&self.ctrl_break), - _ => None, - } - } +#[cfg(not(docsrs))] +#[path = "windows/sys.rs"] +mod imp; +#[cfg(not(docsrs))] +pub(crate) use self::imp::{OsExtraData, OsStorage}; - fn for_each<'a, F>(&'a self, mut f: F) - where - F: FnMut(&'a EventInfo), - { - f(&self.ctrl_c); - f(&self.ctrl_break); - } -} - -#[derive(Debug)] -pub(crate) struct OsExtraData {} - -impl Init for OsExtraData { - fn init() -> Self { - Self {} - } -} +#[cfg(docsrs)] +#[path = "windows/stub.rs"] +mod imp; -/// Stream of events discovered via `SetConsoleCtrlHandler`. -/// -/// This structure can be used to listen for events of the type `CTRL_C_EVENT` -/// and `CTRL_BREAK_EVENT`. The `Stream` trait is implemented for this struct -/// and will resolve for each notification received by the process. Note that -/// there are few limitations with this as well: -/// -/// * A notification to this process notifies *all* `Event` streams for that -/// event type. -/// * Notifications to an `Event` stream **are coalesced** if they aren't -/// processed quickly enough. This means that if two notifications are -/// received back-to-back, then the stream may only receive one item about the -/// two notifications. -#[must_use = "streams do nothing unless polled"] -#[derive(Debug)] -pub(crate) struct Event { - inner: RxFuture, -} - -impl Event { - fn new(signum: DWORD) -> io::Result<Self> { - global_init()?; - - let rx = globals().register_listener(signum as EventId); - - Ok(Self { - inner: RxFuture::new(rx), - }) - } -} - -fn global_init() -> io::Result<()> { - static INIT: Once = Once::new(); - - let mut init = None; - - INIT.call_once(|| unsafe { - let rc = SetConsoleCtrlHandler(Some(handler), TRUE); - let ret = if rc == 0 { - Err(io::Error::last_os_error()) - } else { - Ok(()) - }; - - init = Some(ret); - }); - - init.unwrap_or_else(|| Ok(())) -} - -unsafe extern "system" fn handler(ty: DWORD) -> BOOL { - let globals = globals(); - globals.record_event(ty as EventId); - - // According to https://docs.microsoft.com/en-us/windows/console/handlerroutine - // the handler routine is always invoked in a new thread, thus we don't - // have the same restrictions as in Unix signal handlers, meaning we can - // go ahead and perform the broadcast here. - if globals.broadcast() { - TRUE - } else { - // No one is listening for this notification any more - // let the OS fire the next (possibly the default) handler. - FALSE - } -} - -/// Creates a new stream which receives "ctrl-c" notifications sent to the +/// Creates a new listener which receives "ctrl-c" notifications sent to the /// process. /// /// # Examples @@ -137,12 +32,12 @@ unsafe extern "system" fn handler(ty: DWORD) -> BOOL { /// /// #[tokio::main] /// async fn main() -> Result<(), Box<dyn std::error::Error>> { -/// // An infinite stream of CTRL-C events. -/// let mut stream = ctrl_c()?; +/// // A listener of CTRL-C events. +/// let mut signal = ctrl_c()?; /// /// // Print whenever a CTRL-C event is received. /// for countdown in (0..3).rev() { -/// stream.recv().await; +/// signal.recv().await; /// println!("got CTRL-C. {} more to exit", countdown); /// } /// @@ -150,26 +45,32 @@ unsafe extern "system" fn handler(ty: DWORD) -> BOOL { /// } /// ``` pub fn ctrl_c() -> io::Result<CtrlC> { - Event::new(CTRL_C_EVENT).map(|inner| CtrlC { inner }) + Ok(CtrlC { + inner: self::imp::ctrl_c()?, + }) } -/// Represents a stream which receives "ctrl-c" notifications sent to the process +/// Represents a listener which receives "ctrl-c" notifications sent to the process /// via `SetConsoleCtrlHandler`. /// -/// A notification to this process notifies *all* streams listening for +/// This event can be turned into a `Stream` using [`CtrlCStream`]. +/// +/// [`CtrlCStream`]: https://docs.rs/tokio-stream/latest/tokio_stream/wrappers/struct.CtrlCStream.html +/// +/// A notification to this process notifies *all* receivers for /// this event. Moreover, the notifications **are coalesced** if they aren't processed /// quickly enough. This means that if two notifications are received back-to-back, -/// then the stream may only receive one item about the two notifications. -#[must_use = "streams do nothing unless polled"] +/// then the listener may only receive one item about the two notifications. +#[must_use = "listeners do nothing unless polled"] #[derive(Debug)] pub struct CtrlC { - inner: Event, + inner: RxFuture, } impl CtrlC { /// Receives the next signal notification event. /// - /// `None` is returned if no more events can be received by this stream. + /// `None` is returned if no more events can be received by the listener. /// /// # Examples /// @@ -178,12 +79,11 @@ impl CtrlC { /// /// #[tokio::main] /// async fn main() -> Result<(), Box<dyn std::error::Error>> { - /// // An infinite stream of CTRL-C events. - /// let mut stream = ctrl_c()?; + /// let mut signal = ctrl_c()?; /// /// // Print whenever a CTRL-C event is received. /// for countdown in (0..3).rev() { - /// stream.recv().await; + /// signal.recv().await; /// println!("got CTRL-C. {} more to exit", countdown); /// } /// @@ -191,13 +91,13 @@ impl CtrlC { /// } /// ``` pub async fn recv(&mut self) -> Option<()> { - self.inner.inner.recv().await + self.inner.recv().await } /// Polls to receive the next signal notification event, outside of an /// `async` context. /// - /// `None` is returned if no more events can be received by this stream. + /// `None` is returned if no more events can be received. /// /// # Examples /// @@ -223,27 +123,31 @@ impl CtrlC { /// } /// ``` pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<()>> { - self.inner.inner.poll_recv(cx) + self.inner.poll_recv(cx) } } -/// Represents a stream which receives "ctrl-break" notifications sent to the process +/// Represents a listener which receives "ctrl-break" notifications sent to the process /// via `SetConsoleCtrlHandler`. /// -/// A notification to this process notifies *all* streams listening for +/// This listener can be turned into a `Stream` using [`CtrlBreakStream`]. +/// +/// [`CtrlBreakStream`]: https://docs.rs/tokio-stream/latest/tokio_stream/wrappers/struct.CtrlBreakStream.html +/// +/// A notification to this process notifies *all* receivers for /// this event. Moreover, the notifications **are coalesced** if they aren't processed /// quickly enough. This means that if two notifications are received back-to-back, -/// then the stream may only receive one item about the two notifications. -#[must_use = "streams do nothing unless polled"] +/// then the listener may only receive one item about the two notifications. +#[must_use = "listeners do nothing unless polled"] #[derive(Debug)] pub struct CtrlBreak { - inner: Event, + inner: RxFuture, } impl CtrlBreak { /// Receives the next signal notification event. /// - /// `None` is returned if no more events can be received by this stream. + /// `None` is returned if no more events can be received by this listener. /// /// # Examples /// @@ -252,24 +156,24 @@ impl CtrlBreak { /// /// #[tokio::main] /// async fn main() -> Result<(), Box<dyn std::error::Error>> { - /// // An infinite stream of CTRL-BREAK events. - /// let mut stream = ctrl_break()?; + /// // A listener of CTRL-BREAK events. + /// let mut signal = ctrl_break()?; /// /// // Print whenever a CTRL-BREAK event is received. /// loop { - /// stream.recv().await; + /// signal.recv().await; /// println!("got signal CTRL-BREAK"); /// } /// } /// ``` pub async fn recv(&mut self) -> Option<()> { - self.inner.inner.recv().await + self.inner.recv().await } /// Polls to receive the next signal notification event, outside of an /// `async` context. /// - /// `None` is returned if no more events can be received by this stream. + /// `None` is returned if no more events can be received by this listener. /// /// # Examples /// @@ -295,11 +199,11 @@ impl CtrlBreak { /// } /// ``` pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<()>> { - self.inner.inner.poll_recv(cx) + self.inner.poll_recv(cx) } } -/// Creates a new stream which receives "ctrl-break" notifications sent to the +/// Creates a new listener which receives "ctrl-break" notifications sent to the /// process. /// /// # Examples @@ -309,67 +213,312 @@ impl CtrlBreak { /// /// #[tokio::main] /// async fn main() -> Result<(), Box<dyn std::error::Error>> { -/// // An infinite stream of CTRL-BREAK events. -/// let mut stream = ctrl_break()?; +/// // A listener of CTRL-BREAK events. +/// let mut signal = ctrl_break()?; /// /// // Print whenever a CTRL-BREAK event is received. /// loop { -/// stream.recv().await; +/// signal.recv().await; /// println!("got signal CTRL-BREAK"); /// } /// } /// ``` pub fn ctrl_break() -> io::Result<CtrlBreak> { - Event::new(CTRL_BREAK_EVENT).map(|inner| CtrlBreak { inner }) + Ok(CtrlBreak { + inner: self::imp::ctrl_break()?, + }) } -#[cfg(all(test, not(loom)))] -mod tests { - use super::*; - use crate::runtime::Runtime; +/// Creates a new listener which receives "ctrl-close" notifications sent to the +/// process. +/// +/// # Examples +/// +/// ```rust,no_run +/// use tokio::signal::windows::ctrl_close; +/// +/// #[tokio::main] +/// async fn main() -> Result<(), Box<dyn std::error::Error>> { +/// // A listener of CTRL-CLOSE events. +/// let mut signal = ctrl_close()?; +/// +/// // Print whenever a CTRL-CLOSE event is received. +/// for countdown in (0..3).rev() { +/// signal.recv().await; +/// println!("got CTRL-CLOSE. {} more to exit", countdown); +/// } +/// +/// Ok(()) +/// } +/// ``` +pub fn ctrl_close() -> io::Result<CtrlClose> { + Ok(CtrlClose { + inner: self::imp::ctrl_close()?, + }) +} - use tokio_test::{assert_ok, assert_pending, assert_ready_ok, task}; +/// Represents a listener which receives "ctrl-close" notitifications sent to the process +/// via 'SetConsoleCtrlHandler'. +/// +/// A notification to this process notifies *all* listeners listening for +/// this event. Moreover, the notifications **are coalesced** if they aren't processed +/// quickly enough. This means that if two notifications are received back-to-back, +/// then the listener may only receive one item about the two notifications. +#[must_use = "listeners do nothing unless polled"] +#[derive(Debug)] +pub struct CtrlClose { + inner: RxFuture, +} - #[test] - fn ctrl_c() { - let rt = rt(); - let _enter = rt.enter(); +impl CtrlClose { + /// Receives the next signal notification event. + /// + /// `None` is returned if no more events can be received by this listener. + /// + /// # Examples + /// + /// ```rust,no_run + /// use tokio::signal::windows::ctrl_close; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn std::error::Error>> { + /// // A listener of CTRL-CLOSE events. + /// let mut signal = ctrl_close()?; + /// + /// // Print whenever a CTRL-CLOSE event is received. + /// signal.recv().await; + /// println!("got CTRL-CLOSE. Cleaning up before exiting"); + /// + /// Ok(()) + /// } + /// ``` + pub async fn recv(&mut self) -> Option<()> { + self.inner.recv().await + } - let mut ctrl_c = task::spawn(crate::signal::ctrl_c()); + /// Polls to receive the next signal notification event, outside of an + /// `async` context. + /// + /// `None` is returned if no more events can be received by this listener. + /// + /// # Examples + /// + /// Polling from a manually implemented future + /// + /// ```rust,no_run + /// use std::pin::Pin; + /// use std::future::Future; + /// use std::task::{Context, Poll}; + /// use tokio::signal::windows::CtrlClose; + /// + /// struct MyFuture { + /// ctrl_close: CtrlClose, + /// } + /// + /// impl Future for MyFuture { + /// type Output = Option<()>; + /// + /// fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + /// println!("polling MyFuture"); + /// self.ctrl_close.poll_recv(cx) + /// } + /// } + /// ``` + pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<()>> { + self.inner.poll_recv(cx) + } +} - assert_pending!(ctrl_c.poll()); +/// Creates a new listener which receives "ctrl-shutdown" notifications sent to the +/// process. +/// +/// # Examples +/// +/// ```rust,no_run +/// use tokio::signal::windows::ctrl_shutdown; +/// +/// #[tokio::main] +/// async fn main() -> Result<(), Box<dyn std::error::Error>> { +/// // A listener of CTRL-SHUTDOWN events. +/// let mut signal = ctrl_shutdown()?; +/// +/// signal.recv().await; +/// println!("got CTRL-SHUTDOWN. Cleaning up before exiting"); +/// +/// Ok(()) +/// } +/// ``` +pub fn ctrl_shutdown() -> io::Result<CtrlShutdown> { + Ok(CtrlShutdown { + inner: self::imp::ctrl_shutdown()?, + }) +} - // Windows doesn't have a good programmatic way of sending events - // like sending signals on Unix, so we'll stub out the actual OS - // integration and test that our handling works. - unsafe { - super::handler(CTRL_C_EVENT); - } +/// Represents a listener which receives "ctrl-shutdown" notitifications sent to the process +/// via 'SetConsoleCtrlHandler'. +/// +/// A notification to this process notifies *all* listeners listening for +/// this event. Moreover, the notifications **are coalesced** if they aren't processed +/// quickly enough. This means that if two notifications are received back-to-back, +/// then the listener may only receive one item about the two notifications. +#[must_use = "listeners do nothing unless polled"] +#[derive(Debug)] +pub struct CtrlShutdown { + inner: RxFuture, +} - assert_ready_ok!(ctrl_c.poll()); +impl CtrlShutdown { + /// Receives the next signal notification event. + /// + /// `None` is returned if no more events can be received by this listener. + /// + /// # Examples + /// + /// ```rust,no_run + /// use tokio::signal::windows::ctrl_shutdown; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn std::error::Error>> { + /// // A listener of CTRL-SHUTDOWN events. + /// let mut signal = ctrl_shutdown()?; + /// + /// // Print whenever a CTRL-SHUTDOWN event is received. + /// signal.recv().await; + /// println!("got CTRL-SHUTDOWN. Cleaning up before exiting"); + /// + /// Ok(()) + /// } + /// ``` + pub async fn recv(&mut self) -> Option<()> { + self.inner.recv().await } - #[test] - fn ctrl_break() { - let rt = rt(); + /// Polls to receive the next signal notification event, outside of an + /// `async` context. + /// + /// `None` is returned if no more events can be received by this listener. + /// + /// # Examples + /// + /// Polling from a manually implemented future + /// + /// ```rust,no_run + /// use std::pin::Pin; + /// use std::future::Future; + /// use std::task::{Context, Poll}; + /// use tokio::signal::windows::CtrlShutdown; + /// + /// struct MyFuture { + /// ctrl_shutdown: CtrlShutdown, + /// } + /// + /// impl Future for MyFuture { + /// type Output = Option<()>; + /// + /// fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + /// println!("polling MyFuture"); + /// self.ctrl_shutdown.poll_recv(cx) + /// } + /// } + /// ``` + pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<()>> { + self.inner.poll_recv(cx) + } +} - rt.block_on(async { - let mut ctrl_break = assert_ok!(super::ctrl_break()); +/// Creates a new listener which receives "ctrl-logoff" notifications sent to the +/// process. +/// +/// # Examples +/// +/// ```rust,no_run +/// use tokio::signal::windows::ctrl_logoff; +/// +/// #[tokio::main] +/// async fn main() -> Result<(), Box<dyn std::error::Error>> { +/// // A listener of CTRL-LOGOFF events. +/// let mut signal = ctrl_logoff()?; +/// +/// signal.recv().await; +/// println!("got CTRL-LOGOFF. Cleaning up before exiting"); +/// +/// Ok(()) +/// } +/// ``` +pub fn ctrl_logoff() -> io::Result<CtrlLogoff> { + Ok(CtrlLogoff { + inner: self::imp::ctrl_logoff()?, + }) +} - // Windows doesn't have a good programmatic way of sending events - // like sending signals on Unix, so we'll stub out the actual OS - // integration and test that our handling works. - unsafe { - super::handler(CTRL_BREAK_EVENT); - } +/// Represents a listener which receives "ctrl-logoff" notitifications sent to the process +/// via 'SetConsoleCtrlHandler'. +/// +/// A notification to this process notifies *all* listeners listening for +/// this event. Moreover, the notifications **are coalesced** if they aren't processed +/// quickly enough. This means that if two notifications are received back-to-back, +/// then the listener may only receive one item about the two notifications. +#[must_use = "listeners do nothing unless polled"] +#[derive(Debug)] +pub struct CtrlLogoff { + inner: RxFuture, +} - ctrl_break.recv().await.unwrap(); - }); +impl CtrlLogoff { + /// Receives the next signal notification event. + /// + /// `None` is returned if no more events can be received by this listener. + /// + /// # Examples + /// + /// ```rust,no_run + /// use tokio::signal::windows::ctrl_logoff; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn std::error::Error>> { + /// // An listener of CTRL-LOGOFF events. + /// let mut signal = ctrl_logoff()?; + /// + /// // Print whenever a CTRL-LOGOFF event is received. + /// signal.recv().await; + /// println!("got CTRL-LOGOFF. Cleaning up before exiting"); + /// + /// Ok(()) + /// } + /// ``` + pub async fn recv(&mut self) -> Option<()> { + self.inner.recv().await } - fn rt() -> Runtime { - crate::runtime::Builder::new_current_thread() - .build() - .unwrap() + /// Polls to receive the next signal notification event, outside of an + /// `async` context. + /// + /// `None` is returned if no more events can be received by this listener. + /// + /// # Examples + /// + /// Polling from a manually implemented future + /// + /// ```rust,no_run + /// use std::pin::Pin; + /// use std::future::Future; + /// use std::task::{Context, Poll}; + /// use tokio::signal::windows::CtrlLogoff; + /// + /// struct MyFuture { + /// ctrl_logoff: CtrlLogoff, + /// } + /// + /// impl Future for MyFuture { + /// type Output = Option<()>; + /// + /// fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + /// println!("polling MyFuture"); + /// self.ctrl_logoff.poll_recv(cx) + /// } + /// } + /// ``` + pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<()>> { + self.inner.poll_recv(cx) } } diff --git a/vendor/tokio/src/signal/windows/stub.rs b/vendor/tokio/src/signal/windows/stub.rs new file mode 100644 index 000000000..61df30979 --- /dev/null +++ b/vendor/tokio/src/signal/windows/stub.rs @@ -0,0 +1,25 @@ +//! Stub implementations for the platform API so that rustdoc can build linkable +//! documentation on non-windows platforms. + +use crate::signal::RxFuture; +use std::io; + +pub(super) fn ctrl_break() -> io::Result<RxFuture> { + panic!() +} + +pub(super) fn ctrl_close() -> io::Result<RxFuture> { + panic!() +} + +pub(super) fn ctrl_c() -> io::Result<RxFuture> { + panic!() +} + +pub(super) fn ctrl_logoff() -> io::Result<RxFuture> { + panic!() +} + +pub(super) fn ctrl_shutdown() -> io::Result<RxFuture> { + panic!() +} diff --git a/vendor/tokio/src/signal/windows/sys.rs b/vendor/tokio/src/signal/windows/sys.rs new file mode 100644 index 000000000..26e6bdf81 --- /dev/null +++ b/vendor/tokio/src/signal/windows/sys.rs @@ -0,0 +1,229 @@ +use std::io; +use std::sync::Once; + +use crate::signal::registry::{globals, EventId, EventInfo, Init, Storage}; +use crate::signal::RxFuture; + +use windows_sys::Win32::Foundation::BOOL; +use windows_sys::Win32::System::Console as console; + +pub(super) fn ctrl_break() -> io::Result<RxFuture> { + new(console::CTRL_BREAK_EVENT) +} + +pub(super) fn ctrl_close() -> io::Result<RxFuture> { + new(console::CTRL_CLOSE_EVENT) +} + +pub(super) fn ctrl_c() -> io::Result<RxFuture> { + new(console::CTRL_C_EVENT) +} + +pub(super) fn ctrl_logoff() -> io::Result<RxFuture> { + new(console::CTRL_LOGOFF_EVENT) +} + +pub(super) fn ctrl_shutdown() -> io::Result<RxFuture> { + new(console::CTRL_SHUTDOWN_EVENT) +} + +fn new(signum: u32) -> io::Result<RxFuture> { + global_init()?; + let rx = globals().register_listener(signum as EventId); + Ok(RxFuture::new(rx)) +} + +#[derive(Debug)] +pub(crate) struct OsStorage { + ctrl_break: EventInfo, + ctrl_close: EventInfo, + ctrl_c: EventInfo, + ctrl_logoff: EventInfo, + ctrl_shutdown: EventInfo, +} + +impl Init for OsStorage { + fn init() -> Self { + Self { + ctrl_break: Default::default(), + ctrl_close: Default::default(), + ctrl_c: Default::default(), + ctrl_logoff: Default::default(), + ctrl_shutdown: Default::default(), + } + } +} + +impl Storage for OsStorage { + fn event_info(&self, id: EventId) -> Option<&EventInfo> { + match u32::try_from(id) { + Ok(console::CTRL_BREAK_EVENT) => Some(&self.ctrl_break), + Ok(console::CTRL_CLOSE_EVENT) => Some(&self.ctrl_close), + Ok(console::CTRL_C_EVENT) => Some(&self.ctrl_c), + Ok(console::CTRL_LOGOFF_EVENT) => Some(&self.ctrl_logoff), + Ok(console::CTRL_SHUTDOWN_EVENT) => Some(&self.ctrl_shutdown), + _ => None, + } + } + + fn for_each<'a, F>(&'a self, mut f: F) + where + F: FnMut(&'a EventInfo), + { + f(&self.ctrl_break); + f(&self.ctrl_close); + f(&self.ctrl_c); + f(&self.ctrl_logoff); + f(&self.ctrl_shutdown); + } +} + +#[derive(Debug)] +pub(crate) struct OsExtraData {} + +impl Init for OsExtraData { + fn init() -> Self { + Self {} + } +} + +fn global_init() -> io::Result<()> { + static INIT: Once = Once::new(); + + let mut init = None; + + INIT.call_once(|| unsafe { + let rc = console::SetConsoleCtrlHandler(Some(handler), 1); + let ret = if rc == 0 { + Err(io::Error::last_os_error()) + } else { + Ok(()) + }; + + init = Some(ret); + }); + + init.unwrap_or_else(|| Ok(())) +} + +unsafe extern "system" fn handler(ty: u32) -> BOOL { + let globals = globals(); + globals.record_event(ty as EventId); + + // According to https://docs.microsoft.com/en-us/windows/console/handlerroutine + // the handler routine is always invoked in a new thread, thus we don't + // have the same restrictions as in Unix signal handlers, meaning we can + // go ahead and perform the broadcast here. + if globals.broadcast() { + 1 + } else { + // No one is listening for this notification any more + // let the OS fire the next (possibly the default) handler. + 0 + } +} + +#[cfg(all(test, not(loom)))] +mod tests { + use super::*; + use crate::runtime::Runtime; + + use tokio_test::{assert_ok, assert_pending, assert_ready_ok, task}; + + #[test] + fn ctrl_c() { + let rt = rt(); + let _enter = rt.enter(); + + let mut ctrl_c = task::spawn(crate::signal::ctrl_c()); + + assert_pending!(ctrl_c.poll()); + + // Windows doesn't have a good programmatic way of sending events + // like sending signals on Unix, so we'll stub out the actual OS + // integration and test that our handling works. + unsafe { + super::handler(console::CTRL_C_EVENT); + } + + assert_ready_ok!(ctrl_c.poll()); + } + + #[test] + fn ctrl_break() { + let rt = rt(); + + rt.block_on(async { + let mut ctrl_break = assert_ok!(crate::signal::windows::ctrl_break()); + + // Windows doesn't have a good programmatic way of sending events + // like sending signals on Unix, so we'll stub out the actual OS + // integration and test that our handling works. + unsafe { + super::handler(console::CTRL_BREAK_EVENT); + } + + ctrl_break.recv().await.unwrap(); + }); + } + + #[test] + fn ctrl_close() { + let rt = rt(); + + rt.block_on(async { + let mut ctrl_close = assert_ok!(crate::signal::windows::ctrl_close()); + + // Windows doesn't have a good programmatic way of sending events + // like sending signals on Unix, so we'll stub out the actual OS + // integration and test that our handling works. + unsafe { + super::handler(console::CTRL_CLOSE_EVENT); + } + + ctrl_close.recv().await.unwrap(); + }); + } + + #[test] + fn ctrl_shutdown() { + let rt = rt(); + + rt.block_on(async { + let mut ctrl_shutdown = assert_ok!(crate::signal::windows::ctrl_shutdown()); + + // Windows doesn't have a good programmatic way of sending events + // like sending signals on Unix, so we'll stub out the actual OS + // integration and test that our handling works. + unsafe { + super::handler(console::CTRL_SHUTDOWN_EVENT); + } + + ctrl_shutdown.recv().await.unwrap(); + }); + } + + #[test] + fn ctrl_logoff() { + let rt = rt(); + + rt.block_on(async { + let mut ctrl_logoff = assert_ok!(crate::signal::windows::ctrl_logoff()); + + // Windows doesn't have a good programmatic way of sending events + // like sending signals on Unix, so we'll stub out the actual OS + // integration and test that our handling works. + unsafe { + super::handler(console::CTRL_LOGOFF_EVENT); + } + + ctrl_logoff.recv().await.unwrap(); + }); + } + + fn rt() -> Runtime { + crate::runtime::Builder::new_current_thread() + .build() + .unwrap() + } +} diff --git a/vendor/tokio/src/sync/barrier.rs b/vendor/tokio/src/sync/barrier.rs index 0e39dac8b..29b6d4e48 100644 --- a/vendor/tokio/src/sync/barrier.rs +++ b/vendor/tokio/src/sync/barrier.rs @@ -1,5 +1,7 @@ use crate::loom::sync::Mutex; use crate::sync::watch; +#[cfg(all(tokio_unstable, feature = "tracing"))] +use crate::util::trace; /// A barrier enables multiple tasks to synchronize the beginning of some computation. /// @@ -41,6 +43,8 @@ pub struct Barrier { state: Mutex<BarrierState>, wait: watch::Receiver<usize>, n: usize, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, } #[derive(Debug)] @@ -55,6 +59,7 @@ impl Barrier { /// /// A barrier will block `n`-1 tasks which call [`Barrier::wait`] and then wake up all /// tasks at once when the `n`th task calls `wait`. + #[track_caller] pub fn new(mut n: usize) -> Barrier { let (waker, wait) = crate::sync::watch::channel(0); @@ -65,6 +70,32 @@ impl Barrier { n = 1; } + #[cfg(all(tokio_unstable, feature = "tracing"))] + let resource_span = { + let location = std::panic::Location::caller(); + let resource_span = tracing::trace_span!( + "runtime.resource", + concrete_type = "Barrier", + kind = "Sync", + loc.file = location.file(), + loc.line = location.line(), + loc.col = location.column(), + ); + + resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + size = n, + ); + + tracing::trace!( + target: "runtime::resource::state_update", + arrived = 0, + ) + }); + resource_span + }; + Barrier { state: Mutex::new(BarrierState { waker, @@ -73,6 +104,8 @@ impl Barrier { }), n, wait, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span, } } @@ -85,10 +118,26 @@ impl Barrier { /// [`BarrierWaitResult::is_leader`] when returning from this function, and all other tasks /// will receive a result that will return `false` from `is_leader`. pub async fn wait(&self) -> BarrierWaitResult { + #[cfg(all(tokio_unstable, feature = "tracing"))] + return trace::async_op( + || self.wait_internal(), + self.resource_span.clone(), + "Barrier::wait", + "poll", + false, + ) + .await; + + #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] + return self.wait_internal().await; + } + async fn wait_internal(&self) -> BarrierWaitResult { + crate::trace::async_trace_leaf().await; + // NOTE: we are taking a _synchronous_ lock here. // It is okay to do so because the critical section is fast and never yields, so it cannot // deadlock even if another future is concurrently holding the lock. - // It is _desireable_ to do so as synchronous Mutexes are, at least in theory, faster than + // It is _desirable_ to do so as synchronous Mutexes are, at least in theory, faster than // the asynchronous counter-parts, so we should use them where possible [citation needed]. // NOTE: the extra scope here is so that the compiler doesn't think `state` is held across // a yield point, and thus marks the returned future as !Send. @@ -96,7 +145,23 @@ impl Barrier { let mut state = self.state.lock(); let generation = state.generation; state.arrived += 1; + #[cfg(all(tokio_unstable, feature = "tracing"))] + tracing::trace!( + target: "runtime::resource::state_update", + arrived = 1, + arrived.op = "add", + ); + #[cfg(all(tokio_unstable, feature = "tracing"))] + tracing::trace!( + target: "runtime::resource::async_op::state_update", + arrived = true, + ); if state.arrived == self.n { + #[cfg(all(tokio_unstable, feature = "tracing"))] + tracing::trace!( + target: "runtime::resource::async_op::state_update", + is_leader = true, + ); // we are the leader for this generation // wake everyone, increment the generation, and return state diff --git a/vendor/tokio/src/sync/batch_semaphore.rs b/vendor/tokio/src/sync/batch_semaphore.rs index a0bf5ef94..a762f799d 100644 --- a/vendor/tokio/src/sync/batch_semaphore.rs +++ b/vendor/tokio/src/sync/batch_semaphore.rs @@ -1,5 +1,5 @@ #![cfg_attr(not(feature = "sync"), allow(unreachable_pub, dead_code))] -//! # Implementation Details +//! # Implementation Details. //! //! The semaphore is implemented using an intrusive linked list of waiters. An //! atomic counter tracks the number of available permits. If the semaphore does @@ -19,6 +19,9 @@ use crate::loom::cell::UnsafeCell; use crate::loom::sync::atomic::AtomicUsize; use crate::loom::sync::{Mutex, MutexGuard}; use crate::util::linked_list::{self, LinkedList}; +#[cfg(all(tokio_unstable, feature = "tracing"))] +use crate::util::trace; +use crate::util::WakeList; use std::future::Future; use std::marker::PhantomPinned; @@ -34,6 +37,8 @@ pub(crate) struct Semaphore { waiters: Mutex<Waitlist>, /// The current number of available permits in the semaphore. permits: AtomicUsize, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, } struct Waitlist { @@ -44,7 +49,7 @@ struct Waitlist { /// Error returned from the [`Semaphore::try_acquire`] function. /// /// [`Semaphore::try_acquire`]: crate::sync::Semaphore::try_acquire -#[derive(Debug, PartialEq)] +#[derive(Debug, PartialEq, Eq)] pub enum TryAcquireError { /// The semaphore has been [closed] and cannot issue new permits. /// @@ -100,10 +105,21 @@ struct Waiter { /// use `UnsafeCell` internally. pointers: linked_list::Pointers<Waiter>, + #[cfg(all(tokio_unstable, feature = "tracing"))] + ctx: trace::AsyncOpTracingCtx, + /// Should not be `Unpin`. _p: PhantomPinned, } +generate_addr_of_methods! { + impl<> Waiter { + unsafe fn addr_of_pointers(self: NonNull<Self>) -> NonNull<linked_list::Pointers<Waiter>> { + &self.pointers + } + } +} + impl Semaphore { /// The maximum number of permits which a semaphore can hold. /// @@ -128,16 +144,38 @@ impl Semaphore { "a semaphore may not have more than MAX_PERMITS permits ({})", Self::MAX_PERMITS ); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + let resource_span = { + let resource_span = tracing::trace_span!( + "runtime.resource", + concrete_type = "Semaphore", + kind = "Sync", + is_internal = true + ); + + resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + permits = permits, + permits.op = "override", + ) + }); + resource_span + }; + Self { permits: AtomicUsize::new(permits << Self::PERMIT_SHIFT), waiters: Mutex::new(Waitlist { queue: LinkedList::new(), closed: false, }), + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span, } } - /// Creates a new semaphore with the initial number of permits + /// Creates a new semaphore with the initial number of permits. /// /// Maximum number of permits on 32-bit platforms is `1<<29`. /// @@ -155,10 +193,12 @@ impl Semaphore { queue: LinkedList::new(), closed: false, }), + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span::none(), } } - /// Returns the current number of available permits + /// Returns the current number of available permits. pub(crate) fn available_permits(&self) -> usize { self.permits.load(Acquire) >> Self::PERMIT_SHIFT } @@ -196,7 +236,7 @@ impl Semaphore { } } - /// Returns true if the semaphore is closed + /// Returns true if the semaphore is closed. pub(crate) fn is_closed(&self) -> bool { self.permits.load(Acquire) & Self::CLOSED == Self::CLOSED } @@ -223,7 +263,10 @@ impl Semaphore { let next = curr - num_permits; match self.permits.compare_exchange(curr, next, AcqRel, Acquire) { - Ok(_) => return Ok(()), + Ok(_) => { + // TODO: Instrument once issue has been solved + return Ok(()); + } Err(actual) => curr = actual, } } @@ -239,12 +282,12 @@ impl Semaphore { /// If `rem` exceeds the number of permits needed by the wait list, the /// remainder are assigned back to the semaphore. fn add_permits_locked(&self, mut rem: usize, waiters: MutexGuard<'_, Waitlist>) { - let mut wakers: [Option<Waker>; 8] = Default::default(); + let mut wakers = WakeList::new(); let mut lock = Some(waiters); let mut is_empty = false; while rem > 0 { let mut waiters = lock.take().unwrap_or_else(|| self.waiters.lock()); - 'inner: for slot in &mut wakers[..] { + 'inner: while wakers.can_push() { // Was the waiter assigned enough permits to wake it? match waiters.queue.last() { Some(waiter) => { @@ -260,7 +303,11 @@ impl Semaphore { } }; let mut waiter = waiters.queue.pop_back().unwrap(); - *slot = unsafe { waiter.as_mut().waker.with_mut(|waker| (*waker).take()) }; + if let Some(waker) = + unsafe { waiter.as_mut().waker.with_mut(|waker| (*waker).take()) } + { + wakers.push(waker); + } } if rem > 0 && is_empty { @@ -278,15 +325,23 @@ impl Semaphore { rem, Self::MAX_PERMITS ); + + // add remaining permits back + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + permits = rem, + permits.op = "add", + ) + }); + rem = 0; } drop(waiters); // release the lock - wakers - .iter_mut() - .filter_map(Option::take) - .for_each(Waker::wake); + wakers.wake_all(); } assert_eq!(rem, 0); @@ -345,6 +400,20 @@ impl Semaphore { acquired += acq; if remaining == 0 { if !queued { + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + permits = acquired, + permits.op = "sub", + ); + tracing::trace!( + target: "runtime::resource::async_op::state_update", + permits_obtained = acquired, + permits.op = "add", + ) + }); + return Ready(Ok(())); } else if lock.is_none() { break self.waiters.lock(); @@ -360,12 +429,22 @@ impl Semaphore { return Ready(Err(AcquireError::closed())); } + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + permits = acquired, + permits.op = "sub", + ) + }); + if node.assign_permits(&mut acquired) { self.add_permits_locked(acquired, waiters); return Ready(Ok(())); } assert_eq!(acquired, 0); + let mut old_waker = None; // Otherwise, register the waker & enqueue the node. node.waker.with_mut(|waker| { @@ -377,7 +456,7 @@ impl Semaphore { .map(|waker| !waker.will_wake(cx.waker())) .unwrap_or(true) { - *waker = Some(cx.waker().clone()); + old_waker = std::mem::replace(waker, Some(cx.waker().clone())); } }); @@ -390,6 +469,8 @@ impl Semaphore { waiters.queue.push_front(node); } + drop(waiters); + drop(old_waker); Pending } @@ -404,11 +485,16 @@ impl fmt::Debug for Semaphore { } impl Waiter { - fn new(num_permits: u32) -> Self { + fn new( + num_permits: u32, + #[cfg(all(tokio_unstable, feature = "tracing"))] ctx: trace::AsyncOpTracingCtx, + ) -> Self { Waiter { waker: UnsafeCell::new(None), state: AtomicUsize::new(num_permits as usize), pointers: linked_list::Pointers::new(), + #[cfg(all(tokio_unstable, feature = "tracing"))] + ctx, _p: PhantomPinned, } } @@ -424,6 +510,14 @@ impl Waiter { match self.state.compare_exchange(curr, next, AcqRel, Acquire) { Ok(_) => { *n -= assign; + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.ctx.async_op_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::async_op::state_update", + permits_obtained = assign, + permits.op = "add", + ); + }); return next == 0; } Err(actual) => curr = actual, @@ -436,12 +530,26 @@ impl Future for Acquire<'_> { type Output = Result<(), AcquireError>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { - // First, ensure the current task has enough budget to proceed. - let coop = ready!(crate::coop::poll_proceed(cx)); + #[cfg(all(tokio_unstable, feature = "tracing"))] + let _resource_span = self.node.ctx.resource_span.clone().entered(); + #[cfg(all(tokio_unstable, feature = "tracing"))] + let _async_op_span = self.node.ctx.async_op_span.clone().entered(); + #[cfg(all(tokio_unstable, feature = "tracing"))] + let _async_op_poll_span = self.node.ctx.async_op_poll_span.clone().entered(); let (node, semaphore, needed, queued) = self.project(); - match semaphore.poll_acquire(cx, needed, node, *queued) { + // First, ensure the current task has enough budget to proceed. + #[cfg(all(tokio_unstable, feature = "tracing"))] + let coop = ready!(trace_poll_op!( + "poll_acquire", + crate::runtime::coop::poll_proceed(cx), + )); + + #[cfg(not(all(tokio_unstable, feature = "tracing")))] + let coop = ready!(crate::runtime::coop::poll_proceed(cx)); + + let result = match semaphore.poll_acquire(cx, needed, node, *queued) { Pending => { *queued = true; Pending @@ -452,18 +560,59 @@ impl Future for Acquire<'_> { *queued = false; Ready(Ok(())) } - } + }; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + return trace_poll_op!("poll_acquire", result); + + #[cfg(not(all(tokio_unstable, feature = "tracing")))] + return result; } } impl<'a> Acquire<'a> { fn new(semaphore: &'a Semaphore, num_permits: u32) -> Self { - Self { + #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] + return Self { node: Waiter::new(num_permits), semaphore, num_permits, queued: false, - } + }; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + return semaphore.resource_span.in_scope(|| { + let async_op_span = + tracing::trace_span!("runtime.resource.async_op", source = "Acquire::new"); + let async_op_poll_span = async_op_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::async_op::state_update", + permits_requested = num_permits, + permits.op = "override", + ); + + tracing::trace!( + target: "runtime::resource::async_op::state_update", + permits_obtained = 0usize, + permits.op = "override", + ); + + tracing::trace_span!("runtime.resource.async_op.poll") + }); + + let ctx = trace::AsyncOpTracingCtx { + async_op_span, + async_op_poll_span, + resource_span: semaphore.resource_span.clone(), + }; + + Self { + node: Waiter::new(num_permits, ctx), + semaphore, + num_permits, + queued: false, + } + }); } fn project(self: Pin<&mut Self>) -> (Pin<&mut Waiter>, &Semaphore, u32, &mut bool) { @@ -478,7 +627,7 @@ impl<'a> Acquire<'a> { let this = self.get_unchecked_mut(); ( Pin::new_unchecked(&mut this.node), - &this.semaphore, + this.semaphore, this.num_permits, &mut this.queued, ) @@ -566,12 +715,6 @@ impl std::error::Error for TryAcquireError {} /// /// `Waiter` is forced to be !Unpin. unsafe impl linked_list::Link for Waiter { - // XXX: ideally, we would be able to use `Pin` here, to enforce the - // invariant that list entries may not move while in the list. However, we - // can't do this currently, as using `Pin<&'a mut Waiter>` as the `Handle` - // type would require `Semaphore` to be generic over a lifetime. We can't - // use `Pin<*mut Waiter>`, as raw pointers are `Unpin` regardless of whether - // or not they dereference to an `!Unpin` target. type Handle = NonNull<Waiter>; type Target = Waiter; @@ -583,7 +726,7 @@ unsafe impl linked_list::Link for Waiter { ptr } - unsafe fn pointers(mut target: NonNull<Waiter>) -> NonNull<linked_list::Pointers<Waiter>> { - NonNull::from(&mut target.as_mut().pointers) + unsafe fn pointers(target: NonNull<Waiter>) -> NonNull<linked_list::Pointers<Waiter>> { + Waiter::addr_of_pointers(target) } } diff --git a/vendor/tokio/src/sync/broadcast.rs b/vendor/tokio/src/sync/broadcast.rs index a2ca4459e..4b36452ce 100644 --- a/vendor/tokio/src/sync/broadcast.rs +++ b/vendor/tokio/src/sync/broadcast.rs @@ -4,7 +4,7 @@ //! A [`Sender`] is used to broadcast values to **all** connected [`Receiver`] //! values. [`Sender`] handles are clone-able, allowing concurrent send and //! receive actions. [`Sender`] and [`Receiver`] are both `Send` and `Sync` as -//! long as `T` is also `Send` or `Sync` respectively. +//! long as `T` is `Send`. //! //! When a value is sent, **all** [`Receiver`] handles are notified and will //! receive the value. The value is stored once inside the channel and cloned on @@ -18,6 +18,9 @@ //! returned [`Receiver`] will receive values sent **after** the call to //! `subscribe`. //! +//! This channel is also suitable for the single-producer multi-consumer +//! use-case, where a single sender broadcasts values to many receivers. +//! //! ## Lagging //! //! As sent messages must be retained until **all** [`Receiver`] handles receive @@ -51,6 +54,10 @@ //! all values retained by the channel, the next call to [`recv`] will return //! with [`RecvError::Closed`]. //! +//! When a [`Receiver`] handle is dropped, any messages not read by the receiver +//! will be marked as read. If this receiver was the only one not to have read +//! that message, the message will be dropped at this point. +//! //! [`Sender`]: crate::sync::broadcast::Sender //! [`Sender::subscribe`]: crate::sync::broadcast::Sender::subscribe //! [`Receiver`]: crate::sync::broadcast::Receiver @@ -111,8 +118,9 @@ use crate::loom::cell::UnsafeCell; use crate::loom::sync::atomic::AtomicUsize; -use crate::loom::sync::{Arc, Mutex, RwLock, RwLockReadGuard}; +use crate::loom::sync::{Arc, Mutex, MutexGuard, RwLock, RwLockReadGuard}; use crate::util::linked_list::{self, LinkedList}; +use crate::util::WakeList; use std::fmt; use std::future::Future; @@ -230,7 +238,7 @@ pub mod error { /// /// [`recv`]: crate::sync::broadcast::Receiver::recv /// [`Receiver`]: crate::sync::broadcast::Receiver - #[derive(Debug, PartialEq)] + #[derive(Debug, PartialEq, Eq, Clone)] pub enum RecvError { /// There are no more active senders implying no further messages will ever /// be sent. @@ -258,7 +266,7 @@ pub mod error { /// /// [`try_recv`]: crate::sync::broadcast::Receiver::try_recv /// [`Receiver`]: crate::sync::broadcast::Receiver - #[derive(Debug, PartialEq)] + #[derive(Debug, PartialEq, Eq, Clone)] pub enum TryRecvError { /// The channel is currently empty. There are still active /// [`Sender`] handles, so data may yet become available. @@ -293,37 +301,37 @@ pub mod error { use self::error::*; -/// Data shared between senders and receivers +/// Data shared between senders and receivers. struct Shared<T> { - /// slots in the channel + /// slots in the channel. buffer: Box<[RwLock<Slot<T>>]>, - /// Mask a position -> index + /// Mask a position -> index. mask: usize, /// Tail of the queue. Includes the rx wait list. tail: Mutex<Tail>, - /// Number of outstanding Sender handles + /// Number of outstanding Sender handles. num_tx: AtomicUsize, } -/// Next position to write a value +/// Next position to write a value. struct Tail { - /// Next position to write to + /// Next position to write to. pos: u64, - /// Number of active receivers + /// Number of active receivers. rx_cnt: usize, - /// True if the channel is closed + /// True if the channel is closed. closed: bool, - /// Receivers waiting for a value + /// Receivers waiting for a value. waiters: LinkedList<Waiter, <Waiter as linked_list::Link>::Target>, } -/// Slot in the buffer +/// Slot in the buffer. struct Slot<T> { /// Remaining number of receivers that are expected to see this value. /// @@ -333,12 +341,9 @@ struct Slot<T> { /// acquired. rem: AtomicUsize, - /// Uniquely identifies the `send` stored in the slot + /// Uniquely identifies the `send` stored in the slot. pos: u64, - /// True signals the channel is closed. - closed: bool, - /// The value being broadcast. /// /// The value is set by `send` when the write lock is held. When a reader @@ -346,9 +351,9 @@ struct Slot<T> { val: UnsafeCell<Option<T>>, } -/// An entry in the wait queue +/// An entry in the wait queue. struct Waiter { - /// True if queued + /// True if queued. queued: bool, /// Task waiting on the broadcast channel. @@ -361,16 +366,24 @@ struct Waiter { _p: PhantomPinned, } +generate_addr_of_methods! { + impl<> Waiter { + unsafe fn addr_of_pointers(self: NonNull<Self>) -> NonNull<linked_list::Pointers<Waiter>> { + &self.pointers + } + } +} + struct RecvGuard<'a, T> { slot: RwLockReadGuard<'a, Slot<T>>, } -/// Receive a value future +/// Receive a value future. struct Recv<'a, T> { - /// Receiver being waited on + /// Receiver being waited on. receiver: &'a mut Receiver<T>, - /// Entry in the waiter `LinkedList` + /// Entry in the waiter `LinkedList`. waiter: UnsafeCell<Waiter>, } @@ -425,6 +438,12 @@ const MAX_RECEIVERS: usize = usize::MAX >> 2; /// tx.send(20).unwrap(); /// } /// ``` +/// +/// # Panics +/// +/// This will panic if `capacity` is equal to `0` or larger +/// than `usize::MAX / 2`. +#[track_caller] pub fn channel<T: Clone>(mut capacity: usize) -> (Sender<T>, Receiver<T>) { assert!(capacity > 0, "capacity is empty"); assert!(capacity <= usize::MAX >> 1, "requested capacity too large"); @@ -438,7 +457,6 @@ pub fn channel<T: Clone>(mut capacity: usize) -> (Sender<T>, Receiver<T>) { buffer.push(RwLock::new(Slot { rem: AtomicUsize::new(0), pos: (i as u64).wrapping_sub(capacity as u64), - closed: false, val: UnsafeCell::new(None), })); } @@ -523,8 +541,41 @@ impl<T> Sender<T> { /// } /// ``` pub fn send(&self, value: T) -> Result<usize, SendError<T>> { - self.send2(Some(value)) - .map_err(|SendError(maybe_v)| SendError(maybe_v.unwrap())) + let mut tail = self.shared.tail.lock(); + + if tail.rx_cnt == 0 { + return Err(SendError(value)); + } + + // Position to write into + let pos = tail.pos; + let rem = tail.rx_cnt; + let idx = (pos & self.shared.mask as u64) as usize; + + // Update the tail position + tail.pos = tail.pos.wrapping_add(1); + + // Get the slot + let mut slot = self.shared.buffer[idx].write().unwrap(); + + // Track the position + slot.pos = pos; + + // Set remaining receivers + slot.rem.with_mut(|v| *v = rem); + + // Write the value + slot.val = UnsafeCell::new(Some(value)); + + // Release the slot lock before notifying the receivers. + drop(slot); + + // Notify and release the mutex. This must happen after the slot lock is + // released, otherwise the writer lock bit could be cleared while another + // thread is in the critical section. + self.shared.notify_rx(tail); + + Ok(rem) } /// Creates a new [`Receiver`] handle that will receive values sent **after** @@ -555,6 +606,97 @@ impl<T> Sender<T> { new_receiver(shared) } + /// Returns the number of queued values. + /// + /// A value is queued until it has either been seen by all receivers that were alive at the time + /// it was sent, or has been evicted from the queue by subsequent sends that exceeded the + /// queue's capacity. + /// + /// # Note + /// + /// In contrast to [`Receiver::len`], this method only reports queued values and not values that + /// have been evicted from the queue before being seen by all receivers. + /// + /// # Examples + /// + /// ``` + /// use tokio::sync::broadcast; + /// + /// #[tokio::main] + /// async fn main() { + /// let (tx, mut rx1) = broadcast::channel(16); + /// let mut rx2 = tx.subscribe(); + /// + /// tx.send(10).unwrap(); + /// tx.send(20).unwrap(); + /// tx.send(30).unwrap(); + /// + /// assert_eq!(tx.len(), 3); + /// + /// rx1.recv().await.unwrap(); + /// + /// // The len is still 3 since rx2 hasn't seen the first value yet. + /// assert_eq!(tx.len(), 3); + /// + /// rx2.recv().await.unwrap(); + /// + /// assert_eq!(tx.len(), 2); + /// } + /// ``` + pub fn len(&self) -> usize { + let tail = self.shared.tail.lock(); + + let base_idx = (tail.pos & self.shared.mask as u64) as usize; + let mut low = 0; + let mut high = self.shared.buffer.len(); + while low < high { + let mid = low + (high - low) / 2; + let idx = base_idx.wrapping_add(mid) & self.shared.mask; + if self.shared.buffer[idx].read().unwrap().rem.load(SeqCst) == 0 { + low = mid + 1; + } else { + high = mid; + } + } + + self.shared.buffer.len() - low + } + + /// Returns true if there are no queued values. + /// + /// # Examples + /// + /// ``` + /// use tokio::sync::broadcast; + /// + /// #[tokio::main] + /// async fn main() { + /// let (tx, mut rx1) = broadcast::channel(16); + /// let mut rx2 = tx.subscribe(); + /// + /// assert!(tx.is_empty()); + /// + /// tx.send(10).unwrap(); + /// + /// assert!(!tx.is_empty()); + /// + /// rx1.recv().await.unwrap(); + /// + /// // The queue is still not empty since rx2 hasn't seen the value. + /// assert!(!tx.is_empty()); + /// + /// rx2.recv().await.unwrap(); + /// + /// assert!(tx.is_empty()); + /// } + /// ``` + pub fn is_empty(&self) -> bool { + let tail = self.shared.tail.lock(); + + let idx = (tail.pos.wrapping_sub(1) & self.shared.mask as u64) as usize; + self.shared.buffer[idx].read().unwrap().rem.load(SeqCst) == 0 + } + /// Returns the number of active receivers /// /// An active receiver is a [`Receiver`] handle returned from [`channel`] or @@ -596,52 +738,38 @@ impl<T> Sender<T> { tail.rx_cnt } - fn send2(&self, value: Option<T>) -> Result<usize, SendError<Option<T>>> { - let mut tail = self.shared.tail.lock(); - - if tail.rx_cnt == 0 { - return Err(SendError(value)); - } - - // Position to write into - let pos = tail.pos; - let rem = tail.rx_cnt; - let idx = (pos & self.shared.mask as u64) as usize; - - // Update the tail position - tail.pos = tail.pos.wrapping_add(1); - - // Get the slot - let mut slot = self.shared.buffer[idx].write().unwrap(); - - // Track the position - slot.pos = pos; - - // Set remaining receivers - slot.rem.with_mut(|v| *v = rem); - - // Set the closed bit if the value is `None`; otherwise write the value - if value.is_none() { - tail.closed = true; - slot.closed = true; - } else { - slot.val.with_mut(|ptr| unsafe { *ptr = value }); - } - - // Release the slot lock before notifying the receivers. - drop(slot); - - tail.notify_rx(); + /// Returns `true` if senders belong to the same channel. + /// + /// # Examples + /// + /// ``` + /// use tokio::sync::broadcast; + /// + /// #[tokio::main] + /// async fn main() { + /// let (tx, _rx) = broadcast::channel::<()>(16); + /// let tx2 = tx.clone(); + /// + /// assert!(tx.same_channel(&tx2)); + /// + /// let (tx3, _rx3) = broadcast::channel::<()>(16); + /// + /// assert!(!tx3.same_channel(&tx2)); + /// } + /// ``` + pub fn same_channel(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.shared, &other.shared) + } - // Release the mutex. This must happen after the slot lock is released, - // otherwise the writer lock bit could be cleared while another thread - // is in the critical section. - drop(tail); + fn close_channel(&self) { + let mut tail = self.shared.tail.lock(); + tail.closed = true; - Ok(rem) + self.shared.notify_rx(tail); } } +/// Create a new `Receiver` which reads starting from the tail. fn new_receiver<T>(shared: Arc<Shared<T>>) -> Receiver<T> { let mut tail = shared.tail.lock(); @@ -658,18 +786,47 @@ fn new_receiver<T>(shared: Arc<Shared<T>>) -> Receiver<T> { Receiver { shared, next } } -impl Tail { - fn notify_rx(&mut self) { - while let Some(mut waiter) = self.waiters.pop_back() { - // Safety: `waiters` lock is still held. - let waiter = unsafe { waiter.as_mut() }; +impl<T> Shared<T> { + fn notify_rx<'a, 'b: 'a>(&'b self, mut tail: MutexGuard<'a, Tail>) { + let mut wakers = WakeList::new(); + 'outer: loop { + while wakers.can_push() { + match tail.waiters.pop_back() { + Some(mut waiter) => { + // Safety: `tail` lock is still held. + let waiter = unsafe { waiter.as_mut() }; + + assert!(waiter.queued); + waiter.queued = false; + + if let Some(waker) = waiter.waker.take() { + wakers.push(waker); + } + } + None => { + break 'outer; + } + } + } + + // Release the lock before waking. + drop(tail); - assert!(waiter.queued); - waiter.queued = false; + // Before we acquire the lock again all sorts of things can happen: + // some waiters may remove themselves from the list and new waiters + // may be added. This is fine since at worst we will unnecessarily + // wake up waiters which will then queue themselves again. - let waker = waiter.waker.take().unwrap(); - waker.wake(); + wakers.wake_all(); + + // Acquire the lock again. + tail = self.tail.lock(); } + + // Release the lock before waking. + drop(tail); + + wakers.wake_all(); } } @@ -685,12 +842,102 @@ impl<T> Clone for Sender<T> { impl<T> Drop for Sender<T> { fn drop(&mut self) { if 1 == self.shared.num_tx.fetch_sub(1, SeqCst) { - let _ = self.send2(None); + self.close_channel(); } } } impl<T> Receiver<T> { + /// Returns the number of messages that were sent into the channel and that + /// this [`Receiver`] has yet to receive. + /// + /// If the returned value from `len` is larger than the next largest power of 2 + /// of the capacity of the channel any call to [`recv`] will return an + /// `Err(RecvError::Lagged)` and any call to [`try_recv`] will return an + /// `Err(TryRecvError::Lagged)`, e.g. if the capacity of the channel is 10, + /// [`recv`] will start to return `Err(RecvError::Lagged)` once `len` returns + /// values larger than 16. + /// + /// [`Receiver`]: crate::sync::broadcast::Receiver + /// [`recv`]: crate::sync::broadcast::Receiver::recv + /// [`try_recv`]: crate::sync::broadcast::Receiver::try_recv + /// + /// # Examples + /// + /// ``` + /// use tokio::sync::broadcast; + /// + /// #[tokio::main] + /// async fn main() { + /// let (tx, mut rx1) = broadcast::channel(16); + /// + /// tx.send(10).unwrap(); + /// tx.send(20).unwrap(); + /// + /// assert_eq!(rx1.len(), 2); + /// assert_eq!(rx1.recv().await.unwrap(), 10); + /// assert_eq!(rx1.len(), 1); + /// assert_eq!(rx1.recv().await.unwrap(), 20); + /// assert_eq!(rx1.len(), 0); + /// } + /// ``` + pub fn len(&self) -> usize { + let next_send_pos = self.shared.tail.lock().pos; + (next_send_pos - self.next) as usize + } + + /// Returns true if there aren't any messages in the channel that the [`Receiver`] + /// has yet to receive. + /// + /// [`Receiver]: create::sync::broadcast::Receiver + /// + /// # Examples + /// + /// ``` + /// use tokio::sync::broadcast; + /// + /// #[tokio::main] + /// async fn main() { + /// let (tx, mut rx1) = broadcast::channel(16); + /// + /// assert!(rx1.is_empty()); + /// + /// tx.send(10).unwrap(); + /// tx.send(20).unwrap(); + /// + /// assert!(!rx1.is_empty()); + /// assert_eq!(rx1.recv().await.unwrap(), 10); + /// assert_eq!(rx1.recv().await.unwrap(), 20); + /// assert!(rx1.is_empty()); + /// } + /// ``` + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Returns `true` if receivers belong to the same channel. + /// + /// # Examples + /// + /// ``` + /// use tokio::sync::broadcast; + /// + /// #[tokio::main] + /// async fn main() { + /// let (tx, rx) = broadcast::channel::<()>(16); + /// let rx2 = tx.subscribe(); + /// + /// assert!(rx.same_channel(&rx2)); + /// + /// let (_tx3, rx3) = broadcast::channel::<()>(16); + /// + /// assert!(!rx3.same_channel(&rx2)); + /// } + /// ``` + pub fn same_channel(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.shared, &other.shared) + } + /// Locks the next value if there is one. fn recv_ref( &mut self, @@ -702,14 +949,6 @@ impl<T> Receiver<T> { let mut slot = self.shared.buffer[idx].read().unwrap(); if slot.pos != self.next { - let next_pos = slot.pos.wrapping_add(self.shared.buffer.len() as u64); - - // The receiver has read all current values in the channel and there - // is no waiter to register - if waiter.is_none() && next_pos == self.next { - return Err(TryRecvError::Empty); - } - // Release the `slot` lock before attempting to acquire the `tail` // lock. This is required because `send2` acquires the tail lock // first followed by the slot lock. Acquiring the locks in reverse @@ -719,6 +958,8 @@ impl<T> Receiver<T> { // the slot lock. drop(slot); + let mut old_waker = None; + let mut tail = self.shared.tail.lock(); // Acquire slot lock again @@ -731,6 +972,13 @@ impl<T> Receiver<T> { let next_pos = slot.pos.wrapping_add(self.shared.buffer.len() as u64); if next_pos == self.next { + // At this point the channel is empty for *this* receiver. If + // it's been closed, then that's what we return, otherwise we + // set a waker and return empty. + if tail.closed { + return Err(TryRecvError::Closed); + } + // Store the waker if let Some((waiter, waker)) = waiter { // Safety: called while locked. @@ -744,7 +992,10 @@ impl<T> Receiver<T> { match (*ptr).waker { Some(ref w) if w.will_wake(waker) => {} _ => { - (*ptr).waker = Some(waker.clone()); + old_waker = std::mem::replace( + &mut (*ptr).waker, + Some(waker.clone()), + ); } } @@ -756,6 +1007,11 @@ impl<T> Receiver<T> { } } + // Drop the old waker after releasing the locks. + drop(slot); + drop(tail); + drop(old_waker); + return Err(TryRecvError::Empty); } @@ -764,22 +1020,7 @@ impl<T> Receiver<T> { // catch up by skipping dropped messages and setting the // internal cursor to the **oldest** message stored by the // channel. - // - // However, finding the oldest position is a bit more - // complicated than `tail-position - buffer-size`. When - // the channel is closed, the tail position is incremented to - // signal a new `None` message, but `None` is not stored in the - // channel itself (see issue #2425 for why). - // - // To account for this, if the channel is closed, the tail - // position is decremented by `buffer-size + 1`. - let mut adjust = 0; - if tail.closed { - adjust = 1 - } - let next = tail - .pos - .wrapping_sub(self.shared.buffer.len() as u64 + adjust); + let next = tail.pos.wrapping_sub(self.shared.buffer.len() as u64); let missed = next.wrapping_sub(self.next); @@ -800,15 +1041,38 @@ impl<T> Receiver<T> { self.next = self.next.wrapping_add(1); - if slot.closed { - return Err(TryRecvError::Closed); - } - Ok(RecvGuard { slot }) } } impl<T: Clone> Receiver<T> { + /// Re-subscribes to the channel starting from the current tail element. + /// + /// This [`Receiver`] handle will receive a clone of all values sent + /// **after** it has resubscribed. This will not include elements that are + /// in the queue of the current receiver. Consider the following example. + /// + /// # Examples + /// + /// ``` + /// use tokio::sync::broadcast; + /// + /// #[tokio::main] + /// async fn main() { + /// let (tx, mut rx) = broadcast::channel(2); + /// + /// tx.send(1).unwrap(); + /// let mut rx2 = rx.resubscribe(); + /// tx.send(2).unwrap(); + /// + /// assert_eq!(rx2.recv().await.unwrap(), 2); + /// assert_eq!(rx.recv().await.unwrap(), 1); + /// } + /// ``` + pub fn resubscribe(&self) -> Self { + let shared = self.shared.clone(); + new_receiver(shared) + } /// Receives the next value for this receiver. /// /// Each [`Receiver`] handle will receive a clone of all values sent @@ -930,6 +1194,33 @@ impl<T: Clone> Receiver<T> { let guard = self.recv_ref(None)?; guard.clone_value().ok_or(TryRecvError::Closed) } + + /// Blocking receive to call outside of asynchronous contexts. + /// + /// # Panics + /// + /// This function panics if called within an asynchronous execution + /// context. + /// + /// # Examples + /// ``` + /// use std::thread; + /// use tokio::sync::broadcast; + /// + /// #[tokio::main] + /// async fn main() { + /// let (tx, mut rx) = broadcast::channel(16); + /// + /// let sync_code = thread::spawn(move || { + /// assert_eq!(rx.blocking_recv(), Ok(10)); + /// }); + /// + /// let _ = tx.send(10); + /// sync_code.join().unwrap(); + /// } + pub fn blocking_recv(&mut self) -> Result<T, RecvError> { + crate::future::block_on(self.recv()) + } } impl<T> Drop for Receiver<T> { @@ -988,6 +1279,8 @@ where type Output = Result<T, RecvError>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<T, RecvError>> { + ready!(crate::trace::trace_leaf(cx)); + let (receiver, waiter) = self.project(); let guard = match receiver.recv_ref(Some((waiter, cx.waker()))) { @@ -1039,8 +1332,8 @@ unsafe impl linked_list::Link for Waiter { ptr } - unsafe fn pointers(mut target: NonNull<Waiter>) -> NonNull<linked_list::Pointers<Waiter>> { - NonNull::from(&mut target.as_mut().pointers) + unsafe fn pointers(target: NonNull<Waiter>) -> NonNull<linked_list::Pointers<Waiter>> { + Waiter::addr_of_pointers(target) } } diff --git a/vendor/tokio/src/sync/mod.rs b/vendor/tokio/src/sync/mod.rs index 457e6ab29..70fd9b9e3 100644 --- a/vendor/tokio/src/sync/mod.rs +++ b/vendor/tokio/src/sync/mod.rs @@ -94,6 +94,10 @@ //! producers to a single consumer. This channel is often used to send work to a //! task or to receive the result of many computations. //! +//! This is also the channel you should use if you want to send many messages +//! from a single producer to a single consumer. There is no dedicated spsc +//! channel. +//! //! **Example:** using an mpsc to incrementally stream the results of a series //! of computations. //! @@ -244,6 +248,10 @@ //! This channel tends to be used less often than `oneshot` and `mpsc` but still //! has its use cases. //! +//! This is also the channel you should use if you want to broadcast values from +//! a single producer to many consumers. There is no dedicated spmc broadcast +//! channel. +//! //! Basic usage //! //! ``` @@ -441,7 +449,7 @@ cfg_sync! { pub mod mpsc; mod mutex; - pub use mutex::{Mutex, MutexGuard, TryLockError, OwnedMutexGuard, MappedMutexGuard}; + pub use mutex::{Mutex, MutexGuard, TryLockError, OwnedMutexGuard, MappedMutexGuard, OwnedMappedMutexGuard}; pub(crate) mod notify; pub use notify::Notify; diff --git a/vendor/tokio/src/sync/mpsc/block.rs b/vendor/tokio/src/sync/mpsc/block.rs index 1c9ab14e9..39c3e1be2 100644 --- a/vendor/tokio/src/sync/mpsc/block.rs +++ b/vendor/tokio/src/sync/mpsc/block.rs @@ -1,7 +1,7 @@ use crate::loom::cell::UnsafeCell; use crate::loom::sync::atomic::{AtomicPtr, AtomicUsize}; -use crate::loom::thread; +use std::alloc::Layout; use std::mem::MaybeUninit; use std::ops; use std::ptr::{self, NonNull}; @@ -11,6 +11,17 @@ use std::sync::atomic::Ordering::{self, AcqRel, Acquire, Release}; /// /// Each block in the list can hold up to `BLOCK_CAP` messages. pub(crate) struct Block<T> { + /// The header fields. + header: BlockHeader<T>, + + /// Array containing values pushed into the block. Values are stored in a + /// continuous array in order to improve cache line behavior when reading. + /// The values must be manually dropped. + values: Values<T>, +} + +/// Extra fields for a `Block<T>`. +struct BlockHeader<T> { /// The start index of this block. /// /// Slots in this block have indices in `start_index .. start_index + BLOCK_CAP`. @@ -25,11 +36,6 @@ pub(crate) struct Block<T> { /// The observed `tail_position` value *after* the block has been passed by /// `block_tail`. observed_tail_position: UnsafeCell<usize>, - - /// Array containing values pushed into the block. Values are stored in a - /// continuous array in order to improve cache line behavior when reading. - /// The values must be manually dropped. - values: Values<T>, } pub(crate) enum Read<T> { @@ -37,11 +43,12 @@ pub(crate) enum Read<T> { Closed, } +#[repr(transparent)] struct Values<T>([UnsafeCell<MaybeUninit<T>>; BLOCK_CAP]); use super::BLOCK_CAP; -/// Masks an index to get the block identifier +/// Masks an index to get the block identifier. const BLOCK_MASK: usize = !(BLOCK_CAP - 1); /// Masks an index to get the value offset in a block. @@ -72,28 +79,56 @@ pub(crate) fn offset(slot_index: usize) -> usize { SLOT_MASK & slot_index } +generate_addr_of_methods! { + impl<T> Block<T> { + unsafe fn addr_of_header(self: NonNull<Self>) -> NonNull<BlockHeader<T>> { + &self.header + } + + unsafe fn addr_of_values(self: NonNull<Self>) -> NonNull<Values<T>> { + &self.values + } + } +} + impl<T> Block<T> { - pub(crate) fn new(start_index: usize) -> Block<T> { - Block { - // The absolute index in the channel of the first slot in the block. - start_index, + pub(crate) fn new(start_index: usize) -> Box<Block<T>> { + unsafe { + // Allocate the block on the heap. + // SAFETY: The size of the Block<T> is non-zero, since it is at least the size of the header. + let block = std::alloc::alloc(Layout::new::<Block<T>>()) as *mut Block<T>; + let block = match NonNull::new(block) { + Some(block) => block, + None => std::alloc::handle_alloc_error(Layout::new::<Block<T>>()), + }; + + // Write the header to the block. + Block::addr_of_header(block).as_ptr().write(BlockHeader { + // The absolute index in the channel of the first slot in the block. + start_index, - // Pointer to the next block in the linked list. - next: AtomicPtr::new(ptr::null_mut()), + // Pointer to the next block in the linked list. + next: AtomicPtr::new(ptr::null_mut()), - ready_slots: AtomicUsize::new(0), + ready_slots: AtomicUsize::new(0), - observed_tail_position: UnsafeCell::new(0), + observed_tail_position: UnsafeCell::new(0), + }); - // Value storage - values: unsafe { Values::uninitialized() }, + // Initialize the values array. + Values::initialize(Block::addr_of_values(block)); + + // Convert the pointer to a `Box`. + // Safety: The raw pointer was allocated using the global allocator, and with + // the layout for a `Block<T>`, so it's valid to convert it to box. + Box::from_raw(block.as_ptr()) } } - /// Returns `true` if the block matches the given index + /// Returns `true` if the block matches the given index. pub(crate) fn is_at_index(&self, index: usize) -> bool { debug_assert!(offset(index) == 0); - self.start_index == index + self.header.start_index == index } /// Returns the number of blocks between `self` and the block at the @@ -102,7 +137,7 @@ impl<T> Block<T> { /// `start_index` must represent a block *after* `self`. pub(crate) fn distance(&self, other_index: usize) -> usize { debug_assert!(offset(other_index) == 0); - other_index.wrapping_sub(self.start_index) / BLOCK_CAP + other_index.wrapping_sub(self.header.start_index) / BLOCK_CAP } /// Reads the value at the given offset. @@ -117,7 +152,7 @@ impl<T> Block<T> { pub(crate) unsafe fn read(&self, slot_index: usize) -> Option<Read<T>> { let offset = offset(slot_index); - let ready_bits = self.ready_slots.load(Acquire); + let ready_bits = self.header.ready_slots.load(Acquire); if !is_ready(ready_bits, offset) { if is_tx_closed(ready_bits) { @@ -157,7 +192,7 @@ impl<T> Block<T> { /// Signal to the receiver that the sender half of the list is closed. pub(crate) unsafe fn tx_close(&self) { - self.ready_slots.fetch_or(TX_CLOSED, Release); + self.header.ready_slots.fetch_or(TX_CLOSED, Release); } /// Resets the block to a blank state. This enables reusing blocks in the @@ -170,9 +205,9 @@ impl<T> Block<T> { /// * All slots are empty. /// * The caller holds a unique pointer to the block. pub(crate) unsafe fn reclaim(&mut self) { - self.start_index = 0; - self.next = AtomicPtr::new(ptr::null_mut()); - self.ready_slots = AtomicUsize::new(0); + self.header.start_index = 0; + self.header.next = AtomicPtr::new(ptr::null_mut()); + self.header.ready_slots = AtomicUsize::new(0); } /// Releases the block to the rx half for freeing. @@ -188,19 +223,20 @@ impl<T> Block<T> { pub(crate) unsafe fn tx_release(&self, tail_position: usize) { // Track the observed tail_position. Any sender targeting a greater // tail_position is guaranteed to not access this block. - self.observed_tail_position + self.header + .observed_tail_position .with_mut(|ptr| *ptr = tail_position); // Set the released bit, signalling to the receiver that it is safe to // free the block's memory as soon as all slots **prior** to // `observed_tail_position` have been filled. - self.ready_slots.fetch_or(RELEASED, Release); + self.header.ready_slots.fetch_or(RELEASED, Release); } /// Mark a slot as ready fn set_ready(&self, slot: usize) { let mask = 1 << slot; - self.ready_slots.fetch_or(mask, Release); + self.header.ready_slots.fetch_or(mask, Release); } /// Returns `true` when all slots have their `ready` bits set. @@ -215,25 +251,31 @@ impl<T> Block<T> { /// single atomic cell. However, this could have negative impact on cache /// behavior as there would be many more mutations to a single slot. pub(crate) fn is_final(&self) -> bool { - self.ready_slots.load(Acquire) & READY_MASK == READY_MASK + self.header.ready_slots.load(Acquire) & READY_MASK == READY_MASK } /// Returns the `observed_tail_position` value, if set pub(crate) fn observed_tail_position(&self) -> Option<usize> { - if 0 == RELEASED & self.ready_slots.load(Acquire) { + if 0 == RELEASED & self.header.ready_slots.load(Acquire) { None } else { - Some(self.observed_tail_position.with(|ptr| unsafe { *ptr })) + Some( + self.header + .observed_tail_position + .with(|ptr| unsafe { *ptr }), + ) } } /// Loads the next block pub(crate) fn load_next(&self, ordering: Ordering) -> Option<NonNull<Block<T>>> { - let ret = NonNull::new(self.next.load(ordering)); + let ret = NonNull::new(self.header.next.load(ordering)); debug_assert!(unsafe { - ret.map(|block| block.as_ref().start_index == self.start_index.wrapping_add(BLOCK_CAP)) - .unwrap_or(true) + ret.map(|block| { + block.as_ref().header.start_index == self.header.start_index.wrapping_add(BLOCK_CAP) + }) + .unwrap_or(true) }); ret @@ -261,9 +303,10 @@ impl<T> Block<T> { success: Ordering, failure: Ordering, ) -> Result<(), NonNull<Block<T>>> { - block.as_mut().start_index = self.start_index.wrapping_add(BLOCK_CAP); + block.as_mut().header.start_index = self.header.start_index.wrapping_add(BLOCK_CAP); let next_ptr = self + .header .next .compare_exchange(ptr::null_mut(), block.as_ptr(), success, failure) .unwrap_or_else(|x| x); @@ -292,7 +335,7 @@ impl<T> Block<T> { // Create the new block. It is assumed that the block will become the // next one after `&self`. If this turns out to not be the case, // `start_index` is updated accordingly. - let new_block = Box::new(Block::new(self.start_index + BLOCK_CAP)); + let new_block = Block::new(self.header.start_index + BLOCK_CAP); let mut new_block = unsafe { NonNull::new_unchecked(Box::into_raw(new_block)) }; @@ -309,7 +352,8 @@ impl<T> Block<T> { // `Release` ensures that the newly allocated block is available to // other threads acquiring the next pointer. let next = NonNull::new( - self.next + self.header + .next .compare_exchange(ptr::null_mut(), new_block.as_ptr(), AcqRel, Acquire) .unwrap_or_else(|x| x), ); @@ -344,8 +388,7 @@ impl<T> Block<T> { Err(curr) => curr, }; - // When running outside of loom, this calls `spin_loop_hint`. - thread::yield_now(); + crate::loom::thread::yield_now(); } } } @@ -362,19 +405,20 @@ fn is_tx_closed(bits: usize) -> bool { } impl<T> Values<T> { - unsafe fn uninitialized() -> Values<T> { - let mut vals = MaybeUninit::uninit(); - + /// Initialize a `Values` struct from a pointer. + /// + /// # Safety + /// + /// The raw pointer must be valid for writing a `Values<T>`. + unsafe fn initialize(_value: NonNull<Values<T>>) { // When fuzzing, `UnsafeCell` needs to be initialized. if_loom! { - let p = vals.as_mut_ptr() as *mut UnsafeCell<MaybeUninit<T>>; + let p = _value.as_ptr() as *mut UnsafeCell<MaybeUninit<T>>; for i in 0..BLOCK_CAP { p.add(i) .write(UnsafeCell::new(MaybeUninit::uninit())); } } - - Values(vals.assume_init()) } } @@ -385,3 +429,20 @@ impl<T> ops::Index<usize> for Values<T> { self.0.index(index) } } + +#[cfg(all(test, not(loom)))] +#[test] +fn assert_no_stack_overflow() { + // https://github.com/tokio-rs/tokio/issues/5293 + + struct Foo { + _a: [u8; 2_000_000], + } + + assert_eq!( + Layout::new::<MaybeUninit<Block<Foo>>>(), + Layout::new::<Block<Foo>>() + ); + + let _block = Block::<Foo>::new(0); +} diff --git a/vendor/tokio/src/sync/mpsc/bounded.rs b/vendor/tokio/src/sync/mpsc/bounded.rs index d7af17251..e870ae5f4 100644 --- a/vendor/tokio/src/sync/mpsc/bounded.rs +++ b/vendor/tokio/src/sync/mpsc/bounded.rs @@ -1,6 +1,7 @@ +use crate::loom::sync::Arc; use crate::sync::batch_semaphore::{self as semaphore, TryAcquireError}; use crate::sync::mpsc::chan; -use crate::sync::mpsc::error::{SendError, TrySendError}; +use crate::sync::mpsc::error::{SendError, TryRecvError, TrySendError}; cfg_time! { use crate::sync::mpsc::error::SendTimeoutError; @@ -10,19 +11,53 @@ cfg_time! { use std::fmt; use std::task::{Context, Poll}; -/// Send values to the associated `Receiver`. +/// Sends values to the associated `Receiver`. /// /// Instances are created by the [`channel`](channel) function. /// -/// To use the `Sender` in a poll function, you can use the [`PollSender`] -/// utility. +/// To convert the `Sender` into a `Sink` or use it in a poll function, you can +/// use the [`PollSender`] utility. /// -/// [`PollSender`]: https://docs.rs/tokio-util/0.6/tokio_util/sync/struct.PollSender.html +/// [`PollSender`]: https://docs.rs/tokio-util/latest/tokio_util/sync/struct.PollSender.html pub struct Sender<T> { chan: chan::Tx<T, Semaphore>, } -/// Permit to send one value into the channel. +/// A sender that does not prevent the channel from being closed. +/// +/// If all [`Sender`] instances of a channel were dropped and only `WeakSender` +/// instances remain, the channel is closed. +/// +/// In order to send messages, the `WeakSender` needs to be upgraded using +/// [`WeakSender::upgrade`], which returns `Option<Sender>`. It returns `None` +/// if all `Sender`s have been dropped, and otherwise it returns a `Sender`. +/// +/// [`Sender`]: Sender +/// [`WeakSender::upgrade`]: WeakSender::upgrade +/// +/// #Examples +/// +/// ``` +/// use tokio::sync::mpsc::channel; +/// +/// #[tokio::main] +/// async fn main() { +/// let (tx, _rx) = channel::<i32>(15); +/// let tx_weak = tx.downgrade(); +/// +/// // Upgrading will succeed because `tx` still exists. +/// assert!(tx_weak.upgrade().is_some()); +/// +/// // If we drop `tx`, then it will fail. +/// drop(tx); +/// assert!(tx_weak.clone().upgrade().is_none()); +/// } +/// ``` +pub struct WeakSender<T> { + chan: Arc<chan::Chan<T, Semaphore>>, +} + +/// Permits to send one value into the channel. /// /// `Permit` values are returned by [`Sender::reserve()`] and [`Sender::try_reserve()`] /// and are used to guarantee channel capacity before generating a message to send. @@ -49,7 +84,7 @@ pub struct OwnedPermit<T> { chan: Option<chan::Tx<T, Semaphore>>, } -/// Receive values from the associated `Sender`. +/// Receives values from the associated `Sender`. /// /// Instances are created by the [`channel`](channel) function. /// @@ -57,7 +92,7 @@ pub struct OwnedPermit<T> { /// /// [`ReceiverStream`]: https://docs.rs/tokio-stream/0.1/tokio_stream/wrappers/struct.ReceiverStream.html pub struct Receiver<T> { - /// The channel receiver + /// The channel receiver. chan: chan::Rx<T, Semaphore>, } @@ -105,9 +140,13 @@ pub struct Receiver<T> { /// } /// } /// ``` +#[track_caller] pub fn channel<T>(buffer: usize) -> (Sender<T>, Receiver<T>) { assert!(buffer > 0, "mpsc bounded channel requires buffer > 0"); - let semaphore = (semaphore::Semaphore::new(buffer), buffer); + let semaphore = Semaphore { + semaphore: semaphore::Semaphore::new(buffer), + bound: buffer, + }; let (tx, rx) = chan::channel(semaphore); let tx = Sender::new(tx); @@ -118,7 +157,11 @@ pub fn channel<T>(buffer: usize) -> (Sender<T>, Receiver<T>) { /// Channel semaphore is a tuple of the semaphore implementation and a `usize` /// representing the channel bound. -type Semaphore = (semaphore::Semaphore, usize); +#[derive(Debug)] +pub(crate) struct Semaphore { + pub(crate) semaphore: semaphore::Semaphore, + pub(crate) bound: usize, +} impl<T> Receiver<T> { pub(crate) fn new(chan: chan::Rx<T, Semaphore>) -> Receiver<T> { @@ -187,6 +230,50 @@ impl<T> Receiver<T> { poll_fn(|cx| self.chan.recv(cx)).await } + /// Tries to receive the next value for this receiver. + /// + /// This method returns the [`Empty`] error if the channel is currently + /// empty, but there are still outstanding [senders] or [permits]. + /// + /// This method returns the [`Disconnected`] error if the channel is + /// currently empty, and there are no outstanding [senders] or [permits]. + /// + /// Unlike the [`poll_recv`] method, this method will never return an + /// [`Empty`] error spuriously. + /// + /// [`Empty`]: crate::sync::mpsc::error::TryRecvError::Empty + /// [`Disconnected`]: crate::sync::mpsc::error::TryRecvError::Disconnected + /// [`poll_recv`]: Self::poll_recv + /// [senders]: crate::sync::mpsc::Sender + /// [permits]: crate::sync::mpsc::Permit + /// + /// # Examples + /// + /// ``` + /// use tokio::sync::mpsc; + /// use tokio::sync::mpsc::error::TryRecvError; + /// + /// #[tokio::main] + /// async fn main() { + /// let (tx, mut rx) = mpsc::channel(100); + /// + /// tx.send("hello").await.unwrap(); + /// + /// assert_eq!(Ok("hello"), rx.try_recv()); + /// assert_eq!(Err(TryRecvError::Empty), rx.try_recv()); + /// + /// tx.send("hello").await.unwrap(); + /// // Drop the last sender, closing the channel. + /// drop(tx); + /// + /// assert_eq!(Ok("hello"), rx.try_recv()); + /// assert_eq!(Err(TryRecvError::Disconnected), rx.try_recv()); + /// } + /// ``` + pub fn try_recv(&mut self) -> Result<T, TryRecvError> { + self.chan.try_recv() + } + /// Blocking receive to call outside of asynchronous contexts. /// /// This method returns `None` if the channel has been closed and there are @@ -237,7 +324,9 @@ impl<T> Receiver<T> { /// sync_code.join().unwrap() /// } /// ``` + #[track_caller] #[cfg(feature = "sync")] + #[cfg_attr(docsrs, doc(alias = "recv_blocking"))] pub fn blocking_recv(&mut self) -> Option<T> { crate::future::block_on(self.recv()) } @@ -291,7 +380,7 @@ impl<T> Receiver<T> { /// This method returns: /// /// * `Poll::Pending` if no messages are available but the channel is not - /// closed. + /// closed, or if a spurious failure happens. /// * `Poll::Ready(Some(message))` if a message is available. /// * `Poll::Ready(None)` if the channel has been closed and all messages /// sent before it was closed have been received. @@ -301,6 +390,12 @@ impl<T> Receiver<T> { /// receiver, or when the channel is closed. Note that on multiple calls to /// `poll_recv`, only the `Waker` from the `Context` passed to the most /// recent call is scheduled to receive a wakeup. + /// + /// If this method returns `Poll::Pending` due to a spurious failure, then + /// the `Waker` will be notified when the situation causing the spurious + /// failure has been resolved. Note that receiving such a wakeup does not + /// guarantee that the next call will succeed — it could fail with another + /// spurious failure. pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<T>> { self.chan.recv(cx) } @@ -485,7 +580,7 @@ impl<T> Sender<T> { /// } /// ``` pub fn try_send(&self, message: T) -> Result<(), TrySendError<T>> { - match self.chan.semaphore().0.try_acquire(1) { + match self.chan.semaphore().semaphore.try_acquire(1) { Ok(_) => {} Err(TryAcquireError::Closed) => return Err(TrySendError::Closed(message)), Err(TryAcquireError::NoPermits) => return Err(TrySendError::Full(message)), @@ -513,6 +608,11 @@ impl<T> Sender<T> { /// [`close`]: Receiver::close /// [`Receiver`]: Receiver /// + /// # Panics + /// + /// This function panics if it is called outside the context of a Tokio + /// runtime [with time enabled](crate::runtime::Builder::enable_time). + /// /// # Examples /// /// In the following example, each call to `send_timeout` will block until the @@ -595,7 +695,9 @@ impl<T> Sender<T> { /// sync_code.join().unwrap() /// } /// ``` + #[track_caller] #[cfg(feature = "sync")] + #[cfg_attr(docsrs, doc(alias = "send_blocking"))] pub fn blocking_send(&self, value: T) -> Result<(), SendError<T>> { crate::future::block_on(self.send(value)) } @@ -622,7 +724,7 @@ impl<T> Sender<T> { self.chan.is_closed() } - /// Wait for channel capacity. Once capacity to send one message is + /// Waits for channel capacity. Once capacity to send one message is /// available, it is reserved for the caller. /// /// If the channel is full, the function waits for the number of unreceived @@ -671,7 +773,7 @@ impl<T> Sender<T> { Ok(Permit { chan: &self.chan }) } - /// Wait for channel capacity, moving the `Sender` and returning an owned + /// Waits for channel capacity, moving the `Sender` and returning an owned /// permit. Once capacity to send one message is available, it is reserved /// for the caller. /// @@ -759,13 +861,15 @@ impl<T> Sender<T> { } async fn reserve_inner(&self) -> Result<(), SendError<()>> { - match self.chan.semaphore().0.acquire(1).await { + crate::trace::async_trace_leaf().await; + + match self.chan.semaphore().semaphore.acquire(1).await { Ok(_) => Ok(()), Err(_) => Err(SendError(())), } } - /// Try to acquire a slot in the channel without waiting for the slot to become + /// Tries to acquire a slot in the channel without waiting for the slot to become /// available. /// /// If the channel is full this function will return [`TrySendError`], otherwise @@ -809,15 +913,16 @@ impl<T> Sender<T> { /// } /// ``` pub fn try_reserve(&self) -> Result<Permit<'_, T>, TrySendError<()>> { - match self.chan.semaphore().0.try_acquire(1) { + match self.chan.semaphore().semaphore.try_acquire(1) { Ok(_) => {} - Err(_) => return Err(TrySendError::Full(())), + Err(TryAcquireError::Closed) => return Err(TrySendError::Closed(())), + Err(TryAcquireError::NoPermits) => return Err(TrySendError::Full(())), } Ok(Permit { chan: &self.chan }) } - /// Try to acquire a slot in the channel without waiting for the slot to become + /// Tries to acquire a slot in the channel without waiting for the slot to become /// available, returning an owned permit. /// /// This moves the sender _by value_, and returns an owned permit that can @@ -873,9 +978,10 @@ impl<T> Sender<T> { /// } /// ``` pub fn try_reserve_owned(self) -> Result<OwnedPermit<T>, TrySendError<Self>> { - match self.chan.semaphore().0.try_acquire(1) { + match self.chan.semaphore().semaphore.try_acquire(1) { Ok(_) => {} - Err(_) => return Err(TrySendError::Full(self)), + Err(TryAcquireError::Closed) => return Err(TrySendError::Closed(self)), + Err(TryAcquireError::NoPermits) => return Err(TrySendError::Full(self)), } Ok(OwnedPermit { @@ -903,6 +1009,8 @@ impl<T> Sender<T> { /// /// The capacity goes down when sending a value by calling [`send`] or by reserving capacity /// with [`reserve`]. The capacity goes up when values are received by the [`Receiver`]. + /// This is distinct from [`max_capacity`], which always returns buffer capacity initially + /// specified when calling [`channel`] /// /// # Examples /// @@ -928,8 +1036,56 @@ impl<T> Sender<T> { /// /// [`send`]: Sender::send /// [`reserve`]: Sender::reserve + /// [`channel`]: channel + /// [`max_capacity`]: Sender::max_capacity pub fn capacity(&self) -> usize { - self.chan.semaphore().0.available_permits() + self.chan.semaphore().semaphore.available_permits() + } + + /// Converts the `Sender` to a [`WeakSender`] that does not count + /// towards RAII semantics, i.e. if all `Sender` instances of the + /// channel were dropped and only `WeakSender` instances remain, + /// the channel is closed. + pub fn downgrade(&self) -> WeakSender<T> { + WeakSender { + chan: self.chan.downgrade(), + } + } + + /// Returns the maximum buffer capacity of the channel. + /// + /// The maximum capacity is the buffer capacity initially specified when calling + /// [`channel`]. This is distinct from [`capacity`], which returns the *current* + /// available buffer capacity: as messages are sent and received, the + /// value returned by [`capacity`] will go up or down, whereas the value + /// returned by `max_capacity` will remain constant. + /// + /// # Examples + /// + /// ``` + /// use tokio::sync::mpsc; + /// + /// #[tokio::main] + /// async fn main() { + /// let (tx, _rx) = mpsc::channel::<()>(5); + /// + /// // both max capacity and capacity are the same at first + /// assert_eq!(tx.max_capacity(), 5); + /// assert_eq!(tx.capacity(), 5); + /// + /// // Making a reservation doesn't change the max capacity. + /// let permit = tx.reserve().await.unwrap(); + /// assert_eq!(tx.max_capacity(), 5); + /// // but drops the capacity by one + /// assert_eq!(tx.capacity(), 4); + /// } + /// ``` + /// + /// [`channel`]: channel + /// [`max_capacity`]: Sender::max_capacity + /// [`capacity`]: Sender::capacity + pub fn max_capacity(&self) -> usize { + self.chan.semaphore().bound } } @@ -949,6 +1105,29 @@ impl<T> fmt::Debug for Sender<T> { } } +impl<T> Clone for WeakSender<T> { + fn clone(&self) -> Self { + WeakSender { + chan: self.chan.clone(), + } + } +} + +impl<T> WeakSender<T> { + /// Tries to convert a WeakSender into a [`Sender`]. This will return `Some` + /// if there are other `Sender` instances alive and the channel wasn't + /// previously dropped, otherwise `None` is returned. + pub fn upgrade(&self) -> Option<Sender<T>> { + chan::Tx::upgrade(self.chan.clone()).map(Sender::new) + } +} + +impl<T> fmt::Debug for WeakSender<T> { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt.debug_struct("WeakSender").finish() + } +} + // ===== impl Permit ===== impl<T> Permit<'_, T> { @@ -1065,7 +1244,7 @@ impl<T> OwnedPermit<T> { Sender { chan } } - /// Release the reserved capacity *without* sending a message, returning the + /// Releases the reserved capacity *without* sending a message, returning the /// [`Sender`]. /// /// # Examples diff --git a/vendor/tokio/src/sync/mpsc/chan.rs b/vendor/tokio/src/sync/mpsc/chan.rs index 554d02284..6f87715dd 100644 --- a/vendor/tokio/src/sync/mpsc/chan.rs +++ b/vendor/tokio/src/sync/mpsc/chan.rs @@ -2,16 +2,18 @@ use crate::loom::cell::UnsafeCell; use crate::loom::future::AtomicWaker; use crate::loom::sync::atomic::AtomicUsize; use crate::loom::sync::Arc; -use crate::sync::mpsc::list; +use crate::runtime::park::CachedParkThread; +use crate::sync::mpsc::error::TryRecvError; +use crate::sync::mpsc::{bounded, list, unbounded}; use crate::sync::notify::Notify; use std::fmt; use std::process; -use std::sync::atomic::Ordering::{AcqRel, Relaxed}; +use std::sync::atomic::Ordering::{AcqRel, Acquire, Relaxed, Release}; use std::task::Poll::{Pending, Ready}; use std::task::{Context, Poll}; -/// Channel sender +/// Channel sender. pub(crate) struct Tx<T, S> { inner: Arc<Chan<T, S>>, } @@ -22,7 +24,7 @@ impl<T, S: fmt::Debug> fmt::Debug for Tx<T, S> { } } -/// Channel receiver +/// Channel receiver. pub(crate) struct Rx<T, S: Semaphore> { inner: Arc<Chan<T, S>>, } @@ -43,8 +45,8 @@ pub(crate) trait Semaphore { fn is_closed(&self) -> bool; } -struct Chan<T, S> { - /// Notifies all tasks listening for the receiver being dropped +pub(super) struct Chan<T, S> { + /// Notifies all tasks listening for the receiver being dropped. notify_rx_closed: Notify, /// Handle to the push half of the lock-free list. @@ -126,6 +128,30 @@ impl<T, S> Tx<T, S> { Tx { inner: chan } } + pub(super) fn downgrade(&self) -> Arc<Chan<T, S>> { + self.inner.clone() + } + + // Returns the upgraded channel or None if the upgrade failed. + pub(super) fn upgrade(chan: Arc<Chan<T, S>>) -> Option<Self> { + let mut tx_count = chan.tx_count.load(Acquire); + + loop { + if tx_count == 0 { + // channel is closed + return None; + } + + match chan + .tx_count + .compare_exchange_weak(tx_count, tx_count + 1, AcqRel, Acquire) + { + Ok(_) => return Some(Tx { inner: chan }), + Err(prev_count) => tx_count = prev_count, + } + } + } + pub(super) fn semaphore(&self) -> &S { &self.inner.semaphore } @@ -216,8 +242,10 @@ impl<T, S: Semaphore> Rx<T, S> { pub(crate) fn recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<T>> { use super::block::Read::*; + ready!(crate::trace::trace_leaf(cx)); + // Keep track of task budget - let coop = ready!(crate::coop::poll_proceed(cx)); + let coop = ready!(crate::runtime::coop::poll_proceed(cx)); self.inner.rx_fields.with_mut(|rx_fields_ptr| { let rx_fields = unsafe { &mut *rx_fields_ptr }; @@ -263,6 +291,51 @@ impl<T, S: Semaphore> Rx<T, S> { } }) } + + /// Try to receive the next value. + pub(crate) fn try_recv(&mut self) -> Result<T, TryRecvError> { + use super::list::TryPopResult; + + self.inner.rx_fields.with_mut(|rx_fields_ptr| { + let rx_fields = unsafe { &mut *rx_fields_ptr }; + + macro_rules! try_recv { + () => { + match rx_fields.list.try_pop(&self.inner.tx) { + TryPopResult::Ok(value) => { + self.inner.semaphore.add_permit(); + return Ok(value); + } + TryPopResult::Closed => return Err(TryRecvError::Disconnected), + TryPopResult::Empty => return Err(TryRecvError::Empty), + TryPopResult::Busy => {} // fall through + } + }; + } + + try_recv!(); + + // If a previous `poll_recv` call has set a waker, we wake it here. + // This allows us to put our own CachedParkThread waker in the + // AtomicWaker slot instead. + // + // This is not a spurious wakeup to `poll_recv` since we just got a + // Busy from `try_pop`, which only happens if there are messages in + // the queue. + self.inner.rx_waker.wake(); + + // Park the thread until the problematic send has completed. + let mut park = CachedParkThread::new(); + let waker = park.waker().unwrap(); + loop { + self.inner.rx_waker.register_by_ref(&waker); + // It is possible that the problematic send has now completed, + // so we have to check for messages again. + try_recv!(); + park.park(); + } + }) + } } impl<T, S: Semaphore> Drop for Rx<T, S> { @@ -297,7 +370,7 @@ impl<T, S> Drop for Chan<T, S> { fn drop(&mut self) { use super::block::Read::Value; - // Safety: the only owner of the rx fields is Chan, and eing + // Safety: the only owner of the rx fields is Chan, and being // inside its own Drop means we're the last ones to touch it. self.rx_fields.with_mut(|rx_fields_ptr| { let rx_fields = unsafe { &mut *rx_fields_ptr }; @@ -310,32 +383,29 @@ impl<T, S> Drop for Chan<T, S> { // ===== impl Semaphore for (::Semaphore, capacity) ===== -impl Semaphore for (crate::sync::batch_semaphore::Semaphore, usize) { +impl Semaphore for bounded::Semaphore { fn add_permit(&self) { - self.0.release(1) + self.semaphore.release(1) } fn is_idle(&self) -> bool { - self.0.available_permits() == self.1 + self.semaphore.available_permits() == self.bound } fn close(&self) { - self.0.close(); + self.semaphore.close(); } fn is_closed(&self) -> bool { - self.0.is_closed() + self.semaphore.is_closed() } } // ===== impl Semaphore for AtomicUsize ===== -use std::sync::atomic::Ordering::{Acquire, Release}; -use std::usize; - -impl Semaphore for AtomicUsize { +impl Semaphore for unbounded::Semaphore { fn add_permit(&self) { - let prev = self.fetch_sub(2, Release); + let prev = self.0.fetch_sub(2, Release); if prev >> 1 == 0 { // Something went wrong @@ -344,14 +414,14 @@ impl Semaphore for AtomicUsize { } fn is_idle(&self) -> bool { - self.load(Acquire) >> 1 == 0 + self.0.load(Acquire) >> 1 == 0 } fn close(&self) { - self.fetch_or(1, Release); + self.0.fetch_or(1, Release); } fn is_closed(&self) -> bool { - self.load(Acquire) & 1 == 1 + self.0.load(Acquire) & 1 == 1 } } diff --git a/vendor/tokio/src/sync/mpsc/error.rs b/vendor/tokio/src/sync/mpsc/error.rs index 0d25ad386..25b4455be 100644 --- a/vendor/tokio/src/sync/mpsc/error.rs +++ b/vendor/tokio/src/sync/mpsc/error.rs @@ -1,25 +1,31 @@ -//! Channel error types +//! Channel error types. use std::error::Error; use std::fmt; /// Error returned by the `Sender`. -#[derive(Debug)] +#[derive(PartialEq, Eq, Clone, Copy)] pub struct SendError<T>(pub T); +impl<T> fmt::Debug for SendError<T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("SendError").finish_non_exhaustive() + } +} + impl<T> fmt::Display for SendError<T> { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { write!(fmt, "channel closed") } } -impl<T: fmt::Debug> std::error::Error for SendError<T> {} +impl<T> std::error::Error for SendError<T> {} // ===== TrySendError ===== /// This enumeration is the list of the possible error outcomes for the /// [try_send](super::Sender::try_send) method. -#[derive(Debug)] +#[derive(PartialEq, Eq, Clone, Copy)] pub enum TrySendError<T> { /// The data could not be sent on the channel because the channel is /// currently full and sending would require blocking. @@ -30,7 +36,14 @@ pub enum TrySendError<T> { Closed(T), } -impl<T: fmt::Debug> Error for TrySendError<T> {} +impl<T> fmt::Debug for TrySendError<T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match *self { + TrySendError::Full(..) => "Full(..)".fmt(f), + TrySendError::Closed(..) => "Closed(..)".fmt(f), + } + } +} impl<T> fmt::Display for TrySendError<T> { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { @@ -45,16 +58,42 @@ impl<T> fmt::Display for TrySendError<T> { } } +impl<T> Error for TrySendError<T> {} + impl<T> From<SendError<T>> for TrySendError<T> { fn from(src: SendError<T>) -> TrySendError<T> { TrySendError::Closed(src.0) } } +// ===== TryRecvError ===== + +/// Error returned by `try_recv`. +#[derive(PartialEq, Eq, Clone, Copy, Debug)] +pub enum TryRecvError { + /// This **channel** is currently empty, but the **Sender**(s) have not yet + /// disconnected, so data may yet become available. + Empty, + /// The **channel**'s sending half has become disconnected, and there will + /// never be any more data received on it. + Disconnected, +} + +impl fmt::Display for TryRecvError { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + match *self { + TryRecvError::Empty => "receiving on an empty channel".fmt(fmt), + TryRecvError::Disconnected => "receiving on a closed channel".fmt(fmt), + } + } +} + +impl Error for TryRecvError {} + // ===== RecvError ===== /// Error returned by `Receiver`. -#[derive(Debug)] +#[derive(Debug, Clone)] #[doc(hidden)] #[deprecated(note = "This type is unused because recv returns an Option.")] pub struct RecvError(()); @@ -72,7 +111,7 @@ impl Error for RecvError {} cfg_time! { // ===== SendTimeoutError ===== - #[derive(Debug)] + #[derive(PartialEq, Eq, Clone, Copy)] /// Error returned by [`Sender::send_timeout`](super::Sender::send_timeout)]. pub enum SendTimeoutError<T> { /// The data could not be sent on the channel because the channel is @@ -84,7 +123,14 @@ cfg_time! { Closed(T), } - impl<T: fmt::Debug> Error for SendTimeoutError<T> {} + impl<T> fmt::Debug for SendTimeoutError<T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match *self { + SendTimeoutError::Timeout(..) => "Timeout(..)".fmt(f), + SendTimeoutError::Closed(..) => "Closed(..)".fmt(f), + } + } + } impl<T> fmt::Display for SendTimeoutError<T> { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { @@ -98,4 +144,6 @@ cfg_time! { ) } } + + impl<T> Error for SendTimeoutError<T> {} } diff --git a/vendor/tokio/src/sync/mpsc/list.rs b/vendor/tokio/src/sync/mpsc/list.rs index 5dad2babf..10b29575b 100644 --- a/vendor/tokio/src/sync/mpsc/list.rs +++ b/vendor/tokio/src/sync/mpsc/list.rs @@ -8,31 +8,43 @@ use std::fmt; use std::ptr::NonNull; use std::sync::atomic::Ordering::{AcqRel, Acquire, Relaxed, Release}; -/// List queue transmit handle +/// List queue transmit handle. pub(crate) struct Tx<T> { /// Tail in the `Block` mpmc list. block_tail: AtomicPtr<Block<T>>, - /// Position to push the next message. This reference a block and offset + /// Position to push the next message. This references a block and offset /// into the block. tail_position: AtomicUsize, } /// List queue receive handle pub(crate) struct Rx<T> { - /// Pointer to the block being processed + /// Pointer to the block being processed. head: NonNull<Block<T>>, - /// Next slot index to process + /// Next slot index to process. index: usize, - /// Pointer to the next block pending release + /// Pointer to the next block pending release. free_head: NonNull<Block<T>>, } +/// Return value of `Rx::try_pop`. +pub(crate) enum TryPopResult<T> { + /// Successfully popped a value. + Ok(T), + /// The channel is empty. + Empty, + /// The channel is empty and closed. + Closed, + /// The channel is not empty, but the first value is being written. + Busy, +} + pub(crate) fn channel<T>() -> (Tx<T>, Rx<T>) { // Create the initial block shared between the tx and rx halves. - let initial_block = Box::new(Block::new(0)); + let initial_block = Block::new(0); let initial_block_ptr = Box::into_raw(initial_block); let tx = Tx { @@ -67,7 +79,7 @@ impl<T> Tx<T> { } } - /// Closes the send half of the list + /// Closes the send half of the list. /// /// Similar process as pushing a value, but instead of writing the value & /// setting the ready flag, the TX_CLOSED flag is set on the block. @@ -218,7 +230,7 @@ impl<T> fmt::Debug for Tx<T> { } impl<T> Rx<T> { - /// Pops the next value off the queue + /// Pops the next value off the queue. pub(crate) fn pop(&mut self, tx: &Tx<T>) -> Option<block::Read<T>> { // Advance `head`, if needed if !self.try_advancing_head() { @@ -240,6 +252,26 @@ impl<T> Rx<T> { } } + /// Pops the next value off the queue, detecting whether the block + /// is busy or empty on failure. + /// + /// This function exists because `Rx::pop` can return `None` even if the + /// channel's queue contains a message that has been completely written. + /// This can happen if the fully delivered message is behind another message + /// that is in the middle of being written to the block, since the channel + /// can't return the messages out of order. + pub(crate) fn try_pop(&mut self, tx: &Tx<T>) -> TryPopResult<T> { + let tail_position = tx.tail_position.load(Acquire); + let result = self.pop(tx); + + match result { + Some(block::Read::Value(t)) => TryPopResult::Ok(t), + Some(block::Read::Closed) => TryPopResult::Closed, + None if tail_position == self.index => TryPopResult::Empty, + None => TryPopResult::Busy, + } + } + /// Tries advancing the block pointer to the block referenced by `self.index`. /// /// Returns `true` if successful, `false` if there is no next block to load. diff --git a/vendor/tokio/src/sync/mpsc/mod.rs b/vendor/tokio/src/sync/mpsc/mod.rs index 879e3dcfc..b2af084b2 100644 --- a/vendor/tokio/src/sync/mpsc/mod.rs +++ b/vendor/tokio/src/sync/mpsc/mod.rs @@ -21,6 +21,9 @@ //! when additional capacity is available. In other words, the channel provides //! backpressure. //! +//! This channel is also suitable for the single-producer single-consumer +//! use-case. (Unless you only need to send one message, in which case you +//! should use the [oneshot] channel.) //! //! # Disconnection //! @@ -30,7 +33,8 @@ //! //! If the [`Receiver`] handle is dropped, then messages can no longer //! be read out of the channel. In this case, all further attempts to send will -//! result in an error. +//! result in an error. Additionally, all unread messages will be drained from the +//! channel and dropped. //! //! # Clean Shutdown //! @@ -58,6 +62,22 @@ //! [crossbeam][crossbeam-unbounded]. Similarly, for sending a message _from sync //! to async_, you should use an unbounded Tokio `mpsc` channel. //! +//! Please be aware that the above remarks were written with the `mpsc` channel +//! in mind, but they can also be generalized to other kinds of channels. In +//! general, any channel method that isn't marked async can be called anywhere, +//! including outside of the runtime. For example, sending a message on a +//! [oneshot] channel from outside the runtime is perfectly fine. +//! +//! # Multiple runtimes +//! +//! The mpsc channel does not care about which runtime you use it in, and can be +//! used to send messages from one runtime to another. It can also be used in +//! non-Tokio runtimes. +//! +//! There is one exception to the above: the [`send_timeout`] must be used from +//! within a Tokio runtime, however it is still not tied to one specific Tokio +//! runtime, and the sender may be moved from one Tokio runtime to another. +//! //! [`Sender`]: crate::sync::mpsc::Sender //! [`Receiver`]: crate::sync::mpsc::Receiver //! [bounded-send]: crate::sync::mpsc::Sender::send() @@ -66,21 +86,25 @@ //! [blocking-recv]: crate::sync::mpsc::Receiver::blocking_recv() //! [`UnboundedSender`]: crate::sync::mpsc::UnboundedSender //! [`UnboundedReceiver`]: crate::sync::mpsc::UnboundedReceiver +//! [oneshot]: crate::sync::oneshot //! [`Handle::block_on`]: crate::runtime::Handle::block_on() //! [std-unbounded]: std::sync::mpsc::channel //! [crossbeam-unbounded]: https://docs.rs/crossbeam/*/crossbeam/channel/fn.unbounded.html +//! [`send_timeout`]: crate::sync::mpsc::Sender::send_timeout pub(super) mod block; mod bounded; -pub use self::bounded::{channel, OwnedPermit, Permit, Receiver, Sender}; +pub use self::bounded::{channel, OwnedPermit, Permit, Receiver, Sender, WeakSender}; mod chan; pub(super) mod list; mod unbounded; -pub use self::unbounded::{unbounded_channel, UnboundedReceiver, UnboundedSender}; +pub use self::unbounded::{ + unbounded_channel, UnboundedReceiver, UnboundedSender, WeakUnboundedSender, +}; pub mod error; diff --git a/vendor/tokio/src/sync/mpsc/unbounded.rs b/vendor/tokio/src/sync/mpsc/unbounded.rs index 23c80f60a..cd83fc125 100644 --- a/vendor/tokio/src/sync/mpsc/unbounded.rs +++ b/vendor/tokio/src/sync/mpsc/unbounded.rs @@ -1,6 +1,6 @@ -use crate::loom::sync::atomic::AtomicUsize; +use crate::loom::sync::{atomic::AtomicUsize, Arc}; use crate::sync::mpsc::chan; -use crate::sync::mpsc::error::SendError; +use crate::sync::mpsc::error::{SendError, TryRecvError}; use std::fmt; use std::task::{Context, Poll}; @@ -13,6 +13,40 @@ pub struct UnboundedSender<T> { chan: chan::Tx<T, Semaphore>, } +/// An unbounded sender that does not prevent the channel from being closed. +/// +/// If all [`UnboundedSender`] instances of a channel were dropped and only +/// `WeakUnboundedSender` instances remain, the channel is closed. +/// +/// In order to send messages, the `WeakUnboundedSender` needs to be upgraded using +/// [`WeakUnboundedSender::upgrade`], which returns `Option<UnboundedSender>`. It returns `None` +/// if all `UnboundedSender`s have been dropped, and otherwise it returns an `UnboundedSender`. +/// +/// [`UnboundedSender`]: UnboundedSender +/// [`WeakUnboundedSender::upgrade`]: WeakUnboundedSender::upgrade +/// +/// #Examples +/// +/// ``` +/// use tokio::sync::mpsc::unbounded_channel; +/// +/// #[tokio::main] +/// async fn main() { +/// let (tx, _rx) = unbounded_channel::<i32>(); +/// let tx_weak = tx.downgrade(); +/// +/// // Upgrading will succeed because `tx` still exists. +/// assert!(tx_weak.upgrade().is_some()); +/// +/// // If we drop `tx`, then it will fail. +/// drop(tx); +/// assert!(tx_weak.clone().upgrade().is_none()); +/// } +/// ``` +pub struct WeakUnboundedSender<T> { + chan: Arc<chan::Chan<T, Semaphore>>, +} + impl<T> Clone for UnboundedSender<T> { fn clone(&self) -> Self { UnboundedSender { @@ -61,7 +95,7 @@ impl<T> fmt::Debug for UnboundedReceiver<T> { /// the channel. Using an `unbounded` channel has the ability of causing the /// process to run out of memory. In this case, the process will be aborted. pub fn unbounded_channel<T>() -> (UnboundedSender<T>, UnboundedReceiver<T>) { - let (tx, rx) = chan::channel(AtomicUsize::new(0)); + let (tx, rx) = chan::channel(Semaphore(AtomicUsize::new(0))); let tx = UnboundedSender::new(tx); let rx = UnboundedReceiver::new(rx); @@ -70,7 +104,8 @@ pub fn unbounded_channel<T>() -> (UnboundedSender<T>, UnboundedReceiver<T>) { } /// No capacity -type Semaphore = AtomicUsize; +#[derive(Debug)] +pub(crate) struct Semaphore(pub(crate) AtomicUsize); impl<T> UnboundedReceiver<T> { pub(crate) fn new(chan: chan::Rx<T, Semaphore>) -> UnboundedReceiver<T> { @@ -79,8 +114,14 @@ impl<T> UnboundedReceiver<T> { /// Receives the next value for this receiver. /// - /// `None` is returned when all `Sender` halves have dropped, indicating - /// that no further values can be sent on the channel. + /// This method returns `None` if the channel has been closed and there are + /// no remaining messages in the channel's buffer. This indicates that no + /// further values can ever be received from this `Receiver`. The channel is + /// closed when all senders have been dropped, or when [`close`] is called. + /// + /// If there are no messages in the channel's buffer, but the channel has + /// not yet been closed, this method will sleep until a message is sent or + /// the channel is closed. /// /// # Cancel safety /// @@ -89,6 +130,8 @@ impl<T> UnboundedReceiver<T> { /// completes first, it is guaranteed that no messages were received on this /// channel. /// + /// [`close`]: Self::close + /// /// # Examples /// /// ``` @@ -129,6 +172,50 @@ impl<T> UnboundedReceiver<T> { poll_fn(|cx| self.poll_recv(cx)).await } + /// Tries to receive the next value for this receiver. + /// + /// This method returns the [`Empty`] error if the channel is currently + /// empty, but there are still outstanding [senders] or [permits]. + /// + /// This method returns the [`Disconnected`] error if the channel is + /// currently empty, and there are no outstanding [senders] or [permits]. + /// + /// Unlike the [`poll_recv`] method, this method will never return an + /// [`Empty`] error spuriously. + /// + /// [`Empty`]: crate::sync::mpsc::error::TryRecvError::Empty + /// [`Disconnected`]: crate::sync::mpsc::error::TryRecvError::Disconnected + /// [`poll_recv`]: Self::poll_recv + /// [senders]: crate::sync::mpsc::Sender + /// [permits]: crate::sync::mpsc::Permit + /// + /// # Examples + /// + /// ``` + /// use tokio::sync::mpsc; + /// use tokio::sync::mpsc::error::TryRecvError; + /// + /// #[tokio::main] + /// async fn main() { + /// let (tx, mut rx) = mpsc::unbounded_channel(); + /// + /// tx.send("hello").unwrap(); + /// + /// assert_eq!(Ok("hello"), rx.try_recv()); + /// assert_eq!(Err(TryRecvError::Empty), rx.try_recv()); + /// + /// tx.send("hello").unwrap(); + /// // Drop the last sender, closing the channel. + /// drop(tx); + /// + /// assert_eq!(Ok("hello"), rx.try_recv()); + /// assert_eq!(Err(TryRecvError::Disconnected), rx.try_recv()); + /// } + /// ``` + pub fn try_recv(&mut self) -> Result<T, TryRecvError> { + self.chan.try_recv() + } + /// Blocking receive to call outside of asynchronous contexts. /// /// # Panics @@ -154,7 +241,9 @@ impl<T> UnboundedReceiver<T> { /// sync_code.join().unwrap(); /// } /// ``` + #[track_caller] #[cfg(feature = "sync")] + #[cfg_attr(docsrs, doc(alias = "recv_blocking"))] pub fn blocking_recv(&mut self) -> Option<T> { crate::future::block_on(self.recv()) } @@ -163,6 +252,9 @@ impl<T> UnboundedReceiver<T> { /// /// This prevents any further messages from being sent on the channel while /// still enabling the receiver to drain messages that are buffered. + /// + /// To guarantee that no messages are dropped, after calling `close()`, + /// `recv()` must be called until `None` is returned. pub fn close(&mut self) { self.chan.close(); } @@ -172,7 +264,7 @@ impl<T> UnboundedReceiver<T> { /// This method returns: /// /// * `Poll::Pending` if no messages are available but the channel is not - /// closed. + /// closed, or if a spurious failure happens. /// * `Poll::Ready(Some(message))` if a message is available. /// * `Poll::Ready(None)` if the channel has been closed and all messages /// sent before it was closed have been received. @@ -182,6 +274,12 @@ impl<T> UnboundedReceiver<T> { /// receiver, or when the channel is closed. Note that on multiple calls to /// `poll_recv`, only the `Waker` from the `Context` passed to the most /// recent call is scheduled to receive a wakeup. + /// + /// If this method returns `Poll::Pending` due to a spurious failure, then + /// the `Waker` will be notified when the situation causing the spurious + /// failure has been resolved. Note that receiving such a wakeup does not + /// guarantee that the next call will succeed — it could fail with another + /// spurious failure. pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<T>> { self.chan.recv(cx) } @@ -217,7 +315,7 @@ impl<T> UnboundedSender<T> { use std::process; use std::sync::atomic::Ordering::{AcqRel, Acquire}; - let mut curr = self.chan.semaphore().load(Acquire); + let mut curr = self.chan.semaphore().0.load(Acquire); loop { if curr & 1 == 1 { @@ -233,6 +331,7 @@ impl<T> UnboundedSender<T> { match self .chan .semaphore() + .0 .compare_exchange(curr, curr + 2, AcqRel, Acquire) { Ok(_) => return true, @@ -320,4 +419,37 @@ impl<T> UnboundedSender<T> { pub fn same_channel(&self, other: &Self) -> bool { self.chan.same_channel(&other.chan) } + + /// Converts the `UnboundedSender` to a [`WeakUnboundedSender`] that does not count + /// towards RAII semantics, i.e. if all `UnboundedSender` instances of the + /// channel were dropped and only `WeakUnboundedSender` instances remain, + /// the channel is closed. + pub fn downgrade(&self) -> WeakUnboundedSender<T> { + WeakUnboundedSender { + chan: self.chan.downgrade(), + } + } +} + +impl<T> Clone for WeakUnboundedSender<T> { + fn clone(&self) -> Self { + WeakUnboundedSender { + chan: self.chan.clone(), + } + } +} + +impl<T> WeakUnboundedSender<T> { + /// Tries to convert a WeakUnboundedSender into an [`UnboundedSender`]. + /// This will return `Some` if there are other `Sender` instances alive and + /// the channel wasn't previously dropped, otherwise `None` is returned. + pub fn upgrade(&self) -> Option<UnboundedSender<T>> { + chan::Tx::upgrade(self.chan.clone()).map(UnboundedSender::new) + } +} + +impl<T> fmt::Debug for WeakUnboundedSender<T> { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt.debug_struct("WeakUnboundedSender").finish() + } } diff --git a/vendor/tokio/src/sync/mutex.rs b/vendor/tokio/src/sync/mutex.rs index 8ae824770..549c77b32 100644 --- a/vendor/tokio/src/sync/mutex.rs +++ b/vendor/tokio/src/sync/mutex.rs @@ -1,12 +1,15 @@ #![cfg_attr(not(feature = "sync"), allow(unreachable_pub, dead_code))] use crate::sync::batch_semaphore as semaphore; +#[cfg(all(tokio_unstable, feature = "tracing"))] +use crate::util::trace; use std::cell::UnsafeCell; use std::error::Error; +use std::marker::PhantomData; use std::ops::{Deref, DerefMut}; use std::sync::Arc; -use std::{fmt, marker, mem}; +use std::{fmt, mem, ptr}; /// An asynchronous `Mutex`-like type. /// @@ -124,6 +127,8 @@ use std::{fmt, marker, mem}; /// [`Send`]: trait@std::marker::Send /// [`lock`]: method@Mutex::lock pub struct Mutex<T: ?Sized> { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, s: semaphore::Semaphore, c: UnsafeCell<T>, } @@ -137,7 +142,13 @@ pub struct Mutex<T: ?Sized> { /// /// The lock is automatically released whenever the guard is dropped, at which /// point `lock` will succeed yet again. +#[clippy::has_significant_drop] +#[must_use = "if unused the Mutex will immediately unlock"] pub struct MutexGuard<'a, T: ?Sized> { + // When changing the fields in this struct, make sure to update the + // `skip_drop` method. + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, lock: &'a Mutex<T>, } @@ -156,7 +167,12 @@ pub struct MutexGuard<'a, T: ?Sized> { /// point `lock` will succeed yet again. /// /// [`Arc`]: std::sync::Arc +#[clippy::has_significant_drop] pub struct OwnedMutexGuard<T: ?Sized> { + // When changing the fields in this struct, make sure to update the + // `skip_drop` method. + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, lock: Arc<Mutex<T>>, } @@ -165,12 +181,71 @@ pub struct OwnedMutexGuard<T: ?Sized> { /// This can be used to hold a subfield of the protected data. /// /// [`MutexGuard::map`]: method@MutexGuard::map +#[clippy::has_significant_drop] #[must_use = "if unused the Mutex will immediately unlock"] pub struct MappedMutexGuard<'a, T: ?Sized> { + // When changing the fields in this struct, make sure to update the + // `skip_drop` method. + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, s: &'a semaphore::Semaphore, data: *mut T, // Needed to tell the borrow checker that we are holding a `&mut T` - marker: marker::PhantomData<&'a mut T>, + marker: PhantomData<&'a mut T>, +} + +/// A owned handle to a held `Mutex` that has had a function applied to it via +/// [`OwnedMutexGuard::map`]. +/// +/// This can be used to hold a subfield of the protected data. +/// +/// [`OwnedMutexGuard::map`]: method@OwnedMutexGuard::map +#[clippy::has_significant_drop] +#[must_use = "if unused the Mutex will immediately unlock"] +pub struct OwnedMappedMutexGuard<T: ?Sized, U: ?Sized = T> { + // When changing the fields in this struct, make sure to update the + // `skip_drop` method. + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, + data: *mut U, + lock: Arc<Mutex<T>>, +} + +/// A helper type used when taking apart a `MutexGuard` without running its +/// Drop implementation. +#[allow(dead_code)] // Unused fields are still used in Drop. +struct MutexGuardInner<'a, T: ?Sized> { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, + lock: &'a Mutex<T>, +} + +/// A helper type used when taking apart a `OwnedMutexGuard` without running +/// its Drop implementation. +struct OwnedMutexGuardInner<T: ?Sized> { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, + lock: Arc<Mutex<T>>, +} + +/// A helper type used when taking apart a `MappedMutexGuard` without running +/// its Drop implementation. +#[allow(dead_code)] // Unused fields are still used in Drop. +struct MappedMutexGuardInner<'a, T: ?Sized> { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, + s: &'a semaphore::Semaphore, + data: *mut T, +} + +/// A helper type used when taking apart a `OwnedMappedMutexGuard` without running +/// its Drop implementation. +#[allow(dead_code)] // Unused fields are still used in Drop. +struct OwnedMappedMutexGuardInner<T: ?Sized, U: ?Sized> { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, + data: *mut U, + lock: Arc<Mutex<T>>, } // As long as T: Send, it's fine to send and share Mutex<T> between threads. @@ -183,6 +258,19 @@ unsafe impl<T> Sync for OwnedMutexGuard<T> where T: ?Sized + Send + Sync {} unsafe impl<'a, T> Sync for MappedMutexGuard<'a, T> where T: ?Sized + Sync + 'a {} unsafe impl<'a, T> Send for MappedMutexGuard<'a, T> where T: ?Sized + Send + 'a {} +unsafe impl<T, U> Sync for OwnedMappedMutexGuard<T, U> +where + T: ?Sized + Send + Sync, + U: ?Sized + Send + Sync, +{ +} +unsafe impl<T, U> Send for OwnedMappedMutexGuard<T, U> +where + T: ?Sized + Send, + U: ?Sized + Send, +{ +} + /// Error returned from the [`Mutex::try_lock`], [`RwLock::try_read`] and /// [`RwLock::try_write`] functions. /// @@ -191,8 +279,8 @@ unsafe impl<'a, T> Send for MappedMutexGuard<'a, T> where T: ?Sized + Send + 'a /// `RwLock::try_read` operation will only fail if the lock is currently held /// by an exclusive writer. /// -/// `RwLock::try_write` operation will if lock is held by any reader or by an -/// exclusive writer. +/// `RwLock::try_write` operation will only fail if the lock is currently held +/// by any reader or by an exclusive writer. /// /// [`Mutex::try_lock`]: Mutex::try_lock /// [`RwLock::try_read`]: fn@super::RwLock::try_read @@ -242,13 +330,42 @@ impl<T: ?Sized> Mutex<T> { /// /// let lock = Mutex::new(5); /// ``` + #[track_caller] pub fn new(t: T) -> Self where T: Sized, { + #[cfg(all(tokio_unstable, feature = "tracing"))] + let resource_span = { + let location = std::panic::Location::caller(); + + tracing::trace_span!( + "runtime.resource", + concrete_type = "Mutex", + kind = "Sync", + loc.file = location.file(), + loc.line = location.line(), + loc.col = location.column(), + ) + }; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + let s = resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + locked = false, + ); + semaphore::Semaphore::new(1) + }); + + #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] + let s = semaphore::Semaphore::new(1); + Self { c: UnsafeCell::new(t), - s: semaphore::Semaphore::new(1), + s, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span, } } @@ -270,6 +387,8 @@ impl<T: ?Sized> Mutex<T> { Self { c: UnsafeCell::new(t), s: semaphore::Semaphore::const_new(1), + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span::none(), } } @@ -297,8 +416,147 @@ impl<T: ?Sized> Mutex<T> { /// } /// ``` pub async fn lock(&self) -> MutexGuard<'_, T> { - self.acquire().await; - MutexGuard { lock: self } + let acquire_fut = async { + self.acquire().await; + + MutexGuard { + lock: self, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: self.resource_span.clone(), + } + }; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + let acquire_fut = trace::async_op( + move || acquire_fut, + self.resource_span.clone(), + "Mutex::lock", + "poll", + false, + ); + + #[allow(clippy::let_and_return)] // this lint triggers when disabling tracing + let guard = acquire_fut.await; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + locked = true, + ); + }); + + guard + } + + /// Blockingly locks this `Mutex`. When the lock has been acquired, function returns a + /// [`MutexGuard`]. + /// + /// This method is intended for use cases where you + /// need to use this mutex in asynchronous code as well as in synchronous code. + /// + /// # Panics + /// + /// This function panics if called within an asynchronous execution context. + /// + /// - If you find yourself in an asynchronous execution context and needing + /// to call some (synchronous) function which performs one of these + /// `blocking_` operations, then consider wrapping that call inside + /// [`spawn_blocking()`][crate::runtime::Handle::spawn_blocking] + /// (or [`block_in_place()`][crate::task::block_in_place]). + /// + /// # Examples + /// + /// ``` + /// use std::sync::Arc; + /// use tokio::sync::Mutex; + /// + /// #[tokio::main] + /// async fn main() { + /// let mutex = Arc::new(Mutex::new(1)); + /// let lock = mutex.lock().await; + /// + /// let mutex1 = Arc::clone(&mutex); + /// let blocking_task = tokio::task::spawn_blocking(move || { + /// // This shall block until the `lock` is released. + /// let mut n = mutex1.blocking_lock(); + /// *n = 2; + /// }); + /// + /// assert_eq!(*lock, 1); + /// // Release the lock. + /// drop(lock); + /// + /// // Await the completion of the blocking task. + /// blocking_task.await.unwrap(); + /// + /// // Assert uncontended. + /// let n = mutex.try_lock().unwrap(); + /// assert_eq!(*n, 2); + /// } + /// + /// ``` + #[track_caller] + #[cfg(feature = "sync")] + #[cfg_attr(docsrs, doc(alias = "lock_blocking"))] + pub fn blocking_lock(&self) -> MutexGuard<'_, T> { + crate::future::block_on(self.lock()) + } + + /// Blockingly locks this `Mutex`. When the lock has been acquired, function returns an + /// [`OwnedMutexGuard`]. + /// + /// This method is identical to [`Mutex::blocking_lock`], except that the returned + /// guard references the `Mutex` with an [`Arc`] rather than by borrowing + /// it. Therefore, the `Mutex` must be wrapped in an `Arc` to call this + /// method, and the guard will live for the `'static` lifetime, as it keeps + /// the `Mutex` alive by holding an `Arc`. + /// + /// # Panics + /// + /// This function panics if called within an asynchronous execution context. + /// + /// - If you find yourself in an asynchronous execution context and needing + /// to call some (synchronous) function which performs one of these + /// `blocking_` operations, then consider wrapping that call inside + /// [`spawn_blocking()`][crate::runtime::Handle::spawn_blocking] + /// (or [`block_in_place()`][crate::task::block_in_place]). + /// + /// # Examples + /// + /// ``` + /// use std::sync::Arc; + /// use tokio::sync::Mutex; + /// + /// #[tokio::main] + /// async fn main() { + /// let mutex = Arc::new(Mutex::new(1)); + /// let lock = mutex.lock().await; + /// + /// let mutex1 = Arc::clone(&mutex); + /// let blocking_task = tokio::task::spawn_blocking(move || { + /// // This shall block until the `lock` is released. + /// let mut n = mutex1.blocking_lock_owned(); + /// *n = 2; + /// }); + /// + /// assert_eq!(*lock, 1); + /// // Release the lock. + /// drop(lock); + /// + /// // Await the completion of the blocking task. + /// blocking_task.await.unwrap(); + /// + /// // Assert uncontended. + /// let n = mutex.try_lock().unwrap(); + /// assert_eq!(*n, 2); + /// } + /// + /// ``` + #[track_caller] + #[cfg(feature = "sync")] + pub fn blocking_lock_owned(self: Arc<Self>) -> OwnedMutexGuard<T> { + crate::future::block_on(self.lock_owned()) } /// Locks this mutex, causing the current task to yield until the lock has @@ -334,11 +592,45 @@ impl<T: ?Sized> Mutex<T> { /// /// [`Arc`]: std::sync::Arc pub async fn lock_owned(self: Arc<Self>) -> OwnedMutexGuard<T> { - self.acquire().await; - OwnedMutexGuard { lock: self } + #[cfg(all(tokio_unstable, feature = "tracing"))] + let resource_span = self.resource_span.clone(); + + let acquire_fut = async { + self.acquire().await; + + OwnedMutexGuard { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: self.resource_span.clone(), + lock: self, + } + }; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + let acquire_fut = trace::async_op( + move || acquire_fut, + resource_span, + "Mutex::lock_owned", + "poll", + false, + ); + + #[allow(clippy::let_and_return)] // this lint triggers when disabling tracing + let guard = acquire_fut.await; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + guard.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + locked = true, + ); + }); + + guard } async fn acquire(&self) { + crate::trace::async_trace_leaf().await; + self.s.acquire(1).await.unwrap_or_else(|_| { // The semaphore was closed. but, we never explicitly close it, and // we own it exclusively, which means that this can never happen. @@ -365,7 +657,23 @@ impl<T: ?Sized> Mutex<T> { /// ``` pub fn try_lock(&self) -> Result<MutexGuard<'_, T>, TryLockError> { match self.s.try_acquire(1) { - Ok(_) => Ok(MutexGuard { lock: self }), + Ok(_) => { + let guard = MutexGuard { + lock: self, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: self.resource_span.clone(), + }; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + locked = true, + ); + }); + + Ok(guard) + } Err(_) => Err(TryLockError(())), } } @@ -420,7 +728,23 @@ impl<T: ?Sized> Mutex<T> { /// # } pub fn try_lock_owned(self: Arc<Self>) -> Result<OwnedMutexGuard<T>, TryLockError> { match self.s.try_acquire(1) { - Ok(_) => Ok(OwnedMutexGuard { lock: self }), + Ok(_) => { + let guard = OwnedMutexGuard { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: self.resource_span.clone(), + lock: self, + }; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + guard.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + locked = true, + ); + }); + + Ok(guard) + } Err(_) => Err(TryLockError(())), } } @@ -462,14 +786,14 @@ where } } -impl<T> std::fmt::Debug for Mutex<T> +impl<T: ?Sized> std::fmt::Debug for Mutex<T> where T: std::fmt::Debug, { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let mut d = f.debug_struct("Mutex"); match self.try_lock() { - Ok(inner) => d.field("data", &*inner), + Ok(inner) => d.field("data", &&*inner), Err(_) => d.field("data", &format_args!("<locked>")), }; d.finish() @@ -479,6 +803,17 @@ where // === impl MutexGuard === impl<'a, T: ?Sized> MutexGuard<'a, T> { + fn skip_drop(self) -> MutexGuardInner<'a, T> { + let me = mem::ManuallyDrop::new(self); + // SAFETY: This duplicates the `resource_span` and then forgets the + // original. In the end, we have not duplicated or forgotten any values. + MutexGuardInner { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: unsafe { std::ptr::read(&me.resource_span) }, + lock: me.lock, + } + } + /// Makes a new [`MappedMutexGuard`] for a component of the locked data. /// /// This operation cannot fail as the [`MutexGuard`] passed in already locked the mutex. @@ -515,12 +850,13 @@ impl<'a, T: ?Sized> MutexGuard<'a, T> { F: FnOnce(&mut T) -> &mut U, { let data = f(&mut *this) as *mut U; - let s = &this.lock.s; - mem::forget(this); + let inner = this.skip_drop(); MappedMutexGuard { - s, + s: &inner.lock.s, data, - marker: marker::PhantomData, + marker: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: inner.resource_span, } } @@ -565,19 +901,54 @@ impl<'a, T: ?Sized> MutexGuard<'a, T> { Some(data) => data as *mut U, None => return Err(this), }; - let s = &this.lock.s; - mem::forget(this); + let inner = this.skip_drop(); Ok(MappedMutexGuard { - s, + s: &inner.lock.s, data, - marker: marker::PhantomData, + marker: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: inner.resource_span, }) } + + /// Returns a reference to the original `Mutex`. + /// + /// ``` + /// use tokio::sync::{Mutex, MutexGuard}; + /// + /// async fn unlock_and_relock<'l>(guard: MutexGuard<'l, u32>) -> MutexGuard<'l, u32> { + /// println!("1. contains: {:?}", *guard); + /// let mutex = MutexGuard::mutex(&guard); + /// drop(guard); + /// let guard = mutex.lock().await; + /// println!("2. contains: {:?}", *guard); + /// guard + /// } + /// # + /// # #[tokio::main] + /// # async fn main() { + /// # let mutex = Mutex::new(0u32); + /// # let guard = mutex.lock().await; + /// # let _guard = unlock_and_relock(guard).await; + /// # } + /// ``` + #[inline] + pub fn mutex(this: &Self) -> &'a Mutex<T> { + this.lock + } } impl<T: ?Sized> Drop for MutexGuard<'_, T> { fn drop(&mut self) { - self.lock.s.release(1) + self.lock.s.release(1); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + locked = false, + ); + }); } } @@ -608,9 +979,156 @@ impl<T: ?Sized + fmt::Display> fmt::Display for MutexGuard<'_, T> { // === impl OwnedMutexGuard === +impl<T: ?Sized> OwnedMutexGuard<T> { + fn skip_drop(self) -> OwnedMutexGuardInner<T> { + let me = mem::ManuallyDrop::new(self); + // SAFETY: This duplicates the values in every field of the guard, then + // forgets the originals, so in the end no value is duplicated. + unsafe { + OwnedMutexGuardInner { + lock: ptr::read(&me.lock), + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: ptr::read(&me.resource_span), + } + } + } + + /// Makes a new [`OwnedMappedMutexGuard`] for a component of the locked data. + /// + /// This operation cannot fail as the [`OwnedMutexGuard`] passed in already locked the mutex. + /// + /// This is an associated function that needs to be used as `OwnedMutexGuard::map(...)`. A method + /// would interfere with methods of the same name on the contents of the locked data. + /// + /// # Examples + /// + /// ``` + /// use tokio::sync::{Mutex, OwnedMutexGuard}; + /// use std::sync::Arc; + /// + /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] + /// struct Foo(u32); + /// + /// # #[tokio::main] + /// # async fn main() { + /// let foo = Arc::new(Mutex::new(Foo(1))); + /// + /// { + /// let mut mapped = OwnedMutexGuard::map(foo.clone().lock_owned().await, |f| &mut f.0); + /// *mapped = 2; + /// } + /// + /// assert_eq!(Foo(2), *foo.lock().await); + /// # } + /// ``` + /// + /// [`OwnedMutexGuard`]: struct@OwnedMutexGuard + /// [`OwnedMappedMutexGuard`]: struct@OwnedMappedMutexGuard + #[inline] + pub fn map<U, F>(mut this: Self, f: F) -> OwnedMappedMutexGuard<T, U> + where + F: FnOnce(&mut T) -> &mut U, + { + let data = f(&mut *this) as *mut U; + let inner = this.skip_drop(); + OwnedMappedMutexGuard { + data, + lock: inner.lock, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: inner.resource_span, + } + } + + /// Attempts to make a new [`OwnedMappedMutexGuard`] for a component of the locked data. The + /// original guard is returned if the closure returns `None`. + /// + /// This operation cannot fail as the [`OwnedMutexGuard`] passed in already locked the mutex. + /// + /// This is an associated function that needs to be used as `OwnedMutexGuard::try_map(...)`. A + /// method would interfere with methods of the same name on the contents of the locked data. + /// + /// # Examples + /// + /// ``` + /// use tokio::sync::{Mutex, OwnedMutexGuard}; + /// use std::sync::Arc; + /// + /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] + /// struct Foo(u32); + /// + /// # #[tokio::main] + /// # async fn main() { + /// let foo = Arc::new(Mutex::new(Foo(1))); + /// + /// { + /// let mut mapped = OwnedMutexGuard::try_map(foo.clone().lock_owned().await, |f| Some(&mut f.0)) + /// .expect("should not fail"); + /// *mapped = 2; + /// } + /// + /// assert_eq!(Foo(2), *foo.lock().await); + /// # } + /// ``` + /// + /// [`OwnedMutexGuard`]: struct@OwnedMutexGuard + /// [`OwnedMappedMutexGuard`]: struct@OwnedMappedMutexGuard + #[inline] + pub fn try_map<U, F>(mut this: Self, f: F) -> Result<OwnedMappedMutexGuard<T, U>, Self> + where + F: FnOnce(&mut T) -> Option<&mut U>, + { + let data = match f(&mut *this) { + Some(data) => data as *mut U, + None => return Err(this), + }; + let inner = this.skip_drop(); + Ok(OwnedMappedMutexGuard { + data, + lock: inner.lock, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: inner.resource_span, + }) + } + + /// Returns a reference to the original `Arc<Mutex>`. + /// + /// ``` + /// use std::sync::Arc; + /// use tokio::sync::{Mutex, OwnedMutexGuard}; + /// + /// async fn unlock_and_relock(guard: OwnedMutexGuard<u32>) -> OwnedMutexGuard<u32> { + /// println!("1. contains: {:?}", *guard); + /// let mutex: Arc<Mutex<u32>> = OwnedMutexGuard::mutex(&guard).clone(); + /// drop(guard); + /// let guard = mutex.lock_owned().await; + /// println!("2. contains: {:?}", *guard); + /// guard + /// } + /// # + /// # #[tokio::main] + /// # async fn main() { + /// # let mutex = Arc::new(Mutex::new(0u32)); + /// # let guard = mutex.lock_owned().await; + /// # unlock_and_relock(guard).await; + /// # } + /// ``` + #[inline] + pub fn mutex(this: &Self) -> &Arc<Mutex<T>> { + &this.lock + } +} + impl<T: ?Sized> Drop for OwnedMutexGuard<T> { fn drop(&mut self) { - self.lock.s.release(1) + self.lock.s.release(1); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + locked = false, + ); + }); } } @@ -642,6 +1160,16 @@ impl<T: ?Sized + fmt::Display> fmt::Display for OwnedMutexGuard<T> { // === impl MappedMutexGuard === impl<'a, T: ?Sized> MappedMutexGuard<'a, T> { + fn skip_drop(self) -> MappedMutexGuardInner<'a, T> { + let me = mem::ManuallyDrop::new(self); + MappedMutexGuardInner { + s: me.s, + data: me.data, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: unsafe { std::ptr::read(&me.resource_span) }, + } + } + /// Makes a new [`MappedMutexGuard`] for a component of the locked data. /// /// This operation cannot fail as the [`MappedMutexGuard`] passed in already locked the mutex. @@ -656,12 +1184,13 @@ impl<'a, T: ?Sized> MappedMutexGuard<'a, T> { F: FnOnce(&mut T) -> &mut U, { let data = f(&mut *this) as *mut U; - let s = this.s; - mem::forget(this); + let inner = this.skip_drop(); MappedMutexGuard { - s, + s: inner.s, data, - marker: marker::PhantomData, + marker: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: inner.resource_span, } } @@ -683,19 +1212,28 @@ impl<'a, T: ?Sized> MappedMutexGuard<'a, T> { Some(data) => data as *mut U, None => return Err(this), }; - let s = this.s; - mem::forget(this); + let inner = this.skip_drop(); Ok(MappedMutexGuard { - s, + s: inner.s, data, - marker: marker::PhantomData, + marker: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: inner.resource_span, }) } } impl<'a, T: ?Sized> Drop for MappedMutexGuard<'a, T> { fn drop(&mut self) { - self.s.release(1) + self.s.release(1); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + locked = false, + ); + }); } } @@ -723,3 +1261,111 @@ impl<'a, T: ?Sized + fmt::Display> fmt::Display for MappedMutexGuard<'a, T> { fmt::Display::fmt(&**self, f) } } + +// === impl OwnedMappedMutexGuard === + +impl<T: ?Sized, U: ?Sized> OwnedMappedMutexGuard<T, U> { + fn skip_drop(self) -> OwnedMappedMutexGuardInner<T, U> { + let me = mem::ManuallyDrop::new(self); + // SAFETY: This duplicates the values in every field of the guard, then + // forgets the originals, so in the end no value is duplicated. + unsafe { + OwnedMappedMutexGuardInner { + data: me.data, + lock: ptr::read(&me.lock), + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: ptr::read(&me.resource_span), + } + } + } + + /// Makes a new [`OwnedMappedMutexGuard`] for a component of the locked data. + /// + /// This operation cannot fail as the [`OwnedMappedMutexGuard`] passed in already locked the mutex. + /// + /// This is an associated function that needs to be used as `OwnedMappedMutexGuard::map(...)`. A method + /// would interfere with methods of the same name on the contents of the locked data. + /// + /// [`OwnedMappedMutexGuard`]: struct@OwnedMappedMutexGuard + #[inline] + pub fn map<S, F>(mut this: Self, f: F) -> OwnedMappedMutexGuard<T, S> + where + F: FnOnce(&mut U) -> &mut S, + { + let data = f(&mut *this) as *mut S; + let inner = this.skip_drop(); + OwnedMappedMutexGuard { + data, + lock: inner.lock, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: inner.resource_span, + } + } + + /// Attempts to make a new [`OwnedMappedMutexGuard`] for a component of the locked data. The + /// original guard is returned if the closure returns `None`. + /// + /// This operation cannot fail as the [`OwnedMutexGuard`] passed in already locked the mutex. + /// + /// This is an associated function that needs to be used as `OwnedMutexGuard::try_map(...)`. A + /// method would interfere with methods of the same name on the contents of the locked data. + /// + /// [`OwnedMutexGuard`]: struct@OwnedMutexGuard + /// [`OwnedMappedMutexGuard`]: struct@OwnedMappedMutexGuard + #[inline] + pub fn try_map<S, F>(mut this: Self, f: F) -> Result<OwnedMappedMutexGuard<T, S>, Self> + where + F: FnOnce(&mut U) -> Option<&mut S>, + { + let data = match f(&mut *this) { + Some(data) => data as *mut S, + None => return Err(this), + }; + let inner = this.skip_drop(); + Ok(OwnedMappedMutexGuard { + data, + lock: inner.lock, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: inner.resource_span, + }) + } +} + +impl<T: ?Sized, U: ?Sized> Drop for OwnedMappedMutexGuard<T, U> { + fn drop(&mut self) { + self.lock.s.release(1); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + locked = false, + ); + }); + } +} + +impl<T: ?Sized, U: ?Sized> Deref for OwnedMappedMutexGuard<T, U> { + type Target = U; + fn deref(&self) -> &Self::Target { + unsafe { &*self.data } + } +} + +impl<T: ?Sized, U: ?Sized> DerefMut for OwnedMappedMutexGuard<T, U> { + fn deref_mut(&mut self) -> &mut Self::Target { + unsafe { &mut *self.data } + } +} + +impl<T: ?Sized, U: ?Sized + fmt::Debug> fmt::Debug for OwnedMappedMutexGuard<T, U> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Debug::fmt(&**self, f) + } +} + +impl<T: ?Sized, U: ?Sized + fmt::Display> fmt::Display for OwnedMappedMutexGuard<T, U> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Display::fmt(&**self, f) + } +} diff --git a/vendor/tokio/src/sync/notify.rs b/vendor/tokio/src/sync/notify.rs index af7b9423a..0f104b71a 100644 --- a/vendor/tokio/src/sync/notify.rs +++ b/vendor/tokio/src/sync/notify.rs @@ -5,42 +5,46 @@ // triggers this warning but it is safe to ignore in this case. #![cfg_attr(not(feature = "sync"), allow(unreachable_pub, dead_code))] +use crate::loom::cell::UnsafeCell; use crate::loom::sync::atomic::AtomicUsize; use crate::loom::sync::Mutex; -use crate::util::linked_list::{self, LinkedList}; +use crate::util::linked_list::{self, GuardedLinkedList, LinkedList}; +use crate::util::WakeList; -use std::cell::UnsafeCell; use std::future::Future; use std::marker::PhantomPinned; +use std::panic::{RefUnwindSafe, UnwindSafe}; use std::pin::Pin; use std::ptr::NonNull; -use std::sync::atomic::Ordering::SeqCst; +use std::sync::atomic::Ordering::{self, Acquire, Relaxed, Release, SeqCst}; use std::task::{Context, Poll, Waker}; type WaitList = LinkedList<Waiter, <Waiter as linked_list::Link>::Target>; +type GuardedWaitList = GuardedLinkedList<Waiter, <Waiter as linked_list::Link>::Target>; -/// Notify a single task to wake up. +/// Notifies a single task to wake up. /// /// `Notify` provides a basic mechanism to notify a single task of an event. /// `Notify` itself does not carry any data. Instead, it is to be used to signal /// another task to perform an operation. /// -/// `Notify` can be thought of as a [`Semaphore`] starting with 0 permits. -/// [`notified().await`] waits for a permit to become available, and [`notify_one()`] -/// sets a permit **if there currently are no available permits**. +/// A `Notify` can be thought of as a [`Semaphore`] starting with 0 permits. The +/// [`notified().await`] method waits for a permit to become available, and +/// [`notify_one()`] sets a permit **if there currently are no available +/// permits**. /// /// The synchronization details of `Notify` are similar to /// [`thread::park`][park] and [`Thread::unpark`][unpark] from std. A [`Notify`] /// value contains a single permit. [`notified().await`] waits for the permit to -/// be made available, consumes the permit, and resumes. [`notify_one()`] sets the -/// permit, waking a pending task if there is one. +/// be made available, consumes the permit, and resumes. [`notify_one()`] sets +/// the permit, waking a pending task if there is one. /// -/// If `notify_one()` is called **before** `notified().await`, then the next call to -/// `notified().await` will complete immediately, consuming the permit. Any -/// subsequent calls to `notified().await` will wait for a new permit. +/// If `notify_one()` is called **before** `notified().await`, then the next +/// call to `notified().await` will complete immediately, consuming the permit. +/// Any subsequent calls to `notified().await` will wait for a new permit. /// -/// If `notify_one()` is called **multiple** times before `notified().await`, only a -/// **single** permit is stored. The next call to `notified().await` will +/// If `notify_one()` is called **multiple** times before `notified().await`, +/// only a **single** permit is stored. The next call to `notified().await` will /// complete immediately, but the one after will wait for a new permit. /// /// # Examples @@ -56,17 +60,24 @@ type WaitList = LinkedList<Waiter, <Waiter as linked_list::Link>::Target>; /// let notify = Arc::new(Notify::new()); /// let notify2 = notify.clone(); /// -/// tokio::spawn(async move { +/// let handle = tokio::spawn(async move { /// notify2.notified().await; /// println!("received notification"); /// }); /// /// println!("sending notification"); /// notify.notify_one(); +/// +/// // Wait for task to receive notification. +/// handle.await.unwrap(); /// } /// ``` /// -/// Unbound mpsc channel. +/// Unbound multi-producer single-consumer (mpsc) channel. +/// +/// No wakeups can be lost when using this channel because the call to +/// `notify_one()` will store a permit in the `Notify`, which the following call +/// to `notified()` will consume. /// /// ``` /// use tokio::sync::Notify; @@ -88,6 +99,8 @@ type WaitList = LinkedList<Waiter, <Waiter as linked_list::Link>::Target>; /// self.notify.notify_one(); /// } /// +/// // This is a single-consumer channel, so several concurrent calls to +/// // `recv` are not allowed. /// pub async fn recv(&self) -> T { /// loop { /// // Drain values @@ -102,45 +115,250 @@ type WaitList = LinkedList<Waiter, <Waiter as linked_list::Link>::Target>; /// } /// ``` /// +/// Unbound multi-producer multi-consumer (mpmc) channel. +/// +/// The call to [`enable`] is important because otherwise if you have two +/// calls to `recv` and two calls to `send` in parallel, the following could +/// happen: +/// +/// 1. Both calls to `try_recv` return `None`. +/// 2. Both new elements are added to the vector. +/// 3. The `notify_one` method is called twice, adding only a single +/// permit to the `Notify`. +/// 4. Both calls to `recv` reach the `Notified` future. One of them +/// consumes the permit, and the other sleeps forever. +/// +/// By adding the `Notified` futures to the list by calling `enable` before +/// `try_recv`, the `notify_one` calls in step three would remove the +/// futures from the list and mark them notified instead of adding a permit +/// to the `Notify`. This ensures that both futures are woken. +/// +/// Notice that this failure can only happen if there are two concurrent calls +/// to `recv`. This is why the mpsc example above does not require a call to +/// `enable`. +/// +/// ``` +/// use tokio::sync::Notify; +/// +/// use std::collections::VecDeque; +/// use std::sync::Mutex; +/// +/// struct Channel<T> { +/// messages: Mutex<VecDeque<T>>, +/// notify_on_sent: Notify, +/// } +/// +/// impl<T> Channel<T> { +/// pub fn send(&self, msg: T) { +/// let mut locked_queue = self.messages.lock().unwrap(); +/// locked_queue.push_back(msg); +/// drop(locked_queue); +/// +/// // Send a notification to one of the calls currently +/// // waiting in a call to `recv`. +/// self.notify_on_sent.notify_one(); +/// } +/// +/// pub fn try_recv(&self) -> Option<T> { +/// let mut locked_queue = self.messages.lock().unwrap(); +/// locked_queue.pop_front() +/// } +/// +/// pub async fn recv(&self) -> T { +/// let future = self.notify_on_sent.notified(); +/// tokio::pin!(future); +/// +/// loop { +/// // Make sure that no wakeup is lost if we get +/// // `None` from `try_recv`. +/// future.as_mut().enable(); +/// +/// if let Some(msg) = self.try_recv() { +/// return msg; +/// } +/// +/// // Wait for a call to `notify_one`. +/// // +/// // This uses `.as_mut()` to avoid consuming the future, +/// // which lets us call `Pin::set` below. +/// future.as_mut().await; +/// +/// // Reset the future in case another call to +/// // `try_recv` got the message before us. +/// future.set(self.notify_on_sent.notified()); +/// } +/// } +/// } +/// ``` +/// /// [park]: std::thread::park /// [unpark]: std::thread::Thread::unpark /// [`notified().await`]: Notify::notified() /// [`notify_one()`]: Notify::notify_one() +/// [`enable`]: Notified::enable() /// [`Semaphore`]: crate::sync::Semaphore #[derive(Debug)] pub struct Notify { - // This uses 2 bits to store one of `EMPTY`, + // `state` uses 2 bits to store one of `EMPTY`, // `WAITING` or `NOTIFIED`. The rest of the bits // are used to store the number of times `notify_waiters` // was called. + // + // Throughout the code there are two assumptions: + // - state can be transitioned *from* `WAITING` only if + // `waiters` lock is held + // - number of times `notify_waiters` was called can + // be modified only if `waiters` lock is held state: AtomicUsize, waiters: Mutex<WaitList>, } -#[derive(Debug, Clone, Copy)] -enum NotificationType { - // Notification triggered by calling `notify_waiters` - AllWaiters, - // Notification triggered by calling `notify_one` - OneWaiter, -} - #[derive(Debug)] struct Waiter { - /// Intrusive linked-list pointers + /// Intrusive linked-list pointers. pointers: linked_list::Pointers<Waiter>, - /// Waiting task's waker - waker: Option<Waker>, + /// Waiting task's waker. Depending on the value of `notification`, + /// this field is either protected by the `waiters` lock in + /// `Notify`, or it is exclusively owned by the enclosing `Waiter`. + waker: UnsafeCell<Option<Waker>>, - /// `true` if the notification has been assigned to this waiter. - notified: Option<NotificationType>, + /// Notification for this waiter. + /// * if it's `None`, then `waker` is protected by the `waiters` lock. + /// * if it's `Some`, then `waker` is exclusively owned by the + /// enclosing `Waiter` and can be accessed without locking. + notification: AtomicNotification, /// Should not be `Unpin`. _p: PhantomPinned, } -/// Future returned from [`Notify::notified()`] +impl Waiter { + fn new() -> Waiter { + Waiter { + pointers: linked_list::Pointers::new(), + waker: UnsafeCell::new(None), + notification: AtomicNotification::none(), + _p: PhantomPinned, + } + } +} + +generate_addr_of_methods! { + impl<> Waiter { + unsafe fn addr_of_pointers(self: NonNull<Self>) -> NonNull<linked_list::Pointers<Waiter>> { + &self.pointers + } + } +} + +// No notification. +const NOTIFICATION_NONE: usize = 0; + +// Notification type used by `notify_one`. +const NOTIFICATION_ONE: usize = 1; + +// Notification type used by `notify_waiters`. +const NOTIFICATION_ALL: usize = 2; + +/// Notification for a `Waiter`. +/// This struct is equivalent to `Option<Notification>`, but uses +/// `AtomicUsize` inside for atomic operations. +#[derive(Debug)] +struct AtomicNotification(AtomicUsize); + +impl AtomicNotification { + fn none() -> Self { + AtomicNotification(AtomicUsize::new(NOTIFICATION_NONE)) + } + + /// Store-release a notification. + /// This method should be called exactly once. + fn store_release(&self, notification: Notification) { + self.0.store(notification as usize, Release); + } + + fn load(&self, ordering: Ordering) -> Option<Notification> { + match self.0.load(ordering) { + NOTIFICATION_NONE => None, + NOTIFICATION_ONE => Some(Notification::One), + NOTIFICATION_ALL => Some(Notification::All), + _ => unreachable!(), + } + } + + /// Clears the notification. + /// This method is used by a `Notified` future to consume the + /// notification. It uses relaxed ordering and should be only + /// used once the atomic notification is no longer shared. + fn clear(&self) { + self.0.store(NOTIFICATION_NONE, Relaxed); + } +} + +#[derive(Debug, PartialEq, Eq)] +#[repr(usize)] +enum Notification { + One = NOTIFICATION_ONE, + All = NOTIFICATION_ALL, +} + +/// List used in `Notify::notify_waiters`. It wraps a guarded linked list +/// and gates the access to it on `notify.waiters` mutex. It also empties +/// the list on drop. +struct NotifyWaitersList<'a> { + list: GuardedWaitList, + is_empty: bool, + notify: &'a Notify, +} + +impl<'a> NotifyWaitersList<'a> { + fn new( + unguarded_list: WaitList, + guard: Pin<&'a Waiter>, + notify: &'a Notify, + ) -> NotifyWaitersList<'a> { + let guard_ptr = NonNull::from(guard.get_ref()); + let list = unguarded_list.into_guarded(guard_ptr); + NotifyWaitersList { + list, + is_empty: false, + notify, + } + } + + /// Removes the last element from the guarded list. Modifying this list + /// requires an exclusive access to the main list in `Notify`. + fn pop_back_locked(&mut self, _waiters: &mut WaitList) -> Option<NonNull<Waiter>> { + let result = self.list.pop_back(); + if result.is_none() { + // Save information about emptiness to avoid waiting for lock + // in the destructor. + self.is_empty = true; + } + result + } +} + +impl Drop for NotifyWaitersList<'_> { + fn drop(&mut self) { + // If the list is not empty, we unlink all waiters from it. + // We do not wake the waiters to avoid double panics. + if !self.is_empty { + let _lock_guard = self.notify.waiters.lock(); + while let Some(waiter) = self.list.pop_back() { + // Safety: we never make mutable references to waiters. + let waiter = unsafe { waiter.as_ref() }; + waiter.notification.store_release(Notification::All); + } + } + } +} + +/// Future returned from [`Notify::notified()`]. +/// +/// This future is fused, so once it has completed, any future calls to poll +/// will immediately return `Poll::Ready`. #[derive(Debug)] pub struct Notified<'a> { /// The `Notify` being received on. @@ -149,8 +367,11 @@ pub struct Notified<'a> { /// The current state of the receiving process. state: State, + /// Number of calls to `notify_waiters` at the time of creation. + notify_waiters_calls: usize, + /// Entry in the waiter `LinkedList`. - waiter: UnsafeCell<Waiter>, + waiter: Waiter, } unsafe impl<'a> Send for Notified<'a> {} @@ -158,7 +379,7 @@ unsafe impl<'a> Sync for Notified<'a> {} #[derive(Debug)] enum State { - Init(usize), + Init, Waiting, Done, } @@ -167,13 +388,13 @@ const NOTIFY_WAITERS_SHIFT: usize = 2; const STATE_MASK: usize = (1 << NOTIFY_WAITERS_SHIFT) - 1; const NOTIFY_WAITERS_CALLS_MASK: usize = !STATE_MASK; -/// Initial "idle" state +/// Initial "idle" state. const EMPTY: usize = 0; /// One or more threads are currently waiting to be notified. const WAITING: usize = 1; -/// Pending notification +/// Pending notification. const NOTIFIED: usize = 2; fn set_state(data: usize, state: usize) -> usize { @@ -244,7 +465,16 @@ impl Notify { /// immediately, consuming that permit. Otherwise, `notified().await` waits /// for a permit to be made available by the next call to `notify_one()`. /// + /// The `Notified` future is not guaranteed to receive wakeups from calls to + /// `notify_one()` if it has not yet been polled. See the documentation for + /// [`Notified::enable()`] for more details. + /// + /// The `Notified` future is guaranteed to receive wakeups from + /// `notify_waiters()` as soon as it has been created, even if it has not + /// yet been polled. + /// /// [`notify_one()`]: Notify::notify_one + /// [`Notified::enable()`]: Notified::enable /// /// # Cancel safety /// @@ -274,21 +504,17 @@ impl Notify { /// ``` pub fn notified(&self) -> Notified<'_> { // we load the number of times notify_waiters - // was called and store that in our initial state + // was called and store that in the future. let state = self.state.load(SeqCst); Notified { notify: self, - state: State::Init(state >> NOTIFY_WAITERS_SHIFT), - waiter: UnsafeCell::new(Waiter { - pointers: linked_list::Pointers::new(), - waker: None, - notified: None, - _p: PhantomPinned, - }), + state: State::Init, + notify_waiters_calls: get_num_notify_waiters_calls(state), + waiter: Waiter::new(), } } - /// Notifies a waiting task + /// Notifies a waiting task. /// /// If a task is currently waiting, that task is notified. Otherwise, a /// permit is stored in this `Notify` value and the **next** call to @@ -358,7 +584,7 @@ impl Notify { } } - /// Notifies all waiting tasks + /// Notifies all waiting tasks. /// /// If a task is currently waiting, that task is notified. Unlike with /// `notify_one()`, no permit is stored to be used by the next call to @@ -391,43 +617,56 @@ impl Notify { /// } /// ``` pub fn notify_waiters(&self) { - const NUM_WAKERS: usize = 32; - - let mut wakers: [Option<Waker>; NUM_WAKERS] = Default::default(); - let mut curr_waker = 0; - - // There are waiters, the lock must be acquired to notify. let mut waiters = self.waiters.lock(); - // The state must be reloaded while the lock is held. The state may only + // The state must be loaded while the lock is held. The state may only // transition out of WAITING while the lock is held. let curr = self.state.load(SeqCst); - if let EMPTY | NOTIFIED = get_state(curr) { + if matches!(get_state(curr), EMPTY | NOTIFIED) { // There are no waiting tasks. All we need to do is increment the // number of times this method was called. atomic_inc_num_notify_waiters_calls(&self.state); return; } - // At this point, it is guaranteed that the state will not - // concurrently change, as holding the lock is required to - // transition **out** of `WAITING`. + // Increment the number of times this method was called + // and transition to empty. + let new_state = set_state(inc_num_notify_waiters_calls(curr), EMPTY); + self.state.store(new_state, SeqCst); + + // It is critical for `GuardedLinkedList` safety that the guard node is + // pinned in memory and is not dropped until the guarded list is dropped. + let guard = Waiter::new(); + pin!(guard); + + // We move all waiters to a secondary list. It uses a `GuardedLinkedList` + // underneath to allow every waiter to safely remove itself from it. + // + // * This list will be still guarded by the `waiters` lock. + // `NotifyWaitersList` wrapper makes sure we hold the lock to modify it. + // * This wrapper will empty the list on drop. It is critical for safety + // that we will not leave any list entry with a pointer to the local + // guard node after this function returns / panics. + let mut list = NotifyWaitersList::new(std::mem::take(&mut *waiters), guard.as_ref(), self); + + let mut wakers = WakeList::new(); 'outer: loop { - while curr_waker < NUM_WAKERS { - match waiters.pop_back() { - Some(mut waiter) => { - // Safety: `waiters` lock is still held. - let waiter = unsafe { waiter.as_mut() }; - - assert!(waiter.notified.is_none()); - - waiter.notified = Some(NotificationType::AllWaiters); - - if let Some(waker) = waiter.waker.take() { - wakers[curr_waker] = Some(waker); - curr_waker += 1; + while wakers.can_push() { + match list.pop_back_locked(&mut waiters) { + Some(waiter) => { + // Safety: we never make mutable references to waiters. + let waiter = unsafe { waiter.as_ref() }; + + // Safety: we hold the lock, so we can access the waker. + if let Some(waker) = + unsafe { waiter.waker.with_mut(|waker| (*waker).take()) } + { + wakers.push(waker); } + + // This waiter is unlinked and will not be shared ever again, release it. + waiter.notification.store_release(Notification::All); } None => { break 'outer; @@ -435,30 +674,21 @@ impl Notify { } } + // Release the lock before notifying. drop(waiters); - for waker in wakers.iter_mut().take(curr_waker) { - waker.take().unwrap().wake(); - } - - curr_waker = 0; + // One of the wakers may panic, but the remaining waiters will still + // be unlinked from the list in `NotifyWaitersList` destructor. + wakers.wake_all(); // Acquire the lock again. waiters = self.waiters.lock(); } - // All waiters will be notified, the state must be transitioned to - // `EMPTY`. As transitioning **from** `WAITING` requires the lock to be - // held, a `store` is sufficient. - let new = set_state(inc_num_notify_waiters_calls(curr), EMPTY); - self.state.store(new, SeqCst); - // Release the lock before notifying drop(waiters); - for waker in wakers.iter_mut().take(curr_waker) { - waker.take().unwrap().wake(); - } + wakers.wake_all(); } } @@ -468,6 +698,9 @@ impl Default for Notify { } } +impl UnwindSafe for Notify {} +impl RefUnwindSafe for Notify {} + fn notify_locked(waiters: &mut WaitList, state: &AtomicUsize, curr: usize) -> Option<Waker> { loop { match get_state(curr) { @@ -490,15 +723,16 @@ fn notify_locked(waiters: &mut WaitList, state: &AtomicUsize, curr: usize) -> Op // transition **out** of `WAITING`. // // Get a pending waiter - let mut waiter = waiters.pop_back().unwrap(); + let waiter = waiters.pop_back().unwrap(); - // Safety: `waiters` lock is still held. - let waiter = unsafe { waiter.as_mut() }; + // Safety: we never make mutable references to waiters. + let waiter = unsafe { waiter.as_ref() }; - assert!(waiter.notified.is_none()); + // Safety: we hold the lock, so we can access the waker. + let waker = unsafe { waiter.waker.with_mut(|waker| (*waker).take()) }; - waiter.notified = Some(NotificationType::OneWaiter); - let waker = waiter.waker.take(); + // This waiter is unlinked and will not be shared ever again, release it. + waiter.notification.store_release(Notification::One); if waiters.is_empty() { // As this the **final** waiter in the list, the state @@ -518,32 +752,142 @@ fn notify_locked(waiters: &mut WaitList, state: &AtomicUsize, curr: usize) -> Op // ===== impl Notified ===== impl Notified<'_> { + /// Adds this future to the list of futures that are ready to receive + /// wakeups from calls to [`notify_one`]. + /// + /// Polling the future also adds it to the list, so this method should only + /// be used if you want to add the future to the list before the first call + /// to `poll`. (In fact, this method is equivalent to calling `poll` except + /// that no `Waker` is registered.) + /// + /// This has no effect on notifications sent using [`notify_waiters`], which + /// are received as long as they happen after the creation of the `Notified` + /// regardless of whether `enable` or `poll` has been called. + /// + /// This method returns true if the `Notified` is ready. This happens in the + /// following situations: + /// + /// 1. The `notify_waiters` method was called between the creation of the + /// `Notified` and the call to this method. + /// 2. This is the first call to `enable` or `poll` on this future, and the + /// `Notify` was holding a permit from a previous call to `notify_one`. + /// The call consumes the permit in that case. + /// 3. The future has previously been enabled or polled, and it has since + /// then been marked ready by either consuming a permit from the + /// `Notify`, or by a call to `notify_one` or `notify_waiters` that + /// removed it from the list of futures ready to receive wakeups. + /// + /// If this method returns true, any future calls to poll on the same future + /// will immediately return `Poll::Ready`. + /// + /// # Examples + /// + /// Unbound multi-producer multi-consumer (mpmc) channel. + /// + /// The call to `enable` is important because otherwise if you have two + /// calls to `recv` and two calls to `send` in parallel, the following could + /// happen: + /// + /// 1. Both calls to `try_recv` return `None`. + /// 2. Both new elements are added to the vector. + /// 3. The `notify_one` method is called twice, adding only a single + /// permit to the `Notify`. + /// 4. Both calls to `recv` reach the `Notified` future. One of them + /// consumes the permit, and the other sleeps forever. + /// + /// By adding the `Notified` futures to the list by calling `enable` before + /// `try_recv`, the `notify_one` calls in step three would remove the + /// futures from the list and mark them notified instead of adding a permit + /// to the `Notify`. This ensures that both futures are woken. + /// + /// ``` + /// use tokio::sync::Notify; + /// + /// use std::collections::VecDeque; + /// use std::sync::Mutex; + /// + /// struct Channel<T> { + /// messages: Mutex<VecDeque<T>>, + /// notify_on_sent: Notify, + /// } + /// + /// impl<T> Channel<T> { + /// pub fn send(&self, msg: T) { + /// let mut locked_queue = self.messages.lock().unwrap(); + /// locked_queue.push_back(msg); + /// drop(locked_queue); + /// + /// // Send a notification to one of the calls currently + /// // waiting in a call to `recv`. + /// self.notify_on_sent.notify_one(); + /// } + /// + /// pub fn try_recv(&self) -> Option<T> { + /// let mut locked_queue = self.messages.lock().unwrap(); + /// locked_queue.pop_front() + /// } + /// + /// pub async fn recv(&self) -> T { + /// let future = self.notify_on_sent.notified(); + /// tokio::pin!(future); + /// + /// loop { + /// // Make sure that no wakeup is lost if we get + /// // `None` from `try_recv`. + /// future.as_mut().enable(); + /// + /// if let Some(msg) = self.try_recv() { + /// return msg; + /// } + /// + /// // Wait for a call to `notify_one`. + /// // + /// // This uses `.as_mut()` to avoid consuming the future, + /// // which lets us call `Pin::set` below. + /// future.as_mut().await; + /// + /// // Reset the future in case another call to + /// // `try_recv` got the message before us. + /// future.set(self.notify_on_sent.notified()); + /// } + /// } + /// } + /// ``` + /// + /// [`notify_one`]: Notify::notify_one() + /// [`notify_waiters`]: Notify::notify_waiters() + pub fn enable(self: Pin<&mut Self>) -> bool { + self.poll_notified(None).is_ready() + } + /// A custom `project` implementation is used in place of `pin-project-lite` /// as a custom drop implementation is needed. - fn project(self: Pin<&mut Self>) -> (&Notify, &mut State, &UnsafeCell<Waiter>) { + fn project(self: Pin<&mut Self>) -> (&Notify, &mut State, &usize, &Waiter) { unsafe { - // Safety: both `notify` and `state` are `Unpin`. + // Safety: `notify`, `state` and `notify_waiters_calls` are `Unpin`. is_unpin::<&Notify>(); - is_unpin::<AtomicUsize>(); + is_unpin::<State>(); + is_unpin::<usize>(); let me = self.get_unchecked_mut(); - (&me.notify, &mut me.state, &me.waiter) + ( + me.notify, + &mut me.state, + &me.notify_waiters_calls, + &me.waiter, + ) } } -} - -impl Future for Notified<'_> { - type Output = (); - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { + fn poll_notified(self: Pin<&mut Self>, waker: Option<&Waker>) -> Poll<()> { use State::*; - let (notify, state, waiter) = self.project(); + let (notify, state, notify_waiters_calls, waiter) = self.project(); - loop { + 'outer_loop: loop { match *state { - Init(initial_notify_waiters_calls) => { + Init => { let curr = notify.state.load(SeqCst); // Optimistically try acquiring a pending notification @@ -557,9 +901,13 @@ impl Future for Notified<'_> { if res.is_ok() { // Acquired the notification *state = Done; - return Poll::Ready(()); + continue 'outer_loop; } + // Clone the waker before locking, a waker clone can be + // triggering arbitrary code. + let waker = waker.cloned(); + // Acquire the lock and attempt to transition to the waiting // state. let mut waiters = notify.waiters.lock(); @@ -569,9 +917,9 @@ impl Future for Notified<'_> { // if notify_waiters has been called after the future // was created, then we are done - if get_num_notify_waiters_calls(curr) != initial_notify_waiters_calls { + if get_num_notify_waiters_calls(curr) != *notify_waiters_calls { *state = Done; - return Poll::Ready(()); + continue 'outer_loop; } // Transition the state to WAITING. @@ -607,7 +955,7 @@ impl Future for Notified<'_> { Ok(_) => { // Acquired the notification *state = Done; - return Poll::Ready(()); + continue 'outer_loop; } Err(actual) => { assert_eq!(get_state(actual), EMPTY); @@ -619,44 +967,109 @@ impl Future for Notified<'_> { } } - // Safety: called while locked. - unsafe { - (*waiter.get()).waker = Some(cx.waker().clone()); + let mut old_waker = None; + if waker.is_some() { + // Safety: called while locked. + // + // The use of `old_waiter` here is not necessary, as the field is always + // None when we reach this line. + unsafe { + old_waker = + waiter.waker.with_mut(|v| std::mem::replace(&mut *v, waker)); + } } // Insert the waiter into the linked list - // - // safety: pointers from `UnsafeCell` are never null. - waiters.push_front(unsafe { NonNull::new_unchecked(waiter.get()) }); + waiters.push_front(NonNull::from(waiter)); *state = Waiting; + drop(waiters); + drop(old_waker); + return Poll::Pending; } Waiting => { - // Currently in the "Waiting" state, implying the caller has - // a waiter stored in the waiter list (guarded by - // `notify.waiters`). In order to access the waker fields, - // we must hold the lock. + #[cfg(tokio_taskdump)] + if let Some(waker) = waker { + let mut ctx = Context::from_waker(waker); + ready!(crate::trace::trace_leaf(&mut ctx)); + } + + if waiter.notification.load(Acquire).is_some() { + // Safety: waiter is already unlinked and will not be shared again, + // so we have an exclusive access to `waker`. + drop(unsafe { waiter.waker.with_mut(|waker| (*waker).take()) }); + + waiter.notification.clear(); + *state = Done; + return Poll::Ready(()); + } + + // Our waiter was not notified, implying it is still stored in a waiter + // list (guarded by `notify.waiters`). In order to access the waker + // fields, we must acquire the lock. + + let mut old_waker = None; + let mut waiters = notify.waiters.lock(); + + // We hold the lock and notifications are set only with the lock held, + // so this can be relaxed, because the happens-before relationship is + // established through the mutex. + if waiter.notification.load(Relaxed).is_some() { + // Safety: waiter is already unlinked and will not be shared again, + // so we have an exclusive access to `waker`. + old_waker = unsafe { waiter.waker.with_mut(|waker| (*waker).take()) }; + + waiter.notification.clear(); + + // Drop the old waker after releasing the lock. + drop(waiters); + drop(old_waker); - let waiters = notify.waiters.lock(); + *state = Done; + return Poll::Ready(()); + } + + // Load the state with the lock held. + let curr = notify.state.load(SeqCst); + + if get_num_notify_waiters_calls(curr) != *notify_waiters_calls { + // Before we add a waiter to the list we check if these numbers are + // different while holding the lock. If these numbers are different now, + // it means that there is a call to `notify_waiters` in progress and this + // waiter must be contained by a guarded list used in `notify_waiters`. + // We can treat the waiter as notified and remove it from the list, as + // it would have been notified in the `notify_waiters` call anyways. - // Safety: called while locked - let w = unsafe { &mut *waiter.get() }; + // Safety: we hold the lock, so we can modify the waker. + old_waker = unsafe { waiter.waker.with_mut(|waker| (*waker).take()) }; - if w.notified.is_some() { - // Our waker has been notified. Reset the fields and - // remove it from the list. - w.waker = None; - w.notified = None; + // Safety: we hold the lock, so we have an exclusive access to the list. + // The list is used in `notify_waiters`, so it must be guarded. + unsafe { waiters.remove(NonNull::from(waiter)) }; *state = Done; } else { - // Update the waker, if necessary. - if !w.waker.as_ref().unwrap().will_wake(cx.waker()) { - w.waker = Some(cx.waker().clone()); + // Safety: we hold the lock, so we can modify the waker. + unsafe { + waiter.waker.with_mut(|v| { + if let Some(waker) = waker { + let should_update = match &*v { + Some(current_waker) => !current_waker.will_wake(waker), + None => true, + }; + if should_update { + old_waker = std::mem::replace(&mut *v, Some(waker.clone())); + } + } + }); } + // Drop the old waker after releasing the lock. + drop(waiters); + drop(old_waker); + return Poll::Pending; } @@ -666,8 +1079,16 @@ impl Future for Notified<'_> { // is helpful to visualize the scope of the critical // section. drop(waiters); + + // Drop the old waker after releasing the lock. + drop(old_waker); } Done => { + #[cfg(tokio_taskdump)] + if let Some(waker) = waker { + let mut ctx = Context::from_waker(waker); + ready!(crate::trace::trace_leaf(&mut ctx)); + } return Poll::Ready(()); } } @@ -675,40 +1096,48 @@ impl Future for Notified<'_> { } } +impl Future for Notified<'_> { + type Output = (); + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { + self.poll_notified(Some(cx.waker())) + } +} + impl Drop for Notified<'_> { fn drop(&mut self) { use State::*; // Safety: The type only transitions to a "Waiting" state when pinned. - let (notify, state, waiter) = unsafe { Pin::new_unchecked(self).project() }; + let (notify, state, _, waiter) = unsafe { Pin::new_unchecked(self).project() }; // This is where we ensure safety. The `Notified` value is being // dropped, which means we must ensure that the waiter entry is no // longer stored in the linked list. - if let Waiting = *state { + if matches!(*state, Waiting) { let mut waiters = notify.waiters.lock(); let mut notify_state = notify.state.load(SeqCst); + // We hold the lock, so this field is not concurrently accessed by + // `notify_*` functions and we can use the relaxed ordering. + let notification = waiter.notification.load(Relaxed); + // remove the entry from the list (if not already removed) // - // safety: the waiter is only added to `waiters` by virtue of it - // being the only `LinkedList` available to the type. - unsafe { waiters.remove(NonNull::new_unchecked(waiter.get())) }; - - if waiters.is_empty() { - if let WAITING = get_state(notify_state) { - notify_state = set_state(notify_state, EMPTY); - notify.state.store(notify_state, SeqCst); - } + // Safety: we hold the lock, so we have an exclusive access to every list the + // waiter may be contained in. If the node is not contained in the `waiters` + // list, then it is contained by a guarded list used by `notify_waiters`. + unsafe { waiters.remove(NonNull::from(waiter)) }; + + if waiters.is_empty() && get_state(notify_state) == WAITING { + notify_state = set_state(notify_state, EMPTY); + notify.state.store(notify_state, SeqCst); } // See if the node was notified but not received. In this case, if // the notification was triggered via `notify_one`, it must be sent // to the next waiter. - // - // Safety: with the entry removed from the linked list, there can be - // no concurrent access to the entry - if let Some(NotificationType::OneWaiter) = unsafe { (*waiter.get()).notified } { + if notification == Some(Notification::One) { if let Some(waker) = notify_locked(&mut waiters, ¬ify.state, notify_state) { drop(waiters); waker.wake(); @@ -733,8 +1162,8 @@ unsafe impl linked_list::Link for Waiter { ptr } - unsafe fn pointers(mut target: NonNull<Waiter>) -> NonNull<linked_list::Pointers<Waiter>> { - NonNull::from(&mut target.as_mut().pointers) + unsafe fn pointers(target: NonNull<Waiter>) -> NonNull<linked_list::Pointers<Waiter>> { + Waiter::addr_of_pointers(target) } } diff --git a/vendor/tokio/src/sync/once_cell.rs b/vendor/tokio/src/sync/once_cell.rs index ce55d9e35..90ea5cd68 100644 --- a/vendor/tokio/src/sync/once_cell.rs +++ b/vendor/tokio/src/sync/once_cell.rs @@ -1,4 +1,4 @@ -use super::Semaphore; +use super::{Semaphore, SemaphorePermit, TryAcquireError}; use crate::loom::cell::UnsafeCell; use std::error::Error; use std::fmt; @@ -8,15 +8,30 @@ use std::ops::Drop; use std::ptr; use std::sync::atomic::{AtomicBool, Ordering}; -/// A thread-safe cell which can be written to only once. +// This file contains an implementation of an OnceCell. The principle +// behind the safety the of the cell is that any thread with an `&OnceCell` may +// access the `value` field according the following rules: +// +// 1. When `value_set` is false, the `value` field may be modified by the +// thread holding the permit on the semaphore. +// 2. When `value_set` is true, the `value` field may be accessed immutably by +// any thread. +// +// It is an invariant that if the semaphore is closed, then `value_set` is true. +// The reverse does not necessarily hold — but if not, the semaphore may not +// have any available permits. +// +// A thread with a `&mut OnceCell` may modify the value in any way it wants as +// long as the invariants are upheld. + +/// A thread-safe cell that can be written to only once. /// -/// Provides the functionality to either set the value, in case `OnceCell` -/// is uninitialized, or get the already initialized value by using an async -/// function via [`OnceCell::get_or_init`]. -/// -/// [`OnceCell::get_or_init`]: crate::sync::OnceCell::get_or_init +/// A `OnceCell` is typically used for global variables that need to be +/// initialized once on first use, but need no further changes. The `OnceCell` +/// in Tokio allows the initialization procedure to be asynchronous. /// /// # Examples +/// /// ``` /// use tokio::sync::OnceCell; /// @@ -28,8 +43,28 @@ use std::sync::atomic::{AtomicBool, Ordering}; /// /// #[tokio::main] /// async fn main() { -/// let result1 = ONCE.get_or_init(some_computation).await; -/// assert_eq!(*result1, 2); +/// let result = ONCE.get_or_init(some_computation).await; +/// assert_eq!(*result, 2); +/// } +/// ``` +/// +/// It is often useful to write a wrapper method for accessing the value. +/// +/// ``` +/// use tokio::sync::OnceCell; +/// +/// static ONCE: OnceCell<u32> = OnceCell::const_new(); +/// +/// async fn get_global_integer() -> &'static u32 { +/// ONCE.get_or_init(|| async { +/// 1 + 1 +/// }).await +/// } +/// +/// #[tokio::main] +/// async fn main() { +/// let result = get_global_integer().await; +/// assert_eq!(*result, 2); /// } /// ``` pub struct OnceCell<T> { @@ -68,10 +103,10 @@ impl<T: Eq> Eq for OnceCell<T> {} impl<T> Drop for OnceCell<T> { fn drop(&mut self) { - if self.initialized() { + if self.initialized_mut() { unsafe { self.value - .with_mut(|ptr| ptr::drop_in_place((&mut *ptr).as_mut_ptr())); + .with_mut(|ptr| ptr::drop_in_place((*ptr).as_mut_ptr())); }; } } @@ -90,7 +125,7 @@ impl<T> From<T> for OnceCell<T> { } impl<T> OnceCell<T> { - /// Creates a new uninitialized OnceCell instance. + /// Creates a new empty `OnceCell` instance. pub fn new() -> Self { OnceCell { value_set: AtomicBool::new(false), @@ -99,8 +134,9 @@ impl<T> OnceCell<T> { } } - /// Creates a new initialized OnceCell instance if `value` is `Some`, otherwise - /// has the same functionality as [`OnceCell::new`]. + /// Creates a new `OnceCell` that contains the provided value, if any. + /// + /// If the `Option` is `None`, this is equivalent to `OnceCell::new`. /// /// [`OnceCell::new`]: crate::sync::OnceCell::new pub fn new_with(value: Option<T>) -> Self { @@ -111,8 +147,31 @@ impl<T> OnceCell<T> { } } - /// Creates a new uninitialized OnceCell instance. - #[cfg(all(feature = "parking_lot", not(all(loom, test)),))] + /// Creates a new empty `OnceCell` instance. + /// + /// Equivalent to `OnceCell::new`, except that it can be used in static + /// variables. + /// + /// # Example + /// + /// ``` + /// use tokio::sync::OnceCell; + /// + /// static ONCE: OnceCell<u32> = OnceCell::const_new(); + /// + /// async fn get_global_integer() -> &'static u32 { + /// ONCE.get_or_init(|| async { + /// 1 + 1 + /// }).await + /// } + /// + /// #[tokio::main] + /// async fn main() { + /// let result = get_global_integer().await; + /// assert_eq!(*result, 2); + /// } + /// ``` + #[cfg(all(feature = "parking_lot", not(all(loom, test))))] #[cfg_attr(docsrs, doc(cfg(feature = "parking_lot")))] pub const fn const_new() -> Self { OnceCell { @@ -122,33 +181,48 @@ impl<T> OnceCell<T> { } } - /// Whether the value of the OnceCell is set or not. + /// Returns `true` if the `OnceCell` currently contains a value, and `false` + /// otherwise. pub fn initialized(&self) -> bool { + // Using acquire ordering so any threads that read a true from this + // atomic is able to read the value. self.value_set.load(Ordering::Acquire) } - // SAFETY: safe to call only once self.initialized() is true + /// Returns `true` if the `OnceCell` currently contains a value, and `false` + /// otherwise. + fn initialized_mut(&mut self) -> bool { + *self.value_set.get_mut() + } + + // SAFETY: The OnceCell must not be empty. unsafe fn get_unchecked(&self) -> &T { &*self.value.with(|ptr| (*ptr).as_ptr()) } - // SAFETY: safe to call only once self.initialized() is true. Safe because - // because of the mutable reference. + // SAFETY: The OnceCell must not be empty. unsafe fn get_unchecked_mut(&mut self) -> &mut T { &mut *self.value.with_mut(|ptr| (*ptr).as_mut_ptr()) } - // SAFETY: safe to call only once a permit on the semaphore has been - // acquired - unsafe fn set_value(&self, value: T) { - self.value.with_mut(|ptr| (*ptr).as_mut_ptr().write(value)); + fn set_value(&self, value: T, permit: SemaphorePermit<'_>) -> &T { + // SAFETY: We are holding the only permit on the semaphore. + unsafe { + self.value.with_mut(|ptr| (*ptr).as_mut_ptr().write(value)); + } + + // Using release ordering so any threads that read a true from this + // atomic is able to read the value we just stored. self.value_set.store(true, Ordering::Release); self.semaphore.close(); + permit.forget(); + + // SAFETY: We just initialized the cell. + unsafe { self.get_unchecked() } } - /// Tries to get a reference to the value of the OnceCell. - /// - /// Returns None if the value of the OnceCell hasn't previously been initialized. + /// Returns a reference to the value currently stored in the `OnceCell`, or + /// `None` if the `OnceCell` is empty. pub fn get(&self) -> Option<&T> { if self.initialized() { Some(unsafe { self.get_unchecked() }) @@ -157,179 +231,165 @@ impl<T> OnceCell<T> { } } - /// Tries to return a mutable reference to the value of the cell. + /// Returns a mutable reference to the value currently stored in the + /// `OnceCell`, or `None` if the `OnceCell` is empty. /// - /// Returns None if the cell hasn't previously been initialized. + /// Since this call borrows the `OnceCell` mutably, it is safe to mutate the + /// value inside the `OnceCell` — the mutable borrow statically guarantees + /// no other references exist. pub fn get_mut(&mut self) -> Option<&mut T> { - if self.initialized() { + if self.initialized_mut() { Some(unsafe { self.get_unchecked_mut() }) } else { None } } - /// Sets the value of the OnceCell to the argument value. + /// Sets the value of the `OnceCell` to the given value if the `OnceCell` is + /// empty. + /// + /// If the `OnceCell` already has a value, this call will fail with an + /// [`SetError::AlreadyInitializedError`]. /// - /// If the value of the OnceCell was already set prior to this call - /// then [`SetError::AlreadyInitializedError`] is returned. If another thread - /// is initializing the cell while this method is called, - /// [`SetError::InitializingError`] is returned. In order to wait - /// for an ongoing initialization to finish, call - /// [`OnceCell::get_or_init`] instead. + /// If the `OnceCell` is empty, but some other task is currently trying to + /// set the value, this call will fail with [`SetError::InitializingError`]. /// /// [`SetError::AlreadyInitializedError`]: crate::sync::SetError::AlreadyInitializedError /// [`SetError::InitializingError`]: crate::sync::SetError::InitializingError - /// ['OnceCell::get_or_init`]: crate::sync::OnceCell::get_or_init pub fn set(&self, value: T) -> Result<(), SetError<T>> { - if !self.initialized() { - // Another thread might be initializing the cell, in which case `try_acquire` will - // return an error - match self.semaphore.try_acquire() { - Ok(_permit) => { - if !self.initialized() { - // SAFETY: There is only one permit on the semaphore, hence only one - // mutable reference is created - unsafe { self.set_value(value) }; - - return Ok(()); - } else { - unreachable!( - "acquired the permit after OnceCell value was already initialized." - ); - } - } - _ => { - // Couldn't acquire the permit, look if initializing process is already completed - if !self.initialized() { - return Err(SetError::InitializingError(value)); - } - } - } + if self.initialized() { + return Err(SetError::AlreadyInitializedError(value)); } - Err(SetError::AlreadyInitializedError(value)) + // Another task might be initializing the cell, in which case + // `try_acquire` will return an error. If we succeed to acquire the + // permit, then we can set the value. + match self.semaphore.try_acquire() { + Ok(permit) => { + debug_assert!(!self.initialized()); + self.set_value(value, permit); + Ok(()) + } + Err(TryAcquireError::NoPermits) => { + // Some other task is holding the permit. That task is + // currently trying to initialize the value. + Err(SetError::InitializingError(value)) + } + Err(TryAcquireError::Closed) => { + // The semaphore was closed. Some other task has initialized + // the value. + Err(SetError::AlreadyInitializedError(value)) + } + } } - /// Tries to initialize the value of the OnceCell using the async function `f`. - /// If the value of the OnceCell was already initialized prior to this call, - /// a reference to that initialized value is returned. If some other thread - /// initiated the initialization prior to this call and the initialization - /// hasn't completed, this call waits until the initialization is finished. + /// Gets the value currently in the `OnceCell`, or initialize it with the + /// given asynchronous operation. + /// + /// If some other task is currently working on initializing the `OnceCell`, + /// this call will wait for that other task to finish, then return the value + /// that the other task produced. + /// + /// If the provided operation is cancelled or panics, the initialization + /// attempt is cancelled. If there are other tasks waiting for the value to + /// be initialized, one of them will start another attempt at initializing + /// the value. /// - /// This will deadlock if `f` tries to initialize the cell itself. + /// This will deadlock if `f` tries to initialize the cell recursively. pub async fn get_or_init<F, Fut>(&self, f: F) -> &T where F: FnOnce() -> Fut, Fut: Future<Output = T>, { + crate::trace::async_trace_leaf().await; + if self.initialized() { - // SAFETY: once the value is initialized, no mutable references are given out, so - // we can give out arbitrarily many immutable references + // SAFETY: The OnceCell has been fully initialized. unsafe { self.get_unchecked() } } else { - // After acquire().await we have either acquired a permit while self.value - // is still uninitialized, or the current thread is awoken after another thread - // has initialized the value and closed the semaphore, in which case self.initialized - // is true and we don't set the value here + // Here we try to acquire the semaphore permit. Holding the permit + // will allow us to set the value of the OnceCell, and prevents + // other tasks from initializing the OnceCell while we are holding + // it. match self.semaphore.acquire().await { - Ok(_permit) => { - if !self.initialized() { - // If `f()` panics or `select!` is called, this `get_or_init` call - // is aborted and the semaphore permit is dropped. - let value = f().await; - - // SAFETY: There is only one permit on the semaphore, hence only one - // mutable reference is created - unsafe { self.set_value(value) }; - - // SAFETY: once the value is initialized, no mutable references are given out, so - // we can give out arbitrarily many immutable references - unsafe { self.get_unchecked() } - } else { - unreachable!("acquired semaphore after value was already initialized."); - } + Ok(permit) => { + debug_assert!(!self.initialized()); + + // If `f()` panics or `select!` is called, this + // `get_or_init` call is aborted and the semaphore permit is + // dropped. + let value = f().await; + + self.set_value(value, permit) } Err(_) => { - if self.initialized() { - // SAFETY: once the value is initialized, no mutable references are given out, so - // we can give out arbitrarily many immutable references - unsafe { self.get_unchecked() } - } else { - unreachable!( - "Semaphore closed, but the OnceCell has not been initialized." - ); - } + debug_assert!(self.initialized()); + + // SAFETY: The semaphore has been closed. This only happens + // when the OnceCell is fully initialized. + unsafe { self.get_unchecked() } } } } } - /// Tries to initialize the value of the OnceCell using the async function `f`. - /// If the value of the OnceCell was already initialized prior to this call, - /// a reference to that initialized value is returned. If some other thread - /// initiated the initialization prior to this call and the initialization - /// hasn't completed, this call waits until the initialization is finished. - /// If the function argument `f` returns an error, `get_or_try_init` - /// returns that error, otherwise the result of `f` will be stored in the cell. + /// Gets the value currently in the `OnceCell`, or initialize it with the + /// given asynchronous operation. + /// + /// If some other task is currently working on initializing the `OnceCell`, + /// this call will wait for that other task to finish, then return the value + /// that the other task produced. /// - /// This will deadlock if `f` tries to initialize the cell itself. + /// If the provided operation returns an error, is cancelled or panics, the + /// initialization attempt is cancelled. If there are other tasks waiting + /// for the value to be initialized, one of them will start another attempt + /// at initializing the value. + /// + /// This will deadlock if `f` tries to initialize the cell recursively. pub async fn get_or_try_init<E, F, Fut>(&self, f: F) -> Result<&T, E> where F: FnOnce() -> Fut, Fut: Future<Output = Result<T, E>>, { + crate::trace::async_trace_leaf().await; + if self.initialized() { - // SAFETY: once the value is initialized, no mutable references are given out, so - // we can give out arbitrarily many immutable references + // SAFETY: The OnceCell has been fully initialized. unsafe { Ok(self.get_unchecked()) } } else { - // After acquire().await we have either acquired a permit while self.value - // is still uninitialized, or the current thread is awoken after another thread - // has initialized the value and closed the semaphore, in which case self.initialized - // is true and we don't set the value here + // Here we try to acquire the semaphore permit. Holding the permit + // will allow us to set the value of the OnceCell, and prevents + // other tasks from initializing the OnceCell while we are holding + // it. match self.semaphore.acquire().await { - Ok(_permit) => { - if !self.initialized() { - // If `f()` panics or `select!` is called, this `get_or_try_init` call - // is aborted and the semaphore permit is dropped. - let value = f().await; - - match value { - Ok(value) => { - // SAFETY: There is only one permit on the semaphore, hence only one - // mutable reference is created - unsafe { self.set_value(value) }; - - // SAFETY: once the value is initialized, no mutable references are given out, so - // we can give out arbitrarily many immutable references - unsafe { Ok(self.get_unchecked()) } - } - Err(e) => Err(e), - } - } else { - unreachable!("acquired semaphore after value was already initialized."); + Ok(permit) => { + debug_assert!(!self.initialized()); + + // If `f()` panics or `select!` is called, this + // `get_or_try_init` call is aborted and the semaphore + // permit is dropped. + let value = f().await; + + match value { + Ok(value) => Ok(self.set_value(value, permit)), + Err(e) => Err(e), } } Err(_) => { - if self.initialized() { - // SAFETY: once the value is initialized, no mutable references are given out, so - // we can give out arbitrarily many immutable references - unsafe { Ok(self.get_unchecked()) } - } else { - unreachable!( - "Semaphore closed, but the OnceCell has not been initialized." - ); - } + debug_assert!(self.initialized()); + + // SAFETY: The semaphore has been closed. This only happens + // when the OnceCell is fully initialized. + unsafe { Ok(self.get_unchecked()) } } } } } - /// Moves the value out of the cell, destroying the cell in the process. - /// - /// Returns `None` if the cell is uninitialized. + /// Takes the value from the cell, destroying the cell in the process. + /// Returns `None` if the cell is empty. pub fn into_inner(mut self) -> Option<T> { - if self.initialized() { + if self.initialized_mut() { // Set to uninitialized for the destructor of `OnceCell` to work properly *self.value_set.get_mut() = false; Some(unsafe { self.value.with(|ptr| ptr::read(ptr).assume_init()) }) @@ -338,20 +398,18 @@ impl<T> OnceCell<T> { } } - /// Takes ownership of the current value, leaving the cell uninitialized. - /// - /// Returns `None` if the cell is uninitialized. + /// Takes ownership of the current value, leaving the cell empty. Returns + /// `None` if the cell is empty. pub fn take(&mut self) -> Option<T> { std::mem::take(self).into_inner() } } -// Since `get` gives us access to immutable references of the -// OnceCell, OnceCell can only be Sync if T is Sync, otherwise -// OnceCell would allow sharing references of !Sync values across -// threads. We need T to be Send in order for OnceCell to by Sync -// because we can use `set` on `&OnceCell<T>` to send -// values (of type T) across threads. +// Since `get` gives us access to immutable references of the OnceCell, OnceCell +// can only be Sync if T is Sync, otherwise OnceCell would allow sharing +// references of !Sync values across threads. We need T to be Send in order for +// OnceCell to by Sync because we can use `set` on `&OnceCell<T>` to send values +// (of type T) across threads. unsafe impl<T: Sync + Send> Sync for OnceCell<T> {} // Access to OnceCell's value is guarded by the semaphore permit @@ -359,20 +417,17 @@ unsafe impl<T: Sync + Send> Sync for OnceCell<T> {} // it's safe to send it to another thread unsafe impl<T: Send> Send for OnceCell<T> {} -/// Errors that can be returned from [`OnceCell::set`] +/// Errors that can be returned from [`OnceCell::set`]. /// /// [`OnceCell::set`]: crate::sync::OnceCell::set -#[derive(Debug, PartialEq)] +#[derive(Debug, PartialEq, Eq)] pub enum SetError<T> { - /// Error resulting from [`OnceCell::set`] calls if the cell was previously initialized. + /// The cell was already initialized when [`OnceCell::set`] was called. /// /// [`OnceCell::set`]: crate::sync::OnceCell::set AlreadyInitializedError(T), - /// Error resulting from [`OnceCell::set`] calls when the cell is currently being - /// initialized during the calls to that method. - /// - /// [`OnceCell::set`]: crate::sync::OnceCell::set + /// The cell is currently being initialized. InitializingError(T), } diff --git a/vendor/tokio/src/sync/oneshot.rs b/vendor/tokio/src/sync/oneshot.rs index cb4649d86..af3cc854f 100644 --- a/vendor/tokio/src/sync/oneshot.rs +++ b/vendor/tokio/src/sync/oneshot.rs @@ -9,6 +9,13 @@ //! //! Each handle can be used on separate tasks. //! +//! Since the `send` method is not async, it can be used anywhere. This includes +//! sending between two runtimes, and using it from non-async code. +//! +//! If the [`Receiver`] is closed before receiving a message which has already +//! been sent, the message will remain in the channel until the receiver is +//! dropped, at which point the message will be dropped immediately. +//! //! # Examples //! //! ``` @@ -51,10 +58,76 @@ //! } //! } //! ``` +//! +//! To use a oneshot channel in a `tokio::select!` loop, add `&mut` in front of +//! the channel. +//! +//! ``` +//! use tokio::sync::oneshot; +//! use tokio::time::{interval, sleep, Duration}; +//! +//! #[tokio::main] +//! # async fn _doc() {} +//! # #[tokio::main(flavor = "current_thread", start_paused = true)] +//! async fn main() { +//! let (send, mut recv) = oneshot::channel(); +//! let mut interval = interval(Duration::from_millis(100)); +//! +//! # let handle = +//! tokio::spawn(async move { +//! sleep(Duration::from_secs(1)).await; +//! send.send("shut down").unwrap(); +//! }); +//! +//! loop { +//! tokio::select! { +//! _ = interval.tick() => println!("Another 100ms"), +//! msg = &mut recv => { +//! println!("Got message: {}", msg.unwrap()); +//! break; +//! } +//! } +//! } +//! # handle.await.unwrap(); +//! } +//! ``` +//! +//! To use a `Sender` from a destructor, put it in an [`Option`] and call +//! [`Option::take`]. +//! +//! ``` +//! use tokio::sync::oneshot; +//! +//! struct SendOnDrop { +//! sender: Option<oneshot::Sender<&'static str>>, +//! } +//! impl Drop for SendOnDrop { +//! fn drop(&mut self) { +//! if let Some(sender) = self.sender.take() { +//! // Using `let _ =` to ignore send errors. +//! let _ = sender.send("I got dropped!"); +//! } +//! } +//! } +//! +//! #[tokio::main] +//! # async fn _doc() {} +//! # #[tokio::main(flavor = "current_thread")] +//! async fn main() { +//! let (send, recv) = oneshot::channel(); +//! +//! let send_on_drop = SendOnDrop { sender: Some(send) }; +//! drop(send_on_drop); +//! +//! assert_eq!(recv.await, Ok("I got dropped!")); +//! } +//! ``` use crate::loom::cell::UnsafeCell; use crate::loom::sync::atomic::AtomicUsize; use crate::loom::sync::Arc; +#[cfg(all(tokio_unstable, feature = "tracing"))] +use crate::util::trace; use std::fmt; use std::future::Future; @@ -68,16 +141,108 @@ use std::task::{Context, Poll, Waker}; /// /// A pair of both a [`Sender`] and a [`Receiver`] are created by the /// [`channel`](fn@channel) function. +/// +/// # Examples +/// +/// ``` +/// use tokio::sync::oneshot; +/// +/// #[tokio::main] +/// async fn main() { +/// let (tx, rx) = oneshot::channel(); +/// +/// tokio::spawn(async move { +/// if let Err(_) = tx.send(3) { +/// println!("the receiver dropped"); +/// } +/// }); +/// +/// match rx.await { +/// Ok(v) => println!("got = {:?}", v), +/// Err(_) => println!("the sender dropped"), +/// } +/// } +/// ``` +/// +/// If the sender is dropped without sending, the receiver will fail with +/// [`error::RecvError`]: +/// +/// ``` +/// use tokio::sync::oneshot; +/// +/// #[tokio::main] +/// async fn main() { +/// let (tx, rx) = oneshot::channel::<u32>(); +/// +/// tokio::spawn(async move { +/// drop(tx); +/// }); +/// +/// match rx.await { +/// Ok(_) => panic!("This doesn't happen"), +/// Err(_) => println!("the sender dropped"), +/// } +/// } +/// ``` +/// +/// To use a `Sender` from a destructor, put it in an [`Option`] and call +/// [`Option::take`]. +/// +/// ``` +/// use tokio::sync::oneshot; +/// +/// struct SendOnDrop { +/// sender: Option<oneshot::Sender<&'static str>>, +/// } +/// impl Drop for SendOnDrop { +/// fn drop(&mut self) { +/// if let Some(sender) = self.sender.take() { +/// // Using `let _ =` to ignore send errors. +/// let _ = sender.send("I got dropped!"); +/// } +/// } +/// } +/// +/// #[tokio::main] +/// # async fn _doc() {} +/// # #[tokio::main(flavor = "current_thread")] +/// async fn main() { +/// let (send, recv) = oneshot::channel(); +/// +/// let send_on_drop = SendOnDrop { sender: Some(send) }; +/// drop(send_on_drop); +/// +/// assert_eq!(recv.await, Ok("I got dropped!")); +/// } +/// ``` +/// +/// [`Option`]: std::option::Option +/// [`Option::take`]: std::option::Option::take #[derive(Debug)] pub struct Sender<T> { inner: Option<Arc<Inner<T>>>, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, } -/// Receive a value from the associated [`Sender`]. +/// Receives a value from the associated [`Sender`]. /// /// A pair of both a [`Sender`] and a [`Receiver`] are created by the /// [`channel`](fn@channel) function. /// +/// This channel has no `recv` method because the receiver itself implements the +/// [`Future`] trait. To receive a `Result<T, `[`error::RecvError`]`>`, `.await` the `Receiver` object directly. +/// +/// The `poll` method on the `Future` trait is allowed to spuriously return +/// `Poll::Pending` even if the message has been sent. If such a spurious +/// failure happens, then the caller will be woken when the spurious failure has +/// been resolved so that the caller can attempt to receive the message again. +/// Note that receiving such a wakeup does not guarantee that the next call will +/// succeed — it could fail with another spurious failure. (A spurious failure +/// does not mean that the message is lost. It is just delayed.) +/// +/// [`Future`]: trait@std::future::Future +/// /// # Examples /// /// ``` @@ -120,22 +285,63 @@ pub struct Sender<T> { /// } /// } /// ``` +/// +/// To use a `Receiver` in a `tokio::select!` loop, add `&mut` in front of the +/// channel. +/// +/// ``` +/// use tokio::sync::oneshot; +/// use tokio::time::{interval, sleep, Duration}; +/// +/// #[tokio::main] +/// # async fn _doc() {} +/// # #[tokio::main(flavor = "current_thread", start_paused = true)] +/// async fn main() { +/// let (send, mut recv) = oneshot::channel(); +/// let mut interval = interval(Duration::from_millis(100)); +/// +/// # let handle = +/// tokio::spawn(async move { +/// sleep(Duration::from_secs(1)).await; +/// send.send("shut down").unwrap(); +/// }); +/// +/// loop { +/// tokio::select! { +/// _ = interval.tick() => println!("Another 100ms"), +/// msg = &mut recv => { +/// println!("Got message: {}", msg.unwrap()); +/// break; +/// } +/// } +/// } +/// # handle.await.unwrap(); +/// } +/// ``` #[derive(Debug)] pub struct Receiver<T> { inner: Option<Arc<Inner<T>>>, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, + #[cfg(all(tokio_unstable, feature = "tracing"))] + async_op_span: tracing::Span, + #[cfg(all(tokio_unstable, feature = "tracing"))] + async_op_poll_span: tracing::Span, } pub mod error { - //! Oneshot error types + //! Oneshot error types. use std::fmt; /// Error returned by the `Future` implementation for `Receiver`. - #[derive(Debug, Eq, PartialEq)] + /// + /// This error is returned by the receiver when the sender is dropped without sending. + #[derive(Debug, Eq, PartialEq, Clone)] pub struct RecvError(pub(super) ()); /// Error returned by the `try_recv` function on `Receiver`. - #[derive(Debug, Eq, PartialEq)] + #[derive(Debug, Eq, PartialEq, Clone)] pub enum TryRecvError { /// The send half of the channel has not yet sent a value. Empty, @@ -171,7 +377,7 @@ pub mod error { use self::error::*; struct Inner<T> { - /// Manages the state of the inner cell + /// Manages the state of the inner cell. state: AtomicUsize, /// The value. This is set by `Sender` and read by `Receiver`. The state of @@ -207,21 +413,21 @@ impl Task { F: FnOnce(&Waker) -> R, { self.0.with(|ptr| { - let waker: *const Waker = (&*ptr).as_ptr(); + let waker: *const Waker = (*ptr).as_ptr(); f(&*waker) }) } unsafe fn drop_task(&self) { self.0.with_mut(|ptr| { - let ptr: *mut Waker = (&mut *ptr).as_mut_ptr(); + let ptr: *mut Waker = (*ptr).as_mut_ptr(); ptr.drop_in_place(); }); } unsafe fn set_task(&self, cx: &mut Context<'_>) { self.0.with_mut(|ptr| { - let ptr: *mut Waker = (&mut *ptr).as_mut_ptr(); + let ptr: *mut Waker = (*ptr).as_mut_ptr(); ptr.write(cx.waker().clone()); }); } @@ -230,7 +436,7 @@ impl Task { #[derive(Clone, Copy)] struct State(usize); -/// Create a new one-shot channel for sending single values across asynchronous +/// Creates a new one-shot channel for sending single values across asynchronous /// tasks. /// /// The function returns separate "send" and "receive" handles. The `Sender` @@ -260,7 +466,56 @@ struct State(usize); /// } /// } /// ``` +#[track_caller] pub fn channel<T>() -> (Sender<T>, Receiver<T>) { + #[cfg(all(tokio_unstable, feature = "tracing"))] + let resource_span = { + let location = std::panic::Location::caller(); + + let resource_span = tracing::trace_span!( + "runtime.resource", + concrete_type = "Sender|Receiver", + kind = "Sync", + loc.file = location.file(), + loc.line = location.line(), + loc.col = location.column(), + ); + + resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + tx_dropped = false, + tx_dropped.op = "override", + ) + }); + + resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + rx_dropped = false, + rx_dropped.op = "override", + ) + }); + + resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + value_sent = false, + value_sent.op = "override", + ) + }); + + resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + value_received = false, + value_received.op = "override", + ) + }); + + resource_span + }; + let inner = Arc::new(Inner { state: AtomicUsize::new(State::new().as_usize()), value: UnsafeCell::new(None), @@ -270,8 +525,27 @@ pub fn channel<T>() -> (Sender<T>, Receiver<T>) { let tx = Sender { inner: Some(inner.clone()), + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: resource_span.clone(), + }; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + let async_op_span = resource_span + .in_scope(|| tracing::trace_span!("runtime.resource.async_op", source = "Receiver::await")); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + let async_op_poll_span = + async_op_span.in_scope(|| tracing::trace_span!("runtime.resource.async_op.poll")); + + let rx = Receiver { + inner: Some(inner), + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span, + #[cfg(all(tokio_unstable, feature = "tracing"))] + async_op_span, + #[cfg(all(tokio_unstable, feature = "tracing"))] + async_op_poll_span, }; - let rx = Receiver { inner: Some(inner) }; (tx, rx) } @@ -343,6 +617,15 @@ impl<T> Sender<T> { } } + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + value_sent = true, + value_sent.op = "override", + ) + }); + Ok(()) } @@ -416,7 +699,20 @@ impl<T> Sender<T> { pub async fn closed(&mut self) { use crate::future::poll_fn; - poll_fn(|cx| self.poll_closed(cx)).await + #[cfg(all(tokio_unstable, feature = "tracing"))] + let resource_span = self.resource_span.clone(); + #[cfg(all(tokio_unstable, feature = "tracing"))] + let closed = trace::async_op( + || poll_fn(|cx| self.poll_closed(cx)), + resource_span, + "Sender::closed", + "poll_closed", + false, + ); + #[cfg(not(all(tokio_unstable, feature = "tracing")))] + let closed = poll_fn(|cx| self.poll_closed(cx)); + + closed.await } /// Returns `true` if the associated [`Receiver`] handle has been dropped. @@ -453,7 +749,7 @@ impl<T> Sender<T> { state.is_closed() } - /// Check whether the oneshot channel has been closed, and if not, schedules the + /// Checks whether the oneshot channel has been closed, and if not, schedules the /// `Waker` in the provided `Context` to receive a notification when the channel is /// closed. /// @@ -494,8 +790,10 @@ impl<T> Sender<T> { /// } /// ``` pub fn poll_closed(&mut self, cx: &mut Context<'_>) -> Poll<()> { + ready!(crate::trace::trace_leaf(cx)); + // Keep track of task budget - let coop = ready!(crate::coop::poll_proceed(cx)); + let coop = ready!(crate::runtime::coop::poll_proceed(cx)); let inner = self.inner.as_ref().unwrap(); @@ -503,7 +801,7 @@ impl<T> Sender<T> { if state.is_closed() { coop.made_progress(); - return Poll::Ready(()); + return Ready(()); } if state.is_tx_task_set() { @@ -546,6 +844,14 @@ impl<T> Drop for Sender<T> { fn drop(&mut self) { if let Some(inner) = self.inner.as_ref() { inner.complete(); + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + tx_dropped = true, + tx_dropped.op = "override", + ) + }); } } } @@ -613,6 +919,14 @@ impl<T> Receiver<T> { pub fn close(&mut self) { if let Some(inner) = self.inner.as_ref() { inner.close(); + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + rx_dropped = true, + rx_dropped.op = "override", + ) + }); } } @@ -625,12 +939,16 @@ impl<T> Receiver<T> { /// This function is useful to call from outside the context of an /// asynchronous task. /// + /// Note that unlike the `poll` method, the `try_recv` method cannot fail + /// spuriously. Any send or close event that happens before this call to + /// `try_recv` will be correctly returned to the caller. + /// /// # Return /// /// - `Ok(T)` if a value is pending in the channel. /// - `Err(TryRecvError::Empty)` if no value has been sent yet. /// - `Err(TryRecvError::Closed)` if the sender has dropped without sending - /// a value. + /// a value, or if the message has already been received. /// /// # Examples /// @@ -690,7 +1008,17 @@ impl<T> Receiver<T> { // `UnsafeCell`. Therefore, it is now safe for us to access the // cell. match unsafe { inner.consume_value() } { - Some(value) => Ok(value), + Some(value) => { + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + value_received = true, + value_received.op = "override", + ) + }); + Ok(value) + } None => Err(TryRecvError::Closed), } } else if state.is_closed() { @@ -706,12 +1034,52 @@ impl<T> Receiver<T> { self.inner = None; result } + + /// Blocking receive to call outside of asynchronous contexts. + /// + /// # Panics + /// + /// This function panics if called within an asynchronous execution + /// context. + /// + /// # Examples + /// + /// ``` + /// use std::thread; + /// use tokio::sync::oneshot; + /// + /// #[tokio::main] + /// async fn main() { + /// let (tx, rx) = oneshot::channel::<u8>(); + /// + /// let sync_code = thread::spawn(move || { + /// assert_eq!(Ok(10), rx.blocking_recv()); + /// }); + /// + /// let _ = tx.send(10); + /// sync_code.join().unwrap(); + /// } + /// ``` + #[track_caller] + #[cfg(feature = "sync")] + #[cfg_attr(docsrs, doc(alias = "recv_blocking"))] + pub fn blocking_recv(self) -> Result<T, RecvError> { + crate::future::block_on(self) + } } impl<T> Drop for Receiver<T> { fn drop(&mut self) { if let Some(inner) = self.inner.as_ref() { inner.close(); + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + rx_dropped = true, + rx_dropped.op = "override", + ) + }); } } } @@ -721,8 +1089,21 @@ impl<T> Future for Receiver<T> { fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { // If `inner` is `None`, then `poll()` has already completed. + #[cfg(all(tokio_unstable, feature = "tracing"))] + let _res_span = self.resource_span.clone().entered(); + #[cfg(all(tokio_unstable, feature = "tracing"))] + let _ao_span = self.async_op_span.clone().entered(); + #[cfg(all(tokio_unstable, feature = "tracing"))] + let _ao_poll_span = self.async_op_poll_span.clone().entered(); + let ret = if let Some(inner) = self.as_ref().get_ref().inner.as_ref() { - ready!(inner.poll_recv(cx))? + #[cfg(all(tokio_unstable, feature = "tracing"))] + let res = ready!(trace_poll_op!("poll_recv", inner.poll_recv(cx)))?; + + #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] + let res = ready!(inner.poll_recv(cx))?; + + res } else { panic!("called after complete"); }; @@ -751,8 +1132,9 @@ impl<T> Inner<T> { } fn poll_recv(&self, cx: &mut Context<'_>) -> Poll<Result<T, RecvError>> { + ready!(crate::trace::trace_leaf(cx)); // Keep track of task budget - let coop = ready!(crate::coop::poll_proceed(cx)); + let coop = ready!(crate::runtime::coop::poll_proceed(cx)); // Load the state let mut state = State::load(&self.state, Acquire); diff --git a/vendor/tokio/src/sync/rwlock.rs b/vendor/tokio/src/sync/rwlock.rs index 120bc72b8..dd4928546 100644 --- a/vendor/tokio/src/sync/rwlock.rs +++ b/vendor/tokio/src/sync/rwlock.rs @@ -1,9 +1,10 @@ use crate::sync::batch_semaphore::{Semaphore, TryAcquireError}; use crate::sync::mutex::TryLockError; +#[cfg(all(tokio_unstable, feature = "tracing"))] +use crate::util::trace; use std::cell::UnsafeCell; use std::marker; use std::marker::PhantomData; -use std::mem::ManuallyDrop; use std::sync::Arc; pub(crate) mod owned_read_guard; @@ -84,8 +85,10 @@ const MAX_READS: u32 = 10; /// [`RwLockWriteGuard`]: struct@RwLockWriteGuard /// [`Send`]: trait@std::marker::Send /// [_write-preferring_]: https://en.wikipedia.org/wiki/Readers%E2%80%93writer_lock#Priority_policies -#[derive(Debug)] pub struct RwLock<T: ?Sized> { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, + // maximum number of concurrent readers mr: u32, @@ -197,14 +200,55 @@ impl<T: ?Sized> RwLock<T> { /// /// let lock = RwLock::new(5); /// ``` + #[track_caller] pub fn new(value: T) -> RwLock<T> where T: Sized, { + #[cfg(all(tokio_unstable, feature = "tracing"))] + let resource_span = { + let location = std::panic::Location::caller(); + let resource_span = tracing::trace_span!( + "runtime.resource", + concrete_type = "RwLock", + kind = "Sync", + loc.file = location.file(), + loc.line = location.line(), + loc.col = location.column(), + ); + + resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + max_readers = MAX_READS, + ); + + tracing::trace!( + target: "runtime::resource::state_update", + write_locked = false, + ); + + tracing::trace!( + target: "runtime::resource::state_update", + current_readers = 0, + ); + }); + + resource_span + }; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + let s = resource_span.in_scope(|| Semaphore::new(MAX_READS as usize)); + + #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] + let s = Semaphore::new(MAX_READS as usize); + RwLock { mr: MAX_READS, c: UnsafeCell::new(value), - s: Semaphore::new(MAX_READS as usize), + s, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span, } } @@ -222,6 +266,7 @@ impl<T: ?Sized> RwLock<T> { /// # Panics /// /// Panics if `max_reads` is more than `u32::MAX >> 3`. + #[track_caller] pub fn with_max_readers(value: T, max_reads: u32) -> RwLock<T> where T: Sized, @@ -231,10 +276,52 @@ impl<T: ?Sized> RwLock<T> { "a RwLock may not be created with more than {} readers", MAX_READS ); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + let resource_span = { + let location = std::panic::Location::caller(); + + let resource_span = tracing::trace_span!( + "runtime.resource", + concrete_type = "RwLock", + kind = "Sync", + loc.file = location.file(), + loc.line = location.line(), + loc.col = location.column(), + ); + + resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + max_readers = max_reads, + ); + + tracing::trace!( + target: "runtime::resource::state_update", + write_locked = false, + ); + + tracing::trace!( + target: "runtime::resource::state_update", + current_readers = 0, + ); + }); + + resource_span + }; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + let s = resource_span.in_scope(|| Semaphore::new(max_reads as usize)); + + #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] + let s = Semaphore::new(max_reads as usize); + RwLock { mr: max_reads, c: UnsafeCell::new(value), - s: Semaphore::new(max_reads as usize), + s, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span, } } @@ -257,6 +344,8 @@ impl<T: ?Sized> RwLock<T> { mr: MAX_READS, c: UnsafeCell::new(value), s: Semaphore::const_new(MAX_READS as usize), + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span::none(), } } @@ -281,6 +370,8 @@ impl<T: ?Sized> RwLock<T> { mr: max_reads, c: UnsafeCell::new(value), s: Semaphore::const_new(max_reads as usize), + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span::none(), } } @@ -327,19 +418,100 @@ impl<T: ?Sized> RwLock<T> { /// /// // Drop the guard after the spawned task finishes. /// drop(n); - ///} + /// } /// ``` pub async fn read(&self) -> RwLockReadGuard<'_, T> { - self.s.acquire(1).await.unwrap_or_else(|_| { - // The semaphore was closed. but, we never explicitly close it, and we have a - // handle to it through the Arc, which means that this can never happen. - unreachable!() + let acquire_fut = async { + self.s.acquire(1).await.unwrap_or_else(|_| { + // The semaphore was closed. but, we never explicitly close it, and we have a + // handle to it through the Arc, which means that this can never happen. + unreachable!() + }); + + RwLockReadGuard { + s: &self.s, + data: self.c.get(), + marker: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: self.resource_span.clone(), + } + }; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + let acquire_fut = trace::async_op( + move || acquire_fut, + self.resource_span.clone(), + "RwLock::read", + "poll", + false, + ); + + #[allow(clippy::let_and_return)] // this lint triggers when disabling tracing + let guard = acquire_fut.await; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + current_readers = 1, + current_readers.op = "add", + ) }); - RwLockReadGuard { - s: &self.s, - data: self.c.get(), - marker: marker::PhantomData, - } + + guard + } + + /// Blockingly locks this `RwLock` with shared read access. + /// + /// This method is intended for use cases where you + /// need to use this rwlock in asynchronous code as well as in synchronous code. + /// + /// Returns an RAII guard which will drop the read access of this `RwLock` when dropped. + /// + /// # Panics + /// + /// This function panics if called within an asynchronous execution context. + /// + /// - If you find yourself in an asynchronous execution context and needing + /// to call some (synchronous) function which performs one of these + /// `blocking_` operations, then consider wrapping that call inside + /// [`spawn_blocking()`][crate::runtime::Handle::spawn_blocking] + /// (or [`block_in_place()`][crate::task::block_in_place]). + /// + /// # Examples + /// + /// ``` + /// use std::sync::Arc; + /// use tokio::sync::RwLock; + /// + /// #[tokio::main] + /// async fn main() { + /// let rwlock = Arc::new(RwLock::new(1)); + /// let mut write_lock = rwlock.write().await; + /// + /// let blocking_task = tokio::task::spawn_blocking({ + /// let rwlock = Arc::clone(&rwlock); + /// move || { + /// // This shall block until the `write_lock` is released. + /// let read_lock = rwlock.blocking_read(); + /// assert_eq!(*read_lock, 0); + /// } + /// }); + /// + /// *write_lock -= 1; + /// drop(write_lock); // release the lock. + /// + /// // Await the completion of the blocking task. + /// blocking_task.await.unwrap(); + /// + /// // Assert uncontended. + /// assert!(rwlock.try_write().is_ok()); + /// } + /// ``` + #[track_caller] + #[cfg(feature = "sync")] + pub fn blocking_read(&self) -> RwLockReadGuard<'_, T> { + crate::future::block_on(self.read()) } /// Locks this `RwLock` with shared read access, causing the current task @@ -394,16 +566,47 @@ impl<T: ?Sized> RwLock<T> { ///} /// ``` pub async fn read_owned(self: Arc<Self>) -> OwnedRwLockReadGuard<T> { - self.s.acquire(1).await.unwrap_or_else(|_| { - // The semaphore was closed. but, we never explicitly close it, and we have a - // handle to it through the Arc, which means that this can never happen. - unreachable!() + #[cfg(all(tokio_unstable, feature = "tracing"))] + let resource_span = self.resource_span.clone(); + + let acquire_fut = async { + self.s.acquire(1).await.unwrap_or_else(|_| { + // The semaphore was closed. but, we never explicitly close it, and we have a + // handle to it through the Arc, which means that this can never happen. + unreachable!() + }); + + OwnedRwLockReadGuard { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: self.resource_span.clone(), + data: self.c.get(), + lock: self, + _p: PhantomData, + } + }; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + let acquire_fut = trace::async_op( + move || acquire_fut, + resource_span, + "RwLock::read_owned", + "poll", + false, + ); + + #[allow(clippy::let_and_return)] // this lint triggers when disabling tracing + let guard = acquire_fut.await; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + guard.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + current_readers = 1, + current_readers.op = "add", + ) }); - OwnedRwLockReadGuard { - data: self.c.get(), - lock: ManuallyDrop::new(self), - _p: PhantomData, - } + + guard } /// Attempts to acquire this `RwLock` with shared read access. @@ -445,11 +648,24 @@ impl<T: ?Sized> RwLock<T> { Err(TryAcquireError::Closed) => unreachable!(), } - Ok(RwLockReadGuard { + let guard = RwLockReadGuard { s: &self.s, data: self.c.get(), marker: marker::PhantomData, - }) + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: self.resource_span.clone(), + }; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + current_readers = 1, + current_readers.op = "add", + ) + }); + + Ok(guard) } /// Attempts to acquire this `RwLock` with shared read access. @@ -497,11 +713,24 @@ impl<T: ?Sized> RwLock<T> { Err(TryAcquireError::Closed) => unreachable!(), } - Ok(OwnedRwLockReadGuard { + let guard = OwnedRwLockReadGuard { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: self.resource_span.clone(), data: self.c.get(), - lock: ManuallyDrop::new(self), + lock: self, _p: PhantomData, - }) + }; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + guard.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + current_readers = 1, + current_readers.op = "add", + ) + }); + + Ok(guard) } /// Locks this `RwLock` with exclusive write access, causing the current @@ -533,17 +762,100 @@ impl<T: ?Sized> RwLock<T> { ///} /// ``` pub async fn write(&self) -> RwLockWriteGuard<'_, T> { - self.s.acquire(self.mr).await.unwrap_or_else(|_| { - // The semaphore was closed. but, we never explicitly close it, and we have a - // handle to it through the Arc, which means that this can never happen. - unreachable!() + let acquire_fut = async { + self.s.acquire(self.mr).await.unwrap_or_else(|_| { + // The semaphore was closed. but, we never explicitly close it, and we have a + // handle to it through the Arc, which means that this can never happen. + unreachable!() + }); + + RwLockWriteGuard { + permits_acquired: self.mr, + s: &self.s, + data: self.c.get(), + marker: marker::PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: self.resource_span.clone(), + } + }; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + let acquire_fut = trace::async_op( + move || acquire_fut, + self.resource_span.clone(), + "RwLock::write", + "poll", + false, + ); + + #[allow(clippy::let_and_return)] // this lint triggers when disabling tracing + let guard = acquire_fut.await; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + write_locked = true, + write_locked.op = "override", + ) }); - RwLockWriteGuard { - permits_acquired: self.mr, - s: &self.s, - data: self.c.get(), - marker: marker::PhantomData, - } + + guard + } + + /// Blockingly locks this `RwLock` with exclusive write access. + /// + /// This method is intended for use cases where you + /// need to use this rwlock in asynchronous code as well as in synchronous code. + /// + /// Returns an RAII guard which will drop the write access of this `RwLock` when dropped. + /// + /// # Panics + /// + /// This function panics if called within an asynchronous execution context. + /// + /// - If you find yourself in an asynchronous execution context and needing + /// to call some (synchronous) function which performs one of these + /// `blocking_` operations, then consider wrapping that call inside + /// [`spawn_blocking()`][crate::runtime::Handle::spawn_blocking] + /// (or [`block_in_place()`][crate::task::block_in_place]). + /// + /// # Examples + /// + /// ``` + /// use std::sync::Arc; + /// use tokio::{sync::RwLock}; + /// + /// #[tokio::main] + /// async fn main() { + /// let rwlock = Arc::new(RwLock::new(1)); + /// let read_lock = rwlock.read().await; + /// + /// let blocking_task = tokio::task::spawn_blocking({ + /// let rwlock = Arc::clone(&rwlock); + /// move || { + /// // This shall block until the `read_lock` is released. + /// let mut write_lock = rwlock.blocking_write(); + /// *write_lock = 2; + /// } + /// }); + /// + /// assert_eq!(*read_lock, 1); + /// // Release the last outstanding read lock. + /// drop(read_lock); + /// + /// // Await the completion of the blocking task. + /// blocking_task.await.unwrap(); + /// + /// // Assert uncontended. + /// let read_lock = rwlock.try_read().unwrap(); + /// assert_eq!(*read_lock, 2); + /// } + /// ``` + #[track_caller] + #[cfg(feature = "sync")] + pub fn blocking_write(&self) -> RwLockWriteGuard<'_, T> { + crate::future::block_on(self.write()) } /// Locks this `RwLock` with exclusive write access, causing the current @@ -582,17 +894,48 @@ impl<T: ?Sized> RwLock<T> { ///} /// ``` pub async fn write_owned(self: Arc<Self>) -> OwnedRwLockWriteGuard<T> { - self.s.acquire(self.mr).await.unwrap_or_else(|_| { - // The semaphore was closed. but, we never explicitly close it, and we have a - // handle to it through the Arc, which means that this can never happen. - unreachable!() + #[cfg(all(tokio_unstable, feature = "tracing"))] + let resource_span = self.resource_span.clone(); + + let acquire_fut = async { + self.s.acquire(self.mr).await.unwrap_or_else(|_| { + // The semaphore was closed. but, we never explicitly close it, and we have a + // handle to it through the Arc, which means that this can never happen. + unreachable!() + }); + + OwnedRwLockWriteGuard { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: self.resource_span.clone(), + permits_acquired: self.mr, + data: self.c.get(), + lock: self, + _p: PhantomData, + } + }; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + let acquire_fut = trace::async_op( + move || acquire_fut, + resource_span, + "RwLock::write_owned", + "poll", + false, + ); + + #[allow(clippy::let_and_return)] // this lint triggers when disabling tracing + let guard = acquire_fut.await; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + guard.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + write_locked = true, + write_locked.op = "override", + ) }); - OwnedRwLockWriteGuard { - permits_acquired: self.mr, - data: self.c.get(), - lock: ManuallyDrop::new(self), - _p: PhantomData, - } + + guard } /// Attempts to acquire this `RwLock` with exclusive write access. @@ -625,12 +968,25 @@ impl<T: ?Sized> RwLock<T> { Err(TryAcquireError::Closed) => unreachable!(), } - Ok(RwLockWriteGuard { + let guard = RwLockWriteGuard { permits_acquired: self.mr, s: &self.s, data: self.c.get(), marker: marker::PhantomData, - }) + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: self.resource_span.clone(), + }; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + write_locked = true, + write_locked.op = "override", + ) + }); + + Ok(guard) } /// Attempts to acquire this `RwLock` with exclusive write access. @@ -670,12 +1026,25 @@ impl<T: ?Sized> RwLock<T> { Err(TryAcquireError::Closed) => unreachable!(), } - Ok(OwnedRwLockWriteGuard { + let guard = OwnedRwLockWriteGuard { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: self.resource_span.clone(), permits_acquired: self.mr, data: self.c.get(), - lock: ManuallyDrop::new(self), + lock: self, _p: PhantomData, - }) + }; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + guard.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + write_locked = true, + write_locked.op = "override", + ) + }); + + Ok(guard) } /// Returns a mutable reference to the underlying data. @@ -725,3 +1094,17 @@ where Self::new(T::default()) } } + +impl<T: ?Sized> std::fmt::Debug for RwLock<T> +where + T: std::fmt::Debug, +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let mut d = f.debug_struct("RwLock"); + match self.try_read() { + Ok(inner) => d.field("data", &&*inner), + Err(_) => d.field("data", &format_args!("<locked>")), + }; + d.finish() + } +} diff --git a/vendor/tokio/src/sync/rwlock/owned_read_guard.rs b/vendor/tokio/src/sync/rwlock/owned_read_guard.rs index b7f3926a4..273e7b86f 100644 --- a/vendor/tokio/src/sync/rwlock/owned_read_guard.rs +++ b/vendor/tokio/src/sync/rwlock/owned_read_guard.rs @@ -1,10 +1,7 @@ use crate::sync::rwlock::RwLock; -use std::fmt; use std::marker::PhantomData; -use std::mem; -use std::mem::ManuallyDrop; -use std::ops; use std::sync::Arc; +use std::{fmt, mem, ops, ptr}; /// Owned RAII structure used to release the shared read access of a lock when /// dropped. @@ -14,15 +11,41 @@ use std::sync::Arc; /// /// [`read_owned`]: method@crate::sync::RwLock::read_owned /// [`RwLock`]: struct@crate::sync::RwLock +#[clippy::has_significant_drop] pub struct OwnedRwLockReadGuard<T: ?Sized, U: ?Sized = T> { - // ManuallyDrop allows us to destructure into this field without running the destructor. - pub(super) lock: ManuallyDrop<Arc<RwLock<T>>>, + // When changing the fields in this struct, make sure to update the + // `skip_drop` method. + #[cfg(all(tokio_unstable, feature = "tracing"))] + pub(super) resource_span: tracing::Span, + pub(super) lock: Arc<RwLock<T>>, pub(super) data: *const U, pub(super) _p: PhantomData<T>, } +#[allow(dead_code)] // Unused fields are still used in Drop. +struct Inner<T: ?Sized, U: ?Sized> { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, + lock: Arc<RwLock<T>>, + data: *const U, +} + impl<T: ?Sized, U: ?Sized> OwnedRwLockReadGuard<T, U> { - /// Make a new `OwnedRwLockReadGuard` for a component of the locked data. + fn skip_drop(self) -> Inner<T, U> { + let me = mem::ManuallyDrop::new(self); + // SAFETY: This duplicates the values in every field of the guard, then + // forgets the originals, so in the end no value is duplicated. + unsafe { + Inner { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: ptr::read(&me.resource_span), + lock: ptr::read(&me.lock), + data: me.data, + } + } + } + + /// Makes a new `OwnedRwLockReadGuard` for a component of the locked data. /// This operation cannot fail as the `OwnedRwLockReadGuard` passed in /// already locked the data. /// @@ -50,18 +73,19 @@ impl<T: ?Sized, U: ?Sized> OwnedRwLockReadGuard<T, U> { /// # } /// ``` #[inline] - pub fn map<F, V: ?Sized>(mut this: Self, f: F) -> OwnedRwLockReadGuard<T, V> + pub fn map<F, V: ?Sized>(this: Self, f: F) -> OwnedRwLockReadGuard<T, V> where F: FnOnce(&U) -> &V, { let data = f(&*this) as *const V; - let lock = unsafe { ManuallyDrop::take(&mut this.lock) }; - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); + OwnedRwLockReadGuard { - lock: ManuallyDrop::new(lock), + lock: this.lock, data, _p: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: this.resource_span, } } @@ -96,7 +120,7 @@ impl<T: ?Sized, U: ?Sized> OwnedRwLockReadGuard<T, U> { /// # } /// ``` #[inline] - pub fn try_map<F, V: ?Sized>(mut this: Self, f: F) -> Result<OwnedRwLockReadGuard<T, V>, Self> + pub fn try_map<F, V: ?Sized>(this: Self, f: F) -> Result<OwnedRwLockReadGuard<T, V>, Self> where F: FnOnce(&U) -> Option<&V>, { @@ -104,13 +128,14 @@ impl<T: ?Sized, U: ?Sized> OwnedRwLockReadGuard<T, U> { Some(data) => data as *const V, None => return Err(this), }; - let lock = unsafe { ManuallyDrop::take(&mut this.lock) }; - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); + Ok(OwnedRwLockReadGuard { - lock: ManuallyDrop::new(lock), + lock: this.lock, data, _p: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: this.resource_span, }) } } @@ -144,6 +169,14 @@ where impl<T: ?Sized, U: ?Sized> Drop for OwnedRwLockReadGuard<T, U> { fn drop(&mut self) { self.lock.s.release(1); - unsafe { ManuallyDrop::drop(&mut self.lock) }; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + current_readers = 1, + current_readers.op = "sub", + ) + }); } } diff --git a/vendor/tokio/src/sync/rwlock/owned_write_guard.rs b/vendor/tokio/src/sync/rwlock/owned_write_guard.rs index 91b659524..a8ce4a160 100644 --- a/vendor/tokio/src/sync/rwlock/owned_write_guard.rs +++ b/vendor/tokio/src/sync/rwlock/owned_write_guard.rs @@ -1,11 +1,9 @@ use crate::sync::rwlock::owned_read_guard::OwnedRwLockReadGuard; use crate::sync::rwlock::owned_write_guard_mapped::OwnedRwLockMappedWriteGuard; use crate::sync::rwlock::RwLock; -use std::fmt; use std::marker::PhantomData; -use std::mem::{self, ManuallyDrop}; -use std::ops; use std::sync::Arc; +use std::{fmt, mem, ops, ptr}; /// Owned RAII structure used to release the exclusive write access of a lock when /// dropped. @@ -15,16 +13,44 @@ use std::sync::Arc; /// /// [`write_owned`]: method@crate::sync::RwLock::write_owned /// [`RwLock`]: struct@crate::sync::RwLock +#[clippy::has_significant_drop] pub struct OwnedRwLockWriteGuard<T: ?Sized> { + // When changing the fields in this struct, make sure to update the + // `skip_drop` method. + #[cfg(all(tokio_unstable, feature = "tracing"))] + pub(super) resource_span: tracing::Span, pub(super) permits_acquired: u32, - // ManuallyDrop allows us to destructure into this field without running the destructor. - pub(super) lock: ManuallyDrop<Arc<RwLock<T>>>, + pub(super) lock: Arc<RwLock<T>>, pub(super) data: *mut T, pub(super) _p: PhantomData<T>, } +#[allow(dead_code)] // Unused fields are still used in Drop. +struct Inner<T: ?Sized> { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, + permits_acquired: u32, + lock: Arc<RwLock<T>>, + data: *const T, +} + impl<T: ?Sized> OwnedRwLockWriteGuard<T> { - /// Make a new [`OwnedRwLockMappedWriteGuard`] for a component of the locked + fn skip_drop(self) -> Inner<T> { + let me = mem::ManuallyDrop::new(self); + // SAFETY: This duplicates the values in every field of the guard, then + // forgets the originals, so in the end no value is duplicated. + unsafe { + Inner { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: ptr::read(&me.resource_span), + permits_acquired: me.permits_acquired, + lock: ptr::read(&me.lock), + data: me.data, + } + } + } + + /// Makes a new [`OwnedRwLockMappedWriteGuard`] for a component of the locked /// data. /// /// This operation cannot fail as the `OwnedRwLockWriteGuard` passed in @@ -62,19 +88,90 @@ impl<T: ?Sized> OwnedRwLockWriteGuard<T> { F: FnOnce(&mut T) -> &mut U, { let data = f(&mut *this) as *mut U; - let lock = unsafe { ManuallyDrop::take(&mut this.lock) }; - let permits_acquired = this.permits_acquired; - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); + OwnedRwLockMappedWriteGuard { - permits_acquired, - lock: ManuallyDrop::new(lock), + permits_acquired: this.permits_acquired, + lock: this.lock, data, _p: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: this.resource_span, } } - /// Attempts to make a new [`OwnedRwLockMappedWriteGuard`] for a component + /// Makes a new [`OwnedRwLockReadGuard`] for a component of the locked data. + /// + /// This operation cannot fail as the `OwnedRwLockWriteGuard` passed in already + /// locked the data. + /// + /// This is an associated function that needs to be used as + /// `OwnedRwLockWriteGuard::downgrade_map(..)`. A method would interfere with methods of + /// the same name on the contents of the locked data. + /// + /// Inside of `f`, you retain exclusive access to the data, despite only being given a `&T`. Handing out a + /// `&mut T` would result in unsoundness, as you could use interior mutability. + /// + /// # Examples + /// + /// ``` + /// use std::sync::Arc; + /// use tokio::sync::{RwLock, OwnedRwLockWriteGuard}; + /// + /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] + /// struct Foo(u32); + /// + /// # #[tokio::main] + /// # async fn main() { + /// let lock = Arc::new(RwLock::new(Foo(1))); + /// + /// let guard = Arc::clone(&lock).write_owned().await; + /// let mapped = OwnedRwLockWriteGuard::downgrade_map(guard, |f| &f.0); + /// let foo = lock.read_owned().await; + /// assert_eq!(foo.0, *mapped); + /// # } + /// ``` + #[inline] + pub fn downgrade_map<F, U: ?Sized>(this: Self, f: F) -> OwnedRwLockReadGuard<T, U> + where + F: FnOnce(&T) -> &U, + { + let data = f(&*this) as *const U; + let this = this.skip_drop(); + let guard = OwnedRwLockReadGuard { + lock: this.lock, + data, + _p: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: this.resource_span, + }; + + // Release all but one of the permits held by the write guard + let to_release = (this.permits_acquired - 1) as usize; + guard.lock.s.release(to_release); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + guard.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + write_locked = false, + write_locked.op = "override", + ) + }); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + guard.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + current_readers = 1, + current_readers.op = "add", + ) + }); + + guard + } + + /// Attempts to make a new [`OwnedRwLockMappedWriteGuard`] for a component /// of the locked data. The original guard is returned if the closure /// returns `None`. /// @@ -121,18 +218,99 @@ impl<T: ?Sized> OwnedRwLockWriteGuard<T> { Some(data) => data as *mut U, None => return Err(this), }; - let permits_acquired = this.permits_acquired; - let lock = unsafe { ManuallyDrop::take(&mut this.lock) }; - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); + Ok(OwnedRwLockMappedWriteGuard { - permits_acquired, - lock: ManuallyDrop::new(lock), + permits_acquired: this.permits_acquired, + lock: this.lock, data, _p: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: this.resource_span, }) } + /// Attempts to make a new [`OwnedRwLockReadGuard`] for a component of + /// the locked data. The original guard is returned if the closure returns + /// `None`. + /// + /// This operation cannot fail as the `OwnedRwLockWriteGuard` passed in already + /// locked the data. + /// + /// This is an associated function that needs to be + /// used as `OwnedRwLockWriteGuard::try_downgrade_map(...)`. A method would interfere with + /// methods of the same name on the contents of the locked data. + /// + /// Inside of `f`, you retain exclusive access to the data, despite only being given a `&T`. Handing out a + /// `&mut T` would result in unsoundness, as you could use interior mutability. + /// + /// If this function returns `Err(...)`, the lock is never unlocked nor downgraded. + /// + /// # Examples + /// + /// ``` + /// use std::sync::Arc; + /// use tokio::sync::{RwLock, OwnedRwLockWriteGuard}; + /// + /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] + /// struct Foo(u32); + /// + /// # #[tokio::main] + /// # async fn main() { + /// let lock = Arc::new(RwLock::new(Foo(1))); + /// + /// let guard = Arc::clone(&lock).write_owned().await; + /// let guard = OwnedRwLockWriteGuard::try_downgrade_map(guard, |f| Some(&f.0)).expect("should not fail"); + /// let foo = lock.read_owned().await; + /// assert_eq!(foo.0, *guard); + /// # } + /// ``` + #[inline] + pub fn try_downgrade_map<F, U: ?Sized>( + this: Self, + f: F, + ) -> Result<OwnedRwLockReadGuard<T, U>, Self> + where + F: FnOnce(&T) -> Option<&U>, + { + let data = match f(&*this) { + Some(data) => data as *const U, + None => return Err(this), + }; + let this = this.skip_drop(); + let guard = OwnedRwLockReadGuard { + lock: this.lock, + data, + _p: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: this.resource_span, + }; + + // Release all but one of the permits held by the write guard + let to_release = (this.permits_acquired - 1) as usize; + guard.lock.s.release(to_release); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + guard.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + write_locked = false, + write_locked.op = "override", + ) + }); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + guard.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + current_readers = 1, + current_readers.op = "add", + ) + }); + + Ok(guard) + } + /// Converts this `OwnedRwLockWriteGuard` into an /// `OwnedRwLockMappedWriteGuard`. This method can be used to store a /// non-mapped guard in a struct field that expects a mapped guard. @@ -178,19 +356,39 @@ impl<T: ?Sized> OwnedRwLockWriteGuard<T> { /// assert_eq!(*lock.read().await, 2, "second writer obtained write lock"); /// # } /// ``` - pub fn downgrade(mut self) -> OwnedRwLockReadGuard<T> { - let lock = unsafe { ManuallyDrop::take(&mut self.lock) }; - let data = self.data; + pub fn downgrade(self) -> OwnedRwLockReadGuard<T> { + let this = self.skip_drop(); + let guard = OwnedRwLockReadGuard { + lock: this.lock, + data: this.data, + _p: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: this.resource_span, + }; // Release all but one of the permits held by the write guard - lock.s.release((self.permits_acquired - 1) as usize); - // NB: Forget to avoid drop impl from being called. - mem::forget(self); - OwnedRwLockReadGuard { - lock: ManuallyDrop::new(lock), - data, - _p: PhantomData, - } + let to_release = (this.permits_acquired - 1) as usize; + guard.lock.s.release(to_release); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + guard.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + write_locked = false, + write_locked.op = "override", + ) + }); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + guard.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + current_readers = 1, + current_readers.op = "add", + ) + }); + + guard } } @@ -229,6 +427,14 @@ where impl<T: ?Sized> Drop for OwnedRwLockWriteGuard<T> { fn drop(&mut self) { self.lock.s.release(self.permits_acquired as usize); - unsafe { ManuallyDrop::drop(&mut self.lock) }; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + write_locked = false, + write_locked.op = "override", + ) + }); } } diff --git a/vendor/tokio/src/sync/rwlock/owned_write_guard_mapped.rs b/vendor/tokio/src/sync/rwlock/owned_write_guard_mapped.rs index 6453236eb..9f4952100 100644 --- a/vendor/tokio/src/sync/rwlock/owned_write_guard_mapped.rs +++ b/vendor/tokio/src/sync/rwlock/owned_write_guard_mapped.rs @@ -1,9 +1,7 @@ use crate::sync::rwlock::RwLock; -use std::fmt; use std::marker::PhantomData; -use std::mem::{self, ManuallyDrop}; -use std::ops; use std::sync::Arc; +use std::{fmt, mem, ops, ptr}; /// Owned RAII structure used to release the exclusive write access of a lock when /// dropped. @@ -14,16 +12,44 @@ use std::sync::Arc; /// /// [mapping]: method@crate::sync::OwnedRwLockWriteGuard::map /// [`OwnedRwLockWriteGuard`]: struct@crate::sync::OwnedRwLockWriteGuard +#[clippy::has_significant_drop] pub struct OwnedRwLockMappedWriteGuard<T: ?Sized, U: ?Sized = T> { + // When changing the fields in this struct, make sure to update the + // `skip_drop` method. + #[cfg(all(tokio_unstable, feature = "tracing"))] + pub(super) resource_span: tracing::Span, pub(super) permits_acquired: u32, - // ManuallyDrop allows us to destructure into this field without running the destructor. - pub(super) lock: ManuallyDrop<Arc<RwLock<T>>>, + pub(super) lock: Arc<RwLock<T>>, pub(super) data: *mut U, pub(super) _p: PhantomData<T>, } +#[allow(dead_code)] // Unused fields are still used in Drop. +struct Inner<T: ?Sized, U: ?Sized> { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, + permits_acquired: u32, + lock: Arc<RwLock<T>>, + data: *const U, +} + impl<T: ?Sized, U: ?Sized> OwnedRwLockMappedWriteGuard<T, U> { - /// Make a new `OwnedRwLockMappedWriteGuard` for a component of the locked + fn skip_drop(self) -> Inner<T, U> { + let me = mem::ManuallyDrop::new(self); + // SAFETY: This duplicates the values in every field of the guard, then + // forgets the originals, so in the end no value is duplicated. + unsafe { + Inner { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: ptr::read(&me.resource_span), + permits_acquired: me.permits_acquired, + lock: ptr::read(&me.lock), + data: me.data, + } + } + } + + /// Makes a new `OwnedRwLockMappedWriteGuard` for a component of the locked /// data. /// /// This operation cannot fail as the `OwnedRwLockMappedWriteGuard` passed @@ -61,15 +87,15 @@ impl<T: ?Sized, U: ?Sized> OwnedRwLockMappedWriteGuard<T, U> { F: FnOnce(&mut U) -> &mut V, { let data = f(&mut *this) as *mut V; - let lock = unsafe { ManuallyDrop::take(&mut this.lock) }; - let permits_acquired = this.permits_acquired; - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); + OwnedRwLockMappedWriteGuard { - permits_acquired, - lock: ManuallyDrop::new(lock), + permits_acquired: this.permits_acquired, + lock: this.lock, data, _p: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: this.resource_span, } } @@ -118,15 +144,15 @@ impl<T: ?Sized, U: ?Sized> OwnedRwLockMappedWriteGuard<T, U> { Some(data) => data as *mut V, None => return Err(this), }; - let lock = unsafe { ManuallyDrop::take(&mut this.lock) }; - let permits_acquired = this.permits_acquired; - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); + Ok(OwnedRwLockMappedWriteGuard { - permits_acquired, - lock: ManuallyDrop::new(lock), + permits_acquired: this.permits_acquired, + lock: this.lock, data, _p: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: this.resource_span, }) } } @@ -166,6 +192,14 @@ where impl<T: ?Sized, U: ?Sized> Drop for OwnedRwLockMappedWriteGuard<T, U> { fn drop(&mut self) { self.lock.s.release(self.permits_acquired as usize); - unsafe { ManuallyDrop::drop(&mut self.lock) }; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + write_locked = false, + write_locked.op = "override", + ) + }); } } diff --git a/vendor/tokio/src/sync/rwlock/read_guard.rs b/vendor/tokio/src/sync/rwlock/read_guard.rs index 38eec7727..a04b59588 100644 --- a/vendor/tokio/src/sync/rwlock/read_guard.rs +++ b/vendor/tokio/src/sync/rwlock/read_guard.rs @@ -1,8 +1,6 @@ use crate::sync::batch_semaphore::Semaphore; -use std::fmt; -use std::marker; -use std::mem; -use std::ops; +use std::marker::PhantomData; +use std::{fmt, mem, ops}; /// RAII structure used to release the shared read access of a lock when /// dropped. @@ -12,14 +10,40 @@ use std::ops; /// /// [`read`]: method@crate::sync::RwLock::read /// [`RwLock`]: struct@crate::sync::RwLock +#[clippy::has_significant_drop] +#[must_use = "if unused the RwLock will immediately unlock"] pub struct RwLockReadGuard<'a, T: ?Sized> { + // When changing the fields in this struct, make sure to update the + // `skip_drop` method. + #[cfg(all(tokio_unstable, feature = "tracing"))] + pub(super) resource_span: tracing::Span, pub(super) s: &'a Semaphore, pub(super) data: *const T, - pub(super) marker: marker::PhantomData<&'a T>, + pub(super) marker: PhantomData<&'a T>, +} + +#[allow(dead_code)] // Unused fields are still used in Drop. +struct Inner<'a, T: ?Sized> { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, + s: &'a Semaphore, + data: *const T, } impl<'a, T: ?Sized> RwLockReadGuard<'a, T> { - /// Make a new `RwLockReadGuard` for a component of the locked data. + fn skip_drop(self) -> Inner<'a, T> { + let me = mem::ManuallyDrop::new(self); + // SAFETY: This duplicates the values in every field of the guard, then + // forgets the originals, so in the end no value is duplicated. + Inner { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: unsafe { std::ptr::read(&me.resource_span) }, + s: me.s, + data: me.data, + } + } + + /// Makes a new `RwLockReadGuard` for a component of the locked data. /// /// This operation cannot fail as the `RwLockReadGuard` passed in already /// locked the data. @@ -58,13 +82,14 @@ impl<'a, T: ?Sized> RwLockReadGuard<'a, T> { F: FnOnce(&T) -> &U, { let data = f(&*this) as *const U; - let s = this.s; - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); + RwLockReadGuard { - s, + s: this.s, data, - marker: marker::PhantomData, + marker: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: this.resource_span, } } @@ -112,13 +137,14 @@ impl<'a, T: ?Sized> RwLockReadGuard<'a, T> { Some(data) => data as *const U, None => return Err(this), }; - let s = this.s; - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); + Ok(RwLockReadGuard { - s, + s: this.s, data, - marker: marker::PhantomData, + marker: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: this.resource_span, }) } } @@ -152,5 +178,14 @@ where impl<'a, T: ?Sized> Drop for RwLockReadGuard<'a, T> { fn drop(&mut self) { self.s.release(1); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + current_readers = 1, + current_readers.op = "sub", + ) + }); } } diff --git a/vendor/tokio/src/sync/rwlock/write_guard.rs b/vendor/tokio/src/sync/rwlock/write_guard.rs index 865a121ed..d405fc2b3 100644 --- a/vendor/tokio/src/sync/rwlock/write_guard.rs +++ b/vendor/tokio/src/sync/rwlock/write_guard.rs @@ -1,10 +1,8 @@ use crate::sync::batch_semaphore::Semaphore; use crate::sync::rwlock::read_guard::RwLockReadGuard; use crate::sync::rwlock::write_guard_mapped::RwLockMappedWriteGuard; -use std::fmt; -use std::marker; -use std::mem; -use std::ops; +use std::marker::PhantomData; +use std::{fmt, mem, ops}; /// RAII structure used to release the exclusive write access of a lock when /// dropped. @@ -14,15 +12,43 @@ use std::ops; /// /// [`write`]: method@crate::sync::RwLock::write /// [`RwLock`]: struct@crate::sync::RwLock +#[clippy::has_significant_drop] +#[must_use = "if unused the RwLock will immediately unlock"] pub struct RwLockWriteGuard<'a, T: ?Sized> { + // When changing the fields in this struct, make sure to update the + // `skip_drop` method. + #[cfg(all(tokio_unstable, feature = "tracing"))] + pub(super) resource_span: tracing::Span, pub(super) permits_acquired: u32, pub(super) s: &'a Semaphore, pub(super) data: *mut T, - pub(super) marker: marker::PhantomData<&'a mut T>, + pub(super) marker: PhantomData<&'a mut T>, +} + +#[allow(dead_code)] // Unused fields are still used in Drop. +struct Inner<'a, T: ?Sized> { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, + permits_acquired: u32, + s: &'a Semaphore, + data: *mut T, } impl<'a, T: ?Sized> RwLockWriteGuard<'a, T> { - /// Make a new [`RwLockMappedWriteGuard`] for a component of the locked data. + fn skip_drop(self) -> Inner<'a, T> { + let me = mem::ManuallyDrop::new(self); + // SAFETY: This duplicates the values in every field of the guard, then + // forgets the originals, so in the end no value is duplicated. + Inner { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: unsafe { std::ptr::read(&me.resource_span) }, + permits_acquired: me.permits_acquired, + s: me.s, + data: me.data, + } + } + + /// Makes a new [`RwLockMappedWriteGuard`] for a component of the locked data. /// /// This operation cannot fail as the `RwLockWriteGuard` passed in already /// locked the data. @@ -64,19 +90,96 @@ impl<'a, T: ?Sized> RwLockWriteGuard<'a, T> { F: FnOnce(&mut T) -> &mut U, { let data = f(&mut *this) as *mut U; - let s = this.s; - let permits_acquired = this.permits_acquired; - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); + RwLockMappedWriteGuard { - permits_acquired, - s, + permits_acquired: this.permits_acquired, + s: this.s, data, - marker: marker::PhantomData, + marker: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: this.resource_span, } } - /// Attempts to make a new [`RwLockMappedWriteGuard`] for a component of + /// Makes a new [`RwLockReadGuard`] for a component of the locked data. + /// + /// This operation cannot fail as the `RwLockWriteGuard` passed in already + /// locked the data. + /// + /// This is an associated function that needs to be used as + /// `RwLockWriteGuard::downgrade_map(..)`. A method would interfere with methods of + /// the same name on the contents of the locked data. + /// + /// This is equivalent to a combination of asynchronous [`RwLockWriteGuard::map`] and [`RwLockWriteGuard::downgrade`] + /// from the [`parking_lot` crate]. + /// + /// Inside of `f`, you retain exclusive access to the data, despite only being given a `&T`. Handing out a + /// `&mut T` would result in unsoundness, as you could use interior mutability. + /// + /// [`RwLockMappedWriteGuard`]: struct@crate::sync::RwLockMappedWriteGuard + /// [`RwLockWriteGuard::map`]: https://docs.rs/lock_api/latest/lock_api/struct.RwLockWriteGuard.html#method.map + /// [`RwLockWriteGuard::downgrade`]: https://docs.rs/lock_api/latest/lock_api/struct.RwLockWriteGuard.html#method.downgrade + /// [`parking_lot` crate]: https://crates.io/crates/parking_lot + /// + /// # Examples + /// + /// ``` + /// use tokio::sync::{RwLock, RwLockWriteGuard}; + /// + /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] + /// struct Foo(u32); + /// + /// # #[tokio::main] + /// # async fn main() { + /// let lock = RwLock::new(Foo(1)); + /// + /// let mapped = RwLockWriteGuard::downgrade_map(lock.write().await, |f| &f.0); + /// let foo = lock.read().await; + /// assert_eq!(foo.0, *mapped); + /// # } + /// ``` + #[inline] + pub fn downgrade_map<F, U: ?Sized>(this: Self, f: F) -> RwLockReadGuard<'a, U> + where + F: FnOnce(&T) -> &U, + { + let data = f(&*this) as *const U; + let this = this.skip_drop(); + let guard = RwLockReadGuard { + s: this.s, + data, + marker: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: this.resource_span, + }; + + // Release all but one of the permits held by the write guard + let to_release = (this.permits_acquired - 1) as usize; + this.s.release(to_release); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + guard.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + write_locked = false, + write_locked.op = "override", + ) + }); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + guard.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + current_readers = 1, + current_readers.op = "add", + ) + }); + + guard + } + + /// Attempts to make a new [`RwLockMappedWriteGuard`] for a component of /// the locked data. The original guard is returned if the closure returns /// `None`. /// @@ -127,18 +230,102 @@ impl<'a, T: ?Sized> RwLockWriteGuard<'a, T> { Some(data) => data as *mut U, None => return Err(this), }; - let s = this.s; - let permits_acquired = this.permits_acquired; - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); + Ok(RwLockMappedWriteGuard { - permits_acquired, - s, + permits_acquired: this.permits_acquired, + s: this.s, data, - marker: marker::PhantomData, + marker: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: this.resource_span, }) } + /// Attempts to make a new [`RwLockReadGuard`] for a component of + /// the locked data. The original guard is returned if the closure returns + /// `None`. + /// + /// This operation cannot fail as the `RwLockWriteGuard` passed in already + /// locked the data. + /// + /// This is an associated function that needs to be + /// used as `RwLockWriteGuard::try_downgrade_map(...)`. A method would interfere with + /// methods of the same name on the contents of the locked data. + /// + /// This is equivalent to a combination of asynchronous [`RwLockWriteGuard::try_map`] and [`RwLockWriteGuard::downgrade`] + /// from the [`parking_lot` crate]. + /// + /// Inside of `f`, you retain exclusive access to the data, despite only being given a `&T`. Handing out a + /// `&mut T` would result in unsoundness, as you could use interior mutability. + /// + /// If this function returns `Err(...)`, the lock is never unlocked nor downgraded. + /// + /// [`RwLockMappedWriteGuard`]: struct@crate::sync::RwLockMappedWriteGuard + /// [`RwLockWriteGuard::map`]: https://docs.rs/lock_api/latest/lock_api/struct.RwLockWriteGuard.html#method.map + /// [`RwLockWriteGuard::downgrade`]: https://docs.rs/lock_api/latest/lock_api/struct.RwLockWriteGuard.html#method.downgrade + /// [`parking_lot` crate]: https://crates.io/crates/parking_lot + /// + /// # Examples + /// + /// ``` + /// use tokio::sync::{RwLock, RwLockWriteGuard}; + /// + /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] + /// struct Foo(u32); + /// + /// # #[tokio::main] + /// # async fn main() { + /// let lock = RwLock::new(Foo(1)); + /// + /// let guard = RwLockWriteGuard::try_downgrade_map(lock.write().await, |f| Some(&f.0)).expect("should not fail"); + /// let foo = lock.read().await; + /// assert_eq!(foo.0, *guard); + /// # } + /// ``` + #[inline] + pub fn try_downgrade_map<F, U: ?Sized>(this: Self, f: F) -> Result<RwLockReadGuard<'a, U>, Self> + where + F: FnOnce(&T) -> Option<&U>, + { + let data = match f(&*this) { + Some(data) => data as *const U, + None => return Err(this), + }; + let this = this.skip_drop(); + let guard = RwLockReadGuard { + s: this.s, + data, + marker: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: this.resource_span, + }; + + // Release all but one of the permits held by the write guard + let to_release = (this.permits_acquired - 1) as usize; + this.s.release(to_release); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + guard.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + write_locked = false, + write_locked.op = "override", + ) + }); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + guard.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + current_readers = 1, + current_readers.op = "add", + ) + }); + + Ok(guard) + } + /// Converts this `RwLockWriteGuard` into an `RwLockMappedWriteGuard`. This /// method can be used to store a non-mapped guard in a struct field that /// expects a mapped guard. @@ -187,17 +374,38 @@ impl<'a, T: ?Sized> RwLockWriteGuard<'a, T> { /// /// [`RwLock`]: struct@crate::sync::RwLock pub fn downgrade(self) -> RwLockReadGuard<'a, T> { - let RwLockWriteGuard { s, data, .. } = self; + let this = self.skip_drop(); + let guard = RwLockReadGuard { + s: this.s, + data: this.data, + marker: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: this.resource_span, + }; // Release all but one of the permits held by the write guard - s.release((self.permits_acquired - 1) as usize); - // NB: Forget to avoid drop impl from being called. - mem::forget(self); - RwLockReadGuard { - s, - data, - marker: marker::PhantomData, - } + let to_release = (this.permits_acquired - 1) as usize; + this.s.release(to_release); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + guard.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + write_locked = false, + write_locked.op = "override", + ) + }); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + guard.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + current_readers = 1, + current_readers.op = "add", + ) + }); + + guard } } @@ -236,5 +444,14 @@ where impl<'a, T: ?Sized> Drop for RwLockWriteGuard<'a, T> { fn drop(&mut self) { self.s.release(self.permits_acquired as usize); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + write_locked = false, + write_locked.op = "override", + ) + }); } } diff --git a/vendor/tokio/src/sync/rwlock/write_guard_mapped.rs b/vendor/tokio/src/sync/rwlock/write_guard_mapped.rs index 9c5b1e7c3..7705189e7 100644 --- a/vendor/tokio/src/sync/rwlock/write_guard_mapped.rs +++ b/vendor/tokio/src/sync/rwlock/write_guard_mapped.rs @@ -1,8 +1,6 @@ use crate::sync::batch_semaphore::Semaphore; -use std::fmt; -use std::marker; -use std::mem; -use std::ops; +use std::marker::PhantomData; +use std::{fmt, mem, ops}; /// RAII structure used to release the exclusive write access of a lock when /// dropped. @@ -13,15 +11,42 @@ use std::ops; /// /// [mapping]: method@crate::sync::RwLockWriteGuard::map /// [`RwLockWriteGuard`]: struct@crate::sync::RwLockWriteGuard +#[clippy::has_significant_drop] pub struct RwLockMappedWriteGuard<'a, T: ?Sized> { + // When changing the fields in this struct, make sure to update the + // `skip_drop` method. + #[cfg(all(tokio_unstable, feature = "tracing"))] + pub(super) resource_span: tracing::Span, pub(super) permits_acquired: u32, pub(super) s: &'a Semaphore, pub(super) data: *mut T, - pub(super) marker: marker::PhantomData<&'a mut T>, + pub(super) marker: PhantomData<&'a mut T>, +} + +#[allow(dead_code)] // Unused fields are still used in Drop. +struct Inner<'a, T: ?Sized> { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, + permits_acquired: u32, + s: &'a Semaphore, + data: *mut T, } impl<'a, T: ?Sized> RwLockMappedWriteGuard<'a, T> { - /// Make a new `RwLockMappedWriteGuard` for a component of the locked data. + fn skip_drop(self) -> Inner<'a, T> { + let me = mem::ManuallyDrop::new(self); + // SAFETY: This duplicates the values in every field of the guard, then + // forgets the originals, so in the end no value is duplicated. + Inner { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: unsafe { std::ptr::read(&me.resource_span) }, + permits_acquired: me.permits_acquired, + s: me.s, + data: me.data, + } + } + + /// Makes a new `RwLockMappedWriteGuard` for a component of the locked data. /// /// This operation cannot fail as the `RwLockMappedWriteGuard` passed in already /// locked the data. @@ -62,15 +87,15 @@ impl<'a, T: ?Sized> RwLockMappedWriteGuard<'a, T> { F: FnOnce(&mut T) -> &mut U, { let data = f(&mut *this) as *mut U; - let s = this.s; - let permits_acquired = this.permits_acquired; - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); + RwLockMappedWriteGuard { - permits_acquired, - s, + permits_acquired: this.permits_acquired, + s: this.s, data, - marker: marker::PhantomData, + marker: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: this.resource_span, } } @@ -124,17 +149,20 @@ impl<'a, T: ?Sized> RwLockMappedWriteGuard<'a, T> { Some(data) => data as *mut U, None => return Err(this), }; - let s = this.s; - let permits_acquired = this.permits_acquired; - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); + Ok(RwLockMappedWriteGuard { - permits_acquired, - s, + permits_acquired: this.permits_acquired, + s: this.s, data, - marker: marker::PhantomData, + marker: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: this.resource_span, }) } + + // Note: No `downgrade`, `downgrade_map` nor `try_downgrade_map` because they would be unsound, as we're already + // potentially been mapped with internal mutability. } impl<T: ?Sized> ops::Deref for RwLockMappedWriteGuard<'_, T> { @@ -172,5 +200,14 @@ where impl<'a, T: ?Sized> Drop for RwLockMappedWriteGuard<'a, T> { fn drop(&mut self) { self.s.release(self.permits_acquired as usize); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + write_locked = false, + write_locked.op = "override", + ) + }); } } diff --git a/vendor/tokio/src/sync/semaphore.rs b/vendor/tokio/src/sync/semaphore.rs index 4b697a9bf..e679d0e6b 100644 --- a/vendor/tokio/src/sync/semaphore.rs +++ b/vendor/tokio/src/sync/semaphore.rs @@ -1,5 +1,7 @@ use super::batch_semaphore as ll; // low level implementation use super::{AcquireError, TryAcquireError}; +#[cfg(all(tokio_unstable, feature = "tracing"))] +use crate::util::trace; use std::sync::Arc; /// Counting semaphore performing asynchronous permit acquisition. @@ -71,12 +73,14 @@ use std::sync::Arc; /// } /// ``` /// -/// [`PollSemaphore`]: https://docs.rs/tokio-util/0.6/tokio_util/sync/struct.PollSemaphore.html +/// [`PollSemaphore`]: https://docs.rs/tokio-util/latest/tokio_util/sync/struct.PollSemaphore.html /// [`Semaphore::acquire_owned`]: crate::sync::Semaphore::acquire_owned #[derive(Debug)] pub struct Semaphore { /// The low level semaphore ll_sem: ll::Semaphore, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, } /// A permit from the semaphore. @@ -85,6 +89,7 @@ pub struct Semaphore { /// /// [`acquire`]: crate::sync::Semaphore::acquire() #[must_use] +#[clippy::has_significant_drop] #[derive(Debug)] pub struct SemaphorePermit<'a> { sem: &'a Semaphore, @@ -97,6 +102,7 @@ pub struct SemaphorePermit<'a> { /// /// [`acquire_owned`]: crate::sync::Semaphore::acquire_owned() #[must_use] +#[clippy::has_significant_drop] #[derive(Debug)] pub struct OwnedSemaphorePermit { sem: Arc<Semaphore>, @@ -119,10 +125,41 @@ fn bounds() { } impl Semaphore { + /// The maximum number of permits which a semaphore can hold. It is `usize::MAX >> 3`. + /// + /// Exceeding this limit typically results in a panic. + pub const MAX_PERMITS: usize = super::batch_semaphore::Semaphore::MAX_PERMITS; + /// Creates a new semaphore with the initial number of permits. + /// + /// Panics if `permits` exceeds [`Semaphore::MAX_PERMITS`]. + #[track_caller] pub fn new(permits: usize) -> Self { + #[cfg(all(tokio_unstable, feature = "tracing"))] + let resource_span = { + let location = std::panic::Location::caller(); + + tracing::trace_span!( + "runtime.resource", + concrete_type = "Semaphore", + kind = "Sync", + loc.file = location.file(), + loc.line = location.line(), + loc.col = location.column(), + inherits_child_attrs = true, + ) + }; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + let ll_sem = resource_span.in_scope(|| ll::Semaphore::new(permits)); + + #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] + let ll_sem = ll::Semaphore::new(permits); + Self { - ll_sem: ll::Semaphore::new(permits), + ll_sem, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span, } } @@ -139,9 +176,16 @@ impl Semaphore { #[cfg(all(feature = "parking_lot", not(all(loom, test))))] #[cfg_attr(docsrs, doc(cfg(feature = "parking_lot")))] pub const fn const_new(permits: usize) -> Self { - Self { + #[cfg(all(tokio_unstable, feature = "tracing"))] + return Self { ll_sem: ll::Semaphore::const_new(permits), - } + resource_span: tracing::Span::none(), + }; + + #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] + return Self { + ll_sem: ll::Semaphore::const_new(permits), + }; } /// Returns the current number of available permits. @@ -151,7 +195,7 @@ impl Semaphore { /// Adds `n` new permits to the semaphore. /// - /// The maximum number of permits is `usize::MAX >> 3`, and this function will panic if the limit is exceeded. + /// The maximum number of permits is [`Semaphore::MAX_PERMITS`], and this function will panic if the limit is exceeded. pub fn add_permits(&self, n: usize) { self.ll_sem.release(n); } @@ -191,9 +235,20 @@ impl Semaphore { /// [`AcquireError`]: crate::sync::AcquireError /// [`SemaphorePermit`]: crate::sync::SemaphorePermit pub async fn acquire(&self) -> Result<SemaphorePermit<'_>, AcquireError> { - self.ll_sem.acquire(1).await?; + #[cfg(all(tokio_unstable, feature = "tracing"))] + let inner = trace::async_op( + || self.ll_sem.acquire(1), + self.resource_span.clone(), + "Semaphore::acquire", + "poll", + true, + ); + #[cfg(not(all(tokio_unstable, feature = "tracing")))] + let inner = self.ll_sem.acquire(1); + + inner.await?; Ok(SemaphorePermit { - sem: &self, + sem: self, permits: 1, }) } @@ -227,9 +282,21 @@ impl Semaphore { /// [`AcquireError`]: crate::sync::AcquireError /// [`SemaphorePermit`]: crate::sync::SemaphorePermit pub async fn acquire_many(&self, n: u32) -> Result<SemaphorePermit<'_>, AcquireError> { + #[cfg(all(tokio_unstable, feature = "tracing"))] + trace::async_op( + || self.ll_sem.acquire(n), + self.resource_span.clone(), + "Semaphore::acquire_many", + "poll", + true, + ) + .await?; + + #[cfg(not(all(tokio_unstable, feature = "tracing")))] self.ll_sem.acquire(n).await?; + Ok(SemaphorePermit { - sem: &self, + sem: self, permits: n, }) } @@ -350,7 +417,18 @@ impl Semaphore { /// [`AcquireError`]: crate::sync::AcquireError /// [`OwnedSemaphorePermit`]: crate::sync::OwnedSemaphorePermit pub async fn acquire_owned(self: Arc<Self>) -> Result<OwnedSemaphorePermit, AcquireError> { - self.ll_sem.acquire(1).await?; + #[cfg(all(tokio_unstable, feature = "tracing"))] + let inner = trace::async_op( + || self.ll_sem.acquire(1), + self.resource_span.clone(), + "Semaphore::acquire_owned", + "poll", + true, + ); + #[cfg(not(all(tokio_unstable, feature = "tracing")))] + let inner = self.ll_sem.acquire(1); + + inner.await?; Ok(OwnedSemaphorePermit { sem: self, permits: 1, @@ -403,7 +481,18 @@ impl Semaphore { self: Arc<Self>, n: u32, ) -> Result<OwnedSemaphorePermit, AcquireError> { - self.ll_sem.acquire(n).await?; + #[cfg(all(tokio_unstable, feature = "tracing"))] + let inner = trace::async_op( + || self.ll_sem.acquire(n), + self.resource_span.clone(), + "Semaphore::acquire_many_owned", + "poll", + true, + ); + #[cfg(not(all(tokio_unstable, feature = "tracing")))] + let inner = self.ll_sem.acquire(n); + + inner.await?; Ok(OwnedSemaphorePermit { sem: self, permits: n, @@ -540,6 +629,25 @@ impl<'a> SemaphorePermit<'a> { pub fn forget(mut self) { self.permits = 0; } + + /// Merge two [`SemaphorePermit`] instances together, consuming `other` + /// without releasing the permits it holds. + /// + /// Permits held by both `self` and `other` are released when `self` drops. + /// + /// # Panics + /// + /// This function panics if permits from different [`Semaphore`] instances + /// are merged. + #[track_caller] + pub fn merge(&mut self, mut other: Self) { + assert!( + std::ptr::eq(self.sem, other.sem), + "merging permits from different semaphore instances" + ); + self.permits += other.permits; + other.permits = 0; + } } impl OwnedSemaphorePermit { @@ -549,9 +657,33 @@ impl OwnedSemaphorePermit { pub fn forget(mut self) { self.permits = 0; } + + /// Merge two [`OwnedSemaphorePermit`] instances together, consuming `other` + /// without releasing the permits it holds. + /// + /// Permits held by both `self` and `other` are released when `self` drops. + /// + /// # Panics + /// + /// This function panics if permits from different [`Semaphore`] instances + /// are merged. + #[track_caller] + pub fn merge(&mut self, mut other: Self) { + assert!( + Arc::ptr_eq(&self.sem, &other.sem), + "merging permits from different semaphore instances" + ); + self.permits += other.permits; + other.permits = 0; + } + + /// Returns the [`Semaphore`] from which this permit was acquired. + pub fn semaphore(&self) -> &Arc<Semaphore> { + &self.sem + } } -impl<'a> Drop for SemaphorePermit<'_> { +impl Drop for SemaphorePermit<'_> { fn drop(&mut self) { self.sem.add_permits(self.permits as usize); } diff --git a/vendor/tokio/src/sync/task/atomic_waker.rs b/vendor/tokio/src/sync/task/atomic_waker.rs index 8616007a3..13aba3544 100644 --- a/vendor/tokio/src/sync/task/atomic_waker.rs +++ b/vendor/tokio/src/sync/task/atomic_waker.rs @@ -1,9 +1,11 @@ #![cfg_attr(any(loom, not(feature = "sync")), allow(dead_code, unreachable_pub))] use crate::loom::cell::UnsafeCell; -use crate::loom::sync::atomic::{self, AtomicUsize}; +use crate::loom::hint; +use crate::loom::sync::atomic::AtomicUsize; use std::fmt; +use std::panic::{resume_unwind, AssertUnwindSafe, RefUnwindSafe, UnwindSafe}; use std::sync::atomic::Ordering::{AcqRel, Acquire, Release}; use std::task::Waker; @@ -27,6 +29,9 @@ pub(crate) struct AtomicWaker { waker: UnsafeCell<Option<Waker>>, } +impl RefUnwindSafe for AtomicWaker {} +impl UnwindSafe for AtomicWaker {} + // `AtomicWaker` is a multi-consumer, single-producer transfer cell. The cell // stores a `Waker` value produced by calls to `register` and many threads can // race to take the waker by calling `wake`. @@ -84,7 +89,7 @@ pub(crate) struct AtomicWaker { // back to `WAITING`. This transition must succeed as, at this point, the state // cannot be transitioned by another thread. // -// If the thread is unable to obtain the lock, the `WAKING` bit is still. +// If the thread is unable to obtain the lock, the `WAKING` bit is still set. // This is because it has either been set by the current thread but the previous // value included the `REGISTERING` bit **or** a concurrent thread is in the // `WAKING` critical section. Either way, no action must be taken. @@ -123,7 +128,7 @@ pub(crate) struct AtomicWaker { // Thread A still holds the `wake` lock, the call to `register` will result // in the task waking itself and get scheduled again. -/// Idle state +/// Idle state. const WAITING: usize = 0; /// A new waker value is being registered with the `AtomicWaker` cell. @@ -171,6 +176,10 @@ impl AtomicWaker { where W: WakerRef, { + fn catch_unwind<F: FnOnce() -> R, R>(f: F) -> std::thread::Result<R> { + std::panic::catch_unwind(AssertUnwindSafe(f)) + } + match self .state .compare_exchange(WAITING, REGISTERING, Acquire, Acquire) @@ -178,8 +187,24 @@ impl AtomicWaker { { WAITING => { unsafe { - // Locked acquired, update the waker cell - self.waker.with_mut(|t| *t = Some(waker.into_waker())); + // If `into_waker` panics (because it's code outside of + // AtomicWaker) we need to prime a guard that is called on + // unwind to restore the waker to a WAITING state. Otherwise + // any future calls to register will incorrectly be stuck + // believing it's being updated by someone else. + let new_waker_or_panic = catch_unwind(move || waker.into_waker()); + + // Set the field to contain the new waker, or if + // `into_waker` panicked, leave the old value. + let mut maybe_panic = None; + let mut old_waker = None; + match new_waker_or_panic { + Ok(new_waker) => { + old_waker = self.waker.with_mut(|t| (*t).take()); + self.waker.with_mut(|t| *t = Some(new_waker)); + } + Err(panic) => maybe_panic = Some(panic), + } // Release the lock. If the state transitioned to include // the `WAKING` bit, this means that a wake has been @@ -193,39 +218,71 @@ impl AtomicWaker { .compare_exchange(REGISTERING, WAITING, AcqRel, Acquire); match res { - Ok(_) => {} + Ok(_) => { + // We don't want to give the caller the panic if it + // was someone else who put in that waker. + let _ = catch_unwind(move || { + drop(old_waker); + }); + } Err(actual) => { // This branch can only be reached if a // concurrent thread called `wake`. In this // case, `actual` **must** be `REGISTERING | - // `WAKING`. + // WAKING`. debug_assert_eq!(actual, REGISTERING | WAKING); // Take the waker to wake once the atomic operation has // completed. - let waker = self.waker.with_mut(|t| (*t).take()).unwrap(); + let mut waker = self.waker.with_mut(|t| (*t).take()); // Just swap, because no one could change state // while state == `Registering | `Waking` self.state.swap(WAITING, AcqRel); - // The atomic swap was complete, now - // wake the waker and return. - waker.wake(); + // If `into_waker` panicked, then the waker in the + // waker slot is actually the old waker. + if maybe_panic.is_some() { + old_waker = waker.take(); + } + + // We don't want to give the caller the panic if it + // was someone else who put in that waker. + if let Some(old_waker) = old_waker { + let _ = catch_unwind(move || { + old_waker.wake(); + }); + } + + // The atomic swap was complete, now wake the waker + // and return. + // + // If this panics, we end up in a consumed state and + // return the panic to the caller. + if let Some(waker) = waker { + debug_assert!(maybe_panic.is_none()); + waker.wake(); + } } } + + if let Some(panic) = maybe_panic { + // If `into_waker` panicked, return the panic to the caller. + resume_unwind(panic); + } } } WAKING => { // Currently in the process of waking the task, i.e., // `wake` is currently being called on the old waker. // So, we call wake on the new waker. + // + // If this panics, someone else is responsible for restoring the + // state of the waker. waker.wake(); // This is equivalent to a spin lock, so use a spin hint. - // TODO: once we bump MSRV to 1.49+, use `hint::spin_loop` instead. - #[allow(deprecated)] - atomic::spin_loop_hint(); + hint::spin_loop(); } state => { // In this case, a concurrent thread is holding the @@ -245,6 +302,8 @@ impl AtomicWaker { /// If `register` has not been called yet, then this does nothing. pub(crate) fn wake(&self) { if let Some(waker) = self.take_waker() { + // If wake panics, we've consumed the waker which is a legitimate + // outcome. waker.wake(); } } diff --git a/vendor/tokio/src/sync/tests/atomic_waker.rs b/vendor/tokio/src/sync/tests/atomic_waker.rs index c832d62e9..8ebfb915f 100644 --- a/vendor/tokio/src/sync/tests/atomic_waker.rs +++ b/vendor/tokio/src/sync/tests/atomic_waker.rs @@ -4,7 +4,7 @@ use tokio_test::task; use std::task::Waker; trait AssertSend: Send {} -trait AssertSync: Send {} +trait AssertSync: Sync {} impl AssertSend for AtomicWaker {} impl AssertSync for AtomicWaker {} @@ -12,6 +12,9 @@ impl AssertSync for AtomicWaker {} impl AssertSend for Waker {} impl AssertSync for Waker {} +#[cfg(tokio_wasm_not_wasi)] +use wasm_bindgen_test::wasm_bindgen_test as test; + #[test] fn basic_usage() { let mut waker = task::spawn(AtomicWaker::new()); @@ -32,3 +35,43 @@ fn wake_without_register() { assert!(!waker.is_woken()); } + +#[test] +#[cfg(not(tokio_wasm))] // wasm currently doesn't support unwinding +fn atomic_waker_panic_safe() { + use std::panic; + use std::ptr; + use std::task::{RawWaker, RawWakerVTable, Waker}; + + static PANICKING_VTABLE: RawWakerVTable = RawWakerVTable::new( + |_| panic!("clone"), + |_| unimplemented!("wake"), + |_| unimplemented!("wake_by_ref"), + |_| (), + ); + + static NONPANICKING_VTABLE: RawWakerVTable = RawWakerVTable::new( + |_| RawWaker::new(ptr::null(), &NONPANICKING_VTABLE), + |_| unimplemented!("wake"), + |_| unimplemented!("wake_by_ref"), + |_| (), + ); + + let panicking = unsafe { Waker::from_raw(RawWaker::new(ptr::null(), &PANICKING_VTABLE)) }; + let nonpanicking = unsafe { Waker::from_raw(RawWaker::new(ptr::null(), &NONPANICKING_VTABLE)) }; + + let atomic_waker = AtomicWaker::new(); + + let panicking = panic::AssertUnwindSafe(&panicking); + + let result = panic::catch_unwind(|| { + let panic::AssertUnwindSafe(panicking) = panicking; + atomic_waker.register_by_ref(panicking); + }); + + assert!(result.is_err()); + assert!(atomic_waker.take_waker().is_none()); + + atomic_waker.register_by_ref(&nonpanicking); + assert!(atomic_waker.take_waker().is_some()); +} diff --git a/vendor/tokio/src/sync/tests/loom_atomic_waker.rs b/vendor/tokio/src/sync/tests/loom_atomic_waker.rs index c148bcbe1..f8bae65d1 100644 --- a/vendor/tokio/src/sync/tests/loom_atomic_waker.rs +++ b/vendor/tokio/src/sync/tests/loom_atomic_waker.rs @@ -43,3 +43,58 @@ fn basic_notification() { })); }); } + +#[test] +fn test_panicky_waker() { + use std::panic; + use std::ptr; + use std::task::{RawWaker, RawWakerVTable, Waker}; + + static PANICKING_VTABLE: RawWakerVTable = + RawWakerVTable::new(|_| panic!("clone"), |_| (), |_| (), |_| ()); + + let panicking = unsafe { Waker::from_raw(RawWaker::new(ptr::null(), &PANICKING_VTABLE)) }; + + // If you're working with this test (and I sure hope you never have to!), + // uncomment the following section because there will be a lot of panics + // which would otherwise log. + // + // We can't however leaved it uncommented, because it's global. + // panic::set_hook(Box::new(|_| ())); + + const NUM_NOTIFY: usize = 2; + + loom::model(move || { + let chan = Arc::new(Chan { + num: AtomicUsize::new(0), + task: AtomicWaker::new(), + }); + + for _ in 0..NUM_NOTIFY { + let chan = chan.clone(); + + thread::spawn(move || { + chan.num.fetch_add(1, Relaxed); + chan.task.wake(); + }); + } + + // Note: this panic should have no effect on the overall state of the + // waker and it should proceed as normal. + // + // A thread above might race to flag a wakeup, and a WAKING state will + // be preserved if this expected panic races with that so the below + // procedure should be allowed to continue uninterrupted. + let _ = panic::catch_unwind(|| chan.task.register_by_ref(&panicking)); + + block_on(poll_fn(move |cx| { + chan.task.register_by_ref(cx.waker()); + + if NUM_NOTIFY == chan.num.load(Relaxed) { + return Ready(()); + } + + Pending + })); + }); +} diff --git a/vendor/tokio/src/sync/tests/loom_mpsc.rs b/vendor/tokio/src/sync/tests/loom_mpsc.rs index c12313bd3..f165e7076 100644 --- a/vendor/tokio/src/sync/tests/loom_mpsc.rs +++ b/vendor/tokio/src/sync/tests/loom_mpsc.rs @@ -132,3 +132,59 @@ fn dropping_unbounded_tx() { assert!(v.is_none()); }); } + +#[test] +fn try_recv() { + loom::model(|| { + use crate::sync::{mpsc, Semaphore}; + use loom::sync::{Arc, Mutex}; + + const PERMITS: usize = 2; + const TASKS: usize = 2; + const CYCLES: usize = 1; + + struct Context { + sem: Arc<Semaphore>, + tx: mpsc::Sender<()>, + rx: Mutex<mpsc::Receiver<()>>, + } + + fn run(ctx: &Context) { + block_on(async { + let permit = ctx.sem.acquire().await; + assert_ok!(ctx.rx.lock().unwrap().try_recv()); + crate::task::yield_now().await; + assert_ok!(ctx.tx.clone().try_send(())); + drop(permit); + }); + } + + let (tx, rx) = mpsc::channel(PERMITS); + let sem = Arc::new(Semaphore::new(PERMITS)); + let ctx = Arc::new(Context { + sem, + tx, + rx: Mutex::new(rx), + }); + + for _ in 0..PERMITS { + assert_ok!(ctx.tx.clone().try_send(())); + } + + let mut ths = Vec::new(); + + for _ in 0..TASKS { + let ctx = ctx.clone(); + + ths.push(thread::spawn(move || { + run(&ctx); + })); + } + + run(&ctx); + + for th in ths { + th.join().unwrap(); + } + }); +} diff --git a/vendor/tokio/src/sync/tests/loom_notify.rs b/vendor/tokio/src/sync/tests/loom_notify.rs index d484a7581..a4ded1d35 100644 --- a/vendor/tokio/src/sync/tests/loom_notify.rs +++ b/vendor/tokio/src/sync/tests/loom_notify.rs @@ -4,6 +4,11 @@ use loom::future::block_on; use loom::sync::Arc; use loom::thread; +use tokio_test::{assert_pending, assert_ready}; + +/// `util::wake_list::NUM_WAKERS` +const WAKE_LIST_SIZE: usize = 32; + #[test] fn notify_one() { loom::model(|| { @@ -138,3 +143,189 @@ fn notify_drop() { th2.join().unwrap(); }); } + +/// Polls two `Notified` futures and checks if poll results are consistent +/// with each other. If the first future is notified by a `notify_waiters` +/// call, then the second one must be notified as well. +#[test] +fn notify_waiters_poll_consistency() { + fn notify_waiters_poll_consistency_variant(poll_setting: [bool; 2]) { + let notify = Arc::new(Notify::new()); + let mut notified = [ + tokio_test::task::spawn(notify.notified()), + tokio_test::task::spawn(notify.notified()), + ]; + for i in 0..2 { + if poll_setting[i] { + assert_pending!(notified[i].poll()); + } + } + + let tx = notify.clone(); + let th = thread::spawn(move || { + tx.notify_waiters(); + }); + + let res1 = notified[0].poll(); + let res2 = notified[1].poll(); + + // If res1 is ready, then res2 must also be ready. + assert!(res1.is_pending() || res2.is_ready()); + + th.join().unwrap(); + } + + // We test different scenarios in which pending futures had or had not + // been polled before the call to `notify_waiters`. + loom::model(|| notify_waiters_poll_consistency_variant([false, false])); + loom::model(|| notify_waiters_poll_consistency_variant([true, false])); + loom::model(|| notify_waiters_poll_consistency_variant([false, true])); + loom::model(|| notify_waiters_poll_consistency_variant([true, true])); +} + +/// Polls two `Notified` futures and checks if poll results are consistent +/// with each other. If the first future is notified by a `notify_waiters` +/// call, then the second one must be notified as well. +/// +/// Here we also add other `Notified` futures in between to force the two +/// tested futures to end up in different chunks. +#[test] +fn notify_waiters_poll_consistency_many() { + fn notify_waiters_poll_consistency_many_variant(order: [usize; 2]) { + let notify = Arc::new(Notify::new()); + + let mut futs = (0..WAKE_LIST_SIZE + 1) + .map(|_| tokio_test::task::spawn(notify.notified())) + .collect::<Vec<_>>(); + + assert_pending!(futs[order[0]].poll()); + for i in 2..futs.len() { + assert_pending!(futs[i].poll()); + } + assert_pending!(futs[order[1]].poll()); + + let tx = notify.clone(); + let th = thread::spawn(move || { + tx.notify_waiters(); + }); + + let res1 = futs[0].poll(); + let res2 = futs[1].poll(); + + // If res1 is ready, then res2 must also be ready. + assert!(res1.is_pending() || res2.is_ready()); + + th.join().unwrap(); + } + + // We test different scenarios in which futures are polled in different order. + loom::model(|| notify_waiters_poll_consistency_many_variant([0, 1])); + loom::model(|| notify_waiters_poll_consistency_many_variant([1, 0])); +} + +/// Checks if a call to `notify_waiters` is observed as atomic when combined +/// with a concurrent call to `notify_one`. +#[test] +fn notify_waiters_is_atomic() { + fn notify_waiters_is_atomic_variant(tested_fut_index: usize) { + let notify = Arc::new(Notify::new()); + + let mut futs = (0..WAKE_LIST_SIZE + 1) + .map(|_| tokio_test::task::spawn(notify.notified())) + .collect::<Vec<_>>(); + + for fut in &mut futs { + assert_pending!(fut.poll()); + } + + let tx = notify.clone(); + let th = thread::spawn(move || { + tx.notify_waiters(); + }); + + block_on(async { + // If awaiting one of the futures completes, then we should be + // able to assume that all pending futures are notified. Therefore + // a notification from a subsequent `notify_one` call should not + // be consumed by an old future. + futs.remove(tested_fut_index).await; + + let mut new_fut = tokio_test::task::spawn(notify.notified()); + assert_pending!(new_fut.poll()); + + notify.notify_one(); + + // `new_fut` must consume the notification from `notify_one`. + assert_ready!(new_fut.poll()); + }); + + th.join().unwrap(); + } + + // We test different scenarios in which the tested future is at the beginning + // or at the end of the waiters queue used by `Notify`. + loom::model(|| notify_waiters_is_atomic_variant(0)); + loom::model(|| notify_waiters_is_atomic_variant(32)); +} + +/// Checks if a single call to `notify_waiters` does not get through two `Notified` +/// futures created and awaited sequentially like this: +/// ```ignore +/// notify.notified().await; +/// notify.notified().await; +/// ``` +#[test] +fn notify_waiters_sequential_notified_await() { + use crate::sync::oneshot; + + loom::model(|| { + let notify = Arc::new(Notify::new()); + + let (tx_fst, rx_fst) = oneshot::channel(); + let (tx_snd, rx_snd) = oneshot::channel(); + + let receiver = thread::spawn({ + let notify = notify.clone(); + move || { + block_on(async { + // Poll the first `Notified` to put it as the first waiter + // in the queue. + let mut first_notified = tokio_test::task::spawn(notify.notified()); + assert_pending!(first_notified.poll()); + + // Create additional waiters to force `notify_waiters` to + // release the lock at least once. + let _task_pile = (0..WAKE_LIST_SIZE + 1) + .map(|_| { + let mut fut = tokio_test::task::spawn(notify.notified()); + assert_pending!(fut.poll()); + fut + }) + .collect::<Vec<_>>(); + + // We are ready for the notify_waiters call. + tx_fst.send(()).unwrap(); + + first_notified.await; + + // Poll the second `Notified` future to try to insert + // it to the waiters queue. + let mut second_notified = tokio_test::task::spawn(notify.notified()); + assert_pending!(second_notified.poll()); + + // Wait for the `notify_waiters` to end and check if we + // are woken up. + rx_snd.await.unwrap(); + assert_pending!(second_notified.poll()); + }); + } + }); + + // Wait for the signal and call `notify_waiters`. + block_on(rx_fst).unwrap(); + notify.notify_waiters(); + tx_snd.send(()).unwrap(); + + receiver.join().unwrap(); + }); +} diff --git a/vendor/tokio/src/sync/tests/loom_watch.rs b/vendor/tokio/src/sync/tests/loom_watch.rs index c575b5b66..51589cd80 100644 --- a/vendor/tokio/src/sync/tests/loom_watch.rs +++ b/vendor/tokio/src/sync/tests/loom_watch.rs @@ -2,6 +2,7 @@ use crate::sync::watch; use loom::future::block_on; use loom::thread; +use std::sync::Arc; #[test] fn smoke() { @@ -34,3 +35,56 @@ fn smoke() { th.join().unwrap(); }) } + +#[test] +fn wait_for_test() { + loom::model(move || { + let (tx, mut rx) = watch::channel(false); + + let tx_arc = Arc::new(tx); + let tx1 = tx_arc.clone(); + let tx2 = tx_arc.clone(); + + let th1 = thread::spawn(move || { + for _ in 0..2 { + tx1.send_modify(|_x| {}); + } + }); + + let th2 = thread::spawn(move || { + tx2.send(true).unwrap(); + }); + + assert_eq!(*block_on(rx.wait_for(|x| *x)).unwrap(), true); + + th1.join().unwrap(); + th2.join().unwrap(); + }); +} + +#[test] +fn wait_for_returns_correct_value() { + loom::model(move || { + let (tx, mut rx) = watch::channel(0); + + let jh = thread::spawn(move || { + tx.send(1).unwrap(); + tx.send(2).unwrap(); + tx.send(3).unwrap(); + }); + + // Stop at the first value we are called at. + let mut stopped_at = usize::MAX; + let returned = *block_on(rx.wait_for(|x| { + stopped_at = *x; + true + })) + .unwrap(); + + // Check that it returned the same value as the one we returned + // `true` for. + assert_eq!(stopped_at, returned); + + jh.join().unwrap(); + }); +} diff --git a/vendor/tokio/src/sync/tests/mod.rs b/vendor/tokio/src/sync/tests/mod.rs index c5d560196..ee76418ac 100644 --- a/vendor/tokio/src/sync/tests/mod.rs +++ b/vendor/tokio/src/sync/tests/mod.rs @@ -1,5 +1,6 @@ cfg_not_loom! { mod atomic_waker; + mod notify; mod semaphore_batch; } diff --git a/vendor/tokio/src/sync/tests/notify.rs b/vendor/tokio/src/sync/tests/notify.rs new file mode 100644 index 000000000..4b5989597 --- /dev/null +++ b/vendor/tokio/src/sync/tests/notify.rs @@ -0,0 +1,120 @@ +use crate::sync::Notify; +use std::future::Future; +use std::mem::ManuallyDrop; +use std::sync::Arc; +use std::task::{Context, RawWaker, RawWakerVTable, Waker}; + +#[cfg(tokio_wasm_not_wasi)] +use wasm_bindgen_test::wasm_bindgen_test as test; + +#[test] +fn notify_clones_waker_before_lock() { + const VTABLE: &RawWakerVTable = &RawWakerVTable::new(clone_w, wake, wake_by_ref, drop_w); + + unsafe fn clone_w(data: *const ()) -> RawWaker { + let arc = ManuallyDrop::new(Arc::<Notify>::from_raw(data as *const Notify)); + // Or some other arbitrary code that shouldn't be executed while the + // Notify wait list is locked. + arc.notify_one(); + let _arc_clone: ManuallyDrop<_> = arc.clone(); + RawWaker::new(data, VTABLE) + } + + unsafe fn drop_w(data: *const ()) { + let _ = Arc::<Notify>::from_raw(data as *const Notify); + } + + unsafe fn wake(_data: *const ()) { + unreachable!() + } + + unsafe fn wake_by_ref(_data: *const ()) { + unreachable!() + } + + let notify = Arc::new(Notify::new()); + let notify2 = notify.clone(); + + let waker = + unsafe { Waker::from_raw(RawWaker::new(Arc::into_raw(notify2) as *const _, VTABLE)) }; + let mut cx = Context::from_waker(&waker); + + let future = notify.notified(); + pin!(future); + + // The result doesn't matter, we're just testing that we don't deadlock. + let _ = future.poll(&mut cx); +} + +#[cfg(panic = "unwind")] +#[test] +fn notify_waiters_handles_panicking_waker() { + use futures::task::ArcWake; + + let notify = Arc::new(Notify::new()); + + struct PanickingWaker(Arc<Notify>); + + impl ArcWake for PanickingWaker { + fn wake_by_ref(_arc_self: &Arc<Self>) { + panic!("waker panicked"); + } + } + + let bad_fut = notify.notified(); + pin!(bad_fut); + + let waker = futures::task::waker(Arc::new(PanickingWaker(notify.clone()))); + let mut cx = Context::from_waker(&waker); + let _ = bad_fut.poll(&mut cx); + + let mut futs = Vec::new(); + for _ in 0..32 { + let mut fut = tokio_test::task::spawn(notify.notified()); + assert!(fut.poll().is_pending()); + futs.push(fut); + } + + assert!(std::panic::catch_unwind(|| { + notify.notify_waiters(); + }) + .is_err()); + + for mut fut in futs { + assert!(fut.poll().is_ready()); + } +} + +#[test] +fn notify_simple() { + let notify = Notify::new(); + + let mut fut1 = tokio_test::task::spawn(notify.notified()); + assert!(fut1.poll().is_pending()); + + let mut fut2 = tokio_test::task::spawn(notify.notified()); + assert!(fut2.poll().is_pending()); + + notify.notify_waiters(); + + assert!(fut1.poll().is_ready()); + assert!(fut2.poll().is_ready()); +} + +#[test] +#[cfg(not(tokio_wasm))] +fn watch_test() { + let rt = crate::runtime::Builder::new_current_thread() + .build() + .unwrap(); + + rt.block_on(async { + let (tx, mut rx) = crate::sync::watch::channel(()); + + crate::spawn(async move { + let _ = tx.send(()); + }); + + let _ = rx.changed().await; + }); +} diff --git a/vendor/tokio/src/sync/tests/semaphore_batch.rs b/vendor/tokio/src/sync/tests/semaphore_batch.rs index 9342cd1cb..85089cd22 100644 --- a/vendor/tokio/src/sync/tests/semaphore_batch.rs +++ b/vendor/tokio/src/sync/tests/semaphore_batch.rs @@ -1,6 +1,11 @@ use crate::sync::batch_semaphore::Semaphore; use tokio_test::*; +const MAX_PERMITS: usize = crate::sync::Semaphore::MAX_PERMITS; + +#[cfg(tokio_wasm_not_wasi)] +use wasm_bindgen_test::wasm_bindgen_test as test; + #[test] fn poll_acquire_one_available() { let s = Semaphore::new(100); @@ -166,10 +171,15 @@ fn poll_acquire_one_zero_permits() { } #[test] +fn max_permits_doesnt_panic() { + Semaphore::new(MAX_PERMITS); +} + +#[test] #[should_panic] +#[cfg(not(tokio_wasm))] // wasm currently doesn't support unwinding fn validates_max_permits() { - use std::usize; - Semaphore::new((usize::MAX >> 2) + 1); + Semaphore::new(MAX_PERMITS + 1); } #[test] @@ -248,3 +258,32 @@ fn cancel_acquire_releases_permits() { assert_eq!(6, s.available_permits()); assert_ok!(s.try_acquire(6)); } + +#[test] +fn release_permits_at_drop() { + use crate::sync::semaphore::*; + use futures::task::ArcWake; + use std::future::Future; + use std::sync::Arc; + + let sem = Arc::new(Semaphore::new(1)); + + struct ReleaseOnDrop(Option<OwnedSemaphorePermit>); + + impl ArcWake for ReleaseOnDrop { + fn wake_by_ref(_arc_self: &Arc<Self>) {} + } + + let mut fut = Box::pin(async { + let _permit = sem.acquire().await.unwrap(); + }); + + // Second iteration shouldn't deadlock. + for _ in 0..=1 { + let waker = futures::task::waker(Arc::new(ReleaseOnDrop( + sem.clone().try_acquire_owned().ok(), + ))); + let mut cx = std::task::Context::from_waker(&waker); + assert!(fut.as_mut().poll(&mut cx).is_pending()); + } +} diff --git a/vendor/tokio/src/sync/watch.rs b/vendor/tokio/src/sync/watch.rs index 7852b0cb1..61049b71e 100644 --- a/vendor/tokio/src/sync/watch.rs +++ b/vendor/tokio/src/sync/watch.rs @@ -9,10 +9,10 @@ //! # Usage //! //! [`channel`] returns a [`Sender`] / [`Receiver`] pair. These are the producer -//! and sender halves of the channel. The channel is created with an initial +//! and consumer halves of the channel. The channel is created with an initial //! value. The **latest** value stored in the channel is accessed with //! [`Receiver::borrow()`]. Awaiting [`Receiver::changed()`] waits for a new -//! value to sent by the [`Sender`] half. +//! value to be sent by the [`Sender`] half. //! //! # Examples //! @@ -20,15 +20,15 @@ //! use tokio::sync::watch; //! //! # async fn dox() -> Result<(), Box<dyn std::error::Error>> { -//! let (tx, mut rx) = watch::channel("hello"); +//! let (tx, mut rx) = watch::channel("hello"); //! -//! tokio::spawn(async move { -//! while rx.changed().await.is_ok() { -//! println!("received = {:?}", *rx.borrow()); -//! } -//! }); +//! tokio::spawn(async move { +//! while rx.changed().await.is_ok() { +//! println!("received = {:?}", *rx.borrow()); +//! } +//! }); //! -//! tx.send("world")?; +//! tx.send("world")?; //! # Ok(()) //! # } //! ``` @@ -39,6 +39,9 @@ //! when all [`Receiver`] handles have been dropped. This indicates that there //! is no further interest in the values being produced and work can be stopped. //! +//! The value in the channel will not be dropped until the sender and all receivers +//! have been dropped. +//! //! # Thread safety //! //! Both [`Sender`] and [`Receiver`] are thread safe. They can be moved to other @@ -56,9 +59,12 @@ use crate::sync::notify::Notify; use crate::loom::sync::atomic::AtomicUsize; -use crate::loom::sync::atomic::Ordering::{Relaxed, SeqCst}; +use crate::loom::sync::atomic::Ordering::Relaxed; use crate::loom::sync::{Arc, RwLock, RwLockReadGuard}; +use std::fmt; +use std::mem; use std::ops; +use std::panic; /// Receives values from the associated [`Sender`](struct@Sender). /// @@ -74,7 +80,7 @@ pub struct Receiver<T> { shared: Arc<Shared<T>>, /// Last observed version - version: usize, + version: Version, } /// Sends values to the associated [`Receiver`](struct@Receiver). @@ -85,60 +91,144 @@ pub struct Sender<T> { shared: Arc<Shared<T>>, } -/// Returns a reference to the inner value +/// Returns a reference to the inner value. /// /// Outstanding borrows hold a read lock on the inner value. This means that -/// long lived borrows could cause the produce half to block. It is recommended -/// to keep the borrow as short lived as possible. +/// long-lived borrows could cause the producer half to block. It is recommended +/// to keep the borrow as short-lived as possible. Additionally, if you are +/// running in an environment that allows `!Send` futures, you must ensure that +/// the returned `Ref` type is never held alive across an `.await` point, +/// otherwise, it can lead to a deadlock. +/// +/// The priority policy of the lock is dependent on the underlying lock +/// implementation, and this type does not guarantee that any particular policy +/// will be used. In particular, a producer which is waiting to acquire the lock +/// in `send` might or might not block concurrent calls to `borrow`, e.g.: +/// +/// <details><summary>Potential deadlock example</summary> +/// +/// ```text +/// // Task 1 (on thread A) | // Task 2 (on thread B) +/// let _ref1 = rx.borrow(); | +/// | // will block +/// | let _ = tx.send(()); +/// // may deadlock | +/// let _ref2 = rx.borrow(); | +/// ``` +/// </details> #[derive(Debug)] pub struct Ref<'a, T> { inner: RwLockReadGuard<'a, T>, + has_changed: bool, +} + +impl<'a, T> Ref<'a, T> { + /// Indicates if the borrowed value is considered as _changed_ since the last + /// time it has been marked as seen. + /// + /// Unlike [`Receiver::has_changed()`], this method does not fail if the channel is closed. + /// + /// When borrowed from the [`Sender`] this function will always return `false`. + /// + /// # Examples + /// + /// ``` + /// use tokio::sync::watch; + /// + /// #[tokio::main] + /// async fn main() { + /// let (tx, mut rx) = watch::channel("hello"); + /// + /// tx.send("goodbye").unwrap(); + /// // The sender does never consider the value as changed. + /// assert!(!tx.borrow().has_changed()); + /// + /// // Drop the sender immediately, just for testing purposes. + /// drop(tx); + /// + /// // Even if the sender has already been dropped... + /// assert!(rx.has_changed().is_err()); + /// // ...the modified value is still readable and detected as changed. + /// assert_eq!(*rx.borrow(), "goodbye"); + /// assert!(rx.borrow().has_changed()); + /// + /// // Read the changed value and mark it as seen. + /// { + /// let received = rx.borrow_and_update(); + /// assert_eq!(*received, "goodbye"); + /// assert!(received.has_changed()); + /// // Release the read lock when leaving this scope. + /// } + /// + /// // Now the value has already been marked as seen and could + /// // never be modified again (after the sender has been dropped). + /// assert!(!rx.borrow().has_changed()); + /// } + /// ``` + pub fn has_changed(&self) -> bool { + self.has_changed + } } -#[derive(Debug)] struct Shared<T> { - /// The most recent value + /// The most recent value. value: RwLock<T>, - /// The current version + /// The current version. /// /// The lowest bit represents a "closed" state. The rest of the bits /// represent the current version. - version: AtomicUsize, + state: AtomicState, - /// Tracks the number of `Receiver` instances + /// Tracks the number of `Receiver` instances. ref_count_rx: AtomicUsize, /// Notifies waiting receivers that the value changed. - notify_rx: Notify, + notify_rx: big_notify::BigNotify, - /// Notifies any task listening for `Receiver` dropped events + /// Notifies any task listening for `Receiver` dropped events. notify_tx: Notify, } +impl<T: fmt::Debug> fmt::Debug for Shared<T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let state = self.state.load(); + f.debug_struct("Shared") + .field("value", &self.value) + .field("version", &state.version()) + .field("is_closed", &state.is_closed()) + .field("ref_count_rx", &self.ref_count_rx) + .finish() + } +} + pub mod error { - //! Watch error types + //! Watch error types. use std::fmt; /// Error produced when sending a value fails. - #[derive(Debug)] - pub struct SendError<T> { - pub(crate) inner: T, - } + #[derive(PartialEq, Eq, Clone, Copy)] + pub struct SendError<T>(pub T); // ===== impl SendError ===== - impl<T: fmt::Debug> fmt::Display for SendError<T> { + impl<T> fmt::Debug for SendError<T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("SendError").finish_non_exhaustive() + } + } + + impl<T> fmt::Display for SendError<T> { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { write!(fmt, "channel closed") } } - impl<T: fmt::Debug> std::error::Error for SendError<T> {} + impl<T> std::error::Error for SendError<T> {} /// Error produced when receiving a change notification. - #[derive(Debug)] + #[derive(Debug, Clone)] pub struct RecvError(pub(super) ()); // ===== impl RecvError ===== @@ -152,7 +242,128 @@ pub mod error { impl std::error::Error for RecvError {} } -const CLOSED: usize = 1; +mod big_notify { + use super::*; + use crate::sync::notify::Notified; + + // To avoid contention on the lock inside the `Notify`, we store multiple + // copies of it. Then, we use either circular access or randomness to spread + // out threads over different `Notify` objects. + // + // Some simple benchmarks show that randomness performs slightly better than + // circular access (probably due to contention on `next`), so we prefer to + // use randomness when Tokio is compiled with a random number generator. + // + // When the random number generator is not available, we fall back to + // circular access. + + pub(super) struct BigNotify { + #[cfg(not(all(not(loom), feature = "sync", any(feature = "rt", feature = "macros"))))] + next: AtomicUsize, + inner: [Notify; 8], + } + + impl BigNotify { + pub(super) fn new() -> Self { + Self { + #[cfg(not(all( + not(loom), + feature = "sync", + any(feature = "rt", feature = "macros") + )))] + next: AtomicUsize::new(0), + inner: Default::default(), + } + } + + pub(super) fn notify_waiters(&self) { + for notify in &self.inner { + notify.notify_waiters(); + } + } + + /// This function implements the case where randomness is not available. + #[cfg(not(all(not(loom), feature = "sync", any(feature = "rt", feature = "macros"))))] + pub(super) fn notified(&self) -> Notified<'_> { + let i = self.next.fetch_add(1, Relaxed) % 8; + self.inner[i].notified() + } + + /// This function implements the case where randomness is available. + #[cfg(all(not(loom), feature = "sync", any(feature = "rt", feature = "macros")))] + pub(super) fn notified(&self) -> Notified<'_> { + let i = crate::runtime::context::thread_rng_n(8) as usize; + self.inner[i].notified() + } + } +} + +use self::state::{AtomicState, Version}; +mod state { + use crate::loom::sync::atomic::AtomicUsize; + use crate::loom::sync::atomic::Ordering::SeqCst; + + const CLOSED: usize = 1; + + /// The version part of the state. The lowest bit is always zero. + #[derive(Copy, Clone, Debug, Eq, PartialEq)] + pub(super) struct Version(usize); + + /// Snapshot of the state. The first bit is used as the CLOSED bit. + /// The remaining bits are used as the version. + /// + /// The CLOSED bit tracks whether the Sender has been dropped. Dropping all + /// receivers does not set it. + #[derive(Copy, Clone, Debug)] + pub(super) struct StateSnapshot(usize); + + /// The state stored in an atomic integer. + #[derive(Debug)] + pub(super) struct AtomicState(AtomicUsize); + + impl Version { + /// Get the initial version when creating the channel. + pub(super) fn initial() -> Self { + Version(0) + } + } + + impl StateSnapshot { + /// Extract the version from the state. + pub(super) fn version(self) -> Version { + Version(self.0 & !CLOSED) + } + + /// Is the closed bit set? + pub(super) fn is_closed(self) -> bool { + (self.0 & CLOSED) == CLOSED + } + } + + impl AtomicState { + /// Create a new `AtomicState` that is not closed and which has the + /// version set to `Version::initial()`. + pub(super) fn new() -> Self { + AtomicState(AtomicUsize::new(0)) + } + + /// Load the current value of the state. + pub(super) fn load(&self) -> StateSnapshot { + StateSnapshot(self.0.load(SeqCst)) + } + + /// Increment the version counter. + pub(super) fn increment_version(&self) { + // Increment by two to avoid touching the CLOSED bit. + self.0.fetch_add(2, SeqCst); + } + + /// Set the closed bit in the state. + pub(super) fn set_closed(&self) { + self.0.fetch_or(CLOSED, SeqCst); + } + } +} /// Creates a new watch channel, returning the "send" and "receive" handles. /// @@ -184,9 +395,9 @@ const CLOSED: usize = 1; pub fn channel<T>(init: T) -> (Sender<T>, Receiver<T>) { let shared = Arc::new(Shared { value: RwLock::new(init), - version: AtomicUsize::new(0), + state: AtomicState::new(), ref_count_rx: AtomicUsize::new(1), - notify_rx: Notify::new(), + notify_rx: big_notify::BigNotify::new(), notify_tx: Notify::new(), }); @@ -194,13 +405,16 @@ pub fn channel<T>(init: T) -> (Sender<T>, Receiver<T>) { shared: shared.clone(), }; - let rx = Receiver { shared, version: 0 }; + let rx = Receiver { + shared, + version: Version::initial(), + }; (tx, rx) } impl<T> Receiver<T> { - fn from_shared(version: usize, shared: Arc<Shared<T>>) -> Self { + fn from_shared(version: Version, shared: Arc<Shared<T>>) -> Self { // No synchronization necessary as this is only used as a counter and // not memory access. shared.ref_count_rx.fetch_add(1, Relaxed); @@ -214,9 +428,29 @@ impl<T> Receiver<T> { /// [`changed`] may return immediately even if you have already seen the /// value with a call to `borrow`. /// - /// Outstanding borrows hold a read lock. This means that long lived borrows - /// could cause the send half to block. It is recommended to keep the borrow - /// as short lived as possible. + /// Outstanding borrows hold a read lock on the inner value. This means that + /// long-lived borrows could cause the producer half to block. It is recommended + /// to keep the borrow as short-lived as possible. Additionally, if you are + /// running in an environment that allows `!Send` futures, you must ensure that + /// the returned `Ref` type is never held alive across an `.await` point, + /// otherwise, it can lead to a deadlock. + /// + /// The priority policy of the lock is dependent on the underlying lock + /// implementation, and this type does not guarantee that any particular policy + /// will be used. In particular, a producer which is waiting to acquire the lock + /// in `send` might or might not block concurrent calls to `borrow`, e.g.: + /// + /// <details><summary>Potential deadlock example</summary> + /// + /// ```text + /// // Task 1 (on thread A) | // Task 2 (on thread B) + /// let _ref1 = rx.borrow(); | + /// | // will block + /// | let _ = tx.send(()); + /// // may deadlock | + /// let _ref2 = rx.borrow(); | + /// ``` + /// </details> /// /// [`changed`]: Receiver::changed /// @@ -230,28 +464,104 @@ impl<T> Receiver<T> { /// ``` pub fn borrow(&self) -> Ref<'_, T> { let inner = self.shared.value.read().unwrap(); - Ref { inner } + + // After obtaining a read-lock no concurrent writes could occur + // and the loaded version matches that of the borrowed reference. + let new_version = self.shared.state.load().version(); + let has_changed = self.version != new_version; + + Ref { inner, has_changed } } - /// Returns a reference to the most recently sent value and mark that value + /// Returns a reference to the most recently sent value and marks that value /// as seen. /// - /// This method marks the value as seen, so [`changed`] will not return - /// immediately if the newest value is one previously returned by - /// `borrow_and_update`. + /// This method marks the current value as seen. Subsequent calls to [`changed`] + /// will not return immediately until the [`Sender`] has modified the shared + /// value again. + /// + /// Outstanding borrows hold a read lock on the inner value. This means that + /// long-lived borrows could cause the producer half to block. It is recommended + /// to keep the borrow as short-lived as possible. Additionally, if you are + /// running in an environment that allows `!Send` futures, you must ensure that + /// the returned `Ref` type is never held alive across an `.await` point, + /// otherwise, it can lead to a deadlock. + /// + /// The priority policy of the lock is dependent on the underlying lock + /// implementation, and this type does not guarantee that any particular policy + /// will be used. In particular, a producer which is waiting to acquire the lock + /// in `send` might or might not block concurrent calls to `borrow`, e.g.: /// - /// Outstanding borrows hold a read lock. This means that long lived borrows - /// could cause the send half to block. It is recommended to keep the borrow - /// as short lived as possible. + /// <details><summary>Potential deadlock example</summary> + /// + /// ```text + /// // Task 1 (on thread A) | // Task 2 (on thread B) + /// let _ref1 = rx1.borrow_and_update(); | + /// | // will block + /// | let _ = tx.send(()); + /// // may deadlock | + /// let _ref2 = rx2.borrow_and_update(); | + /// ``` + /// </details> /// /// [`changed`]: Receiver::changed pub fn borrow_and_update(&mut self) -> Ref<'_, T> { let inner = self.shared.value.read().unwrap(); - self.version = self.shared.version.load(SeqCst) & !CLOSED; - Ref { inner } + + // After obtaining a read-lock no concurrent writes could occur + // and the loaded version matches that of the borrowed reference. + let new_version = self.shared.state.load().version(); + let has_changed = self.version != new_version; + + // Mark the shared value as seen by updating the version + self.version = new_version; + + Ref { inner, has_changed } + } + + /// Checks if this channel contains a message that this receiver has not yet + /// seen. The new value is not marked as seen. + /// + /// Although this method is called `has_changed`, it does not check new + /// messages for equality, so this call will return true even if the new + /// message is equal to the old message. + /// + /// Returns an error if the channel has been closed. + /// # Examples + /// + /// ``` + /// use tokio::sync::watch; + /// + /// #[tokio::main] + /// async fn main() { + /// let (tx, mut rx) = watch::channel("hello"); + /// + /// tx.send("goodbye").unwrap(); + /// + /// assert!(rx.has_changed().unwrap()); + /// assert_eq!(*rx.borrow_and_update(), "goodbye"); + /// + /// // The value has been marked as seen + /// assert!(!rx.has_changed().unwrap()); + /// + /// drop(tx); + /// // The `tx` handle has been dropped + /// assert!(rx.has_changed().is_err()); + /// } + /// ``` + pub fn has_changed(&self) -> Result<bool, error::RecvError> { + // Load the version from the state + let state = self.shared.state.load(); + if state.is_closed() { + // The sender has dropped. + return Err(error::RecvError(())); + } + let new_version = state.version(); + + Ok(self.version != new_version) } - /// Wait for a change notification, then mark the newest value as seen. + /// Waits for a change notification, then marks the newest value as seen. /// /// If the newest value in the channel has not yet been marked seen when /// this method is called, the method marks that value seen and returns @@ -291,21 +601,112 @@ impl<T> Receiver<T> { /// } /// ``` pub async fn changed(&mut self) -> Result<(), error::RecvError> { + changed_impl(&self.shared, &mut self.version).await + } + + /// Waits for a value that satisifes the provided condition. + /// + /// This method will call the provided closure whenever something is sent on + /// the channel. Once the closure returns `true`, this method will return a + /// reference to the value that was passed to the closure. + /// + /// Before `wait_for` starts waiting for changes, it will call the closure + /// on the current value. If the closure returns `true` when given the + /// current value, then `wait_for` will immediately return a reference to + /// the current value. This is the case even if the current value is already + /// considered seen. + /// + /// The watch channel only keeps track of the most recent value, so if + /// several messages are sent faster than `wait_for` is able to call the + /// closure, then it may skip some updates. Whenever the closure is called, + /// it will be called with the most recent value. + /// + /// When this function returns, the value that was passed to the closure + /// when it returned `true` will be considered seen. + /// + /// If the channel is closed, then `wait_for` will return a `RecvError`. + /// Once this happens, no more messages can ever be sent on the channel. + /// When an error is returned, it is guaranteed that the closure has been + /// called on the last value, and that it returned `false` for that value. + /// (If the closure returned `true`, then the last value would have been + /// returned instead of the error.) + /// + /// Like the `borrow` method, the returned borrow holds a read lock on the + /// inner value. This means that long-lived borrows could cause the producer + /// half to block. It is recommended to keep the borrow as short-lived as + /// possible. See the documentation of `borrow` for more information on + /// this. + /// + /// [`Receiver::changed()`]: crate::sync::watch::Receiver::changed + /// + /// # Examples + /// + /// ``` + /// use tokio::sync::watch; + /// + /// #[tokio::main] + /// + /// async fn main() { + /// let (tx, _rx) = watch::channel("hello"); + /// + /// tx.send("goodbye").unwrap(); + /// + /// // here we subscribe to a second receiver + /// // now in case of using `changed` we would have + /// // to first check the current value and then wait + /// // for changes or else `changed` would hang. + /// let mut rx2 = tx.subscribe(); + /// + /// // in place of changed we have use `wait_for` + /// // which would automatically check the current value + /// // and wait for changes until the closure returns true. + /// assert!(rx2.wait_for(|val| *val == "goodbye").await.is_ok()); + /// assert_eq!(*rx2.borrow(), "goodbye"); + /// } + /// ``` + pub async fn wait_for( + &mut self, + mut f: impl FnMut(&T) -> bool, + ) -> Result<Ref<'_, T>, error::RecvError> { + let mut closed = false; loop { - // In order to avoid a race condition, we first request a notification, - // **then** check the current value's version. If a new version exists, - // the notification request is dropped. - let notified = self.shared.notify_rx.notified(); + { + let inner = self.shared.value.read().unwrap(); - if let Some(ret) = maybe_changed(&self.shared, &mut self.version) { - return ret; + let new_version = self.shared.state.load().version(); + let has_changed = self.version != new_version; + self.version = new_version; + + if (!closed || has_changed) && f(&inner) { + return Ok(Ref { inner, has_changed }); + } } - notified.await; - // loop around again in case the wake-up was spurious + if closed { + return Err(error::RecvError(())); + } + + // Wait for the value to change. + closed = changed_impl(&self.shared, &mut self.version).await.is_err(); } } + /// Returns `true` if receivers belong to the same channel. + /// + /// # Examples + /// + /// ``` + /// let (tx, rx) = tokio::sync::watch::channel(true); + /// let rx2 = rx.clone(); + /// assert!(rx.same_channel(&rx2)); + /// + /// let (tx3, rx3) = tokio::sync::watch::channel(true); + /// assert!(!rx3.same_channel(&rx2)); + /// ``` + pub fn same_channel(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.shared, &other.shared) + } + cfg_process_driver! { pub(crate) fn try_has_changed(&mut self) -> Option<Result<(), error::RecvError>> { maybe_changed(&self.shared, &mut self.version) @@ -315,11 +716,11 @@ impl<T> Receiver<T> { fn maybe_changed<T>( shared: &Shared<T>, - version: &mut usize, + version: &mut Version, ) -> Option<Result<(), error::RecvError>> { // Load the version from the state - let state = shared.version.load(SeqCst); - let new_version = state & !CLOSED; + let state = shared.state.load(); + let new_version = state.version(); if *version != new_version { // Observe the new version and return @@ -327,7 +728,7 @@ fn maybe_changed<T>( return Some(Ok(())); } - if CLOSED == state & CLOSED { + if state.is_closed() { // All receivers have dropped. return Some(Err(error::RecvError(()))); } @@ -335,6 +736,27 @@ fn maybe_changed<T>( None } +async fn changed_impl<T>( + shared: &Shared<T>, + version: &mut Version, +) -> Result<(), error::RecvError> { + crate::trace::async_trace_leaf().await; + + loop { + // In order to avoid a race condition, we first request a notification, + // **then** check the current value's version. If a new version exists, + // the notification request is dropped. + let notified = shared.notify_rx.notified(); + + if let Some(ret) = maybe_changed(shared, version) { + return ret; + } + + notified.await; + // loop around again in case the wake-up was spurious + } +} + impl<T> Clone for Receiver<T> { fn clone(&self) -> Self { let version = self.version; @@ -357,19 +779,158 @@ impl<T> Drop for Receiver<T> { impl<T> Sender<T> { /// Sends a new value via the channel, notifying all receivers. + /// + /// This method fails if the channel is closed, which is the case when + /// every receiver has been dropped. It is possible to reopen the channel + /// using the [`subscribe`] method. However, when `send` fails, the value + /// isn't made available for future receivers (but returned with the + /// [`SendError`]). + /// + /// To always make a new value available for future receivers, even if no + /// receiver currently exists, one of the other send methods + /// ([`send_if_modified`], [`send_modify`], or [`send_replace`]) can be + /// used instead. + /// + /// [`subscribe`]: Sender::subscribe + /// [`SendError`]: error::SendError + /// [`send_if_modified`]: Sender::send_if_modified + /// [`send_modify`]: Sender::send_modify + /// [`send_replace`]: Sender::send_replace pub fn send(&self, value: T) -> Result<(), error::SendError<T>> { // This is pretty much only useful as a hint anyway, so synchronization isn't critical. - if 0 == self.shared.ref_count_rx.load(Relaxed) { - return Err(error::SendError { inner: value }); + if 0 == self.receiver_count() { + return Err(error::SendError(value)); } + self.send_replace(value); + Ok(()) + } + + /// Modifies the watched value **unconditionally** in-place, + /// notifying all receivers. + /// + /// This can be useful for modifying the watched value, without + /// having to allocate a new instance. Additionally, this + /// method permits sending values even when there are no receivers. + /// + /// Prefer to use the more versatile function [`Self::send_if_modified()`] + /// if the value is only modified conditionally during the mutable borrow + /// to prevent unneeded change notifications for unmodified values. + /// + /// # Panics + /// + /// This function panics when the invocation of the `modify` closure panics. + /// No receivers are notified when panicking. All changes of the watched + /// value applied by the closure before panicking will be visible in + /// subsequent calls to `borrow`. + /// + /// # Examples + /// + /// ``` + /// use tokio::sync::watch; + /// + /// struct State { + /// counter: usize, + /// } + /// let (state_tx, state_rx) = watch::channel(State { counter: 0 }); + /// state_tx.send_modify(|state| state.counter += 1); + /// assert_eq!(state_rx.borrow().counter, 1); + /// ``` + pub fn send_modify<F>(&self, modify: F) + where + F: FnOnce(&mut T), + { + self.send_if_modified(|value| { + modify(value); + true + }); + } + + /// Modifies the watched value **conditionally** in-place, + /// notifying all receivers only if modified. + /// + /// This can be useful for modifying the watched value, without + /// having to allocate a new instance. Additionally, this + /// method permits sending values even when there are no receivers. + /// + /// The `modify` closure must return `true` if the value has actually + /// been modified during the mutable borrow. It should only return `false` + /// if the value is guaranteed to be unmodified despite the mutable + /// borrow. + /// + /// Receivers are only notified if the closure returned `true`. If the + /// closure has modified the value but returned `false` this results + /// in a *silent modification*, i.e. the modified value will be visible + /// in subsequent calls to `borrow`, but receivers will not receive + /// a change notification. + /// + /// Returns the result of the closure, i.e. `true` if the value has + /// been modified and `false` otherwise. + /// + /// # Panics + /// + /// This function panics when the invocation of the `modify` closure panics. + /// No receivers are notified when panicking. All changes of the watched + /// value applied by the closure before panicking will be visible in + /// subsequent calls to `borrow`. + /// + /// # Examples + /// + /// ``` + /// use tokio::sync::watch; + /// + /// struct State { + /// counter: usize, + /// } + /// let (state_tx, mut state_rx) = watch::channel(State { counter: 1 }); + /// let inc_counter_if_odd = |state: &mut State| { + /// if state.counter % 2 == 1 { + /// state.counter += 1; + /// return true; + /// } + /// false + /// }; + /// + /// assert_eq!(state_rx.borrow().counter, 1); + /// + /// assert!(!state_rx.has_changed().unwrap()); + /// assert!(state_tx.send_if_modified(inc_counter_if_odd)); + /// assert!(state_rx.has_changed().unwrap()); + /// assert_eq!(state_rx.borrow_and_update().counter, 2); + /// + /// assert!(!state_rx.has_changed().unwrap()); + /// assert!(!state_tx.send_if_modified(inc_counter_if_odd)); + /// assert!(!state_rx.has_changed().unwrap()); + /// assert_eq!(state_rx.borrow_and_update().counter, 2); + /// ``` + pub fn send_if_modified<F>(&self, modify: F) -> bool + where + F: FnOnce(&mut T) -> bool, + { { // Acquire the write lock and update the value. let mut lock = self.shared.value.write().unwrap(); - *lock = value; - // Update the version. 2 is used so that the CLOSED bit is not set. - self.shared.version.fetch_add(2, SeqCst); + // Update the value and catch possible panic inside func. + let result = panic::catch_unwind(panic::AssertUnwindSafe(|| modify(&mut lock))); + match result { + Ok(modified) => { + if !modified { + // Abort, i.e. don't notify receivers if unmodified + return false; + } + // Continue if modified + } + Err(panicked) => { + // Drop the lock to avoid poisoning it. + drop(lock); + // Forward the panic to the caller. + panic::resume_unwind(panicked); + // Unreachable + } + }; + + self.shared.state.increment_version(); // Release the write lock. // @@ -379,17 +940,42 @@ impl<T> Sender<T> { drop(lock); } - // Notify all watchers self.shared.notify_rx.notify_waiters(); - Ok(()) + true + } + + /// Sends a new value via the channel, notifying all receivers and returning + /// the previous value in the channel. + /// + /// This can be useful for reusing the buffers inside a watched value. + /// Additionally, this method permits sending values even when there are no + /// receivers. + /// + /// # Examples + /// + /// ``` + /// use tokio::sync::watch; + /// + /// let (tx, _rx) = watch::channel(1); + /// assert_eq!(tx.send_replace(2), 1); + /// assert_eq!(tx.send_replace(3), 2); + /// ``` + pub fn send_replace(&self, mut value: T) -> T { + // swap old watched value with the new one + self.send_modify(|old| mem::swap(old, &mut value)); + + value } /// Returns a reference to the most recently sent value /// - /// Outstanding borrows hold a read lock. This means that long lived borrows - /// could cause the send half to block. It is recommended to keep the borrow - /// as short lived as possible. + /// Outstanding borrows hold a read lock on the inner value. This means that + /// long-lived borrows could cause the producer half to block. It is recommended + /// to keep the borrow as short-lived as possible. Additionally, if you are + /// running in an environment that allows `!Send` futures, you must ensure that + /// the returned `Ref` type is never held alive across an `.await` point, + /// otherwise, it can lead to a deadlock. /// /// # Examples /// @@ -401,7 +987,11 @@ impl<T> Sender<T> { /// ``` pub fn borrow(&self) -> Ref<'_, T> { let inner = self.shared.value.read().unwrap(); - Ref { inner } + + // The sender/producer always sees the current version + let has_changed = false; + + Ref { inner, has_changed } } /// Checks if the channel has been closed. This happens when all receivers @@ -417,7 +1007,7 @@ impl<T> Sender<T> { /// assert!(tx.is_closed()); /// ``` pub fn is_closed(&self) -> bool { - self.shared.ref_count_rx.load(Relaxed) == 0 + self.receiver_count() == 0 } /// Completes when all receivers have dropped. @@ -450,26 +1040,86 @@ impl<T> Sender<T> { /// } /// ``` pub async fn closed(&self) { - let notified = self.shared.notify_tx.notified(); + crate::trace::async_trace_leaf().await; - if self.shared.ref_count_rx.load(Relaxed) == 0 { - return; - } + while self.receiver_count() > 0 { + let notified = self.shared.notify_tx.notified(); - notified.await; - debug_assert_eq!(0, self.shared.ref_count_rx.load(Relaxed)); + if self.receiver_count() == 0 { + return; + } + + notified.await; + // The channel could have been reopened in the meantime by calling + // `subscribe`, so we loop again. + } } - cfg_signal_internal! { - pub(crate) fn subscribe(&self) -> Receiver<T> { - let shared = self.shared.clone(); - let version = shared.version.load(SeqCst); + /// Creates a new [`Receiver`] connected to this `Sender`. + /// + /// All messages sent before this call to `subscribe` are initially marked + /// as seen by the new `Receiver`. + /// + /// This method can be called even if there are no other receivers. In this + /// case, the channel is reopened. + /// + /// # Examples + /// + /// The new channel will receive messages sent on this `Sender`. + /// + /// ``` + /// use tokio::sync::watch; + /// + /// #[tokio::main] + /// async fn main() { + /// let (tx, _rx) = watch::channel(0u64); + /// + /// tx.send(5).unwrap(); + /// + /// let rx = tx.subscribe(); + /// assert_eq!(5, *rx.borrow()); + /// + /// tx.send(10).unwrap(); + /// assert_eq!(10, *rx.borrow()); + /// } + /// ``` + /// + /// The most recent message is considered seen by the channel, so this test + /// is guaranteed to pass. + /// + /// ``` + /// use tokio::sync::watch; + /// use tokio::time::Duration; + /// + /// #[tokio::main] + /// async fn main() { + /// let (tx, _rx) = watch::channel(0u64); + /// tx.send(5).unwrap(); + /// let mut rx = tx.subscribe(); + /// + /// tokio::spawn(async move { + /// // by spawning and sleeping, the message is sent after `main` + /// // hits the call to `changed`. + /// # if false { + /// tokio::time::sleep(Duration::from_millis(10)).await; + /// # } + /// tx.send(100).unwrap(); + /// }); + /// + /// rx.changed().await.unwrap(); + /// assert_eq!(100, *rx.borrow()); + /// } + /// ``` + pub fn subscribe(&self) -> Receiver<T> { + let shared = self.shared.clone(); + let version = shared.state.load().version(); - Receiver::from_shared(version, shared) - } + // The CLOSED bit in the state tracks only whether the sender is + // dropped, so we do not need to unset it if this reopens the channel. + Receiver::from_shared(version, shared) } - /// Returns the number of receivers that currently exist + /// Returns the number of receivers that currently exist. /// /// # Examples /// @@ -494,7 +1144,7 @@ impl<T> Sender<T> { impl<T> Drop for Sender<T> { fn drop(&mut self) { - self.shared.version.fetch_or(CLOSED, SeqCst); + self.shared.state.set_closed(); self.shared.notify_rx.notify_waiters(); } } diff --git a/vendor/tokio/src/task/blocking.rs b/vendor/tokio/src/task/blocking.rs index e4fe254a0..9bd15ebd5 100644 --- a/vendor/tokio/src/task/blocking.rs +++ b/vendor/tokio/src/task/blocking.rs @@ -70,11 +70,12 @@ cfg_rt_multi_thread! { /// This function panics if called from a [`current_thread`] runtime. /// /// [`current_thread`]: fn@crate::runtime::Builder::new_current_thread + #[track_caller] pub fn block_in_place<F, R>(f: F) -> R where F: FnOnce() -> R, { - crate::runtime::thread_pool::block_in_place(f) + crate::runtime::scheduler::multi_thread::block_in_place(f) } } @@ -89,50 +90,117 @@ cfg_rt! { /// /// Tokio will spawn more blocking threads when they are requested through this /// function until the upper limit configured on the [`Builder`] is reached. - /// This limit is very large by default, because `spawn_blocking` is often used - /// for various kinds of IO operations that cannot be performed asynchronously. - /// When you run CPU-bound code using `spawn_blocking`, you should keep this - /// large upper limit in mind. When running many CPU-bound computations, a - /// semaphore or some other synchronization primitive should be used to limit - /// the number of computation executed in parallel. Specialized CPU-bound - /// executors, such as [rayon], may also be a good fit. + /// After reaching the upper limit, the tasks are put in a queue. + /// The thread limit is very large by default, because `spawn_blocking` is often + /// used for various kinds of IO operations that cannot be performed + /// asynchronously. When you run CPU-bound code using `spawn_blocking`, you + /// should keep this large upper limit in mind. When running many CPU-bound + /// computations, a semaphore or some other synchronization primitive should be + /// used to limit the number of computation executed in parallel. Specialized + /// CPU-bound executors, such as [rayon], may also be a good fit. /// /// This function is intended for non-async operations that eventually finish on /// their own. If you want to spawn an ordinary thread, you should use /// [`thread::spawn`] instead. /// - /// Closures spawned using `spawn_blocking` cannot be cancelled. When you shut - /// down the executor, it will wait indefinitely for all blocking operations to + /// Closures spawned using `spawn_blocking` cannot be cancelled abruptly; there + /// is no standard low level API to cause a thread to stop running. However, + /// a useful pattern is to pass some form of "cancellation token" into + /// the thread. This could be an [`AtomicBool`] that the task checks periodically. + /// Another approach is to have the thread primarily read or write from a channel, + /// and to exit when the channel closes; assuming the other side of the channel is dropped + /// when cancellation occurs, this will cause the blocking task thread to exit + /// soon after as well. + /// + /// When you shut down the executor, it will wait indefinitely for all blocking operations to /// finish. You can use [`shutdown_timeout`] to stop waiting for them after a /// certain timeout. Be aware that this will still not cancel the tasks — they - /// are simply allowed to keep running after the method returns. + /// are simply allowed to keep running after the method returns. It is possible + /// for a blocking task to be cancelled if it has not yet started running, but this + /// is not guaranteed. /// /// Note that if you are using the single threaded runtime, this function will - /// still spawn additional threads for blocking operations. The basic + /// still spawn additional threads for blocking operations. The current-thread /// scheduler's single thread is only used for asynchronous code. /// + /// # Related APIs and patterns for bridging asynchronous and blocking code + /// + /// In simple cases, it is sufficient to have the closure accept input + /// parameters at creation time and return a single value (or struct/tuple, etc.). + /// + /// For more complex situations in which it is desirable to stream data to or from + /// the synchronous context, the [`mpsc channel`] has `blocking_send` and + /// `blocking_recv` methods for use in non-async code such as the thread created + /// by `spawn_blocking`. + /// + /// Another option is [`SyncIoBridge`] for cases where the synchronous context + /// is operating on byte streams. For example, you might use an asynchronous + /// HTTP client such as [hyper] to fetch data, but perform complex parsing + /// of the payload body using a library written for synchronous I/O. + /// + /// Finally, see also [Bridging with sync code][bridgesync] for discussions + /// around the opposite case of using Tokio as part of a larger synchronous + /// codebase. + /// /// [`Builder`]: struct@crate::runtime::Builder /// [blocking]: ../index.html#cpu-bound-tasks-and-blocking-code /// [rayon]: https://docs.rs/rayon + /// [`mpsc channel`]: crate::sync::mpsc + /// [`SyncIoBridge`]: https://docs.rs/tokio-util/latest/tokio_util/io/struct.SyncIoBridge.html + /// [hyper]: https://docs.rs/hyper /// [`thread::spawn`]: fn@std::thread::spawn /// [`shutdown_timeout`]: fn@crate::runtime::Runtime::shutdown_timeout + /// [bridgesync]: https://tokio.rs/tokio/topics/bridging + /// [`AtomicBool`]: struct@std::sync::atomic::AtomicBool /// /// # Examples /// + /// Pass an input value and receive result of computation: + /// /// ``` /// use tokio::task; /// /// # async fn docs() -> Result<(), Box<dyn std::error::Error>>{ + /// // Initial input + /// let mut v = "Hello, ".to_string(); /// let res = task::spawn_blocking(move || { - /// // do some compute-heavy work or call synchronous code - /// "done computing" + /// // Stand-in for compute-heavy work or using synchronous APIs + /// v.push_str("world"); + /// // Pass ownership of the value back to the asynchronous context + /// v /// }).await?; /// - /// assert_eq!(res, "done computing"); + /// // `res` is the value returned from the thread + /// assert_eq!(res.as_str(), "Hello, world"); /// # Ok(()) /// # } /// ``` - #[cfg_attr(tokio_track_caller, track_caller)] + /// + /// Use a channel: + /// + /// ``` + /// use tokio::task; + /// use tokio::sync::mpsc; + /// + /// # async fn docs() { + /// let (tx, mut rx) = mpsc::channel(2); + /// let start = 5; + /// let worker = task::spawn_blocking(move || { + /// for x in 0..10 { + /// // Stand in for complex computation + /// tx.blocking_send(start + x).unwrap(); + /// } + /// }); + /// + /// let mut acc = 0; + /// while let Some(v) = rx.recv().await { + /// acc += v; + /// } + /// assert_eq!(acc, 95); + /// worker.await.unwrap(); + /// # } + /// ``` + #[track_caller] pub fn spawn_blocking<F, R>(f: F) -> JoinHandle<R> where F: FnOnce() -> R + Send + 'static, diff --git a/vendor/tokio/src/task/builder.rs b/vendor/tokio/src/task/builder.rs index e46bdefe9..306cc39d1 100644 --- a/vendor/tokio/src/task/builder.rs +++ b/vendor/tokio/src/task/builder.rs @@ -1,10 +1,16 @@ #![allow(unreachable_pub)] -use crate::util::error::CONTEXT_MISSING_ERROR; -use crate::{runtime::context, task::JoinHandle}; -use std::future::Future; +use crate::{ + runtime::Handle, + task::{JoinHandle, LocalSet}, +}; +use std::{future::Future, io}; /// Factory which is used to configure the properties of a new task. /// +/// **Note**: This is an [unstable API][unstable]. The public API of this type +/// may break in 1.x releases. See [the documentation on unstable +/// features][unstable] for details. +/// /// Methods can be chained in order to configure it. /// /// Currently, there is only one configuration option: @@ -42,11 +48,17 @@ use std::future::Future; /// .spawn(async move { /// // Process each socket concurrently. /// process(socket).await -/// }); +/// })?; /// } /// } /// ``` +/// [unstable]: crate#unstable-features +/// [`name`]: Builder::name +/// [`spawn_local`]: Builder::spawn_local +/// [`spawn`]: Builder::spawn +/// [`spawn_blocking`]: Builder::spawn_blocking #[derive(Default, Debug)] +#[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "tracing"))))] pub struct Builder<'a> { name: Option<&'a str>, } @@ -62,44 +74,128 @@ impl<'a> Builder<'a> { Self { name: Some(name) } } - /// Spawns a task on the executor. + /// Spawns a task with this builder's settings on the current runtime. + /// + /// # Panics + /// + /// This method panics if called outside of a Tokio runtime. /// - /// See [`task::spawn`](crate::task::spawn) for + /// See [`task::spawn`](crate::task::spawn()) for /// more details. - #[cfg_attr(tokio_track_caller, track_caller)] - pub fn spawn<Fut>(self, future: Fut) -> JoinHandle<Fut::Output> + #[track_caller] + pub fn spawn<Fut>(self, future: Fut) -> io::Result<JoinHandle<Fut::Output>> where Fut: Future + Send + 'static, Fut::Output: Send + 'static, { - super::spawn::spawn_inner(future, self.name) + Ok(super::spawn::spawn_inner(future, self.name)) } - /// Spawns a task on the current thread. + /// Spawn a task with this builder's settings on the provided [runtime + /// handle]. /// - /// See [`task::spawn_local`](crate::task::spawn_local) - /// for more details. - #[cfg_attr(tokio_track_caller, track_caller)] - pub fn spawn_local<Fut>(self, future: Fut) -> JoinHandle<Fut::Output> + /// See [`Handle::spawn`] for more details. + /// + /// [runtime handle]: crate::runtime::Handle + /// [`Handle::spawn`]: crate::runtime::Handle::spawn + #[track_caller] + pub fn spawn_on<Fut>(self, future: Fut, handle: &Handle) -> io::Result<JoinHandle<Fut::Output>> + where + Fut: Future + Send + 'static, + Fut::Output: Send + 'static, + { + Ok(handle.spawn_named(future, self.name)) + } + + /// Spawns `!Send` a task on the current [`LocalSet`] with this builder's + /// settings. + /// + /// The spawned future will be run on the same thread that called `spawn_local`. + /// This may only be called from the context of a [local task set][`LocalSet`]. + /// + /// # Panics + /// + /// This function panics if called outside of a [local task set][`LocalSet`]. + /// + /// See [`task::spawn_local`] for more details. + /// + /// [`task::spawn_local`]: crate::task::spawn_local + /// [`LocalSet`]: crate::task::LocalSet + #[track_caller] + pub fn spawn_local<Fut>(self, future: Fut) -> io::Result<JoinHandle<Fut::Output>> where Fut: Future + 'static, Fut::Output: 'static, { - super::local::spawn_local_inner(future, self.name) + Ok(super::local::spawn_local_inner(future, self.name)) + } + + /// Spawns `!Send` a task on the provided [`LocalSet`] with this builder's + /// settings. + /// + /// See [`LocalSet::spawn_local`] for more details. + /// + /// [`LocalSet::spawn_local`]: crate::task::LocalSet::spawn_local + /// [`LocalSet`]: crate::task::LocalSet + #[track_caller] + pub fn spawn_local_on<Fut>( + self, + future: Fut, + local_set: &LocalSet, + ) -> io::Result<JoinHandle<Fut::Output>> + where + Fut: Future + 'static, + Fut::Output: 'static, + { + Ok(local_set.spawn_named(future, self.name)) } /// Spawns blocking code on the blocking threadpool. /// + /// # Panics + /// + /// This method panics if called outside of a Tokio runtime. + /// /// See [`task::spawn_blocking`](crate::task::spawn_blocking) /// for more details. - #[cfg_attr(tokio_track_caller, track_caller)] - pub fn spawn_blocking<Function, Output>(self, function: Function) -> JoinHandle<Output> + #[track_caller] + pub fn spawn_blocking<Function, Output>( + self, + function: Function, + ) -> io::Result<JoinHandle<Output>> where Function: FnOnce() -> Output + Send + 'static, Output: Send + 'static, { - context::current() - .expect(CONTEXT_MISSING_ERROR) - .spawn_blocking_inner(function, self.name) + let handle = Handle::current(); + self.spawn_blocking_on(function, &handle) + } + + /// Spawns blocking code on the provided [runtime handle]'s blocking threadpool. + /// + /// See [`Handle::spawn_blocking`] for more details. + /// + /// [runtime handle]: crate::runtime::Handle + /// [`Handle::spawn_blocking`]: crate::runtime::Handle::spawn_blocking + #[track_caller] + pub fn spawn_blocking_on<Function, Output>( + self, + function: Function, + handle: &Handle, + ) -> io::Result<JoinHandle<Output>> + where + Function: FnOnce() -> Output + Send + 'static, + Output: Send + 'static, + { + use crate::runtime::Mandatory; + let (join_handle, spawn_result) = handle.inner.blocking_spawner().spawn_blocking_inner( + function, + Mandatory::NonMandatory, + self.name, + handle, + ); + + spawn_result?; + Ok(join_handle) } } diff --git a/vendor/tokio/src/task/consume_budget.rs b/vendor/tokio/src/task/consume_budget.rs new file mode 100644 index 000000000..e7432ffe7 --- /dev/null +++ b/vendor/tokio/src/task/consume_budget.rs @@ -0,0 +1,46 @@ +use std::task::Poll; + +/// Consumes a unit of budget and returns the execution back to the Tokio +/// runtime *if* the task's coop budget was exhausted. +/// +/// The task will only yield if its entire coop budget has been exhausted. +/// This function can be used in order to insert optional yield points into long +/// computations that do not use Tokio resources like sockets or semaphores, +/// without redundantly yielding to the runtime each time. +/// +/// **Note**: This is an [unstable API][unstable]. The public API of this type +/// may break in 1.x releases. See [the documentation on unstable +/// features][unstable] for details. +/// +/// # Examples +/// +/// Make sure that a function which returns a sum of (potentially lots of) +/// iterated values is cooperative. +/// +/// ``` +/// async fn sum_iterator(input: &mut impl std::iter::Iterator<Item=i64>) -> i64 { +/// let mut sum: i64 = 0; +/// while let Some(i) = input.next() { +/// sum += i; +/// tokio::task::consume_budget().await +/// } +/// sum +/// } +/// ``` +/// [unstable]: crate#unstable-features +#[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "rt"))))] +pub async fn consume_budget() { + let mut status = Poll::Pending; + + crate::future::poll_fn(move |cx| { + ready!(crate::trace::trace_leaf(cx)); + if status.is_ready() { + return status; + } + status = crate::runtime::coop::poll_proceed(cx).map(|restore| { + restore.made_progress(); + }); + status + }) + .await +} diff --git a/vendor/tokio/src/task/join_set.rs b/vendor/tokio/src/task/join_set.rs new file mode 100644 index 000000000..4eb15a24d --- /dev/null +++ b/vendor/tokio/src/task/join_set.rs @@ -0,0 +1,584 @@ +//! A collection of tasks spawned on a Tokio runtime. +//! +//! This module provides the [`JoinSet`] type, a collection which stores a set +//! of spawned tasks and allows asynchronously awaiting the output of those +//! tasks as they complete. See the documentation for the [`JoinSet`] type for +//! details. +use std::fmt; +use std::future::Future; +use std::pin::Pin; +use std::task::{Context, Poll}; + +use crate::runtime::Handle; +#[cfg(tokio_unstable)] +use crate::task::Id; +use crate::task::{AbortHandle, JoinError, JoinHandle, LocalSet}; +use crate::util::IdleNotifiedSet; + +/// A collection of tasks spawned on a Tokio runtime. +/// +/// A `JoinSet` can be used to await the completion of some or all of the tasks +/// in the set. The set is not ordered, and the tasks will be returned in the +/// order they complete. +/// +/// All of the tasks must have the same return type `T`. +/// +/// When the `JoinSet` is dropped, all tasks in the `JoinSet` are immediately aborted. +/// +/// # Examples +/// +/// Spawn multiple tasks and wait for them. +/// +/// ``` +/// use tokio::task::JoinSet; +/// +/// #[tokio::main] +/// async fn main() { +/// let mut set = JoinSet::new(); +/// +/// for i in 0..10 { +/// set.spawn(async move { i }); +/// } +/// +/// let mut seen = [false; 10]; +/// while let Some(res) = set.join_next().await { +/// let idx = res.unwrap(); +/// seen[idx] = true; +/// } +/// +/// for i in 0..10 { +/// assert!(seen[i]); +/// } +/// } +/// ``` +#[cfg_attr(docsrs, doc(cfg(feature = "rt")))] +pub struct JoinSet<T> { + inner: IdleNotifiedSet<JoinHandle<T>>, +} + +/// A variant of [`task::Builder`] that spawns tasks on a [`JoinSet`] rather +/// than on the current default runtime. +/// +/// [`task::Builder`]: crate::task::Builder +#[cfg(all(tokio_unstable, feature = "tracing"))] +#[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "tracing"))))] +#[must_use = "builders do nothing unless used to spawn a task"] +pub struct Builder<'a, T> { + joinset: &'a mut JoinSet<T>, + builder: super::Builder<'a>, +} + +impl<T> JoinSet<T> { + /// Create a new `JoinSet`. + pub fn new() -> Self { + Self { + inner: IdleNotifiedSet::new(), + } + } + + /// Returns the number of tasks currently in the `JoinSet`. + pub fn len(&self) -> usize { + self.inner.len() + } + + /// Returns whether the `JoinSet` is empty. + pub fn is_empty(&self) -> bool { + self.inner.is_empty() + } +} + +impl<T: 'static> JoinSet<T> { + /// Returns a [`Builder`] that can be used to configure a task prior to + /// spawning it on this `JoinSet`. + /// + /// # Examples + /// + /// ``` + /// use tokio::task::JoinSet; + /// + /// #[tokio::main] + /// async fn main() -> std::io::Result<()> { + /// let mut set = JoinSet::new(); + /// + /// // Use the builder to configure a task's name before spawning it. + /// set.build_task() + /// .name("my_task") + /// .spawn(async { /* ... */ })?; + /// + /// Ok(()) + /// } + /// ``` + #[cfg(all(tokio_unstable, feature = "tracing"))] + #[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "tracing"))))] + pub fn build_task(&mut self) -> Builder<'_, T> { + Builder { + builder: super::Builder::new(), + joinset: self, + } + } + + /// Spawn the provided task on the `JoinSet`, returning an [`AbortHandle`] + /// that can be used to remotely cancel the task. + /// + /// The provided future will start running in the background immediately + /// when this method is called, even if you don't await anything on this + /// `JoinSet`. + /// + /// # Panics + /// + /// This method panics if called outside of a Tokio runtime. + /// + /// [`AbortHandle`]: crate::task::AbortHandle + #[track_caller] + pub fn spawn<F>(&mut self, task: F) -> AbortHandle + where + F: Future<Output = T>, + F: Send + 'static, + T: Send, + { + self.insert(crate::spawn(task)) + } + + /// Spawn the provided task on the provided runtime and store it in this + /// `JoinSet` returning an [`AbortHandle`] that can be used to remotely + /// cancel the task. + /// + /// The provided future will start running in the background immediately + /// when this method is called, even if you don't await anything on this + /// `JoinSet`. + /// + /// [`AbortHandle`]: crate::task::AbortHandle + #[track_caller] + pub fn spawn_on<F>(&mut self, task: F, handle: &Handle) -> AbortHandle + where + F: Future<Output = T>, + F: Send + 'static, + T: Send, + { + self.insert(handle.spawn(task)) + } + + /// Spawn the provided task on the current [`LocalSet`] and store it in this + /// `JoinSet`, returning an [`AbortHandle`] that can be used to remotely + /// cancel the task. + /// + /// The provided future will start running in the background immediately + /// when this method is called, even if you don't await anything on this + /// `JoinSet`. + /// + /// # Panics + /// + /// This method panics if it is called outside of a `LocalSet`. + /// + /// [`LocalSet`]: crate::task::LocalSet + /// [`AbortHandle`]: crate::task::AbortHandle + #[track_caller] + pub fn spawn_local<F>(&mut self, task: F) -> AbortHandle + where + F: Future<Output = T>, + F: 'static, + { + self.insert(crate::task::spawn_local(task)) + } + + /// Spawn the provided task on the provided [`LocalSet`] and store it in + /// this `JoinSet`, returning an [`AbortHandle`] that can be used to + /// remotely cancel the task. + /// + /// Unlike the [`spawn_local`] method, this method may be used to spawn local + /// tasks on a `LocalSet` that is _not_ currently running. The provided + /// future will start running whenever the `LocalSet` is next started. + /// + /// [`LocalSet`]: crate::task::LocalSet + /// [`AbortHandle`]: crate::task::AbortHandle + /// [`spawn_local`]: Self::spawn_local + #[track_caller] + pub fn spawn_local_on<F>(&mut self, task: F, local_set: &LocalSet) -> AbortHandle + where + F: Future<Output = T>, + F: 'static, + { + self.insert(local_set.spawn_local(task)) + } + + /// Spawn the blocking code on the blocking threadpool and store + /// it in this `JoinSet`, returning an [`AbortHandle`] that can be + /// used to remotely cancel the task. + /// + /// # Examples + /// + /// Spawn multiple blocking tasks and wait for them. + /// + /// ``` + /// use tokio::task::JoinSet; + /// + /// #[tokio::main] + /// async fn main() { + /// let mut set = JoinSet::new(); + /// + /// for i in 0..10 { + /// set.spawn_blocking(move || { i }); + /// } + /// + /// let mut seen = [false; 10]; + /// while let Some(res) = set.join_next().await { + /// let idx = res.unwrap(); + /// seen[idx] = true; + /// } + /// + /// for i in 0..10 { + /// assert!(seen[i]); + /// } + /// } + /// ``` + /// + /// # Panics + /// + /// This method panics if called outside of a Tokio runtime. + /// + /// [`AbortHandle`]: crate::task::AbortHandle + #[track_caller] + pub fn spawn_blocking<F>(&mut self, f: F) -> AbortHandle + where + F: FnOnce() -> T, + F: Send + 'static, + T: Send, + { + self.insert(crate::runtime::spawn_blocking(f)) + } + + /// Spawn the blocking code on the blocking threadpool of the + /// provided runtime and store it in this `JoinSet`, returning an + /// [`AbortHandle`] that can be used to remotely cancel the task. + /// + /// [`AbortHandle`]: crate::task::AbortHandle + #[track_caller] + pub fn spawn_blocking_on<F>(&mut self, f: F, handle: &Handle) -> AbortHandle + where + F: FnOnce() -> T, + F: Send + 'static, + T: Send, + { + self.insert(handle.spawn_blocking(f)) + } + + fn insert(&mut self, jh: JoinHandle<T>) -> AbortHandle { + let abort = jh.abort_handle(); + let mut entry = self.inner.insert_idle(jh); + + // Set the waker that is notified when the task completes. + entry.with_value_and_context(|jh, ctx| jh.set_join_waker(ctx.waker())); + abort + } + + /// Waits until one of the tasks in the set completes and returns its output. + /// + /// Returns `None` if the set is empty. + /// + /// # Cancel Safety + /// + /// This method is cancel safe. If `join_next` is used as the event in a `tokio::select!` + /// statement and some other branch completes first, it is guaranteed that no tasks were + /// removed from this `JoinSet`. + pub async fn join_next(&mut self) -> Option<Result<T, JoinError>> { + crate::future::poll_fn(|cx| self.poll_join_next(cx)).await + } + + /// Waits until one of the tasks in the set completes and returns its + /// output, along with the [task ID] of the completed task. + /// + /// Returns `None` if the set is empty. + /// + /// When this method returns an error, then the id of the task that failed can be accessed + /// using the [`JoinError::id`] method. + /// + /// # Cancel Safety + /// + /// This method is cancel safe. If `join_next_with_id` is used as the event in a `tokio::select!` + /// statement and some other branch completes first, it is guaranteed that no tasks were + /// removed from this `JoinSet`. + /// + /// [task ID]: crate::task::Id + /// [`JoinError::id`]: fn@crate::task::JoinError::id + #[cfg(tokio_unstable)] + #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))] + pub async fn join_next_with_id(&mut self) -> Option<Result<(Id, T), JoinError>> { + crate::future::poll_fn(|cx| self.poll_join_next_with_id(cx)).await + } + + /// Aborts all tasks and waits for them to finish shutting down. + /// + /// Calling this method is equivalent to calling [`abort_all`] and then calling [`join_next`] in + /// a loop until it returns `None`. + /// + /// This method ignores any panics in the tasks shutting down. When this call returns, the + /// `JoinSet` will be empty. + /// + /// [`abort_all`]: fn@Self::abort_all + /// [`join_next`]: fn@Self::join_next + pub async fn shutdown(&mut self) { + self.abort_all(); + while self.join_next().await.is_some() {} + } + + /// Aborts all tasks on this `JoinSet`. + /// + /// This does not remove the tasks from the `JoinSet`. To wait for the tasks to complete + /// cancellation, you should call `join_next` in a loop until the `JoinSet` is empty. + pub fn abort_all(&mut self) { + self.inner.for_each(|jh| jh.abort()); + } + + /// Removes all tasks from this `JoinSet` without aborting them. + /// + /// The tasks removed by this call will continue to run in the background even if the `JoinSet` + /// is dropped. + pub fn detach_all(&mut self) { + self.inner.drain(drop); + } + + /// Polls for one of the tasks in the set to complete. + /// + /// If this returns `Poll::Ready(Some(_))`, then the task that completed is removed from the set. + /// + /// When the method returns `Poll::Pending`, the `Waker` in the provided `Context` is scheduled + /// to receive a wakeup when a task in the `JoinSet` completes. Note that on multiple calls to + /// `poll_join_next`, only the `Waker` from the `Context` passed to the most recent call is + /// scheduled to receive a wakeup. + /// + /// # Returns + /// + /// This function returns: + /// + /// * `Poll::Pending` if the `JoinSet` is not empty but there is no task whose output is + /// available right now. + /// * `Poll::Ready(Some(Ok(value)))` if one of the tasks in this `JoinSet` has completed. + /// The `value` is the return value of one of the tasks that completed. + /// * `Poll::Ready(Some(Err(err)))` if one of the tasks in this `JoinSet` has panicked or been + /// aborted. The `err` is the `JoinError` from the panicked/aborted task. + /// * `Poll::Ready(None)` if the `JoinSet` is empty. + /// + /// Note that this method may return `Poll::Pending` even if one of the tasks has completed. + /// This can happen if the [coop budget] is reached. + /// + /// [coop budget]: crate::task#cooperative-scheduling + pub fn poll_join_next(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<T, JoinError>>> { + // The call to `pop_notified` moves the entry to the `idle` list. It is moved back to + // the `notified` list if the waker is notified in the `poll` call below. + let mut entry = match self.inner.pop_notified(cx.waker()) { + Some(entry) => entry, + None => { + if self.is_empty() { + return Poll::Ready(None); + } else { + // The waker was set by `pop_notified`. + return Poll::Pending; + } + } + }; + + let res = entry.with_value_and_context(|jh, ctx| Pin::new(jh).poll(ctx)); + + if let Poll::Ready(res) = res { + let _entry = entry.remove(); + Poll::Ready(Some(res)) + } else { + // A JoinHandle generally won't emit a wakeup without being ready unless + // the coop limit has been reached. We yield to the executor in this + // case. + cx.waker().wake_by_ref(); + Poll::Pending + } + } + + /// Polls for one of the tasks in the set to complete. + /// + /// If this returns `Poll::Ready(Some(_))`, then the task that completed is removed from the set. + /// + /// When the method returns `Poll::Pending`, the `Waker` in the provided `Context` is scheduled + /// to receive a wakeup when a task in the `JoinSet` completes. Note that on multiple calls to + /// `poll_join_next`, only the `Waker` from the `Context` passed to the most recent call is + /// scheduled to receive a wakeup. + /// + /// # Returns + /// + /// This function returns: + /// + /// * `Poll::Pending` if the `JoinSet` is not empty but there is no task whose output is + /// available right now. + /// * `Poll::Ready(Some(Ok((id, value))))` if one of the tasks in this `JoinSet` has completed. + /// The `value` is the return value of one of the tasks that completed, and + /// `id` is the [task ID] of that task. + /// * `Poll::Ready(Some(Err(err)))` if one of the tasks in this `JoinSet` has panicked or been + /// aborted. The `err` is the `JoinError` from the panicked/aborted task. + /// * `Poll::Ready(None)` if the `JoinSet` is empty. + /// + /// Note that this method may return `Poll::Pending` even if one of the tasks has completed. + /// This can happen if the [coop budget] is reached. + /// + /// [coop budget]: crate::task#cooperative-scheduling + /// [task ID]: crate::task::Id + #[cfg(tokio_unstable)] + #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))] + pub fn poll_join_next_with_id( + &mut self, + cx: &mut Context<'_>, + ) -> Poll<Option<Result<(Id, T), JoinError>>> { + // The call to `pop_notified` moves the entry to the `idle` list. It is moved back to + // the `notified` list if the waker is notified in the `poll` call below. + let mut entry = match self.inner.pop_notified(cx.waker()) { + Some(entry) => entry, + None => { + if self.is_empty() { + return Poll::Ready(None); + } else { + // The waker was set by `pop_notified`. + return Poll::Pending; + } + } + }; + + let res = entry.with_value_and_context(|jh, ctx| Pin::new(jh).poll(ctx)); + + if let Poll::Ready(res) = res { + let entry = entry.remove(); + // If the task succeeded, add the task ID to the output. Otherwise, the + // `JoinError` will already have the task's ID. + Poll::Ready(Some(res.map(|output| (entry.id(), output)))) + } else { + // A JoinHandle generally won't emit a wakeup without being ready unless + // the coop limit has been reached. We yield to the executor in this + // case. + cx.waker().wake_by_ref(); + Poll::Pending + } + } +} + +impl<T> Drop for JoinSet<T> { + fn drop(&mut self) { + self.inner.drain(|join_handle| join_handle.abort()); + } +} + +impl<T> fmt::Debug for JoinSet<T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("JoinSet").field("len", &self.len()).finish() + } +} + +impl<T> Default for JoinSet<T> { + fn default() -> Self { + Self::new() + } +} + +// === impl Builder === + +#[cfg(all(tokio_unstable, feature = "tracing"))] +#[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "tracing"))))] +impl<'a, T: 'static> Builder<'a, T> { + /// Assigns a name to the task which will be spawned. + pub fn name(self, name: &'a str) -> Self { + let builder = self.builder.name(name); + Self { builder, ..self } + } + + /// Spawn the provided task with this builder's settings and store it in the + /// [`JoinSet`], returning an [`AbortHandle`] that can be used to remotely + /// cancel the task. + /// + /// # Returns + /// + /// An [`AbortHandle`] that can be used to remotely cancel the task. + /// + /// # Panics + /// + /// This method panics if called outside of a Tokio runtime. + /// + /// [`AbortHandle`]: crate::task::AbortHandle + #[track_caller] + pub fn spawn<F>(self, future: F) -> std::io::Result<AbortHandle> + where + F: Future<Output = T>, + F: Send + 'static, + T: Send, + { + Ok(self.joinset.insert(self.builder.spawn(future)?)) + } + + /// Spawn the provided task on the provided [runtime handle] with this + /// builder's settings, and store it in the [`JoinSet`]. + /// + /// # Returns + /// + /// An [`AbortHandle`] that can be used to remotely cancel the task. + /// + /// + /// [`AbortHandle`]: crate::task::AbortHandle + /// [runtime handle]: crate::runtime::Handle + #[track_caller] + pub fn spawn_on<F>(self, future: F, handle: &Handle) -> std::io::Result<AbortHandle> + where + F: Future<Output = T>, + F: Send + 'static, + T: Send, + { + Ok(self.joinset.insert(self.builder.spawn_on(future, handle)?)) + } + + /// Spawn the provided task on the current [`LocalSet`] with this builder's + /// settings, and store it in the [`JoinSet`]. + /// + /// # Returns + /// + /// An [`AbortHandle`] that can be used to remotely cancel the task. + /// + /// # Panics + /// + /// This method panics if it is called outside of a `LocalSet`. + /// + /// [`LocalSet`]: crate::task::LocalSet + /// [`AbortHandle`]: crate::task::AbortHandle + #[track_caller] + pub fn spawn_local<F>(self, future: F) -> std::io::Result<AbortHandle> + where + F: Future<Output = T>, + F: 'static, + { + Ok(self.joinset.insert(self.builder.spawn_local(future)?)) + } + + /// Spawn the provided task on the provided [`LocalSet`] with this builder's + /// settings, and store it in the [`JoinSet`]. + /// + /// # Returns + /// + /// An [`AbortHandle`] that can be used to remotely cancel the task. + /// + /// [`LocalSet`]: crate::task::LocalSet + /// [`AbortHandle`]: crate::task::AbortHandle + #[track_caller] + pub fn spawn_local_on<F>(self, future: F, local_set: &LocalSet) -> std::io::Result<AbortHandle> + where + F: Future<Output = T>, + F: 'static, + { + Ok(self + .joinset + .insert(self.builder.spawn_local_on(future, local_set)?)) + } +} + +// Manual `Debug` impl so that `Builder` is `Debug` regardless of whether `T` is +// `Debug`. +#[cfg(all(tokio_unstable, feature = "tracing"))] +#[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "tracing"))))] +impl<'a, T> fmt::Debug for Builder<'a, T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("join_set::Builder") + .field("joinset", &self.joinset) + .field("builder", &self.builder) + .finish() + } +} diff --git a/vendor/tokio/src/task/local.rs b/vendor/tokio/src/task/local.rs index 37c2c508a..734b95587 100644 --- a/vendor/tokio/src/task/local.rs +++ b/vendor/tokio/src/task/local.rs @@ -1,15 +1,18 @@ //! Runs `!Send` futures on the current thread. +use crate::loom::cell::UnsafeCell; use crate::loom::sync::{Arc, Mutex}; -use crate::runtime::task::{self, JoinHandle, Task}; +use crate::runtime::task::{self, JoinHandle, LocalOwnedTasks, Task}; +use crate::runtime::{context, ThreadId}; use crate::sync::AtomicWaker; -use crate::util::linked_list::{Link, LinkedList}; +use crate::util::RcCell; -use std::cell::{Cell, RefCell}; +use std::cell::Cell; use std::collections::VecDeque; use std::fmt; use std::future::Future; use std::marker::PhantomData; use std::pin::Pin; +use std::rc::Rc; use std::task::Poll; use pin_project_lite::pin_project; @@ -31,14 +34,14 @@ cfg_rt! { /// async fn main() { /// // `Rc` does not implement `Send`, and thus may not be sent between /// // threads safely. - /// let unsend_data = Rc::new("my unsend data..."); + /// let nonsend_data = Rc::new("my nonsend data..."); /// - /// let unsend_data = unsend_data.clone(); - /// // Because the `async` block here moves `unsend_data`, the future is `!Send`. + /// let nonsend_data = nonsend_data.clone(); + /// // Because the `async` block here moves `nonsend_data`, the future is `!Send`. /// // Since `tokio::spawn` requires the spawned future to implement `Send`, this /// // will not compile. /// tokio::spawn(async move { - /// println!("{}", unsend_data); + /// println!("{}", nonsend_data); /// // ... /// }).await.unwrap(); /// } @@ -57,18 +60,18 @@ cfg_rt! { /// /// #[tokio::main] /// async fn main() { - /// let unsend_data = Rc::new("my unsend data..."); + /// let nonsend_data = Rc::new("my nonsend data..."); /// /// // Construct a local task set that can run `!Send` futures. /// let local = task::LocalSet::new(); /// /// // Run the local task set. /// local.run_until(async move { - /// let unsend_data = unsend_data.clone(); + /// let nonsend_data = nonsend_data.clone(); /// // `spawn_local` ensures that the future is spawned on the local /// // task set. /// task::spawn_local(async move { - /// println!("{}", unsend_data); + /// println!("{}", nonsend_data); /// // ... /// }).await.unwrap(); /// }).await; @@ -91,18 +94,18 @@ cfg_rt! { /// /// #[tokio::main] /// async fn main() { - /// let unsend_data = Rc::new("world"); + /// let nonsend_data = Rc::new("world"); /// let local = task::LocalSet::new(); /// - /// let unsend_data2 = unsend_data.clone(); + /// let nonsend_data2 = nonsend_data.clone(); /// local.spawn_local(async move { /// // ... - /// println!("hello {}", unsend_data2) + /// println!("hello {}", nonsend_data2) /// }); /// /// local.spawn_local(async move { /// time::sleep(time::Duration::from_millis(100)).await; - /// println!("goodbye {}", unsend_data) + /// println!("goodbye {}", nonsend_data) /// }); /// /// // ... @@ -159,7 +162,7 @@ cfg_rt! { /// tokio::task::spawn_local(run_task(new_task)); /// } /// // If the while loop returns, then all the LocalSpawner - /// // objects have have been dropped. + /// // objects have been dropped. /// }); /// /// // This will return once all senders are dropped and all @@ -211,41 +214,57 @@ cfg_rt! { /// [`task::spawn_local`]: fn@spawn_local /// [`mpsc`]: mod@crate::sync::mpsc pub struct LocalSet { - /// Current scheduler tick + /// Current scheduler tick. tick: Cell<u8>, - /// State available from thread-local - context: Context, + /// State available from thread-local. + context: Rc<Context>, /// This type should not be Send. _not_send: PhantomData<*const ()>, } } -/// State available from the thread-local +/// State available from the thread-local. struct Context { - /// Owned task set and local run queue - tasks: RefCell<Tasks>, - /// State shared between threads. shared: Arc<Shared>, -} -struct Tasks { - /// Collection of all active tasks spawned onto this executor. - owned: LinkedList<Task<Arc<Shared>>, <Task<Arc<Shared>> as Link>::Target>, - - /// Local run queue sender and receiver. - queue: VecDeque<task::Notified<Arc<Shared>>>, + /// True if a task panicked without being handled and the local set is + /// configured to shutdown on unhandled panic. + unhandled_panic: Cell<bool>, } /// LocalSet state shared between threads. struct Shared { - /// Remote run queue sender + /// # Safety + /// + /// This field must *only* be accessed from the thread that owns the + /// `LocalSet` (i.e., `Thread::current().id() == owner`). + local_state: LocalState, + + /// Remote run queue sender. queue: Mutex<Option<VecDeque<task::Notified<Arc<Shared>>>>>, - /// Wake the `LocalSet` task + /// Wake the `LocalSet` task. waker: AtomicWaker, + + /// How to respond to unhandled task panics. + #[cfg(tokio_unstable)] + pub(crate) unhandled_panic: crate::runtime::UnhandledPanic, +} + +/// Tracks the `LocalSet` state that must only be accessed from the thread that +/// created the `LocalSet`. +struct LocalState { + /// The `ThreadId` of the thread that owns the `LocalSet`. + owner: ThreadId, + + /// Local run queue sender and receiver. + local_queue: UnsafeCell<VecDeque<task::Notified<Arc<Shared>>>>, + + /// Collection of all active tasks spawned onto this executor. + owned: LocalOwnedTasks<Arc<Shared>>, } pin_project! { @@ -257,17 +276,30 @@ pin_project! { } } -scoped_thread_local!(static CURRENT: Context); +tokio_thread_local!(static CURRENT: LocalData = const { LocalData { + ctx: RcCell::new(), +} }); + +struct LocalData { + ctx: RcCell<Context>, +} cfg_rt! { - /// Spawns a `!Send` future on the local task set. + /// Spawns a `!Send` future on the current [`LocalSet`]. + /// + /// The spawned future will run on the same thread that called `spawn_local`. /// - /// The spawned future will be run on the same thread that called `spawn_local.` - /// This may only be called from the context of a local task set. + /// The provided future will start running in the background immediately + /// when `spawn_local` is called, even if you don't await the returned + /// `JoinHandle`. /// /// # Panics /// - /// - This function panics if called outside of a local task set. + /// This function panics if called outside of a [`LocalSet`]. + /// + /// Note that if [`tokio::spawn`] is used from within a `LocalSet`, the + /// resulting new task will _not_ be inside the `LocalSet`, so you must use + /// `spawn_local` if you want to stay within the `LocalSet`. /// /// # Examples /// @@ -277,21 +309,24 @@ cfg_rt! { /// /// #[tokio::main] /// async fn main() { - /// let unsend_data = Rc::new("my unsend data..."); + /// let nonsend_data = Rc::new("my nonsend data..."); /// /// let local = task::LocalSet::new(); /// /// // Run the local task set. /// local.run_until(async move { - /// let unsend_data = unsend_data.clone(); + /// let nonsend_data = nonsend_data.clone(); /// task::spawn_local(async move { - /// println!("{}", unsend_data); + /// println!("{}", nonsend_data); /// // ... /// }).await.unwrap(); /// }).await; /// } /// ``` - #[cfg_attr(tokio_track_caller, track_caller)] + /// + /// [`LocalSet`]: struct@crate::task::LocalSet + /// [`tokio::spawn`]: fn@crate::task::spawn + #[track_caller] pub fn spawn_local<F>(future: F) -> JoinHandle<F::Output> where F: Future + 'static, @@ -300,58 +335,94 @@ cfg_rt! { spawn_local_inner(future, None) } + + #[track_caller] pub(super) fn spawn_local_inner<F>(future: F, name: Option<&str>) -> JoinHandle<F::Output> where F: Future + 'static, F::Output: 'static { - let future = crate::util::trace::task(future, "local", name); - CURRENT.with(|maybe_cx| { - let cx = maybe_cx - .expect("`spawn_local` called from outside of a `task::LocalSet`"); - - // Safety: Tasks are only polled and dropped from the thread that - // spawns them. - let (task, handle) = unsafe { task::joinable_local(future) }; - cx.tasks.borrow_mut().queue.push_back(task); - handle - }) + match CURRENT.with(|LocalData { ctx, .. }| ctx.get()) { + None => panic!("`spawn_local` called from outside of a `task::LocalSet`"), + Some(cx) => cx.spawn(future, name) + } } } -/// Initial queue capacity +/// Initial queue capacity. const INITIAL_CAPACITY: usize = 64; /// Max number of tasks to poll per tick. const MAX_TASKS_PER_TICK: usize = 61; -/// How often it check the remote queue first +/// How often it check the remote queue first. const REMOTE_FIRST_INTERVAL: u8 = 31; +/// Context guard for LocalSet +pub struct LocalEnterGuard(Option<Rc<Context>>); + +impl Drop for LocalEnterGuard { + fn drop(&mut self) { + CURRENT.with(|LocalData { ctx, .. }| { + ctx.set(self.0.take()); + }) + } +} + +impl fmt::Debug for LocalEnterGuard { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("LocalEnterGuard").finish() + } +} + impl LocalSet { /// Returns a new local task set. pub fn new() -> LocalSet { + let owner = context::thread_id().expect("cannot create LocalSet during thread shutdown"); + LocalSet { tick: Cell::new(0), - context: Context { - tasks: RefCell::new(Tasks { - owned: LinkedList::new(), - queue: VecDeque::with_capacity(INITIAL_CAPACITY), - }), + context: Rc::new(Context { shared: Arc::new(Shared { + local_state: LocalState { + owner, + owned: LocalOwnedTasks::new(), + local_queue: UnsafeCell::new(VecDeque::with_capacity(INITIAL_CAPACITY)), + }, queue: Mutex::new(Some(VecDeque::with_capacity(INITIAL_CAPACITY))), waker: AtomicWaker::new(), + #[cfg(tokio_unstable)] + unhandled_panic: crate::runtime::UnhandledPanic::Ignore, }), - }, + unhandled_panic: Cell::new(false), + }), _not_send: PhantomData, } } + /// Enters the context of this `LocalSet`. + /// + /// The [`spawn_local`] method will spawn tasks on the `LocalSet` whose + /// context you are inside. + /// + /// [`spawn_local`]: fn@crate::task::spawn_local + pub fn enter(&self) -> LocalEnterGuard { + CURRENT.with(|LocalData { ctx, .. }| { + let old = ctx.replace(Some(self.context.clone())); + LocalEnterGuard(old) + }) + } + /// Spawns a `!Send` task onto the local task set. /// /// This task is guaranteed to be run on the current thread. /// /// Unlike the free function [`spawn_local`], this method may be used to - /// spawn local tasks when the task set is _not_ running. For example: + /// spawn local tasks when the `LocalSet` is _not_ running. The provided + /// future will start running once the `LocalSet` is next started, even if + /// you don't await the returned `JoinHandle`. + /// + /// # Examples + /// /// ```rust /// use tokio::task; /// @@ -382,17 +453,13 @@ impl LocalSet { /// } /// ``` /// [`spawn_local`]: fn@spawn_local - #[cfg_attr(tokio_track_caller, track_caller)] + #[track_caller] pub fn spawn_local<F>(&self, future: F) -> JoinHandle<F::Output> where F: Future + 'static, F::Output: 'static, { - let future = crate::util::trace::task(future, "local", None); - let (task, handle) = unsafe { task::joinable_local(future) }; - self.context.tasks.borrow_mut().queue.push_back(task); - self.context.shared.waker.wake(); - handle + self.spawn_named(future, None) } /// Runs a future to completion on the provided runtime, driving any local @@ -457,6 +524,7 @@ impl LocalSet { /// [`Runtime::block_on`]: method@crate::runtime::Runtime::block_on /// [in-place blocking]: fn@crate::task::block_in_place /// [`spawn_blocking`]: fn@crate::task::spawn_blocking + #[track_caller] #[cfg(feature = "rt")] #[cfg_attr(docsrs, doc(cfg(feature = "rt")))] pub fn block_on<F>(&self, rt: &crate::runtime::Runtime, future: F) -> F::Output @@ -466,7 +534,7 @@ impl LocalSet { rt.block_on(self.run_until(future)) } - /// Run a future to completion on the local set, returning its output. + /// Runs a future to completion on the local set, returning its output. /// /// This returns a future that runs the given future with a local set, /// allowing it to call [`spawn_local`] to spawn additional `!Send` futures. @@ -505,10 +573,36 @@ impl LocalSet { run_until.await } - /// Tick the scheduler, returning whether the local future needs to be + pub(in crate::task) fn spawn_named<F>( + &self, + future: F, + name: Option<&str>, + ) -> JoinHandle<F::Output> + where + F: Future + 'static, + F::Output: 'static, + { + let handle = self.context.spawn(future, name); + + // Because a task was spawned from *outside* the `LocalSet`, wake the + // `LocalSet` future to execute the new task, if it hasn't been woken. + // + // Spawning via the free fn `spawn` does not require this, as it can + // only be called from *within* a future executing on the `LocalSet` — + // in that case, the `LocalSet` must already be awake. + self.context.shared.waker.wake(); + handle + } + + /// Ticks the scheduler, returning whether the local future needs to be /// notified again. fn tick(&self) -> bool { for _ in 0..MAX_TASKS_PER_TICK { + // Make sure we didn't hit an unhandled panic + if self.context.unhandled_panic.get() { + panic!("a spawned task panicked and the LocalSet is configured to shutdown on unhandled panic"); + } + match self.next_task() { // Run the task // @@ -518,7 +612,7 @@ impl LocalSet { // task initially. Because `LocalSet` itself is `!Send`, and // `spawn_local` spawns into the `LocalSet` on the current // thread, the invariant is maintained. - Some(task) => crate::coop::budget(|| task.run()), + Some(task) => crate::runtime::coop::budget(|| task.run()), // We have fully drained the queue of notified tasks, so the // local future doesn't need to be notified again — it can wait // until something else wakes a task in the local set. @@ -529,37 +623,168 @@ impl LocalSet { true } - fn next_task(&self) -> Option<task::Notified<Arc<Shared>>> { + fn next_task(&self) -> Option<task::LocalNotified<Arc<Shared>>> { let tick = self.tick.get(); self.tick.set(tick.wrapping_add(1)); - if tick % REMOTE_FIRST_INTERVAL == 0 { + let task = if tick % REMOTE_FIRST_INTERVAL == 0 { self.context .shared .queue .lock() .as_mut() .and_then(|queue| queue.pop_front()) - .or_else(|| self.context.tasks.borrow_mut().queue.pop_front()) + .or_else(|| self.pop_local()) } else { - self.context - .tasks - .borrow_mut() - .queue - .pop_front() - .or_else(|| { - self.context - .shared - .queue - .lock() - .as_mut() - .and_then(|queue| queue.pop_front()) - }) + self.pop_local().or_else(|| { + self.context + .shared + .queue + .lock() + .as_mut() + .and_then(|queue| queue.pop_front()) + }) + }; + + task.map(|task| unsafe { + // Safety: because the `LocalSet` itself is `!Send`, we know we are + // on the same thread if we have access to the `LocalSet`, and can + // therefore access the local run queue. + self.context.shared.local_state.assert_owner(task) + }) + } + + fn pop_local(&self) -> Option<task::Notified<Arc<Shared>>> { + unsafe { + // Safety: because the `LocalSet` itself is `!Send`, we know we are + // on the same thread if we have access to the `LocalSet`, and can + // therefore access the local run queue. + self.context.shared.local_state.task_pop_front() } } fn with<T>(&self, f: impl FnOnce() -> T) -> T { - CURRENT.set(&self.context, f) + CURRENT.with(|LocalData { ctx, .. }| { + struct Reset<'a> { + ctx_ref: &'a RcCell<Context>, + val: Option<Rc<Context>>, + } + impl<'a> Drop for Reset<'a> { + fn drop(&mut self) { + self.ctx_ref.set(self.val.take()); + } + } + let old = ctx.replace(Some(self.context.clone())); + + let _reset = Reset { + ctx_ref: ctx, + val: old, + }; + + f() + }) + } + + /// This method is like `with`, but it just calls `f` without setting the thread-local if that + /// fails. + fn with_if_possible<T>(&self, f: impl FnOnce() -> T) -> T { + let mut f = Some(f); + + let res = CURRENT.try_with(|LocalData { ctx, .. }| { + struct Reset<'a> { + ctx_ref: &'a RcCell<Context>, + val: Option<Rc<Context>>, + } + impl<'a> Drop for Reset<'a> { + fn drop(&mut self) { + self.ctx_ref.replace(self.val.take()); + } + } + let old = ctx.replace(Some(self.context.clone())); + + let _reset = Reset { + ctx_ref: ctx, + val: old, + }; + + (f.take().unwrap())() + }); + + match res { + Ok(res) => res, + Err(_access_error) => (f.take().unwrap())(), + } + } +} + +cfg_unstable! { + impl LocalSet { + /// Configure how the `LocalSet` responds to an unhandled panic on a + /// spawned task. + /// + /// By default, an unhandled panic (i.e. a panic not caught by + /// [`std::panic::catch_unwind`]) has no impact on the `LocalSet`'s + /// execution. The panic is error value is forwarded to the task's + /// [`JoinHandle`] and all other spawned tasks continue running. + /// + /// The `unhandled_panic` option enables configuring this behavior. + /// + /// * `UnhandledPanic::Ignore` is the default behavior. Panics on + /// spawned tasks have no impact on the `LocalSet`'s execution. + /// * `UnhandledPanic::ShutdownRuntime` will force the `LocalSet` to + /// shutdown immediately when a spawned task panics even if that + /// task's `JoinHandle` has not been dropped. All other spawned tasks + /// will immediately terminate and further calls to + /// [`LocalSet::block_on`] and [`LocalSet::run_until`] will panic. + /// + /// # Panics + /// + /// This method panics if called after the `LocalSet` has started + /// running. + /// + /// # Unstable + /// + /// This option is currently unstable and its implementation is + /// incomplete. The API may change or be removed in the future. See + /// tokio-rs/tokio#4516 for more details. + /// + /// # Examples + /// + /// The following demonstrates a `LocalSet` configured to shutdown on + /// panic. The first spawned task panics and results in the `LocalSet` + /// shutting down. The second spawned task never has a chance to + /// execute. The call to `run_until` will panic due to the runtime being + /// forcibly shutdown. + /// + /// ```should_panic + /// use tokio::runtime::UnhandledPanic; + /// + /// # #[tokio::main] + /// # async fn main() { + /// tokio::task::LocalSet::new() + /// .unhandled_panic(UnhandledPanic::ShutdownRuntime) + /// .run_until(async { + /// tokio::task::spawn_local(async { panic!("boom"); }); + /// tokio::task::spawn_local(async { + /// // This task never completes + /// }); + /// + /// // Do some work, but `run_until` will panic before it completes + /// # loop { tokio::task::yield_now().await; } + /// }) + /// .await; + /// # } + /// ``` + /// + /// [`JoinHandle`]: struct@crate::task::JoinHandle + pub fn unhandled_panic(&mut self, behavior: crate::runtime::UnhandledPanic) -> &mut Self { + // TODO: This should be set as a builder + Rc::get_mut(&mut self.context) + .and_then(|ctx| Arc::get_mut(&mut ctx.shared)) + .expect("Unhandled Panic behavior modified after starting LocalSet") + .unhandled_panic = behavior; + self + } } } @@ -581,7 +806,10 @@ impl Future for LocalSet { // there are still tasks remaining in the run queue. cx.waker().wake_by_ref(); Poll::Pending - } else if self.context.tasks.borrow().owned.is_empty() { + + // Safety: called from the thread that owns `LocalSet`. Because + // `LocalSet` is `!Send`, this is safe. + } else if unsafe { self.context.shared.local_state.owned_is_empty() } { // If the scheduler has no remaining futures, we're done! Poll::Ready(()) } else { @@ -601,35 +829,79 @@ impl Default for LocalSet { impl Drop for LocalSet { fn drop(&mut self) { - self.with(|| { - // Loop required here to ensure borrow is dropped between iterations - #[allow(clippy::while_let_loop)] - loop { - let task = match self.context.tasks.borrow_mut().owned.pop_back() { - Some(task) => task, - None => break, - }; - - // Safety: same as `run_unchecked`. - task.shutdown(); + self.with_if_possible(|| { + // Shut down all tasks in the LocalOwnedTasks and close it to + // prevent new tasks from ever being added. + unsafe { + // Safety: called from the thread that owns `LocalSet` + self.context.shared.local_state.close_and_shutdown_all(); } - for task in self.context.tasks.borrow_mut().queue.drain(..) { - task.shutdown(); + // We already called shutdown on all tasks above, so there is no + // need to call shutdown. + + // Safety: note that this *intentionally* bypasses the unsafe + // `Shared::local_queue()` method. This is in order to avoid the + // debug assertion that we are on the thread that owns the + // `LocalSet`, because on some systems (e.g. at least some macOS + // versions), attempting to get the current thread ID can panic due + // to the thread's local data that stores the thread ID being + // dropped *before* the `LocalSet`. + // + // Despite avoiding the assertion here, it is safe for us to access + // the local queue in `Drop`, because the `LocalSet` itself is + // `!Send`, so we can reasonably guarantee that it will not be + // `Drop`ped from another thread. + let local_queue = unsafe { + // Safety: called from the thread that owns `LocalSet` + self.context.shared.local_state.take_local_queue() + }; + for task in local_queue { + drop(task); } // Take the queue from the Shared object to prevent pushing // notifications to it in the future. let queue = self.context.shared.queue.lock().take().unwrap(); for task in queue { - task.shutdown(); + drop(task); } - assert!(self.context.tasks.borrow().owned.is_empty()); + // Safety: called from the thread that owns `LocalSet` + assert!(unsafe { self.context.shared.local_state.owned_is_empty() }); }); } } +// === impl Context === + +impl Context { + #[track_caller] + fn spawn<F>(&self, future: F, name: Option<&str>) -> JoinHandle<F::Output> + where + F: Future + 'static, + F::Output: 'static, + { + let id = crate::runtime::task::Id::next(); + let future = crate::util::trace::task(future, "local", name, id.as_u64()); + + // Safety: called from the thread that owns the `LocalSet` + let (handle, notified) = { + self.shared.local_state.assert_called_from_owner_thread(); + self.shared + .local_state + .owned + .bind(future, self.shared.clone(), id) + }; + + if let Some(notified) = notified { + self.shared.schedule(notified); + } + + handle + } +} + // === impl LocalFuture === impl<T: Future> Future for RunUntil<'_, T> { @@ -645,10 +917,10 @@ impl<T: Future> Future for RunUntil<'_, T> { .waker .register_by_ref(cx.waker()); - let _no_blocking = crate::runtime::enter::disallow_blocking(); + let _no_blocking = crate::runtime::context::disallow_block_in_place(); let f = me.future; - if let Poll::Ready(output) = crate::coop::budget(|| f.poll(cx)) { + if let Poll::Ready(output) = f.poll(cx) { return Poll::Ready(output); } @@ -666,21 +938,41 @@ impl<T: Future> Future for RunUntil<'_, T> { impl Shared { /// Schedule the provided task on the scheduler. fn schedule(&self, task: task::Notified<Arc<Self>>) { - CURRENT.with(|maybe_cx| match maybe_cx { - Some(cx) if cx.shared.ptr_eq(self) => { - cx.tasks.borrow_mut().queue.push_back(task); - } - _ => { - // First check whether the queue is still there (if not, the - // LocalSet is dropped). Then push to it if so, and if not, - // do nothing. - let mut lock = self.queue.lock(); - - if let Some(queue) = lock.as_mut() { - queue.push_back(task); - drop(lock); + CURRENT.with(|localdata| { + match localdata.ctx.get() { + Some(cx) if cx.shared.ptr_eq(self) => unsafe { + // Safety: if the current `LocalSet` context points to this + // `LocalSet`, then we are on the thread that owns it. + cx.shared.local_state.task_push_back(task); + }, + + // We are on the thread that owns the `LocalSet`, so we can + // wake to the local queue. + _ if context::thread_id().ok() == Some(self.local_state.owner) => { + unsafe { + // Safety: we just checked that the thread ID matches + // the localset's owner, so this is safe. + self.local_state.task_push_back(task); + } + // We still have to wake the `LocalSet`, because it isn't + // currently being polled. self.waker.wake(); } + + // We are *not* on the thread that owns the `LocalSet`, so we + // have to wake to the remote queue. + _ => { + // First, check whether the queue is still there (if not, the + // LocalSet is dropped). Then push to it if so, and if not, + // do nothing. + let mut lock = self.queue.lock(); + + if let Some(queue) = lock.as_mut() { + queue.push_back(task); + drop(lock); + self.waker.wake(); + } + } } }); } @@ -690,31 +982,194 @@ impl Shared { } } +// This is safe because (and only because) we *pinky pwomise* to never touch the +// local run queue except from the thread that owns the `LocalSet`. +unsafe impl Sync for Shared {} + impl task::Schedule for Arc<Shared> { - fn bind(task: Task<Self>) -> Arc<Shared> { - CURRENT.with(|maybe_cx| { - let cx = maybe_cx.expect("scheduler context missing"); - cx.tasks.borrow_mut().owned.push_front(task); - cx.shared.clone() - }) + fn release(&self, task: &Task<Self>) -> Option<Task<Self>> { + // Safety, this is always called from the thread that owns `LocalSet` + unsafe { self.local_state.task_remove(task) } } - fn release(&self, task: &Task<Self>) -> Option<Task<Self>> { - use std::ptr::NonNull; + fn schedule(&self, task: task::Notified<Self>) { + Shared::schedule(self, task); + } + + cfg_unstable! { + fn unhandled_panic(&self) { + use crate::runtime::UnhandledPanic; + + match self.unhandled_panic { + UnhandledPanic::Ignore => { + // Do nothing + } + UnhandledPanic::ShutdownRuntime => { + // This hook is only called from within the runtime, so + // `CURRENT` should match with `&self`, i.e. there is no + // opportunity for a nested scheduler to be called. + CURRENT.with(|LocalData { ctx, .. }| match ctx.get() { + Some(cx) if Arc::ptr_eq(self, &cx.shared) => { + cx.unhandled_panic.set(true); + // Safety: this is always called from the thread that owns `LocalSet` + unsafe { cx.shared.local_state.close_and_shutdown_all(); } + } + _ => unreachable!("runtime core not set in CURRENT thread-local"), + }) + } + } + } + } +} - CURRENT.with(|maybe_cx| { - let cx = maybe_cx.expect("scheduler context missing"); +impl LocalState { + unsafe fn task_pop_front(&self) -> Option<task::Notified<Arc<Shared>>> { + // The caller ensures it is called from the same thread that owns + // the LocalSet. + self.assert_called_from_owner_thread(); - assert!(cx.shared.ptr_eq(self)); + self.local_queue.with_mut(|ptr| (*ptr).pop_front()) + } - let ptr = NonNull::from(task.header()); - // safety: task must be contained by list. It is inserted into the - // list in `bind`. - unsafe { cx.tasks.borrow_mut().owned.remove(ptr) } - }) + unsafe fn task_push_back(&self, task: task::Notified<Arc<Shared>>) { + // The caller ensures it is called from the same thread that owns + // the LocalSet. + self.assert_called_from_owner_thread(); + + self.local_queue.with_mut(|ptr| (*ptr).push_back(task)) } - fn schedule(&self, task: task::Notified<Self>) { - Shared::schedule(self, task); + unsafe fn take_local_queue(&self) -> VecDeque<task::Notified<Arc<Shared>>> { + // The caller ensures it is called from the same thread that owns + // the LocalSet. + self.assert_called_from_owner_thread(); + + self.local_queue.with_mut(|ptr| std::mem::take(&mut (*ptr))) + } + + unsafe fn task_remove(&self, task: &Task<Arc<Shared>>) -> Option<Task<Arc<Shared>>> { + // The caller ensures it is called from the same thread that owns + // the LocalSet. + self.assert_called_from_owner_thread(); + + self.owned.remove(task) + } + + /// Returns true if the `LocalSet` does not have any spawned tasks + unsafe fn owned_is_empty(&self) -> bool { + // The caller ensures it is called from the same thread that owns + // the LocalSet. + self.assert_called_from_owner_thread(); + + self.owned.is_empty() + } + + unsafe fn assert_owner( + &self, + task: task::Notified<Arc<Shared>>, + ) -> task::LocalNotified<Arc<Shared>> { + // The caller ensures it is called from the same thread that owns + // the LocalSet. + self.assert_called_from_owner_thread(); + + self.owned.assert_owner(task) + } + + unsafe fn close_and_shutdown_all(&self) { + // The caller ensures it is called from the same thread that owns + // the LocalSet. + self.assert_called_from_owner_thread(); + + self.owned.close_and_shutdown_all() + } + + #[track_caller] + fn assert_called_from_owner_thread(&self) { + // FreeBSD has some weirdness around thread-local destruction. + // TODO: remove this hack when thread id is cleaned up + #[cfg(not(any(target_os = "openbsd", target_os = "freebsd")))] + debug_assert!( + // if we couldn't get the thread ID because we're dropping the local + // data, skip the assertion --- the `Drop` impl is not going to be + // called from another thread, because `LocalSet` is `!Send` + context::thread_id() + .map(|id| id == self.owner) + .unwrap_or(true), + "`LocalSet`'s local run queue must not be accessed by another thread!" + ); + } +} + +// This is `Send` because it is stored in `Shared`. It is up to the caller to +// ensure they are on the same thread that owns the `LocalSet`. +unsafe impl Send for LocalState {} + +#[cfg(all(test, not(loom)))] +mod tests { + use super::*; + + // Does a `LocalSet` running on a current-thread runtime...basically work? + // + // This duplicates a test in `tests/task_local_set.rs`, but because this is + // a lib test, it wil run under Miri, so this is necessary to catch stacked + // borrows violations in the `LocalSet` implementation. + #[test] + fn local_current_thread_scheduler() { + let f = async { + LocalSet::new() + .run_until(async { + spawn_local(async {}).await.unwrap(); + }) + .await; + }; + crate::runtime::Builder::new_current_thread() + .build() + .expect("rt") + .block_on(f) + } + + // Tests that when a task on a `LocalSet` is woken by an io driver on the + // same thread, the task is woken to the localset's local queue rather than + // its remote queue. + // + // This test has to be defined in the `local.rs` file as a lib test, rather + // than in `tests/`, because it makes assertions about the local set's + // internal state. + #[test] + fn wakes_to_local_queue() { + use super::*; + use crate::sync::Notify; + let rt = crate::runtime::Builder::new_current_thread() + .build() + .expect("rt"); + rt.block_on(async { + let local = LocalSet::new(); + let notify = Arc::new(Notify::new()); + let task = local.spawn_local({ + let notify = notify.clone(); + async move { + notify.notified().await; + } + }); + let mut run_until = Box::pin(local.run_until(async move { + task.await.unwrap(); + })); + + // poll the run until future once + crate::future::poll_fn(|cx| { + let _ = run_until.as_mut().poll(cx); + Poll::Ready(()) + }) + .await; + + notify.notify_one(); + let task = unsafe { local.context.shared.local_state.task_pop_front() }; + // TODO(eliza): it would be nice to be able to assert that this is + // the local task. + assert!( + task.is_some(), + "task should have been notified to the LocalSet's local queue" + ); + }) } } diff --git a/vendor/tokio/src/task/mod.rs b/vendor/tokio/src/task/mod.rs index ae4c35c9c..9b7537018 100644 --- a/vendor/tokio/src/task/mod.rs +++ b/vendor/tokio/src/task/mod.rs @@ -243,7 +243,7 @@ //! //! #### unconstrained //! -//! If necessary, [`task::unconstrained`] lets you opt out a future of Tokio's cooperative +//! If necessary, [`task::unconstrained`] lets you opt a future out of of Tokio's cooperative //! scheduling. When a future is wrapped with `unconstrained`, it will never be forced to yield to //! Tokio. For example: //! @@ -278,8 +278,10 @@ cfg_rt! { pub use crate::runtime::task::{JoinError, JoinHandle}; - mod blocking; - pub use blocking::spawn_blocking; + cfg_not_wasi! { + mod blocking; + pub use blocking::spawn_blocking; + } mod spawn; pub use spawn::spawn; @@ -291,8 +293,13 @@ cfg_rt! { mod yield_now; pub use yield_now::yield_now; + cfg_unstable! { + mod consume_budget; + pub use consume_budget::consume_budget; + } + mod local; - pub use local::{spawn_local, LocalSet}; + pub use local::{spawn_local, LocalSet, LocalEnterGuard}; mod task_local; pub use task_local::LocalKey; @@ -300,8 +307,27 @@ cfg_rt! { mod unconstrained; pub use unconstrained::{unconstrained, Unconstrained}; + #[doc(inline)] + pub use join_set::JoinSet; + pub use crate::runtime::task::AbortHandle; + + // Uses #[cfg(...)] instead of macro since the macro adds docsrs annotations. + #[cfg(not(tokio_unstable))] + mod join_set; + #[cfg(tokio_unstable)] + pub mod join_set; + + cfg_unstable! { + pub use crate::runtime::task::{Id, id, try_id}; + } + cfg_trace! { mod builder; pub use builder::Builder; } + + /// Task-related futures. + pub mod futures { + pub use super::task_local::TaskLocalFuture; + } } diff --git a/vendor/tokio/src/task/spawn.rs b/vendor/tokio/src/task/spawn.rs index 3c577b82d..23a46d745 100644 --- a/vendor/tokio/src/task/spawn.rs +++ b/vendor/tokio/src/task/spawn.rs @@ -1,4 +1,4 @@ -use crate::{task::JoinHandle, util::error::CONTEXT_MISSING_ERROR}; +use crate::task::JoinHandle; use std::future::Future; @@ -6,11 +6,19 @@ cfg_rt! { /// Spawns a new asynchronous task, returning a /// [`JoinHandle`](super::JoinHandle) for it. /// + /// The provided future will start running in the background immediately + /// when `spawn` is called, even if you don't await the returned + /// `JoinHandle`. + /// /// Spawning a task enables the task to execute concurrently to other tasks. The /// spawned task may execute on the current thread, or it may be sent to a /// different thread to be executed. The specifics depend on the current /// [`Runtime`](crate::runtime::Runtime) configuration. /// + /// It is guaranteed that spawn will not synchronously poll the task being spawned. + /// This means that calling spawn while holding a lock does not pose a risk of + /// deadlocking with the spawned task. + /// /// There is no guarantee that a spawned task will execute to completion. /// When a runtime is shutdown, all outstanding tasks are dropped, /// regardless of the lifecycle of that task. @@ -49,6 +57,37 @@ cfg_rt! { /// } /// ``` /// + /// To run multiple tasks in parallel and receive their results, join + /// handles can be stored in a vector. + /// ``` + /// # #[tokio::main(flavor = "current_thread")] async fn main() { + /// async fn my_background_op(id: i32) -> String { + /// let s = format!("Starting background task {}.", id); + /// println!("{}", s); + /// s + /// } + /// + /// let ops = vec![1, 2, 3]; + /// let mut tasks = Vec::with_capacity(ops.len()); + /// for op in ops { + /// // This call will make them start running in the background + /// // immediately. + /// tasks.push(tokio::spawn(my_background_op(op))); + /// } + /// + /// let mut outputs = Vec::with_capacity(tasks.len()); + /// for task in tasks { + /// outputs.push(task.await.unwrap()); + /// } + /// println!("{:?}", outputs); + /// # } + /// ``` + /// This example pushes the tasks to `outputs` in the order they were + /// started in. If you do not care about the ordering of the outputs, then + /// you can also use a [`JoinSet`]. + /// + /// [`JoinSet`]: struct@crate::task::JoinSet + /// /// # Panics /// /// Panics if called from **outside** of the Tokio runtime. @@ -121,23 +160,47 @@ cfg_rt! { /// ```text /// error[E0391]: cycle detected when processing `main` /// ``` - #[cfg_attr(tokio_track_caller, track_caller)] + #[track_caller] pub fn spawn<T>(future: T) -> JoinHandle<T::Output> where T: Future + Send + 'static, T::Output: Send + 'static, { - spawn_inner(future, None) + // preventing stack overflows on debug mode, by quickly sending the + // task to the heap. + if cfg!(debug_assertions) && std::mem::size_of::<T>() > 2048 { + spawn_inner(Box::pin(future), None) + } else { + spawn_inner(future, None) + } } - #[cfg_attr(tokio_track_caller, track_caller)] + #[track_caller] pub(super) fn spawn_inner<T>(future: T, name: Option<&str>) -> JoinHandle<T::Output> where T: Future + Send + 'static, T::Output: Send + 'static, { - let spawn_handle = crate::runtime::context::spawn_handle().expect(CONTEXT_MISSING_ERROR); - let task = crate::util::trace::task(future, "task", name); - spawn_handle.spawn(task) + use crate::runtime::{context, task}; + + #[cfg(all( + tokio_unstable, + tokio_taskdump, + feature = "rt", + target_os = "linux", + any( + target_arch = "aarch64", + target_arch = "x86", + target_arch = "x86_64" + ) + ))] + let future = task::trace::Trace::root(future); + let id = task::Id::next(); + let task = crate::util::trace::task(future, "task", name, id.as_u64()); + + match context::with_current(|handle| handle.spawn(task, id)) { + Ok(join_handle) => join_handle, + Err(e) => panic!("{}", e), + } } } diff --git a/vendor/tokio/src/task/task_local.rs b/vendor/tokio/src/task/task_local.rs index 6571ffd7b..eeadfbd3e 100644 --- a/vendor/tokio/src/task/task_local.rs +++ b/vendor/tokio/src/task/task_local.rs @@ -2,9 +2,10 @@ use pin_project_lite::pin_project; use std::cell::RefCell; use std::error::Error; use std::future::Future; +use std::marker::PhantomPinned; use std::pin::Pin; use std::task::{Context, Poll}; -use std::{fmt, thread}; +use std::{fmt, mem, thread}; /// Declares a new task-local key of type [`tokio::task::LocalKey`]. /// @@ -47,9 +48,27 @@ macro_rules! task_local { } #[doc(hidden)] +#[cfg(not(tokio_no_const_thread_local))] #[macro_export] macro_rules! __task_local_inner { ($(#[$attr:meta])* $vis:vis $name:ident, $t:ty) => { + $(#[$attr])* + $vis static $name: $crate::task::LocalKey<$t> = { + std::thread_local! { + static __KEY: std::cell::RefCell<Option<$t>> = const { std::cell::RefCell::new(None) }; + } + + $crate::task::LocalKey { inner: __KEY } + }; + }; +} + +#[doc(hidden)] +#[cfg(tokio_no_const_thread_local)] +#[macro_export] +macro_rules! __task_local_inner { + ($(#[$attr:meta])* $vis:vis $name:ident, $t:ty) => { + $(#[$attr])* $vis static $name: $crate::task::LocalKey<$t> = { std::thread_local! { static __KEY: std::cell::RefCell<Option<$t>> = std::cell::RefCell::new(None); @@ -62,7 +81,7 @@ macro_rules! __task_local_inner { /// A key for task-local data. /// -/// This type is generated by the `task_local!` macro. +/// This type is generated by the [`task_local!`] macro. /// /// Unlike [`std::thread::LocalKey`], `tokio::task::LocalKey` will /// _not_ lazily initialize the value on first access. Instead, the @@ -90,7 +109,9 @@ macro_rules! __task_local_inner { /// }).await; /// # } /// ``` +/// /// [`std::thread::LocalKey`]: struct@std::thread::LocalKey +/// [`task_local!`]: ../macro.task_local.html #[cfg_attr(docsrs, doc(cfg(feature = "rt")))] pub struct LocalKey<T: 'static> { #[doc(hidden)] @@ -102,6 +123,11 @@ impl<T: 'static> LocalKey<T> { /// /// On completion of `scope`, the task-local will be dropped. /// + /// ### Panics + /// + /// If you poll the returned future inside a call to [`with`] or + /// [`try_with`] on the same `LocalKey`, then the call to `poll` will panic. + /// /// ### Examples /// /// ``` @@ -115,21 +141,29 @@ impl<T: 'static> LocalKey<T> { /// }).await; /// # } /// ``` - pub async fn scope<F>(&'static self, value: T, f: F) -> F::Output + /// + /// [`with`]: fn@Self::with + /// [`try_with`]: fn@Self::try_with + pub fn scope<F>(&'static self, value: T, f: F) -> TaskLocalFuture<T, F> where F: Future, { TaskLocalFuture { - local: &self, + local: self, slot: Some(value), - future: f, + future: Some(f), + _pinned: PhantomPinned, } - .await } /// Sets a value `T` as the task-local value for the closure `F`. /// - /// On completion of `scope`, the task-local will be dropped. + /// On completion of `sync_scope`, the task-local will be dropped. + /// + /// ### Panics + /// + /// This method panics if called inside a call to [`with`] or [`try_with`] + /// on the same `LocalKey`. /// /// ### Examples /// @@ -144,32 +178,80 @@ impl<T: 'static> LocalKey<T> { /// }); /// # } /// ``` + /// + /// [`with`]: fn@Self::with + /// [`try_with`]: fn@Self::try_with + #[track_caller] pub fn sync_scope<F, R>(&'static self, value: T, f: F) -> R where F: FnOnce() -> R, { - let mut scope = TaskLocalFuture { - local: &self, - slot: Some(value), - future: (), - }; - Pin::new(&mut scope).with_task(|_| f()) + let mut value = Some(value); + match self.scope_inner(&mut value, f) { + Ok(res) => res, + Err(err) => err.panic(), + } + } + + fn scope_inner<F, R>(&'static self, slot: &mut Option<T>, f: F) -> Result<R, ScopeInnerErr> + where + F: FnOnce() -> R, + { + struct Guard<'a, T: 'static> { + local: &'static LocalKey<T>, + slot: &'a mut Option<T>, + } + + impl<'a, T: 'static> Drop for Guard<'a, T> { + fn drop(&mut self) { + // This should not panic. + // + // We know that the RefCell was not borrowed before the call to + // `scope_inner`, so the only way for this to panic is if the + // closure has created but not destroyed a RefCell guard. + // However, we never give user-code access to the guards, so + // there's no way for user-code to forget to destroy a guard. + // + // The call to `with` also should not panic, since the + // thread-local wasn't destroyed when we first called + // `scope_inner`, and it shouldn't have gotten destroyed since + // then. + self.local.inner.with(|inner| { + let mut ref_mut = inner.borrow_mut(); + mem::swap(self.slot, &mut *ref_mut); + }); + } + } + + self.inner.try_with(|inner| { + inner + .try_borrow_mut() + .map(|mut ref_mut| mem::swap(slot, &mut *ref_mut)) + })??; + + let guard = Guard { local: self, slot }; + + let res = f(); + + drop(guard); + + Ok(res) } /// Accesses the current task-local and runs the provided closure. /// /// # Panics /// - /// This function will panic if not called within the context - /// of a future containing a task-local with the corresponding key. + /// This function will panic if the task local doesn't have a value set. + #[track_caller] pub fn with<F, R>(&'static self, f: F) -> R where F: FnOnce(&T) -> R, { - self.try_with(f).expect( - "cannot access a Task Local Storage value \ - without setting it via `LocalKey::set`", - ) + match self.try_with(f) { + Ok(res) => res, + Err(_) => panic!("cannot access a task-local storage value without setting it first"), + } } /// Accesses the current task-local and runs the provided closure. @@ -181,19 +263,32 @@ impl<T: 'static> LocalKey<T> { where F: FnOnce(&T) -> R, { - self.inner.with(|v| { - if let Some(val) = v.borrow().as_ref() { - Ok(f(val)) - } else { - Err(AccessError { _private: () }) - } - }) + // If called after the thread-local storing the task-local is destroyed, + // then we are outside of a closure where the task-local is set. + // + // Therefore, it is correct to return an AccessError if `try_with` + // returns an error. + let try_with_res = self.inner.try_with(|v| { + // This call to `borrow` cannot panic because no user-defined code + // runs while a `borrow_mut` call is active. + v.borrow().as_ref().map(f) + }); + + match try_with_res { + Ok(Some(res)) => Ok(res), + Ok(None) | Err(_) => Err(AccessError { _private: () }), + } } } impl<T: Copy + 'static> LocalKey<T> { /// Returns a copy of the task-local value /// if the task-local value implements `Copy`. + /// + /// # Panics + /// + /// This function will panic if the task local doesn't have a value set. + #[track_caller] pub fn get(&'static self) -> T { self.with(|v| *v) } @@ -206,55 +301,107 @@ impl<T: 'static> fmt::Debug for LocalKey<T> { } pin_project! { - struct TaskLocalFuture<T: StaticLifetime, F> { + /// A future that sets a value `T` of a task local for the future `F` during + /// its execution. + /// + /// The value of the task-local must be `'static` and will be dropped on the + /// completion of the future. + /// + /// Created by the function [`LocalKey::scope`](self::LocalKey::scope). + /// + /// ### Examples + /// + /// ``` + /// # async fn dox() { + /// tokio::task_local! { + /// static NUMBER: u32; + /// } + /// + /// NUMBER.scope(1, async move { + /// println!("task local value: {}", NUMBER.get()); + /// }).await; + /// # } + /// ``` + pub struct TaskLocalFuture<T, F> + where + T: 'static, + { local: &'static LocalKey<T>, slot: Option<T>, #[pin] - future: F, + future: Option<F>, + #[pin] + _pinned: PhantomPinned, } -} -impl<T: 'static, F> TaskLocalFuture<T, F> { - fn with_task<F2: FnOnce(Pin<&mut F>) -> R, R>(self: Pin<&mut Self>, f: F2) -> R { - struct Guard<'a, T: 'static> { - local: &'static LocalKey<T>, - slot: &'a mut Option<T>, - prev: Option<T>, - } - - impl<T> Drop for Guard<'_, T> { - fn drop(&mut self) { - let value = self.local.inner.with(|c| c.replace(self.prev.take())); - *self.slot = value; + impl<T: 'static, F> PinnedDrop for TaskLocalFuture<T, F> { + fn drop(this: Pin<&mut Self>) { + let this = this.project(); + if mem::needs_drop::<F>() && this.future.is_some() { + // Drop the future while the task-local is set, if possible. Otherwise + // the future is dropped normally when the `Option<F>` field drops. + let mut future = this.future; + let _ = this.local.scope_inner(this.slot, || { + future.set(None); + }); } } - - let mut project = self.project(); - let val = project.slot.take(); - - let prev = project.local.inner.with(|c| c.replace(val)); - - let _guard = Guard { - prev, - slot: &mut project.slot, - local: *project.local, - }; - - f(project.future) } } impl<T: 'static, F: Future> Future for TaskLocalFuture<T, F> { type Output = F::Output; + #[track_caller] fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { - self.with_task(|f| f.poll(cx)) + let this = self.project(); + let mut future_opt = this.future; + + let res = this + .local + .scope_inner(this.slot, || match future_opt.as_mut().as_pin_mut() { + Some(fut) => { + let res = fut.poll(cx); + if res.is_ready() { + future_opt.set(None); + } + Some(res) + } + None => None, + }); + + match res { + Ok(Some(res)) => res, + Ok(None) => panic!("`TaskLocalFuture` polled after completion"), + Err(err) => err.panic(), + } } } -// Required to make `pin_project` happy. -trait StaticLifetime: 'static {} -impl<T: 'static> StaticLifetime for T {} +impl<T: 'static, F> fmt::Debug for TaskLocalFuture<T, F> +where + T: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + /// Format the Option without Some. + struct TransparentOption<'a, T> { + value: &'a Option<T>, + } + impl<'a, T: fmt::Debug> fmt::Debug for TransparentOption<'a, T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self.value.as_ref() { + Some(value) => value.fmt(f), + // Hitting the None branch should not be possible. + None => f.pad("<missing>"), + } + } + } + + f.debug_struct("TaskLocalFuture") + .field("value", &TransparentOption { value: &self.slot }) + .finish() + } +} /// An error returned by [`LocalKey::try_with`](method@LocalKey::try_with). #[derive(Clone, Copy, Eq, PartialEq)] @@ -275,3 +422,30 @@ impl fmt::Display for AccessError { } impl Error for AccessError {} + +enum ScopeInnerErr { + BorrowError, + AccessError, +} + +impl ScopeInnerErr { + #[track_caller] + fn panic(&self) -> ! { + match self { + Self::BorrowError => panic!("cannot enter a task-local scope while the task-local storage is borrowed"), + Self::AccessError => panic!("cannot enter a task-local scope during or after destruction of the underlying thread-local"), + } + } +} + +impl From<std::cell::BorrowMutError> for ScopeInnerErr { + fn from(_: std::cell::BorrowMutError) -> Self { + Self::BorrowError + } +} + +impl From<std::thread::AccessError> for ScopeInnerErr { + fn from(_: std::thread::AccessError) -> Self { + Self::AccessError + } +} diff --git a/vendor/tokio/src/task/unconstrained.rs b/vendor/tokio/src/task/unconstrained.rs index 31c732bfc..40384c870 100644 --- a/vendor/tokio/src/task/unconstrained.rs +++ b/vendor/tokio/src/task/unconstrained.rs @@ -22,7 +22,7 @@ where cfg_coop! { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { let inner = self.project().inner; - crate::coop::with_unconstrained(|| inner.poll(cx)) + crate::runtime::coop::with_unconstrained(|| inner.poll(cx)) } } diff --git a/vendor/tokio/src/task/yield_now.rs b/vendor/tokio/src/task/yield_now.rs index 251cb931b..428d124c3 100644 --- a/vendor/tokio/src/task/yield_now.rs +++ b/vendor/tokio/src/task/yield_now.rs @@ -1,38 +1,64 @@ +use crate::runtime::context; + use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; -cfg_rt! { - /// Yields execution back to the Tokio runtime. - /// - /// A task yields by awaiting on `yield_now()`, and may resume when that - /// future completes (with no output.) The current task will be re-added as - /// a pending task at the _back_ of the pending queue. Any other pending - /// tasks will be scheduled. No other waking is required for the task to - /// continue. - /// - /// See also the usage example in the [task module](index.html#yield_now). - #[must_use = "yield_now does nothing unless polled/`await`-ed"] - pub async fn yield_now() { - /// Yield implementation - struct YieldNow { - yielded: bool, - } +/// Yields execution back to the Tokio runtime. +/// +/// A task yields by awaiting on `yield_now()`, and may resume when that future +/// completes (with no output.) The current task will be re-added as a pending +/// task at the _back_ of the pending queue. Any other pending tasks will be +/// scheduled. No other waking is required for the task to continue. +/// +/// See also the usage example in the [task module](index.html#yield_now). +/// +/// ## Non-guarantees +/// +/// This function may not yield all the way up to the executor if there are any +/// special combinators above it in the call stack. For example, if a +/// [`tokio::select!`] has another branch complete during the same poll as the +/// `yield_now()`, then the yield is not propagated all the way up to the +/// runtime. +/// +/// It is generally not guaranteed that the runtime behaves like you expect it +/// to when deciding which task to schedule next after a call to `yield_now()`. +/// In particular, the runtime may choose to poll the task that just ran +/// `yield_now()` again immediately without polling any other tasks first. For +/// example, the runtime will not drive the IO driver between every poll of a +/// task, and this could result in the runtime polling the current task again +/// immediately even if there is another task that could make progress if that +/// other task is waiting for a notification from the IO driver. +/// +/// In general, changes to the order in which the runtime polls tasks is not +/// considered a breaking change, and your program should be correct no matter +/// which order the runtime polls your tasks in. +/// +/// [`tokio::select!`]: macro@crate::select +#[cfg_attr(docsrs, doc(cfg(feature = "rt")))] +pub async fn yield_now() { + /// Yield implementation + struct YieldNow { + yielded: bool, + } - impl Future for YieldNow { - type Output = (); + impl Future for YieldNow { + type Output = (); - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { - if self.yielded { - return Poll::Ready(()); - } + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { + ready!(crate::trace::trace_leaf(cx)); - self.yielded = true; - cx.waker().wake_by_ref(); - Poll::Pending + if self.yielded { + return Poll::Ready(()); } - } - YieldNow { yielded: false }.await + self.yielded = true; + + context::defer(cx.waker()); + + Poll::Pending + } } + + YieldNow { yielded: false }.await } diff --git a/vendor/tokio/src/time/clock.rs b/vendor/tokio/src/time/clock.rs index a44d75f3c..091cf4b19 100644 --- a/vendor/tokio/src/time/clock.rs +++ b/vendor/tokio/src/time/clock.rs @@ -29,26 +29,53 @@ cfg_not_test_util! { cfg_test_util! { use crate::time::{Duration, Instant}; - use crate::loom::sync::{Arc, Mutex}; + use crate::loom::sync::Mutex; + use crate::loom::sync::atomic::Ordering; + use std::sync::atomic::AtomicBool as StdAtomicBool; cfg_rt! { - fn clock() -> Option<Clock> { - crate::runtime::context::clock() + #[track_caller] + fn with_clock<R>(f: impl FnOnce(Option<&Clock>) -> Result<R, &'static str>) -> R { + use crate::runtime::Handle; + + let res = match Handle::try_current() { + Ok(handle) => f(Some(handle.inner.driver().clock())), + Err(ref e) if e.is_missing_context() => f(None), + Err(_) => panic!("{}", crate::util::error::THREAD_LOCAL_DESTROYED_ERROR), + }; + + match res { + Ok(ret) => ret, + Err(msg) => panic!("{}", msg), + } } } cfg_not_rt! { - fn clock() -> Option<Clock> { - None + #[track_caller] + fn with_clock<R>(f: impl FnOnce(Option<&Clock>) -> Result<R, &'static str>) -> R { + match f(None) { + Ok(ret) => ret, + Err(msg) => panic!("{}", msg), + } } } /// A handle to a source of time. - #[derive(Debug, Clone)] + #[derive(Debug)] pub(crate) struct Clock { - inner: Arc<Mutex<Inner>>, + inner: Mutex<Inner>, } + // Used to track if the clock was ever paused. This is an optimization to + // avoid touching the mutex if `test-util` was accidentally enabled in + // release mode. + // + // A static is used so we can avoid accessing the thread-local as well. The + // `std` AtomicBool is used directly because loom does not support static + // atomics. + static DID_PAUSE_CLOCK: StdAtomicBool = StdAtomicBool::new(false); + #[derive(Debug)] struct Inner { /// True if the ability to pause time is enabled. @@ -57,22 +84,34 @@ cfg_test_util! { /// Instant to use as the clock's base instant. base: std::time::Instant, - /// Instant at which the clock was last unfrozen + /// Instant at which the clock was last unfrozen. unfrozen: Option<std::time::Instant>, + + /// Number of `inhibit_auto_advance` calls still in effect. + auto_advance_inhibit_count: usize, } - /// Pause time + /// Pauses time. /// /// The current value of `Instant::now()` is saved and all subsequent calls - /// to `Instant::now()` until the timer wheel is checked again will return - /// the saved value. Once the timer wheel is checked, time will immediately - /// advance to the next registered `Sleep`. This is useful for running tests - /// that depend on time. + /// to `Instant::now()` will return the saved value. The saved value can be + /// changed by [`advance`] or by the time auto-advancing once the runtime + /// has no work to do. This only affects the `Instant` type in Tokio, and + /// the `Instant` in std continues to work as normal. /// /// Pausing time requires the `current_thread` Tokio runtime. This is the /// default runtime used by `#[tokio::test]`. The runtime can be initialized /// with time in a paused state using the `Builder::start_paused` method. /// + /// For cases where time is immediately paused, it is better to pause + /// the time using the `main` or `test` macro: + /// ``` + /// #[tokio::main(flavor = "current_thread", start_paused = true)] + /// async fn main() { + /// println!("Hello world"); + /// } + /// ``` + /// /// # Panics /// /// Panics if time is already frozen or if called from outside of a @@ -86,12 +125,18 @@ cfg_test_util! { /// current time when awaited. /// /// [`Sleep`]: crate::time::Sleep + /// [`advance`]: crate::time::advance + #[track_caller] pub fn pause() { - let clock = clock().expect("time cannot be frozen from outside the Tokio runtime"); - clock.pause(); + with_clock(|maybe_clock| { + match maybe_clock { + Some(clock) => clock.pause(), + None => Err("time cannot be frozen from outside the Tokio runtime"), + } + }) } - /// Resume time + /// Resumes time. /// /// Clears the saved `Instant::now()` value. Subsequent calls to /// `Instant::now()` will return the value returned by the system call. @@ -100,22 +145,44 @@ cfg_test_util! { /// /// Panics if time is not frozen or if called from outside of the Tokio /// runtime. + #[track_caller] pub fn resume() { - let clock = clock().expect("time cannot be frozen from outside the Tokio runtime"); - let mut inner = clock.inner.lock(); + with_clock(|maybe_clock| { + let clock = match maybe_clock { + Some(clock) => clock, + None => return Err("time cannot be frozen from outside the Tokio runtime"), + }; - if inner.unfrozen.is_some() { - panic!("time is not frozen"); - } + let mut inner = clock.inner.lock(); + + if inner.unfrozen.is_some() { + return Err("time is not frozen"); + } - inner.unfrozen = Some(std::time::Instant::now()); + inner.unfrozen = Some(std::time::Instant::now()); + Ok(()) + }) } - /// Advance time + /// Advances time. /// /// Increments the saved `Instant::now()` value by `duration`. Subsequent /// calls to `Instant::now()` will return the result of the increment. /// + /// This function will make the current time jump forward by the given + /// duration in one jump. This means that all `sleep` calls with a deadline + /// before the new time will immediately complete "at the same time", and + /// the runtime is free to poll them in any order. Additionally, this + /// method will not wait for the `sleep` calls it advanced past to complete. + /// If you want to do that, you should instead call [`sleep`] and rely on + /// the runtime's auto-advance feature. + /// + /// Note that calls to `sleep` are not guaranteed to complete the first time + /// they are polled after a call to `advance`. For example, this can happen + /// if the runtime has not yet touched the timer driver after the call to + /// `advance`. However if they don't, the runtime will poll the task again + /// shortly. + /// /// # Panics /// /// Panics if time is not frozen or if called from outside of the Tokio @@ -126,70 +193,107 @@ cfg_test_util! { /// If the time is paused and there is no work to do, the runtime advances /// time to the next timer. See [`pause`](pause#auto-advance) for more /// details. + /// + /// [`sleep`]: fn@crate::time::sleep pub async fn advance(duration: Duration) { - let clock = clock().expect("time cannot be frozen from outside the Tokio runtime"); - clock.advance(duration); + with_clock(|maybe_clock| { + let clock = match maybe_clock { + Some(clock) => clock, + None => return Err("time cannot be frozen from outside the Tokio runtime"), + }; + + clock.advance(duration) + }); crate::task::yield_now().await; } - /// Return the current instant, factoring in frozen time. + /// Returns the current instant, factoring in frozen time. pub(crate) fn now() -> Instant { - if let Some(clock) = clock() { - clock.now() - } else { - Instant::from_std(std::time::Instant::now()) + if !DID_PAUSE_CLOCK.load(Ordering::Acquire) { + return Instant::from_std(std::time::Instant::now()); } + + with_clock(|maybe_clock| { + Ok(if let Some(clock) = maybe_clock { + clock.now() + } else { + Instant::from_std(std::time::Instant::now()) + }) + }) } impl Clock { - /// Return a new `Clock` instance that uses the current execution context's + /// Returns a new `Clock` instance that uses the current execution context's /// source of time. pub(crate) fn new(enable_pausing: bool, start_paused: bool) -> Clock { let now = std::time::Instant::now(); let clock = Clock { - inner: Arc::new(Mutex::new(Inner { + inner: Mutex::new(Inner { enable_pausing, base: now, unfrozen: Some(now), - })), + auto_advance_inhibit_count: 0, + }), }; if start_paused { - clock.pause(); + if let Err(msg) = clock.pause() { + panic!("{}", msg); + } } clock } - pub(crate) fn pause(&self) { + pub(crate) fn pause(&self) -> Result<(), &'static str> { let mut inner = self.inner.lock(); if !inner.enable_pausing { drop(inner); // avoid poisoning the lock - panic!("`time::pause()` requires the `current_thread` Tokio runtime. \ + return Err("`time::pause()` requires the `current_thread` Tokio runtime. \ This is the default Runtime used by `#[tokio::test]."); } - let elapsed = inner.unfrozen.as_ref().expect("time is already frozen").elapsed(); + // Track that we paused the clock + DID_PAUSE_CLOCK.store(true, Ordering::Release); + + let elapsed = match inner.unfrozen.as_ref() { + Some(v) => v.elapsed(), + None => return Err("time is already frozen") + }; inner.base += elapsed; inner.unfrozen = None; + + Ok(()) + } + + /// Temporarily stop auto-advancing the clock (see `tokio::time::pause`). + pub(crate) fn inhibit_auto_advance(&self) { + let mut inner = self.inner.lock(); + inner.auto_advance_inhibit_count += 1; + } + + pub(crate) fn allow_auto_advance(&self) { + let mut inner = self.inner.lock(); + inner.auto_advance_inhibit_count -= 1; } - pub(crate) fn is_paused(&self) -> bool { + pub(crate) fn can_auto_advance(&self) -> bool { let inner = self.inner.lock(); - inner.unfrozen.is_none() + inner.unfrozen.is_none() && inner.auto_advance_inhibit_count == 0 } - pub(crate) fn advance(&self, duration: Duration) { + pub(crate) fn advance(&self, duration: Duration) -> Result<(), &'static str> { let mut inner = self.inner.lock(); if inner.unfrozen.is_some() { - panic!("time is not frozen"); + return Err("time is not frozen"); } inner.base += duration; + Ok(()) } pub(crate) fn now(&self) -> Instant { diff --git a/vendor/tokio/src/time/driver/handle.rs b/vendor/tokio/src/time/driver/handle.rs deleted file mode 100644 index 77b435873..000000000 --- a/vendor/tokio/src/time/driver/handle.rs +++ /dev/null @@ -1,88 +0,0 @@ -use crate::loom::sync::Arc; -use crate::time::driver::ClockTime; -use std::fmt; - -/// Handle to time driver instance. -#[derive(Clone)] -pub(crate) struct Handle { - time_source: ClockTime, - inner: Arc<super::Inner>, -} - -impl Handle { - /// Creates a new timer `Handle` from a shared `Inner` timer state. - pub(super) fn new(inner: Arc<super::Inner>) -> Self { - let time_source = inner.state.lock().time_source.clone(); - Handle { time_source, inner } - } - - /// Returns the time source associated with this handle - pub(super) fn time_source(&self) -> &ClockTime { - &self.time_source - } - - /// Access the driver's inner structure - pub(super) fn get(&self) -> &super::Inner { - &*self.inner - } - - // Check whether the driver has been shutdown - pub(super) fn is_shutdown(&self) -> bool { - self.inner.is_shutdown() - } -} - -cfg_rt! { - impl Handle { - /// Tries to get a handle to the current timer. - /// - /// # Panics - /// - /// This function panics if there is no current timer set. - /// - /// It can be triggered when `Builder::enable_time()` or - /// `Builder::enable_all()` are not included in the builder. - /// - /// It can also panic whenever a timer is created outside of a - /// Tokio runtime. That is why `rt.block_on(delay_for(...))` will panic, - /// since the function is executed outside of the runtime. - /// Whereas `rt.block_on(async {delay_for(...).await})` doesn't panic. - /// And this is because wrapping the function on an async makes it lazy, - /// and so gets executed inside the runtime successfully without - /// panicking. - pub(crate) fn current() -> Self { - crate::runtime::context::time_handle() - .expect("A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers.") - } - } -} - -cfg_not_rt! { - impl Handle { - /// Tries to get a handle to the current timer. - /// - /// # Panics - /// - /// This function panics if there is no current timer set. - /// - /// It can be triggered when `Builder::enable_time()` or - /// `Builder::enable_all()` are not included in the builder. - /// - /// It can also panic whenever a timer is created outside of a Tokio - /// runtime. That is why `rt.block_on(delay_for(...))` will panic, - /// since the function is executed outside of the runtime. - /// Whereas `rt.block_on(async {delay_for(...).await})` doesn't - /// panic. And this is because wrapping the function on an async makes it - /// lazy, and so outside executed inside the runtime successfully without - /// panicking. - pub(crate) fn current() -> Self { - panic!("{}", crate::util::error::CONTEXT_MISSING_ERROR) - } - } -} - -impl fmt::Debug for Handle { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "Handle") - } -} diff --git a/vendor/tokio/src/time/driver/sleep.rs b/vendor/tokio/src/time/driver/sleep.rs deleted file mode 100644 index 40f745ad7..000000000 --- a/vendor/tokio/src/time/driver/sleep.rs +++ /dev/null @@ -1,257 +0,0 @@ -use crate::time::driver::{Handle, TimerEntry}; -use crate::time::{error::Error, Duration, Instant}; - -use pin_project_lite::pin_project; -use std::future::Future; -use std::pin::Pin; -use std::task::{self, Poll}; - -/// Waits until `deadline` is reached. -/// -/// No work is performed while awaiting on the sleep future to complete. `Sleep` -/// operates at millisecond granularity and should not be used for tasks that -/// require high-resolution timers. -/// -/// # Cancellation -/// -/// Canceling a sleep instance is done by dropping the returned future. No additional -/// cleanup work is required. -// Alias for old name in 0.x -#[cfg_attr(docsrs, doc(alias = "delay_until"))] -pub fn sleep_until(deadline: Instant) -> Sleep { - Sleep::new_timeout(deadline) -} - -/// Waits until `duration` has elapsed. -/// -/// Equivalent to `sleep_until(Instant::now() + duration)`. An asynchronous -/// analog to `std::thread::sleep`. -/// -/// No work is performed while awaiting on the sleep future to complete. `Sleep` -/// operates at millisecond granularity and should not be used for tasks that -/// require high-resolution timers. -/// -/// To run something regularly on a schedule, see [`interval`]. -/// -/// The maximum duration for a sleep is 68719476734 milliseconds (approximately 2.2 years). -/// -/// # Cancellation -/// -/// Canceling a sleep instance is done by dropping the returned future. No additional -/// cleanup work is required. -/// -/// # Examples -/// -/// Wait 100ms and print "100 ms have elapsed". -/// -/// ``` -/// use tokio::time::{sleep, Duration}; -/// -/// #[tokio::main] -/// async fn main() { -/// sleep(Duration::from_millis(100)).await; -/// println!("100 ms have elapsed"); -/// } -/// ``` -/// -/// [`interval`]: crate::time::interval() -// Alias for old name in 0.x -#[cfg_attr(docsrs, doc(alias = "delay_for"))] -#[cfg_attr(docsrs, doc(alias = "wait"))] -pub fn sleep(duration: Duration) -> Sleep { - match Instant::now().checked_add(duration) { - Some(deadline) => sleep_until(deadline), - None => sleep_until(Instant::far_future()), - } -} - -pin_project! { - /// Future returned by [`sleep`](sleep) and [`sleep_until`](sleep_until). - /// - /// This type does not implement the `Unpin` trait, which means that if you - /// use it with [`select!`] or by calling `poll`, you have to pin it first. - /// If you use it with `.await`, this does not apply. - /// - /// # Examples - /// - /// Wait 100ms and print "100 ms have elapsed". - /// - /// ``` - /// use tokio::time::{sleep, Duration}; - /// - /// #[tokio::main] - /// async fn main() { - /// sleep(Duration::from_millis(100)).await; - /// println!("100 ms have elapsed"); - /// } - /// ``` - /// - /// Use with [`select!`]. Pinning the `Sleep` with [`tokio::pin!`] is - /// necessary when the same `Sleep` is selected on multiple times. - /// ```no_run - /// use tokio::time::{self, Duration, Instant}; - /// - /// #[tokio::main] - /// async fn main() { - /// let sleep = time::sleep(Duration::from_millis(10)); - /// tokio::pin!(sleep); - /// - /// loop { - /// tokio::select! { - /// () = &mut sleep => { - /// println!("timer elapsed"); - /// sleep.as_mut().reset(Instant::now() + Duration::from_millis(50)); - /// }, - /// } - /// } - /// } - /// ``` - /// Use in a struct with boxing. By pinning the `Sleep` with a `Box`, the - /// `HasSleep` struct implements `Unpin`, even though `Sleep` does not. - /// ``` - /// use std::future::Future; - /// use std::pin::Pin; - /// use std::task::{Context, Poll}; - /// use tokio::time::Sleep; - /// - /// struct HasSleep { - /// sleep: Pin<Box<Sleep>>, - /// } - /// - /// impl Future for HasSleep { - /// type Output = (); - /// - /// fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { - /// self.sleep.as_mut().poll(cx) - /// } - /// } - /// ``` - /// Use in a struct with pin projection. This method avoids the `Box`, but - /// the `HasSleep` struct will not be `Unpin` as a consequence. - /// ``` - /// use std::future::Future; - /// use std::pin::Pin; - /// use std::task::{Context, Poll}; - /// use tokio::time::Sleep; - /// use pin_project_lite::pin_project; - /// - /// pin_project! { - /// struct HasSleep { - /// #[pin] - /// sleep: Sleep, - /// } - /// } - /// - /// impl Future for HasSleep { - /// type Output = (); - /// - /// fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { - /// self.project().sleep.poll(cx) - /// } - /// } - /// ``` - /// - /// [`select!`]: ../macro.select.html - /// [`tokio::pin!`]: ../macro.pin.html - // Alias for old name in 0.2 - #[cfg_attr(docsrs, doc(alias = "Delay"))] - #[derive(Debug)] - #[must_use = "futures do nothing unless you `.await` or poll them"] - pub struct Sleep { - deadline: Instant, - - // The link between the `Sleep` instance and the timer that drives it. - #[pin] - entry: TimerEntry, - } -} - -impl Sleep { - pub(crate) fn new_timeout(deadline: Instant) -> Sleep { - let handle = Handle::current(); - let entry = TimerEntry::new(&handle, deadline); - - Sleep { deadline, entry } - } - - pub(crate) fn far_future() -> Sleep { - Self::new_timeout(Instant::far_future()) - } - - /// Returns the instant at which the future will complete. - pub fn deadline(&self) -> Instant { - self.deadline - } - - /// Returns `true` if `Sleep` has elapsed. - /// - /// A `Sleep` instance is elapsed when the requested duration has elapsed. - pub fn is_elapsed(&self) -> bool { - self.entry.is_elapsed() - } - - /// Resets the `Sleep` instance to a new deadline. - /// - /// Calling this function allows changing the instant at which the `Sleep` - /// future completes without having to create new associated state. - /// - /// This function can be called both before and after the future has - /// completed. - /// - /// To call this method, you will usually combine the call with - /// [`Pin::as_mut`], which lets you call the method without consuming the - /// `Sleep` itself. - /// - /// # Example - /// - /// ``` - /// use tokio::time::{Duration, Instant}; - /// - /// # #[tokio::main(flavor = "current_thread")] - /// # async fn main() { - /// let sleep = tokio::time::sleep(Duration::from_millis(10)); - /// tokio::pin!(sleep); - /// - /// sleep.as_mut().reset(Instant::now() + Duration::from_millis(20)); - /// # } - /// ``` - /// - /// [`Pin::as_mut`]: fn@std::pin::Pin::as_mut - pub fn reset(self: Pin<&mut Self>, deadline: Instant) { - let me = self.project(); - me.entry.reset(deadline); - *me.deadline = deadline; - } - - fn poll_elapsed(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<(), Error>> { - let me = self.project(); - - // Keep track of task budget - let coop = ready!(crate::coop::poll_proceed(cx)); - - me.entry.poll_elapsed(cx).map(move |r| { - coop.made_progress(); - r - }) - } -} - -impl Future for Sleep { - type Output = (); - - fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> { - // `poll_elapsed` can return an error in two cases: - // - // - AtCapacity: this is a pathological case where far too many - // sleep instances have been scheduled. - // - Shutdown: No timer has been setup, which is a mis-use error. - // - // Both cases are extremely rare, and pretty accurately fit into - // "logic errors", so we just panic in this case. A user couldn't - // really do much better if we passed the error onwards. - match ready!(self.as_mut().poll_elapsed(cx)) { - Ok(()) => Poll::Ready(()), - Err(e) => panic!("timer error: {}", e), - } - } -} diff --git a/vendor/tokio/src/time/driver/wheel/stack.rs b/vendor/tokio/src/time/driver/wheel/stack.rs deleted file mode 100644 index e7ed137f5..000000000 --- a/vendor/tokio/src/time/driver/wheel/stack.rs +++ /dev/null @@ -1,112 +0,0 @@ -use super::{Item, OwnedItem}; -use crate::time::driver::Entry; - -use std::ptr; - -/// A doubly linked stack -#[derive(Debug)] -pub(crate) struct Stack { - head: Option<OwnedItem>, -} - -impl Default for Stack { - fn default() -> Stack { - Stack { head: None } - } -} - -impl Stack { - pub(crate) fn is_empty(&self) -> bool { - self.head.is_none() - } - - pub(crate) fn push(&mut self, entry: OwnedItem) { - // Get a pointer to the entry to for the prev link - let ptr: *const Entry = &*entry as *const _; - - // Remove the old head entry - let old = self.head.take(); - - unsafe { - // Ensure the entry is not already in a stack. - debug_assert!((*entry.next_stack.get()).is_none()); - debug_assert!((*entry.prev_stack.get()).is_null()); - - if let Some(ref entry) = old.as_ref() { - debug_assert!({ - // The head is not already set to the entry - ptr != &***entry as *const _ - }); - - // Set the previous link on the old head - *entry.prev_stack.get() = ptr; - } - - // Set this entry's next pointer - *entry.next_stack.get() = old; - } - - // Update the head pointer - self.head = Some(entry); - } - - /// Pops an item from the stack - pub(crate) fn pop(&mut self) -> Option<OwnedItem> { - let entry = self.head.take(); - - unsafe { - if let Some(entry) = entry.as_ref() { - self.head = (*entry.next_stack.get()).take(); - - if let Some(entry) = self.head.as_ref() { - *entry.prev_stack.get() = ptr::null(); - } - - *entry.prev_stack.get() = ptr::null(); - } - } - - entry - } - - pub(crate) fn remove(&mut self, entry: &Item) { - unsafe { - // Ensure that the entry is in fact contained by the stack - debug_assert!({ - // This walks the full linked list even if an entry is found. - let mut next = self.head.as_ref(); - let mut contains = false; - - while let Some(n) = next { - if entry as *const _ == &**n as *const _ { - debug_assert!(!contains); - contains = true; - } - - next = (*n.next_stack.get()).as_ref(); - } - - contains - }); - - // Unlink `entry` from the next node - let next = (*entry.next_stack.get()).take(); - - if let Some(next) = next.as_ref() { - (*next.prev_stack.get()) = *entry.prev_stack.get(); - } - - // Unlink `entry` from the prev node - - if let Some(prev) = (*entry.prev_stack.get()).as_ref() { - *prev.next_stack.get() = next; - } else { - // It is the head - self.head = next; - } - - // Unset the prev pointer - *entry.prev_stack.get() = ptr::null(); - } - } -} diff --git a/vendor/tokio/src/time/error.rs b/vendor/tokio/src/time/error.rs index 8674febe9..71344d434 100644 --- a/vendor/tokio/src/time/error.rs +++ b/vendor/tokio/src/time/error.rs @@ -40,8 +40,11 @@ impl From<Kind> for Error { } } -/// Error returned by `Timeout`. -#[derive(Debug, PartialEq)] +/// Errors returned by `Timeout`. +/// +/// This error is returned when a timeout expires before the function was able +/// to finish. +#[derive(Debug, PartialEq, Eq)] pub struct Elapsed(()); #[derive(Debug)] @@ -72,7 +75,7 @@ impl Error { matches!(self.0, Kind::AtCapacity) } - /// Create an error representing a misconfigured timer. + /// Creates an error representing a misconfigured timer. pub fn invalid() -> Error { Error(Invalid) } diff --git a/vendor/tokio/src/time/instant.rs b/vendor/tokio/src/time/instant.rs index f7cf12d4a..14cf6e567 100644 --- a/vendor/tokio/src/time/instant.rs +++ b/vendor/tokio/src/time/instant.rs @@ -67,13 +67,10 @@ impl Instant { self.std } - /// Returns the amount of time elapsed from another instant to this one. - /// - /// # Panics - /// - /// This function will panic if `earlier` is later than `self`. + /// Returns the amount of time elapsed from another instant to this one, or + /// zero duration if that instant is later than this one. pub fn duration_since(&self, earlier: Instant) -> Duration { - self.std.duration_since(earlier.std) + self.std.saturating_duration_since(earlier.std) } /// Returns the amount of time elapsed from another instant to this one, or @@ -118,13 +115,8 @@ impl Instant { self.std.saturating_duration_since(earlier.std) } - /// Returns the amount of time elapsed since this instant was created. - /// - /// # Panics - /// - /// This function may panic if the current time is earlier than this - /// instant, which is something that can happen if an `Instant` is - /// produced synthetically. + /// Returns the amount of time elapsed since this instant was created, + /// or zero duration if that this instant is in the future. /// /// # Examples /// @@ -140,7 +132,7 @@ impl Instant { /// } /// ``` pub fn elapsed(&self) -> Duration { - Instant::now() - *self + Instant::now().saturating_duration_since(*self) } /// Returns `Some(t)` where `t` is the time `self + duration` if `t` can be @@ -188,7 +180,7 @@ impl ops::Sub for Instant { type Output = Duration; fn sub(self, rhs: Instant) -> Duration { - self.std - rhs.std + self.std.saturating_duration_since(rhs.std) } } @@ -196,7 +188,7 @@ impl ops::Sub<Duration> for Instant { type Output = Instant; fn sub(self, rhs: Duration) -> Instant { - Instant::from_std(self.std - rhs) + Instant::from_std(std::time::Instant::sub(self.std, rhs)) } } diff --git a/vendor/tokio/src/time/interval.rs b/vendor/tokio/src/time/interval.rs index a63e47b6e..2e153fdbe 100644 --- a/vendor/tokio/src/time/interval.rs +++ b/vendor/tokio/src/time/interval.rs @@ -1,9 +1,11 @@ use crate::future::poll_fn; use crate::time::{sleep_until, Duration, Instant, Sleep}; +use crate::util::trace; +use std::future::Future; +use std::panic::Location; use std::pin::Pin; use std::task::{Context, Poll}; -use std::{convert::TryInto, future::Future}; /// Creates new [`Interval`] that yields with interval of `period`. The first /// tick completes immediately. The default [`MissedTickBehavior`] is @@ -68,10 +70,10 @@ use std::{convert::TryInto, future::Future}; /// /// [`sleep`]: crate::time::sleep() /// [`.tick().await`]: Interval::tick +#[track_caller] pub fn interval(period: Duration) -> Interval { assert!(period > Duration::new(0, 0), "`period` must be non-zero."); - - interval_at(Instant::now(), period) + internal_interval_at(Instant::now(), period, trace::caller_location()) } /// Creates new [`Interval`] that yields with interval of `period` with the @@ -103,13 +105,44 @@ pub fn interval(period: Duration) -> Interval { /// // approximately 70ms have elapsed. /// } /// ``` +#[track_caller] pub fn interval_at(start: Instant, period: Duration) -> Interval { assert!(period > Duration::new(0, 0), "`period` must be non-zero."); + internal_interval_at(start, period, trace::caller_location()) +} + +#[cfg_attr(not(all(tokio_unstable, feature = "tracing")), allow(unused_variables))] +fn internal_interval_at( + start: Instant, + period: Duration, + location: Option<&'static Location<'static>>, +) -> Interval { + #[cfg(all(tokio_unstable, feature = "tracing"))] + let resource_span = { + let location = location.expect("should have location if tracing"); + + tracing::trace_span!( + "runtime.resource", + concrete_type = "Interval", + kind = "timer", + loc.file = location.file(), + loc.line = location.line(), + loc.col = location.column(), + ) + }; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + let delay = resource_span.in_scope(|| Box::pin(sleep_until(start))); + + #[cfg(not(all(tokio_unstable, feature = "tracing")))] + let delay = Box::pin(sleep_until(start)); Interval { - delay: Box::pin(sleep_until(start)), + delay, period, missed_tick_behavior: Default::default(), + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span, } } @@ -124,7 +157,7 @@ pub fn interval_at(start: Instant, period: Duration) -> Interval { /// /// #[tokio::main] /// async fn main() { -/// // ticks every 2 seconds +/// // ticks every 2 milliseconds /// let mut interval = time::interval(Duration::from_millis(2)); /// for _ in 0..5 { /// interval.tick().await; @@ -147,7 +180,7 @@ pub fn interval_at(start: Instant, period: Duration) -> Interval { /// milliseconds. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum MissedTickBehavior { - /// Tick as fast as possible until caught up. + /// Ticks as fast as possible until caught up. /// /// When this strategy is used, [`Interval`] schedules ticks "normally" (the /// same as it would have if the ticks hadn't been delayed), which results @@ -174,6 +207,9 @@ pub enum MissedTickBehavior { /// # async fn main() { /// let mut interval = interval(Duration::from_millis(50)); /// + /// // First tick resolves immediately after creation + /// interval.tick().await; + /// /// task_that_takes_200_millis().await; /// // The `Interval` has missed a tick /// @@ -242,7 +278,7 @@ pub enum MissedTickBehavior { /// // 50ms after the call to `tick` up above. That is, in `tick`, when we /// // recognize that we missed a tick, we schedule the next tick to happen /// // 50ms (or whatever the `period` is) from right then, not from when - /// // were were *supposed* to tick + /// // were *supposed* to tick /// interval.tick().await; /// # } /// ``` @@ -252,7 +288,7 @@ pub enum MissedTickBehavior { /// [`tick`]: Interval::tick Delay, - /// Skip missed ticks and tick on the next multiple of `period` from + /// Skips missed ticks and tick on the next multiple of `period` from /// `start`. /// /// When this strategy is used, [`Interval`] schedules the next tick to fire @@ -342,7 +378,7 @@ impl Default for MissedTickBehavior { } } -/// Interval returned by [`interval`] and [`interval_at`] +/// Interval returned by [`interval`] and [`interval_at`]. /// /// This type allows you to wait on a sequence of instants with a certain /// duration between each instant. Unlike calling [`sleep`] in a loop, this lets @@ -351,7 +387,7 @@ impl Default for MissedTickBehavior { /// An `Interval` can be turned into a `Stream` with [`IntervalStream`]. /// /// [`IntervalStream`]: https://docs.rs/tokio-stream/latest/tokio_stream/wrappers/struct.IntervalStream.html -/// [`sleep`]: crate::time::sleep +/// [`sleep`]: crate::time::sleep() #[derive(Debug)] pub struct Interval { /// Future that completes the next time the `Interval` yields a value. @@ -362,11 +398,19 @@ pub struct Interval { /// The strategy `Interval` should use when a tick is missed. missed_tick_behavior: MissedTickBehavior, + + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, } impl Interval { /// Completes when the next instant in the interval has been reached. /// + /// # Cancel safety + /// + /// This method is cancellation safe. If `tick` is used as the branch in a `tokio::select!` and + /// another branch completes first, then no tick has been consumed. + /// /// # Examples /// /// ``` @@ -379,6 +423,7 @@ impl Interval { /// let mut interval = time::interval(Duration::from_millis(10)); /// /// interval.tick().await; + /// // approximately 0ms have elapsed. The first tick completes immediately. /// interval.tick().await; /// interval.tick().await; /// @@ -386,10 +431,23 @@ impl Interval { /// } /// ``` pub async fn tick(&mut self) -> Instant { - poll_fn(|cx| self.poll_tick(cx)).await + #[cfg(all(tokio_unstable, feature = "tracing"))] + let resource_span = self.resource_span.clone(); + #[cfg(all(tokio_unstable, feature = "tracing"))] + let instant = trace::async_op( + || poll_fn(|cx| self.poll_tick(cx)), + resource_span, + "Interval::tick", + "poll_tick", + false, + ); + #[cfg(not(all(tokio_unstable, feature = "tracing")))] + let instant = poll_fn(|cx| self.poll_tick(cx)); + + instant.await } - /// Poll for the next instant in the interval to be reached. + /// Polls for the next instant in the interval to be reached. /// /// This method can return the following values: /// @@ -424,12 +482,45 @@ impl Interval { timeout + self.period }; - self.delay.as_mut().reset(next); + // When we arrive here, the internal delay returned `Poll::Ready`. + // Reset the delay but do not register it. It should be registered with + // the next call to [`poll_tick`]. + self.delay.as_mut().reset_without_reregister(next); // Return the time when we were scheduled to tick Poll::Ready(timeout) } + /// Resets the interval to complete one period after the current time. + /// + /// This method ignores [`MissedTickBehavior`] strategy. + /// + /// # Examples + /// + /// ``` + /// use tokio::time; + /// + /// use std::time::Duration; + /// + /// #[tokio::main] + /// async fn main() { + /// let mut interval = time::interval(Duration::from_millis(100)); + /// + /// interval.tick().await; + /// + /// time::sleep(Duration::from_millis(50)).await; + /// interval.reset(); + /// + /// interval.tick().await; + /// interval.tick().await; + /// + /// // approximately 250ms have elapsed. + /// } + /// ``` + pub fn reset(&mut self) { + self.delay.as_mut().reset(Instant::now() + self.period); + } + /// Returns the [`MissedTickBehavior`] strategy currently being used. pub fn missed_tick_behavior(&self) -> MissedTickBehavior { self.missed_tick_behavior diff --git a/vendor/tokio/src/time/mod.rs b/vendor/tokio/src/time/mod.rs index 281990ef9..a1f27b839 100644 --- a/vendor/tokio/src/time/mod.rs +++ b/vendor/tokio/src/time/mod.rs @@ -82,17 +82,13 @@ //! ``` //! //! [`interval`]: crate::time::interval() +//! [`sleep`]: sleep() mod clock; pub(crate) use self::clock::Clock; #[cfg(feature = "test-util")] pub use clock::{advance, pause, resume}; -pub(crate) mod driver; - -#[doc(inline)] -pub use driver::sleep::{sleep, sleep_until, Sleep}; - pub mod error; mod instant; @@ -101,14 +97,13 @@ pub use self::instant::Instant; mod interval; pub use interval::{interval, interval_at, Interval, MissedTickBehavior}; +mod sleep; +pub use sleep::{sleep, sleep_until, Sleep}; + mod timeout; #[doc(inline)] pub use timeout::{timeout, timeout_at, Timeout}; -#[cfg(test)] -#[cfg(not(loom))] -mod tests; - // Re-export for convenience #[doc(no_inline)] pub use std::time::Duration; diff --git a/vendor/tokio/src/time/sleep.rs b/vendor/tokio/src/time/sleep.rs new file mode 100644 index 000000000..6c9b33793 --- /dev/null +++ b/vendor/tokio/src/time/sleep.rs @@ -0,0 +1,452 @@ +use crate::runtime::time::TimerEntry; +use crate::time::{error::Error, Duration, Instant}; +use crate::util::trace; + +use pin_project_lite::pin_project; +use std::future::Future; +use std::panic::Location; +use std::pin::Pin; +use std::task::{self, Poll}; + +/// Waits until `deadline` is reached. +/// +/// No work is performed while awaiting on the sleep future to complete. `Sleep` +/// operates at millisecond granularity and should not be used for tasks that +/// require high-resolution timers. +/// +/// To run something regularly on a schedule, see [`interval`]. +/// +/// # Cancellation +/// +/// Canceling a sleep instance is done by dropping the returned future. No additional +/// cleanup work is required. +/// +/// # Examples +/// +/// Wait 100ms and print "100 ms have elapsed". +/// +/// ``` +/// use tokio::time::{sleep_until, Instant, Duration}; +/// +/// #[tokio::main] +/// async fn main() { +/// sleep_until(Instant::now() + Duration::from_millis(100)).await; +/// println!("100 ms have elapsed"); +/// } +/// ``` +/// +/// See the documentation for the [`Sleep`] type for more examples. +/// +/// # Panics +/// +/// This function panics if there is no current timer set. +/// +/// It can be triggered when [`Builder::enable_time`] or +/// [`Builder::enable_all`] are not included in the builder. +/// +/// It can also panic whenever a timer is created outside of a +/// Tokio runtime. That is why `rt.block_on(sleep(...))` will panic, +/// since the function is executed outside of the runtime. +/// Whereas `rt.block_on(async {sleep(...).await})` doesn't panic. +/// And this is because wrapping the function on an async makes it lazy, +/// and so gets executed inside the runtime successfully without +/// panicking. +/// +/// [`Sleep`]: struct@crate::time::Sleep +/// [`interval`]: crate::time::interval() +/// [`Builder::enable_time`]: crate::runtime::Builder::enable_time +/// [`Builder::enable_all`]: crate::runtime::Builder::enable_all +// Alias for old name in 0.x +#[cfg_attr(docsrs, doc(alias = "delay_until"))] +#[track_caller] +pub fn sleep_until(deadline: Instant) -> Sleep { + return Sleep::new_timeout(deadline, trace::caller_location()); +} + +/// Waits until `duration` has elapsed. +/// +/// Equivalent to `sleep_until(Instant::now() + duration)`. An asynchronous +/// analog to `std::thread::sleep`. +/// +/// No work is performed while awaiting on the sleep future to complete. `Sleep` +/// operates at millisecond granularity and should not be used for tasks that +/// require high-resolution timers. The implementation is platform specific, +/// and some platforms (specifically Windows) will provide timers with a +/// larger resolution than 1 ms. +/// +/// To run something regularly on a schedule, see [`interval`]. +/// +/// The maximum duration for a sleep is 68719476734 milliseconds (approximately 2.2 years). +/// +/// # Cancellation +/// +/// Canceling a sleep instance is done by dropping the returned future. No additional +/// cleanup work is required. +/// +/// # Examples +/// +/// Wait 100ms and print "100 ms have elapsed". +/// +/// ``` +/// use tokio::time::{sleep, Duration}; +/// +/// #[tokio::main] +/// async fn main() { +/// sleep(Duration::from_millis(100)).await; +/// println!("100 ms have elapsed"); +/// } +/// ``` +/// +/// See the documentation for the [`Sleep`] type for more examples. +/// +/// # Panics +/// +/// This function panics if there is no current timer set. +/// +/// It can be triggered when [`Builder::enable_time`] or +/// [`Builder::enable_all`] are not included in the builder. +/// +/// It can also panic whenever a timer is created outside of a +/// Tokio runtime. That is why `rt.block_on(sleep(...))` will panic, +/// since the function is executed outside of the runtime. +/// Whereas `rt.block_on(async {sleep(...).await})` doesn't panic. +/// And this is because wrapping the function on an async makes it lazy, +/// and so gets executed inside the runtime successfully without +/// panicking. +/// +/// [`Sleep`]: struct@crate::time::Sleep +/// [`interval`]: crate::time::interval() +/// [`Builder::enable_time`]: crate::runtime::Builder::enable_time +/// [`Builder::enable_all`]: crate::runtime::Builder::enable_all +// Alias for old name in 0.x +#[cfg_attr(docsrs, doc(alias = "delay_for"))] +#[cfg_attr(docsrs, doc(alias = "wait"))] +#[track_caller] +pub fn sleep(duration: Duration) -> Sleep { + let location = trace::caller_location(); + + match Instant::now().checked_add(duration) { + Some(deadline) => Sleep::new_timeout(deadline, location), + None => Sleep::new_timeout(Instant::far_future(), location), + } +} + +pin_project! { + /// Future returned by [`sleep`](sleep) and [`sleep_until`](sleep_until). + /// + /// This type does not implement the `Unpin` trait, which means that if you + /// use it with [`select!`] or by calling `poll`, you have to pin it first. + /// If you use it with `.await`, this does not apply. + /// + /// # Examples + /// + /// Wait 100ms and print "100 ms have elapsed". + /// + /// ``` + /// use tokio::time::{sleep, Duration}; + /// + /// #[tokio::main] + /// async fn main() { + /// sleep(Duration::from_millis(100)).await; + /// println!("100 ms have elapsed"); + /// } + /// ``` + /// + /// Use with [`select!`]. Pinning the `Sleep` with [`tokio::pin!`] is + /// necessary when the same `Sleep` is selected on multiple times. + /// ```no_run + /// use tokio::time::{self, Duration, Instant}; + /// + /// #[tokio::main] + /// async fn main() { + /// let sleep = time::sleep(Duration::from_millis(10)); + /// tokio::pin!(sleep); + /// + /// loop { + /// tokio::select! { + /// () = &mut sleep => { + /// println!("timer elapsed"); + /// sleep.as_mut().reset(Instant::now() + Duration::from_millis(50)); + /// }, + /// } + /// } + /// } + /// ``` + /// Use in a struct with boxing. By pinning the `Sleep` with a `Box`, the + /// `HasSleep` struct implements `Unpin`, even though `Sleep` does not. + /// ``` + /// use std::future::Future; + /// use std::pin::Pin; + /// use std::task::{Context, Poll}; + /// use tokio::time::Sleep; + /// + /// struct HasSleep { + /// sleep: Pin<Box<Sleep>>, + /// } + /// + /// impl Future for HasSleep { + /// type Output = (); + /// + /// fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { + /// self.sleep.as_mut().poll(cx) + /// } + /// } + /// ``` + /// Use in a struct with pin projection. This method avoids the `Box`, but + /// the `HasSleep` struct will not be `Unpin` as a consequence. + /// ``` + /// use std::future::Future; + /// use std::pin::Pin; + /// use std::task::{Context, Poll}; + /// use tokio::time::Sleep; + /// use pin_project_lite::pin_project; + /// + /// pin_project! { + /// struct HasSleep { + /// #[pin] + /// sleep: Sleep, + /// } + /// } + /// + /// impl Future for HasSleep { + /// type Output = (); + /// + /// fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { + /// self.project().sleep.poll(cx) + /// } + /// } + /// ``` + /// + /// [`select!`]: ../macro.select.html + /// [`tokio::pin!`]: ../macro.pin.html + // Alias for old name in 0.2 + #[cfg_attr(docsrs, doc(alias = "Delay"))] + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct Sleep { + inner: Inner, + + // The link between the `Sleep` instance and the timer that drives it. + #[pin] + entry: TimerEntry, + } +} + +cfg_trace! { + #[derive(Debug)] + struct Inner { + ctx: trace::AsyncOpTracingCtx, + } +} + +cfg_not_trace! { + #[derive(Debug)] + struct Inner { + } +} + +impl Sleep { + #[cfg_attr(not(all(tokio_unstable, feature = "tracing")), allow(unused_variables))] + #[track_caller] + pub(crate) fn new_timeout( + deadline: Instant, + location: Option<&'static Location<'static>>, + ) -> Sleep { + use crate::runtime::scheduler; + + let handle = scheduler::Handle::current(); + let entry = TimerEntry::new(&handle, deadline); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + let inner = { + let clock = handle.driver().clock(); + let handle = &handle.driver().time(); + let time_source = handle.time_source(); + let deadline_tick = time_source.deadline_to_tick(deadline); + let duration = deadline_tick.saturating_sub(time_source.now(clock)); + + let location = location.expect("should have location if tracing"); + let resource_span = tracing::trace_span!( + "runtime.resource", + concrete_type = "Sleep", + kind = "timer", + loc.file = location.file(), + loc.line = location.line(), + loc.col = location.column(), + ); + + let async_op_span = resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + duration = duration, + duration.unit = "ms", + duration.op = "override", + ); + + tracing::trace_span!("runtime.resource.async_op", source = "Sleep::new_timeout") + }); + + let async_op_poll_span = + async_op_span.in_scope(|| tracing::trace_span!("runtime.resource.async_op.poll")); + + let ctx = trace::AsyncOpTracingCtx { + async_op_span, + async_op_poll_span, + resource_span, + }; + + Inner { ctx } + }; + + #[cfg(not(all(tokio_unstable, feature = "tracing")))] + let inner = Inner {}; + + Sleep { inner, entry } + } + + pub(crate) fn far_future(location: Option<&'static Location<'static>>) -> Sleep { + Self::new_timeout(Instant::far_future(), location) + } + + /// Returns the instant at which the future will complete. + pub fn deadline(&self) -> Instant { + self.entry.deadline() + } + + /// Returns `true` if `Sleep` has elapsed. + /// + /// A `Sleep` instance is elapsed when the requested duration has elapsed. + pub fn is_elapsed(&self) -> bool { + self.entry.is_elapsed() + } + + /// Resets the `Sleep` instance to a new deadline. + /// + /// Calling this function allows changing the instant at which the `Sleep` + /// future completes without having to create new associated state. + /// + /// This function can be called both before and after the future has + /// completed. + /// + /// To call this method, you will usually combine the call with + /// [`Pin::as_mut`], which lets you call the method without consuming the + /// `Sleep` itself. + /// + /// # Example + /// + /// ``` + /// use tokio::time::{Duration, Instant}; + /// + /// # #[tokio::main(flavor = "current_thread")] + /// # async fn main() { + /// let sleep = tokio::time::sleep(Duration::from_millis(10)); + /// tokio::pin!(sleep); + /// + /// sleep.as_mut().reset(Instant::now() + Duration::from_millis(20)); + /// # } + /// ``` + /// + /// See also the top-level examples. + /// + /// [`Pin::as_mut`]: fn@std::pin::Pin::as_mut + pub fn reset(self: Pin<&mut Self>, deadline: Instant) { + self.reset_inner(deadline) + } + + /// Resets the `Sleep` instance to a new deadline without reregistering it + /// to be woken up. + /// + /// Calling this function allows changing the instant at which the `Sleep` + /// future completes without having to create new associated state and + /// without having it registered. This is required in e.g. the + /// [crate::time::Interval] where we want to reset the internal [Sleep] + /// without having it wake up the last task that polled it. + pub(crate) fn reset_without_reregister(self: Pin<&mut Self>, deadline: Instant) { + let mut me = self.project(); + me.entry.as_mut().reset(deadline, false); + } + + fn reset_inner(self: Pin<&mut Self>, deadline: Instant) { + let mut me = self.project(); + me.entry.as_mut().reset(deadline, true); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + { + let _resource_enter = me.inner.ctx.resource_span.enter(); + me.inner.ctx.async_op_span = + tracing::trace_span!("runtime.resource.async_op", source = "Sleep::reset"); + let _async_op_enter = me.inner.ctx.async_op_span.enter(); + + me.inner.ctx.async_op_poll_span = + tracing::trace_span!("runtime.resource.async_op.poll"); + + let duration = { + let clock = me.entry.clock(); + let time_source = me.entry.driver().time_source(); + let now = time_source.now(clock); + let deadline_tick = time_source.deadline_to_tick(deadline); + deadline_tick.saturating_sub(now) + }; + + tracing::trace!( + target: "runtime::resource::state_update", + duration = duration, + duration.unit = "ms", + duration.op = "override", + ); + } + } + + fn poll_elapsed(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<(), Error>> { + let me = self.project(); + + ready!(crate::trace::trace_leaf(cx)); + + // Keep track of task budget + #[cfg(all(tokio_unstable, feature = "tracing"))] + let coop = ready!(trace_poll_op!( + "poll_elapsed", + crate::runtime::coop::poll_proceed(cx), + )); + + #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] + let coop = ready!(crate::runtime::coop::poll_proceed(cx)); + + let result = me.entry.poll_elapsed(cx).map(move |r| { + coop.made_progress(); + r + }); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + return trace_poll_op!("poll_elapsed", result); + + #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] + return result; + } +} + +impl Future for Sleep { + type Output = (); + + // `poll_elapsed` can return an error in two cases: + // + // - AtCapacity: this is a pathological case where far too many + // sleep instances have been scheduled. + // - Shutdown: No timer has been setup, which is a mis-use error. + // + // Both cases are extremely rare, and pretty accurately fit into + // "logic errors", so we just panic in this case. A user couldn't + // really do much better if we passed the error onwards. + fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> { + #[cfg(all(tokio_unstable, feature = "tracing"))] + let _res_span = self.inner.ctx.resource_span.clone().entered(); + #[cfg(all(tokio_unstable, feature = "tracing"))] + let _ao_span = self.inner.ctx.async_op_span.clone().entered(); + #[cfg(all(tokio_unstable, feature = "tracing"))] + let _ao_poll_span = self.inner.ctx.async_op_poll_span.clone().entered(); + match ready!(self.as_mut().poll_elapsed(cx)) { + Ok(()) => Poll::Ready(()), + Err(e) => panic!("timer error: {}", e), + } + } +} diff --git a/vendor/tokio/src/time/tests/mod.rs b/vendor/tokio/src/time/tests/mod.rs deleted file mode 100644 index 35e1060ac..000000000 --- a/vendor/tokio/src/time/tests/mod.rs +++ /dev/null @@ -1,22 +0,0 @@ -mod test_sleep; - -use crate::time::{self, Instant}; -use std::time::Duration; - -fn assert_send<T: Send>() {} -fn assert_sync<T: Sync>() {} - -#[test] -fn registration_is_send_and_sync() { - use crate::time::Sleep; - - assert_send::<Sleep>(); - assert_sync::<Sleep>(); -} - -#[test] -#[should_panic] -fn sleep_is_eager() { - let when = Instant::now() + Duration::from_millis(100); - let _ = time::sleep_until(when); -} diff --git a/vendor/tokio/src/time/tests/test_sleep.rs b/vendor/tokio/src/time/tests/test_sleep.rs deleted file mode 100644 index 77ca07e31..000000000 --- a/vendor/tokio/src/time/tests/test_sleep.rs +++ /dev/null @@ -1,443 +0,0 @@ -//use crate::time::driver::{Driver, Entry, Handle}; - -/* -macro_rules! poll { - ($e:expr) => { - $e.enter(|cx, e| e.poll_elapsed(cx)) - }; -} - -#[test] -fn frozen_utility_returns_correct_advanced_duration() { - let clock = Clock::new(); - clock.pause(); - let start = clock.now(); - - clock.advance(ms(10)); - assert_eq!(clock.now() - start, ms(10)); -} - -#[test] -fn immediate_sleep() { - let (mut driver, clock, handle) = setup(); - let start = clock.now(); - - let when = clock.now(); - let mut e = task::spawn(sleep_until(&handle, when)); - - assert_ready_ok!(poll!(e)); - - assert_ok!(driver.park_timeout(Duration::from_millis(1000))); - - // The time has not advanced. The `turn` completed immediately. - assert_eq!(clock.now() - start, ms(1000)); -} - -#[test] -fn delayed_sleep_level_0() { - let (mut driver, clock, handle) = setup(); - let start = clock.now(); - - for &i in &[1, 10, 60] { - // Create a `Sleep` that elapses in the future - let mut e = task::spawn(sleep_until(&handle, start + ms(i))); - - // The sleep instance has not elapsed. - assert_pending!(poll!(e)); - - assert_ok!(driver.park()); - assert_eq!(clock.now() - start, ms(i)); - - assert_ready_ok!(poll!(e)); - } -} - -#[test] -fn sub_ms_delayed_sleep() { - let (mut driver, clock, handle) = setup(); - - for _ in 0..5 { - let deadline = clock.now() + ms(1) + Duration::new(0, 1); - - let mut e = task::spawn(sleep_until(&handle, deadline)); - - assert_pending!(poll!(e)); - - assert_ok!(driver.park()); - assert_ready_ok!(poll!(e)); - - assert!(clock.now() >= deadline); - - clock.advance(Duration::new(0, 1)); - } -} - -#[test] -fn delayed_sleep_wrapping_level_0() { - let (mut driver, clock, handle) = setup(); - let start = clock.now(); - - assert_ok!(driver.park_timeout(ms(5))); - assert_eq!(clock.now() - start, ms(5)); - - let mut e = task::spawn(sleep_until(&handle, clock.now() + ms(60))); - - assert_pending!(poll!(e)); - - assert_ok!(driver.park()); - assert_eq!(clock.now() - start, ms(64)); - assert_pending!(poll!(e)); - - assert_ok!(driver.park()); - assert_eq!(clock.now() - start, ms(65)); - - assert_ready_ok!(poll!(e)); -} - -#[test] -fn timer_wrapping_with_higher_levels() { - let (mut driver, clock, handle) = setup(); - let start = clock.now(); - - // Set sleep to hit level 1 - let mut e1 = task::spawn(sleep_until(&handle, clock.now() + ms(64))); - assert_pending!(poll!(e1)); - - // Turn a bit - assert_ok!(driver.park_timeout(ms(5))); - - // Set timeout such that it will hit level 0, but wrap - let mut e2 = task::spawn(sleep_until(&handle, clock.now() + ms(60))); - assert_pending!(poll!(e2)); - - // This should result in s1 firing - assert_ok!(driver.park()); - assert_eq!(clock.now() - start, ms(64)); - - assert_ready_ok!(poll!(e1)); - assert_pending!(poll!(e2)); - - assert_ok!(driver.park()); - assert_eq!(clock.now() - start, ms(65)); - - assert_ready_ok!(poll!(e1)); -} - -#[test] -fn sleep_with_deadline_in_past() { - let (mut driver, clock, handle) = setup(); - let start = clock.now(); - - // Create `Sleep` that elapsed immediately. - let mut e = task::spawn(sleep_until(&handle, clock.now() - ms(100))); - - // Even though the `Sleep` expires in the past, it is not ready yet - // because the timer must observe it. - assert_ready_ok!(poll!(e)); - - // Turn the timer, it runs for the elapsed time - assert_ok!(driver.park_timeout(ms(1000))); - - // The time has not advanced. The `turn` completed immediately. - assert_eq!(clock.now() - start, ms(1000)); -} - -#[test] -fn delayed_sleep_level_1() { - let (mut driver, clock, handle) = setup(); - let start = clock.now(); - - // Create a `Sleep` that elapses in the future - let mut e = task::spawn(sleep_until(&handle, clock.now() + ms(234))); - - // The sleep has not elapsed. - assert_pending!(poll!(e)); - - // Turn the timer, this will wake up to cascade the timer down. - assert_ok!(driver.park_timeout(ms(1000))); - assert_eq!(clock.now() - start, ms(192)); - - // The sleep has not elapsed. - assert_pending!(poll!(e)); - - // Turn the timer again - assert_ok!(driver.park_timeout(ms(1000))); - assert_eq!(clock.now() - start, ms(234)); - - // The sleep has elapsed. - assert_ready_ok!(poll!(e)); - - let (mut driver, clock, handle) = setup(); - let start = clock.now(); - - // Create a `Sleep` that elapses in the future - let mut e = task::spawn(sleep_until(&handle, clock.now() + ms(234))); - - // The sleep has not elapsed. - assert_pending!(poll!(e)); - - // Turn the timer with a smaller timeout than the cascade. - assert_ok!(driver.park_timeout(ms(100))); - assert_eq!(clock.now() - start, ms(100)); - - assert_pending!(poll!(e)); - - // Turn the timer, this will wake up to cascade the timer down. - assert_ok!(driver.park_timeout(ms(1000))); - assert_eq!(clock.now() - start, ms(192)); - - // The sleep has not elapsed. - assert_pending!(poll!(e)); - - // Turn the timer again - assert_ok!(driver.park_timeout(ms(1000))); - assert_eq!(clock.now() - start, ms(234)); - - // The sleep has elapsed. - assert_ready_ok!(poll!(e)); -} - -#[test] -fn concurrently_set_two_timers_second_one_shorter() { - let (mut driver, clock, handle) = setup(); - let start = clock.now(); - - let mut e1 = task::spawn(sleep_until(&handle, clock.now() + ms(500))); - let mut e2 = task::spawn(sleep_until(&handle, clock.now() + ms(200))); - - // The sleep has not elapsed - assert_pending!(poll!(e1)); - assert_pending!(poll!(e2)); - - // Sleep until a cascade - assert_ok!(driver.park()); - assert_eq!(clock.now() - start, ms(192)); - - // Sleep until the second timer. - assert_ok!(driver.park()); - assert_eq!(clock.now() - start, ms(200)); - - // The shorter sleep fires - assert_ready_ok!(poll!(e2)); - assert_pending!(poll!(e1)); - - assert_ok!(driver.park()); - assert_eq!(clock.now() - start, ms(448)); - - assert_pending!(poll!(e1)); - - // Turn again, this time the time will advance to the second sleep - assert_ok!(driver.park()); - assert_eq!(clock.now() - start, ms(500)); - - assert_ready_ok!(poll!(e1)); -} - -#[test] -fn short_sleep() { - let (mut driver, clock, handle) = setup(); - let start = clock.now(); - - // Create a `Sleep` that elapses in the future - let mut e = task::spawn(sleep_until(&handle, clock.now() + ms(1))); - - // The sleep has not elapsed. - assert_pending!(poll!(e)); - - // Turn the timer, but not enough time will go by. - assert_ok!(driver.park()); - - // The sleep has elapsed. - assert_ready_ok!(poll!(e)); - - // The time has advanced to the point of the sleep elapsing. - assert_eq!(clock.now() - start, ms(1)); -} - -#[test] -fn sorta_long_sleep_until() { - const MIN_5: u64 = 5 * 60 * 1000; - - let (mut driver, clock, handle) = setup(); - let start = clock.now(); - - // Create a `Sleep` that elapses in the future - let mut e = task::spawn(sleep_until(&handle, clock.now() + ms(MIN_5))); - - // The sleep has not elapsed. - assert_pending!(poll!(e)); - - let cascades = &[262_144, 262_144 + 9 * 4096, 262_144 + 9 * 4096 + 15 * 64]; - - for &elapsed in cascades { - assert_ok!(driver.park()); - assert_eq!(clock.now() - start, ms(elapsed)); - - assert_pending!(poll!(e)); - } - - assert_ok!(driver.park()); - assert_eq!(clock.now() - start, ms(MIN_5)); - - // The sleep has elapsed. - assert_ready_ok!(poll!(e)); -} - -#[test] -fn very_long_sleep() { - const MO_5: u64 = 5 * 30 * 24 * 60 * 60 * 1000; - - let (mut driver, clock, handle) = setup(); - let start = clock.now(); - - // Create a `Sleep` that elapses in the future - let mut e = task::spawn(sleep_until(&handle, clock.now() + ms(MO_5))); - - // The sleep has not elapsed. - assert_pending!(poll!(e)); - - let cascades = &[ - 12_884_901_888, - 12_952_010_752, - 12_959_875_072, - 12_959_997_952, - ]; - - for &elapsed in cascades { - assert_ok!(driver.park()); - assert_eq!(clock.now() - start, ms(elapsed)); - - assert_pending!(poll!(e)); - } - - // Turn the timer, but not enough time will go by. - assert_ok!(driver.park()); - - // The time has advanced to the point of the sleep elapsing. - assert_eq!(clock.now() - start, ms(MO_5)); - - // The sleep has elapsed. - assert_ready_ok!(poll!(e)); -} - -#[test] -fn unpark_is_delayed() { - // A special park that will take much longer than the requested duration - struct MockPark(Clock); - - struct MockUnpark; - - impl Park for MockPark { - type Unpark = MockUnpark; - type Error = (); - - fn unpark(&self) -> Self::Unpark { - MockUnpark - } - - fn park(&mut self) -> Result<(), Self::Error> { - panic!("parking forever"); - } - - fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> { - assert_eq!(duration, ms(0)); - self.0.advance(ms(436)); - Ok(()) - } - - fn shutdown(&mut self) {} - } - - impl Unpark for MockUnpark { - fn unpark(&self) {} - } - - let clock = Clock::new(); - clock.pause(); - let start = clock.now(); - let mut driver = Driver::new(MockPark(clock.clone()), clock.clone()); - let handle = driver.handle(); - - let mut e1 = task::spawn(sleep_until(&handle, clock.now() + ms(100))); - let mut e2 = task::spawn(sleep_until(&handle, clock.now() + ms(101))); - let mut e3 = task::spawn(sleep_until(&handle, clock.now() + ms(200))); - - assert_pending!(poll!(e1)); - assert_pending!(poll!(e2)); - assert_pending!(poll!(e3)); - - assert_ok!(driver.park()); - - assert_eq!(clock.now() - start, ms(500)); - - assert_ready_ok!(poll!(e1)); - assert_ready_ok!(poll!(e2)); - assert_ready_ok!(poll!(e3)); -} - -#[test] -fn set_timeout_at_deadline_greater_than_max_timer() { - const YR_1: u64 = 365 * 24 * 60 * 60 * 1000; - const YR_5: u64 = 5 * YR_1; - - let (mut driver, clock, handle) = setup(); - let start = clock.now(); - - for _ in 0..5 { - assert_ok!(driver.park_timeout(ms(YR_1))); - } - - let mut e = task::spawn(sleep_until(&handle, clock.now() + ms(1))); - assert_pending!(poll!(e)); - - assert_ok!(driver.park_timeout(ms(1000))); - assert_eq!(clock.now() - start, ms(YR_5) + ms(1)); - - assert_ready_ok!(poll!(e)); -} - -fn setup() -> (Driver<MockPark>, Clock, Handle) { - let clock = Clock::new(); - clock.pause(); - let driver = Driver::new(MockPark(clock.clone()), clock.clone()); - let handle = driver.handle(); - - (driver, clock, handle) -} - -fn sleep_until(handle: &Handle, when: Instant) -> Arc<Entry> { - Entry::new(&handle, when, ms(0)) -} - -struct MockPark(Clock); - -struct MockUnpark; - -impl Park for MockPark { - type Unpark = MockUnpark; - type Error = (); - - fn unpark(&self) -> Self::Unpark { - MockUnpark - } - - fn park(&mut self) -> Result<(), Self::Error> { - panic!("parking forever"); - } - - fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> { - self.0.advance(duration); - Ok(()) - } - - fn shutdown(&mut self) {} -} - -impl Unpark for MockUnpark { - fn unpark(&self) {} -} - -fn ms(n: u64) -> Duration { - Duration::from_millis(n) -} -*/ diff --git a/vendor/tokio/src/time/timeout.rs b/vendor/tokio/src/time/timeout.rs index 61964ad24..52ab9891c 100644 --- a/vendor/tokio/src/time/timeout.rs +++ b/vendor/tokio/src/time/timeout.rs @@ -4,20 +4,39 @@ //! //! [`Timeout`]: struct@Timeout -use crate::time::{error::Elapsed, sleep_until, Duration, Instant, Sleep}; +use crate::{ + runtime::coop, + time::{error::Elapsed, sleep_until, Duration, Instant, Sleep}, + util::trace, +}; use pin_project_lite::pin_project; use std::future::Future; use std::pin::Pin; use std::task::{self, Poll}; -/// Require a `Future` to complete before the specified duration has elapsed. +/// Requires a `Future` to complete before the specified duration has elapsed. /// /// If the future completes before the duration has elapsed, then the completed /// value is returned. Otherwise, an error is returned and the future is /// canceled. /// -/// # Cancelation +/// Note that the timeout is checked before polling the future, so if the future +/// does not yield during execution then it is possible for the future to complete +/// and exceed the timeout _without_ returning an error. +/// +/// This function returns a future whose return type is [`Result`]`<T,`[`Elapsed`]`>`, where `T` is the +/// return type of the provided future. +/// +/// If the provided future completes immediately, then the future returned from +/// this function is guaranteed to complete immediately with an [`Ok`] variant +/// no matter the provided duration. +/// +/// [`Ok`]: std::result::Result::Ok +/// [`Result`]: std::result::Result +/// [`Elapsed`]: crate::time::error::Elapsed +/// +/// # Cancellation /// /// Cancelling a timeout is done by dropping the future. No additional cleanup /// or other work is required. @@ -45,24 +64,56 @@ use std::task::{self, Poll}; /// } /// # } /// ``` -pub fn timeout<T>(duration: Duration, future: T) -> Timeout<T> +/// +/// # Panics +/// +/// This function panics if there is no current timer set. +/// +/// It can be triggered when [`Builder::enable_time`] or +/// [`Builder::enable_all`] are not included in the builder. +/// +/// It can also panic whenever a timer is created outside of a +/// Tokio runtime. That is why `rt.block_on(sleep(...))` will panic, +/// since the function is executed outside of the runtime. +/// Whereas `rt.block_on(async {sleep(...).await})` doesn't panic. +/// And this is because wrapping the function on an async makes it lazy, +/// and so gets executed inside the runtime successfully without +/// panicking. +/// +/// [`Builder::enable_time`]: crate::runtime::Builder::enable_time +/// [`Builder::enable_all`]: crate::runtime::Builder::enable_all +#[track_caller] +pub fn timeout<F>(duration: Duration, future: F) -> Timeout<F> where - T: Future, + F: Future, { + let location = trace::caller_location(); + let deadline = Instant::now().checked_add(duration); let delay = match deadline { - Some(deadline) => Sleep::new_timeout(deadline), - None => Sleep::far_future(), + Some(deadline) => Sleep::new_timeout(deadline, location), + None => Sleep::far_future(location), }; Timeout::new_with_delay(future, delay) } -/// Require a `Future` to complete before the specified instant in time. +/// Requires a `Future` to complete before the specified instant in time. /// /// If the future completes before the instant is reached, then the completed /// value is returned. Otherwise, an error is returned. /// -/// # Cancelation +/// This function returns a future whose return type is [`Result`]`<T,`[`Elapsed`]`>`, where `T` is the +/// return type of the provided future. +/// +/// If the provided future completes immediately, then the future returned from +/// this function is guaranteed to complete immediately with an [`Ok`] variant +/// no matter the provided deadline. +/// +/// [`Ok`]: std::result::Result::Ok +/// [`Result`]: std::result::Result +/// [`Elapsed`]: crate::time::error::Elapsed +/// +/// # Cancellation /// /// Cancelling a timeout is done by dropping the future. No additional cleanup /// or other work is required. @@ -91,9 +142,9 @@ where /// } /// # } /// ``` -pub fn timeout_at<T>(deadline: Instant, future: T) -> Timeout<T> +pub fn timeout_at<F>(deadline: Instant, future: F) -> Timeout<F> where - T: Future, + F: Future, { let delay = sleep_until(deadline); @@ -145,15 +196,33 @@ where fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> { let me = self.project(); + let had_budget_before = coop::has_budget_remaining(); + // First, try polling the future if let Poll::Ready(v) = me.value.poll(cx) { return Poll::Ready(Ok(v)); } - // Now check the timer - match me.delay.poll(cx) { - Poll::Ready(()) => Poll::Ready(Err(Elapsed::new())), - Poll::Pending => Poll::Pending, + let has_budget_now = coop::has_budget_remaining(); + + let delay = me.delay; + + let poll_delay = || -> Poll<Self::Output> { + match delay.poll(cx) { + Poll::Ready(()) => Poll::Ready(Err(Elapsed::new())), + Poll::Pending => Poll::Pending, + } + }; + + if let (true, false) = (had_budget_before, has_budget_now) { + // if it is the underlying future that exhausted the budget, we poll + // the `delay` with an unconstrained one. This prevents pathological + // cases where the underlying future always exhausts the budget and + // we never get a chance to evaluate whether the timeout was hit or + // not. + coop::with_unconstrained(poll_delay) + } else { + poll_delay() } } } diff --git a/vendor/tokio/src/runtime/thread_pool/atomic_cell.rs b/vendor/tokio/src/util/atomic_cell.rs index 98847e6ff..07e37303a 100644 --- a/vendor/tokio/src/runtime/thread_pool/atomic_cell.rs +++ b/vendor/tokio/src/util/atomic_cell.rs @@ -3,7 +3,7 @@ use crate::loom::sync::atomic::AtomicPtr; use std::ptr; use std::sync::atomic::Ordering::AcqRel; -pub(super) struct AtomicCell<T> { +pub(crate) struct AtomicCell<T> { data: AtomicPtr<T>, } @@ -11,22 +11,22 @@ unsafe impl<T: Send> Send for AtomicCell<T> {} unsafe impl<T: Send> Sync for AtomicCell<T> {} impl<T> AtomicCell<T> { - pub(super) fn new(data: Option<Box<T>>) -> AtomicCell<T> { + pub(crate) fn new(data: Option<Box<T>>) -> AtomicCell<T> { AtomicCell { data: AtomicPtr::new(to_raw(data)), } } - pub(super) fn swap(&self, val: Option<Box<T>>) -> Option<Box<T>> { + pub(crate) fn swap(&self, val: Option<Box<T>>) -> Option<Box<T>> { let old = self.data.swap(to_raw(val), AcqRel); from_raw(old) } - pub(super) fn set(&self, val: Box<T>) { + pub(crate) fn set(&self, val: Box<T>) { let _ = self.swap(Some(val)); } - pub(super) fn take(&self) -> Option<Box<T>> { + pub(crate) fn take(&self) -> Option<Box<T>> { self.swap(None) } } diff --git a/vendor/tokio/src/util/bit.rs b/vendor/tokio/src/util/bit.rs index 392a0e8b0..a43c2c2d3 100644 --- a/vendor/tokio/src/util/bit.rs +++ b/vendor/tokio/src/util/bit.rs @@ -27,7 +27,7 @@ impl Pack { pointer_width() - (self.mask >> self.shift).leading_zeros() } - /// Max representable value + /// Max representable value. pub(crate) const fn max_value(&self) -> usize { (1 << self.width()) - 1 } @@ -60,7 +60,7 @@ impl fmt::Debug for Pack { } } -/// Returns the width of a pointer in bits +/// Returns the width of a pointer in bits. pub(crate) const fn pointer_width() -> u32 { std::mem::size_of::<usize>() as u32 * 8 } @@ -71,7 +71,7 @@ pub(crate) const fn mask_for(n: u32) -> usize { shift | (shift - 1) } -/// Unpack a value using a mask & shift +/// Unpacks a value using a mask & shift. pub(crate) const fn unpack(src: usize, mask: usize, shift: u32) -> usize { (src & mask) >> shift } diff --git a/vendor/tokio/src/util/error.rs b/vendor/tokio/src/util/error.rs index 0e52364a7..ebb27f638 100644 --- a/vendor/tokio/src/util/error.rs +++ b/vendor/tokio/src/util/error.rs @@ -1,9 +1,16 @@ +// Some combinations of features may not use these constants. +#![cfg_attr(not(feature = "full"), allow(dead_code))] + /// Error string explaining that the Tokio context hasn't been instantiated. pub(crate) const CONTEXT_MISSING_ERROR: &str = "there is no reactor running, must be called from the context of a Tokio 1.x runtime"; -// some combinations of features might not use this -#[allow(dead_code)] /// Error string explaining that the Tokio context is shutting down and cannot drive timers. pub(crate) const RUNTIME_SHUTTING_DOWN_ERROR: &str = "A Tokio 1.x context was found, but it is being shutdown."; + +/// Error string explaining that the Tokio context is not available because the +/// thread-local storing it has been destroyed. This usually only happens during +/// destructors of other thread-locals. +pub(crate) const THREAD_LOCAL_DESTROYED_ERROR: &str = + "The Tokio context thread-local variable has been destroyed."; diff --git a/vendor/tokio/src/util/idle_notified_set.rs b/vendor/tokio/src/util/idle_notified_set.rs new file mode 100644 index 000000000..19b81e28b --- /dev/null +++ b/vendor/tokio/src/util/idle_notified_set.rs @@ -0,0 +1,481 @@ +//! This module defines an `IdleNotifiedSet`, which is a collection of elements. +//! Each element is intended to correspond to a task, and the collection will +//! keep track of which tasks have had their waker notified, and which have not. +//! +//! Each entry in the set holds some user-specified value. The value's type is +//! specified using the `T` parameter. It will usually be a `JoinHandle` or +//! similar. + +use std::marker::PhantomPinned; +use std::mem::ManuallyDrop; +use std::ptr::NonNull; +use std::task::{Context, Waker}; + +use crate::loom::cell::UnsafeCell; +use crate::loom::sync::{Arc, Mutex}; +use crate::util::linked_list::{self, Link}; +use crate::util::{waker_ref, Wake}; + +type LinkedList<T> = + linked_list::LinkedList<ListEntry<T>, <ListEntry<T> as linked_list::Link>::Target>; + +/// This is the main handle to the collection. +pub(crate) struct IdleNotifiedSet<T> { + lists: Arc<Lists<T>>, + length: usize, +} + +/// A handle to an entry that is guaranteed to be stored in the idle or notified +/// list of its `IdleNotifiedSet`. This value borrows the `IdleNotifiedSet` +/// mutably to prevent the entry from being moved to the `Neither` list, which +/// only the `IdleNotifiedSet` may do. +/// +/// The main consequence of being stored in one of the lists is that the `value` +/// field has not yet been consumed. +/// +/// Note: This entry can be moved from the idle to the notified list while this +/// object exists by waking its waker. +pub(crate) struct EntryInOneOfTheLists<'a, T> { + entry: Arc<ListEntry<T>>, + set: &'a mut IdleNotifiedSet<T>, +} + +type Lists<T> = Mutex<ListsInner<T>>; + +/// The linked lists hold strong references to the ListEntry items, and the +/// ListEntry items also hold a strong reference back to the Lists object, but +/// the destructor of the `IdleNotifiedSet` will clear the two lists, so once +/// that object is destroyed, no ref-cycles will remain. +struct ListsInner<T> { + notified: LinkedList<T>, + idle: LinkedList<T>, + /// Whenever an element in the `notified` list is woken, this waker will be + /// notified and consumed, if it exists. + waker: Option<Waker>, +} + +/// Which of the two lists in the shared Lists object is this entry stored in? +/// +/// If the value is `Idle`, then an entry's waker may move it to the notified +/// list. Otherwise, only the `IdleNotifiedSet` may move it. +/// +/// If the value is `Neither`, then it is still possible that the entry is in +/// some third external list (this happens in `drain`). +#[derive(Copy, Clone, Eq, PartialEq)] +enum List { + Notified, + Idle, + Neither, +} + +/// An entry in the list. +/// +/// # Safety +/// +/// The `my_list` field must only be accessed while holding the mutex in +/// `parent`. It is an invariant that the value of `my_list` corresponds to +/// which linked list in the `parent` holds this entry. Once this field takes +/// the value `Neither`, then it may never be modified again. +/// +/// If the value of `my_list` is `Notified` or `Idle`, then the `pointers` field +/// must only be accessed while holding the mutex. If the value of `my_list` is +/// `Neither`, then the `pointers` field may be accessed by the +/// `IdleNotifiedSet` (this happens inside `drain`). +/// +/// The `value` field is owned by the `IdleNotifiedSet` and may only be accessed +/// by the `IdleNotifiedSet`. The operation that sets the value of `my_list` to +/// `Neither` assumes ownership of the `value`, and it must either drop it or +/// move it out from this entry to prevent it from getting leaked. (Since the +/// two linked lists are emptied in the destructor of `IdleNotifiedSet`, the +/// value should not be leaked.) +struct ListEntry<T> { + /// The linked list pointers of the list this entry is in. + pointers: linked_list::Pointers<ListEntry<T>>, + /// Pointer to the shared `Lists` struct. + parent: Arc<Lists<T>>, + /// The value stored in this entry. + value: UnsafeCell<ManuallyDrop<T>>, + /// Used to remember which list this entry is in. + my_list: UnsafeCell<List>, + /// Required by the `linked_list::Pointers` field. + _pin: PhantomPinned, +} + +generate_addr_of_methods! { + impl<T> ListEntry<T> { + unsafe fn addr_of_pointers(self: NonNull<Self>) -> NonNull<linked_list::Pointers<ListEntry<T>>> { + &self.pointers + } + } +} + +// With mutable access to the `IdleNotifiedSet`, you can get mutable access to +// the values. +unsafe impl<T: Send> Send for IdleNotifiedSet<T> {} +// With the current API we strictly speaking don't even need `T: Sync`, but we +// require it anyway to support adding &self APIs that access the values in the +// future. +unsafe impl<T: Sync> Sync for IdleNotifiedSet<T> {} + +// These impls control when it is safe to create a Waker. Since the waker does +// not allow access to the value in any way (including its destructor), it is +// not necessary for `T` to be Send or Sync. +unsafe impl<T> Send for ListEntry<T> {} +unsafe impl<T> Sync for ListEntry<T> {} + +impl<T> IdleNotifiedSet<T> { + /// Create a new IdleNotifiedSet. + pub(crate) fn new() -> Self { + let lists = Mutex::new(ListsInner { + notified: LinkedList::new(), + idle: LinkedList::new(), + waker: None, + }); + + IdleNotifiedSet { + lists: Arc::new(lists), + length: 0, + } + } + + pub(crate) fn len(&self) -> usize { + self.length + } + + pub(crate) fn is_empty(&self) -> bool { + self.length == 0 + } + + /// Insert the given value into the `idle` list. + pub(crate) fn insert_idle(&mut self, value: T) -> EntryInOneOfTheLists<'_, T> { + self.length += 1; + + let entry = Arc::new(ListEntry { + parent: self.lists.clone(), + value: UnsafeCell::new(ManuallyDrop::new(value)), + my_list: UnsafeCell::new(List::Idle), + pointers: linked_list::Pointers::new(), + _pin: PhantomPinned, + }); + + { + let mut lock = self.lists.lock(); + lock.idle.push_front(entry.clone()); + } + + // Safety: We just put the entry in the idle list, so it is in one of the lists. + EntryInOneOfTheLists { entry, set: self } + } + + /// Pop an entry from the notified list to poll it. The entry is moved to + /// the idle list atomically. + pub(crate) fn pop_notified(&mut self, waker: &Waker) -> Option<EntryInOneOfTheLists<'_, T>> { + // We don't decrement the length because this call moves the entry to + // the idle list rather than removing it. + if self.length == 0 { + // Fast path. + return None; + } + + let mut lock = self.lists.lock(); + + let should_update_waker = match lock.waker.as_mut() { + Some(cur_waker) => !waker.will_wake(cur_waker), + None => true, + }; + if should_update_waker { + lock.waker = Some(waker.clone()); + } + + // Pop the entry, returning None if empty. + let entry = lock.notified.pop_back()?; + + lock.idle.push_front(entry.clone()); + + // Safety: We are holding the lock. + entry.my_list.with_mut(|ptr| unsafe { + *ptr = List::Idle; + }); + + drop(lock); + + // Safety: We just put the entry in the idle list, so it is in one of the lists. + Some(EntryInOneOfTheLists { entry, set: self }) + } + + /// Call a function on every element in this list. + pub(crate) fn for_each<F: FnMut(&mut T)>(&mut self, mut func: F) { + fn get_ptrs<T>(list: &mut LinkedList<T>, ptrs: &mut Vec<*mut T>) { + let mut node = list.last(); + + while let Some(entry) = node { + ptrs.push(entry.value.with_mut(|ptr| { + let ptr: *mut ManuallyDrop<T> = ptr; + let ptr: *mut T = ptr.cast(); + ptr + })); + + let prev = entry.pointers.get_prev(); + node = prev.map(|prev| unsafe { &*prev.as_ptr() }); + } + } + + // Atomically get a raw pointer to the value of every entry. + // + // Since this only locks the mutex once, it is not possible for a value + // to get moved from the idle list to the notified list during the + // operation, which would otherwise result in some value being listed + // twice. + let mut ptrs = Vec::with_capacity(self.len()); + { + let mut lock = self.lists.lock(); + + get_ptrs(&mut lock.idle, &mut ptrs); + get_ptrs(&mut lock.notified, &mut ptrs); + } + debug_assert_eq!(ptrs.len(), ptrs.capacity()); + + for ptr in ptrs { + // Safety: When we grabbed the pointers, the entries were in one of + // the two lists. This means that their value was valid at the time, + // and it must still be valid because we are the IdleNotifiedSet, + // and only we can remove an entry from the two lists. (It's + // possible that an entry is moved from one list to the other during + // this loop, but that is ok.) + func(unsafe { &mut *ptr }); + } + } + + /// Remove all entries in both lists, applying some function to each element. + /// + /// The closure is called on all elements even if it panics. Having it panic + /// twice is a double-panic, and will abort the application. + pub(crate) fn drain<F: FnMut(T)>(&mut self, func: F) { + if self.length == 0 { + // Fast path. + return; + } + self.length = 0; + + // The LinkedList is not cleared on panic, so we use a bomb to clear it. + // + // This value has the invariant that any entry in its `all_entries` list + // has `my_list` set to `Neither` and that the value has not yet been + // dropped. + struct AllEntries<T, F: FnMut(T)> { + all_entries: LinkedList<T>, + func: F, + } + + impl<T, F: FnMut(T)> AllEntries<T, F> { + fn pop_next(&mut self) -> bool { + if let Some(entry) = self.all_entries.pop_back() { + // Safety: We just took this value from the list, so we can + // destroy the value in the entry. + entry + .value + .with_mut(|ptr| unsafe { (self.func)(ManuallyDrop::take(&mut *ptr)) }); + true + } else { + false + } + } + } + + impl<T, F: FnMut(T)> Drop for AllEntries<T, F> { + fn drop(&mut self) { + while self.pop_next() {} + } + } + + let mut all_entries = AllEntries { + all_entries: LinkedList::new(), + func, + }; + + // Atomically move all entries to the new linked list in the AllEntries + // object. + { + let mut lock = self.lists.lock(); + unsafe { + // Safety: We are holding the lock and `all_entries` is a new + // LinkedList. + move_to_new_list(&mut lock.idle, &mut all_entries.all_entries); + move_to_new_list(&mut lock.notified, &mut all_entries.all_entries); + } + } + + // Keep destroying entries in the list until it is empty. + // + // If the closure panics, then the destructor of the `AllEntries` bomb + // ensures that we keep running the destructor on the remaining values. + // A second panic will abort the program. + while all_entries.pop_next() {} + } +} + +/// # Safety +/// +/// The mutex for the entries must be held, and the target list must be such +/// that setting `my_list` to `Neither` is ok. +unsafe fn move_to_new_list<T>(from: &mut LinkedList<T>, to: &mut LinkedList<T>) { + while let Some(entry) = from.pop_back() { + entry.my_list.with_mut(|ptr| { + *ptr = List::Neither; + }); + to.push_front(entry); + } +} + +impl<'a, T> EntryInOneOfTheLists<'a, T> { + /// Remove this entry from the list it is in, returning the value associated + /// with the entry. + /// + /// This consumes the value, since it is no longer guaranteed to be in a + /// list. + pub(crate) fn remove(self) -> T { + self.set.length -= 1; + + { + let mut lock = self.set.lists.lock(); + + // Safety: We are holding the lock so there is no race, and we will + // remove the entry afterwards to uphold invariants. + let old_my_list = self.entry.my_list.with_mut(|ptr| unsafe { + let old_my_list = *ptr; + *ptr = List::Neither; + old_my_list + }); + + let list = match old_my_list { + List::Idle => &mut lock.idle, + List::Notified => &mut lock.notified, + // An entry in one of the lists is in one of the lists. + List::Neither => unreachable!(), + }; + + unsafe { + // Safety: We just checked that the entry is in this particular + // list. + list.remove(ListEntry::as_raw(&self.entry)).unwrap(); + } + } + + // By setting `my_list` to `Neither`, we have taken ownership of the + // value. We return it to the caller. + // + // Safety: We have a mutable reference to the `IdleNotifiedSet` that + // owns this entry, so we can use its permission to access the value. + self.entry + .value + .with_mut(|ptr| unsafe { ManuallyDrop::take(&mut *ptr) }) + } + + /// Access the value in this entry together with a context for its waker. + pub(crate) fn with_value_and_context<F, U>(&mut self, func: F) -> U + where + F: FnOnce(&mut T, &mut Context<'_>) -> U, + T: 'static, + { + let waker = waker_ref(&self.entry); + + let mut context = Context::from_waker(&waker); + + // Safety: We have a mutable reference to the `IdleNotifiedSet` that + // owns this entry, so we can use its permission to access the value. + self.entry + .value + .with_mut(|ptr| unsafe { func(&mut *ptr, &mut context) }) + } +} + +impl<T> Drop for IdleNotifiedSet<T> { + fn drop(&mut self) { + // Clear both lists. + self.drain(drop); + + #[cfg(debug_assertions)] + if !std::thread::panicking() { + let lock = self.lists.lock(); + assert!(lock.idle.is_empty()); + assert!(lock.notified.is_empty()); + } + } +} + +impl<T: 'static> Wake for ListEntry<T> { + fn wake_by_ref(me: &Arc<Self>) { + let mut lock = me.parent.lock(); + + // Safety: We are holding the lock and we will update the lists to + // maintain invariants. + let old_my_list = me.my_list.with_mut(|ptr| unsafe { + let old_my_list = *ptr; + if old_my_list == List::Idle { + *ptr = List::Notified; + } + old_my_list + }); + + if old_my_list == List::Idle { + // We move ourself to the notified list. + let me = unsafe { + // Safety: We just checked that we are in this particular list. + lock.idle.remove(ListEntry::as_raw(me)).unwrap() + }; + lock.notified.push_front(me); + + if let Some(waker) = lock.waker.take() { + drop(lock); + waker.wake(); + } + } + } + + fn wake(me: Arc<Self>) { + Self::wake_by_ref(&me) + } +} + +/// # Safety +/// +/// `ListEntry` is forced to be !Unpin. +unsafe impl<T> linked_list::Link for ListEntry<T> { + type Handle = Arc<ListEntry<T>>; + type Target = ListEntry<T>; + + fn as_raw(handle: &Self::Handle) -> NonNull<ListEntry<T>> { + let ptr: *const ListEntry<T> = Arc::as_ptr(handle); + // Safety: We can't get a null pointer from `Arc::as_ptr`. + unsafe { NonNull::new_unchecked(ptr as *mut ListEntry<T>) } + } + + unsafe fn from_raw(ptr: NonNull<ListEntry<T>>) -> Arc<ListEntry<T>> { + Arc::from_raw(ptr.as_ptr()) + } + + unsafe fn pointers( + target: NonNull<ListEntry<T>>, + ) -> NonNull<linked_list::Pointers<ListEntry<T>>> { + ListEntry::addr_of_pointers(target) + } +} + +#[cfg(all(test, not(loom)))] +mod tests { + use crate::runtime::Builder; + use crate::task::JoinSet; + + // A test that runs under miri. + // + // https://github.com/tokio-rs/tokio/pull/5693 + #[test] + fn join_set_test() { + let rt = Builder::new_current_thread().build().unwrap(); + + let mut set = JoinSet::new(); + set.spawn_on(futures::future::ready(()), rt.handle()); + + rt.block_on(set.join_next()).unwrap().unwrap(); + } +} diff --git a/vendor/tokio/src/util/linked_list.rs b/vendor/tokio/src/util/linked_list.rs index a74f56215..ca6ef0e7b 100644 --- a/vendor/tokio/src/util/linked_list.rs +++ b/vendor/tokio/src/util/linked_list.rs @@ -1,6 +1,6 @@ #![cfg_attr(not(feature = "full"), allow(dead_code))] -//! An intrusive double linked list of data +//! An intrusive double linked list of data. //! //! The data structure supports tracking pinned nodes. Most of the data //! structure's APIs are `unsafe` as they require the caller to ensure the @@ -46,10 +46,10 @@ pub(crate) unsafe trait Link { /// This is usually a pointer-ish type. type Handle; - /// Node type + /// Node type. type Target; - /// Convert the handle to a raw pointer without consuming the handle + /// Convert the handle to a raw pointer without consuming the handle. #[allow(clippy::wrong_self_convention)] fn as_raw(handle: &Self::Handle) -> NonNull<Self::Target>; @@ -57,10 +57,17 @@ pub(crate) unsafe trait Link { unsafe fn from_raw(ptr: NonNull<Self::Target>) -> Self::Handle; /// Return the pointers for a node + /// + /// # Safety + /// + /// The resulting pointer should have the same tag in the stacked-borrows + /// stack as the argument. In particular, the method may not create an + /// intermediate reference in the process of creating the resulting raw + /// pointer. unsafe fn pointers(target: NonNull<Self::Target>) -> NonNull<Pointers<Self::Target>>; } -/// Previous / next pointers +/// Previous / next pointers. pub(crate) struct Pointers<T> { inner: UnsafeCell<PointersInner<T>>, } @@ -119,7 +126,7 @@ impl<L: Link> LinkedList<L, L::Target> { pub(crate) fn push_front(&mut self, val: L::Handle) { // The value should not be dropped, it is being inserted into the list let val = ManuallyDrop::new(val); - let ptr = L::as_raw(&*val); + let ptr = L::as_raw(&val); assert_ne!(self.head, Some(ptr)); unsafe { L::pointers(ptr).as_mut().set_next(self.head); @@ -171,8 +178,12 @@ impl<L: Link> LinkedList<L, L::Target> { /// /// # Safety /// - /// The caller **must** ensure that `node` is currently contained by - /// `self` or not contained by any other list. + /// The caller **must** ensure that exactly one of the following is true: + /// - `node` is currently contained by `self`, + /// - `node` is not contained by any list, + /// - `node` is currently contained by some other `GuardedLinkedList` **and** + /// the caller has an exclusive access to that list. This condition is + /// used by the linked list in `sync::Notify`. pub(crate) unsafe fn remove(&mut self, node: NonNull<L::Target>) -> Option<L::Handle> { if let Some(prev) = L::pointers(node).as_ref().get_prev() { debug_assert_eq!(L::pointers(prev).as_ref().get_next(), Some(node)); @@ -217,8 +228,56 @@ impl<L: Link> fmt::Debug for LinkedList<L, L::Target> { } } +// ===== impl CountedLinkedList ==== + +// Delegates operations to the base LinkedList implementation, and adds a counter to the elements +// in the list. +pub(crate) struct CountedLinkedList<L: Link, T> { + list: LinkedList<L, T>, + count: usize, +} + +impl<L: Link> CountedLinkedList<L, L::Target> { + pub(crate) fn new() -> CountedLinkedList<L, L::Target> { + CountedLinkedList { + list: LinkedList::new(), + count: 0, + } + } + + pub(crate) fn push_front(&mut self, val: L::Handle) { + self.list.push_front(val); + self.count += 1; + } + + pub(crate) fn pop_back(&mut self) -> Option<L::Handle> { + let val = self.list.pop_back(); + if val.is_some() { + self.count -= 1; + } + val + } + + pub(crate) fn is_empty(&self) -> bool { + self.list.is_empty() + } + + pub(crate) unsafe fn remove(&mut self, node: NonNull<L::Target>) -> Option<L::Handle> { + let val = self.list.remove(node); + if val.is_some() { + self.count -= 1; + } + val + } + + pub(crate) fn count(&self) -> usize { + self.count + } +} + #[cfg(any( feature = "fs", + feature = "rt", all(unix, feature = "process"), feature = "signal", feature = "sync", @@ -236,37 +295,6 @@ impl<L: Link> Default for LinkedList<L, L::Target> { } } -// ===== impl Iter ===== - -cfg_rt_multi_thread! { - pub(crate) struct Iter<'a, T: Link> { - curr: Option<NonNull<T::Target>>, - _p: core::marker::PhantomData<&'a T>, - } - - impl<L: Link> LinkedList<L, L::Target> { - pub(crate) fn iter(&self) -> Iter<'_, L> { - Iter { - curr: self.head, - _p: core::marker::PhantomData, - } - } - } - - impl<'a, T: Link> Iterator for Iter<'a, T> { - type Item = &'a T::Target; - - fn next(&mut self) -> Option<&'a T::Target> { - let curr = self.curr?; - // safety: the pointer references data contained by the list - self.curr = unsafe { T::pointers(curr).as_ref() }.get_next(); - - // safety: the value is still owned by the linked list. - Some(unsafe { &*curr.as_ptr() }) - } - } -} - // ===== impl DrainFilter ===== cfg_io_readiness! { @@ -313,6 +341,126 @@ cfg_io_readiness! { } } +cfg_taskdump! { + impl<T: Link> CountedLinkedList<T, T::Target> { + pub(crate) fn for_each<F>(&mut self, f: F) + where + F: FnMut(&T::Handle), + { + self.list.for_each(f) + } + } + + impl<T: Link> LinkedList<T, T::Target> { + pub(crate) fn for_each<F>(&mut self, mut f: F) + where + F: FnMut(&T::Handle), + { + use std::mem::ManuallyDrop; + + let mut next = self.head; + + while let Some(curr) = next { + unsafe { + let handle = ManuallyDrop::new(T::from_raw(curr)); + f(&handle); + next = T::pointers(curr).as_ref().get_next(); + } + } + } + } +} + +// ===== impl GuardedLinkedList ===== + +feature! { + #![any( + feature = "process", + feature = "sync", + feature = "rt", + feature = "signal", + )] + + /// An intrusive linked list, but instead of keeping pointers to the head + /// and tail nodes, it uses a special guard node linked with those nodes. + /// It means that the list is circular and every pointer of a node from + /// the list is not `None`, including pointers from the guard node. + /// + /// If a list is empty, then both pointers of the guard node are pointing + /// at the guard node itself. + pub(crate) struct GuardedLinkedList<L, T> { + /// Pointer to the guard node. + guard: NonNull<T>, + + /// Node type marker. + _marker: PhantomData<*const L>, + } + + impl<U, L: Link<Handle = NonNull<U>>> LinkedList<L, L::Target> { + /// Turns a linked list into the guarded version by linking the guard node + /// with the head and tail nodes. Like with other nodes, you should guarantee + /// that the guard node is pinned in memory. + pub(crate) fn into_guarded(self, guard_handle: L::Handle) -> GuardedLinkedList<L, L::Target> { + // `guard_handle` is a NonNull pointer, we don't have to care about dropping it. + let guard = L::as_raw(&guard_handle); + + unsafe { + if let Some(head) = self.head { + debug_assert!(L::pointers(head).as_ref().get_prev().is_none()); + L::pointers(head).as_mut().set_prev(Some(guard)); + L::pointers(guard).as_mut().set_next(Some(head)); + + // The list is not empty, so the tail cannot be `None`. + let tail = self.tail.unwrap(); + debug_assert!(L::pointers(tail).as_ref().get_next().is_none()); + L::pointers(tail).as_mut().set_next(Some(guard)); + L::pointers(guard).as_mut().set_prev(Some(tail)); + } else { + // The list is empty. + L::pointers(guard).as_mut().set_prev(Some(guard)); + L::pointers(guard).as_mut().set_next(Some(guard)); + } + } + + GuardedLinkedList { guard, _marker: PhantomData } + } + } + + impl<L: Link> GuardedLinkedList<L, L::Target> { + fn tail(&self) -> Option<NonNull<L::Target>> { + let tail_ptr = unsafe { + L::pointers(self.guard).as_ref().get_prev().unwrap() + }; + + // Compare the tail pointer with the address of the guard node itself. + // If the guard points at itself, then there are no other nodes and + // the list is considered empty. + if tail_ptr != self.guard { + Some(tail_ptr) + } else { + None + } + } + + /// Removes the last element from a list and returns it, or None if it is + /// empty. + pub(crate) fn pop_back(&mut self) -> Option<L::Handle> { + unsafe { + let last = self.tail()?; + let before_last = L::pointers(last).as_ref().get_prev().unwrap(); + + L::pointers(self.guard).as_mut().set_prev(Some(before_last)); + L::pointers(before_last).as_mut().set_next(Some(self.guard)); + + L::pointers(last).as_mut().set_prev(None); + L::pointers(last).as_mut().set_next(None); + + Some(L::from_raw(last)) + } + } + } +} + // ===== impl Pointers ===== impl<T> Pointers<T> { @@ -327,7 +475,7 @@ impl<T> Pointers<T> { } } - fn get_prev(&self) -> Option<NonNull<T>> { + pub(crate) fn get_prev(&self) -> Option<NonNull<T>> { // SAFETY: prev is the first field in PointersInner, which is #[repr(C)]. unsafe { let inner = self.inner.get(); @@ -335,7 +483,7 @@ impl<T> Pointers<T> { ptr::read(prev) } } - fn get_next(&self) -> Option<NonNull<T>> { + pub(crate) fn get_next(&self) -> Option<NonNull<T>> { // SAFETY: next is the second field in PointersInner, which is #[repr(C)]. unsafe { let inner = self.inner.get(); @@ -375,14 +523,15 @@ impl<T> fmt::Debug for Pointers<T> { } } -#[cfg(test)] +#[cfg(any(test, fuzzing))] #[cfg(not(loom))] -mod tests { +pub(crate) mod tests { use super::*; use std::pin::Pin; #[derive(Debug)] + #[repr(C)] struct Entry { pointers: Pointers<Entry>, val: i32, @@ -400,8 +549,8 @@ mod tests { Pin::new_unchecked(&*ptr.as_ptr()) } - unsafe fn pointers(mut target: NonNull<Entry>) -> NonNull<Pointers<Entry>> { - NonNull::from(&mut target.as_mut().pointers) + unsafe fn pointers(target: NonNull<Entry>) -> NonNull<Pointers<Entry>> { + target.cast() } } @@ -435,6 +584,7 @@ mod tests { } } + #[cfg(test)] macro_rules! assert_clean { ($e:ident) => {{ assert!($e.pointers.get_next().is_none()); @@ -442,6 +592,7 @@ mod tests { }}; } + #[cfg(test)] macro_rules! assert_ptr_eq { ($a:expr, $b:expr) => {{ // Deal with mapping a Pin<&mut T> -> Option<NonNull<T>> @@ -646,46 +797,41 @@ mod tests { } #[test] - fn iter() { + fn count() { + let mut list = CountedLinkedList::<&Entry, <&Entry as Link>::Target>::new(); + assert_eq!(0, list.count()); + let a = entry(5); let b = entry(7); - - let mut list = LinkedList::<&Entry, <&Entry as Link>::Target>::new(); - - assert_eq!(0, list.iter().count()); - list.push_front(a.as_ref()); list.push_front(b.as_ref()); + assert_eq!(2, list.count()); - let mut i = list.iter(); - assert_eq!(7, i.next().unwrap().val); - assert_eq!(5, i.next().unwrap().val); - assert!(i.next().is_none()); - } + list.pop_back(); + assert_eq!(1, list.count()); - proptest::proptest! { - #[test] - fn fuzz_linked_list(ops: Vec<usize>) { - run_fuzz(ops); + unsafe { + list.remove(ptr(&b)); } + assert_eq!(0, list.count()); } - fn run_fuzz(ops: Vec<usize>) { - use std::collections::VecDeque; - - #[derive(Debug)] + /// This is a fuzz test. You run it by entering `cargo fuzz run fuzz_linked_list` in CLI in `/tokio/` module. + #[cfg(fuzzing)] + pub fn fuzz_linked_list(ops: &[u8]) { enum Op { Push, Pop, Remove(usize), } + use std::collections::VecDeque; let ops = ops .iter() - .map(|i| match i % 3 { + .map(|i| match i % 3u8 { 0 => Op::Push, 1 => Op::Pop, - 2 => Op::Remove(i / 3), + 2 => Op::Remove((i / 3u8) as usize), _ => unreachable!(), }) .collect::<Vec<_>>(); diff --git a/vendor/tokio/src/util/markers.rs b/vendor/tokio/src/util/markers.rs new file mode 100644 index 000000000..7031fb6bc --- /dev/null +++ b/vendor/tokio/src/util/markers.rs @@ -0,0 +1,8 @@ +/// Marker for types that are `Sync` but not `Send` +pub(crate) struct SyncNotSend(*mut ()); + +unsafe impl Sync for SyncNotSend {} + +cfg_rt! { + pub(crate) struct NotSendOrSync(*mut ()); +} diff --git a/vendor/tokio/src/util/memchr.rs b/vendor/tokio/src/util/memchr.rs new file mode 100644 index 000000000..44fd2da37 --- /dev/null +++ b/vendor/tokio/src/util/memchr.rs @@ -0,0 +1,74 @@ +//! Search for a byte in a byte array using libc. +//! +//! When nothing pulls in libc, then just use a trivial implementation. Note +//! that we only depend on libc on unix. + +#[cfg(not(all(unix, feature = "libc")))] +pub(crate) fn memchr(needle: u8, haystack: &[u8]) -> Option<usize> { + haystack.iter().position(|val| needle == *val) +} + +#[cfg(all(unix, feature = "libc"))] +pub(crate) fn memchr(needle: u8, haystack: &[u8]) -> Option<usize> { + let start = haystack.as_ptr(); + + // SAFETY: `start` is valid for `haystack.len()` bytes. + let ptr = unsafe { libc::memchr(start.cast(), needle as _, haystack.len()) }; + + if ptr.is_null() { + None + } else { + Some(ptr as usize - start as usize) + } +} + +#[cfg(test)] +mod tests { + use super::memchr; + + #[test] + fn memchr_test() { + let haystack = b"123abc456\0\xffabc\n"; + + assert_eq!(memchr(b'1', haystack), Some(0)); + assert_eq!(memchr(b'2', haystack), Some(1)); + assert_eq!(memchr(b'3', haystack), Some(2)); + assert_eq!(memchr(b'4', haystack), Some(6)); + assert_eq!(memchr(b'5', haystack), Some(7)); + assert_eq!(memchr(b'6', haystack), Some(8)); + assert_eq!(memchr(b'7', haystack), None); + assert_eq!(memchr(b'a', haystack), Some(3)); + assert_eq!(memchr(b'b', haystack), Some(4)); + assert_eq!(memchr(b'c', haystack), Some(5)); + assert_eq!(memchr(b'd', haystack), None); + assert_eq!(memchr(b'A', haystack), None); + assert_eq!(memchr(0, haystack), Some(9)); + assert_eq!(memchr(0xff, haystack), Some(10)); + assert_eq!(memchr(0xfe, haystack), None); + assert_eq!(memchr(1, haystack), None); + assert_eq!(memchr(b'\n', haystack), Some(14)); + assert_eq!(memchr(b'\r', haystack), None); + } + + #[test] + fn memchr_all() { + let mut arr = Vec::new(); + for b in 0..=255 { + arr.push(b); + } + for b in 0..=255 { + assert_eq!(memchr(b, &arr), Some(b as usize)); + } + arr.reverse(); + for b in 0..=255 { + assert_eq!(memchr(b, &arr), Some(255 - b as usize)); + } + } + + #[test] + fn memchr_empty() { + for b in 0..=255 { + assert_eq!(memchr(b, b""), None); + } + } +} diff --git a/vendor/tokio/src/util/mod.rs b/vendor/tokio/src/util/mod.rs index b267125b1..6a7d4b103 100644 --- a/vendor/tokio/src/util/mod.rs +++ b/vendor/tokio/src/util/mod.rs @@ -3,6 +3,40 @@ cfg_io_driver! { pub(crate) mod slab; } +#[cfg(feature = "rt")] +pub(crate) mod atomic_cell; + +#[cfg(any( + feature = "rt", + feature = "signal", + feature = "process", + tokio_no_const_mutex_new, +))] +pub(crate) mod once_cell; + +#[cfg(any( + // io driver uses `WakeList` directly + feature = "net", + feature = "process", + // `sync` enables `Notify` and `batch_semaphore`, which require `WakeList`. + feature = "sync", + // `fs` uses `batch_semaphore`, which requires `WakeList`. + feature = "fs", + // rt and signal use `Notify`, which requires `WakeList`. + feature = "rt", + feature = "signal", +))] +mod wake_list; +#[cfg(any( + feature = "net", + feature = "process", + feature = "sync", + feature = "fs", + feature = "rt", + feature = "signal", +))] +pub(crate) use wake_list::WakeList; + #[cfg(any( feature = "fs", feature = "net", @@ -14,33 +48,36 @@ cfg_io_driver! { ))] pub(crate) mod linked_list; -#[cfg(any(feature = "rt-multi-thread", feature = "macros"))] -mod rand; +#[cfg(any(feature = "rt", feature = "macros"))] +pub(crate) mod rand; cfg_rt! { + mod idle_notified_set; + pub(crate) use idle_notified_set::IdleNotifiedSet; + + pub(crate) use self::rand::RngSeedGenerator; + mod wake; pub(crate) use wake::WakerRef; pub(crate) use wake::{waker_ref, Wake}; + + mod sync_wrapper; + pub(crate) use sync_wrapper::SyncWrapper; + + mod rc_cell; + pub(crate) use rc_cell::RcCell; } cfg_rt_multi_thread! { - pub(crate) use self::rand::FastRand; - mod try_lock; pub(crate) use try_lock::TryLock; } pub(crate) mod trace; -#[cfg(any(feature = "macros"))] -#[cfg_attr(not(feature = "macros"), allow(unreachable_pub))] -pub use self::rand::thread_rng_n; - -#[cfg(any( - feature = "rt", - feature = "time", - feature = "net", - feature = "process", - all(unix, feature = "signal") -))] pub(crate) mod error; + +#[cfg(feature = "io-util")] +pub(crate) mod memchr; + +pub(crate) mod markers; diff --git a/vendor/tokio/src/util/once_cell.rs b/vendor/tokio/src/util/once_cell.rs new file mode 100644 index 000000000..71fc00758 --- /dev/null +++ b/vendor/tokio/src/util/once_cell.rs @@ -0,0 +1,70 @@ +#![allow(dead_code)] +use std::cell::UnsafeCell; +use std::mem::MaybeUninit; +use std::sync::Once; + +pub(crate) struct OnceCell<T> { + once: Once, + value: UnsafeCell<MaybeUninit<T>>, +} + +unsafe impl<T: Send + Sync> Send for OnceCell<T> {} +unsafe impl<T: Send + Sync> Sync for OnceCell<T> {} + +impl<T> OnceCell<T> { + pub(crate) const fn new() -> Self { + Self { + once: Once::new(), + value: UnsafeCell::new(MaybeUninit::uninit()), + } + } + + /// Get the value inside this cell, initializing it using the provided + /// function if necessary. + /// + /// If the `init` closure panics, then the `OnceCell` is poisoned and all + /// future calls to `get` will panic. + #[inline] + pub(crate) fn get(&self, init: impl FnOnce() -> T) -> &T { + if !self.once.is_completed() { + self.do_init(init); + } + + // Safety: The `std::sync::Once` guarantees that we can only reach this + // line if a `call_once` closure has been run exactly once and without + // panicking. Thus, the value is not uninitialized. + // + // There is also no race because the only `&self` method that modifies + // `value` is `do_init`, but if the `call_once` closure is still + // running, then no thread has gotten past the `call_once`. + unsafe { &*(self.value.get() as *const T) } + } + + #[cold] + fn do_init(&self, init: impl FnOnce() -> T) { + let value_ptr = self.value.get() as *mut T; + + self.once.call_once(|| { + let set_to = init(); + + // Safety: The `std::sync::Once` guarantees that this initialization + // will run at most once, and that no thread can get past the + // `call_once` until it has run exactly once. Thus, we have + // exclusive access to `value`. + unsafe { + std::ptr::write(value_ptr, set_to); + } + }); + } +} + +impl<T> Drop for OnceCell<T> { + fn drop(&mut self) { + if self.once.is_completed() { + let value_ptr = self.value.get() as *mut T; + unsafe { + std::ptr::drop_in_place(value_ptr); + } + } + } +} diff --git a/vendor/tokio/src/util/pad.rs b/vendor/tokio/src/util/pad.rs deleted file mode 100644 index bf0913ca8..000000000 --- a/vendor/tokio/src/util/pad.rs +++ /dev/null @@ -1,52 +0,0 @@ -use core::fmt; -use core::ops::{Deref, DerefMut}; - -#[derive(Clone, Copy, Default, Hash, PartialEq, Eq)] -// Starting from Intel's Sandy Bridge, spatial prefetcher is now pulling pairs of 64-byte cache -// lines at a time, so we have to align to 128 bytes rather than 64. -// -// Sources: -// - https://www.intel.com/content/dam/www/public/us/en/documents/manuals/64-ia-32-architectures-optimization-manual.pdf -// - https://github.com/facebook/folly/blob/1b5288e6eea6df074758f877c849b6e73bbb9fbb/folly/lang/Align.h#L107 -#[cfg_attr(target_arch = "x86_64", repr(align(128)))] -#[cfg_attr(not(target_arch = "x86_64"), repr(align(64)))] -pub(crate) struct CachePadded<T> { - value: T, -} - -unsafe impl<T: Send> Send for CachePadded<T> {} -unsafe impl<T: Sync> Sync for CachePadded<T> {} - -impl<T> CachePadded<T> { - pub(crate) fn new(t: T) -> CachePadded<T> { - CachePadded::<T> { value: t } - } -} - -impl<T> Deref for CachePadded<T> { - type Target = T; - - fn deref(&self) -> &T { - &self.value - } -} - -impl<T> DerefMut for CachePadded<T> { - fn deref_mut(&mut self) -> &mut T { - &mut self.value - } -} - -impl<T: fmt::Debug> fmt::Debug for CachePadded<T> { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("CachePadded") - .field("value", &self.value) - .finish() - } -} - -impl<T> From<T> for CachePadded<T> { - fn from(t: T) -> Self { - CachePadded::new(t) - } -} diff --git a/vendor/tokio/src/util/rand.rs b/vendor/tokio/src/util/rand.rs index 17b3ec1ae..d96c8d37e 100644 --- a/vendor/tokio/src/util/rand.rs +++ b/vendor/tokio/src/util/rand.rs @@ -1,21 +1,43 @@ -use std::cell::Cell; +cfg_rt! { + mod rt; + pub(crate) use rt::RngSeedGenerator; -/// Fast random number generate + cfg_unstable! { + mod rt_unstable; + } +} + +/// A seed for random number generation. +/// +/// In order to make certain functions within a runtime deterministic, a seed +/// can be specified at the time of creation. +#[allow(unreachable_pub)] +#[derive(Clone, Debug)] +pub struct RngSeed { + s: u32, + r: u32, +} + +/// Fast random number generate. /// /// Implement xorshift64+: 2 32-bit xorshift sequences added together. /// Shift triplet `[17,7,16]` was calculated as indicated in Marsaglia's /// Xorshift paper: <https://www.jstatsoft.org/article/view/v008i14/xorshift.pdf> /// This generator passes the SmallCrush suite, part of TestU01 framework: /// <http://simul.iro.umontreal.ca/testu01/tu01.html> -#[derive(Debug)] +#[derive(Clone, Copy, Debug)] pub(crate) struct FastRand { - one: Cell<u32>, - two: Cell<u32>, + one: u32, + two: u32, } -impl FastRand { - /// Initialize a new, thread-local, fast random number generator. - pub(crate) fn new(seed: u64) -> FastRand { +impl RngSeed { + /// Creates a random seed using loom internally. + pub(crate) fn new() -> Self { + Self::from_u64(crate::loom::rand::seed()) + } + + fn from_u64(seed: u64) -> Self { let one = (seed >> 32) as u32; let mut two = seed as u32; @@ -24,41 +46,50 @@ impl FastRand { two = 1; } + Self::from_pair(one, two) + } + + fn from_pair(s: u32, r: u32) -> Self { + Self { s, r } + } +} + +impl FastRand { + /// Initialize a new fast random number generator using the default source of entropy. + pub(crate) fn new() -> FastRand { + FastRand::from_seed(RngSeed::new()) + } + + /// Initializes a new, thread-local, fast random number generator. + pub(crate) fn from_seed(seed: RngSeed) -> FastRand { FastRand { - one: Cell::new(one), - two: Cell::new(two), + one: seed.s, + two: seed.r, } } - pub(crate) fn fastrand_n(&self, n: u32) -> u32 { + #[cfg(any( + feature = "macros", + feature = "rt-multi-thread", + all(feature = "sync", feature = "rt") + ))] + pub(crate) fn fastrand_n(&mut self, n: u32) -> u32 { // This is similar to fastrand() % n, but faster. // See https://lemire.me/blog/2016/06/27/a-fast-alternative-to-the-modulo-reduction/ let mul = (self.fastrand() as u64).wrapping_mul(n as u64); (mul >> 32) as u32 } - fn fastrand(&self) -> u32 { - let mut s1 = self.one.get(); - let s0 = self.two.get(); + fn fastrand(&mut self) -> u32 { + let mut s1 = self.one; + let s0 = self.two; s1 ^= s1 << 17; s1 = s1 ^ s0 ^ s1 >> 7 ^ s0 >> 16; - self.one.set(s0); - self.two.set(s1); + self.one = s0; + self.two = s1; s0.wrapping_add(s1) } } - -// Used by the select macro and `StreamMap` -#[cfg(any(feature = "macros"))] -#[doc(hidden)] -#[cfg_attr(not(feature = "macros"), allow(unreachable_pub))] -pub fn thread_rng_n(n: u32) -> u32 { - thread_local! { - static THREAD_RNG: FastRand = FastRand::new(crate::loom::rand::seed()); - } - - THREAD_RNG.with(|rng| rng.fastrand_n(n)) -} diff --git a/vendor/tokio/src/util/rand/rt.rs b/vendor/tokio/src/util/rand/rt.rs new file mode 100644 index 000000000..6584ba4f2 --- /dev/null +++ b/vendor/tokio/src/util/rand/rt.rs @@ -0,0 +1,61 @@ +use super::{FastRand, RngSeed}; + +use std::sync::Mutex; + +/// A deterministic generator for seeds (and other generators). +/// +/// Given the same initial seed, the generator will output the same sequence of seeds. +/// +/// Since the seed generator will be kept in a runtime handle, we need to wrap `FastRand` +/// in a Mutex to make it thread safe. Different to the `FastRand` that we keep in a +/// thread local store, the expectation is that seed generation will not need to happen +/// very frequently, so the cost of the mutex should be minimal. +#[derive(Debug)] +pub(crate) struct RngSeedGenerator { + /// Internal state for the seed generator. We keep it in a Mutex so that we can safely + /// use it across multiple threads. + state: Mutex<FastRand>, +} + +impl RngSeedGenerator { + /// Returns a new generator from the provided seed. + pub(crate) fn new(seed: RngSeed) -> Self { + Self { + state: Mutex::new(FastRand::from_seed(seed)), + } + } + + /// Returns the next seed in the sequence. + pub(crate) fn next_seed(&self) -> RngSeed { + let mut rng = self + .state + .lock() + .expect("RNG seed generator is internally corrupt"); + + let s = rng.fastrand(); + let r = rng.fastrand(); + + RngSeed::from_pair(s, r) + } + + /// Directly creates a generator using the next seed. + pub(crate) fn next_generator(&self) -> Self { + RngSeedGenerator::new(self.next_seed()) + } +} + +impl FastRand { + /// Replaces the state of the random number generator with the provided seed, returning + /// the seed that represents the previous state of the random number generator. + /// + /// The random number generator will become equivalent to one created with + /// the same seed. + pub(crate) fn replace_seed(&mut self, seed: RngSeed) -> RngSeed { + let old_seed = RngSeed::from_pair(self.one, self.two); + + self.one = seed.s; + self.two = seed.r; + + old_seed + } +} diff --git a/vendor/tokio/src/util/rand/rt_unstable.rs b/vendor/tokio/src/util/rand/rt_unstable.rs new file mode 100644 index 000000000..3a8613eb0 --- /dev/null +++ b/vendor/tokio/src/util/rand/rt_unstable.rs @@ -0,0 +1,20 @@ +use super::RngSeed; + +use std::collections::hash_map::DefaultHasher; +use std::hash::Hasher; + +impl RngSeed { + /// Generates a seed from the provided byte slice. + /// + /// # Example + /// + /// ``` + /// # use tokio::runtime::RngSeed; + /// let seed = RngSeed::from_bytes(b"make me a seed"); + /// ``` + pub fn from_bytes(bytes: &[u8]) -> Self { + let mut hasher = DefaultHasher::default(); + hasher.write(bytes); + Self::from_u64(hasher.finish()) + } +} diff --git a/vendor/tokio/src/util/rc_cell.rs b/vendor/tokio/src/util/rc_cell.rs new file mode 100644 index 000000000..447249268 --- /dev/null +++ b/vendor/tokio/src/util/rc_cell.rs @@ -0,0 +1,57 @@ +use crate::loom::cell::UnsafeCell; + +use std::rc::Rc; + +/// This is exactly like `Cell<Option<Rc<T>>>`, except that it provides a `get` +/// method even though `Rc` is not `Copy`. +pub(crate) struct RcCell<T> { + inner: UnsafeCell<Option<Rc<T>>>, +} + +impl<T> RcCell<T> { + #[cfg(not(all(loom, test)))] + pub(crate) const fn new() -> Self { + Self { + inner: UnsafeCell::new(None), + } + } + + // The UnsafeCell in loom does not have a const `new` fn. + #[cfg(all(loom, test))] + pub(crate) fn new() -> Self { + Self { + inner: UnsafeCell::new(None), + } + } + + /// Safety: This method may not be called recursively. + #[inline] + unsafe fn with_inner<F, R>(&self, f: F) -> R + where + F: FnOnce(&mut Option<Rc<T>>) -> R, + { + // safety: This type is not Sync, so concurrent calls of this method + // cannot happen. Furthermore, the caller guarantees that the method is + // not called recursively. Finally, this is the only place that can + // create mutable references to the inner Rc. This ensures that any + // mutable references created here are exclusive. + self.inner.with_mut(|ptr| f(&mut *ptr)) + } + + pub(crate) fn get(&self) -> Option<Rc<T>> { + // safety: The `Rc::clone` method will not call any unknown user-code, + // so it will not result in a recursive call to `with_inner`. + unsafe { self.with_inner(|rc| rc.clone()) } + } + + pub(crate) fn replace(&self, val: Option<Rc<T>>) -> Option<Rc<T>> { + // safety: No destructors or other unknown user-code will run inside the + // `with_inner` call, so no recursive call to `with_inner` can happen. + unsafe { self.with_inner(|rc| std::mem::replace(rc, val)) } + } + + pub(crate) fn set(&self, val: Option<Rc<T>>) { + let old = self.replace(val); + drop(old); + } +} diff --git a/vendor/tokio/src/util/slab.rs b/vendor/tokio/src/util/slab.rs index 2ddaa6c74..cc1208ebe 100644 --- a/vendor/tokio/src/util/slab.rs +++ b/vendor/tokio/src/util/slab.rs @@ -85,11 +85,11 @@ pub(crate) struct Address(usize); /// An entry in the slab. pub(crate) trait Entry: Default { - /// Reset the entry's value and track the generation. + /// Resets the entry's value and track the generation. fn reset(&self); } -/// A reference to a value stored in the slab +/// A reference to a value stored in the slab. pub(crate) struct Ref<T> { value: *const Value<T>, } @@ -101,9 +101,9 @@ const NUM_PAGES: usize = 19; const PAGE_INITIAL_SIZE: usize = 32; const PAGE_INDEX_SHIFT: u32 = PAGE_INITIAL_SIZE.trailing_zeros() + 1; -/// A page in the slab +/// A page in the slab. struct Page<T> { - /// Slots + /// Slots. slots: Mutex<Slots<T>>, // Number of slots currently being used. This is not guaranteed to be up to @@ -116,7 +116,7 @@ struct Page<T> { // The number of slots the page can hold. len: usize, - // Length of all previous pages combined + // Length of all previous pages combined. prev_len: usize, } @@ -128,9 +128,9 @@ struct CachedPage<T> { init: usize, } -/// Page state +/// Page state. struct Slots<T> { - /// Slots + /// Slots. slots: Vec<Slot<T>>, head: usize, @@ -157,11 +157,17 @@ struct Slot<T> { /// Next entry in the free list. next: u32, + + /// Makes miri happy by making mutable references not take exclusive access. + /// + /// Could probably also be fixed by replacing `slots` with a raw-pointer + /// based equivalent. + _pin: std::marker::PhantomPinned, } -/// Value paired with a reference to the page +/// Value paired with a reference to the page. struct Value<T> { - /// Value stored in the value + /// Value stored in the value. value: T, /// Pointer to the page containing the slot. @@ -171,7 +177,7 @@ struct Value<T> { } impl<T> Slab<T> { - /// Create a new, empty, slab + /// Create a new, empty, slab. pub(crate) fn new() -> Slab<T> { // Initializing arrays is a bit annoying. Instead of manually writing // out an array and every single entry, `Default::default()` is used to @@ -409,7 +415,7 @@ impl<T: Entry> Page<T> { slot.value.with(|ptr| unsafe { (*ptr).value.reset() }); // Return a reference to the slot - Some((me.addr(idx), slot.gen_ref(me))) + Some((me.addr(idx), locked.gen_ref(idx, me))) } else if me.len == locked.slots.len() { // The page is full None @@ -428,9 +434,10 @@ impl<T: Entry> Page<T> { locked.slots.push(Slot { value: UnsafeCell::new(Value { value: Default::default(), - page: &**me as *const _, + page: Arc::as_ptr(me), }), next: 0, + _pin: std::marker::PhantomPinned, }); // Increment the head to indicate the free stack is empty @@ -443,7 +450,7 @@ impl<T: Entry> Page<T> { debug_assert_eq!(locked.slots.len(), locked.head); - Some((me.addr(idx), locked.slots[idx].gen_ref(me))) + Some((me.addr(idx), locked.gen_ref(idx, me))) } } } @@ -455,7 +462,7 @@ impl<T> Page<T> { addr.0 - self.prev_len } - /// Returns the address for the given slot + /// Returns the address for the given slot. fn addr(&self, slot: usize) -> Address { Address(slot + self.prev_len) } @@ -478,7 +485,7 @@ impl<T> Default for Page<T> { } impl<T> Page<T> { - /// Release a slot into the page's free list + /// Release a slot into the page's free list. fn release(&self, value: *const Value<T>) { let mut locked = self.slots.lock(); @@ -492,7 +499,7 @@ impl<T> Page<T> { } impl<T> CachedPage<T> { - /// Refresh the cache + /// Refreshes the cache. fn refresh(&mut self, page: &Page<T>) { let slots = page.slots.lock(); @@ -502,7 +509,7 @@ impl<T> CachedPage<T> { } } - // Get a value by index + /// Gets a value by index. fn get(&self, idx: usize) -> &T { assert!(idx < self.init); @@ -544,39 +551,35 @@ impl<T> Slots<T> { fn index_for(&self, slot: *const Value<T>) -> usize { use std::mem; - let base = &self.slots[0] as *const _ as usize; - - assert!(base != 0, "page is unallocated"); + assert_ne!(self.slots.capacity(), 0, "page is unallocated"); + let base = self.slots.as_ptr() as usize; let slot = slot as usize; let width = mem::size_of::<Slot<T>>(); assert!(slot >= base, "unexpected pointer"); let idx = (slot - base) / width; - assert!(idx < self.slots.len() as usize); + assert!(idx < self.slots.len()); idx } -} -impl<T: Entry> Slot<T> { - /// Generates a `Ref` for the slot. This involves bumping the page's ref count. - fn gen_ref(&self, page: &Arc<Page<T>>) -> Ref<T> { - // The ref holds a ref on the page. The `Arc` is forgotten here and is - // resurrected in `release` when the `Ref` is dropped. By avoiding to - // hold on to an explicit `Arc` value, the struct size of `Ref` is - // reduced. + /// Generates a `Ref` for the slot at the given index. This involves bumping the page's ref count. + fn gen_ref(&self, idx: usize, page: &Arc<Page<T>>) -> Ref<T> { + assert!(idx < self.slots.len()); mem::forget(page.clone()); - let slot = self as *const Slot<T>; - let value = slot as *const Value<T>; + + let vec_ptr = self.slots.as_ptr(); + let slot: *const Slot<T> = unsafe { vec_ptr.add(idx) }; + let value: *const Value<T> = slot as *const Value<T>; Ref { value } } } impl<T> Value<T> { - // Release the slot, returning the `Arc<Page<T>>` logically owned by the ref. + /// Releases the slot, returning the `Arc<Page<T>>` logically owned by the ref. fn release(&self) -> Arc<Page<T>> { // Safety: called by `Ref`, which owns an `Arc<Page<T>>` instance. let page = unsafe { Arc::from_raw(self.page) }; @@ -691,11 +694,13 @@ mod test { #[test] fn insert_many() { + const MANY: usize = normal_or_miri(10_000, 50); + let mut slab = Slab::<Foo>::new(); let alloc = slab.allocator(); let mut entries = vec![]; - for i in 0..10_000 { + for i in 0..MANY { let (addr, val) = alloc.allocate().unwrap(); val.id.store(i, SeqCst); entries.push((addr, val)); @@ -708,15 +713,15 @@ mod test { entries.clear(); - for i in 0..10_000 { + for i in 0..MANY { let (addr, val) = alloc.allocate().unwrap(); - val.id.store(10_000 - i, SeqCst); + val.id.store(MANY - i, SeqCst); entries.push((addr, val)); } for (i, (addr, v)) in entries.iter().enumerate() { - assert_eq!(10_000 - i, v.id.load(SeqCst)); - assert_eq!(10_000 - i, slab.get(*addr).unwrap().id.load(SeqCst)); + assert_eq!(MANY - i, v.id.load(SeqCst)); + assert_eq!(MANY - i, slab.get(*addr).unwrap().id.load(SeqCst)); } } @@ -726,7 +731,7 @@ mod test { let alloc = slab.allocator(); let mut entries = vec![]; - for i in 0..10_000 { + for i in 0..normal_or_miri(10_000, 100) { let (addr, val) = alloc.allocate().unwrap(); val.id.store(i, SeqCst); entries.push((addr, val)); @@ -734,7 +739,7 @@ mod test { for _ in 0..10 { // Drop 1000 in reverse - for _ in 0..1_000 { + for _ in 0..normal_or_miri(1_000, 10) { entries.pop(); } @@ -753,7 +758,7 @@ mod test { let mut entries1 = vec![]; let mut entries2 = vec![]; - for i in 0..10_000 { + for i in 0..normal_or_miri(10_000, 100) { let (addr, val) = alloc.allocate().unwrap(); val.id.store(i, SeqCst); @@ -771,6 +776,14 @@ mod test { } } + const fn normal_or_miri(normal: usize, miri: usize) -> usize { + if cfg!(miri) { + miri + } else { + normal + } + } + #[test] fn compact_all() { let mut slab = Slab::<Foo>::new(); @@ -780,7 +793,7 @@ mod test { for _ in 0..2 { entries.clear(); - for i in 0..10_000 { + for i in 0..normal_or_miri(10_000, 100) { let (addr, val) = alloc.allocate().unwrap(); val.id.store(i, SeqCst); @@ -808,7 +821,7 @@ mod test { let alloc = slab.allocator(); let mut entries = vec![]; - for _ in 0..5 { + for _ in 0..normal_or_miri(5, 2) { entries.clear(); // Allocate a few pages + 1 diff --git a/vendor/tokio/src/util/sync_wrapper.rs b/vendor/tokio/src/util/sync_wrapper.rs new file mode 100644 index 000000000..5ffc8f96b --- /dev/null +++ b/vendor/tokio/src/util/sync_wrapper.rs @@ -0,0 +1,26 @@ +//! This module contains a type that can make `Send + !Sync` types `Sync` by +//! disallowing all immutable access to the value. +//! +//! A similar primitive is provided in the `sync_wrapper` crate. + +pub(crate) struct SyncWrapper<T> { + value: T, +} + +// safety: The SyncWrapper being send allows you to send the inner value across +// thread boundaries. +unsafe impl<T: Send> Send for SyncWrapper<T> {} + +// safety: An immutable reference to a SyncWrapper is useless, so moving such an +// immutable reference across threads is safe. +unsafe impl<T> Sync for SyncWrapper<T> {} + +impl<T> SyncWrapper<T> { + pub(crate) fn new(value: T) -> Self { + Self { value } + } + + pub(crate) fn into_inner(self) -> T { + self.value + } +} diff --git a/vendor/tokio/src/util/trace.rs b/vendor/tokio/src/util/trace.rs index c51a5a72b..76e8a6cbf 100644 --- a/vendor/tokio/src/util/trace.rs +++ b/vendor/tokio/src/util/trace.rs @@ -1,37 +1,98 @@ cfg_trace! { cfg_rt! { + use core::{ + pin::Pin, + task::{Context, Poll}, + }; + use pin_project_lite::pin_project; + use std::future::Future; pub(crate) use tracing::instrument::Instrumented; #[inline] - #[cfg_attr(tokio_track_caller, track_caller)] - pub(crate) fn task<F>(task: F, kind: &'static str, name: Option<&str>) -> Instrumented<F> { + #[track_caller] + pub(crate) fn task<F>(task: F, kind: &'static str, name: Option<&str>, id: u64) -> Instrumented<F> { use tracing::instrument::Instrument; - #[cfg(tokio_track_caller)] let location = std::panic::Location::caller(); - #[cfg(tokio_track_caller)] let span = tracing::trace_span!( target: "tokio::task", - "task", + "runtime.spawn", %kind, - spawn.location = %format_args!("{}:{}:{}", location.file(), location.line(), location.column()), - task.name = %name.unwrap_or_default() - ); - #[cfg(not(tokio_track_caller))] - let span = tracing::trace_span!( - target: "tokio::task", - "task", - %kind, - task.name = %name.unwrap_or_default() + task.name = %name.unwrap_or_default(), + task.id = id, + loc.file = location.file(), + loc.line = location.line(), + loc.col = location.column(), ); task.instrument(span) } + + pub(crate) fn async_op<P,F>(inner: P, resource_span: tracing::Span, source: &str, poll_op_name: &'static str, inherits_child_attrs: bool) -> InstrumentedAsyncOp<F> + where P: FnOnce() -> F { + resource_span.in_scope(|| { + let async_op_span = tracing::trace_span!("runtime.resource.async_op", source = source, inherits_child_attrs = inherits_child_attrs); + let enter = async_op_span.enter(); + let async_op_poll_span = tracing::trace_span!("runtime.resource.async_op.poll"); + let inner = inner(); + drop(enter); + let tracing_ctx = AsyncOpTracingCtx { + async_op_span, + async_op_poll_span, + resource_span: resource_span.clone(), + }; + InstrumentedAsyncOp { + inner, + tracing_ctx, + poll_op_name, + } + }) + } + + #[derive(Debug, Clone)] + pub(crate) struct AsyncOpTracingCtx { + pub(crate) async_op_span: tracing::Span, + pub(crate) async_op_poll_span: tracing::Span, + pub(crate) resource_span: tracing::Span, + } + + + pin_project! { + #[derive(Debug, Clone)] + pub(crate) struct InstrumentedAsyncOp<F> { + #[pin] + pub(crate) inner: F, + pub(crate) tracing_ctx: AsyncOpTracingCtx, + pub(crate) poll_op_name: &'static str + } + } + + impl<F: Future> Future for InstrumentedAsyncOp<F> { + type Output = F::Output; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let this = self.project(); + let poll_op_name = &*this.poll_op_name; + let _res_enter = this.tracing_ctx.resource_span.enter(); + let _async_op_enter = this.tracing_ctx.async_op_span.enter(); + let _async_op_poll_enter = this.tracing_ctx.async_op_poll_span.enter(); + trace_poll_op!(poll_op_name, this.inner.poll(cx)) + } + } + } +} +cfg_time! { + #[track_caller] + pub(crate) fn caller_location() -> Option<&'static std::panic::Location<'static>> { + #[cfg(all(tokio_unstable, feature = "tracing"))] + return Some(std::panic::Location::caller()); + #[cfg(not(all(tokio_unstable, feature = "tracing")))] + None } } cfg_not_trace! { cfg_rt! { #[inline] - pub(crate) fn task<F>(task: F, _: &'static str, _name: Option<&str>) -> F { + pub(crate) fn task<F>(task: F, _: &'static str, _name: Option<&str>, _: u64) -> F { // nop task } diff --git a/vendor/tokio/src/util/wake.rs b/vendor/tokio/src/util/wake.rs index 57739371d..5526cbc63 100644 --- a/vendor/tokio/src/util/wake.rs +++ b/vendor/tokio/src/util/wake.rs @@ -1,15 +1,16 @@ +use crate::loom::sync::Arc; + use std::marker::PhantomData; use std::mem::ManuallyDrop; use std::ops::Deref; -use std::sync::Arc; use std::task::{RawWaker, RawWakerVTable, Waker}; -/// Simplified waking interface based on Arcs -pub(crate) trait Wake: Send + Sync { - /// Wake by value - fn wake(self: Arc<Self>); +/// Simplified waking interface based on Arcs. +pub(crate) trait Wake: Send + Sync + Sized + 'static { + /// Wake by value. + fn wake(arc_self: Arc<Self>); - /// Wake by reference + /// Wake by reference. fn wake_by_ref(arc_self: &Arc<Self>); } @@ -30,7 +31,7 @@ impl Deref for WakerRef<'_> { /// Creates a reference to a `Waker` from a reference to `Arc<impl Wake>`. pub(crate) fn waker_ref<W: Wake>(wake: &Arc<W>) -> WakerRef<'_> { - let ptr = &**wake as *const _ as *const (); + let ptr = Arc::as_ptr(wake) as *const (); let waker = unsafe { Waker::from_raw(RawWaker::new(ptr, waker_vtable::<W>())) }; diff --git a/vendor/tokio/src/util/wake_list.rs b/vendor/tokio/src/util/wake_list.rs new file mode 100644 index 000000000..aa569dd17 --- /dev/null +++ b/vendor/tokio/src/util/wake_list.rs @@ -0,0 +1,53 @@ +use core::mem::MaybeUninit; +use core::ptr; +use std::task::Waker; + +const NUM_WAKERS: usize = 32; + +pub(crate) struct WakeList { + inner: [MaybeUninit<Waker>; NUM_WAKERS], + curr: usize, +} + +impl WakeList { + pub(crate) fn new() -> Self { + Self { + inner: unsafe { + // safety: Create an uninitialized array of `MaybeUninit`. The + // `assume_init` is safe because the type we are claiming to + // have initialized here is a bunch of `MaybeUninit`s, which do + // not require initialization. + MaybeUninit::uninit().assume_init() + }, + curr: 0, + } + } + + #[inline] + pub(crate) fn can_push(&self) -> bool { + self.curr < NUM_WAKERS + } + + pub(crate) fn push(&mut self, val: Waker) { + debug_assert!(self.can_push()); + + self.inner[self.curr] = MaybeUninit::new(val); + self.curr += 1; + } + + pub(crate) fn wake_all(&mut self) { + assert!(self.curr <= NUM_WAKERS); + while self.curr > 0 { + self.curr -= 1; + let waker = unsafe { ptr::read(self.inner[self.curr].as_mut_ptr()) }; + waker.wake(); + } + } +} + +impl Drop for WakeList { + fn drop(&mut self) { + let slice = ptr::slice_from_raw_parts_mut(self.inner.as_mut_ptr() as *mut Waker, self.curr); + unsafe { ptr::drop_in_place(slice) }; + } +} |