async-stdio 0.3.0-alpha.3

//! # Adapter for using async read/write streams in [std::io] contexts
//!
//! Sometimes, you'll come across an interface that only takes [std::io] `Read + Write`
//! types, but also needs to be adapted for an async/await application. The
//! [AsStdIo] adapter allows an [AsyncRead] + [AsyncWrite] stream to be used
//! as its counterpart from [std::io]. Assuming that whatever is consuming the
//! wrapped stream will bubble up [io::ErrorKind::WouldBlock] errors and allows
//! operations to be resumed, this provides a way to both use an async stream
//! with the [std::io]-only interface, and to write an async wrapper around it.
//!
//! ## Example
//!
//! ```rust
//! # use std::collections::VecDeque;
//! # use std::io::{self, Read};
//! # use std::pin::Pin;
//! #
//! # use futures::prelude::*;
//! # use futures::ready;
//! # use futures::task::{Context, Poll};
//! #
//! use async_stdio::*;
//!
//! struct ChunkReader<R> {
//!     // ...
//!     # reader: R,
//!     # chunk_size: usize,
//!     # buffer: VecDeque<u8>,
//! }
//!
//! impl<R: Read> ChunkReader<R> {
//!     fn new(reader: R, chunk_size: usize) -> Self {
//!         // ...
//!         # ChunkReader {
//!         #     reader,
//!         #     chunk_size,
//!         #     buffer: Default::default(),
//!         # }
//!     }
//!
//!     /// Reads a chunk from the stream
//!     ///
//!     /// If the stream ends before a full chunk is read, may return a smaller
//!     /// chunk. Returns an empty chunk if there is no more to be read.
//!     fn read_chunk(&mut self) -> io::Result<Vec<u8>> {
//!         // ...
//!         # let mut tmp = vec![0u8; self.chunk_size];
//!         # let mut bytes = self.chunk_size;
//!         # loop {
//!         #     if self.buffer.len() >= self.chunk_size || bytes == 0 {
//!         #         let end = self.buffer.len().min(self.chunk_size);
//!         #         tmp.truncate(0);
//!         #         return Ok(self.buffer.drain(..end).fold(tmp, |mut out, b| {
//!         #             out.push(b);
//!         #             out
//!         #         }));
//!         #     }
//!         #     bytes = self.reader.read(&mut tmp)?;
//!         #     self.buffer.extend(&tmp[..bytes]);
//!         # }
//!     }
//! }
//!
//! /// Wrapper around the std-only `ChunkReader` to turn it
//! /// into an async `Stream`
//! struct AsyncChunked<S> {
//!     inner: ChunkReader<AsStdIo<S>>,
//!     waker_ctrl: WakerCtrl,
//! }
//!
//! impl<S: AsyncRead + Unpin> AsyncChunked<S> {
//!     fn new(stream: S, chunk_size: usize) -> AsyncChunked<S> {
//!         let (stream, waker_ctrl) = AsStdIo::new(stream, None);
//!         let inner = ChunkReader::new(stream, chunk_size);
//!         AsyncChunked { inner, waker_ctrl }
//!     }
//! }
//!
//! impl<S: AsyncRead + Unpin> Stream for AsyncChunked<S> {
//!     type Item = io::Result<Vec<u8>>;
//!
//!     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
//!         let this = self.get_mut();
//!         // Make sure the waker is set before the calls to `std::io::Read::read`
//!         this.waker_ctrl.register(cx.waker());
//!         // `into_poll` (from `ResultExt`) converts `WouldBlock` into `Pending`
//!         let chunk_res = ready!(this.inner.read_chunk().into_poll());
//!
//!         Poll::Ready(
//!             chunk_res
//!                 .map(|chunk| if chunk.is_empty() { None } else { Some(chunk) })
//!                 .transpose(),
//!         )
//!     }
//! }
//!
//! // Pretend this doesn't already implement `io::Read`
//! let stream = io::Cursor::new(vec![0, 1, 2, 3, 4, 5]);
//! let mut async_chunked = AsyncChunked::new(stream, 2);
//!
//! # use futures::executor::block_on;
//! let chunks: Vec<Vec<u8>> = block_on(async_chunked.map(|chunk| chunk.unwrap()).collect());
//!
//! let expected: Vec<Vec<u8>> = vec![vec![0, 1], vec![2, 3], vec![4, 5]];
//!
//! assert_eq!(chunks, expected,);
//! ```

#![warn(missing_docs)]

use std::cell::RefCell;
use std::io;
use std::marker::Unpin;
use std::ops::{Deref, DerefMut};
use std::pin::Pin;
use std::sync::Arc;

use thread_local::ThreadLocal;

use futures::io::{AsyncRead, AsyncWrite};
use futures::task::{Context, Poll, Waker};

/// Adapter to use an async byte stream as a `std::io::{Read, Write}` type
///
/// For correct operation, `WakerCtrl::register` must be called before any calls
/// to `read` or `write`. If the async stream would return `Poll::Pending`, it
/// will be converted to an `io::ErrorKind::WouldBlock` error.
#[derive(Debug)]
pub struct AsStdIo<S> {
    inner: S,
    waker: Arc<ThreadLocal<RefCell<Waker>>>,
}

/// Waker control for the [AsStdIo] wrapper
#[derive(Debug)]
pub struct WakerCtrl(Arc<ThreadLocal<RefCell<Waker>>>);

impl WakerCtrl {
    /// Register a waker to be used the next time the corresponding `AsStdIo`
    /// wrapper performs an IO operation.
    ///
    /// This *must* be called before every operation that may perform IO. It is
    /// an error to call `register` from a different thread than the one
    /// attempting the IO operation.
    pub fn register(&self, waker: &Waker) {
        *self
            .0
            .get_or(|| Box::new(RefCell::new(futures::task::noop_waker())))
            .borrow_mut() = waker.clone();
    }
}

/// Extension traits for `Poll` and `io::Result`
pub mod ext {
    use super::*;

    /// Extension trait for converting an `io::Result<T>` to a `Poll<io::Result<T>>`
    pub trait ResultExt<T> {
        /// Perform the conversion
        ///
        /// An IO error with the kind [io::ErrorKind::WouldBlock] will get converted
        /// to [Poll::Pending]. All other results will become [Poll::Ready]
        fn into_poll(self) -> Poll<io::Result<T>>;
    }

    impl<T> ResultExt<T> for io::Result<T> {
        fn into_poll(self) -> Poll<io::Result<T>> {
            if let Err(ref e) = self {
                if e.kind() == io::ErrorKind::WouldBlock {
                    return Poll::Pending;
                }
            }

            Poll::Ready(self)
        }
    }

    /// Extension trait for converting a `Poll<io::Result<T>>` into an `io::Result<T>`
    pub trait PollExt<T> {
        /// Perform the conversion
        ///
        /// A `Poll::Pending` will be turned into an IO error with the kind
        /// [io::ErrorKind::WouldBlock]. [Poll::Ready] will simply be unwrapped.
        fn into_io_result(self) -> io::Result<T>;
    }

    impl<T> PollExt<T> for Poll<io::Result<T>> {
        fn into_io_result(self) -> io::Result<T> {
            match self {
                Poll::Ready(res) => res,
                Poll::Pending => Err(io::ErrorKind::WouldBlock.into()),
            }
        }
    }
}

pub use ext::{PollExt as _, ResultExt as _};

impl<S> Deref for AsStdIo<S> {
    type Target = S;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<S> DerefMut for AsStdIo<S> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

impl<S> AsStdIo<S> {
    /// Wrap the async stream and initialize the wrapper with the provided waker.
    ///
    /// If no waker is provided, it will be initialized with a no-op waker.
    ///
    /// Also returns a handle that can be used to control the waker.
    pub fn new(stream: S, waker: Option<&Waker>) -> (Self, WakerCtrl) {
        let tls_waker = ThreadLocal::<RefCell<Waker>>::new();
        let tls_handle = Arc::new(tls_waker);
        let waker_ctrl = WakerCtrl(tls_handle.clone());
        if let Some(waker) = waker {
            waker_ctrl.register(waker);
        }
        (
            AsStdIo {
                inner: stream,
                waker: tls_handle,
            },
            waker_ctrl,
        )
    }

    /// Get a pinned reference to the internal stream and a context with which to
    /// poll it.
    ///
    /// Sometimes useful if you need to perform some operation such as
    /// [AsyncWrite::close] on it which doesn't have an `std` counterpart.
    pub fn with_context<F, T>(&mut self, f: F) -> T
    where
        F: for<'c> FnOnce(Pin<&mut S>, &mut Context<'c>) -> T,
        S: Unpin,
    {
        let waker = self
            .waker
            .get_or(|| Box::new(RefCell::new(futures::task::noop_waker())))
            .borrow();
        f(Pin::new(&mut self.inner), &mut Context::from_waker(&*waker))
    }
}

impl<S> io::Read for AsStdIo<S>
where
    S: AsyncRead + Unpin,
{
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.with_context(|stream, cx| stream.poll_read(cx, buf).into_io_result())
    }
}

impl<S> io::Write for AsStdIo<S>
where
    S: AsyncWrite + Unpin,
{
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.with_context(|stream, cx| stream.poll_write(cx, buf).into_io_result())
    }

    fn flush(&mut self) -> io::Result<()> {
        self.with_context(|stream, cx| stream.poll_flush(cx).into_io_result())
    }
}