diff options
Diffstat (limited to 'third_party/rust/tokio/src/io')
59 files changed, 11874 insertions, 0 deletions
diff --git a/third_party/rust/tokio/src/io/async_buf_read.rs b/third_party/rust/tokio/src/io/async_buf_read.rs new file mode 100644 index 0000000000..ecaafba4c2 --- /dev/null +++ b/third_party/rust/tokio/src/io/async_buf_read.rs @@ -0,0 +1,117 @@ +use crate::io::AsyncRead; + +use std::io; +use std::ops::DerefMut; +use std::pin::Pin; +use std::task::{Context, Poll}; + +/// Reads bytes asynchronously. +/// +/// This trait is analogous to [`std::io::BufRead`], but integrates with +/// the asynchronous task system. In particular, the [`poll_fill_buf`] method, +/// unlike [`BufRead::fill_buf`], will automatically queue the current task for wakeup +/// and return if data is not yet available, rather than blocking the calling +/// thread. +/// +/// Utilities for working with `AsyncBufRead` values are provided by +/// [`AsyncBufReadExt`]. +/// +/// [`std::io::BufRead`]: std::io::BufRead +/// [`poll_fill_buf`]: AsyncBufRead::poll_fill_buf +/// [`BufRead::fill_buf`]: std::io::BufRead::fill_buf +/// [`AsyncBufReadExt`]: crate::io::AsyncBufReadExt +pub trait AsyncBufRead: AsyncRead { + /// Attempts to return the contents of the internal buffer, filling it with more data + /// from the inner reader if it is empty. + /// + /// On success, returns `Poll::Ready(Ok(buf))`. + /// + /// If no data is available for reading, the method returns + /// `Poll::Pending` and arranges for the current task (via + /// `cx.waker().wake_by_ref()`) to receive a notification when the object becomes + /// readable or is closed. + /// + /// 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 [`poll_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. + /// + /// [`poll_read`]: AsyncRead::poll_read + /// [`consume`]: AsyncBufRead::consume + fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>>; + + /// Tells this buffer that `amt` bytes have been consumed from the buffer, + /// so they should no longer be returned in calls to [`poll_read`]. + /// + /// This function is a lower-level call. It needs to be paired with the + /// [`poll_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 [`poll_fill_buf`], has been consumed and should + /// no longer be returned. As such, this function may do odd things if + /// [`poll_fill_buf`] isn't called before calling it. + /// + /// The `amt` must be `<=` the number of bytes in the buffer returned by + /// [`poll_fill_buf`]. + /// + /// [`poll_read`]: AsyncRead::poll_read + /// [`poll_fill_buf`]: AsyncBufRead::poll_fill_buf + fn consume(self: Pin<&mut Self>, amt: usize); +} + +macro_rules! deref_async_buf_read { + () => { + fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> { + Pin::new(&mut **self.get_mut()).poll_fill_buf(cx) + } + + fn consume(mut self: Pin<&mut Self>, amt: usize) { + Pin::new(&mut **self).consume(amt) + } + }; +} + +impl<T: ?Sized + AsyncBufRead + Unpin> AsyncBufRead for Box<T> { + deref_async_buf_read!(); +} + +impl<T: ?Sized + AsyncBufRead + Unpin> AsyncBufRead for &mut T { + deref_async_buf_read!(); +} + +impl<P> AsyncBufRead for Pin<P> +where + P: DerefMut + Unpin, + P::Target: AsyncBufRead, +{ + fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> { + self.get_mut().as_mut().poll_fill_buf(cx) + } + + fn consume(self: Pin<&mut Self>, amt: usize) { + self.get_mut().as_mut().consume(amt) + } +} + +impl AsyncBufRead for &[u8] { + fn poll_fill_buf(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> { + Poll::Ready(Ok(*self)) + } + + fn consume(mut self: Pin<&mut Self>, amt: usize) { + *self = &self[amt..]; + } +} + +impl<T: AsRef<[u8]> + Unpin> AsyncBufRead for io::Cursor<T> { + fn poll_fill_buf(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> { + Poll::Ready(io::BufRead::fill_buf(self.get_mut())) + } + + fn consume(self: Pin<&mut Self>, amt: usize) { + io::BufRead::consume(self.get_mut(), amt) + } +} diff --git a/third_party/rust/tokio/src/io/async_fd.rs b/third_party/rust/tokio/src/io/async_fd.rs new file mode 100644 index 0000000000..93f9cb458a --- /dev/null +++ b/third_party/rust/tokio/src/io/async_fd.rs @@ -0,0 +1,660 @@ +use crate::io::driver::{Handle, Interest, ReadyEvent, Registration}; + +use mio::unix::SourceFd; +use std::io; +use std::os::unix::io::{AsRawFd, RawFd}; +use std::{task::Context, task::Poll}; + +/// Associates an IO object backed by a Unix file descriptor with the tokio +/// reactor, allowing for readiness to be polled. The file descriptor must be of +/// a type that can be used with the OS polling facilities (ie, `poll`, `epoll`, +/// `kqueue`, etc), such as a network socket or pipe, and the file descriptor +/// must have the nonblocking mode set to true. +/// +/// Creating an AsyncFd registers the file descriptor with the current tokio +/// Reactor, allowing you to directly await the file descriptor being readable +/// or writable. Once registered, the file descriptor remains registered until +/// the AsyncFd is dropped. +/// +/// The AsyncFd takes ownership of an arbitrary object to represent the IO +/// object. It is intended that this object will handle closing the file +/// descriptor when it is dropped, avoiding resource leaks and ensuring that the +/// AsyncFd can clean up the registration before closing the file descriptor. +/// The [`AsyncFd::into_inner`] function can be used to extract the inner object +/// to retake control from the tokio IO reactor. +/// +/// The inner object is required to implement [`AsRawFd`]. This file descriptor +/// must not change while [`AsyncFd`] owns the inner object, i.e. the +/// [`AsRawFd::as_raw_fd`] method on the inner type must always return the same +/// file descriptor when called multiple times. Failure to uphold this results +/// in unspecified behavior in the IO driver, which may include breaking +/// notifications for other sockets/etc. +/// +/// Polling for readiness is done by calling the async functions [`readable`] +/// and [`writable`]. These functions complete when the associated readiness +/// condition is observed. Any number of tasks can query the same `AsyncFd` in +/// parallel, on the same or different conditions. +/// +/// On some platforms, the readiness detecting mechanism relies on +/// edge-triggered notifications. This means that the OS will only notify Tokio +/// when the file descriptor transitions from not-ready to ready. For this to +/// work you should first try to read or write and only poll for readiness +/// if that fails with an error of [`std::io::ErrorKind::WouldBlock`]. +/// +/// Tokio internally tracks when it has received a ready notification, and when +/// readiness checking functions like [`readable`] and [`writable`] are called, +/// if the readiness flag is set, these async functions will complete +/// immediately. This however does mean that it is critical to ensure that this +/// ready flag is cleared when (and only when) the file descriptor ceases to be +/// ready. The [`AsyncFdReadyGuard`] returned from readiness checking functions +/// serves this function; after calling a readiness-checking async function, +/// you must use this [`AsyncFdReadyGuard`] to signal to tokio whether the file +/// descriptor is no longer in a ready state. +/// +/// ## Use with to a poll-based API +/// +/// In some cases it may be desirable to use `AsyncFd` from APIs similar to +/// [`TcpStream::poll_read_ready`]. The [`AsyncFd::poll_read_ready`] and +/// [`AsyncFd::poll_write_ready`] functions are provided for this purpose. +/// Because these functions don't create a future to hold their state, they have +/// the limitation that only one task can wait on each direction (read or write) +/// at a time. +/// +/// # 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. +/// +/// ```no_run +/// use futures::ready; +/// use std::io::{self, Read, Write}; +/// use std::net::TcpStream; +/// use std::pin::Pin; +/// use std::task::{Context, Poll}; +/// use tokio::io::AsyncWrite; +/// use tokio::io::unix::AsyncFd; +/// +/// pub struct AsyncTcpStream { +/// inner: AsyncFd<TcpStream>, +/// } +/// +/// impl AsyncTcpStream { +/// pub fn new(tcp: TcpStream) -> io::Result<Self> { +/// tcp.set_nonblocking(true)?; +/// Ok(Self { +/// inner: AsyncFd::new(tcp)?, +/// }) +/// } +/// +/// pub async fn read(&self, out: &mut [u8]) -> io::Result<usize> { +/// loop { +/// let mut guard = self.inner.readable().await?; +/// +/// match guard.try_io(|inner| inner.get_ref().read(out)) { +/// Ok(result) => return result, +/// Err(_would_block) => continue, +/// } +/// } +/// } +/// } +/// +/// impl AsyncWrite for AsyncTcpStream { +/// fn poll_write( +/// self: Pin<&mut Self>, +/// cx: &mut Context<'_>, +/// buf: &[u8] +/// ) -> Poll<io::Result<usize>> { +/// loop { +/// let mut guard = ready!(self.inner.poll_write_ready(cx))?; +/// +/// match guard.try_io(|inner| inner.get_ref().write(buf)) { +/// Ok(result) => return Poll::Ready(result), +/// Err(_would_block) => continue, +/// } +/// } +/// } +/// +/// fn poll_flush( +/// self: Pin<&mut Self>, +/// cx: &mut Context<'_>, +/// ) -> Poll<io::Result<()>> { +/// // tcp flush is a no-op +/// Poll::Ready(Ok(())) +/// } +/// +/// fn poll_shutdown( +/// self: Pin<&mut Self>, +/// cx: &mut Context<'_>, +/// ) -> Poll<io::Result<()>> { +/// self.inner.get_ref().shutdown(std::net::Shutdown::Write)?; +/// Poll::Ready(Ok(())) +/// } +/// } +/// ``` +/// +/// [`readable`]: method@Self::readable +/// [`writable`]: method@Self::writable +/// [`AsyncFdReadyGuard`]: struct@self::AsyncFdReadyGuard +/// [`TcpStream::poll_read_ready`]: struct@crate::net::TcpStream +pub struct AsyncFd<T: AsRawFd> { + registration: Registration, + inner: Option<T>, +} + +/// Represents an IO-ready event detected on a particular file descriptor that +/// has not yet been acknowledged. This is a `must_use` structure to help ensure +/// that you do not forget to explicitly clear (or not clear) the event. +/// +/// This type exposes an immutable reference to the underlying IO object. +#[must_use = "You must explicitly choose whether to clear the readiness state by calling a method on ReadyGuard"] +pub struct AsyncFdReadyGuard<'a, T: AsRawFd> { + async_fd: &'a AsyncFd<T>, + event: Option<ReadyEvent>, +} + +/// Represents an IO-ready event detected on a particular file descriptor that +/// has not yet been acknowledged. This is a `must_use` structure to help ensure +/// that you do not forget to explicitly clear (or not clear) the event. +/// +/// This type exposes a mutable reference to the underlying IO object. +#[must_use = "You must explicitly choose whether to clear the readiness state by calling a method on ReadyGuard"] +pub struct AsyncFdReadyMutGuard<'a, T: AsRawFd> { + async_fd: &'a mut AsyncFd<T>, + 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. + /// + /// This method must be called in the context of a tokio runtime. + pub fn new(inner: T) -> io::Result<Self> + where + T: AsRawFd, + { + Self::with_interest(inner, ALL_INTEREST) + } + + #[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. + pub fn with_interest(inner: T, interest: Interest) -> io::Result<Self> + where + T: AsRawFd, + { + Self::new_with_handle_and_interest(inner, Handle::current(), interest) + } + + pub(crate) fn new_with_handle_and_interest( + inner: T, + handle: Handle, + interest: Interest, + ) -> io::Result<Self> { + let fd = inner.as_raw_fd(); + + let registration = + Registration::new_with_interest_and_handle(&mut SourceFd(&fd), interest, handle)?; + + Ok(AsyncFd { + registration, + inner: Some(inner), + }) + } + + /// 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`]. + #[inline] + pub fn get_mut(&mut self) -> &mut T { + self.inner.as_mut().unwrap() + } + + fn take_inner(&mut self) -> Option<T> { + let fd = self.inner.as_ref().map(AsRawFd::as_raw_fd); + + if let Some(fd) = fd { + let _ = self.registration.deregister(&mut SourceFd(&fd)); + } + + self.inner.take() + } + + /// Deregisters this file descriptor and returns ownership of the backing + /// object. + pub fn into_inner(mut self) -> T { + self.take_inner().unwrap() + } + + /// Polls for read readiness. + /// + /// If the file descriptor is not currently ready for reading, this method + /// will store a clone of the [`Waker`] from the provided [`Context`]. When the + /// file descriptor becomes ready for reading, [`Waker::wake`] will be called. + /// + /// Note that on multiple calls to [`poll_read_ready`] or + /// [`poll_read_ready_mut`], only the `Waker` from the `Context` passed to the + /// most recent call is scheduled to receive a wakeup. (However, + /// [`poll_write_ready`] retains a second, independent waker). + /// + /// This method 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. + /// + /// 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`]. + /// + /// [`poll_read_ready`]: method@Self::poll_read_ready + /// [`poll_read_ready_mut`]: method@Self::poll_read_ready_mut + /// [`poll_write_ready`]: method@Self::poll_write_ready + /// [`readable`]: method@Self::readable + /// [`Context`]: struct@std::task::Context + /// [`Waker`]: struct@std::task::Waker + /// [`Waker::wake`]: method@std::task::Waker::wake + pub fn poll_read_ready<'a>( + &'a self, + cx: &mut Context<'_>, + ) -> Poll<io::Result<AsyncFdReadyGuard<'a, T>>> { + let event = ready!(self.registration.poll_read_ready(cx))?; + + Ok(AsyncFdReadyGuard { + async_fd: self, + event: Some(event), + }) + .into() + } + + /// Polls for read readiness. + /// + /// If the file descriptor is not currently ready for reading, this method + /// will store a clone of the [`Waker`] from the provided [`Context`]. When the + /// file descriptor becomes ready for reading, [`Waker::wake`] will be called. + /// + /// Note that on multiple calls to [`poll_read_ready`] or + /// [`poll_read_ready_mut`], only the `Waker` from the `Context` passed to the + /// most recent call is scheduled to receive a wakeup. (However, + /// [`poll_write_ready`] retains a second, independent waker). + /// + /// This method 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. + /// + /// This method takes `&mut self`, so it is possible to access the inner IO + /// resource mutably when handling the [`AsyncFdReadyMutGuard`]. + /// + /// [`poll_read_ready`]: method@Self::poll_read_ready + /// [`poll_read_ready_mut`]: method@Self::poll_read_ready_mut + /// [`poll_write_ready`]: method@Self::poll_write_ready + /// [`readable`]: method@Self::readable + /// [`Context`]: struct@std::task::Context + /// [`Waker`]: struct@std::task::Waker + /// [`Waker::wake`]: method@std::task::Waker::wake + pub fn poll_read_ready_mut<'a>( + &'a mut self, + cx: &mut Context<'_>, + ) -> Poll<io::Result<AsyncFdReadyMutGuard<'a, T>>> { + let event = ready!(self.registration.poll_read_ready(cx))?; + + Ok(AsyncFdReadyMutGuard { + async_fd: self, + event: Some(event), + }) + .into() + } + + /// Polls for write readiness. + /// + /// If the file descriptor is not currently ready for writing, this method + /// will store a clone of the [`Waker`] from the provided [`Context`]. When the + /// file descriptor becomes ready for writing, [`Waker::wake`] will be called. + /// + /// Note that on multiple calls to [`poll_write_ready`] or + /// [`poll_write_ready_mut`], only the `Waker` from the `Context` passed to the + /// most recent call is scheduled to receive a wakeup. (However, + /// [`poll_read_ready`] retains a second, independent waker). + /// + /// This method 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. + /// + /// 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`]. + /// + /// [`poll_read_ready`]: method@Self::poll_read_ready + /// [`poll_write_ready`]: method@Self::poll_write_ready + /// [`poll_write_ready_mut`]: method@Self::poll_write_ready_mut + /// [`writable`]: method@Self::readable + /// [`Context`]: struct@std::task::Context + /// [`Waker`]: struct@std::task::Waker + /// [`Waker::wake`]: method@std::task::Waker::wake + pub fn poll_write_ready<'a>( + &'a self, + cx: &mut Context<'_>, + ) -> Poll<io::Result<AsyncFdReadyGuard<'a, T>>> { + let event = ready!(self.registration.poll_write_ready(cx))?; + + Ok(AsyncFdReadyGuard { + async_fd: self, + event: Some(event), + }) + .into() + } + + /// Polls for write readiness. + /// + /// If the file descriptor is not currently ready for writing, this method + /// will store a clone of the [`Waker`] from the provided [`Context`]. When the + /// file descriptor becomes ready for writing, [`Waker::wake`] will be called. + /// + /// Note that on multiple calls to [`poll_write_ready`] or + /// [`poll_write_ready_mut`], only the `Waker` from the `Context` passed to the + /// most recent call is scheduled to receive a wakeup. (However, + /// [`poll_read_ready`] retains a second, independent waker). + /// + /// This method 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. + /// + /// This method takes `&mut self`, so it is possible to access the inner IO + /// resource mutably when handling the [`AsyncFdReadyMutGuard`]. + /// + /// [`poll_read_ready`]: method@Self::poll_read_ready + /// [`poll_write_ready`]: method@Self::poll_write_ready + /// [`poll_write_ready_mut`]: method@Self::poll_write_ready_mut + /// [`writable`]: method@Self::readable + /// [`Context`]: struct@std::task::Context + /// [`Waker`]: struct@std::task::Waker + /// [`Waker::wake`]: method@std::task::Waker::wake + pub fn poll_write_ready_mut<'a>( + &'a mut self, + cx: &mut Context<'_>, + ) -> Poll<io::Result<AsyncFdReadyMutGuard<'a, T>>> { + let event = ready!(self.registration.poll_write_ready(cx))?; + + Ok(AsyncFdReadyMutGuard { + async_fd: self, + event: Some(event), + }) + .into() + } + + async fn readiness(&self, interest: Interest) -> io::Result<AsyncFdReadyGuard<'_, T>> { + let event = self.registration.readiness(interest).await?; + + Ok(AsyncFdReadyGuard { + async_fd: self, + event: Some(event), + }) + } + + async fn readiness_mut( + &mut self, + interest: Interest, + ) -> io::Result<AsyncFdReadyMutGuard<'_, T>> { + let event = self.registration.readiness(interest).await?; + + Ok(AsyncFdReadyMutGuard { + async_fd: self, + event: Some(event), + }) + } + + /// Waits for the file descriptor to become readable, returning a + /// [`AsyncFdReadyGuard`] that must be dropped to resume read-readiness + /// polling. + /// + /// 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`]. + #[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 + } + + /// Waits for the file descriptor to become readable, returning a + /// [`AsyncFdReadyMutGuard`] that must be dropped to resume read-readiness + /// polling. + /// + /// This method takes `&mut self`, so it is possible to access the inner IO + /// 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 + } + + /// Waits for the file descriptor to become writable, returning a + /// [`AsyncFdReadyGuard`] that must be dropped to resume write-readiness + /// polling. + /// + /// 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`]. + #[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 + } + + /// Waits for the file descriptor to become writable, returning a + /// [`AsyncFdReadyMutGuard`] that must be dropped to resume write-readiness + /// polling. + /// + /// This method takes `&mut self`, so it is possible to access the inner IO + /// 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 + } +} + +impl<T: AsRawFd> AsRawFd for AsyncFd<T> { + fn as_raw_fd(&self) -> RawFd { + self.inner.as_ref().unwrap().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") + .field("inner", &self.inner) + .finish() + } +} + +impl<T: AsRawFd> Drop for AsyncFd<T> { + fn drop(&mut self) { + let _ = self.take_inner(); + } +} + +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 + /// next edge-triggered readiness notification from the OS. + /// + /// 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); + } + } + + /// This method should be invoked when you intentionally want to keep the + /// ready flag asserted. + /// + /// While this function is itself a no-op, it satisfies the `#[must_use]` + /// constraint on the [`AsyncFdReadyGuard`] type. + pub fn retain_ready(&mut self) { + // no-op + } + + /// Performs the provided IO operation. + /// + /// If `f` returns a [`WouldBlock`] error, the readiness state associated + /// with this file descriptor is cleared, and the method returns + /// `Err(TryIoError::WouldBlock)`. You will typically need to poll the + /// `AsyncFd` again when this happens. + /// + /// This method helps ensure that the readiness state of the underlying file + /// descriptor remains in sync with the tokio-side readiness state, by + /// clearing the tokio-side state only when a [`WouldBlock`] condition + /// occurs. It is the responsibility of the caller to ensure that `f` + /// returns [`WouldBlock`] only if the file descriptor that originated this + /// `AsyncFdReadyGuard` no longer expresses the readiness state that was queried to + /// create this `AsyncFdReadyGuard`. + /// + /// [`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(&'a AsyncFd<Inner>) -> io::Result<R>, + ) -> Result<io::Result<R>, TryIoError> { + let result = f(self.async_fd); + + if let Err(e) = result.as_ref() { + if e.kind() == io::ErrorKind::WouldBlock { + self.clear_ready(); + } + } + + match result { + Err(err) if err.kind() == io::ErrorKind::WouldBlock => Err(TryIoError(())), + result => Ok(result), + } + } + + /// Returns a shared reference to the inner [`AsyncFd`]. + 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) -> &'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 + /// next edge-triggered readiness notification from the OS. + /// + /// 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); + } + } + + /// This method should be invoked when you intentionally want to keep the + /// ready flag asserted. + /// + /// While this function is itself a no-op, it satisfies the `#[must_use]` + /// constraint on the [`AsyncFdReadyGuard`] type. + pub fn retain_ready(&mut self) { + // no-op + } + + /// Performs the provided IO operation. + /// + /// If `f` returns a [`WouldBlock`] error, the readiness state associated + /// with this file descriptor is cleared, and the method returns + /// `Err(TryIoError::WouldBlock)`. You will typically need to poll the + /// `AsyncFd` again when this happens. + /// + /// This method helps ensure that the readiness state of the underlying file + /// descriptor remains in sync with the tokio-side readiness state, by + /// clearing the tokio-side state only when a [`WouldBlock`] condition + /// occurs. It is the responsibility of the caller to ensure that `f` + /// returns [`WouldBlock`] only if the file descriptor that originated this + /// `AsyncFdReadyGuard` no longer expresses the readiness state that was queried to + /// create this `AsyncFdReadyGuard`. + /// + /// [`WouldBlock`]: std::io::ErrorKind::WouldBlock + pub fn try_io<R>( + &mut self, + f: impl FnOnce(&mut AsyncFd<Inner>) -> io::Result<R>, + ) -> Result<io::Result<R>, TryIoError> { + let result = f(self.async_fd); + + if let Err(e) = result.as_ref() { + if e.kind() == io::ErrorKind::WouldBlock { + self.clear_ready(); + } + } + + match result { + Err(err) if err.kind() == io::ErrorKind::WouldBlock => Err(TryIoError(())), + result => Ok(result), + } + } + + /// Returns a shared reference to the inner [`AsyncFd`]. + pub fn get_ref(&self) -> &AsyncFd<Inner> { + self.async_fd + } + + /// Returns a mutable reference to the inner [`AsyncFd`]. + pub fn get_mut(&mut self) -> &mut AsyncFd<Inner> { + self.async_fd + } + + /// Returns a shared reference to the backing object of the inner [`AsyncFd`]. + pub fn get_inner(&self) -> &Inner { + self.get_ref().get_ref() + } + + /// Returns a mutable reference to the backing object of the inner [`AsyncFd`]. + pub fn get_inner_mut(&mut self) -> &mut Inner { + self.get_mut().get_mut() + } +} + +impl<'a, T: std::fmt::Debug + AsRawFd> std::fmt::Debug for AsyncFdReadyGuard<'a, T> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ReadyGuard") + .field("async_fd", &self.async_fd) + .finish() + } +} + +impl<'a, T: std::fmt::Debug + AsRawFd> std::fmt::Debug for AsyncFdReadyMutGuard<'a, T> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("MutReadyGuard") + .field("async_fd", &self.async_fd) + .finish() + } +} + +/// The error type returned by [`try_io`]. +/// +/// This error indicates that the IO resource returned a [`WouldBlock`] error. +/// +/// [`WouldBlock`]: std::io::ErrorKind::WouldBlock +/// [`try_io`]: method@AsyncFdReadyGuard::try_io +#[derive(Debug)] +pub struct TryIoError(()); diff --git a/third_party/rust/tokio/src/io/async_read.rs b/third_party/rust/tokio/src/io/async_read.rs new file mode 100644 index 0000000000..93e5d3e66e --- /dev/null +++ b/third_party/rust/tokio/src/io/async_read.rs @@ -0,0 +1,131 @@ +use super::ReadBuf; +use std::io; +use std::ops::DerefMut; +use std::pin::Pin; +use std::task::{Context, Poll}; + +/// Reads bytes from a source. +/// +/// This trait is analogous to the [`std::io::Read`] trait, but integrates with +/// the asynchronous task system. In particular, the [`poll_read`] method, +/// unlike [`Read::read`], will automatically queue the current task for wakeup +/// and return if data is not yet available, rather than blocking the calling +/// thread. +/// +/// Specifically, this means that the `poll_read` function will return one of +/// the following: +/// +/// * `Poll::Ready(Ok(()))` means that data was immediately read and placed into +/// the output buffer. The amount of data read can be determined by the +/// increase in the length of the slice returned by `ReadBuf::filled`. If the +/// difference is 0, EOF has been reached. +/// +/// * `Poll::Pending` means that no data was read into the buffer +/// provided. The I/O object is not currently readable but may become readable +/// in the future. Most importantly, **the current future's task is scheduled +/// to get unparked when the object is readable**. This means that like +/// `Future::poll` you'll receive a notification when the I/O object is +/// readable again. +/// +/// * `Poll::Ready(Err(e))` for other errors are standard I/O errors coming from the +/// underlying object. +/// +/// This trait importantly means that the `read` method only works in the +/// context of a future's task. The object may panic if used outside of a task. +/// +/// Utilities for working with `AsyncRead` values are provided by +/// [`AsyncReadExt`]. +/// +/// [`poll_read`]: AsyncRead::poll_read +/// [`std::io::Read`]: std::io::Read +/// [`Read::read`]: std::io::Read::read +/// [`AsyncReadExt`]: crate::io::AsyncReadExt +pub trait AsyncRead { + /// Attempts to read from the `AsyncRead` into `buf`. + /// + /// On success, returns `Poll::Ready(Ok(()))` and places data in the + /// unfilled portion of `buf`. If no data was read (`buf.filled().len()` is + /// unchanged), it implies that EOF has been reached. + /// + /// If no data is available for reading, the method returns `Poll::Pending` + /// and arranges for the current task (via `cx.waker()`) to receive a + /// notification when the object becomes readable or is closed. + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<()>>; +} + +macro_rules! deref_async_read { + () => { + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<()>> { + Pin::new(&mut **self).poll_read(cx, buf) + } + }; +} + +impl<T: ?Sized + AsyncRead + Unpin> AsyncRead for Box<T> { + deref_async_read!(); +} + +impl<T: ?Sized + AsyncRead + Unpin> AsyncRead for &mut T { + deref_async_read!(); +} + +impl<P> AsyncRead for Pin<P> +where + P: DerefMut + Unpin, + P::Target: AsyncRead, +{ + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<()>> { + self.get_mut().as_mut().poll_read(cx, buf) + } +} + +impl AsyncRead for &[u8] { + fn poll_read( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<()>> { + let amt = std::cmp::min(self.len(), buf.remaining()); + let (a, b) = self.split_at(amt); + buf.put_slice(a); + *self = b; + Poll::Ready(Ok(())) + } +} + +impl<T: AsRef<[u8]> + Unpin> AsyncRead for io::Cursor<T> { + fn poll_read( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<()>> { + let pos = self.position(); + let slice: &[u8] = (*self).get_ref().as_ref(); + + // The position could technically be out of bounds, so don't panic... + if pos > slice.len() as u64 { + return Poll::Ready(Ok(())); + } + + let start = pos as usize; + let amt = std::cmp::min(slice.len() - start, buf.remaining()); + // Add won't overflow because of pos check above. + let end = start + amt; + buf.put_slice(&slice[start..end]); + self.set_position(end as u64); + + Poll::Ready(Ok(())) + } +} diff --git a/third_party/rust/tokio/src/io/async_seek.rs b/third_party/rust/tokio/src/io/async_seek.rs new file mode 100644 index 0000000000..bd7a992e4d --- /dev/null +++ b/third_party/rust/tokio/src/io/async_seek.rs @@ -0,0 +1,90 @@ +use std::io::{self, SeekFrom}; +use std::ops::DerefMut; +use std::pin::Pin; +use std::task::{Context, Poll}; + +/// Seek bytes asynchronously. +/// +/// This trait is analogous to the [`std::io::Seek`] trait, but integrates +/// with the asynchronous task system. In particular, the `start_seek` +/// method, unlike [`Seek::seek`], will not block the calling thread. +/// +/// Utilities for working with `AsyncSeek` values are provided by +/// [`AsyncSeekExt`]. +/// +/// [`std::io::Seek`]: std::io::Seek +/// [`Seek::seek`]: std::io::Seek::seek() +/// [`AsyncSeekExt`]: crate::io::AsyncSeekExt +pub trait AsyncSeek { + /// Attempts to seek to an offset, in bytes, in a stream. + /// + /// A seek beyond the end of a stream is allowed, but behavior is defined + /// by the implementation. + /// + /// If this function returns successfully, then the job has been submitted. + /// To find out when it completes, call `poll_complete`. + /// + /// # Errors + /// + /// This function can return [`io::ErrorKind::Other`] in case there is + /// another seek in progress. To avoid this, it is advisable that any call + /// to `start_seek` is preceded by a call to `poll_complete` to ensure all + /// pending seeks have completed. + fn start_seek(self: Pin<&mut Self>, position: SeekFrom) -> io::Result<()>; + + /// Waits for a seek operation to complete. + /// + /// If the seek operation completed successfully, + /// this method returns the new position from the start of the stream. + /// That position can be used later with [`SeekFrom::Start`]. Repeatedly + /// calling this function without calling `start_seek` might return the + /// same result. + /// + /// # Errors + /// + /// Seeking to a negative offset is considered an error. + fn poll_complete(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>>; +} + +macro_rules! deref_async_seek { + () => { + fn start_seek(mut self: Pin<&mut Self>, pos: SeekFrom) -> io::Result<()> { + Pin::new(&mut **self).start_seek(pos) + } + + fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> { + Pin::new(&mut **self).poll_complete(cx) + } + }; +} + +impl<T: ?Sized + AsyncSeek + Unpin> AsyncSeek for Box<T> { + deref_async_seek!(); +} + +impl<T: ?Sized + AsyncSeek + Unpin> AsyncSeek for &mut T { + deref_async_seek!(); +} + +impl<P> AsyncSeek for Pin<P> +where + P: DerefMut + Unpin, + P::Target: AsyncSeek, +{ + fn start_seek(self: Pin<&mut Self>, pos: SeekFrom) -> io::Result<()> { + self.get_mut().as_mut().start_seek(pos) + } + + fn poll_complete(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> { + self.get_mut().as_mut().poll_complete(cx) + } +} + +impl<T: AsRef<[u8]> + Unpin> AsyncSeek for io::Cursor<T> { + fn start_seek(mut self: Pin<&mut Self>, pos: SeekFrom) -> io::Result<()> { + io::Seek::seek(&mut *self, pos).map(drop) + } + fn poll_complete(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<u64>> { + Poll::Ready(Ok(self.get_mut().position())) + } +} diff --git a/third_party/rust/tokio/src/io/async_write.rs b/third_party/rust/tokio/src/io/async_write.rs new file mode 100644 index 0000000000..7ec1a302ef --- /dev/null +++ b/third_party/rust/tokio/src/io/async_write.rs @@ -0,0 +1,408 @@ +use std::io::{self, IoSlice}; +use std::ops::DerefMut; +use std::pin::Pin; +use std::task::{Context, Poll}; + +/// Writes bytes asynchronously. +/// +/// The trait inherits from [`std::io::Write`] and indicates that an I/O object is +/// **nonblocking**. All non-blocking I/O objects must return an error when +/// bytes cannot be written instead of blocking the current thread. +/// +/// Specifically, this means that the [`poll_write`] function will return one of +/// the following: +/// +/// * `Poll::Ready(Ok(n))` means that `n` bytes of data was immediately +/// written. +/// +/// * `Poll::Pending` means that no data was written from the buffer +/// provided. The I/O object is not currently writable but may become writable +/// in the future. Most importantly, **the current future's task is scheduled +/// to get unparked when the object is writable**. This means that like +/// `Future::poll` you'll receive a notification when the I/O object is +/// writable again. +/// +/// * `Poll::Ready(Err(e))` for other errors are standard I/O errors coming from the +/// underlying object. +/// +/// This trait importantly means that the [`write`][stdwrite] method only works in +/// the context of a future's task. The object may panic if used outside of a task. +/// +/// Note that this trait also represents that the [`Write::flush`][stdflush] method +/// works very similarly to the `write` method, notably that `Ok(())` means that the +/// writer has successfully been flushed, a "would block" error means that the +/// current task is ready to receive a notification when flushing can make more +/// progress, and otherwise normal errors can happen as well. +/// +/// Utilities for working with `AsyncWrite` values are provided by +/// [`AsyncWriteExt`]. +/// +/// [`std::io::Write`]: std::io::Write +/// [`poll_write`]: AsyncWrite::poll_write() +/// [stdwrite]: std::io::Write::write() +/// [stdflush]: std::io::Write::flush() +/// [`AsyncWriteExt`]: crate::io::AsyncWriteExt +pub trait AsyncWrite { + /// Attempt to write bytes from `buf` into the object. + /// + /// On success, returns `Poll::Ready(Ok(num_bytes_written))`. If successful, + /// then it must be guaranteed that `n <= buf.len()`. A return value of `0` + /// typically means that the underlying object is no longer able to accept + /// bytes and will likely not be able to in the future as well, or that the + /// buffer provided is empty. + /// + /// If the object is not ready for writing, the method returns + /// `Poll::Pending` and arranges for the current task (via + /// `cx.waker()`) to receive a notification when the object becomes + /// writable or is closed. + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<Result<usize, io::Error>>; + + /// Attempts to flush the object, ensuring that any buffered data reach + /// their destination. + /// + /// On success, returns `Poll::Ready(Ok(()))`. + /// + /// If flushing cannot immediately complete, this method returns + /// `Poll::Pending` and arranges for the current task (via + /// `cx.waker()`) to receive a notification when the object can make + /// progress towards flushing. + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>>; + + /// Initiates or attempts to shut down this writer, returning success when + /// the I/O connection has completely shut down. + /// + /// This method is intended to be used for asynchronous shutdown of I/O + /// connections. For example this is suitable for implementing shutdown of a + /// TLS connection or calling `TcpStream::shutdown` on a proxied connection. + /// Protocols sometimes need to flush out final pieces of data or otherwise + /// perform a graceful shutdown handshake, reading/writing more data as + /// appropriate. This method is the hook for such protocols to implement the + /// graceful shutdown logic. + /// + /// This `shutdown` method is required by implementers of the + /// `AsyncWrite` trait. Wrappers typically just want to proxy this call + /// through to the wrapped type, and base types will typically implement + /// shutdown logic here or just return `Ok(().into())`. Note that if you're + /// wrapping an underlying `AsyncWrite` a call to `shutdown` implies that + /// transitively the entire stream has been shut down. After your wrapper's + /// shutdown logic has been executed you should shut down the underlying + /// stream. + /// + /// Invocation of a `shutdown` implies an invocation of `flush`. Once this + /// method returns `Ready` it implies that a flush successfully happened + /// before the shutdown happened. That is, callers don't need to call + /// `flush` before calling `shutdown`. They can rely that by calling + /// `shutdown` any pending buffered data will be written out. + /// + /// # Return value + /// + /// This function returns a `Poll<io::Result<()>>` classified as such: + /// + /// * `Poll::Ready(Ok(()))` - indicates that the connection was + /// successfully shut down and is now safe to deallocate/drop/close + /// resources associated with it. This method means that the current task + /// will no longer receive any notifications due to this method and the + /// I/O object itself is likely no longer usable. + /// + /// * `Poll::Pending` - indicates that shutdown is initiated but could + /// not complete just yet. This may mean that more I/O needs to happen to + /// continue this shutdown operation. The current task is scheduled to + /// receive a notification when it's otherwise ready to continue the + /// shutdown operation. When woken up this method should be called again. + /// + /// * `Poll::Ready(Err(e))` - indicates a fatal error has happened with shutdown, + /// indicating that the shutdown operation did not complete successfully. + /// This typically means that the I/O object is no longer usable. + /// + /// # Errors + /// + /// This function can return normal I/O errors through `Err`, described + /// above. Additionally this method may also render the underlying + /// `Write::write` method no longer usable (e.g. will return errors in the + /// future). It's recommended that once `shutdown` is called the + /// `write` method is no longer called. + /// + /// # Panics + /// + /// This function will panic if not called within the context of a future's + /// task. + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>>; + + /// Like [`poll_write`], except that it writes from a slice of buffers. + /// + /// Data is copied from each buffer in order, with the final buffer + /// read from possibly being only partially consumed. This method must + /// behave as a call to [`write`] with the buffers concatenated would. + /// + /// The default implementation calls [`poll_write`] with either the first nonempty + /// buffer provided, or an empty one if none exists. + /// + /// On success, returns `Poll::Ready(Ok(num_bytes_written))`. + /// + /// If the object is not ready for writing, the method returns + /// `Poll::Pending` and arranges for the current task (via + /// `cx.waker()`) to receive a notification when the object becomes + /// writable or is closed. + /// + /// # Note + /// + /// This should be implemented as a single "atomic" write action. If any + /// data has been partially written, it is wrong to return an error or + /// pending. + /// + /// [`poll_write`]: AsyncWrite::poll_write + fn poll_write_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[IoSlice<'_>], + ) -> Poll<Result<usize, io::Error>> { + let buf = bufs + .iter() + .find(|b| !b.is_empty()) + .map_or(&[][..], |b| &**b); + self.poll_write(cx, buf) + } + + /// Determines if this writer has an efficient [`poll_write_vectored`] + /// implementation. + /// + /// If a writer does not override the default [`poll_write_vectored`] + /// implementation, code using it may want to avoid the method all together + /// and coalesce writes into a single buffer for higher performance. + /// + /// The default implementation returns `false`. + /// + /// [`poll_write_vectored`]: AsyncWrite::poll_write_vectored + fn is_write_vectored(&self) -> bool { + false + } +} + +macro_rules! deref_async_write { + () => { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + Pin::new(&mut **self).poll_write(cx, buf) + } + + fn poll_write_vectored( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[IoSlice<'_>], + ) -> Poll<io::Result<usize>> { + Pin::new(&mut **self).poll_write_vectored(cx, bufs) + } + + fn is_write_vectored(&self) -> bool { + (**self).is_write_vectored() + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + Pin::new(&mut **self).poll_flush(cx) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + Pin::new(&mut **self).poll_shutdown(cx) + } + }; +} + +impl<T: ?Sized + AsyncWrite + Unpin> AsyncWrite for Box<T> { + deref_async_write!(); +} + +impl<T: ?Sized + AsyncWrite + Unpin> AsyncWrite for &mut T { + deref_async_write!(); +} + +impl<P> AsyncWrite for Pin<P> +where + P: DerefMut + Unpin, + P::Target: AsyncWrite, +{ + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + self.get_mut().as_mut().poll_write(cx, buf) + } + + fn poll_write_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[IoSlice<'_>], + ) -> Poll<io::Result<usize>> { + self.get_mut().as_mut().poll_write_vectored(cx, bufs) + } + + fn is_write_vectored(&self) -> bool { + (**self).is_write_vectored() + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + self.get_mut().as_mut().poll_flush(cx) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + self.get_mut().as_mut().poll_shutdown(cx) + } +} + +impl AsyncWrite for Vec<u8> { + fn poll_write( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + self.get_mut().extend_from_slice(buf); + Poll::Ready(Ok(buf.len())) + } + + fn poll_write_vectored( + mut self: Pin<&mut Self>, + _: &mut Context<'_>, + bufs: &[IoSlice<'_>], + ) -> Poll<io::Result<usize>> { + Poll::Ready(io::Write::write_vectored(&mut *self, bufs)) + } + + fn is_write_vectored(&self) -> bool { + true + } + + 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(())) + } +} + +impl AsyncWrite for io::Cursor<&mut [u8]> { + fn poll_write( + mut self: Pin<&mut Self>, + _: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + Poll::Ready(io::Write::write(&mut *self, buf)) + } + + fn poll_write_vectored( + mut self: Pin<&mut Self>, + _: &mut Context<'_>, + bufs: &[IoSlice<'_>], + ) -> Poll<io::Result<usize>> { + Poll::Ready(io::Write::write_vectored(&mut *self, bufs)) + } + + fn is_write_vectored(&self) -> bool { + true + } + + fn poll_flush(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> { + Poll::Ready(io::Write::flush(&mut *self)) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + self.poll_flush(cx) + } +} + +impl AsyncWrite for io::Cursor<&mut Vec<u8>> { + fn poll_write( + mut self: Pin<&mut Self>, + _: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + Poll::Ready(io::Write::write(&mut *self, buf)) + } + + fn poll_write_vectored( + mut self: Pin<&mut Self>, + _: &mut Context<'_>, + bufs: &[IoSlice<'_>], + ) -> Poll<io::Result<usize>> { + Poll::Ready(io::Write::write_vectored(&mut *self, bufs)) + } + + fn is_write_vectored(&self) -> bool { + true + } + + fn poll_flush(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> { + Poll::Ready(io::Write::flush(&mut *self)) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + self.poll_flush(cx) + } +} + +impl AsyncWrite for io::Cursor<Vec<u8>> { + fn poll_write( + mut self: Pin<&mut Self>, + _: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + Poll::Ready(io::Write::write(&mut *self, buf)) + } + + fn poll_write_vectored( + mut self: Pin<&mut Self>, + _: &mut Context<'_>, + bufs: &[IoSlice<'_>], + ) -> Poll<io::Result<usize>> { + Poll::Ready(io::Write::write_vectored(&mut *self, bufs)) + } + + fn is_write_vectored(&self) -> bool { + true + } + + fn poll_flush(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> { + Poll::Ready(io::Write::flush(&mut *self)) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + self.poll_flush(cx) + } +} + +impl AsyncWrite for io::Cursor<Box<[u8]>> { + fn poll_write( + mut self: Pin<&mut Self>, + _: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + Poll::Ready(io::Write::write(&mut *self, buf)) + } + + fn poll_write_vectored( + mut self: Pin<&mut Self>, + _: &mut Context<'_>, + bufs: &[IoSlice<'_>], + ) -> Poll<io::Result<usize>> { + Poll::Ready(io::Write::write_vectored(&mut *self, bufs)) + } + + fn is_write_vectored(&self) -> bool { + true + } + + fn poll_flush(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> { + Poll::Ready(io::Write::flush(&mut *self)) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + self.poll_flush(cx) + } +} diff --git a/third_party/rust/tokio/src/io/blocking.rs b/third_party/rust/tokio/src/io/blocking.rs new file mode 100644 index 0000000000..1d79ee7a27 --- /dev/null +++ b/third_party/rust/tokio/src/io/blocking.rs @@ -0,0 +1,279 @@ +use crate::io::sys; +use crate::io::{AsyncRead, AsyncWrite, ReadBuf}; + +use std::cmp; +use std::future::Future; +use std::io; +use std::io::prelude::*; +use std::pin::Pin; +use std::task::Poll::*; +use std::task::{Context, Poll}; + +use self::State::*; + +/// `T` should not implement _both_ Read and Write. +#[derive(Debug)] +pub(crate) struct Blocking<T> { + inner: Option<T>, + state: State<T>, + /// `true` if the lower IO layer needs flushing. + need_flush: bool, +} + +#[derive(Debug)] +pub(crate) struct Buf { + buf: Vec<u8>, + pos: usize, +} + +pub(crate) const MAX_BUF: usize = 16 * 1024; + +#[derive(Debug)] +enum State<T> { + Idle(Option<Buf>), + Busy(sys::Blocking<(io::Result<usize>, Buf, T)>), +} + +cfg_io_std! { + impl<T> Blocking<T> { + pub(crate) fn new(inner: T) -> Blocking<T> { + Blocking { + inner: Some(inner), + state: State::Idle(Some(Buf::with_capacity(0))), + need_flush: false, + } + } + } +} + +impl<T> AsyncRead for Blocking<T> +where + T: Read + Unpin + Send + 'static, +{ + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + dst: &mut ReadBuf<'_>, + ) -> Poll<io::Result<()>> { + loop { + match self.state { + Idle(ref mut buf_cell) => { + let mut buf = buf_cell.take().unwrap(); + + if !buf.is_empty() { + buf.copy_to(dst); + *buf_cell = Some(buf); + return Ready(Ok(())); + } + + buf.ensure_capacity_for(dst); + let mut inner = self.inner.take().unwrap(); + + self.state = Busy(sys::run(move || { + let res = buf.read_from(&mut inner); + (res, buf, inner) + })); + } + Busy(ref mut rx) => { + let (res, mut buf, inner) = ready!(Pin::new(rx).poll(cx))?; + self.inner = Some(inner); + + match res { + Ok(_) => { + buf.copy_to(dst); + self.state = Idle(Some(buf)); + return Ready(Ok(())); + } + Err(e) => { + assert!(buf.is_empty()); + + self.state = Idle(Some(buf)); + return Ready(Err(e)); + } + } + } + } + } + } +} + +impl<T> AsyncWrite for Blocking<T> +where + T: Write + Unpin + Send + 'static, +{ + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + src: &[u8], + ) -> Poll<io::Result<usize>> { + loop { + match self.state { + Idle(ref mut buf_cell) => { + let mut buf = buf_cell.take().unwrap(); + + assert!(buf.is_empty()); + + let n = buf.copy_from(src); + let mut inner = self.inner.take().unwrap(); + + self.state = Busy(sys::run(move || { + let n = buf.len(); + let res = buf.write_to(&mut inner).map(|_| n); + + (res, buf, inner) + })); + self.need_flush = true; + + return Ready(Ok(n)); + } + Busy(ref mut rx) => { + let (res, buf, inner) = ready!(Pin::new(rx).poll(cx))?; + self.state = Idle(Some(buf)); + self.inner = Some(inner); + + // If error, return + res?; + } + } + } + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> { + loop { + let need_flush = self.need_flush; + match self.state { + // The buffer is not used here + Idle(ref mut buf_cell) => { + if need_flush { + let buf = buf_cell.take().unwrap(); + let mut inner = self.inner.take().unwrap(); + + self.state = Busy(sys::run(move || { + let res = inner.flush().map(|_| 0); + (res, buf, inner) + })); + + self.need_flush = false; + } else { + return Ready(Ok(())); + } + } + Busy(ref mut rx) => { + let (res, buf, inner) = ready!(Pin::new(rx).poll(cx))?; + self.state = Idle(Some(buf)); + self.inner = Some(inner); + + // If error, return + res?; + } + } + } + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> { + Poll::Ready(Ok(())) + } +} + +/// Repeats operations that are interrupted. +macro_rules! uninterruptibly { + ($e:expr) => {{ + loop { + match $e { + Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {} + res => break res, + } + } + }}; +} + +impl Buf { + pub(crate) fn with_capacity(n: usize) -> Buf { + Buf { + buf: Vec::with_capacity(n), + pos: 0, + } + } + + pub(crate) fn is_empty(&self) -> bool { + self.len() == 0 + } + + pub(crate) fn len(&self) -> usize { + self.buf.len() - self.pos + } + + pub(crate) fn copy_to(&mut self, dst: &mut ReadBuf<'_>) -> usize { + let n = cmp::min(self.len(), dst.remaining()); + dst.put_slice(&self.bytes()[..n]); + self.pos += n; + + if self.pos == self.buf.len() { + self.buf.truncate(0); + self.pos = 0; + } + + n + } + + pub(crate) fn copy_from(&mut self, src: &[u8]) -> usize { + assert!(self.is_empty()); + + let n = cmp::min(src.len(), MAX_BUF); + + self.buf.extend_from_slice(&src[..n]); + n + } + + pub(crate) fn bytes(&self) -> &[u8] { + &self.buf[self.pos..] + } + + pub(crate) fn ensure_capacity_for(&mut self, bytes: &ReadBuf<'_>) { + assert!(self.is_empty()); + + let len = cmp::min(bytes.remaining(), MAX_BUF); + + if self.buf.len() < len { + self.buf.reserve(len - self.buf.len()); + } + + unsafe { + self.buf.set_len(len); + } + } + + pub(crate) fn read_from<T: Read>(&mut self, rd: &mut T) -> io::Result<usize> { + let res = uninterruptibly!(rd.read(&mut self.buf)); + + if let Ok(n) = res { + self.buf.truncate(n); + } else { + self.buf.clear(); + } + + assert_eq!(self.pos, 0); + + res + } + + pub(crate) fn write_to<T: Write>(&mut self, wr: &mut T) -> io::Result<()> { + assert_eq!(self.pos, 0); + + // `write_all` already ignores interrupts + let res = wr.write_all(&self.buf); + self.buf.clear(); + res + } +} + +cfg_fs! { + impl Buf { + pub(crate) fn discard_read(&mut self) -> i64 { + let ret = -(self.bytes().len() as i64); + self.pos = 0; + self.buf.truncate(0); + ret + } + } +} diff --git a/third_party/rust/tokio/src/io/bsd/poll_aio.rs b/third_party/rust/tokio/src/io/bsd/poll_aio.rs new file mode 100644 index 0000000000..f1ac4b2d77 --- /dev/null +++ b/third_party/rust/tokio/src/io/bsd/poll_aio.rs @@ -0,0 +1,195 @@ +//! Use POSIX AIO futures with Tokio. + +use crate::io::driver::{Handle, Interest, ReadyEvent, Registration}; +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 = 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/third_party/rust/tokio/src/io/driver/interest.rs b/third_party/rust/tokio/src/io/driver/interest.rs new file mode 100644 index 0000000000..d6b46dfb7c --- /dev/null +++ b/third_party/rust/tokio/src/io/driver/interest.rs @@ -0,0 +1,132 @@ +#![cfg_attr(not(feature = "net"), allow(dead_code, unreachable_pub))] + +use crate::io::driver::Ready; + +use std::fmt; +use std::ops; + +/// Readiness event interest. +/// +/// Specifies the readiness events the caller is interested in when awaiting on +/// I/O resource readiness states. +#[cfg_attr(docsrs, doc(cfg(feature = "net")))] +#[derive(Clone, Copy, Eq, PartialEq)] +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. + /// + /// Writable interest includes write-closed events. + pub const WRITABLE: Interest = Interest(mio::Interest::WRITABLE); + + /// Returns true if the value includes readable interest. + /// + /// # Examples + /// + /// ``` + /// use tokio::io::Interest; + /// + /// assert!(Interest::READABLE.is_readable()); + /// assert!(!Interest::WRITABLE.is_readable()); + /// + /// let both = Interest::READABLE | Interest::WRITABLE; + /// assert!(both.is_readable()); + /// ``` + pub const fn is_readable(self) -> bool { + self.0.is_readable() + } + + /// Returns true if the value includes writable interest. + /// + /// # Examples + /// + /// ``` + /// use tokio::io::Interest; + /// + /// assert!(!Interest::READABLE.is_writable()); + /// assert!(Interest::WRITABLE.is_writable()); + /// + /// let both = Interest::READABLE | Interest::WRITABLE; + /// assert!(both.is_writable()); + /// ``` + pub const fn is_writable(self) -> bool { + self.0.is_writable() + } + + /// Add together two `Interest` values. + /// + /// This function works from a `const` context. + /// + /// # Examples + /// + /// ``` + /// use tokio::io::Interest; + /// + /// const BOTH: Interest = Interest::READABLE.add(Interest::WRITABLE); + /// + /// assert!(BOTH.is_readable()); + /// assert!(BOTH.is_writable()); + pub const fn add(self, other: Interest) -> Interest { + Interest(self.0.add(other.0)) + } + + // This function must be crate-private to avoid exposing a `mio` dependency. + pub(crate) const fn to_mio(self) -> mio::Interest { + self.0 + } + + pub(super) fn mask(self) -> Ready { + match self { + Interest::READABLE => Ready::READABLE | Ready::READ_CLOSED, + Interest::WRITABLE => Ready::WRITABLE | Ready::WRITE_CLOSED, + _ => Ready::EMPTY, + } + } +} + +impl ops::BitOr for Interest { + type Output = Self; + + #[inline] + fn bitor(self, other: Self) -> Self { + self.add(other) + } +} + +impl ops::BitOrAssign for Interest { + #[inline] + fn bitor_assign(&mut self, other: Self) { + self.0 = (*self | other).0; + } +} + +impl fmt::Debug for Interest { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(fmt) + } +} diff --git a/third_party/rust/tokio/src/io/driver/mod.rs b/third_party/rust/tokio/src/io/driver/mod.rs new file mode 100644 index 0000000000..19f67a24e7 --- /dev/null +++ b/third_party/rust/tokio/src/io/driver/mod.rs @@ -0,0 +1,354 @@ +#![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>, +} + +#[derive(Debug)] +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/third_party/rust/tokio/src/io/driver/platform.rs b/third_party/rust/tokio/src/io/driver/platform.rs new file mode 100644 index 0000000000..6b27988ce6 --- /dev/null +++ b/third_party/rust/tokio/src/io/driver/platform.rs @@ -0,0 +1,44 @@ +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/third_party/rust/tokio/src/io/driver/ready.rs b/third_party/rust/tokio/src/io/driver/ready.rs new file mode 100644 index 0000000000..2430d3022f --- /dev/null +++ b/third_party/rust/tokio/src/io/driver/ready.rs @@ -0,0 +1,250 @@ +#![cfg_attr(not(feature = "net"), allow(unreachable_pub))] + +use std::fmt; +use std::ops; + +const READABLE: usize = 0b0_01; +const WRITABLE: usize = 0b0_10; +const READ_CLOSED: usize = 0b0_0100; +const WRITE_CLOSED: usize = 0b0_1000; + +/// 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)] +pub struct Ready(usize); + +impl Ready { + /// Returns the empty `Ready` set. + pub const EMPTY: Ready = Ready(0); + + /// Returns a `Ready` representing readable readiness. + pub const READABLE: Ready = Ready(READABLE); + + /// Returns a `Ready` representing writable readiness. + pub const WRITABLE: Ready = Ready(WRITABLE); + + /// Returns a `Ready` representing read closed readiness. + pub const READ_CLOSED: Ready = Ready(READ_CLOSED); + + /// Returns a `Ready` representing write closed readiness. + pub const WRITE_CLOSED: Ready = Ready(WRITE_CLOSED); + + /// Returns a `Ready` representing readiness for all operations. + 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; + } + + if event.is_writable() { + ready |= Ready::WRITABLE; + } + + if event.is_read_closed() { + ready |= Ready::READ_CLOSED; + } + + if event.is_write_closed() { + ready |= Ready::WRITE_CLOSED; + } + + ready + } + + /// Returns true if `Ready` is the empty set. + /// + /// # Examples + /// + /// ``` + /// use tokio::io::Ready; + /// + /// assert!(Ready::EMPTY.is_empty()); + /// assert!(!Ready::READABLE.is_empty()); + /// ``` + pub fn is_empty(self) -> bool { + self == Ready::EMPTY + } + + /// Returns `true` if the value includes `readable`. + /// + /// # Examples + /// + /// ``` + /// use tokio::io::Ready; + /// + /// assert!(!Ready::EMPTY.is_readable()); + /// assert!(Ready::READABLE.is_readable()); + /// assert!(Ready::READ_CLOSED.is_readable()); + /// assert!(!Ready::WRITABLE.is_readable()); + /// ``` + pub fn is_readable(self) -> bool { + self.contains(Ready::READABLE) || self.is_read_closed() + } + + /// Returns `true` if the value includes writable `readiness`. + /// + /// # Examples + /// + /// ``` + /// use tokio::io::Ready; + /// + /// assert!(!Ready::EMPTY.is_writable()); + /// assert!(!Ready::READABLE.is_writable()); + /// assert!(Ready::WRITABLE.is_writable()); + /// assert!(Ready::WRITE_CLOSED.is_writable()); + /// ``` + pub fn is_writable(self) -> bool { + self.contains(Ready::WRITABLE) || self.is_write_closed() + } + + /// Returns `true` if the value includes read-closed `readiness`. + /// + /// # Examples + /// + /// ``` + /// use tokio::io::Ready; + /// + /// assert!(!Ready::EMPTY.is_read_closed()); + /// assert!(!Ready::READABLE.is_read_closed()); + /// assert!(Ready::READ_CLOSED.is_read_closed()); + /// ``` + pub fn is_read_closed(self) -> bool { + self.contains(Ready::READ_CLOSED) + } + + /// Returns `true` if the value includes write-closed `readiness`. + /// + /// # Examples + /// + /// ``` + /// use tokio::io::Ready; + /// + /// assert!(!Ready::EMPTY.is_write_closed()); + /// assert!(!Ready::WRITABLE.is_write_closed()); + /// assert!(Ready::WRITE_CLOSED.is_write_closed()); + /// ``` + pub fn is_write_closed(self) -> bool { + self.contains(Ready::WRITE_CLOSED) + } + + /// Returns true if `self` is a superset of `other`. + /// + /// `other` may represent more than one readiness operations, in which case + /// the function only returns true if `self` contains all readiness + /// specified in `other`. + pub(crate) fn contains<T: Into<Self>>(self, other: T) -> bool { + let other = other.into(); + (self & other) == other + } + + /// Creates a `Ready` instance using the given `usize` representation. + /// + /// The `usize` representation must have been obtained from a call to + /// `Readiness::as_usize`. + /// + /// This function is mainly provided to allow the caller to get a + /// readiness value from an `AtomicUsize`. + pub(crate) fn from_usize(val: usize) -> Ready { + Ready(val & Ready::ALL.as_usize()) + } + + /// Returns a `usize` representation of the `Ready` value. + /// + /// This function is mainly provided to allow the caller to store a + /// readiness value in an `AtomicUsize`. + pub(crate) fn as_usize(self) -> usize { + self.0 + } +} + +cfg_io_readiness! { + use crate::io::Interest; + + impl Ready { + pub(crate) fn from_interest(interest: Interest) -> Ready { + let mut ready = Ready::EMPTY; + + if interest.is_readable() { + ready |= Ready::READABLE; + ready |= Ready::READ_CLOSED; + } + + if interest.is_writable() { + ready |= Ready::WRITABLE; + ready |= Ready::WRITE_CLOSED; + } + + ready + } + + pub(crate) fn intersection(self, interest: Interest) -> Ready { + Ready(self.0 & Ready::from_interest(interest).0) + } + + pub(crate) fn satisfies(self, interest: Interest) -> bool { + self.0 & Ready::from_interest(interest).0 != 0 + } + } +} + +impl ops::BitOr<Ready> for Ready { + type Output = Ready; + + #[inline] + fn bitor(self, other: Ready) -> Ready { + Ready(self.0 | other.0) + } +} + +impl ops::BitOrAssign<Ready> for Ready { + #[inline] + fn bitor_assign(&mut self, other: Ready) { + self.0 |= other.0; + } +} + +impl ops::BitAnd<Ready> for Ready { + type Output = Ready; + + #[inline] + fn bitand(self, other: Ready) -> Ready { + Ready(self.0 & other.0) + } +} + +impl ops::Sub<Ready> for Ready { + type Output = Ready; + + #[inline] + fn sub(self, other: Ready) -> Ready { + Ready(self.0 & !other.0) + } +} + +impl fmt::Debug for Ready { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt.debug_struct("Ready") + .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() + } +} diff --git a/third_party/rust/tokio/src/io/driver/registration.rs b/third_party/rust/tokio/src/io/driver/registration.rs new file mode 100644 index 0000000000..7350be6345 --- /dev/null +++ b/third_party/rust/tokio/src/io/driver/registration.rs @@ -0,0 +1,262 @@ +#![cfg_attr(not(feature = "net"), allow(dead_code))] + +use crate::io::driver::{Direction, Handle, Interest, ReadyEvent, ScheduledIo}; +use crate::util::slab; + +use mio::event::Source; +use std::io; +use std::task::{Context, Poll}; + +cfg_io_driver! { + /// Associates an I/O resource with the reactor instance that drives it. + /// + /// A registration represents an I/O resource registered with a Reactor such + /// that it will receive task notifications on readiness. This is the lowest + /// level API for integrating with a reactor. + /// + /// The association between an I/O resource is made by calling + /// [`new_with_interest_and_handle`]. + /// Once the association is established, it remains established until the + /// registration instance is dropped. + /// + /// A registration instance represents two separate readiness streams. One + /// for the read readiness and one for write readiness. These streams are + /// independent and can be consumed from separate tasks. + /// + /// **Note**: while `Registration` is `Sync`, the caller must ensure that + /// there are at most two tasks that use a registration instance + /// concurrently. One task for [`poll_read_ready`] and one task for + /// [`poll_write_ready`]. While violating this requirement is "safe" from a + /// Rust memory safety point of view, it will result in unexpected behavior + /// in the form of lost notifications and tasks hanging. + /// + /// ## Platform-specific events + /// + /// `Registration` also allows receiving platform-specific `mio::Ready` + /// events. These events are included as part of the read readiness event + /// stream. The write readiness event stream is only for `Ready::writable()` + /// events. + /// + /// [`new_with_interest_and_handle`]: method@Self::new_with_interest_and_handle + /// [`poll_read_ready`]: method@Self::poll_read_ready` + /// [`poll_write_ready`]: method@Self::poll_write_ready` + #[derive(Debug)] + pub(crate) struct Registration { + /// Handle to the associated driver. + handle: Handle, + + /// Reference to state stored by the driver. + shared: slab::Ref<ScheduledIo>, + } +} + +unsafe impl Send for Registration {} +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 + /// interested in those states, you will need to add them to the readiness + /// state passed to this function. + /// + /// # Return + /// + /// - `Ok` if the registration happened successfully + /// - `Err` if an error was encountered during registration + pub(crate) fn new_with_interest_and_handle( + io: &mut impl Source, + interest: Interest, + handle: 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", + )); + }; + + Ok(Registration { handle, shared }) + } + + /// Deregisters the I/O resource from the reactor it is associated with. + /// + /// This function must be called before the I/O resource associated with the + /// registration is dropped. + /// + /// Note that deregistering does not guarantee that the I/O resource can be + /// registered with a different reactor. Some I/O resource types can only be + /// associated with a single reactor instance for their lifetime. + /// + /// # Return + /// + /// If the deregistration was successful, `Ok` is returned. Any calls to + /// `Reactor::turn` that happen after a successful call to `deregister` will + /// no longer result in notifications getting sent for this 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) + } + + pub(crate) fn clear_readiness(&self, event: ReadyEvent) { + self.shared.clear_readiness(event); + } + + // Uses the poll path, requiring the caller to ensure mutual exclusion for + // correctness. Only the last task to call this function is notified. + pub(crate) fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<ReadyEvent>> { + self.poll_ready(cx, Direction::Read) + } + + // Uses the poll path, requiring the caller to ensure mutual exclusion for + // correctness. Only the last task to call this function is notified. + pub(crate) fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<ReadyEvent>> { + self.poll_ready(cx, Direction::Write) + } + + // Uses the poll path, requiring the caller to ensure mutual exclusion for + // correctness. Only the last task to call this function is notified. + pub(crate) fn poll_read_io<R>( + &self, + cx: &mut Context<'_>, + f: impl FnMut() -> io::Result<R>, + ) -> Poll<io::Result<R>> { + self.poll_io(cx, Direction::Read, f) + } + + // Uses the poll path, requiring the caller to ensure mutual exclusion for + // correctness. Only the last task to call this function is notified. + pub(crate) fn poll_write_io<R>( + &self, + cx: &mut Context<'_>, + f: impl FnMut() -> io::Result<R>, + ) -> Poll<io::Result<R>> { + self.poll_io(cx, Direction::Write, f) + } + + /// Polls for events on the I/O resource's `direction` readiness stream. + /// + /// If called with a task context, notify the task when a new event is + /// received. + fn poll_ready( + &self, + cx: &mut Context<'_>, + direction: Direction, + ) -> Poll<io::Result<ReadyEvent>> { + // Keep track of task budget + let coop = ready!(crate::coop::poll_proceed(cx)); + let ev = ready!(self.shared.poll_readiness(cx, direction)); + + if self.handle.inner().is_none() { + return Poll::Ready(Err(gone())); + } + + coop.made_progress(); + Poll::Ready(Ok(ev)) + } + + fn poll_io<R>( + &self, + cx: &mut Context<'_>, + direction: Direction, + mut f: impl FnMut() -> io::Result<R>, + ) -> Poll<io::Result<R>> { + loop { + let ev = ready!(self.poll_ready(cx, direction))?; + + match f() { + Ok(ret) => { + return Poll::Ready(Ok(ret)); + } + Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + self.clear_readiness(ev); + } + Err(e) => return Poll::Ready(Err(e)), + } + } + } + + pub(crate) fn try_io<R>( + &self, + interest: Interest, + f: impl FnOnce() -> io::Result<R>, + ) -> io::Result<R> { + let ev = self.shared.ready_event(interest); + + // Don't attempt the operation if the resource is not ready. + if ev.ready.is_empty() { + return Err(io::ErrorKind::WouldBlock.into()); + } + + match f() { + Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + self.clear_readiness(ev); + Err(io::ErrorKind::WouldBlock.into()) + } + res => res, + } + } +} + +impl Drop for Registration { + fn drop(&mut self) { + // It is possible for a cycle to be created between wakers stored in + // `ScheduledIo` instances and `Arc<driver::Inner>`. To break this + // cycle, wakers are cleared. This is an imperfect solution as it is + // possible to store a `Registration` in a waker. In this case, the + // cycle would remain. + // + // See tokio-rs/tokio#3481 for more details. + self.shared.clear_wakers(); + } +} + +fn gone() -> io::Error { + io::Error::new(io::ErrorKind::Other, "IO driver has terminated") +} + +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); + + 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 + ))); + } + + Pin::new(&mut fut).poll(cx).map(Ok) + }).await + } + + pub(crate) async fn async_io<R>(&self, interest: Interest, mut f: impl FnMut() -> io::Result<R>) -> io::Result<R> { + loop { + let event = self.readiness(interest).await?; + + match f() { + Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + self.clear_readiness(event); + } + x => return x, + } + } + } + } +} diff --git a/third_party/rust/tokio/src/io/driver/scheduled_io.rs b/third_party/rust/tokio/src/io/driver/scheduled_io.rs new file mode 100644 index 0000000000..76f93431ba --- /dev/null +++ b/third_party/rust/tokio/src/io/driver/scheduled_io.rs @@ -0,0 +1,533 @@ +use super::{Interest, Ready, ReadyEvent, Tick}; +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}; + +use super::Direction; + +cfg_io_readiness! { + use crate::util::linked_list::{self, LinkedList}; + + use std::cell::UnsafeCell; + use std::future::Future; + use std::marker::PhantomPinned; + use std::pin::Pin; + use std::ptr::NonNull; +} + +/// Stored in the I/O driver resource slab. +#[derive(Debug)] +pub(crate) struct ScheduledIo { + /// Packs the resource's readiness with the resource's generation. + readiness: AtomicUsize, + + waiters: Mutex<Waiters>, +} + +cfg_io_readiness! { + type WaitList = LinkedList<Waiter, <Waiter as linked_list::Link>::Target>; +} + +#[derive(Debug, Default)] +struct Waiters { + #[cfg(feature = "net")] + /// List of all current waiters. + list: WaitList, + + /// Waker used for AsyncRead. + reader: Option<Waker>, + + /// 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! { + #[derive(Debug)] + struct Waiter { + pointers: linked_list::Pointers<Waiter>, + + /// The waker for this task. + waker: Option<Waker>, + + /// The interest this waiter is waiting on. + interest: Interest, + + is_ready: bool, + + /// Should never be `!Unpin`. + _p: PhantomPinned, + } + + /// Future returned by `readiness()`. + struct Readiness<'a> { + scheduled_io: &'a ScheduledIo, + + state: State, + + /// Entry in the waiter `LinkedList`. + waiter: UnsafeCell<Waiter>, + } + + enum State { + Init, + Waiting, + Done, + } +} + +// The `ScheduledIo::readiness` (`AtomicUsize`) is packed full of goodness. +// +// | reserved | generation | driver tick | readiness | +// |----------+------------+--------------+-----------| +// | 1 bit | 7 bits + 8 bits + 16 bits | + +const READINESS: bit::Pack = bit::Pack::least_significant(16); + +const TICK: bit::Pack = READINESS.then(8); + +const GENERATION: bit::Pack = TICK.then(7); + +#[test] +fn test_generations_assert_same() { + assert_eq!(super::GENERATION, GENERATION); +} + +// ===== impl ScheduledIo ===== + +impl Entry for ScheduledIo { + fn reset(&self) { + let state = self.readiness.load(Acquire); + + let generation = GENERATION.unpack(state); + let next = GENERATION.pack_lossy(generation + 1, 0); + + self.readiness.store(next, Release); + } +} + +impl Default for ScheduledIo { + fn default() -> ScheduledIo { + ScheduledIo { + readiness: AtomicUsize::new(0), + waiters: Mutex::new(Default::default()), + } + } +} + +impl ScheduledIo { + pub(crate) fn generation(&self) -> usize { + GENERATION.unpack(self.readiness.load(Acquire)) + } + + /// Invoked when the IO driver is shut down; forces this ScheduledIo into a + /// permanently ready state. + pub(super) fn shutdown(&self) { + self.wake0(Ready::ALL, true) + } + + /// Sets the readiness on this `ScheduledIo` by invoking the given closure on + /// the current value, returning the previous readiness value. + /// + /// # Arguments + /// - `token`: the token for this `ScheduledIo`. + /// - `tick`: whether setting the tick or trying to clear readiness for a + /// specific tick. + /// - `f`: a closure returning a new readiness value given the previous + /// readiness. + /// + /// # Returns + /// + /// If the given token's generation no longer matches the `ScheduledIo`'s + /// generation, then the corresponding IO resource has been removed and + /// replaced with a new resource. In that case, this method returns `Err`. + /// Otherwise, this returns the previous readiness. + pub(super) fn set_readiness( + &self, + token: Option<usize>, + tick: Tick, + f: impl Fn(Ready) -> Ready, + ) -> Result<(), ()> { + let mut current = self.readiness.load(Acquire); + + loop { + let current_generation = GENERATION.unpack(current); + + if let Some(token) = token { + // Check that the generation for this access is still the + // current one. + if GENERATION.unpack(token) != current_generation { + return Err(()); + } + } + + // Mask out the tick/generation bits so that the modifying + // function doesn't see them. + let current_readiness = Ready::from_usize(current); + let new = f(current_readiness); + + let packed = match tick { + Tick::Set(t) => TICK.pack(t as usize, new.as_usize()), + Tick::Clear(t) => { + if TICK.unpack(current) as u8 != t { + // Trying to clear readiness with an old event! + return Err(()); + } + + TICK.pack(t as usize, new.as_usize()) + } + }; + + let next = GENERATION.pack(current_generation, packed); + + match self + .readiness + .compare_exchange(current, next, AcqRel, Acquire) + { + Ok(_) => return Ok(()), + // we lost the race, retry! + Err(actual) => current = actual, + } + } + } + + /// Notifies all pending waiters that have registered interest in `ready`. + /// + /// There may be many waiters to notify. Waking the pending task **must** be + /// done from outside of the lock otherwise there is a potential for a + /// deadlock. + /// + /// A stack array of wakers is created and filled with wakers to notify, the + /// lock is released, and the wakers are notified. Because there may be more + /// 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) { + 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.push(waker); + } + } + + // check for AsyncWrite slot + if ready.is_writable() { + if let Some(waker) = waiters.writer.take() { + wakers.push(waker); + } + } + + #[cfg(feature = "net")] + 'outer: loop { + let mut iter = waiters.list.drain_filter(|w| ready.satisfies(w.interest)); + + 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.push(waker); + } + } + None => { + break 'outer; + } + } + } + + drop(waiters); + + wakers.wake_all(); + + // Acquire the lock again. + waiters = self.waiters.lock(); + } + + // Release the lock before notifying + drop(waiters); + + wakers.wake_all(); + } + + pub(super) fn ready_event(&self, interest: Interest) -> ReadyEvent { + let curr = self.readiness.load(Acquire); + + ReadyEvent { + tick: TICK.unpack(curr) as u8, + ready: interest.mask() & Ready::from_usize(READINESS.unpack(curr)), + } + } + + /// 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 + /// and writer slots. + pub(super) fn poll_readiness( + &self, + cx: &mut Context<'_>, + direction: Direction, + ) -> Poll<ReadyEvent> { + let curr = self.readiness.load(Acquire); + + let ready = direction.mask() & Ready::from_usize(READINESS.unpack(curr)); + + if ready.is_empty() { + // Update the task info + let mut waiters = self.waiters.lock(); + let slot = match direction { + Direction::Read => &mut waiters.reader, + Direction::Write => &mut waiters.writer, + }; + + // Avoid cloning the waker if one is already stored that matches the + // current task. + match slot { + Some(existing) => { + if !existing.will_wake(cx.waker()) { + *existing = cx.waker().clone(); + } + } + None => { + *slot = Some(cx.waker().clone()); + } + } + + // Try again, in case the readiness was changed while we were + // taking the waiters lock + let curr = self.readiness.load(Acquire); + let ready = direction.mask() & Ready::from_usize(READINESS.unpack(curr)); + if waiters.is_shutdown { + Poll::Ready(ReadyEvent { + tick: TICK.unpack(curr) as u8, + ready: direction.mask(), + }) + } else if ready.is_empty() { + Poll::Pending + } else { + Poll::Ready(ReadyEvent { + tick: TICK.unpack(curr) as u8, + ready, + }) + } + } else { + Poll::Ready(ReadyEvent { + tick: TICK.unpack(curr) as u8, + ready, + }) + } + } + + pub(crate) fn clear_readiness(&self, event: ReadyEvent) { + // This consumes the current readiness state **except** for closed + // states. Closed states are excluded because they are final states. + let mask_no_closed = event.ready - Ready::READ_CLOSED - Ready::WRITE_CLOSED; + + // result isn't important + let _ = self.set_readiness(None, Tick::Clear(event.tick), |curr| curr - mask_no_closed); + } + + pub(crate) fn clear_wakers(&self) { + let mut waiters = self.waiters.lock(); + waiters.reader.take(); + waiters.writer.take(); + } +} + +impl Drop for ScheduledIo { + fn drop(&mut self) { + self.wake(Ready::ALL); + } +} + +unsafe impl Send for ScheduledIo {} +unsafe impl Sync for ScheduledIo {} + +cfg_io_readiness! { + impl ScheduledIo { + /// 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 + } + + // This is in a separate function so that the borrow checker doesn't think + // we are borrowing the `UnsafeCell` possibly over await boundaries. + // + // Go figure. + fn readiness_fut(&self, interest: Interest) -> Readiness<'_> { + Readiness { + scheduled_io: self, + state: State::Init, + waiter: UnsafeCell::new(Waiter { + pointers: linked_list::Pointers::new(), + waker: None, + is_ready: false, + interest, + _p: PhantomPinned, + }), + } + } + } + + unsafe impl linked_list::Link for Waiter { + type Handle = NonNull<Waiter>; + type Target = Waiter; + + fn as_raw(handle: &NonNull<Waiter>) -> NonNull<Waiter> { + *handle + } + + unsafe fn from_raw(ptr: NonNull<Waiter>) -> NonNull<Waiter> { + ptr + } + + unsafe fn pointers(mut target: NonNull<Waiter>) -> NonNull<linked_list::Pointers<Waiter>> { + NonNull::from(&mut target.as_mut().pointers) + } + } + + // ===== impl Readiness ===== + + impl Future for Readiness<'_> { + type Output = ReadyEvent; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + use std::sync::atomic::Ordering::SeqCst; + + let (scheduled_io, state, waiter) = unsafe { + let me = self.get_unchecked_mut(); + (&me.scheduled_io, &mut me.state, &me.waiter) + }; + + loop { + match *state { + State::Init => { + // Optimistically check existing readiness + let curr = scheduled_io.readiness.load(SeqCst); + let ready = Ready::from_usize(READINESS.unpack(curr)); + + // Safety: `waiter.interest` never changes + let interest = unsafe { (*waiter.get()).interest }; + let ready = ready.intersection(interest); + + if !ready.is_empty() { + // Currently ready! + let tick = TICK.unpack(curr) as u8; + *state = State::Done; + return Poll::Ready(ReadyEvent { tick, ready }); + } + + // Wasn't ready, take the lock (and check again while locked). + let mut waiters = scheduled_io.waiters.lock(); + + let curr = scheduled_io.readiness.load(SeqCst); + let mut ready = Ready::from_usize(READINESS.unpack(curr)); + + if waiters.is_shutdown { + ready = Ready::ALL; + } + + let ready = ready.intersection(interest); + + if !ready.is_empty() { + // Currently ready! + let tick = TICK.unpack(curr) as u8; + *state = State::Done; + return Poll::Ready(ReadyEvent { tick, ready }); + } + + // Not ready even after locked, insert into list... + + // Safety: called while locked + unsafe { + (*waiter.get()).waker = Some(cx.waker().clone()); + } + + // Insert the waiter into the linked list + // + // safety: pointers from `UnsafeCell` are never null. + waiters + .list + .push_front(unsafe { NonNull::new_unchecked(waiter.get()) }); + *state = State::Waiting; + } + State::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. + + let waiters = scheduled_io.waiters.lock(); + + // Safety: called while locked + let w = unsafe { &mut *waiter.get() }; + + if w.is_ready { + // Our waker has been notified. + *state = State::Done; + } else { + // Update the waker, if necessary. + if !w.waker.as_ref().unwrap().will_wake(cx.waker()) { + w.waker = Some(cx.waker().clone()); + } + + return Poll::Pending; + } + + // Explicit drop of the lock to indicate the scope that the + // lock is held. Because holding the lock is required to + // ensure safe access to fields not held within the lock, it + // is helpful to visualize the scope of the critical + // section. + 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() }; + + return Poll::Ready(ReadyEvent { + tick, + ready: Ready::from_interest(w.interest), + }); + } + } + } + } + } + + impl Drop for Readiness<'_> { + fn drop(&mut self) { + let mut waiters = self.scheduled_io.waiters.lock(); + + // Safety: `waiter` is only ever stored in `waiters` + unsafe { + waiters + .list + .remove(NonNull::new_unchecked(self.waiter.get())) + }; + } + } + + unsafe impl Send for Readiness<'_> {} + unsafe impl Sync for Readiness<'_> {} +} diff --git a/third_party/rust/tokio/src/io/mod.rs b/third_party/rust/tokio/src/io/mod.rs new file mode 100644 index 0000000000..cfdda61f69 --- /dev/null +++ b/third_party/rust/tokio/src/io/mod.rs @@ -0,0 +1,276 @@ +#![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 +//! defines two traits, [`AsyncRead`] and [`AsyncWrite`], which are asynchronous +//! versions of the [`Read`] and [`Write`] traits in the standard library. +//! +//! # AsyncRead and AsyncWrite +//! +//! Like the standard library's [`Read`] and [`Write`] traits, [`AsyncRead`] and +//! [`AsyncWrite`] provide the most general interface for reading and writing +//! input and output. Unlike the standard library's traits, however, they are +//! _asynchronous_ — meaning that reading from or writing to a `tokio::io` +//! type will _yield_ to the Tokio scheduler when IO is not ready, rather than +//! blocking. This allows other tasks to run while waiting on IO. +//! +//! Another difference is that `AsyncRead` and `AsyncWrite` only contain +//! core methods needed to provide asynchronous reading and writing +//! functionality. Instead, utility methods are defined in the [`AsyncReadExt`] +//! and [`AsyncWriteExt`] extension traits. These traits are automatically +//! implemented for all values that implement `AsyncRead` and `AsyncWrite` +//! respectively. +//! +//! End users will rarely interact directly with `AsyncRead` and +//! `AsyncWrite`. Instead, they will use the async functions defined in the +//! extension traits. Library authors are expected to implement `AsyncRead` +//! and `AsyncWrite` in order to provide types that behave like byte streams. +//! +//! Even with these differences, Tokio's `AsyncRead` and `AsyncWrite` traits +//! can be used in almost exactly the same manner as the standard library's +//! `Read` and `Write`. Most types in the standard library that implement `Read` +//! and `Write` have asynchronous equivalents in `tokio` that implement +//! `AsyncRead` and `AsyncWrite`, such as [`File`] and [`TcpStream`]. +//! +//! For example, the standard library documentation introduces `Read` by +//! [demonstrating][std_example] reading some bytes from a [`std::fs::File`]. We +//! can do the same with [`tokio::fs::File`][`File`]: +//! +//! ```no_run +//! use tokio::io::{self, AsyncReadExt}; +//! use tokio::fs::File; +//! +//! #[tokio::main] +//! async fn main() -> io::Result<()> { +//! let mut f = File::open("foo.txt").await?; +//! let mut buffer = [0; 10]; +//! +//! // read up to 10 bytes +//! let n = f.read(&mut buffer).await?; +//! +//! println!("The bytes: {:?}", &buffer[..n]); +//! Ok(()) +//! } +//! ``` +//! +//! [`File`]: crate::fs::File +//! [`TcpStream`]: crate::net::TcpStream +//! [`std::fs::File`]: std::fs::File +//! [std_example]: std::io#read-and-write +//! +//! ## Buffered Readers and Writers +//! +//! Byte-based interfaces are unwieldy and can be inefficient, as we'd need to be +//! making near-constant calls to the operating system. To help with this, +//! `std::io` comes with [support for _buffered_ readers and writers][stdbuf], +//! and therefore, `tokio::io` does as well. +//! +//! Tokio provides an async version of the [`std::io::BufRead`] trait, +//! [`AsyncBufRead`]; and async [`BufReader`] and [`BufWriter`] structs, which +//! wrap readers and writers. These wrappers use a buffer, reducing the number +//! of calls and providing nicer methods for accessing exactly what you want. +//! +//! For example, [`BufReader`] works with the [`AsyncBufRead`] trait to add +//! extra methods to any async reader: +//! +//! ```no_run +//! use tokio::io::{self, BufReader, AsyncBufReadExt}; +//! use tokio::fs::File; +//! +//! #[tokio::main] +//! async fn main() -> io::Result<()> { +//! let f = File::open("foo.txt").await?; +//! let mut reader = BufReader::new(f); +//! let mut buffer = String::new(); +//! +//! // read a line into buffer +//! reader.read_line(&mut buffer).await?; +//! +//! println!("{}", buffer); +//! Ok(()) +//! } +//! ``` +//! +//! [`BufWriter`] doesn't add any new ways of writing; it just buffers every call +//! to [`write`](crate::io::AsyncWriteExt::write). However, you **must** flush +//! [`BufWriter`] to ensure that any buffered data is written. +//! +//! ```no_run +//! use tokio::io::{self, BufWriter, AsyncWriteExt}; +//! use tokio::fs::File; +//! +//! #[tokio::main] +//! async fn main() -> io::Result<()> { +//! let f = File::create("foo.txt").await?; +//! { +//! let mut writer = BufWriter::new(f); +//! +//! // Write a byte to the buffer. +//! writer.write(&[42u8]).await?; +//! +//! // Flush the buffer before it goes out of scope. +//! writer.flush().await?; +//! +//! } // Unless flushed or shut down, the contents of the buffer is discarded on drop. +//! +//! Ok(()) +//! } +//! ``` +//! +//! [stdbuf]: std::io#bufreader-and-bufwriter +//! [`std::io::BufRead`]: std::io::BufRead +//! [`AsyncBufRead`]: crate::io::AsyncBufRead +//! [`BufReader`]: crate::io::BufReader +//! [`BufWriter`]: crate::io::BufWriter +//! +//! ## Implementing AsyncRead and AsyncWrite +//! +//! Because they are traits, we can implement [`AsyncRead`] and [`AsyncWrite`] for +//! our own types, as well. Note that these traits must only be implemented for +//! non-blocking I/O types that integrate with the futures type system. In +//! 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 +//! +//! 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. +//! +//! [tokio-util]: https://docs.rs/tokio-util/0.6/tokio_util/codec/index.html +//! +//! # Standard input and output +//! +//! Tokio provides asynchronous APIs to standard [input], [output], and [error]. +//! These APIs are very similar to the ones provided by `std`, but they also +//! implement [`AsyncRead`] and [`AsyncWrite`]. +//! +//! Note that the standard input / output APIs **must** be used from the +//! context of the Tokio runtime, as they require Tokio-specific features to +//! function. Calling these functions outside of a Tokio runtime will panic. +//! +//! [input]: fn@stdin +//! [output]: fn@stdout +//! [error]: fn@stderr +//! +//! # `std` re-exports +//! +//! Additionally, [`Error`], [`ErrorKind`], [`Result`], and [`SeekFrom`] are +//! re-exported from `std::io` for ease of use. +//! +//! [`AsyncRead`]: trait@AsyncRead +//! [`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 +//! [`Error`]: struct@Error +//! [`ErrorKind`]: enum@ErrorKind +//! [`Result`]: type@Result +//! [`Read`]: std::io::Read +//! [`SeekFrom`]: enum@SeekFrom +//! [`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_io_blocking! { + pub(crate) mod blocking; +} + +mod async_buf_read; +pub use self::async_buf_read::AsyncBufRead; + +mod async_read; +pub use self::async_read::AsyncRead; + +mod async_seek; +pub use self::async_seek::AsyncSeek; + +mod async_write; +pub use self::async_write::AsyncWrite; + +mod read_buf; +pub use self::read_buf::ReadBuf; + +// Re-export some types from `std::io` so that users don't have to deal +// with conflicts when `use`ing `tokio::io` and `std::io`. +#[doc(no_inline)] +pub use std::io::{Error, ErrorKind, Result, SeekFrom}; + +cfg_io_driver_impl! { + pub(crate) mod driver; + + cfg_net! { + pub use driver::{Interest, Ready}; + } + + mod poll_evented; + + #[cfg(not(loom))] + 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; + + pub mod unix { + //! Asynchronous IO structures specific to Unix-like operating systems. + pub use super::async_fd::{AsyncFd, AsyncFdReadyGuard, AsyncFdReadyMutGuard, TryIoError}; + } +} + +cfg_io_std! { + mod stdio_common; + + mod stderr; + pub use stderr::{stderr, Stderr}; + + mod stdin; + pub use stdin::{stdin, Stdin}; + + mod stdout; + pub use stdout::{stdout, Stdout}; +} + +cfg_io_util! { + mod split; + pub use split::{split, ReadHalf, WriteHalf}; + + pub(crate) mod seek; + pub(crate) mod util; + pub use util::{ + copy, copy_bidirectional, copy_buf, duplex, empty, repeat, sink, AsyncBufReadExt, AsyncReadExt, AsyncSeekExt, AsyncWriteExt, + BufReader, BufStream, BufWriter, DuplexStream, Empty, Lines, Repeat, Sink, Split, Take, + }; +} + +cfg_not_io_util! { + cfg_process! { + pub(crate) mod util; + } +} + +cfg_io_blocking! { + /// Types in this module can be mocked out in tests. + mod sys { + // TODO: don't rename + pub(crate) use crate::blocking::spawn_blocking as run; + pub(crate) use crate::blocking::JoinHandle as Blocking; + } +} diff --git a/third_party/rust/tokio/src/io/poll_evented.rs b/third_party/rust/tokio/src/io/poll_evented.rs new file mode 100644 index 0000000000..ce4c1426ac --- /dev/null +++ b/third_party/rust/tokio/src/io/poll_evented.rs @@ -0,0 +1,214 @@ +use crate::io::driver::{Handle, Interest, Registration}; + +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 + /// [`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 + /// used from within the future's execution model. As such, the + /// `PollEvented` type provides [`AsyncRead`] and [`AsyncWrite`] + /// implementations using the underlying I/O resource as well as readiness + /// events provided by the reactor. + /// + /// **Note**: While `PollEvented` is `Sync` (if the underlying I/O type is + /// `Sync`), the caller must ensure that there are at most two tasks that + /// use a `PollEvented` instance concurrently. One for reading and one for + /// writing. While violating this requirement is "safe" from a Rust memory + /// model point of view, it will result in unexpected behavior in the form + /// of lost notifications and tasks hanging. + /// + /// ## Readiness events + /// + /// Besides just providing [`AsyncRead`] and [`AsyncWrite`] implementations, + /// this type also supports access to the underlying readiness event stream. + /// While similar in function to what [`Registration`] provides, the + /// semantics are a bit different. + /// + /// Two functions are provided to access the readiness events: + /// [`poll_read_ready`] and [`poll_write_ready`]. These functions return the + /// current readiness state of the `PollEvented` instance. If + /// [`poll_read_ready`] indicates read readiness, immediately calling + /// [`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_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`. + /// + /// ## Platform-specific events + /// + /// `PollEvented` also allows receiving platform-specific `mio::Ready` events. + /// These events are included as part of the read readiness event stream. The + /// write readiness event stream is only for `Ready::writable()` events. + /// + /// [`AsyncRead`]: crate::io::AsyncRead + /// [`AsyncWrite`]: crate::io::AsyncWrite + /// [`TcpListener`]: crate::net::TcpListener + /// [`poll_read_ready`]: Registration::poll_read_ready + /// [`poll_write_ready`]: Registration::poll_write_ready + pub(crate) struct PollEvented<E: Source> { + io: Option<E>, + registration: Registration, + } +} + +// ===== impl PollEvented ===== + +impl<E: Source> PollEvented<E> { + /// Creates a new `PollEvented` associated with the default reactor. + /// + /// # Panics + /// + /// This function panics if thread-local runtime is not set. + /// + /// 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. + #[cfg_attr(feature = "signal", allow(unused))] + pub(crate) fn new(io: E) -> io::Result<Self> { + PollEvented::new_with_interest(io, Interest::READABLE | Interest::WRITABLE) + } + + /// Creates a new `PollEvented` associated with the default reactor, for + /// specific `Interest` state. `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 interested in those states, you will need to add them to the + /// readiness state passed to this function. + /// + /// # Panics + /// + /// This function panics if thread-local runtime is not set. + /// + /// 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. + #[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()) + } + + pub(crate) fn new_with_interest_and_handle( + mut io: E, + interest: Interest, + handle: Handle, + ) -> io::Result<Self> { + let registration = Registration::new_with_interest_and_handle(&mut io, interest, handle)?; + Ok(Self { + io: Some(io), + registration, + }) + } + + /// Returns a reference to the registration. + #[cfg(any( + feature = "net", + all(unix, feature = "process"), + all(unix, feature = "signal"), + ))] + pub(crate) fn registration(&self) -> &Registration { + &self.registration + } + + /// 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. + self.registration.deregister(&mut inner)?; + Ok(inner) + } +} + +feature! { + #![any(feature = "net", feature = "process")] + + use crate::io::ReadBuf; + use std::task::{Context, Poll}; + + impl<E: Source> PollEvented<E> { + // Safety: The caller must ensure that `E` can read into uninitialized memory + pub(crate) unsafe fn poll_read<'a>( + &'a self, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<()>> + where + &'a E: io::Read + 'a, + { + use std::io::Read; + + let n = ready!(self.registration.poll_read_io(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(())) + } + + pub(crate) fn poll_write<'a>(&'a self, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> + where + &'a E: io::Write + 'a, + { + use std::io::Write; + self.registration.poll_write_io(cx, || self.io.as_ref().unwrap().write(buf)) + } + + #[cfg(feature = "net")] + pub(crate) fn poll_write_vectored<'a>( + &'a self, + cx: &mut Context<'_>, + bufs: &[io::IoSlice<'_>], + ) -> Poll<io::Result<usize>> + where + &'a E: io::Write + 'a, + { + use std::io::Write; + self.registration.poll_write_io(cx, || self.io.as_ref().unwrap().write_vectored(bufs)) + } + } +} + +impl<E: Source> UnwindSafe for PollEvented<E> {} + +impl<E: Source> RefUnwindSafe for PollEvented<E> {} + +impl<E: Source> Deref for PollEvented<E> { + type Target = E; + + fn deref(&self) -> &E { + self.io.as_ref().unwrap() + } +} + +impl<E: Source + fmt::Debug> fmt::Debug for PollEvented<E> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("PollEvented").field("io", &self.io).finish() + } +} + +impl<E: Source> Drop for PollEvented<E> { + fn drop(&mut self) { + if let Some(mut io) = self.io.take() { + // Ignore errors + let _ = self.registration.deregister(&mut io); + } + } +} diff --git a/third_party/rust/tokio/src/io/read_buf.rs b/third_party/rust/tokio/src/io/read_buf.rs new file mode 100644 index 0000000000..8c34ae6c81 --- /dev/null +++ b/third_party/rust/tokio/src/io/read_buf.rs @@ -0,0 +1,291 @@ +use std::fmt; +use std::mem::MaybeUninit; + +/// A wrapper around a byte buffer that is incrementally filled and initialized. +/// +/// This type is a sort of "double cursor". It tracks three regions in the +/// buffer: a region at the beginning of the buffer that has been logically +/// filled with data, a region that has been initialized at some point but not +/// yet logically filled, and a region at the end that may be uninitialized. +/// The filled region is guaranteed to be a subset of the initialized region. +/// +/// In summary, the contents of the buffer can be visualized as: +/// +/// ```not_rust +/// [ capacity ] +/// [ filled | unfilled ] +/// [ initialized | uninitialized ] +/// ``` +/// +/// It is undefined behavior to de-initialize any bytes from the uninitialized +/// region, since it is merely unknown whether this region is uninitialized or +/// not, and if part of it turns out to be initialized, it must stay initialized. +pub struct ReadBuf<'a> { + buf: &'a mut [MaybeUninit<u8>], + filled: usize, + initialized: usize, +} + +impl<'a> ReadBuf<'a> { + /// Creates a new `ReadBuf` from a fully initialized buffer. + #[inline] + pub fn new(buf: &'a mut [u8]) -> ReadBuf<'a> { + let initialized = buf.len(); + let buf = unsafe { slice_to_uninit_mut(buf) }; + ReadBuf { + buf, + filled: 0, + initialized, + } + } + + /// Creates a new `ReadBuf` from a fully uninitialized buffer. + /// + /// Use `assume_init` if part of the buffer is known to be already initialized. + #[inline] + pub fn uninit(buf: &'a mut [MaybeUninit<u8>]) -> ReadBuf<'a> { + ReadBuf { + buf, + filled: 0, + initialized: 0, + } + } + + /// Returns the total capacity of the buffer. + #[inline] + pub fn capacity(&self) -> usize { + self.buf.len() + } + + /// Returns a shared reference to the filled portion of the buffer. + #[inline] + pub fn filled(&self) -> &[u8] { + 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. + unsafe { slice_assume_init(slice) } + } + + /// Returns a mutable reference to the filled portion of the buffer. + #[inline] + pub fn filled_mut(&mut self) -> &mut [u8] { + 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. + unsafe { slice_assume_init_mut(slice) } + } + + /// Returns a new `ReadBuf` comprised of the unfilled section up to `n`. + #[inline] + pub fn take(&mut self, n: usize) -> ReadBuf<'_> { + let max = std::cmp::min(self.remaining(), n); + // Safety: We don't set any of the `unfilled_mut` with `MaybeUninit::uninit`. + unsafe { ReadBuf::uninit(&mut self.unfilled_mut()[..max]) } + } + + /// Returns a shared reference to the initialized portion of the buffer. + /// + /// This includes the filled portion. + #[inline] + pub fn initialized(&self) -> &[u8] { + let slice = &self.buf[..self.initialized]; + // safety: initialized describes how far into the buffer that the + // user has at some point initialized with bytes. + unsafe { slice_assume_init(slice) } + } + + /// Returns a mutable reference to the initialized portion of the buffer. + /// + /// This includes the filled portion. + #[inline] + pub fn initialized_mut(&mut self) -> &mut [u8] { + 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. + 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 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 + /// [`assume_init`](Self::assume_init) with the number of bytes initialized. + /// + /// # Safety + /// + /// The caller must not de-initialize portions of the buffer that have already been initialized. + /// This includes any bytes in the region marked as uninitialized by `ReadBuf`. + #[inline] + pub unsafe fn inner_mut(&mut self) -> &mut [MaybeUninit<u8>] { + self.buf + } + + /// Returns a mutable reference to the unfilled part of the buffer without ensuring that it has been fully + /// initialized. + /// + /// # Safety + /// + /// The caller must not de-initialize portions of the buffer that have already been initialized. + /// This includes any bytes in the region marked as uninitialized by `ReadBuf`. + #[inline] + pub unsafe fn unfilled_mut(&mut self) -> &mut [MaybeUninit<u8>] { + &mut self.buf[self.filled..] + } + + /// Returns a mutable reference to the unfilled part of the buffer, ensuring it is fully initialized. + /// + /// Since `ReadBuf` tracks the region of the buffer that has been initialized, this is effectively "free" after + /// the first use. + #[inline] + pub fn initialize_unfilled(&mut self) -> &mut [u8] { + self.initialize_unfilled_to(self.remaining()) + } + + /// Returns a mutable reference to the first `n` bytes of the unfilled part of the buffer, ensuring it is + /// fully initialized. + /// + /// # Panics + /// + /// Panics if `self.remaining()` is less than `n`. + #[inline] + pub fn initialize_unfilled_to(&mut self, n: usize) -> &mut [u8] { + assert!(self.remaining() >= n, "n overflows remaining"); + + // This can't overflow, otherwise the assert above would have failed. + let end = self.filled + n; + + if self.initialized < end { + unsafe { + self.buf[self.initialized..end] + .as_mut_ptr() + .write_bytes(0, end - self.initialized); + } + self.initialized = end; + } + + 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 { slice_assume_init_mut(slice) } + } + + /// Returns the number of bytes at the end of the slice that have not yet been filled. + #[inline] + pub fn remaining(&self) -> usize { + self.capacity() - self.filled + } + + /// Clears the buffer, resetting the filled region to empty. + /// + /// The number of initialized bytes is not changed, and the contents of the buffer are not modified. + #[inline] + pub fn clear(&mut self) { + self.filled = 0; + } + + /// Advances the size of the filled region of the buffer. + /// + /// The number of initialized bytes is not changed. + /// + /// # Panics + /// + /// Panics if the filled region of the buffer would become larger than the initialized region. + #[inline] + pub fn advance(&mut self, n: usize) { + let new = self.filled.checked_add(n).expect("filled overflow"); + self.set_filled(new); + } + + /// Sets the size of the filled region of the buffer. + /// + /// The number of initialized bytes is not changed. + /// + /// Note that this can be used to *shrink* the filled region of the buffer in addition to growing it (for + /// example, by a `AsyncRead` implementation that compresses data in-place). + /// + /// # Panics + /// + /// Panics if the filled region of the buffer would become larger than the initialized region. + #[inline] + pub fn set_filled(&mut self, n: usize) { + assert!( + n <= self.initialized, + "filled must not become larger than initialized" + ); + self.filled = n; + } + + /// Asserts that the first `n` unfilled bytes of the buffer are initialized. + /// + /// `ReadBuf` assumes that bytes are never de-initialized, so this method does nothing when called with fewer + /// bytes than are already known to be initialized. + /// + /// # Safety + /// + /// The caller must ensure that `n` unfilled bytes of the buffer have already been initialized. + #[inline] + pub unsafe fn assume_init(&mut self, n: usize) { + let new = self.filled + n; + if new > self.initialized { + self.initialized = new; + } + } + + /// Appends data to the buffer, advancing the written position and possibly also the initialized position. + /// + /// # Panics + /// + /// Panics if `self.remaining()` is less than `buf.len()`. + #[inline] + pub fn put_slice(&mut self, buf: &[u8]) { + assert!( + self.remaining() >= buf.len(), + "buf.len() must fit in remaining()" + ); + + let amt = buf.len(); + // Cannot overflow, asserted above + let end = self.filled + amt; + + // Safety: the length is asserted above + unsafe { + self.buf[self.filled..end] + .as_mut_ptr() + .cast::<u8>() + .copy_from_nonoverlapping(buf.as_ptr(), amt); + } + + if self.initialized < end { + self.initialized = end; + } + self.filled = end; + } +} + +impl fmt::Debug for ReadBuf<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ReadBuf") + .field("filled", &self.filled) + .field("initialized", &self.initialized) + .field("capacity", &self.capacity()) + .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/third_party/rust/tokio/src/io/seek.rs b/third_party/rust/tokio/src/io/seek.rs new file mode 100644 index 0000000000..e64205d9cf --- /dev/null +++ b/third_party/rust/tokio/src/io/seek.rs @@ -0,0 +1,57 @@ +use crate::io::AsyncSeek; + +use pin_project_lite::pin_project; +use std::future::Future; +use std::io::{self, SeekFrom}; +use std::marker::PhantomPinned; +use std::pin::Pin; +use std::task::{Context, Poll}; + +pin_project! { + /// Future for the [`seek`](crate::io::AsyncSeekExt::seek) method. + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct Seek<'a, S: ?Sized> { + seek: &'a mut S, + pos: Option<SeekFrom>, + // Make this future `!Unpin` for compatibility with async trait methods. + #[pin] + _pin: PhantomPinned, + } +} + +pub(crate) fn seek<S>(seek: &mut S, pos: SeekFrom) -> Seek<'_, S> +where + S: AsyncSeek + ?Sized + Unpin, +{ + Seek { + seek, + pos: Some(pos), + _pin: PhantomPinned, + } +} + +impl<S> Future for Seek<'_, S> +where + S: AsyncSeek + ?Sized + Unpin, +{ + type Output = io::Result<u64>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let me = self.project(); + match me.pos { + Some(pos) => { + // ensure no seek in progress + ready!(Pin::new(&mut *me.seek).poll_complete(cx))?; + match Pin::new(&mut *me.seek).start_seek(*pos) { + Ok(()) => { + *me.pos = None; + Pin::new(&mut *me.seek).poll_complete(cx) + } + Err(e) => Poll::Ready(Err(e)), + } + } + None => Pin::new(&mut *me.seek).poll_complete(cx), + } + } +} diff --git a/third_party/rust/tokio/src/io/split.rs b/third_party/rust/tokio/src/io/split.rs new file mode 100644 index 0000000000..8258a0f7a0 --- /dev/null +++ b/third_party/rust/tokio/src/io/split.rs @@ -0,0 +1,180 @@ +//! Split a single value implementing `AsyncRead + AsyncWrite` into separate +//! `AsyncRead` and `AsyncWrite` handles. +//! +//! To restore this read/write object from its `split::ReadHalf` and +//! `split::WriteHalf` use `unsplit`. + +use crate::io::{AsyncRead, AsyncWrite, ReadBuf}; + +use std::cell::UnsafeCell; +use std::fmt; +use std::io; +use std::pin::Pin; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering::{Acquire, Release}; +use std::sync::Arc; +use std::task::{Context, Poll}; + +cfg_io_util! { + /// The readable half of a value returned from [`split`](split()). + pub struct ReadHalf<T> { + inner: Arc<Inner<T>>, + } + + /// The writable half of a value returned from [`split`](split()). + pub struct WriteHalf<T> { + inner: Arc<Inner<T>>, + } + + /// Splits a single value implementing `AsyncRead + AsyncWrite` into separate + /// `AsyncRead` and `AsyncWrite` handles. + /// + /// To restore this read/write object from its `ReadHalf` and + /// `WriteHalf` use [`unsplit`](ReadHalf::unsplit()). + pub fn split<T>(stream: T) -> (ReadHalf<T>, WriteHalf<T>) + where + T: AsyncRead + AsyncWrite, + { + let inner = Arc::new(Inner { + locked: AtomicBool::new(false), + stream: UnsafeCell::new(stream), + }); + + let rd = ReadHalf { + inner: inner.clone(), + }; + + let wr = WriteHalf { inner }; + + (rd, wr) + } +} + +struct Inner<T> { + locked: AtomicBool, + stream: UnsafeCell<T>, +} + +struct Guard<'a, T> { + inner: &'a Inner<T>, +} + +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) + } + + /// Reunites with a previously split `WriteHalf`. + /// + /// # Panics + /// + /// If this `ReadHalf` and the given `WriteHalf` do not originate from the + /// 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 { + if self.is_pair_of(&wr) { + drop(wr); + + let inner = Arc::try_unwrap(self.inner) + .ok() + .expect("`Arc::try_unwrap` failed"); + + inner.stream.into_inner() + } else { + panic!("Unrelated `split::Write` passed to `split::Read::unsplit`.") + } + } +} + +impl<T> WriteHalf<T> { + /// 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) + } +} + +impl<T: AsyncRead> AsyncRead for ReadHalf<T> { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<()>> { + let mut inner = ready!(self.inner.poll_lock(cx)); + inner.stream_pin().poll_read(cx, buf) + } +} + +impl<T: AsyncWrite> AsyncWrite for WriteHalf<T> { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<Result<usize, io::Error>> { + let mut inner = ready!(self.inner.poll_lock(cx)); + inner.stream_pin().poll_write(cx, buf) + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> { + let mut inner = ready!(self.inner.poll_lock(cx)); + inner.stream_pin().poll_flush(cx) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> { + let mut inner = ready!(self.inner.poll_lock(cx)); + inner.stream_pin().poll_shutdown(cx) + } +} + +impl<T> Inner<T> { + fn poll_lock(&self, cx: &mut Context<'_>) -> Poll<Guard<'_, T>> { + if self + .locked + .compare_exchange(false, true, Acquire, Acquire) + .is_ok() + { + Poll::Ready(Guard { inner: self }) + } else { + // Spin... but investigate a better strategy + + std::thread::yield_now(); + cx.waker().wake_by_ref(); + + Poll::Pending + } + } +} + +impl<T> Guard<'_, T> { + fn stream_pin(&mut self) -> Pin<&mut T> { + // safety: the stream is pinned in `Arc` and the `Guard` ensures mutual + // exclusion. + unsafe { Pin::new_unchecked(&mut *self.inner.stream.get()) } + } +} + +impl<T> Drop for Guard<'_, T> { + fn drop(&mut self) { + self.inner.locked.store(false, Release); + } +} + +unsafe impl<T: Send> Send for ReadHalf<T> {} +unsafe impl<T: Send> Send for WriteHalf<T> {} +unsafe impl<T: Sync> Sync for ReadHalf<T> {} +unsafe impl<T: Sync> Sync for WriteHalf<T> {} + +impl<T: fmt::Debug> fmt::Debug for ReadHalf<T> { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt.debug_struct("split::ReadHalf").finish() + } +} + +impl<T: fmt::Debug> fmt::Debug for WriteHalf<T> { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt.debug_struct("split::WriteHalf").finish() + } +} diff --git a/third_party/rust/tokio/src/io/stderr.rs b/third_party/rust/tokio/src/io/stderr.rs new file mode 100644 index 0000000000..2f624fba9d --- /dev/null +++ b/third_party/rust/tokio/src/io/stderr.rs @@ -0,0 +1,109 @@ +use crate::io::blocking::Blocking; +use crate::io::stdio_common::SplitByUtf8BoundaryIfWindows; +use crate::io::AsyncWrite; + +use std::io; +use std::pin::Pin; +use std::task::Context; +use std::task::Poll; + +cfg_io_std! { + /// A handle to the standard error stream of a process. + /// + /// Concurrent writes to stderr must be executed with care: Only individual + /// writes to this [`AsyncWrite`] are guaranteed to be intact. In particular + /// you should be aware that writes using [`write_all`] are not guaranteed + /// to occur as a single write, so multiple threads writing data with + /// [`write_all`] may result in interleaved output. + /// + /// Created by the [`stderr`] function. + /// + /// [`stderr`]: stderr() + /// [`AsyncWrite`]: AsyncWrite + /// [`write_all`]: crate::io::AsyncWriteExt::write_all() + /// + /// # Examples + /// + /// ``` + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut stderr = io::stdout(); + /// stderr.write_all(b"Print some error here.").await?; + /// Ok(()) + /// } + /// ``` + #[derive(Debug)] + pub struct Stderr { + std: SplitByUtf8BoundaryIfWindows<Blocking<std::io::Stderr>>, + } + + /// Constructs a new handle to the standard error of the current process. + /// + /// The returned handle allows writing to standard error from the within the + /// Tokio runtime. + /// + /// Concurrent writes to stderr must be executed with care: Only individual + /// writes to this [`AsyncWrite`] are guaranteed to be intact. In particular + /// you should be aware that writes using [`write_all`] are not guaranteed + /// to occur as a single write, so multiple threads writing data with + /// [`write_all`] may result in interleaved output. + /// + /// [`AsyncWrite`]: AsyncWrite + /// [`write_all`]: crate::io::AsyncWriteExt::write_all() + /// + /// # Examples + /// + /// ``` + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut stderr = io::stderr(); + /// stderr.write_all(b"Print some error here.").await?; + /// Ok(()) + /// } + /// ``` + pub fn stderr() -> Stderr { + let std = io::stderr(); + Stderr { + std: SplitByUtf8BoundaryIfWindows::new(Blocking::new(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() + } +} + +#[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() + } +} + +impl AsyncWrite for Stderr { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + Pin::new(&mut self.std).poll_write(cx, buf) + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> { + Pin::new(&mut self.std).poll_flush(cx) + } + + fn poll_shutdown( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), io::Error>> { + Pin::new(&mut self.std).poll_shutdown(cx) + } +} diff --git a/third_party/rust/tokio/src/io/stdin.rs b/third_party/rust/tokio/src/io/stdin.rs new file mode 100644 index 0000000000..c9578f17b6 --- /dev/null +++ b/third_party/rust/tokio/src/io/stdin.rs @@ -0,0 +1,73 @@ +use crate::io::blocking::Blocking; +use crate::io::{AsyncRead, ReadBuf}; + +use std::io; +use std::pin::Pin; +use std::task::Context; +use std::task::Poll; + +cfg_io_std! { + /// A handle to the standard input stream of a process. + /// + /// The handle implements the [`AsyncRead`] trait, but beware that concurrent + /// reads of `Stdin` must be executed with care. + /// + /// This handle is best used for non-interactive uses, such as when a file + /// is piped into the application. For technical reasons, `stdin` is + /// implemented by using an ordinary blocking read on a separate thread, and + /// it is impossible to cancel that read. This can make shutdown of the + /// runtime hang until the user presses enter. + /// + /// For interactive uses, it is recommended to spawn a thread dedicated to + /// user input and use blocking IO directly in that thread. + /// + /// Created by the [`stdin`] function. + /// + /// [`stdin`]: fn@stdin + /// [`AsyncRead`]: trait@AsyncRead + #[derive(Debug)] + pub struct Stdin { + std: Blocking<std::io::Stdin>, + } + + /// Constructs a new handle to the standard input of the current process. + /// + /// This handle is best used for non-interactive uses, such as when a file + /// is piped into the application. For technical reasons, `stdin` is + /// implemented by using an ordinary blocking read on a separate thread, and + /// it is impossible to cancel that read. This can make shutdown of the + /// runtime hang until the user presses enter. + /// + /// For interactive uses, it is recommended to spawn a thread dedicated to + /// user input and use blocking IO directly in that thread. + pub fn stdin() -> Stdin { + let std = io::stdin(); + Stdin { + std: Blocking::new(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() + } +} + +#[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() + } +} + +impl AsyncRead for Stdin { + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<()>> { + Pin::new(&mut self.std).poll_read(cx, buf) + } +} diff --git a/third_party/rust/tokio/src/io/stdio_common.rs b/third_party/rust/tokio/src/io/stdio_common.rs new file mode 100644 index 0000000000..7e4a198a82 --- /dev/null +++ b/third_party/rust/tokio/src/io/stdio_common.rs @@ -0,0 +1,220 @@ +//! Contains utilities for stdout and stderr. +use crate::io::AsyncWrite; +use std::pin::Pin; +use std::task::{Context, Poll}; +/// # Windows +/// AsyncWrite adapter that finds last char boundary in given buffer and does not write the rest, +/// 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. +#[derive(Debug)] +pub(crate) struct SplitByUtf8BoundaryIfWindows<W> { + inner: W, +} + +impl<W> SplitByUtf8BoundaryIfWindows<W> { + pub(crate) fn new(inner: W) -> Self { + Self { inner } + } +} + +// this constant is defined by Unicode standard. +const MAX_BYTES_PER_CHAR: usize = 4; + +// Subject for tweaking here +const MAGIC_CONST: usize = 8; + +impl<W> crate::io::AsyncWrite for SplitByUtf8BoundaryIfWindows<W> +where + W: AsyncWrite + Unpin, +{ + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + mut buf: &[u8], + ) -> Poll<Result<usize, std::io::Error>> { + // just a closure to avoid repetitive code + let mut call_inner = move |buf| Pin::new(&mut self.inner).poll_write(cx, buf); + + // 1. Only windows stdio can suffer from non-utf8. + // We also check for `test` so that we can write some tests + // 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. + // 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 + { + return call_inner(buf); + } + + buf = &buf[..crate::io::blocking::MAX_BUF]; + + // Now there are two possibilities. + // If caller gave is binary buffer, we **should not** shrink it + // anymore, because excessive shrinking hits performance. + // If caller gave as binary buffer, we **must** additionally + // shrink it to strip incomplete char at the end of buffer. + // that's why check we will perform now is allowed to have + // false-positive. + + // Now let's look at the first MAX_BYTES_PER_CHAR * MAGIC_CONST bytes. + // if they are (possibly incomplete) utf8, then we can be quite sure + // that input buffer was utf8. + + let have_to_fix_up = match std::str::from_utf8(&buf[..MAX_BYTES_PER_CHAR * MAGIC_CONST]) { + Ok(_) => true, + Err(err) => { + let incomplete_bytes = MAX_BYTES_PER_CHAR * MAGIC_CONST - err.valid_up_to(); + incomplete_bytes < MAX_BYTES_PER_CHAR + } + }; + + if have_to_fix_up { + // We must pop several bytes at the end which form incomplete + // character. To achieve it, we exploit UTF8 encoding: + // for any code point, all bytes except first start with 0b10 prefix. + // see https://en.wikipedia.org/wiki/UTF-8#Encoding for details + let trailing_incomplete_char_size = buf + .iter() + .rev() + .take(MAX_BYTES_PER_CHAR) + .position(|byte| *byte < 0b1000_0000 || *byte >= 0b1100_0000) + .unwrap_or(0) + + 1; + buf = &buf[..buf.len() - trailing_incomplete_char_size]; + } + + call_inner(buf) + } + + fn poll_flush( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), std::io::Error>> { + Pin::new(&mut self.inner).poll_flush(cx) + } + + fn poll_shutdown( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), std::io::Error>> { + Pin::new(&mut self.inner).poll_shutdown(cx) + } +} + +#[cfg(test)] +#[cfg(not(loom))] +mod tests { + 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 { + fn poll_write( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<Result<usize, io::Error>> { + assert!(buf.len() <= MAX_BUF); + assert!(std::str::from_utf8(buf).is_ok()); + Poll::Ready(Ok(buf.len())) + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll<Result<(), io::Error>> { + Poll::Ready(Ok(())) + } + } + + struct LoggingMockWriter { + write_history: Vec<usize>, + } + + impl LoggingMockWriter { + fn new() -> Self { + LoggingMockWriter { + write_history: Vec::new(), + } + } + } + + impl crate::io::AsyncWrite for LoggingMockWriter { + fn poll_write( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<Result<usize, io::Error>> { + assert!(buf.len() <= MAX_BUF); + self.write_history.push(buf.len()); + Poll::Ready(Ok(buf.len())) + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll<Result<(), io::Error>> { + Poll::Ready(Ok(())) + } + } + + #[test] + fn test_splitter() { + let data = str::repeat("â–ˆ", MAX_BUF); + let mut wr = super::SplitByUtf8BoundaryIfWindows::new(TextMockWriter); + let fut = async move { + wr.write_all(data.as_bytes()).await.unwrap(); + }; + crate::runtime::Builder::new_current_thread() + .build() + .unwrap() + .block_on(fut); + } + + #[test] + 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. + 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)); + let mut writer = LoggingMockWriter::new(); + let mut splitter = super::SplitByUtf8BoundaryIfWindows::new(&mut writer); + crate::runtime::Builder::new_current_thread() + .build() + .unwrap() + .block_on(async { + splitter.write_all(&data).await.unwrap(); + }); + // Check that at most two writes were performed + assert!(writer.write_history.len() <= 2); + // Check that all has been written + assert_eq!( + writer.write_history.iter().copied().sum::<usize>(), + data.len() + ); + // Check that at most MAX_BYTES_PER_CHAR + 1 (i.e. 5) bytes were shrinked + // 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/third_party/rust/tokio/src/io/stdout.rs b/third_party/rust/tokio/src/io/stdout.rs new file mode 100644 index 0000000000..a08ed01eed --- /dev/null +++ b/third_party/rust/tokio/src/io/stdout.rs @@ -0,0 +1,108 @@ +use crate::io::blocking::Blocking; +use crate::io::stdio_common::SplitByUtf8BoundaryIfWindows; +use crate::io::AsyncWrite; +use std::io; +use std::pin::Pin; +use std::task::Context; +use std::task::Poll; + +cfg_io_std! { + /// A handle to the standard output stream of a process. + /// + /// Concurrent writes to stdout must be executed with care: Only individual + /// writes to this [`AsyncWrite`] are guaranteed to be intact. In particular + /// you should be aware that writes using [`write_all`] are not guaranteed + /// to occur as a single write, so multiple threads writing data with + /// [`write_all`] may result in interleaved output. + /// + /// Created by the [`stdout`] function. + /// + /// [`stdout`]: stdout() + /// [`AsyncWrite`]: AsyncWrite + /// [`write_all`]: crate::io::AsyncWriteExt::write_all() + /// + /// # Examples + /// + /// ``` + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut stdout = io::stdout(); + /// stdout.write_all(b"Hello world!").await?; + /// Ok(()) + /// } + /// ``` + #[derive(Debug)] + pub struct Stdout { + std: SplitByUtf8BoundaryIfWindows<Blocking<std::io::Stdout>>, + } + + /// Constructs a new handle to the standard output of the current process. + /// + /// The returned handle allows writing to standard out from the within the + /// Tokio runtime. + /// + /// Concurrent writes to stdout must be executed with care: Only individual + /// writes to this [`AsyncWrite`] are guaranteed to be intact. In particular + /// you should be aware that writes using [`write_all`] are not guaranteed + /// to occur as a single write, so multiple threads writing data with + /// [`write_all`] may result in interleaved output. + /// + /// [`AsyncWrite`]: AsyncWrite + /// [`write_all`]: crate::io::AsyncWriteExt::write_all() + /// + /// # Examples + /// + /// ``` + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut stdout = io::stdout(); + /// stdout.write_all(b"Hello world!").await?; + /// Ok(()) + /// } + /// ``` + pub fn stdout() -> Stdout { + let std = io::stdout(); + Stdout { + std: SplitByUtf8BoundaryIfWindows::new(Blocking::new(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() + } +} + +#[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() + } +} + +impl AsyncWrite for Stdout { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + Pin::new(&mut self.std).poll_write(cx, buf) + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> { + Pin::new(&mut self.std).poll_flush(cx) + } + + fn poll_shutdown( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), io::Error>> { + Pin::new(&mut self.std).poll_shutdown(cx) + } +} diff --git a/third_party/rust/tokio/src/io/util/async_buf_read_ext.rs b/third_party/rust/tokio/src/io/util/async_buf_read_ext.rs new file mode 100644 index 0000000000..b241e354ba --- /dev/null +++ b/third_party/rust/tokio/src/io/util/async_buf_read_ext.rs @@ -0,0 +1,351 @@ +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}; +use crate::io::util::split::{split, Split}; +use crate::io::AsyncBufRead; + +cfg_io_util! { + /// An extension trait which adds utility methods to [`AsyncBufRead`] types. + /// + /// [`AsyncBufRead`]: crate::io::AsyncBufRead + pub trait AsyncBufReadExt: AsyncBufRead { + /// Reads all bytes into `buf` until the delimiter `byte` or EOF is reached. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_until(&mut self, byte: u8, buf: &mut Vec<u8>) -> io::Result<usize>; + /// ``` + /// + /// This function will read bytes from the underlying stream until the + /// delimiter or EOF is found. Once found, all bytes up to, and including, + /// the delimiter (if found) will be appended to `buf`. + /// + /// If successful, this function will return the total number of bytes read. + /// + /// If this function returns `Ok(0)`, the stream has reached EOF. + /// + /// # Errors + /// + /// This function will ignore all instances of [`ErrorKind::Interrupted`] and + /// will otherwise return any errors returned by [`fill_buf`]. + /// + /// If an I/O error is encountered then all bytes read so far will be + /// present in `buf` and its length will have been adjusted appropriately. + /// + /// [`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 + /// this example, we use [`Cursor`] to read all the bytes in a byte slice + /// in hyphen delimited segments: + /// + /// [`Cursor`]: std::io::Cursor + /// + /// ``` + /// use tokio::io::AsyncBufReadExt; + /// + /// use std::io::Cursor; + /// + /// #[tokio::main] + /// async fn main() { + /// let mut cursor = Cursor::new(b"lorem-ipsum"); + /// let mut buf = vec![]; + /// + /// // cursor is at 'l' + /// let num_bytes = cursor.read_until(b'-', &mut buf) + /// .await + /// .expect("reading from cursor won't fail"); + /// + /// assert_eq!(num_bytes, 6); + /// assert_eq!(buf, b"lorem-"); + /// buf.clear(); + /// + /// // cursor is at 'i' + /// let num_bytes = cursor.read_until(b'-', &mut buf) + /// .await + /// .expect("reading from cursor won't fail"); + /// + /// assert_eq!(num_bytes, 5); + /// assert_eq!(buf, b"ipsum"); + /// buf.clear(); + /// + /// // cursor is at EOF + /// let num_bytes = cursor.read_until(b'-', &mut buf) + /// .await + /// .expect("reading from cursor won't fail"); + /// assert_eq!(num_bytes, 0); + /// assert_eq!(buf, b""); + /// } + /// ``` + fn read_until<'a>(&'a mut self, byte: u8, buf: &'a mut Vec<u8>) -> ReadUntil<'a, Self> + where + Self: Unpin, + { + read_until(self, byte, buf) + } + + /// Reads all bytes until a newline (the 0xA byte) is reached, and append + /// them to the provided buffer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_line(&mut self, buf: &mut String) -> io::Result<usize>; + /// ``` + /// + /// This function will read bytes from the underlying stream until the + /// newline delimiter (the 0xA byte) or EOF is found. Once found, all bytes + /// up to, and including, the delimiter (if found) will be appended to + /// `buf`. + /// + /// If successful, this function will return the total number of bytes read. + /// + /// If this function returns `Ok(0)`, the stream has reached EOF. + /// + /// # Errors + /// + /// This function has the same error semantics as [`read_until`] and will + /// also return an error if the read bytes are not valid UTF-8. If an I/O + /// error is encountered then `buf` may contain some bytes already read in + /// the event that all data read so far was valid UTF-8. + /// + /// [`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/0.6/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 + /// `AsyncBufRead`. In this example, we use [`Cursor`] to read all the + /// lines in a byte slice: + /// + /// [`Cursor`]: std::io::Cursor + /// + /// ``` + /// use tokio::io::AsyncBufReadExt; + /// + /// use std::io::Cursor; + /// + /// #[tokio::main] + /// async fn main() { + /// let mut cursor = Cursor::new(b"foo\nbar"); + /// let mut buf = String::new(); + /// + /// // cursor is at 'f' + /// let num_bytes = cursor.read_line(&mut buf) + /// .await + /// .expect("reading from cursor won't fail"); + /// + /// assert_eq!(num_bytes, 4); + /// assert_eq!(buf, "foo\n"); + /// buf.clear(); + /// + /// // cursor is at 'b' + /// let num_bytes = cursor.read_line(&mut buf) + /// .await + /// .expect("reading from cursor won't fail"); + /// + /// assert_eq!(num_bytes, 3); + /// assert_eq!(buf, "bar"); + /// buf.clear(); + /// + /// // cursor is at EOF + /// let num_bytes = cursor.read_line(&mut buf) + /// .await + /// .expect("reading from cursor won't fail"); + /// + /// assert_eq!(num_bytes, 0); + /// assert_eq!(buf, ""); + /// } + /// ``` + fn read_line<'a>(&'a mut self, buf: &'a mut String) -> ReadLine<'a, Self> + where + Self: Unpin, + { + read_line(self, buf) + } + + /// Returns a stream of the contents of this reader split on the byte + /// `byte`. + /// + /// This method is the asynchronous equivalent to + /// [`BufRead::split`](std::io::BufRead::split). + /// + /// The stream returned from this function will yield instances of + /// [`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 + /// + /// Each item of the stream has the same error semantics as + /// [`AsyncBufReadExt::read_until`](AsyncBufReadExt::read_until). + /// + /// # Examples + /// + /// ``` + /// # use tokio::io::AsyncBufRead; + /// use tokio::io::AsyncBufReadExt; + /// + /// # async fn dox(my_buf_read: impl AsyncBufRead + Unpin) -> std::io::Result<()> { + /// let mut segments = my_buf_read.split(b'f'); + /// + /// while let Some(segment) = segments.next_segment().await? { + /// println!("length = {}", segment.len()) + /// } + /// # Ok(()) + /// # } + /// ``` + fn split(self, byte: u8) -> Split<Self> + where + Self: Sized + Unpin, + { + 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`]`<`[`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 + /// + /// Each line of the stream has the same error semantics as [`AsyncBufReadExt::read_line`]. + /// + /// # Examples + /// + /// [`std::io::Cursor`][`Cursor`] is a type that implements `BufRead`. In + /// this example, we use [`Cursor`] to iterate over all the lines in a byte + /// slice. + /// + /// [`Cursor`]: std::io::Cursor + /// + /// ``` + /// use tokio::io::AsyncBufReadExt; + /// + /// use std::io::Cursor; + /// + /// #[tokio::main] + /// async fn main() { + /// let cursor = Cursor::new(b"lorem\nipsum\r\ndolor"); + /// + /// let mut lines = cursor.lines(); + /// + /// assert_eq!(lines.next_line().await.unwrap(), Some(String::from("lorem"))); + /// assert_eq!(lines.next_line().await.unwrap(), Some(String::from("ipsum"))); + /// assert_eq!(lines.next_line().await.unwrap(), Some(String::from("dolor"))); + /// assert_eq!(lines.next_line().await.unwrap(), None); + /// } + /// ``` + /// + /// [`AsyncBufReadExt::read_line`]: AsyncBufReadExt::read_line + fn lines(self) -> Lines<Self> + where + Self: Sized, + { + lines(self) + } + } +} + +impl<R: AsyncBufRead + ?Sized> AsyncBufReadExt for R {} diff --git a/third_party/rust/tokio/src/io/util/async_read_ext.rs b/third_party/rust/tokio/src/io/util/async_read_ext.rs new file mode 100644 index 0000000000..df5445c2c6 --- /dev/null +++ b/third_party/rust/tokio/src/io/util/async_read_ext.rs @@ -0,0 +1,1294 @@ +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, +}; +use crate::io::util::read_int::{ + ReadU128, ReadU128Le, ReadU16, ReadU16Le, ReadU32, ReadU32Le, ReadU64, ReadU64Le, ReadU8, +}; +use crate::io::util::read_to_end::{read_to_end, ReadToEnd}; +use crate::io::util::read_to_string::{read_to_string, ReadToString}; +use crate::io::util::take::{take, Take}; +use crate::io::AsyncRead; + +use bytes::BufMut; + +cfg_io_util! { + /// Defines numeric reader + macro_rules! read_impl { + ( + $( + $(#[$outer:meta])* + fn $name:ident(&mut self) -> $($fut:ident)*; + )* + ) => { + $( + $(#[$outer])* + fn $name<'a>(&'a mut self) -> $($fut)*<&'a mut Self> where Self: Unpin { + $($fut)*::new(self) + } + )* + } + } + + /// Reads bytes from a source. + /// + /// Implemented as an extension trait, adding utility methods to all + /// [`AsyncRead`] types. Callers will tend to import this trait instead of + /// [`AsyncRead`]. + /// + /// ```no_run + /// use tokio::fs::File; + /// use tokio::io::{self, AsyncReadExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut f = File::open("foo.txt").await?; + /// let mut buffer = [0; 10]; + /// + /// // The `read` method is defined by this trait. + /// let n = f.read(&mut buffer[..]).await?; + /// + /// Ok(()) + /// } + /// ``` + /// + /// See [module][crate::io] documentation for more details. + /// + /// [`AsyncRead`]: AsyncRead + pub trait AsyncReadExt: AsyncRead { + /// Creates a new `AsyncRead` instance that chains this stream with + /// `next`. + /// + /// The returned `AsyncRead` instance will first read all bytes from this object + /// until EOF is encountered. Afterwards the output is equivalent to the + /// output of `next`. + /// + /// # Examples + /// + /// [`File`][crate::fs::File]s implement `AsyncRead`: + /// + /// ```no_run + /// use tokio::fs::File; + /// use tokio::io::{self, AsyncReadExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let f1 = File::open("foo.txt").await?; + /// let f2 = File::open("bar.txt").await?; + /// + /// let mut handle = f1.chain(f2); + /// let mut buffer = String::new(); + /// + /// // read the value into a String. We could use any AsyncRead + /// // method here, this is just one example. + /// handle.read_to_string(&mut buffer).await?; + /// Ok(()) + /// } + /// ``` + fn chain<R>(self, next: R) -> Chain<Self, R> + where + Self: Sized, + R: AsyncRead, + { + chain(self, next) + } + + /// Pulls some bytes from this source into the specified buffer, + /// returning how many bytes were read. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read(&mut self, buf: &mut [u8]) -> io::Result<usize>; + /// ``` + /// + /// This method does not provide any guarantees about whether it + /// completes immediately or asynchronously. + /// + /// # Return + /// + /// If the return value of this method is `Ok(n)`, then it must be + /// guaranteed that `0 <= n <= buf.len()`. A nonzero `n` value indicates + /// that the buffer `buf` has been filled in with `n` bytes of data from + /// this source. If `n` is `0`, then it can indicate one of two + /// scenarios: + /// + /// 1. This reader has reached its "end of file" and will likely no longer + /// be able to produce bytes. Note that this does not mean that the + /// reader will *always* no longer be able to produce bytes. + /// 2. The buffer specified was 0 bytes in length. + /// + /// No guarantees are provided about the contents of `buf` when this + /// function is called, implementations cannot rely on any property of the + /// contents of `buf` being `true`. It is recommended that *implementations* + /// only write data to `buf` instead of reading its contents. + /// + /// Correspondingly, however, *callers* of this method may not assume + /// any guarantees about how the implementation uses `buf`. It is + /// possible that the code that's supposed to write to the buffer might + /// also read from it. It is your responsibility to make sure that `buf` + /// is initialized before calling `read`. + /// + /// # Errors + /// + /// If this function encounters any form of I/O or other error, an error + /// variant will be returned. If an error is returned then it must be + /// guaranteed that no bytes were read. + /// + /// # 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 data was read. + /// + /// # Examples + /// + /// [`File`][crate::fs::File]s implement `Read`: + /// + /// ```no_run + /// use tokio::fs::File; + /// use tokio::io::{self, AsyncReadExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut f = File::open("foo.txt").await?; + /// let mut buffer = [0; 10]; + /// + /// // read up to 10 bytes + /// let n = f.read(&mut buffer[..]).await?; + /// + /// println!("The bytes: {:?}", &buffer[..n]); + /// Ok(()) + /// } + /// ``` + fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> Read<'a, Self> + where + Self: Unpin, + { + read(self, buf) + } + + /// Pulls some bytes from this source into the specified buffer, + /// advancing the buffer's internal cursor. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_buf<B: BufMut>(&mut self, buf: &mut B) -> io::Result<usize>; + /// ``` + /// + /// Usually, only a single `read` syscall is issued, even if there is + /// more space in the supplied buffer. + /// + /// This method does not provide any guarantees about whether it + /// completes immediately or asynchronously. + /// + /// # Return + /// + /// A nonzero `n` value indicates that the buffer `buf` has been filled + /// in with `n` bytes of data from this source. If `n` is `0`, then it + /// can indicate one of two scenarios: + /// + /// 1. This reader has reached its "end of file" and will likely no longer + /// be able to produce bytes. Note that this does not mean that the + /// reader will *always* no longer be able to produce bytes. + /// 2. The buffer specified had a remaining capacity of zero. + /// + /// # Errors + /// + /// If this function encounters any form of I/O or other error, an error + /// variant will be returned. If an error is returned then it must be + /// guaranteed that no bytes were read. + /// + /// # 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 data was read. + /// + /// # Examples + /// + /// [`File`] implements `Read` and [`BytesMut`] implements [`BufMut`]: + /// + /// [`File`]: crate::fs::File + /// [`BytesMut`]: bytes::BytesMut + /// [`BufMut`]: bytes::BufMut + /// + /// ```no_run + /// use tokio::fs::File; + /// use tokio::io::{self, AsyncReadExt}; + /// + /// use bytes::BytesMut; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut f = File::open("foo.txt").await?; + /// let mut buffer = BytesMut::with_capacity(10); + /// + /// assert!(buffer.is_empty()); + /// + /// // read up to 10 bytes, note that the return value is not needed + /// // to access the data that was read as `buffer`'s internal + /// // cursor is updated. + /// f.read_buf(&mut buffer).await?; + /// + /// println!("The bytes: {:?}", &buffer[..]); + /// Ok(()) + /// } + /// ``` + fn read_buf<'a, B>(&'a mut self, buf: &'a mut B) -> ReadBuf<'a, Self, B> + where + Self: Sized + Unpin, + B: BufMut, + { + read_buf(self, buf) + } + + /// Reads the exact number of bytes required to fill `buf`. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<usize>; + /// ``` + /// + /// This function reads as many bytes as necessary to completely fill + /// the specified buffer `buf`. + /// + /// # Errors + /// + /// If the operation encounters an "end of file" before completely + /// filling the buffer, it returns an error of the kind + /// [`ErrorKind::UnexpectedEof`]. The contents of `buf` are unspecified + /// in this case. + /// + /// If any other read error is encountered then the operation + /// immediately returns. The contents of `buf` are unspecified in this + /// case. + /// + /// If this operation returns an error, it is unspecified how many bytes + /// it has read, but it will never read more than would be necessary to + /// completely fill the buffer. + /// + /// # 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 already have been + /// read into `buf`. + /// + /// # Examples + /// + /// [`File`][crate::fs::File]s implement `Read`: + /// + /// ```no_run + /// use tokio::fs::File; + /// use tokio::io::{self, AsyncReadExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut f = File::open("foo.txt").await?; + /// let mut buffer = [0; 10]; + /// + /// // read exactly 10 bytes + /// f.read_exact(&mut buffer).await?; + /// Ok(()) + /// } + /// ``` + /// + /// [`ErrorKind::UnexpectedEof`]: std::io::ErrorKind::UnexpectedEof + fn read_exact<'a>(&'a mut self, buf: &'a mut [u8]) -> ReadExact<'a, Self> + where + Self: Unpin, + { + read_exact(self, buf) + } + + read_impl! { + /// Reads an unsigned 8 bit integer from the underlying reader. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_u8(&mut self) -> io::Result<u8>; + /// ``` + /// + /// 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 unsigned 8 bit integers from an `AsyncRead`: + /// + /// ```rust + /// use tokio::io::{self, AsyncReadExt}; + /// + /// use std::io::Cursor; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut reader = Cursor::new(vec![2, 5]); + /// + /// assert_eq!(2, reader.read_u8().await?); + /// assert_eq!(5, reader.read_u8().await?); + /// + /// Ok(()) + /// } + /// ``` + fn read_u8(&mut self) -> ReadU8; + + /// Reads a signed 8 bit integer from the underlying reader. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_i8(&mut self) -> io::Result<i8>; + /// ``` + /// + /// 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 unsigned 8 bit integers from an `AsyncRead`: + /// + /// ```rust + /// use tokio::io::{self, AsyncReadExt}; + /// + /// use std::io::Cursor; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut reader = Cursor::new(vec![0x02, 0xfb]); + /// + /// assert_eq!(2, reader.read_i8().await?); + /// assert_eq!(-5, reader.read_i8().await?); + /// + /// Ok(()) + /// } + /// ``` + fn read_i8(&mut self) -> ReadI8; + + /// Reads an unsigned 16-bit integer in big-endian order from the + /// underlying reader. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_u16(&mut self) -> io::Result<u16>; + /// ``` + /// + /// 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 unsigned 16 bit big-endian integers 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![2, 5, 3, 0]); + /// + /// assert_eq!(517, reader.read_u16().await?); + /// assert_eq!(768, reader.read_u16().await?); + /// Ok(()) + /// } + /// ``` + fn read_u16(&mut self) -> ReadU16; + + /// Reads a signed 16-bit integer in big-endian order from the + /// underlying reader. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_i16(&mut self) -> io::Result<i16>; + /// ``` + /// + /// 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 signed 16 bit big-endian integers 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![0x00, 0xc1, 0xff, 0x7c]); + /// + /// assert_eq!(193, reader.read_i16().await?); + /// assert_eq!(-132, reader.read_i16().await?); + /// Ok(()) + /// } + /// ``` + fn read_i16(&mut self) -> ReadI16; + + /// Reads an unsigned 32-bit integer in big-endian order from the + /// underlying reader. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_u32(&mut self) -> io::Result<u32>; + /// ``` + /// + /// 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 unsigned 32-bit big-endian integers 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![0x00, 0x00, 0x01, 0x0b]); + /// + /// assert_eq!(267, reader.read_u32().await?); + /// Ok(()) + /// } + /// ``` + fn read_u32(&mut self) -> ReadU32; + + /// Reads a signed 32-bit integer in big-endian order from the + /// underlying reader. + /// + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_i32(&mut self) -> io::Result<i32>; + /// ``` + /// + /// 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 signed 32-bit big-endian integers 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, 0x7a, 0x33]); + /// + /// assert_eq!(-34253, reader.read_i32().await?); + /// Ok(()) + /// } + /// ``` + fn read_i32(&mut self) -> ReadI32; + + /// Reads an unsigned 64-bit integer in big-endian order from the + /// underlying reader. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_u64(&mut self) -> io::Result<u64>; + /// ``` + /// + /// 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 unsigned 64-bit big-endian integers 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![ + /// 0x00, 0x03, 0x43, 0x95, 0x4d, 0x60, 0x86, 0x83 + /// ]); + /// + /// assert_eq!(918733457491587, reader.read_u64().await?); + /// Ok(()) + /// } + /// ``` + fn read_u64(&mut self) -> ReadU64; + + /// Reads an signed 64-bit integer in big-endian order from the + /// underlying reader. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_i64(&mut self) -> io::Result<i64>; + /// ``` + /// + /// 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 signed 64-bit big-endian integers 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![0x80, 0, 0, 0, 0, 0, 0, 0]); + /// + /// assert_eq!(i64::MIN, reader.read_i64().await?); + /// Ok(()) + /// } + /// ``` + fn read_i64(&mut self) -> ReadI64; + + /// Reads an unsigned 128-bit integer in big-endian order from the + /// underlying reader. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_u128(&mut self) -> io::Result<u128>; + /// ``` + /// + /// 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 unsigned 128-bit big-endian integers 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![ + /// 0x00, 0x03, 0x43, 0x95, 0x4d, 0x60, 0x86, 0x83, + /// 0x00, 0x03, 0x43, 0x95, 0x4d, 0x60, 0x86, 0x83 + /// ]); + /// + /// assert_eq!(16947640962301618749969007319746179, reader.read_u128().await?); + /// Ok(()) + /// } + /// ``` + fn read_u128(&mut self) -> ReadU128; + + /// Reads an signed 128-bit integer in big-endian order from the + /// underlying reader. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_i128(&mut self) -> io::Result<i128>; + /// ``` + /// + /// 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 signed 128-bit big-endian integers 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![ + /// 0x80, 0, 0, 0, 0, 0, 0, 0, + /// 0, 0, 0, 0, 0, 0, 0, 0 + /// ]); + /// + /// assert_eq!(i128::MIN, reader.read_i128().await?); + /// Ok(()) + /// } + /// ``` + 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. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_u16_le(&mut self) -> io::Result<u16>; + /// ``` + /// + /// 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 unsigned 16 bit little-endian integers 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![2, 5, 3, 0]); + /// + /// assert_eq!(1282, reader.read_u16_le().await?); + /// assert_eq!(3, reader.read_u16_le().await?); + /// Ok(()) + /// } + /// ``` + fn read_u16_le(&mut self) -> ReadU16Le; + + /// Reads a signed 16-bit integer in little-endian order from the + /// underlying reader. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_i16_le(&mut self) -> io::Result<i16>; + /// ``` + /// + /// 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 signed 16 bit little-endian integers 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![0x00, 0xc1, 0xff, 0x7c]); + /// + /// assert_eq!(-16128, reader.read_i16_le().await?); + /// assert_eq!(31999, reader.read_i16_le().await?); + /// Ok(()) + /// } + /// ``` + fn read_i16_le(&mut self) -> ReadI16Le; + + /// Reads an unsigned 32-bit integer in little-endian order from the + /// underlying reader. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_u32_le(&mut self) -> io::Result<u32>; + /// ``` + /// + /// 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 unsigned 32-bit little-endian integers 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![0x00, 0x00, 0x01, 0x0b]); + /// + /// assert_eq!(184614912, reader.read_u32_le().await?); + /// Ok(()) + /// } + /// ``` + fn read_u32_le(&mut self) -> ReadU32Le; + + /// Reads a signed 32-bit integer in little-endian order from the + /// underlying reader. + /// + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_i32_le(&mut self) -> io::Result<i32>; + /// ``` + /// + /// 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 signed 32-bit little-endian integers 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, 0x7a, 0x33]); + /// + /// assert_eq!(863698943, reader.read_i32_le().await?); + /// Ok(()) + /// } + /// ``` + fn read_i32_le(&mut self) -> ReadI32Le; + + /// Reads an unsigned 64-bit integer in little-endian order from the + /// underlying reader. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_u64_le(&mut self) -> io::Result<u64>; + /// ``` + /// + /// 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 unsigned 64-bit little-endian integers 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![ + /// 0x00, 0x03, 0x43, 0x95, 0x4d, 0x60, 0x86, 0x83 + /// ]); + /// + /// assert_eq!(9477368352180732672, reader.read_u64_le().await?); + /// Ok(()) + /// } + /// ``` + fn read_u64_le(&mut self) -> ReadU64Le; + + /// Reads an signed 64-bit integer in little-endian order from the + /// underlying reader. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_i64_le(&mut self) -> io::Result<i64>; + /// ``` + /// + /// 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 signed 64-bit little-endian integers 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![0x80, 0, 0, 0, 0, 0, 0, 0]); + /// + /// assert_eq!(128, reader.read_i64_le().await?); + /// Ok(()) + /// } + /// ``` + fn read_i64_le(&mut self) -> ReadI64Le; + + /// Reads an unsigned 128-bit integer in little-endian order from the + /// underlying reader. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_u128_le(&mut self) -> io::Result<u128>; + /// ``` + /// + /// 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 unsigned 128-bit little-endian integers 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![ + /// 0x00, 0x03, 0x43, 0x95, 0x4d, 0x60, 0x86, 0x83, + /// 0x00, 0x03, 0x43, 0x95, 0x4d, 0x60, 0x86, 0x83 + /// ]); + /// + /// assert_eq!(174826588484952389081207917399662330624, reader.read_u128_le().await?); + /// Ok(()) + /// } + /// ``` + fn read_u128_le(&mut self) -> ReadU128Le; + + /// Reads an signed 128-bit integer in little-endian order from the + /// underlying reader. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_i128_le(&mut self) -> io::Result<i128>; + /// ``` + /// + /// 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 signed 128-bit little-endian integers 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![ + /// 0x80, 0, 0, 0, 0, 0, 0, 0, + /// 0, 0, 0, 0, 0, 0, 0, 0 + /// ]); + /// + /// assert_eq!(128, reader.read_i128_le().await?); + /// Ok(()) + /// } + /// ``` + 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`. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize>; + /// ``` + /// + /// All bytes read from this source will be appended to the specified + /// buffer `buf`. This function will continuously call [`read()`] to + /// append more data to `buf` until [`read()`] returns `Ok(0)`. + /// + /// If successful, the total number of bytes read is returned. + /// + /// [`read()`]: AsyncReadExt::read + /// + /// # Errors + /// + /// If a read error is encountered then the `read_to_end` operation + /// immediately completes. Any bytes which have already been read will + /// be appended to `buf`. + /// + /// # Examples + /// + /// [`File`][crate::fs::File]s implement `Read`: + /// + /// ```no_run + /// use tokio::io::{self, AsyncReadExt}; + /// use tokio::fs::File; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut f = File::open("foo.txt").await?; + /// let mut buffer = Vec::new(); + /// + /// // read the whole file + /// f.read_to_end(&mut buffer).await?; + /// Ok(()) + /// } + /// ``` + /// + /// (See also the [`tokio::fs::read`] convenience function for reading from a + /// file.) + /// + /// [`tokio::fs::read`]: fn@crate::fs::read + fn read_to_end<'a>(&'a mut self, buf: &'a mut Vec<u8>) -> ReadToEnd<'a, Self> + where + Self: Unpin, + { + read_to_end(self, buf) + } + + /// Reads all bytes until EOF in this source, appending them to `buf`. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn read_to_string(&mut self, buf: &mut String) -> io::Result<usize>; + /// ``` + /// + /// If successful, the number of bytes which were read and appended to + /// `buf` is returned. + /// + /// # Errors + /// + /// If the data in this stream is *not* valid UTF-8 then an error is + /// returned and `buf` is unchanged. + /// + /// See [`read_to_end`][AsyncReadExt::read_to_end] for other error semantics. + /// + /// # Examples + /// + /// [`File`][crate::fs::File]s implement `Read`: + /// + /// ```no_run + /// use tokio::io::{self, AsyncReadExt}; + /// use tokio::fs::File; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut f = File::open("foo.txt").await?; + /// let mut buffer = String::new(); + /// + /// f.read_to_string(&mut buffer).await?; + /// Ok(()) + /// } + /// ``` + /// + /// (See also the [`crate::fs::read_to_string`] convenience function for + /// reading from a file.) + /// + /// [`crate::fs::read_to_string`]: fn@crate::fs::read_to_string + fn read_to_string<'a>(&'a mut self, dst: &'a mut String) -> ReadToString<'a, Self> + where + Self: Unpin, + { + read_to_string(self, dst) + } + + /// Creates an adaptor which reads at most `limit` bytes from it. + /// + /// This function returns a new instance of `AsyncRead` which will read + /// at most `limit` bytes, after which it will always return EOF + /// (`Ok(0)`). Any read errors will not count towards the number of + /// bytes read and future calls to [`read()`] may succeed. + /// + /// [`read()`]: fn@crate::io::AsyncReadExt::read + /// + /// [read]: AsyncReadExt::read + /// + /// # Examples + /// + /// [`File`][crate::fs::File]s implement `Read`: + /// + /// ```no_run + /// use tokio::io::{self, AsyncReadExt}; + /// use tokio::fs::File; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let f = File::open("foo.txt").await?; + /// let mut buffer = [0; 5]; + /// + /// // read at most five bytes + /// let mut handle = f.take(5); + /// + /// handle.read(&mut buffer).await?; + /// Ok(()) + /// } + /// ``` + fn take(self, limit: u64) -> Take<Self> + where + Self: Sized, + { + take(self, limit) + } + } +} + +impl<R: AsyncRead + ?Sized> AsyncReadExt for R {} diff --git a/third_party/rust/tokio/src/io/util/async_seek_ext.rs b/third_party/rust/tokio/src/io/util/async_seek_ext.rs new file mode 100644 index 0000000000..46b3e6c0d3 --- /dev/null +++ b/third_party/rust/tokio/src/io/util/async_seek_ext.rs @@ -0,0 +1,93 @@ +use crate::io::seek::{seek, Seek}; +use crate::io::AsyncSeek; +use std::io::SeekFrom; + +cfg_io_util! { + /// An extension trait that adds utility methods to [`AsyncSeek`] types. + /// + /// # Examples + /// + /// ``` + /// use std::io::{self, Cursor, SeekFrom}; + /// use tokio::io::{AsyncSeekExt, AsyncReadExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut cursor = Cursor::new(b"abcdefg"); + /// + /// // the `seek` method is defined by this trait + /// cursor.seek(SeekFrom::Start(3)).await?; + /// + /// let mut buf = [0; 1]; + /// let n = cursor.read(&mut buf).await?; + /// assert_eq!(n, 1); + /// assert_eq!(buf, [b'd']); + /// + /// Ok(()) + /// } + /// ``` + /// + /// See [module][crate::io] documentation for more details. + /// + /// [`AsyncSeek`]: AsyncSeek + pub trait AsyncSeekExt: AsyncSeek { + /// Creates a future which will seek an IO object, and then yield the + /// new position in the object and the object itself. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn seek(&mut self, pos: SeekFrom) -> io::Result<u64>; + /// ``` + /// + /// In the case of an error the buffer and the object will be discarded, with + /// the error yielded. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::fs::File; + /// use tokio::io::{AsyncSeekExt, AsyncReadExt}; + /// + /// use std::io::SeekFrom; + /// + /// # async fn dox() -> std::io::Result<()> { + /// let mut file = File::open("foo.txt").await?; + /// file.seek(SeekFrom::Start(6)).await?; + /// + /// let mut contents = vec![0u8; 10]; + /// file.read_exact(&mut contents).await?; + /// # Ok(()) + /// # } + /// ``` + fn seek(&mut self, pos: SeekFrom) -> Seek<'_, Self> + where + Self: Unpin, + { + seek(self, pos) + } + + /// Creates a future which will rewind to the beginning of the stream. + /// + /// This is convenience method, equivalent to 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. + /// + /// This is equivalent to `self.seek(SeekFrom::Current(0))`. + fn stream_position(&mut self) -> Seek<'_, Self> + where + Self: Unpin, + { + self.seek(SeekFrom::Current(0)) + } + } +} + +impl<S: AsyncSeek + ?Sized> AsyncSeekExt for S {} diff --git a/third_party/rust/tokio/src/io/util/async_write_ext.rs b/third_party/rust/tokio/src/io/util/async_write_ext.rs new file mode 100644 index 0000000000..93a318315e --- /dev/null +++ b/third_party/rust/tokio/src/io/util/async_write_ext.rs @@ -0,0 +1,1293 @@ +use crate::io::util::flush::{flush, Flush}; +use crate::io::util::shutdown::{shutdown, Shutdown}; +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, +}; +use crate::io::util::write_int::{ + WriteU128, WriteU128Le, WriteU16, WriteU16Le, WriteU32, WriteU32Le, WriteU64, WriteU64Le, + WriteU8, +}; +use crate::io::util::write_vectored::{write_vectored, WriteVectored}; +use crate::io::AsyncWrite; +use std::io::IoSlice; + +use bytes::Buf; + +cfg_io_util! { + /// Defines numeric writer. + macro_rules! write_impl { + ( + $( + $(#[$outer:meta])* + fn $name:ident(&mut self, n: $ty:ty) -> $($fut:ident)*; + )* + ) => { + $( + $(#[$outer])* + fn $name<'a>(&'a mut self, n: $ty) -> $($fut)*<&'a mut Self> where Self: Unpin { + $($fut)*::new(self, n) + } + )* + } + } + + /// Writes bytes to a sink. + /// + /// Implemented as an extension trait, adding utility methods to all + /// [`AsyncWrite`] types. Callers will tend to import this trait instead of + /// [`AsyncWrite`]. + /// + /// ```no_run + /// use tokio::io::{self, AsyncWriteExt}; + /// use tokio::fs::File; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let data = b"some bytes"; + /// + /// let mut pos = 0; + /// let mut buffer = File::create("foo.txt").await?; + /// + /// while pos < data.len() { + /// let bytes_written = buffer.write(&data[pos..]).await?; + /// pos += bytes_written; + /// } + /// + /// Ok(()) + /// } + /// ``` + /// + /// See [module][crate::io] documentation for more details. + /// + /// [`AsyncWrite`]: AsyncWrite + pub trait AsyncWriteExt: AsyncWrite { + /// Writes a buffer into this writer, returning how many bytes were + /// written. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write(&mut self, buf: &[u8]) -> io::Result<usize>; + /// ``` + /// + /// This function will attempt to write the entire contents of `buf`, but + /// the entire write may not succeed, or the write may also generate an + /// error. A call to `write` represents *at most one* attempt to write to + /// any wrapped object. + /// + /// # Return + /// + /// If the return value is `Ok(n)` then it must be guaranteed that `n <= + /// buf.len()`. A return value of `0` typically means that the + /// underlying object is no longer able to accept bytes and will likely + /// not be able to in the future as well, or that the buffer provided is + /// empty. + /// + /// # Errors + /// + /// Each call to `write` may generate an I/O error indicating that the + /// operation could not be completed. If an error is returned then no bytes + /// in the buffer were written to this writer. + /// + /// It is **not** considered an error if the entire buffer could not be + /// written to this writer. + /// + /// # Cancel safety + /// + /// This method is cancellation safe in the sense that if it is used as + /// the event in a [`tokio::select!`](crate::select) statement and some + /// other branch completes first, then it is guaranteed that no data was + /// written to this `AsyncWrite`. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::io::{self, AsyncWriteExt}; + /// use tokio::fs::File; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut file = File::create("foo.txt").await?; + /// + /// // Writes some prefix of the byte string, not necessarily all of it. + /// file.write(b"some bytes").await?; + /// Ok(()) + /// } + /// ``` + fn write<'a>(&'a mut self, src: &'a [u8]) -> Write<'a, Self> + where + Self: Unpin, + { + write(self, src) + } + + /// Like [`write`], except that it writes from a slice of buffers. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize>; + /// ``` + /// + /// See [`AsyncWrite::poll_write_vectored`] for more details. + /// + /// # Cancel safety + /// + /// This method is cancellation safe in the sense that if it is used as + /// the event in a [`tokio::select!`](crate::select) statement and some + /// other branch completes first, then it is guaranteed that no data was + /// written to this `AsyncWrite`. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::io::{self, AsyncWriteExt}; + /// use tokio::fs::File; + /// use std::io::IoSlice; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut file = File::create("foo.txt").await?; + /// + /// let bufs: &[_] = &[ + /// IoSlice::new(b"hello"), + /// IoSlice::new(b" "), + /// IoSlice::new(b"world"), + /// ]; + /// + /// file.write_vectored(&bufs).await?; + /// + /// Ok(()) + /// } + /// ``` + /// + /// [`write`]: AsyncWriteExt::write + fn write_vectored<'a, 'b>(&'a mut self, bufs: &'a [IoSlice<'b>]) -> WriteVectored<'a, 'b, Self> + where + Self: Unpin, + { + write_vectored(self, bufs) + } + + /// Writes a buffer into this writer, advancing the buffer's internal + /// cursor. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_buf<B: Buf>(&mut self, buf: &mut B) -> io::Result<usize>; + /// ``` + /// + /// This function will attempt to write the entire contents of `buf`, but + /// the entire write may not succeed, or the write may also generate an + /// error. After the operation completes, the buffer's + /// internal cursor is advanced by the number of bytes written. A + /// subsequent call to `write_buf` using the **same** `buf` value will + /// resume from the point that the first call to `write_buf` completed. + /// A call to `write_buf` represents *at most one* attempt to write to any + /// wrapped object. + /// + /// # Return + /// + /// If the return value is `Ok(n)` then it must be guaranteed that `n <= + /// buf.len()`. A return value of `0` typically means that the + /// underlying object is no longer able to accept bytes and will likely + /// not be able to in the future as well, or that the buffer provided is + /// empty. + /// + /// # Errors + /// + /// Each call to `write` may generate an I/O error indicating that the + /// operation could not be completed. If an error is returned then no bytes + /// in the buffer were written to this writer. + /// + /// It is **not** considered an error if the entire buffer could not be + /// written to this writer. + /// + /// # Cancel safety + /// + /// This method is cancellation safe in the sense that if it is used as + /// the event in a [`tokio::select!`](crate::select) statement and some + /// other branch completes first, then it is guaranteed that no data was + /// written to this `AsyncWrite`. + /// + /// # Examples + /// + /// [`File`] implements [`AsyncWrite`] and [`Cursor`]`<&[u8]>` implements [`Buf`]: + /// + /// [`File`]: crate::fs::File + /// [`Buf`]: bytes::Buf + /// [`Cursor`]: std::io::Cursor + /// + /// ```no_run + /// use tokio::io::{self, AsyncWriteExt}; + /// use tokio::fs::File; + /// + /// use bytes::Buf; + /// use std::io::Cursor; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut file = File::create("foo.txt").await?; + /// let mut buffer = Cursor::new(b"data to write"); + /// + /// // Loop until the entire contents of the buffer are written to + /// // the file. + /// while buffer.has_remaining() { + /// // Writes some prefix of the byte string, not necessarily + /// // all of it. + /// file.write_buf(&mut buffer).await?; + /// } + /// + /// Ok(()) + /// } + /// ``` + fn write_buf<'a, B>(&'a mut self, src: &'a mut B) -> WriteBuf<'a, Self, B> + where + Self: Sized + Unpin, + B: Buf, + { + write_buf(self, src) + } + + /// Attempts to write an entire buffer into this writer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_all_buf(&mut self, buf: impl Buf) -> Result<(), io::Error> { + /// while buf.has_remaining() { + /// self.write_buf(&mut buf).await?; + /// } + /// Ok(()) + /// } + /// ``` + /// + /// This method will continuously call [`write`] until + /// [`buf.has_remaining()`](bytes::Buf::has_remaining) returns false. This method will not + /// return until the entire buffer has been successfully written or an error occurs. The + /// first error generated will be returned. + /// + /// The buffer is advanced after each chunk is successfully written. After failure, + /// `src.chunk()` will return the chunk that failed to write. + /// + /// # Cancel safety + /// + /// If `write_all_buf` is used as the event in a + /// [`tokio::select!`](crate::select) statement and some other branch + /// completes first, then the data in the provided buffer may have been + /// partially written. However, it is guaranteed that the provided + /// buffer has been [advanced] by the amount of bytes that have been + /// partially written. + /// + /// # Examples + /// + /// [`File`] implements [`AsyncWrite`] and [`Cursor`]`<&[u8]>` implements [`Buf`]: + /// + /// [`File`]: crate::fs::File + /// [`Buf`]: bytes::Buf + /// [`Cursor`]: std::io::Cursor + /// [advanced]: bytes::Buf::advance + /// + /// ```no_run + /// use tokio::io::{self, AsyncWriteExt}; + /// use tokio::fs::File; + /// + /// use std::io::Cursor; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut file = File::create("foo.txt").await?; + /// let mut buffer = Cursor::new(b"data to write"); + /// + /// file.write_all_buf(&mut buffer).await?; + /// Ok(()) + /// } + /// ``` + /// + /// [`write`]: AsyncWriteExt::write + fn write_all_buf<'a, B>(&'a mut self, src: &'a mut B) -> WriteAllBuf<'a, Self, B> + where + Self: Sized + Unpin, + B: Buf, + { + write_all_buf(self, src) + } + + /// Attempts to write an entire buffer into this writer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_all(&mut self, buf: &[u8]) -> io::Result<()>; + /// ``` + /// + /// This method will continuously call [`write`] until there is no more data + /// to be written. This method will not return until the entire buffer + /// has been successfully written or such an error occurs. The first + /// error generated from this method will be returned. + /// + /// # Cancel safety + /// + /// This method is not cancellation safe. If it is used as the event + /// in a [`tokio::select!`](crate::select) statement and some other + /// branch completes first, then the provided buffer may have been + /// partially written, but future calls to `write_all` will start over + /// from the beginning of the buffer. + /// + /// # Errors + /// + /// This function will return the first error that [`write`] returns. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::io::{self, AsyncWriteExt}; + /// use tokio::fs::File; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut file = File::create("foo.txt").await?; + /// + /// file.write_all(b"some bytes").await?; + /// Ok(()) + /// } + /// ``` + /// + /// [`write`]: AsyncWriteExt::write + fn write_all<'a>(&'a mut self, src: &'a [u8]) -> WriteAll<'a, Self> + where + Self: Unpin, + { + write_all(self, src) + } + + write_impl! { + /// Writes an unsigned 8-bit integer to the underlying writer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_u8(&mut self, n: u8) -> 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 unsigned 8 bit integers to a `AsyncWrite`: + /// + /// ```rust + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut writer = Vec::new(); + /// + /// writer.write_u8(2).await?; + /// writer.write_u8(5).await?; + /// + /// assert_eq!(writer, b"\x02\x05"); + /// Ok(()) + /// } + /// ``` + fn write_u8(&mut self, n: u8) -> WriteU8; + + /// Writes an unsigned 8-bit integer to the underlying writer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_i8(&mut self, n: i8) -> 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 unsigned 8 bit integers to a `AsyncWrite`: + /// + /// ```rust + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut writer = Vec::new(); + /// + /// writer.write_u8(2).await?; + /// writer.write_u8(5).await?; + /// + /// assert_eq!(writer, b"\x02\x05"); + /// Ok(()) + /// } + /// ``` + fn write_i8(&mut self, n: i8) -> WriteI8; + + /// Writes an unsigned 16-bit integer in big-endian order to the + /// underlying writer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_u16(&mut self, n: u16) -> 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 unsigned 16-bit integers to a `AsyncWrite`: + /// + /// ```rust + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut writer = Vec::new(); + /// + /// writer.write_u16(517).await?; + /// writer.write_u16(768).await?; + /// + /// assert_eq!(writer, b"\x02\x05\x03\x00"); + /// Ok(()) + /// } + /// ``` + fn write_u16(&mut self, n: u16) -> WriteU16; + + /// Writes a signed 16-bit integer in big-endian order to the + /// underlying writer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_i16(&mut self, n: i16) -> 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 signed 16-bit integers to a `AsyncWrite`: + /// + /// ```rust + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut writer = Vec::new(); + /// + /// writer.write_i16(193).await?; + /// writer.write_i16(-132).await?; + /// + /// assert_eq!(writer, b"\x00\xc1\xff\x7c"); + /// Ok(()) + /// } + /// ``` + fn write_i16(&mut self, n: i16) -> WriteI16; + + /// Writes an unsigned 32-bit integer in big-endian order to the + /// underlying writer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_u32(&mut self, n: u32) -> 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 unsigned 32-bit integers to a `AsyncWrite`: + /// + /// ```rust + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut writer = Vec::new(); + /// + /// writer.write_u32(267).await?; + /// writer.write_u32(1205419366).await?; + /// + /// assert_eq!(writer, b"\x00\x00\x01\x0b\x47\xd9\x3d\x66"); + /// Ok(()) + /// } + /// ``` + fn write_u32(&mut self, n: u32) -> WriteU32; + + /// Writes a signed 32-bit integer in big-endian order to the + /// underlying writer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_i32(&mut self, n: i32) -> 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 signed 32-bit integers to a `AsyncWrite`: + /// + /// ```rust + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut writer = Vec::new(); + /// + /// writer.write_i32(267).await?; + /// writer.write_i32(1205419366).await?; + /// + /// assert_eq!(writer, b"\x00\x00\x01\x0b\x47\xd9\x3d\x66"); + /// Ok(()) + /// } + /// ``` + fn write_i32(&mut self, n: i32) -> WriteI32; + + /// Writes an unsigned 64-bit integer in big-endian order to the + /// underlying writer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_u64(&mut self, n: u64) -> 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 unsigned 64-bit integers to a `AsyncWrite`: + /// + /// ```rust + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut writer = Vec::new(); + /// + /// writer.write_u64(918733457491587).await?; + /// writer.write_u64(143).await?; + /// + /// assert_eq!(writer, b"\x00\x03\x43\x95\x4d\x60\x86\x83\x00\x00\x00\x00\x00\x00\x00\x8f"); + /// Ok(()) + /// } + /// ``` + fn write_u64(&mut self, n: u64) -> WriteU64; + + /// Writes an signed 64-bit integer in big-endian order to the + /// underlying writer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_i64(&mut self, n: i64) -> 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 signed 64-bit integers to a `AsyncWrite`: + /// + /// ```rust + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut writer = Vec::new(); + /// + /// writer.write_i64(i64::MIN).await?; + /// writer.write_i64(i64::MAX).await?; + /// + /// assert_eq!(writer, b"\x80\x00\x00\x00\x00\x00\x00\x00\x7f\xff\xff\xff\xff\xff\xff\xff"); + /// Ok(()) + /// } + /// ``` + fn write_i64(&mut self, n: i64) -> WriteI64; + + /// Writes an unsigned 128-bit integer in big-endian order to the + /// underlying writer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_u128(&mut self, n: u128) -> 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 unsigned 128-bit integers to a `AsyncWrite`: + /// + /// ```rust + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut writer = Vec::new(); + /// + /// writer.write_u128(16947640962301618749969007319746179).await?; + /// + /// assert_eq!(writer, vec![ + /// 0x00, 0x03, 0x43, 0x95, 0x4d, 0x60, 0x86, 0x83, + /// 0x00, 0x03, 0x43, 0x95, 0x4d, 0x60, 0x86, 0x83 + /// ]); + /// Ok(()) + /// } + /// ``` + fn write_u128(&mut self, n: u128) -> WriteU128; + + /// Writes an signed 128-bit integer in big-endian order to the + /// underlying writer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_i128(&mut self, n: i128) -> 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 signed 128-bit integers to a `AsyncWrite`: + /// + /// ```rust + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut writer = Vec::new(); + /// + /// writer.write_i128(i128::MIN).await?; + /// + /// assert_eq!(writer, vec![ + /// 0x80, 0, 0, 0, 0, 0, 0, 0, + /// 0, 0, 0, 0, 0, 0, 0, 0 + /// ]); + /// Ok(()) + /// } + /// ``` + 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. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_u16_le(&mut self, n: u16) -> 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 unsigned 16-bit integers to a `AsyncWrite`: + /// + /// ```rust + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut writer = Vec::new(); + /// + /// writer.write_u16_le(517).await?; + /// writer.write_u16_le(768).await?; + /// + /// assert_eq!(writer, b"\x05\x02\x00\x03"); + /// Ok(()) + /// } + /// ``` + fn write_u16_le(&mut self, n: u16) -> WriteU16Le; + + /// Writes a signed 16-bit integer in little-endian order to the + /// underlying writer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_i16_le(&mut self, n: i16) -> 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 signed 16-bit integers to a `AsyncWrite`: + /// + /// ```rust + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut writer = Vec::new(); + /// + /// writer.write_i16_le(193).await?; + /// writer.write_i16_le(-132).await?; + /// + /// assert_eq!(writer, b"\xc1\x00\x7c\xff"); + /// Ok(()) + /// } + /// ``` + fn write_i16_le(&mut self, n: i16) -> WriteI16Le; + + /// Writes an unsigned 32-bit integer in little-endian order to the + /// underlying writer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_u32_le(&mut self, n: u32) -> 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 unsigned 32-bit integers to a `AsyncWrite`: + /// + /// ```rust + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut writer = Vec::new(); + /// + /// writer.write_u32_le(267).await?; + /// writer.write_u32_le(1205419366).await?; + /// + /// assert_eq!(writer, b"\x0b\x01\x00\x00\x66\x3d\xd9\x47"); + /// Ok(()) + /// } + /// ``` + fn write_u32_le(&mut self, n: u32) -> WriteU32Le; + + /// Writes a signed 32-bit integer in little-endian order to the + /// underlying writer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_i32_le(&mut self, n: i32) -> 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 signed 32-bit integers to a `AsyncWrite`: + /// + /// ```rust + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut writer = Vec::new(); + /// + /// writer.write_i32_le(267).await?; + /// writer.write_i32_le(1205419366).await?; + /// + /// assert_eq!(writer, b"\x0b\x01\x00\x00\x66\x3d\xd9\x47"); + /// Ok(()) + /// } + /// ``` + fn write_i32_le(&mut self, n: i32) -> WriteI32Le; + + /// Writes an unsigned 64-bit integer in little-endian order to the + /// underlying writer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_u64_le(&mut self, n: u64) -> 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 unsigned 64-bit integers to a `AsyncWrite`: + /// + /// ```rust + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut writer = Vec::new(); + /// + /// writer.write_u64_le(918733457491587).await?; + /// writer.write_u64_le(143).await?; + /// + /// assert_eq!(writer, b"\x83\x86\x60\x4d\x95\x43\x03\x00\x8f\x00\x00\x00\x00\x00\x00\x00"); + /// Ok(()) + /// } + /// ``` + fn write_u64_le(&mut self, n: u64) -> WriteU64Le; + + /// Writes an signed 64-bit integer in little-endian order to the + /// underlying writer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_i64_le(&mut self, n: i64) -> 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 signed 64-bit integers to a `AsyncWrite`: + /// + /// ```rust + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut writer = Vec::new(); + /// + /// writer.write_i64_le(i64::MIN).await?; + /// writer.write_i64_le(i64::MAX).await?; + /// + /// assert_eq!(writer, b"\x00\x00\x00\x00\x00\x00\x00\x80\xff\xff\xff\xff\xff\xff\xff\x7f"); + /// Ok(()) + /// } + /// ``` + fn write_i64_le(&mut self, n: i64) -> WriteI64Le; + + /// Writes an unsigned 128-bit integer in little-endian order to the + /// underlying writer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_u128_le(&mut self, n: u128) -> 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 unsigned 128-bit integers to a `AsyncWrite`: + /// + /// ```rust + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut writer = Vec::new(); + /// + /// writer.write_u128_le(16947640962301618749969007319746179).await?; + /// + /// assert_eq!(writer, vec![ + /// 0x83, 0x86, 0x60, 0x4d, 0x95, 0x43, 0x03, 0x00, + /// 0x83, 0x86, 0x60, 0x4d, 0x95, 0x43, 0x03, 0x00, + /// ]); + /// Ok(()) + /// } + /// ``` + fn write_u128_le(&mut self, n: u128) -> WriteU128Le; + + /// Writes an signed 128-bit integer in little-endian order to the + /// underlying writer. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_i128_le(&mut self, n: i128) -> 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 signed 128-bit integers to a `AsyncWrite`: + /// + /// ```rust + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut writer = Vec::new(); + /// + /// writer.write_i128_le(i128::MIN).await?; + /// + /// assert_eq!(writer, vec![ + /// 0, 0, 0, 0, 0, 0, 0, + /// 0, 0, 0, 0, 0, 0, 0, 0, 0x80 + /// ]); + /// Ok(()) + /// } + /// ``` + 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 + /// contents reach their destination. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn flush(&mut self) -> io::Result<()>; + /// ``` + /// + /// # Errors + /// + /// It is considered an error if not all bytes could be written due to + /// I/O errors or EOF being reached. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::io::{self, BufWriter, AsyncWriteExt}; + /// use tokio::fs::File; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let f = File::create("foo.txt").await?; + /// let mut buffer = BufWriter::new(f); + /// + /// buffer.write_all(b"some bytes").await?; + /// buffer.flush().await?; + /// Ok(()) + /// } + /// ``` + fn flush(&mut self) -> Flush<'_, Self> + where + Self: Unpin, + { + flush(self) + } + + /// Shuts down the output stream, ensuring that the value can be dropped + /// cleanly. + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn shutdown(&mut self) -> io::Result<()>; + /// ``` + /// + /// Similar to [`flush`], all intermediately buffered is written to the + /// underlying stream. Once the operation completes, the caller should + /// no longer attempt to write to the stream. For example, the + /// `TcpStream` implementation will issue a `shutdown(Write)` sys call. + /// + /// [`flush`]: fn@crate::io::AsyncWriteExt::flush + /// + /// # Examples + /// + /// ```no_run + /// use tokio::io::{self, BufWriter, AsyncWriteExt}; + /// use tokio::fs::File; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let f = File::create("foo.txt").await?; + /// let mut buffer = BufWriter::new(f); + /// + /// buffer.write_all(b"some bytes").await?; + /// buffer.shutdown().await?; + /// Ok(()) + /// } + /// ``` + fn shutdown(&mut self) -> Shutdown<'_, Self> + where + Self: Unpin, + { + shutdown(self) + } + } +} + +impl<W: AsyncWrite + ?Sized> AsyncWriteExt for W {} diff --git a/third_party/rust/tokio/src/io/util/buf_reader.rs b/third_party/rust/tokio/src/io/util/buf_reader.rs new file mode 100644 index 0000000000..60879c0fdc --- /dev/null +++ b/third_party/rust/tokio/src/io/util/buf_reader.rs @@ -0,0 +1,311 @@ +use crate::io::util::DEFAULT_BUF_SIZE; +use crate::io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, ReadBuf}; + +use pin_project_lite::pin_project; +use std::io::{self, IoSlice, SeekFrom}; +use std::pin::Pin; +use std::task::{Context, Poll}; +use std::{cmp, fmt, mem}; + +pin_project! { + /// The `BufReader` struct adds buffering to any reader. + /// + /// It can be excessively inefficient to work directly with a [`AsyncRead`] + /// instance. A `BufReader` performs large, infrequent reads on the underlying + /// [`AsyncRead`] and maintains an in-memory buffer of the results. + /// + /// `BufReader` can improve the speed of programs that make *small* and + /// *repeated* read calls to the same file or network socket. It does not + /// help when reading very large amounts at once, or reading just one or a few + /// times. It also provides no advantage when reading from a source that is + /// already in memory, like a `Vec<u8>`. + /// + /// When the `BufReader` is dropped, the contents of its buffer will be + /// discarded. Creating multiple instances of a `BufReader` on the same + /// stream can cause data loss. + #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] + pub struct BufReader<R> { + #[pin] + pub(super) inner: R, + pub(super) buf: Box<[u8]>, + pub(super) pos: usize, + pub(super) cap: usize, + pub(super) seek_state: SeekState, + } +} + +impl<R: AsyncRead> BufReader<R> { + /// Creates a new `BufReader` with a default buffer capacity. The default is currently 8 KB, + /// but may change in the future. + pub fn new(inner: R) -> Self { + Self::with_capacity(DEFAULT_BUF_SIZE, inner) + } + + /// Creates a new `BufReader` with the specified buffer capacity. + pub fn with_capacity(capacity: usize, inner: R) -> Self { + let buffer = vec![0; capacity]; + Self { + inner, + buf: buffer.into_boxed_slice(), + pos: 0, + cap: 0, + seek_state: SeekState::Init, + } + } + + /// Gets a reference to the underlying reader. + /// + /// It is inadvisable to directly read from the underlying reader. + pub fn get_ref(&self) -> &R { + &self.inner + } + + /// Gets a mutable reference to the underlying reader. + /// + /// It is inadvisable to directly read from the underlying reader. + pub fn get_mut(&mut self) -> &mut R { + &mut self.inner + } + + /// Gets a pinned mutable reference to the underlying reader. + /// + /// It is inadvisable to directly read from the underlying reader. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut R> { + self.project().inner + } + + /// Consumes this `BufReader`, returning the underlying reader. + /// + /// Note that any leftover data in the internal buffer is lost. + pub fn into_inner(self) -> R { + self.inner + } + + /// Returns a reference to the internally buffered data. + /// + /// Unlike `fill_buf`, this will not attempt to fill the buffer if it is empty. + pub fn buffer(&self) -> &[u8] { + &self.buf[self.pos..self.cap] + } + + /// Invalidates all data in the internal buffer. + #[inline] + fn discard_buffer(self: Pin<&mut Self>) { + let me = self.project(); + *me.pos = 0; + *me.cap = 0; + } +} + +impl<R: AsyncRead> AsyncRead for BufReader<R> { + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<()>> { + // If we don't have any buffered data and we're doing a massive read + // (larger than our internal buffer), bypass our internal buffer + // entirely. + if self.pos == self.cap && buf.remaining() >= self.buf.len() { + let res = ready!(self.as_mut().get_pin_mut().poll_read(cx, buf)); + self.discard_buffer(); + return Poll::Ready(res); + } + let rem = ready!(self.as_mut().poll_fill_buf(cx))?; + let amt = std::cmp::min(rem.len(), buf.remaining()); + buf.put_slice(&rem[..amt]); + self.consume(amt); + Poll::Ready(Ok(())) + } +} + +impl<R: AsyncRead> AsyncBufRead for BufReader<R> { + fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> { + let me = self.project(); + + // If we've reached the end of our internal buffer then we need to fetch + // some more data from the underlying reader. + // Branch using `>=` instead of the more correct `==` + // to tell the compiler that the pos..cap slice is always valid. + if *me.pos >= *me.cap { + debug_assert!(*me.pos == *me.cap); + let mut buf = ReadBuf::new(me.buf); + ready!(me.inner.poll_read(cx, &mut buf))?; + *me.cap = buf.filled().len(); + *me.pos = 0; + } + Poll::Ready(Ok(&me.buf[*me.pos..*me.cap])) + } + + fn consume(self: Pin<&mut Self>, amt: usize) { + let me = self.project(); + *me.pos = cmp::min(*me.pos + amt, *me.cap); + } +} + +#[derive(Debug, Clone, Copy)] +pub(super) enum SeekState { + /// start_seek has not been called. + Init, + /// start_seek has been called, but poll_complete has not yet been called. + Start(SeekFrom), + /// Waiting for completion of the first poll_complete in the `n.checked_sub(remainder).is_none()` branch. + PendingOverflowed(i64), + /// Waiting for completion of poll_complete. + Pending, +} + +/// 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 +/// internal buffer. +/// +/// Seeking always discards the internal buffer, even if the seek position +/// would otherwise fall within it. This guarantees that calling +/// `.into_inner()` immediately after a seek yields the underlying reader +/// at the same position. +/// +/// See [`AsyncSeek`] for more details. +/// +/// Note: In the edge case where you're seeking with `SeekFrom::Current(n)` +/// where `n` minus the internal buffer length overflows an `i64`, two +/// seeks will be performed instead of one. If the second seek returns +/// `Err`, the underlying reader will be left at the same position it would +/// have if you called `seek` with `SeekFrom::Current(0)`. +impl<R: AsyncRead + AsyncSeek> AsyncSeek for BufReader<R> { + fn start_seek(self: Pin<&mut Self>, pos: SeekFrom) -> io::Result<()> { + // We needs to call seek operation multiple times. + // And we should always call both start_seek and poll_complete, + // as start_seek alone cannot guarantee that the operation will be completed. + // poll_complete receives a Context and returns a Poll, so it cannot be called + // inside start_seek. + *self.project().seek_state = SeekState::Start(pos); + Ok(()) + } + + fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> { + let res = match mem::replace(self.as_mut().project().seek_state, SeekState::Init) { + SeekState::Init => { + // 1.x AsyncSeek recommends calling poll_complete before start_seek. + // We don't have to guarantee that the value returned by + // poll_complete called without start_seek is correct, + // so we'll return 0. + return Poll::Ready(Ok(0)); + } + SeekState::Start(SeekFrom::Current(n)) => { + let remainder = (self.cap - self.pos) as i64; + // it should be safe to assume that remainder fits within an i64 as the alternative + // means we managed to allocate 8 exbibytes and that's absurd. + // But it's not out of the realm of possibility for some weird underlying reader to + // support seeking by i64::MIN so we need to handle underflow when subtracting + // remainder. + if let Some(offset) = n.checked_sub(remainder) { + self.as_mut() + .get_pin_mut() + .start_seek(SeekFrom::Current(offset))?; + } else { + // seek backwards by our remainder, and then by the offset + self.as_mut() + .get_pin_mut() + .start_seek(SeekFrom::Current(-remainder))?; + if self.as_mut().get_pin_mut().poll_complete(cx)?.is_pending() { + *self.as_mut().project().seek_state = SeekState::PendingOverflowed(n); + return Poll::Pending; + } + + // https://github.com/rust-lang/rust/pull/61157#issuecomment-495932676 + self.as_mut().discard_buffer(); + + self.as_mut() + .get_pin_mut() + .start_seek(SeekFrom::Current(n))?; + } + self.as_mut().get_pin_mut().poll_complete(cx)? + } + SeekState::PendingOverflowed(n) => { + if self.as_mut().get_pin_mut().poll_complete(cx)?.is_pending() { + *self.as_mut().project().seek_state = SeekState::PendingOverflowed(n); + return Poll::Pending; + } + + // https://github.com/rust-lang/rust/pull/61157#issuecomment-495932676 + self.as_mut().discard_buffer(); + + self.as_mut() + .get_pin_mut() + .start_seek(SeekFrom::Current(n))?; + self.as_mut().get_pin_mut().poll_complete(cx)? + } + SeekState::Start(pos) => { + // Seeking with Start/End doesn't care about our buffer length. + self.as_mut().get_pin_mut().start_seek(pos)?; + self.as_mut().get_pin_mut().poll_complete(cx)? + } + SeekState::Pending => self.as_mut().get_pin_mut().poll_complete(cx)?, + }; + + match res { + Poll::Ready(res) => { + self.discard_buffer(); + Poll::Ready(Ok(res)) + } + Poll::Pending => { + *self.as_mut().project().seek_state = SeekState::Pending; + Poll::Pending + } + } + } +} + +impl<R: AsyncRead + AsyncWrite> AsyncWrite for BufReader<R> { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + self.get_pin_mut().poll_write(cx, buf) + } + + fn poll_write_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[IoSlice<'_>], + ) -> Poll<io::Result<usize>> { + self.get_pin_mut().poll_write_vectored(cx, bufs) + } + + fn is_write_vectored(&self) -> bool { + self.get_ref().is_write_vectored() + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + self.get_pin_mut().poll_flush(cx) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + self.get_pin_mut().poll_shutdown(cx) + } +} + +impl<R: fmt::Debug> fmt::Debug for BufReader<R> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BufReader") + .field("reader", &self.inner) + .field( + "buffer", + &format_args!("{}/{}", self.cap - self.pos, self.buf.len()), + ) + .finish() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn assert_unpin() { + crate::is_unpin::<BufReader<()>>(); + } +} diff --git a/third_party/rust/tokio/src/io/util/buf_stream.rs b/third_party/rust/tokio/src/io/util/buf_stream.rs new file mode 100644 index 0000000000..595c142aca --- /dev/null +++ b/third_party/rust/tokio/src/io/util/buf_stream.rs @@ -0,0 +1,207 @@ +use crate::io::util::{BufReader, BufWriter}; +use crate::io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, ReadBuf}; + +use pin_project_lite::pin_project; +use std::io::{self, IoSlice, SeekFrom}; +use std::pin::Pin; +use std::task::{Context, Poll}; + +pin_project! { + /// Wraps a type that is [`AsyncWrite`] and [`AsyncRead`], and buffers its input and output. + /// + /// It can be excessively inefficient to work directly with something that implements [`AsyncWrite`] + /// and [`AsyncRead`]. For example, every `write`, however small, has to traverse the syscall + /// interface, and similarly, every read has to do the same. The [`BufWriter`] and [`BufReader`] + /// types aid with these problems respectively, but do so in only one direction. `BufStream` wraps + /// one in the other so that both directions are buffered. See their documentation for details. + #[derive(Debug)] + #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] + pub struct BufStream<RW> { + #[pin] + inner: BufReader<BufWriter<RW>>, + } +} + +impl<RW: AsyncRead + AsyncWrite> BufStream<RW> { + /// Wraps a type in both [`BufWriter`] and [`BufReader`]. + /// + /// See the documentation for those types and [`BufStream`] for details. + pub fn new(stream: RW) -> BufStream<RW> { + BufStream { + inner: BufReader::new(BufWriter::new(stream)), + } + } + + /// Creates a `BufStream` with the specified [`BufReader`] capacity and [`BufWriter`] + /// capacity. + /// + /// See the documentation for those types and [`BufStream`] for details. + pub fn with_capacity( + reader_capacity: usize, + writer_capacity: usize, + stream: RW, + ) -> BufStream<RW> { + BufStream { + inner: BufReader::with_capacity( + reader_capacity, + BufWriter::with_capacity(writer_capacity, stream), + ), + } + } + + /// Gets a reference to the underlying I/O object. + /// + /// It is inadvisable to directly read from the underlying I/O object. + pub fn get_ref(&self) -> &RW { + self.inner.get_ref().get_ref() + } + + /// Gets a mutable reference to the underlying I/O object. + /// + /// It is inadvisable to directly read from the underlying I/O object. + pub fn get_mut(&mut self) -> &mut RW { + self.inner.get_mut().get_mut() + } + + /// Gets a pinned mutable reference to the underlying I/O object. + /// + /// It is inadvisable to directly read from the underlying I/O object. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut RW> { + self.project().inner.get_pin_mut().get_pin_mut() + } + + /// Consumes this `BufStream`, returning the underlying I/O object. + /// + /// Note that any leftover data in the internal buffer is lost. + pub fn into_inner(self) -> RW { + self.inner.into_inner().into_inner() + } +} + +impl<RW> From<BufReader<BufWriter<RW>>> for BufStream<RW> { + fn from(b: BufReader<BufWriter<RW>>) -> Self { + BufStream { inner: b } + } +} + +impl<RW> From<BufWriter<BufReader<RW>>> for BufStream<RW> { + fn from(b: BufWriter<BufReader<RW>>) -> Self { + // we need to "invert" the reader and writer + let BufWriter { + inner: + BufReader { + inner, + buf: rbuf, + pos, + cap, + seek_state: rseek_state, + }, + buf: wbuf, + written, + seek_state: wseek_state, + } = b; + + BufStream { + inner: BufReader { + inner: BufWriter { + inner, + buf: wbuf, + written, + seek_state: wseek_state, + }, + buf: rbuf, + pos, + cap, + seek_state: rseek_state, + }, + } + } +} + +impl<RW: AsyncRead + AsyncWrite> AsyncWrite for BufStream<RW> { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + self.project().inner.poll_write(cx, buf) + } + + fn poll_write_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[IoSlice<'_>], + ) -> Poll<io::Result<usize>> { + self.project().inner.poll_write_vectored(cx, bufs) + } + + fn is_write_vectored(&self) -> bool { + self.inner.is_write_vectored() + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + self.project().inner.poll_flush(cx) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + self.project().inner.poll_shutdown(cx) + } +} + +impl<RW: AsyncRead + AsyncWrite> AsyncRead for BufStream<RW> { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<()>> { + self.project().inner.poll_read(cx, buf) + } +} + +/// Seek to an offset, in bytes, in the underlying stream. +/// +/// The position used for seeking with `SeekFrom::Current(_)` is the +/// position the underlying stream would be at if the `BufStream` had no +/// internal buffer. +/// +/// Seeking always discards the internal buffer, even if the seek position +/// would otherwise fall within it. This guarantees that calling +/// `.into_inner()` immediately after a seek yields the underlying reader +/// at the same position. +/// +/// See [`AsyncSeek`] for more details. +/// +/// Note: In the edge case where you're seeking with `SeekFrom::Current(n)` +/// where `n` minus the internal buffer length overflows an `i64`, two +/// seeks will be performed instead of one. If the second seek returns +/// `Err`, the underlying reader will be left at the same position it would +/// have if you called `seek` with `SeekFrom::Current(0)`. +impl<RW: AsyncRead + AsyncWrite + AsyncSeek> AsyncSeek for BufStream<RW> { + fn start_seek(self: Pin<&mut Self>, position: SeekFrom) -> io::Result<()> { + self.project().inner.start_seek(position) + } + + fn poll_complete(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> { + self.project().inner.poll_complete(cx) + } +} + +impl<RW: AsyncRead + AsyncWrite> AsyncBufRead for BufStream<RW> { + fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> { + self.project().inner.poll_fill_buf(cx) + } + + fn consume(self: Pin<&mut Self>, amt: usize) { + self.project().inner.consume(amt) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn assert_unpin() { + crate::is_unpin::<BufStream<()>>(); + } +} diff --git a/third_party/rust/tokio/src/io/util/buf_writer.rs b/third_party/rust/tokio/src/io/util/buf_writer.rs new file mode 100644 index 0000000000..8dd1bba60a --- /dev/null +++ b/third_party/rust/tokio/src/io/util/buf_writer.rs @@ -0,0 +1,310 @@ +use crate::io::util::DEFAULT_BUF_SIZE; +use crate::io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, ReadBuf}; + +use pin_project_lite::pin_project; +use std::fmt; +use std::io::{self, IoSlice, SeekFrom, Write}; +use std::pin::Pin; +use std::task::{Context, Poll}; + +pin_project! { + /// Wraps a writer and buffers its output. + /// + /// It can be excessively inefficient to work directly with something that + /// implements [`AsyncWrite`]. A `BufWriter` keeps an in-memory buffer of data and + /// writes it to an underlying writer in large, infrequent batches. + /// + /// `BufWriter` can improve the speed of programs that make *small* and + /// *repeated* write calls to the same file or network socket. It does not + /// help when writing very large amounts at once, or writing just one or a few + /// times. It also provides no advantage when writing to a destination that is + /// in memory, like a `Vec<u8>`. + /// + /// When the `BufWriter` is dropped, the contents of its buffer will be + /// discarded. Creating multiple instances of a `BufWriter` on the same + /// stream can cause data loss. If you need to write out the contents of its + /// buffer, you must manually call flush before the writer is dropped. + /// + /// [`AsyncWrite`]: AsyncWrite + /// [`flush`]: super::AsyncWriteExt::flush + /// + #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] + pub struct BufWriter<W> { + #[pin] + pub(super) inner: W, + pub(super) buf: Vec<u8>, + pub(super) written: usize, + pub(super) seek_state: SeekState, + } +} + +impl<W: AsyncWrite> BufWriter<W> { + /// Creates a new `BufWriter` with a default buffer capacity. The default is currently 8 KB, + /// but may change in the future. + pub fn new(inner: W) -> Self { + Self::with_capacity(DEFAULT_BUF_SIZE, inner) + } + + /// Creates a new `BufWriter` with the specified buffer capacity. + pub fn with_capacity(cap: usize, inner: W) -> Self { + Self { + inner, + buf: Vec::with_capacity(cap), + written: 0, + seek_state: SeekState::Init, + } + } + + fn flush_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + let mut me = self.project(); + + let len = me.buf.len(); + let mut ret = Ok(()); + while *me.written < len { + match ready!(me.inner.as_mut().poll_write(cx, &me.buf[*me.written..])) { + Ok(0) => { + ret = Err(io::Error::new( + io::ErrorKind::WriteZero, + "failed to write the buffered data", + )); + break; + } + Ok(n) => *me.written += n, + Err(e) => { + ret = Err(e); + break; + } + } + } + if *me.written > 0 { + me.buf.drain(..*me.written); + } + *me.written = 0; + Poll::Ready(ret) + } + + /// Gets a reference to the underlying writer. + pub fn get_ref(&self) -> &W { + &self.inner + } + + /// Gets a mutable reference to the underlying writer. + /// + /// It is inadvisable to directly write to the underlying writer. + pub fn get_mut(&mut self) -> &mut W { + &mut self.inner + } + + /// Gets a pinned mutable reference to the underlying writer. + /// + /// It is inadvisable to directly write to the underlying writer. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut W> { + self.project().inner + } + + /// Consumes this `BufWriter`, returning the underlying writer. + /// + /// Note that any leftover data in the internal buffer is lost. + pub fn into_inner(self) -> W { + self.inner + } + + /// Returns a reference to the internally buffered data. + pub fn buffer(&self) -> &[u8] { + &self.buf + } +} + +impl<W: AsyncWrite> AsyncWrite for BufWriter<W> { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + if self.buf.len() + buf.len() > self.buf.capacity() { + ready!(self.as_mut().flush_buf(cx))?; + } + + let me = self.project(); + if buf.len() >= me.buf.capacity() { + me.inner.poll_write(cx, buf) + } else { + Poll::Ready(me.buf.write(buf)) + } + } + + fn poll_write_vectored( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + mut bufs: &[IoSlice<'_>], + ) -> Poll<io::Result<usize>> { + if self.inner.is_write_vectored() { + let total_len = bufs + .iter() + .fold(0usize, |acc, b| acc.saturating_add(b.len())); + if total_len > self.buf.capacity() - self.buf.len() { + ready!(self.as_mut().flush_buf(cx))?; + } + let me = self.as_mut().project(); + if total_len >= me.buf.capacity() { + // It's more efficient to pass the slices directly to the + // underlying writer than to buffer them. + // The case when the total_len calculation saturates at + // usize::MAX is also handled here. + me.inner.poll_write_vectored(cx, bufs) + } else { + bufs.iter().for_each(|b| me.buf.extend_from_slice(b)); + Poll::Ready(Ok(total_len)) + } + } else { + // Remove empty buffers at the beginning of bufs. + while bufs.first().map(|buf| buf.len()) == Some(0) { + bufs = &bufs[1..]; + } + if bufs.is_empty() { + return Poll::Ready(Ok(0)); + } + // Flush if the first buffer doesn't fit. + let first_len = bufs[0].len(); + if first_len > self.buf.capacity() - self.buf.len() { + ready!(self.as_mut().flush_buf(cx))?; + debug_assert!(self.buf.is_empty()); + } + let me = self.as_mut().project(); + if first_len >= me.buf.capacity() { + // The slice is at least as large as the buffering capacity, + // so it's better to write it directly, bypassing the buffer. + debug_assert!(me.buf.is_empty()); + return me.inner.poll_write(cx, &bufs[0]); + } else { + me.buf.extend_from_slice(&bufs[0]); + bufs = &bufs[1..]; + } + let mut total_written = first_len; + debug_assert!(total_written != 0); + // Append the buffers that fit in the internal buffer. + for buf in bufs { + if buf.len() > me.buf.capacity() - me.buf.len() { + break; + } else { + me.buf.extend_from_slice(buf); + total_written += buf.len(); + } + } + Poll::Ready(Ok(total_written)) + } + } + + fn is_write_vectored(&self) -> bool { + true + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + ready!(self.as_mut().flush_buf(cx))?; + self.get_pin_mut().poll_flush(cx) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + ready!(self.as_mut().flush_buf(cx))?; + self.get_pin_mut().poll_shutdown(cx) + } +} + +#[derive(Debug, Clone, Copy)] +pub(super) enum SeekState { + /// start_seek has not been called. + Init, + /// start_seek has been called, but poll_complete has not yet been called. + Start(SeekFrom), + /// Waiting for completion of poll_complete. + Pending, +} + +/// Seek to the offset, in bytes, in the underlying writer. +/// +/// Seeking always writes out the internal buffer before seeking. +impl<W: AsyncWrite + AsyncSeek> AsyncSeek for BufWriter<W> { + fn start_seek(self: Pin<&mut Self>, pos: SeekFrom) -> io::Result<()> { + // We need to flush the internal buffer before seeking. + // It receives a `Context` and returns a `Poll`, so it cannot be called + // inside `start_seek`. + *self.project().seek_state = SeekState::Start(pos); + Ok(()) + } + + fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> { + let pos = match self.seek_state { + SeekState::Init => { + return self.project().inner.poll_complete(cx); + } + SeekState::Start(pos) => Some(pos), + SeekState::Pending => None, + }; + + // Flush the internal buffer before seeking. + ready!(self.as_mut().flush_buf(cx))?; + + let mut me = self.project(); + if let Some(pos) = pos { + // Ensure previous seeks have finished before starting a new one + ready!(me.inner.as_mut().poll_complete(cx))?; + if let Err(e) = me.inner.as_mut().start_seek(pos) { + *me.seek_state = SeekState::Init; + return Poll::Ready(Err(e)); + } + } + match me.inner.poll_complete(cx) { + Poll::Ready(res) => { + *me.seek_state = SeekState::Init; + Poll::Ready(res) + } + Poll::Pending => { + *me.seek_state = SeekState::Pending; + Poll::Pending + } + } + } +} + +impl<W: AsyncWrite + AsyncRead> AsyncRead for BufWriter<W> { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<()>> { + self.get_pin_mut().poll_read(cx, buf) + } +} + +impl<W: AsyncWrite + AsyncBufRead> AsyncBufRead for BufWriter<W> { + fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> { + self.get_pin_mut().poll_fill_buf(cx) + } + + fn consume(self: Pin<&mut Self>, amt: usize) { + self.get_pin_mut().consume(amt) + } +} + +impl<W: fmt::Debug> fmt::Debug for BufWriter<W> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BufWriter") + .field("writer", &self.inner) + .field( + "buffer", + &format_args!("{}/{}", self.buf.len(), self.buf.capacity()), + ) + .field("written", &self.written) + .finish() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn assert_unpin() { + crate::is_unpin::<BufWriter<()>>(); + } +} diff --git a/third_party/rust/tokio/src/io/util/chain.rs b/third_party/rust/tokio/src/io/util/chain.rs new file mode 100644 index 0000000000..84f37fc7d4 --- /dev/null +++ b/third_party/rust/tokio/src/io/util/chain.rs @@ -0,0 +1,144 @@ +use crate::io::{AsyncBufRead, AsyncRead, ReadBuf}; + +use pin_project_lite::pin_project; +use std::fmt; +use std::io; +use std::pin::Pin; +use std::task::{Context, Poll}; + +pin_project! { + /// Stream for the [`chain`](super::AsyncReadExt::chain) method. + #[must_use = "streams do nothing unless polled"] + #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] + pub struct Chain<T, U> { + #[pin] + first: T, + #[pin] + second: U, + done_first: bool, + } +} + +pub(super) fn chain<T, U>(first: T, second: U) -> Chain<T, U> +where + T: AsyncRead, + U: AsyncRead, +{ + Chain { + first, + second, + done_first: false, + } +} + +impl<T, U> Chain<T, U> +where + T: AsyncRead, + U: AsyncRead, +{ + /// Gets references to the underlying readers in this `Chain`. + pub fn get_ref(&self) -> (&T, &U) { + (&self.first, &self.second) + } + + /// Gets mutable references to the underlying readers in this `Chain`. + /// + /// Care should be taken to avoid modifying the internal I/O state of the + /// underlying readers as doing so may corrupt the internal state of this + /// `Chain`. + pub fn get_mut(&mut self) -> (&mut T, &mut U) { + (&mut self.first, &mut self.second) + } + + /// Gets pinned mutable references to the underlying readers in this `Chain`. + /// + /// Care should be taken to avoid modifying the internal I/O state of the + /// underlying readers as doing so may corrupt the internal state of this + /// `Chain`. + pub fn get_pin_mut(self: Pin<&mut Self>) -> (Pin<&mut T>, Pin<&mut U>) { + let me = self.project(); + (me.first, me.second) + } + + /// Consumes the `Chain`, returning the wrapped readers. + pub fn into_inner(self) -> (T, U) { + (self.first, self.second) + } +} + +impl<T, U> fmt::Debug for Chain<T, U> +where + T: fmt::Debug, + U: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Chain") + .field("t", &self.first) + .field("u", &self.second) + .finish() + } +} + +impl<T, U> AsyncRead for Chain<T, U> +where + T: AsyncRead, + U: AsyncRead, +{ + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<()>> { + let me = self.project(); + + if !*me.done_first { + let rem = buf.remaining(); + ready!(me.first.poll_read(cx, buf))?; + if buf.remaining() == rem { + *me.done_first = true; + } else { + return Poll::Ready(Ok(())); + } + } + me.second.poll_read(cx, buf) + } +} + +impl<T, U> AsyncBufRead for Chain<T, U> +where + T: AsyncBufRead, + U: AsyncBufRead, +{ + fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> { + let me = self.project(); + + if !*me.done_first { + match ready!(me.first.poll_fill_buf(cx)?) { + buf if buf.is_empty() => { + *me.done_first = true; + } + buf => return Poll::Ready(Ok(buf)), + } + } + me.second.poll_fill_buf(cx) + } + + fn consume(self: Pin<&mut Self>, amt: usize) { + let me = self.project(); + if !*me.done_first { + me.first.consume(amt) + } else { + me.second.consume(amt) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn assert_unpin() { + crate::is_unpin::<Chain<(), ()>>(); + } +} diff --git a/third_party/rust/tokio/src/io/util/copy.rs b/third_party/rust/tokio/src/io/util/copy.rs new file mode 100644 index 0000000000..d0ab7cb140 --- /dev/null +++ b/third_party/rust/tokio/src/io/util/copy.rs @@ -0,0 +1,175 @@ +use crate::io::{AsyncRead, AsyncWrite, ReadBuf}; + +use std::future::Future; +use std::io; +use std::pin::Pin; +use std::task::{Context, Poll}; + +#[derive(Debug)] +pub(super) struct CopyBuffer { + read_done: bool, + need_flush: bool, + pos: usize, + cap: usize, + amt: u64, + buf: Box<[u8]>, +} + +impl CopyBuffer { + pub(super) fn new() -> Self { + Self { + read_done: false, + need_flush: false, + pos: 0, + cap: 0, + amt: 0, + buf: vec![0; super::DEFAULT_BUF_SIZE].into_boxed_slice(), + } + } + + pub(super) fn poll_copy<R, W>( + &mut self, + cx: &mut Context<'_>, + mut reader: Pin<&mut R>, + mut writer: Pin<&mut W>, + ) -> Poll<io::Result<u64>> + where + R: AsyncRead + ?Sized, + W: AsyncWrite + ?Sized, + { + loop { + // 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); + + match reader.as_mut().poll_read(cx, &mut buf) { + 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; + } + } + + let n = buf.filled().len(); + if n == 0 { + self.read_done = true; + } else { + self.pos = 0; + self.cap = n; + } + } + + // 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]))?; + if i == 0 { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::WriteZero, + "write zero byte into writer", + ))); + } 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 { + ready!(writer.as_mut().poll_flush(cx))?; + return Poll::Ready(Ok(self.amt)); + } + } + } +} + +/// A future that asynchronously copies the entire contents of a reader into a +/// writer. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +struct Copy<'a, R: ?Sized, W: ?Sized> { + reader: &'a mut R, + writer: &'a mut W, + buf: CopyBuffer, +} + +cfg_io_util! { + /// Asynchronously copies the entire contents of a reader into a writer. + /// + /// 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. + /// + /// 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]. + /// + /// [std]: std::io::copy + /// + /// # Errors + /// + /// The returned future will return an error immediately if any call to + /// `poll_read` or `poll_write` returns an error. + /// + /// # Examples + /// + /// ``` + /// use tokio::io; + /// + /// # async fn dox() -> std::io::Result<()> { + /// let mut reader: &[u8] = b"hello"; + /// let mut writer: Vec<u8> = vec![]; + /// + /// io::copy(&mut reader, &mut writer).await?; + /// + /// assert_eq!(&b"hello"[..], &writer[..]); + /// # Ok(()) + /// # } + /// ``` + pub async fn copy<'a, R, W>(reader: &'a mut R, writer: &'a mut W) -> io::Result<u64> + where + R: AsyncRead + Unpin + ?Sized, + W: AsyncWrite + Unpin + ?Sized, + { + Copy { + reader, + writer, + buf: CopyBuffer::new() + }.await + } +} + +impl<R, W> Future for Copy<'_, R, W> +where + R: AsyncRead + Unpin + ?Sized, + W: AsyncWrite + Unpin + ?Sized, +{ + type Output = io::Result<u64>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> { + let me = &mut *self; + + me.buf + .poll_copy(cx, Pin::new(&mut *me.reader), Pin::new(&mut *me.writer)) + } +} diff --git a/third_party/rust/tokio/src/io/util/copy_bidirectional.rs b/third_party/rust/tokio/src/io/util/copy_bidirectional.rs new file mode 100644 index 0000000000..c93060b361 --- /dev/null +++ b/third_party/rust/tokio/src/io/util/copy_bidirectional.rs @@ -0,0 +1,120 @@ +use super::copy::CopyBuffer; + +use crate::io::{AsyncRead, AsyncWrite}; + +use std::future::Future; +use std::io; +use std::pin::Pin; +use std::task::{Context, Poll}; + +enum TransferState { + Running(CopyBuffer), + ShuttingDown(u64), + 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, + r: &mut A, + w: &mut B, +) -> Poll<io::Result<u64>> +where + A: AsyncRead + AsyncWrite + Unpin + ?Sized, + B: AsyncRead + AsyncWrite + Unpin + ?Sized, +{ + let mut r = Pin::new(r); + let mut w = Pin::new(w); + + loop { + match state { + TransferState::Running(buf) => { + let count = ready!(buf.poll_copy(cx, r.as_mut(), w.as_mut()))?; + *state = TransferState::ShuttingDown(count); + } + TransferState::ShuttingDown(count) => { + ready!(w.as_mut().poll_shutdown(cx))?; + + *state = TransferState::Done(*count); + } + TransferState::Done(count) => return Poll::Ready(Ok(*count)), + } + } +} + +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, +/// writing any data read to the opposing stream. +/// This happens in both directions concurrently. +/// +/// If an EOF is observed on one stream, [`shutdown()`] will be invoked on +/// the other, and reading from that stream will stop. Copying of data in +/// the other direction will continue. +/// +/// The future will complete successfully once both directions of communication has been shut down. +/// A direction is shut down when the reader reports EOF, +/// at which point [`shutdown()`] is called on the corresponding writer. When finished, +/// it will return a tuple of the number of bytes copied from a to b +/// and the number of bytes copied from b to a, in that order. +/// +/// [`shutdown()`]: crate::io::AsyncWriteExt::shutdown +/// +/// # Errors +/// +/// The future will immediately return an error if any IO operation on `a` +/// or `b` returns an error. Some data read from either stream may be lost (not +/// written to the other stream) in this case. +/// +/// # Return value +/// +/// Returns a tuple of bytes copied `a` to `b` and bytes copied `b` to `a`. +#[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] +pub async fn copy_bidirectional<A, B>(a: &mut A, b: &mut B) -> Result<(u64, u64), std::io::Error> +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()), + } + .await +} diff --git a/third_party/rust/tokio/src/io/util/copy_buf.rs b/third_party/rust/tokio/src/io/util/copy_buf.rs new file mode 100644 index 0000000000..6831580b40 --- /dev/null +++ b/third_party/rust/tokio/src/io/util/copy_buf.rs @@ -0,0 +1,102 @@ +use crate::io::{AsyncBufRead, AsyncWrite}; +use std::future::Future; +use std::io; +use std::pin::Pin; +use std::task::{Context, Poll}; + +cfg_io_util! { + /// A future that asynchronously copies the entire contents of a reader into a + /// writer. + /// + /// This struct is generally created by calling [`copy_buf`][copy_buf]. Please + /// see the documentation of `copy_buf()` for more details. + /// + /// [copy_buf]: copy_buf() + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + struct CopyBuf<'a, R: ?Sized, W: ?Sized> { + reader: &'a mut R, + writer: &'a mut W, + amt: u64, + } + + /// Asynchronously copies the entire contents of a reader into a writer. + /// + /// 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. + /// + /// On success, the total number of bytes that were copied from `reader` to + /// `writer` is returned. + /// + /// + /// # Errors + /// + /// The returned future will finish with an error will return an error + /// immediately if any call to `poll_fill_buf` or `poll_write` returns an + /// error. + /// + /// # Examples + /// + /// ``` + /// use tokio::io; + /// + /// # async fn dox() -> std::io::Result<()> { + /// let mut reader: &[u8] = b"hello"; + /// let mut writer: Vec<u8> = vec![]; + /// + /// io::copy_buf(&mut reader, &mut writer).await?; + /// + /// assert_eq!(b"hello", &writer[..]); + /// # Ok(()) + /// # } + /// ``` + pub async fn copy_buf<'a, R, W>(reader: &'a mut R, writer: &'a mut W) -> io::Result<u64> + where + R: AsyncBufRead + Unpin + ?Sized, + W: AsyncWrite + Unpin + ?Sized, + { + CopyBuf { + reader, + writer, + amt: 0, + }.await + } +} + +impl<R, W> Future for CopyBuf<'_, R, W> +where + R: AsyncBufRead + Unpin + ?Sized, + W: AsyncWrite + Unpin + ?Sized, +{ + type Output = io::Result<u64>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + loop { + let me = &mut *self; + let buffer = ready!(Pin::new(&mut *me.reader).poll_fill_buf(cx))?; + if buffer.is_empty() { + ready!(Pin::new(&mut self.writer).poll_flush(cx))?; + return Poll::Ready(Ok(self.amt)); + } + + let i = ready!(Pin::new(&mut *me.writer).poll_write(cx, buffer))?; + if i == 0 { + return Poll::Ready(Err(std::io::ErrorKind::WriteZero.into())); + } + self.amt += i as u64; + Pin::new(&mut *self.reader).consume(i); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn assert_unpin() { + use std::marker::PhantomPinned; + crate::is_unpin::<CopyBuf<'_, PhantomPinned, PhantomPinned>>(); + } +} diff --git a/third_party/rust/tokio/src/io/util/empty.rs b/third_party/rust/tokio/src/io/util/empty.rs new file mode 100644 index 0000000000..77db60e40b --- /dev/null +++ b/third_party/rust/tokio/src/io/util/empty.rs @@ -0,0 +1,100 @@ +use crate::io::{AsyncBufRead, AsyncRead, ReadBuf}; + +use std::fmt; +use std::io; +use std::pin::Pin; +use std::task::{Context, Poll}; + +cfg_io_util! { + /// An async reader which is always at EOF. + /// + /// This struct is generally created by calling [`empty`]. Please see + /// the documentation of [`empty()`][`empty`] for more details. + /// + /// This is an asynchronous version of [`std::io::empty`][std]. + /// + /// [`empty`]: fn@empty + /// [std]: std::io::empty + pub struct Empty { + _p: (), + } + + /// Creates a new empty async reader. + /// + /// All reads from the returned reader will return `Poll::Ready(Ok(0))`. + /// + /// This is an asynchronous version of [`std::io::empty`][std]. + /// + /// [std]: std::io::empty + /// + /// # Examples + /// + /// A slightly sad example of not reading anything into a buffer: + /// + /// ``` + /// use tokio::io::{self, AsyncReadExt}; + /// + /// #[tokio::main] + /// async fn main() { + /// let mut buffer = String::new(); + /// io::empty().read_to_string(&mut buffer).await.unwrap(); + /// assert!(buffer.is_empty()); + /// } + /// ``` + pub fn empty() -> Empty { + Empty { _p: () } + } +} + +impl AsyncRead for Empty { + #[inline] + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + _: &mut ReadBuf<'_>, + ) -> Poll<io::Result<()>> { + ready!(poll_proceed_and_make_progress(cx)); + Poll::Ready(Ok(())) + } +} + +impl AsyncBufRead for Empty { + #[inline] + fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> { + ready!(poll_proceed_and_make_progress(cx)); + Poll::Ready(Ok(&[])) + } + + #[inline] + fn consume(self: Pin<&mut Self>, _: usize) {} +} + +impl fmt::Debug for Empty { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.pad("Empty { .. }") + } +} + +cfg_coop! { + fn poll_proceed_and_make_progress(cx: &mut Context<'_>) -> Poll<()> { + let coop = ready!(crate::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::*; + + #[test] + fn assert_unpin() { + crate::is_unpin::<Empty>(); + } +} diff --git a/third_party/rust/tokio/src/io/util/fill_buf.rs b/third_party/rust/tokio/src/io/util/fill_buf.rs new file mode 100644 index 0000000000..bb07c766e2 --- /dev/null +++ b/third_party/rust/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/third_party/rust/tokio/src/io/util/flush.rs b/third_party/rust/tokio/src/io/util/flush.rs new file mode 100644 index 0000000000..88d60b868d --- /dev/null +++ b/third_party/rust/tokio/src/io/util/flush.rs @@ -0,0 +1,46 @@ +use crate::io::AsyncWrite; + +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! { + /// A future used to fully flush an I/O object. + /// + /// Created by the [`AsyncWriteExt::flush`][flush] function. + /// [flush]: crate::io::AsyncWriteExt::flush + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct Flush<'a, A: ?Sized> { + a: &'a mut A, + // Make this future `!Unpin` for compatibility with async trait methods. + #[pin] + _pin: PhantomPinned, + } +} + +/// Creates a future which will entirely flush an I/O object. +pub(super) fn flush<A>(a: &mut A) -> Flush<'_, A> +where + A: AsyncWrite + Unpin + ?Sized, +{ + Flush { + a, + _pin: PhantomPinned, + } +} + +impl<A> Future for Flush<'_, A> +where + A: AsyncWrite + Unpin + ?Sized, +{ + type Output = io::Result<()>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let me = self.project(); + Pin::new(&mut *me.a).poll_flush(cx) + } +} diff --git a/third_party/rust/tokio/src/io/util/lines.rs b/third_party/rust/tokio/src/io/util/lines.rs new file mode 100644 index 0000000000..717f633f95 --- /dev/null +++ b/third_party/rust/tokio/src/io/util/lines.rs @@ -0,0 +1,145 @@ +use crate::io::util::read_line::read_line_internal; +use crate::io::AsyncBufRead; + +use pin_project_lite::pin_project; +use std::io; +use std::mem; +use std::pin::Pin; +use std::task::{Context, Poll}; + +pin_project! { + /// Reads lines from an [`AsyncBufRead`]. + /// + /// A `Lines` can be turned into a `Stream` with [`LinesStream`]. + /// + /// This type is usually created using the [`lines`] method. + /// + /// [`AsyncBufRead`]: crate::io::AsyncBufRead + /// [`LinesStream`]: https://docs.rs/tokio-stream/0.1/tokio_stream/wrappers/struct.LinesStream.html + /// [`lines`]: crate::io::AsyncBufReadExt::lines + #[derive(Debug)] + #[must_use = "streams do nothing unless polled"] + #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] + pub struct Lines<R> { + #[pin] + reader: R, + buf: String, + bytes: Vec<u8>, + read: usize, + } +} + +pub(crate) fn lines<R>(reader: R) -> Lines<R> +where + R: AsyncBufRead, +{ + Lines { + reader, + buf: String::new(), + bytes: Vec::new(), + read: 0, + } +} + +impl<R> Lines<R> +where + R: AsyncBufRead + Unpin, +{ + /// Returns the next line in the stream. + /// + /// # Cancel safety + /// + /// This method is cancellation safe. + /// + /// # Examples + /// + /// ``` + /// # use tokio::io::AsyncBufRead; + /// use tokio::io::AsyncBufReadExt; + /// + /// # async fn dox(my_buf_read: impl AsyncBufRead + Unpin) -> std::io::Result<()> { + /// let mut lines = my_buf_read.lines(); + /// + /// while let Some(line) = lines.next_line().await? { + /// println!("length = {}", line.len()) + /// } + /// # Ok(()) + /// # } + /// ``` + pub async fn next_line(&mut self) -> io::Result<Option<String>> { + use crate::future::poll_fn; + + poll_fn(|cx| Pin::new(&mut *self).poll_next_line(cx)).await + } + + /// Obtains a mutable reference to the underlying reader. + pub fn get_mut(&mut self) -> &mut R { + &mut self.reader + } + + /// Obtains a reference to the underlying reader. + pub fn get_ref(&mut self) -> &R { + &self.reader + } + + /// Unwraps this `Lines<R>`, returning the underlying reader. + /// + /// Note that any leftover data in the internal buffer is lost. + /// Therefore, a following read from the underlying reader may lead to data loss. + pub fn into_inner(self) -> R { + self.reader + } +} + +impl<R> Lines<R> +where + R: AsyncBufRead, +{ + /// Polls for the next line in the stream. + /// + /// This method returns: + /// + /// * `Poll::Pending` if the next line is not yet available. + /// * `Poll::Ready(Ok(Some(line)))` if the next line is available. + /// * `Poll::Ready(Ok(None))` if there are no more lines in this stream. + /// * `Poll::Ready(Err(err))` if an IO error occurred while reading the next line. + /// + /// 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. + pub fn poll_next_line( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<io::Result<Option<String>>> { + let me = self.project(); + + let n = ready!(read_line_internal(me.reader, cx, me.buf, me.bytes, me.read))?; + debug_assert_eq!(*me.read, 0); + + if n == 0 && me.buf.is_empty() { + return Poll::Ready(Ok(None)); + } + + if me.buf.ends_with('\n') { + me.buf.pop(); + + if me.buf.ends_with('\r') { + me.buf.pop(); + } + } + + Poll::Ready(Ok(Some(mem::take(me.buf)))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn assert_unpin() { + crate::is_unpin::<Lines<()>>(); + } +} diff --git a/third_party/rust/tokio/src/io/util/mem.rs b/third_party/rust/tokio/src/io/util/mem.rs new file mode 100644 index 0000000000..4019db56ff --- /dev/null +++ b/third_party/rust/tokio/src/io/util/mem.rs @@ -0,0 +1,295 @@ +//! In-process memory IO types. + +use crate::io::{AsyncRead, AsyncWrite, ReadBuf}; +use crate::loom::sync::Mutex; + +use bytes::{Buf, BytesMut}; +use std::{ + pin::Pin, + sync::Arc, + task::{self, Poll, Waker}, +}; + +/// A bidirectional pipe to read and write bytes in memory. +/// +/// A pair of `DuplexStream`s are created together, and they act as a "channel" +/// that can be used as in-memory IO types. Writing to one of the pairs will +/// allow that data to be read from the other, and vice versa. +/// +/// # Closing a `DuplexStream` +/// +/// If one end of the `DuplexStream` channel is dropped, any pending reads on +/// the other side will continue to read data until the buffer is drained, then +/// they will signal EOF by returning 0 bytes. Any writes to the other side, +/// including pending ones (that are waiting for free space in the buffer) will +/// return `Err(BrokenPipe)` immediately. +/// +/// # Example +/// +/// ``` +/// # async fn ex() -> std::io::Result<()> { +/// # use tokio::io::{AsyncReadExt, AsyncWriteExt}; +/// let (mut client, mut server) = tokio::io::duplex(64); +/// +/// client.write_all(b"ping").await?; +/// +/// let mut buf = [0u8; 4]; +/// server.read_exact(&mut buf).await?; +/// assert_eq!(&buf, b"ping"); +/// +/// server.write_all(b"pong").await?; +/// +/// client.read_exact(&mut buf).await?; +/// assert_eq!(&buf, b"pong"); +/// # Ok(()) +/// # } +/// ``` +#[derive(Debug)] +#[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] +pub struct DuplexStream { + read: Arc<Mutex<Pipe>>, + write: Arc<Mutex<Pipe>>, +} + +/// A unidirectional IO over a piece of memory. +/// +/// Data can be written to the pipe, and reading will return that data. +#[derive(Debug)] +struct Pipe { + /// The buffer storing the bytes written, also read from. + /// + /// Using a `BytesMut` because it has efficient `Buf` and `BufMut` + /// functionality already. Additionally, it can try to copy data in the + /// same buffer if there read index has advanced far enough. + buffer: BytesMut, + /// Determines if the write side has been closed. + is_closed: bool, + /// The maximum amount of bytes that can be written before returning + /// `Poll::Pending`. + max_buf_size: usize, + /// If the `read` side has been polled and is pending, this is the waker + /// for that parked task. + read_waker: Option<Waker>, + /// If the `write` side has filled the `max_buf_size` and returned + /// `Poll::Pending`, this is the waker for that parked task. + write_waker: Option<Waker>, +} + +// ===== impl DuplexStream ===== + +/// Create a new pair of `DuplexStream`s that act like a pair of connected sockets. +/// +/// The `max_buf_size` argument is the maximum amount of bytes that can be +/// written to a side before the write returns `Poll::Pending`. +#[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] +pub fn duplex(max_buf_size: usize) -> (DuplexStream, DuplexStream) { + let one = Arc::new(Mutex::new(Pipe::new(max_buf_size))); + let two = Arc::new(Mutex::new(Pipe::new(max_buf_size))); + + ( + DuplexStream { + read: one.clone(), + write: two.clone(), + }, + DuplexStream { + read: two, + write: one, + }, + ) +} + +impl AsyncRead for DuplexStream { + // Previous rustc required this `self` to be `mut`, even though newer + // versions recognize it isn't needed to call `lock()`. So for + // compatibility, we include the `mut` and `allow` the lint. + // + // See https://github.com/rust-lang/rust/issues/73592 + #[allow(unused_mut)] + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut task::Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<std::io::Result<()>> { + Pin::new(&mut *self.read.lock()).poll_read(cx, buf) + } +} + +impl AsyncWrite for DuplexStream { + #[allow(unused_mut)] + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut task::Context<'_>, + buf: &[u8], + ) -> Poll<std::io::Result<usize>> { + Pin::new(&mut *self.write.lock()).poll_write(cx, buf) + } + + #[allow(unused_mut)] + fn poll_flush( + mut self: Pin<&mut Self>, + cx: &mut task::Context<'_>, + ) -> Poll<std::io::Result<()>> { + Pin::new(&mut *self.write.lock()).poll_flush(cx) + } + + #[allow(unused_mut)] + fn poll_shutdown( + mut self: Pin<&mut Self>, + cx: &mut task::Context<'_>, + ) -> Poll<std::io::Result<()>> { + Pin::new(&mut *self.write.lock()).poll_shutdown(cx) + } +} + +impl Drop for DuplexStream { + fn drop(&mut self) { + // notify the other side of the closure + self.write.lock().close_write(); + self.read.lock().close_read(); + } +} + +// ===== impl Pipe ===== + +impl Pipe { + fn new(max_buf_size: usize) -> Self { + Pipe { + buffer: BytesMut::new(), + is_closed: false, + max_buf_size, + read_waker: None, + write_waker: None, + } + } + + fn close_write(&mut self) { + self.is_closed = true; + // needs to notify any readers that no more data will come + if let Some(waker) = self.read_waker.take() { + waker.wake(); + } + } + + fn close_read(&mut self) { + self.is_closed = true; + // needs to notify any writers that they have to abort + if let Some(waker) = self.write_waker.take() { + waker.wake(); + } + } + + fn poll_read_internal( + mut self: Pin<&mut Self>, + cx: &mut task::Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<std::io::Result<()>> { + if self.buffer.has_remaining() { + let max = self.buffer.remaining().min(buf.remaining()); + buf.put_slice(&self.buffer[..max]); + self.buffer.advance(max); + if max > 0 { + // The passed `buf` might have been empty, don't wake up if + // no bytes have been moved. + if let Some(waker) = self.write_waker.take() { + waker.wake(); + } + } + Poll::Ready(Ok(())) + } else if self.is_closed { + Poll::Ready(Ok(())) + } else { + self.read_waker = Some(cx.waker().clone()); + Poll::Pending + } + } + + fn poll_write_internal( + mut self: Pin<&mut Self>, + cx: &mut task::Context<'_>, + buf: &[u8], + ) -> Poll<std::io::Result<usize>> { + if self.is_closed { + return Poll::Ready(Err(std::io::ErrorKind::BrokenPipe.into())); + } + let avail = self.max_buf_size - self.buffer.len(); + if avail == 0 { + self.write_waker = Some(cx.waker().clone()); + return Poll::Pending; + } + + let len = buf.len().min(avail); + self.buffer.extend_from_slice(&buf[..len]); + if let Some(waker) = self.read_waker.take() { + waker.wake(); + } + 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<()>> { + let coop = ready!(crate::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<()>> { + 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>> { + let coop = ready!(crate::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>> { + self.poll_write_internal(cx, buf) + } + } + + fn poll_flush(self: Pin<&mut Self>, _: &mut task::Context<'_>) -> Poll<std::io::Result<()>> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown( + mut self: Pin<&mut Self>, + _: &mut task::Context<'_>, + ) -> Poll<std::io::Result<()>> { + self.close_write(); + Poll::Ready(Ok(())) + } +} diff --git a/third_party/rust/tokio/src/io/util/mod.rs b/third_party/rust/tokio/src/io/util/mod.rs new file mode 100644 index 0000000000..21199d0be8 --- /dev/null +++ b/third_party/rust/tokio/src/io/util/mod.rs @@ -0,0 +1,97 @@ +#![allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 + +cfg_io_util! { + mod async_buf_read_ext; + pub use async_buf_read_ext::AsyncBufReadExt; + + mod async_read_ext; + pub use async_read_ext::AsyncReadExt; + + mod async_seek_ext; + pub use async_seek_ext::AsyncSeekExt; + + mod async_write_ext; + pub use async_write_ext::AsyncWriteExt; + + mod buf_reader; + pub use buf_reader::BufReader; + + mod buf_stream; + pub use buf_stream::BufStream; + + mod buf_writer; + pub use buf_writer::BufWriter; + + mod chain; + + mod copy; + pub use copy::copy; + + mod copy_bidirectional; + pub use copy_bidirectional::copy_bidirectional; + + mod copy_buf; + pub use copy_buf::copy_buf; + + mod empty; + pub use empty::{empty, Empty}; + + mod flush; + + mod lines; + pub use lines::Lines; + + mod mem; + pub use mem::{duplex, DuplexStream}; + + mod read; + mod read_buf; + mod read_exact; + mod read_int; + mod read_line; + mod fill_buf; + + mod read_to_end; + mod vec_with_initialized; + cfg_process! { + pub(crate) use read_to_end::read_to_end; + } + + mod read_to_string; + mod read_until; + + mod repeat; + pub use repeat::{repeat, Repeat}; + + mod shutdown; + + mod sink; + pub use sink::{sink, Sink}; + + mod split; + pub use split::Split; + + mod take; + pub use take::Take; + + mod write; + mod write_vectored; + mod write_all; + mod write_buf; + mod write_all_buf; + mod write_int; + + + // used by `BufReader` and `BufWriter` + // https://github.com/rust-lang/rust/blob/master/library/std/src/sys_common/io.rs#L1 + const DEFAULT_BUF_SIZE: usize = 8 * 1024; +} + +cfg_not_io_util! { + cfg_process! { + mod vec_with_initialized; + mod read_to_end; + // Used by process + pub(crate) use read_to_end::read_to_end; + } +} diff --git a/third_party/rust/tokio/src/io/util/read.rs b/third_party/rust/tokio/src/io/util/read.rs new file mode 100644 index 0000000000..edc9d5a9e6 --- /dev/null +++ b/third_party/rust/tokio/src/io/util/read.rs @@ -0,0 +1,55 @@ +use crate::io::{AsyncRead, ReadBuf}; + +use pin_project_lite::pin_project; +use std::future::Future; +use std::io; +use std::marker::PhantomPinned; +use std::marker::Unpin; +use std::pin::Pin; +use std::task::{Context, Poll}; + +/// Tries to read some bytes directly into the given `buf` in asynchronous +/// manner, returning a future type. +/// +/// The returned future will resolve to both the I/O stream and the buffer +/// as well as the number of bytes read once the read operation is completed. +pub(crate) fn read<'a, R>(reader: &'a mut R, buf: &'a mut [u8]) -> Read<'a, R> +where + R: AsyncRead + Unpin + ?Sized, +{ + Read { + reader, + buf, + _pin: PhantomPinned, + } +} + +pin_project! { + /// A future which can be used to easily read available number of bytes to fill + /// a buffer. + /// + /// Created by the [`read`] function. + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct Read<'a, R: ?Sized> { + reader: &'a mut R, + buf: &'a mut [u8], + // Make this future `!Unpin` for compatibility with async trait methods. + #[pin] + _pin: PhantomPinned, + } +} + +impl<R> Future for Read<'_, R> +where + R: AsyncRead + Unpin + ?Sized, +{ + type Output = io::Result<usize>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> { + let me = self.project(); + 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/third_party/rust/tokio/src/io/util/read_buf.rs b/third_party/rust/tokio/src/io/util/read_buf.rs new file mode 100644 index 0000000000..8ec57c0d6f --- /dev/null +++ b/third_party/rust/tokio/src/io/util/read_buf.rs @@ -0,0 +1,72 @@ +use crate::io::AsyncRead; + +use bytes::BufMut; +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}; + +pub(crate) fn read_buf<'a, R, B>(reader: &'a mut R, buf: &'a mut B) -> ReadBuf<'a, R, B> +where + R: AsyncRead + Unpin, + B: BufMut, +{ + ReadBuf { + reader, + buf, + _pin: PhantomPinned, + } +} + +pin_project! { + /// Future returned by [`read_buf`](crate::io::AsyncReadExt::read_buf). + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct ReadBuf<'a, R, B> { + reader: &'a mut R, + buf: &'a mut B, + #[pin] + _pin: PhantomPinned, + } +} + +impl<R, B> Future for ReadBuf<'_, R, B> +where + R: AsyncRead + Unpin, + B: BufMut, +{ + type Output = io::Result<usize>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> { + use crate::io::ReadBuf; + use std::mem::MaybeUninit; + + let me = self.project(); + + if !me.buf.has_remaining_mut() { + return Poll::Ready(Ok(0)); + } + + let n = { + let dst = me.buf.chunk_mut(); + let dst = unsafe { &mut *(dst as *mut _ as *mut [MaybeUninit<u8>]) }; + let mut buf = ReadBuf::uninit(dst); + let ptr = buf.filled().as_ptr(); + ready!(Pin::new(me.reader).poll_read(cx, &mut buf)?); + + // Ensure the pointer does not change from under us + assert_eq!(ptr, buf.filled().as_ptr()); + buf.filled().len() + }; + + // Safety: This is guaranteed to be the number of initialized (and read) + // bytes due to the invariants provided by `ReadBuf::filled`. + unsafe { + me.buf.advance_mut(n); + } + + Poll::Ready(Ok(n)) + } +} diff --git a/third_party/rust/tokio/src/io/util/read_exact.rs b/third_party/rust/tokio/src/io/util/read_exact.rs new file mode 100644 index 0000000000..dbdd58bae9 --- /dev/null +++ b/third_party/rust/tokio/src/io/util/read_exact.rs @@ -0,0 +1,69 @@ +use crate::io::{AsyncRead, ReadBuf}; + +use pin_project_lite::pin_project; +use std::future::Future; +use std::io; +use std::marker::PhantomPinned; +use std::marker::Unpin; +use std::pin::Pin; +use std::task::{Context, Poll}; + +/// A future which can be used to easily read exactly enough bytes to fill +/// a buffer. +/// +/// Created by the [`AsyncReadExt::read_exact`][read_exact]. +/// [read_exact]: [crate::io::AsyncReadExt::read_exact] +pub(crate) fn read_exact<'a, A>(reader: &'a mut A, buf: &'a mut [u8]) -> ReadExact<'a, A> +where + A: AsyncRead + Unpin + ?Sized, +{ + ReadExact { + reader, + buf: ReadBuf::new(buf), + _pin: PhantomPinned, + } +} + +pin_project! { + /// Creates a future which will read exactly enough bytes to fill `buf`, + /// returning an error if EOF is hit sooner. + /// + /// On success the number of bytes is returned + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct ReadExact<'a, A: ?Sized> { + reader: &'a mut A, + buf: ReadBuf<'a>, + // Make this future `!Unpin` for compatibility with async trait methods. + #[pin] + _pin: PhantomPinned, + } +} + +fn eof() -> io::Error { + io::Error::new(io::ErrorKind::UnexpectedEof, "early eof") +} + +impl<A> Future for ReadExact<'_, A> +where + A: AsyncRead + Unpin + ?Sized, +{ + type Output = io::Result<usize>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> { + 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, me.buf))?; + if me.buf.remaining() == rem { + return Err(eof()).into(); + } + } else { + return Poll::Ready(Ok(me.buf.capacity())); + } + } + } +} diff --git a/third_party/rust/tokio/src/io/util/read_int.rs b/third_party/rust/tokio/src/io/util/read_int.rs new file mode 100644 index 0000000000..164dcf5963 --- /dev/null +++ b/third_party/rust/tokio/src/io/util/read_int.rs @@ -0,0 +1,159 @@ +use crate::io::{AsyncRead, ReadBuf}; + +use bytes::Buf; +use pin_project_lite::pin_project; +use std::future::Future; +use std::io; +use std::io::ErrorKind::UnexpectedEof; +use std::marker::PhantomPinned; +use std::mem::size_of; +use std::pin::Pin; +use std::task::{Context, Poll}; + +macro_rules! reader { + ($name:ident, $ty:ty, $reader:ident) => { + reader!($name, $ty, $reader, size_of::<$ty>()); + }; + ($name:ident, $ty:ty, $reader:ident, $bytes:expr) => { + pin_project! { + #[doc(hidden)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct $name<R> { + #[pin] + src: R, + buf: [u8; $bytes], + read: u8, + // Make this future `!Unpin` for compatibility with async trait methods. + #[pin] + _pin: PhantomPinned, + } + } + + impl<R> $name<R> { + pub(crate) fn new(src: R) -> Self { + $name { + src, + buf: [0; $bytes], + read: 0, + _pin: PhantomPinned, + } + } + } + + impl<R> Future for $name<R> + where + R: AsyncRead, + { + type Output = io::Result<$ty>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let mut me = self.project(); + + if *me.read == $bytes as u8 { + return Poll::Ready(Ok(Buf::$reader(&mut &me.buf[..]))); + } + + while *me.read < $bytes as u8 { + let mut buf = ReadBuf::new(&mut me.buf[*me.read as usize..]); + + *me.read += match me.src.as_mut().poll_read(cx, &mut buf) { + Poll::Pending => return Poll::Pending, + Poll::Ready(Err(e)) => return Poll::Ready(Err(e.into())), + Poll::Ready(Ok(())) => { + let n = buf.filled().len(); + if n == 0 { + return Poll::Ready(Err(UnexpectedEof.into())); + } + + n as u8 + } + }; + } + + let num = Buf::$reader(&mut &me.buf[..]); + + Poll::Ready(Ok(num)) + } + } + }; +} + +macro_rules! reader8 { + ($name:ident, $ty:ty) => { + pin_project! { + /// Future returned from `read_u8` + #[doc(hidden)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct $name<R> { + #[pin] + reader: R, + // Make this future `!Unpin` for compatibility with async trait methods. + #[pin] + _pin: PhantomPinned, + } + } + + impl<R> $name<R> { + pub(crate) fn new(reader: R) -> $name<R> { + $name { + reader, + _pin: PhantomPinned, + } + } + } + + impl<R> Future for $name<R> + where + R: AsyncRead, + { + type Output = io::Result<$ty>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let me = self.project(); + + let mut buf = [0; 1]; + let mut buf = ReadBuf::new(&mut buf); + match me.reader.poll_read(cx, &mut buf) { + Poll::Pending => Poll::Pending, + Poll::Ready(Err(e)) => Poll::Ready(Err(e.into())), + Poll::Ready(Ok(())) => { + if buf.filled().len() == 0 { + return Poll::Ready(Err(UnexpectedEof.into())); + } + + Poll::Ready(Ok(buf.filled()[0] as $ty)) + } + } + } + } + }; +} + +reader8!(ReadU8, u8); +reader8!(ReadI8, i8); + +reader!(ReadU16, u16, get_u16); +reader!(ReadU32, u32, get_u32); +reader!(ReadU64, u64, get_u64); +reader!(ReadU128, u128, get_u128); + +reader!(ReadI16, i16, get_i16); +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); +reader!(ReadU128Le, u128, get_u128_le); + +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/third_party/rust/tokio/src/io/util/read_line.rs b/third_party/rust/tokio/src/io/util/read_line.rs new file mode 100644 index 0000000000..e641f51532 --- /dev/null +++ b/third_party/rust/tokio/src/io/util/read_line.rs @@ -0,0 +1,119 @@ +use crate::io::util::read_until::read_until_internal; +use crate::io::AsyncBufRead; + +use pin_project_lite::pin_project; +use std::future::Future; +use std::io; +use std::marker::PhantomPinned; +use std::mem; +use std::pin::Pin; +use std::string::FromUtf8Error; +use std::task::{Context, Poll}; + +pin_project! { + /// Future for the [`read_line`](crate::io::AsyncBufReadExt::read_line) method. + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct ReadLine<'a, R: ?Sized> { + reader: &'a mut R, + // This is the buffer we were provided. It will be replaced with an empty string + // while reading to postpone utf-8 handling until after reading. + output: &'a mut String, + // The actual allocation of the string is moved into this vector instead. + buf: Vec<u8>, + // The number of bytes appended to buf. This can be less than buf.len() if + // the buffer was not empty when the operation was started. + read: usize, + // Make this future `!Unpin` for compatibility with async trait methods. + #[pin] + _pin: PhantomPinned, + } +} + +pub(crate) fn read_line<'a, R>(reader: &'a mut R, string: &'a mut String) -> ReadLine<'a, R> +where + R: AsyncBufRead + ?Sized + Unpin, +{ + ReadLine { + reader, + buf: mem::take(string).into_bytes(), + output: string, + read: 0, + _pin: PhantomPinned, + } +} + +fn put_back_original_data(output: &mut String, mut vector: Vec<u8>, num_bytes_read: usize) { + let original_len = vector.len() - num_bytes_read; + vector.truncate(original_len); + *output = String::from_utf8(vector).expect("The original data must be valid utf-8."); +} + +/// This handles the various failure cases and puts the string back into `output`. +/// +/// The `truncate_on_io_error` bool is necessary because `read_to_string` and `read_line` +/// disagree on what should happen when an IO error occurs. +pub(super) fn finish_string_read( + io_res: io::Result<usize>, + utf8_res: Result<String, FromUtf8Error>, + read: usize, + output: &mut String, + truncate_on_io_error: bool, +) -> Poll<io::Result<usize>> { + match (io_res, utf8_res) { + (Ok(num_bytes), Ok(string)) => { + debug_assert_eq!(read, 0); + *output = string; + Poll::Ready(Ok(num_bytes)) + } + (Err(io_err), Ok(string)) => { + *output = string; + if truncate_on_io_error { + let original_len = output.len() - read; + output.truncate(original_len); + } + Poll::Ready(Err(io_err)) + } + (Ok(num_bytes), Err(utf8_err)) => { + debug_assert_eq!(read, 0); + put_back_original_data(output, utf8_err.into_bytes(), num_bytes); + + Poll::Ready(Err(io::Error::new( + io::ErrorKind::InvalidData, + "stream did not contain valid UTF-8", + ))) + } + (Err(io_err), Err(utf8_err)) => { + put_back_original_data(output, utf8_err.into_bytes(), read); + + Poll::Ready(Err(io_err)) + } + } +} + +pub(super) fn read_line_internal<R: AsyncBufRead + ?Sized>( + reader: Pin<&mut R>, + cx: &mut Context<'_>, + output: &mut String, + buf: &mut Vec<u8>, + read: &mut usize, +) -> Poll<io::Result<usize>> { + let io_res = ready!(read_until_internal(reader, cx, b'\n', buf, read)); + let utf8_res = String::from_utf8(mem::take(buf)); + + // At this point both buf and output are empty. The allocation is in utf8_res. + + debug_assert!(buf.is_empty()); + debug_assert!(output.is_empty()); + finish_string_read(io_res, utf8_res, *read, output, false) +} + +impl<R: AsyncBufRead + ?Sized + Unpin> Future for ReadLine<'_, R> { + type Output = io::Result<usize>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let me = self.project(); + + read_line_internal(Pin::new(*me.reader), cx, me.output, me.buf, me.read) + } +} diff --git a/third_party/rust/tokio/src/io/util/read_to_end.rs b/third_party/rust/tokio/src/io/util/read_to_end.rs new file mode 100644 index 0000000000..f4a564d7dd --- /dev/null +++ b/third_party/rust/tokio/src/io/util/read_to_end.rs @@ -0,0 +1,112 @@ +use crate::io::util::vec_with_initialized::{into_read_buf_parts, VecU8, VecWithInitialized}; +use crate::io::AsyncRead; + +use pin_project_lite::pin_project; +use std::future::Future; +use std::io; +use std::marker::PhantomPinned; +use std::mem; +use std::pin::Pin; +use std::task::{Context, Poll}; + +pin_project! { + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct ReadToEnd<'a, R: ?Sized> { + reader: &'a mut R, + buf: VecWithInitialized<&'a mut Vec<u8>>, + // The number of bytes appended to buf. This can be less than buf.len() if + // the buffer was not empty when the operation was started. + read: usize, + // Make this future `!Unpin` for compatibility with async trait methods. + #[pin] + _pin: PhantomPinned, + } +} + +pub(crate) fn read_to_end<'a, R>(reader: &'a mut R, buffer: &'a mut Vec<u8>) -> ReadToEnd<'a, R> +where + R: AsyncRead + Unpin + ?Sized, +{ + ReadToEnd { + reader, + buf: VecWithInitialized::new(buffer), + read: 0, + _pin: PhantomPinned, + } +} + +pub(super) fn read_to_end_internal<V: VecU8, R: AsyncRead + ?Sized>( + buf: &mut VecWithInitialized<V>, + mut reader: Pin<&mut R>, + num_read: &mut usize, + cx: &mut Context<'_>, +) -> Poll<io::Result<usize>> { + loop { + let ret = ready!(poll_read_to_end(buf, reader.as_mut(), cx)); + match ret { + Err(err) => return Poll::Ready(Err(err)), + Ok(0) => return Poll::Ready(Ok(mem::replace(num_read, 0))), + Ok(num) => { + *num_read += num; + } + } + } +} + +/// Tries to read from the provided AsyncRead. +/// +/// The length of the buffer is increased by the number of bytes read. +fn poll_read_to_end<V: VecU8, R: AsyncRead + ?Sized>( + buf: &mut VecWithInitialized<V>, + read: Pin<&mut R>, + cx: &mut Context<'_>, +) -> Poll<io::Result<usize>> { + // This uses an adaptive system to extend the vector when it fills. We want to + // avoid paying to allocate and zero a huge chunk of memory if the reader only + // 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); + + // Get a ReadBuf into the vector. + let mut read_buf = buf.get_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; + + // Update the length of the vector using the result of poll_read. + let read_buf_parts = into_read_buf_parts(read_buf); + buf.apply_read_buf(read_buf_parts); + + match poll_result { + Poll::Pending => { + // 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); + Poll::Pending + } + Poll::Ready(Err(err)) => { + debug_assert_eq!(filled_before, filled_after); + Poll::Ready(Err(err)) + } + Poll::Ready(Ok(())) => Poll::Ready(Ok(n)), + } +} + +impl<A> Future for ReadToEnd<'_, A> +where + A: AsyncRead + ?Sized + Unpin, +{ + type Output = io::Result<usize>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let me = self.project(); + + read_to_end_internal(me.buf, Pin::new(*me.reader), me.read, cx) + } +} diff --git a/third_party/rust/tokio/src/io/util/read_to_string.rs b/third_party/rust/tokio/src/io/util/read_to_string.rs new file mode 100644 index 0000000000..b3d82a26ba --- /dev/null +++ b/third_party/rust/tokio/src/io/util/read_to_string.rs @@ -0,0 +1,78 @@ +use crate::io::util::read_line::finish_string_read; +use crate::io::util::read_to_end::read_to_end_internal; +use crate::io::util::vec_with_initialized::VecWithInitialized; +use crate::io::AsyncRead; + +use pin_project_lite::pin_project; +use std::future::Future; +use std::marker::PhantomPinned; +use std::pin::Pin; +use std::task::{Context, Poll}; +use std::{io, mem}; + +pin_project! { + /// Future for the [`read_to_string`](super::AsyncReadExt::read_to_string) method. + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct ReadToString<'a, R: ?Sized> { + reader: &'a mut R, + // This is the buffer we were provided. It will be replaced with an empty string + // while reading to postpone utf-8 handling until after reading. + output: &'a mut String, + // The actual allocation of the string is moved into this vector instead. + buf: VecWithInitialized<Vec<u8>>, + // The number of bytes appended to buf. This can be less than buf.len() if + // the buffer was not empty when the operation was started. + read: usize, + // Make this future `!Unpin` for compatibility with async trait methods. + #[pin] + _pin: PhantomPinned, + } +} + +pub(crate) fn read_to_string<'a, R>( + reader: &'a mut R, + string: &'a mut String, +) -> ReadToString<'a, R> +where + R: AsyncRead + ?Sized + Unpin, +{ + let buf = mem::take(string).into_bytes(); + ReadToString { + reader, + buf: VecWithInitialized::new(buf), + output: string, + read: 0, + _pin: PhantomPinned, + } +} + +fn read_to_string_internal<R: AsyncRead + ?Sized>( + reader: Pin<&mut R>, + output: &mut String, + buf: &mut VecWithInitialized<Vec<u8>>, + read: &mut usize, + cx: &mut Context<'_>, +) -> Poll<io::Result<usize>> { + let io_res = ready!(read_to_end_internal(buf, reader, read, cx)); + let utf8_res = String::from_utf8(buf.take()); + + // At this point both buf and output are empty. The allocation is in utf8_res. + + debug_assert!(buf.is_empty()); + debug_assert!(output.is_empty()); + finish_string_read(io_res, utf8_res, *read, output, true) +} + +impl<A> Future for ReadToString<'_, A> +where + A: AsyncRead + ?Sized + Unpin, +{ + type Output = io::Result<usize>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let me = self.project(); + + read_to_string_internal(Pin::new(*me.reader), me.output, me.buf, me.read, cx) + } +} diff --git a/third_party/rust/tokio/src/io/util/read_until.rs b/third_party/rust/tokio/src/io/util/read_until.rs new file mode 100644 index 0000000000..90a0e8a18d --- /dev/null +++ b/third_party/rust/tokio/src/io/util/read_until.rs @@ -0,0 +1,79 @@ +use crate::io::AsyncBufRead; + +use pin_project_lite::pin_project; +use std::future::Future; +use std::io; +use std::marker::PhantomPinned; +use std::mem; +use std::pin::Pin; +use std::task::{Context, Poll}; + +pin_project! { + /// Future for the [`read_until`](crate::io::AsyncBufReadExt::read_until) method. + /// The delimiter is included in the resulting vector. + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct ReadUntil<'a, R: ?Sized> { + reader: &'a mut R, + delimiter: u8, + buf: &'a mut Vec<u8>, + // The number of bytes appended to buf. This can be less than buf.len() if + // the buffer was not empty when the operation was started. + read: usize, + // Make this future `!Unpin` for compatibility with async trait methods. + #[pin] + _pin: PhantomPinned, + } +} + +pub(crate) fn read_until<'a, R>( + reader: &'a mut R, + delimiter: u8, + buf: &'a mut Vec<u8>, +) -> ReadUntil<'a, R> +where + R: AsyncBufRead + ?Sized + Unpin, +{ + ReadUntil { + reader, + delimiter, + buf, + read: 0, + _pin: PhantomPinned, + } +} + +pub(super) fn read_until_internal<R: AsyncBufRead + ?Sized>( + mut reader: Pin<&mut R>, + cx: &mut Context<'_>, + delimiter: u8, + buf: &mut Vec<u8>, + read: &mut usize, +) -> Poll<io::Result<usize>> { + loop { + let (done, used) = { + let available = ready!(reader.as_mut().poll_fill_buf(cx))?; + if let Some(i) = memchr::memchr(delimiter, available) { + buf.extend_from_slice(&available[..=i]); + (true, i + 1) + } else { + buf.extend_from_slice(available); + (false, available.len()) + } + }; + reader.as_mut().consume(used); + *read += used; + if done || used == 0 { + return Poll::Ready(Ok(mem::replace(read, 0))); + } + } +} + +impl<R: AsyncBufRead + ?Sized + Unpin> Future for ReadUntil<'_, R> { + type Output = io::Result<usize>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let me = self.project(); + read_until_internal(Pin::new(*me.reader), cx, *me.delimiter, me.buf, me.read) + } +} diff --git a/third_party/rust/tokio/src/io/util/repeat.rs b/third_party/rust/tokio/src/io/util/repeat.rs new file mode 100644 index 0000000000..1142765df5 --- /dev/null +++ b/third_party/rust/tokio/src/io/util/repeat.rs @@ -0,0 +1,72 @@ +use crate::io::{AsyncRead, ReadBuf}; + +use std::io; +use std::pin::Pin; +use std::task::{Context, Poll}; + +cfg_io_util! { + /// An async reader which yields one byte over and over and over and over and + /// over and... + /// + /// This struct is generally created by calling [`repeat`][repeat]. Please + /// see the documentation of `repeat()` for more details. + /// + /// This is an asynchronous version of [`std::io::Repeat`][std]. + /// + /// [repeat]: fn@repeat + /// [std]: std::io::Repeat + #[derive(Debug)] + pub struct Repeat { + byte: u8, + } + + /// Creates an instance of an async reader that infinitely repeats one byte. + /// + /// All reads from this reader will succeed by filling the specified buffer with + /// the given byte. + /// + /// This is an asynchronous version of [`std::io::repeat`][std]. + /// + /// [std]: std::io::repeat + /// + /// # Examples + /// + /// ``` + /// use tokio::io::{self, AsyncReadExt}; + /// + /// #[tokio::main] + /// async fn main() { + /// let mut buffer = [0; 3]; + /// io::repeat(0b101).read_exact(&mut buffer).await.unwrap(); + /// assert_eq!(buffer, [0b101, 0b101, 0b101]); + /// } + /// ``` + pub fn repeat(byte: u8) -> Repeat { + Repeat { byte } + } +} + +impl AsyncRead for Repeat { + #[inline] + fn poll_read( + self: Pin<&mut Self>, + _: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<()>> { + // TODO: could be faster, but should we unsafe it? + while buf.remaining() != 0 { + buf.put_slice(&[self.byte]); + } + Poll::Ready(Ok(())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn assert_unpin() { + crate::is_unpin::<Repeat>(); + } +} diff --git a/third_party/rust/tokio/src/io/util/shutdown.rs b/third_party/rust/tokio/src/io/util/shutdown.rs new file mode 100644 index 0000000000..6d30b004b1 --- /dev/null +++ b/third_party/rust/tokio/src/io/util/shutdown.rs @@ -0,0 +1,46 @@ +use crate::io::AsyncWrite; + +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! { + /// A future used to shutdown an I/O object. + /// + /// Created by the [`AsyncWriteExt::shutdown`][shutdown] function. + /// [shutdown]: crate::io::AsyncWriteExt::shutdown + #[must_use = "futures do nothing unless you `.await` or poll them"] + #[derive(Debug)] + pub struct Shutdown<'a, A: ?Sized> { + a: &'a mut A, + // Make this future `!Unpin` for compatibility with async trait methods. + #[pin] + _pin: PhantomPinned, + } +} + +/// Creates a future which will shutdown an I/O object. +pub(super) fn shutdown<A>(a: &mut A) -> Shutdown<'_, A> +where + A: AsyncWrite + Unpin + ?Sized, +{ + Shutdown { + a, + _pin: PhantomPinned, + } +} + +impl<A> Future for Shutdown<'_, A> +where + A: AsyncWrite + Unpin + ?Sized, +{ + type Output = io::Result<()>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let me = self.project(); + Pin::new(me.a).poll_shutdown(cx) + } +} diff --git a/third_party/rust/tokio/src/io/util/sink.rs b/third_party/rust/tokio/src/io/util/sink.rs new file mode 100644 index 0000000000..05ee773fa3 --- /dev/null +++ b/third_party/rust/tokio/src/io/util/sink.rs @@ -0,0 +1,87 @@ +use crate::io::AsyncWrite; + +use std::fmt; +use std::io; +use std::pin::Pin; +use std::task::{Context, Poll}; + +cfg_io_util! { + /// An async writer which will move data into the void. + /// + /// This struct is generally created by calling [`sink`][sink]. Please + /// see the documentation of `sink()` for more details. + /// + /// This is an asynchronous version of [`std::io::Sink`][std]. + /// + /// [sink]: sink() + /// [std]: std::io::Sink + pub struct Sink { + _p: (), + } + + /// Creates an instance of an async writer which will successfully consume all + /// data. + /// + /// All calls to [`poll_write`] on the returned instance will return + /// `Poll::Ready(Ok(buf.len()))` and the contents of the buffer will not be + /// inspected. + /// + /// This is an asynchronous version of [`std::io::sink`][std]. + /// + /// [`poll_write`]: crate::io::AsyncWrite::poll_write() + /// [std]: std::io::sink + /// + /// # Examples + /// + /// ``` + /// use tokio::io::{self, AsyncWriteExt}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let buffer = vec![1, 2, 3, 5, 8]; + /// let num_bytes = io::sink().write(&buffer).await?; + /// assert_eq!(num_bytes, 5); + /// Ok(()) + /// } + /// ``` + pub fn sink() -> Sink { + Sink { _p: () } + } +} + +impl AsyncWrite for Sink { + #[inline] + fn poll_write( + self: Pin<&mut Self>, + _: &mut Context<'_>, + buf: &[u8], + ) -> Poll<Result<usize, io::Error>> { + Poll::Ready(Ok(buf.len())) + } + + #[inline] + fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), io::Error>> { + Poll::Ready(Ok(())) + } + + #[inline] + fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), io::Error>> { + Poll::Ready(Ok(())) + } +} + +impl fmt::Debug for Sink { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.pad("Sink { .. }") + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn assert_unpin() { + crate::is_unpin::<Sink>(); + } +} diff --git a/third_party/rust/tokio/src/io/util/split.rs b/third_party/rust/tokio/src/io/util/split.rs new file mode 100644 index 0000000000..7489c24281 --- /dev/null +++ b/third_party/rust/tokio/src/io/util/split.rs @@ -0,0 +1,121 @@ +use crate::io::util::read_until::read_until_internal; +use crate::io::AsyncBufRead; + +use pin_project_lite::pin_project; +use std::io; +use std::mem; +use std::pin::Pin; +use std::task::{Context, Poll}; + +pin_project! { + /// Splitter for the [`split`](crate::io::AsyncBufReadExt::split) method. + /// + /// A `Split` can be turned into a `Stream` with [`SplitStream`]. + /// + /// [`SplitStream`]: https://docs.rs/tokio-stream/0.1/tokio_stream/wrappers/struct.SplitStream.html + #[derive(Debug)] + #[must_use = "streams do nothing unless polled"] + #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] + pub struct Split<R> { + #[pin] + reader: R, + buf: Vec<u8>, + delim: u8, + read: usize, + } +} + +pub(crate) fn split<R>(reader: R, delim: u8) -> Split<R> +where + R: AsyncBufRead, +{ + Split { + reader, + buf: Vec::new(), + delim, + read: 0, + } +} + +impl<R> Split<R> +where + R: AsyncBufRead + Unpin, +{ + /// Returns the next segment in the stream. + /// + /// # Examples + /// + /// ``` + /// # use tokio::io::AsyncBufRead; + /// use tokio::io::AsyncBufReadExt; + /// + /// # async fn dox(my_buf_read: impl AsyncBufRead + Unpin) -> std::io::Result<()> { + /// let mut segments = my_buf_read.split(b'f'); + /// + /// while let Some(segment) = segments.next_segment().await? { + /// println!("length = {}", segment.len()) + /// } + /// # Ok(()) + /// # } + /// ``` + pub async fn next_segment(&mut self) -> io::Result<Option<Vec<u8>>> { + use crate::future::poll_fn; + + poll_fn(|cx| Pin::new(&mut *self).poll_next_segment(cx)).await + } +} + +impl<R> Split<R> +where + R: AsyncBufRead, +{ + /// Polls for the next segment in the stream. + /// + /// This method returns: + /// + /// * `Poll::Pending` if the next segment is not yet available. + /// * `Poll::Ready(Ok(Some(segment)))` if the next segment is available. + /// * `Poll::Ready(Ok(None))` if there are no more segments in this stream. + /// * `Poll::Ready(Err(err))` if an IO error occurred while reading the + /// next segment. + /// + /// 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_segment`, only the `Waker` + /// from the `Context` passed to the most recent call is scheduled to + /// receive a wakeup. + pub fn poll_next_segment( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<io::Result<Option<Vec<u8>>>> { + let me = self.project(); + + let n = ready!(read_until_internal( + me.reader, cx, *me.delim, me.buf, me.read, + ))?; + // read_until_internal resets me.read to zero once it finds the delimiter + debug_assert_eq!(*me.read, 0); + + if n == 0 && me.buf.is_empty() { + return Poll::Ready(Ok(None)); + } + + if me.buf.last() == Some(me.delim) { + me.buf.pop(); + } + + Poll::Ready(Ok(Some(mem::take(me.buf)))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn assert_unpin() { + crate::is_unpin::<Split<()>>(); + } +} diff --git a/third_party/rust/tokio/src/io/util/take.rs b/third_party/rust/tokio/src/io/util/take.rs new file mode 100644 index 0000000000..df2f61b9e6 --- /dev/null +++ b/third_party/rust/tokio/src/io/util/take.rs @@ -0,0 +1,137 @@ +use crate::io::{AsyncBufRead, AsyncRead, ReadBuf}; + +use pin_project_lite::pin_project; +use std::pin::Pin; +use std::task::{Context, Poll}; +use std::{cmp, io}; + +pin_project! { + /// Stream for the [`take`](super::AsyncReadExt::take) method. + #[derive(Debug)] + #[must_use = "streams do nothing unless you `.await` or poll them"] + #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] + pub struct Take<R> { + #[pin] + inner: R, + // Add '_' to avoid conflicts with `limit` method. + limit_: u64, + } +} + +pub(super) fn take<R: AsyncRead>(inner: R, limit: u64) -> Take<R> { + Take { + inner, + limit_: limit, + } +} + +impl<R: AsyncRead> Take<R> { + /// Returns the remaining number of bytes that can be + /// read before this instance will return EOF. + /// + /// # Note + /// + /// This instance may reach `EOF` after reading fewer bytes than indicated by + /// this method if the underlying [`AsyncRead`] instance reaches EOF. + pub fn limit(&self) -> u64 { + self.limit_ + } + + /// Sets the number of bytes that can be read before this instance will + /// return EOF. This is the same as constructing a new `Take` instance, so + /// the amount of bytes read and the previous limit value don't matter when + /// calling this method. + pub fn set_limit(&mut self, limit: u64) { + self.limit_ = limit + } + + /// Gets a reference to the underlying reader. + pub fn get_ref(&self) -> &R { + &self.inner + } + + /// Gets a mutable reference to the underlying reader. + /// + /// Care should be taken to avoid modifying the internal I/O state of the + /// underlying reader as doing so may corrupt the internal limit of this + /// `Take`. + pub fn get_mut(&mut self) -> &mut R { + &mut self.inner + } + + /// Gets a pinned mutable reference to the underlying reader. + /// + /// Care should be taken to avoid modifying the internal I/O state of the + /// underlying reader as doing so may corrupt the internal limit of this + /// `Take`. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut R> { + self.project().inner + } + + /// Consumes the `Take`, returning the wrapped reader. + pub fn into_inner(self) -> R { + self.inner + } +} + +impl<R: AsyncRead> AsyncRead for Take<R> { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<Result<(), io::Error>> { + if self.limit_ == 0 { + return Poll::Ready(Ok(())); + } + + 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 + unsafe { + buf.assume_init(n); + } + buf.advance(n); + *me.limit_ -= n as u64; + Poll::Ready(Ok(())) + } +} + +impl<R: AsyncBufRead> AsyncBufRead for Take<R> { + fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> { + let me = self.project(); + + // Don't call into inner reader at all at EOF because it may still block + if *me.limit_ == 0 { + return Poll::Ready(Ok(&[])); + } + + let buf = ready!(me.inner.poll_fill_buf(cx)?); + let cap = cmp::min(buf.len() as u64, *me.limit_) as usize; + Poll::Ready(Ok(&buf[..cap])) + } + + fn consume(self: Pin<&mut Self>, amt: usize) { + let me = self.project(); + // Don't let callers reset the limit by passing an overlarge value + let amt = cmp::min(amt as u64, *me.limit_) as usize; + *me.limit_ -= amt as u64; + me.inner.consume(amt); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn assert_unpin() { + crate::is_unpin::<Take<()>>(); + } +} diff --git a/third_party/rust/tokio/src/io/util/vec_with_initialized.rs b/third_party/rust/tokio/src/io/util/vec_with_initialized.rs new file mode 100644 index 0000000000..208cc939c1 --- /dev/null +++ b/third_party/rust/tokio/src/io/util/vec_with_initialized.rs @@ -0,0 +1,132 @@ +use crate::io::ReadBuf; +use std::mem::MaybeUninit; + +mod private { + pub trait Sealed {} + + impl Sealed for Vec<u8> {} + impl Sealed for &mut 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 {} + +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. +/// +/// The purpose of this struct is to remember how many bytes were initialized +/// through a `ReadBuf` from call to call. +/// +/// This struct has the safety invariant that the first `num_initialized` of the +/// vector's allocation must be initialized at any time. +#[derive(Debug)] +pub(crate) struct VecWithInitialized<V> { + vec: V, + // The number of initialized bytes in the vector. + // Always between `vec.len()` and `vec.capacity()`. + num_initialized: usize, +} + +impl VecWithInitialized<Vec<u8>> { + #[cfg(feature = "io-util")] + pub(crate) fn take(&mut self) -> Vec<u8> { + self.num_initialized = 0; + std::mem::take(&mut self.vec) + } +} + +impl<V> VecWithInitialized<V> +where + V: VecU8, +{ + pub(crate) fn new(mut vec: V) -> Self { + // SAFETY: The safety invariants of vector guarantee that the bytes up + // to its length are initialized. + Self { + num_initialized: vec.as_mut().len(), + vec, + } + } + + pub(crate) fn reserve(&mut self, num_bytes: usize) { + let vec = self.vec.as_mut(); + if vec.capacity() - vec.len() >= num_bytes { + return; + } + // SAFETY: Setting num_initialized to `vec.len()` is correct as + // `reserve` does not change the length of the vector. + self.num_initialized = vec.len(); + vec.reserve(num_bytes); + } + + #[cfg(feature = "io-util")] + pub(crate) fn is_empty(&mut self) -> bool { + self.vec.as_mut().is_empty() + } + + pub(crate) fn get_read_buf<'a>(&'a mut self) -> ReadBuf<'a> { + let num_initialized = self.num_initialized; + + // SAFETY: Creating the slice is safe because of the safety invariants + // on Vec<u8>. The safety invariants of `ReadBuf` will further guarantee + // that no bytes in the slice are de-initialized. + let vec = self.vec.as_mut(); + let len = vec.len(); + let cap = vec.capacity(); + let ptr = vec.as_mut_ptr().cast::<MaybeUninit<u8>>(); + let slice = unsafe { std::slice::from_raw_parts_mut::<'a, MaybeUninit<u8>>(ptr, cap) }; + + // SAFETY: This is safe because the safety invariants of + // VecWithInitialized say that the first num_initialized bytes must be + // initialized. + let mut read_buf = ReadBuf::uninit(slice); + unsafe { + read_buf.assume_init(num_initialized); + } + read_buf.set_filled(len); + + read_buf + } + + pub(crate) fn apply_read_buf(&mut self, parts: ReadBufParts) { + let vec = self.vec.as_mut(); + assert_eq!(vec.as_ptr(), parts.ptr); + + // SAFETY: + // The ReadBufParts really does point inside `self.vec` due to the above + // check, and the safety invariants of `ReadBuf` guarantee that the + // first `parts.initialized` bytes of `self.vec` really have been + // initialized. Additionally, `ReadBuf` guarantees that `parts.len` is + // at most `parts.initialized`, so the first `parts.len` bytes are also + // initialized. + // + // Note that this relies on the fact that `V` is either `Vec<u8>` or + // `&mut Vec<u8>`, so the vector returned by `self.vec.as_mut()` cannot + // change from call to call. + unsafe { + self.num_initialized = parts.initialized; + vec.set_len(parts.len); + } + } +} + +pub(crate) struct ReadBufParts { + // Pointer is only used to check that the ReadBuf actually came from the + // right VecWithInitialized. + ptr: *const u8, + len: usize, + initialized: usize, +} + +// This is needed to release the borrow on `VecWithInitialized<V>`. +pub(crate) fn into_read_buf_parts(rb: ReadBuf<'_>) -> ReadBufParts { + ReadBufParts { + ptr: rb.filled().as_ptr(), + len: rb.filled().len(), + initialized: rb.initialized().len(), + } +} diff --git a/third_party/rust/tokio/src/io/util/write.rs b/third_party/rust/tokio/src/io/util/write.rs new file mode 100644 index 0000000000..92169ebc1d --- /dev/null +++ b/third_party/rust/tokio/src/io/util/write.rs @@ -0,0 +1,46 @@ +use crate::io::AsyncWrite; + +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! { + /// A future to write some of the buffer to an `AsyncWrite`. + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct Write<'a, W: ?Sized> { + writer: &'a mut W, + buf: &'a [u8], + // Make this future `!Unpin` for compatibility with async trait methods. + #[pin] + _pin: PhantomPinned, + } +} + +/// Tries to write some bytes from the given `buf` to the writer in an +/// asynchronous manner, returning a future. +pub(crate) fn write<'a, W>(writer: &'a mut W, buf: &'a [u8]) -> Write<'a, W> +where + W: AsyncWrite + Unpin + ?Sized, +{ + Write { + writer, + buf, + _pin: PhantomPinned, + } +} + +impl<W> Future for Write<'_, W> +where + W: AsyncWrite + Unpin + ?Sized, +{ + type Output = io::Result<usize>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> { + let me = self.project(); + Pin::new(&mut *me.writer).poll_write(cx, me.buf) + } +} diff --git a/third_party/rust/tokio/src/io/util/write_all.rs b/third_party/rust/tokio/src/io/util/write_all.rs new file mode 100644 index 0000000000..abd3e39d31 --- /dev/null +++ b/third_party/rust/tokio/src/io/util/write_all.rs @@ -0,0 +1,55 @@ +use crate::io::AsyncWrite; + +use pin_project_lite::pin_project; +use std::future::Future; +use std::io; +use std::marker::PhantomPinned; +use std::mem; +use std::pin::Pin; +use std::task::{Context, Poll}; + +pin_project! { + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct WriteAll<'a, W: ?Sized> { + writer: &'a mut W, + buf: &'a [u8], + // Make this future `!Unpin` for compatibility with async trait methods. + #[pin] + _pin: PhantomPinned, + } +} + +pub(crate) fn write_all<'a, W>(writer: &'a mut W, buf: &'a [u8]) -> WriteAll<'a, W> +where + W: AsyncWrite + Unpin + ?Sized, +{ + WriteAll { + writer, + buf, + _pin: PhantomPinned, + } +} + +impl<W> Future for WriteAll<'_, W> +where + W: AsyncWrite + Unpin + ?Sized, +{ + type Output = io::Result<()>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + let me = self.project(); + while !me.buf.is_empty() { + let n = ready!(Pin::new(&mut *me.writer).poll_write(cx, me.buf))?; + { + let (_, rest) = mem::take(&mut *me.buf).split_at(n); + *me.buf = rest; + } + if n == 0 { + return Poll::Ready(Err(io::ErrorKind::WriteZero.into())); + } + } + + Poll::Ready(Ok(())) + } +} diff --git a/third_party/rust/tokio/src/io/util/write_all_buf.rs b/third_party/rust/tokio/src/io/util/write_all_buf.rs new file mode 100644 index 0000000000..05af7fe99b --- /dev/null +++ b/third_party/rust/tokio/src/io/util/write_all_buf.rs @@ -0,0 +1,56 @@ +use crate::io::AsyncWrite; + +use bytes::Buf; +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! { + /// A future to write some of the buffer to an `AsyncWrite`. + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct WriteAllBuf<'a, W, B> { + writer: &'a mut W, + buf: &'a mut B, + #[pin] + _pin: PhantomPinned, + } +} + +/// Tries to write some bytes from the given `buf` to the writer in an +/// asynchronous manner, returning a future. +pub(crate) fn write_all_buf<'a, W, B>(writer: &'a mut W, buf: &'a mut B) -> WriteAllBuf<'a, W, B> +where + W: AsyncWrite + Unpin, + B: Buf, +{ + WriteAllBuf { + writer, + buf, + _pin: PhantomPinned, + } +} + +impl<W, B> Future for WriteAllBuf<'_, W, B> +where + W: AsyncWrite + Unpin, + B: Buf, +{ + type Output = io::Result<()>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + let me = self.project(); + while me.buf.has_remaining() { + let n = ready!(Pin::new(&mut *me.writer).poll_write(cx, me.buf.chunk())?); + me.buf.advance(n); + if n == 0 { + return Poll::Ready(Err(io::ErrorKind::WriteZero.into())); + } + } + + Poll::Ready(Ok(())) + } +} diff --git a/third_party/rust/tokio/src/io/util/write_buf.rs b/third_party/rust/tokio/src/io/util/write_buf.rs new file mode 100644 index 0000000000..82fd7a759f --- /dev/null +++ b/third_party/rust/tokio/src/io/util/write_buf.rs @@ -0,0 +1,55 @@ +use crate::io::AsyncWrite; + +use bytes::Buf; +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! { + /// A future to write some of the buffer to an `AsyncWrite`. + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct WriteBuf<'a, W, B> { + writer: &'a mut W, + buf: &'a mut B, + #[pin] + _pin: PhantomPinned, + } +} + +/// Tries to write some bytes from the given `buf` to the writer in an +/// asynchronous manner, returning a future. +pub(crate) fn write_buf<'a, W, B>(writer: &'a mut W, buf: &'a mut B) -> WriteBuf<'a, W, B> +where + W: AsyncWrite + Unpin, + B: Buf, +{ + WriteBuf { + writer, + buf, + _pin: PhantomPinned, + } +} + +impl<W, B> Future for WriteBuf<'_, W, B> +where + W: AsyncWrite + Unpin, + B: Buf, +{ + type Output = io::Result<usize>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> { + let me = self.project(); + + if !me.buf.has_remaining() { + return Poll::Ready(Ok(0)); + } + + let n = ready!(Pin::new(me.writer).poll_write(cx, me.buf.chunk()))?; + me.buf.advance(n); + Poll::Ready(Ok(n)) + } +} diff --git a/third_party/rust/tokio/src/io/util/write_int.rs b/third_party/rust/tokio/src/io/util/write_int.rs new file mode 100644 index 0000000000..63cd49126f --- /dev/null +++ b/third_party/rust/tokio/src/io/util/write_int.rs @@ -0,0 +1,152 @@ +use crate::io::AsyncWrite; + +use bytes::BufMut; +use pin_project_lite::pin_project; +use std::future::Future; +use std::io; +use std::marker::PhantomPinned; +use std::mem::size_of; +use std::pin::Pin; +use std::task::{Context, Poll}; + +macro_rules! writer { + ($name:ident, $ty:ty, $writer:ident) => { + writer!($name, $ty, $writer, size_of::<$ty>()); + }; + ($name:ident, $ty:ty, $writer:ident, $bytes:expr) => { + pin_project! { + #[doc(hidden)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct $name<W> { + #[pin] + dst: W, + buf: [u8; $bytes], + written: u8, + // Make this future `!Unpin` for compatibility with async trait methods. + #[pin] + _pin: PhantomPinned, + } + } + + impl<W> $name<W> { + pub(crate) fn new(w: W, value: $ty) -> Self { + let mut writer = Self { + buf: [0; $bytes], + written: 0, + dst: w, + _pin: PhantomPinned, + }; + BufMut::$writer(&mut &mut writer.buf[..], value); + writer + } + } + + impl<W> Future for $name<W> + where + W: AsyncWrite, + { + type Output = io::Result<()>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let mut me = self.project(); + + if *me.written == $bytes as u8 { + return Poll::Ready(Ok(())); + } + + while *me.written < $bytes as u8 { + *me.written += match me + .dst + .as_mut() + .poll_write(cx, &me.buf[*me.written as usize..]) + { + Poll::Pending => return Poll::Pending, + Poll::Ready(Err(e)) => return Poll::Ready(Err(e.into())), + Poll::Ready(Ok(0)) => { + return Poll::Ready(Err(io::ErrorKind::WriteZero.into())); + } + Poll::Ready(Ok(n)) => n as u8, + }; + } + Poll::Ready(Ok(())) + } + } + }; +} + +macro_rules! writer8 { + ($name:ident, $ty:ty) => { + pin_project! { + #[doc(hidden)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct $name<W> { + #[pin] + dst: W, + byte: $ty, + // Make this future `!Unpin` for compatibility with async trait methods. + #[pin] + _pin: PhantomPinned, + } + } + + impl<W> $name<W> { + pub(crate) fn new(dst: W, byte: $ty) -> Self { + Self { + dst, + byte, + _pin: PhantomPinned, + } + } + } + + impl<W> Future for $name<W> + where + W: AsyncWrite, + { + type Output = io::Result<()>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let me = self.project(); + + let buf = [*me.byte as u8]; + + match me.dst.poll_write(cx, &buf[..]) { + Poll::Pending => Poll::Pending, + Poll::Ready(Err(e)) => Poll::Ready(Err(e.into())), + Poll::Ready(Ok(0)) => Poll::Ready(Err(io::ErrorKind::WriteZero.into())), + Poll::Ready(Ok(1)) => Poll::Ready(Ok(())), + Poll::Ready(Ok(_)) => unreachable!(), + } + } + } + }; +} + +writer8!(WriteU8, u8); +writer8!(WriteI8, i8); + +writer!(WriteU16, u16, put_u16); +writer!(WriteU32, u32, put_u32); +writer!(WriteU64, u64, put_u64); +writer!(WriteU128, u128, put_u128); + +writer!(WriteI16, i16, put_i16); +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); +writer!(WriteU128Le, u128, put_u128_le); + +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/third_party/rust/tokio/src/io/util/write_vectored.rs b/third_party/rust/tokio/src/io/util/write_vectored.rs new file mode 100644 index 0000000000..be40322943 --- /dev/null +++ b/third_party/rust/tokio/src/io/util/write_vectored.rs @@ -0,0 +1,47 @@ +use crate::io::AsyncWrite; + +use pin_project_lite::pin_project; +use std::io; +use std::marker::PhantomPinned; +use std::pin::Pin; +use std::task::{Context, Poll}; +use std::{future::Future, io::IoSlice}; + +pin_project! { + /// A future to write a slice of buffers to an `AsyncWrite`. + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct WriteVectored<'a, 'b, W: ?Sized> { + writer: &'a mut W, + bufs: &'a [IoSlice<'b>], + // Make this future `!Unpin` for compatibility with async trait methods. + #[pin] + _pin: PhantomPinned, + } +} + +pub(crate) fn write_vectored<'a, 'b, W>( + writer: &'a mut W, + bufs: &'a [IoSlice<'b>], +) -> WriteVectored<'a, 'b, W> +where + W: AsyncWrite + Unpin + ?Sized, +{ + WriteVectored { + writer, + bufs, + _pin: PhantomPinned, + } +} + +impl<W> Future for WriteVectored<'_, '_, W> +where + W: AsyncWrite + Unpin + ?Sized, +{ + type Output = io::Result<usize>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> { + let me = self.project(); + Pin::new(&mut *me.writer).poll_write_vectored(cx, me.bufs) + } +} |