Using AsyncRead and AsyncWrite directly

So far, we have primarily talked about AsyncRead and AsyncWrite inthe context of I/O combinators provided by Tokio. While these are oftenenough, sometimes you need to implement your own combinators thatwant to perform asynchronous reads and writes directly.

Reading data with AsyncRead

The heart of AsyncRead is the poll_read method. It maps theWouldBlock error that indicates that an I/O read operation wouldhave blocked into NotReady, which in turn lets us interoperate withthe world of futures. When you write a Future (or something like it,such as Stream) that internally contains an AsyncRead, poll_readis likely the method you will be interacting with.

The important thing to keep in mind with poll_read is that it followsthe same contract as Future::poll. Specifically, it cannot returnNotReady unless it has arranged for the current task to be notifiedwhen it can make progress again. This fact is what lets us callpoll_read inside of poll in our own Futures; we know that we areupholding the contract of poll when we forward a NotReady frompoll_read, because poll_read follows that same contract!

The exact mechanism Tokio uses to ensure that poll_read later notifiesthe current task is out of scope for this section, but you can read moreabout it in the non-blocking I/O section of Tokio internals if you’reinterested.

With that all said, let’s look at how we might implement theread_exact method ourselves!

  1. # extern crate tokio;
  2. #[macro_use]
  3. extern crate futures;
  4. # fn main() {}
  5. use std::io;
  6. use tokio::prelude::*;
  8. // This is going to be our Future.
  9. // In the common case, this is set to Some(Reading),
  10. // but we'll set it to None when we return Async::Ready
  11. // so that we can return the reader and the buffer.
  12. struct ReadExact<R, T>(Option<Reading<R, T>>);
  14. struct Reading<R, T> {
  15.     // This is the stream we're reading from.
  16.     reader: R,
  17.     // This is the buffer we're reading into.
  18.     buffer: T,
  19.     // And this is how far into the buffer we've written.
  20.     pos: usize,
  21. }
  23. // We want to be able to construct a ReadExact over anything
  24. // that implements AsyncRead, and any buffer that can be
  25. // thought of as a &mut [u8].
  26. fn read_exact<R, T>(reader: R, buffer: T) -> ReadExact<R, T>
  27. where
  28.     R: AsyncRead,
  29.     T: AsMut<[u8]>,
  30. {
  31.     ReadExact(Some(Reading {
  32.         reader,
  33.         buffer,
  34.         // Initially, we've read no bytes into buffer.
  35.         pos: 0,
  36.     }))
  37. }
  39. impl<R, T> Future for ReadExact<R, T>
  40. where
  41.     R: AsyncRead,
  42.     T: AsMut<[u8]>,
  43. {
  44.     // When we've filled up the buffer, we want to return both the buffer
  45.     // with the data that we read and the reader itself.
  46.     type Item = (R, T);
  47.     type Error = io::Error;
  49.     fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
  50.         match self.0 {
  51.             Some(Reading {
  52.                 ref mut reader,
  53.                 ref mut buffer,
  54.                 ref mut pos,
  55.             }) => {
  56.                 let buffer = buffer.as_mut();
  57.                 // Check that we haven't finished
  58.                 while *pos < buffer.len() {
  59.                     // Try to read data into the remainder of the buffer.
  60.                     // Just like read in std::io::Read, poll_read *can* read
  61.                     // fewer bytes than the length of the buffer it is given,
  62.                     // and we need to handle that by looking at its return
  63.                     // value, which is the number of bytes actually read.
  64.                     //
  65.                     // Notice that we are using try_ready! here, so if poll_read
  66.                     // returns NotReady (or an error), we will do the same!
  67.                     // We uphold the contract that we have arranged to be
  68.                     // notified later because poll_read follows that same
  69.                     // contract, and _it_ returned NotReady.
  70.                     let n = try_ready!(reader.poll_read(&mut buffer[*pos..]));
  71.                     *pos += n;
  73.                     // If no bytes were read, but there was no error, this
  74.                     // generally implies that the reader will provide no more
  75.                     // data (for example, because the TCP connection was closed
  76.                     // by the other side).
  77.                     if n == 0 {
  78.                         return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "early eof"));
  79.                     }
  80.                 }
  81.             }
  82.             None => panic!("poll a ReadExact after it's done"),
  83.         }
  85.         // We need to return the reader and the buffer, which we can only
  86.         // do by moving them out of self. We do this by taking our state
  87.         // and leaving `None`. This _should_ be fine, because poll()
  88.         // requires callers to not call poll() again after Ready has been
  89.         // returned, so we should only ever see Some(Reading) when poll()
  90.         // is called.
  91.         let reading = self.0.take().expect("must have seen Some above");
  92.         Ok(Async::Ready((reading.reader, reading.buffer)))
  93.     }
  94. }

Writing data with AsyncWrite

Just like pollread is the core piece of AsyncRead, poll_write isthe core of AsyncWrite. Like poll_read, it maps the WouldBlockerror that indicates that an I/O write_ operation would have blocked intoNotReady, which again lets us interoperate with the world offutures. AsyncWrite also has a poll_flush, which provides anasynchronous analogue to Write’s flush method. The role ofpoll_flush is to make sure that any bytes previously written bypoll_write are, well, flushed onto the underlying I/O resource(written out in network packets for example). Similar to poll_write, itwraps around Write::flush, and maps a WouldBlock error intoNotReady to indicate that the flushing is still ongoing.

AsyncWrite’s poll_write and poll_flush follow the same contract asFuture::poll and AsyncRead::poll_read, namely that if they returnNotReady, they have arranged for the current task to be notified whenthey can make progress again. Like with poll_read, this means that wecan safely call these methods in our own futures, and know that we arealso following the contract.

Tokio uses the same mechanism to manage notifications for poll_writeand poll_flush as it does for poll_read, and you can read more aboutit in the non-blocking I/O section of Tokio internals.

Shutdown

AsyncWrite also adds one method that is not part of Write:shutdown. From its documentation:

Initiates or attempts to shut down this writer, returning success whenthe I/O connection has completely shut down.This method is intended to be used for asynchronous shutdown of I/Oconnections. For example this is suitable for implementing shutdown ofa TLS connection or calling TcpStream::shutdown on a proxiedconnection. Protocols sometimes need to flush out final pieces of dataor otherwise perform a graceful shutdown handshake, reading/writingmore data as appropriate. This method is the hook for such protocolsto implement the graceful shutdown logic.

This sums shutdown up pretty nicely: it is a way to tell the writerthat no more data is coming, and that it should indicate in whatever waythe underlying I/O protocol requires. For TCP connections, for example,this usually entails closing the writing side of the TCP channel so thatthe other end receives and end-of-file in the form of a read thatreturns 0 bytes. You can often think of shutdown as what you would_have done synchronously in the implementation of Drop; it’s just thatin the asynchronous world, you can’t easily do something in Dropbecause you need to have an executor that keeps polling your writer!

Note that calling shutdown on a write “half” of a type that implementsAsyncWrite _and AsyncRead does not shut down the read “half”. Youcan usually still continue reading data as you please until the otherside shuts down their corresponding write “half”.

An example of using AsyncWrite

Without further ado, let’s take a look at how we might implement

  1. # extern crate tokio;
  2. #[macro_use]
  3. extern crate futures;
  4. # fn main() {}
  5. use std::io;
  6. use tokio::prelude::*;
  8. // This is going to be our Future.
  9. // It'll seem awfully familiar to ReadExact above!
  10. // In the common case, this is set to Some(Writing),
  11. // but we'll set it to None when we return Async::Ready
  12. // so that we can return the writer and the buffer.
  13. struct WriteAll<W, T>(Option<Writing<W, T>>);
  15. struct Writing<W, T> {
  16.     // This is the stream we're writing into.
  17.     writer: W,
  18.     // This is the buffer we're writing from.
  19.     buffer: T,
  20.     // And this is much of the buffer we've written.
  21.     pos: usize,
  22. }
  24. // We want to be able to construct a WriteAll over anything
  25. // that implements AsyncWrite, and any buffer that can be
  26. // thought of as a &[u8].
  27. fn write_all<W, T>(writer: W, buffer: T) -> WriteAll<W, T>
  28. where
  29.     W: AsyncWrite,
  30.     T: AsRef<[u8]>,
  31. {
  32.     WriteAll(Some(Writing {
  33.         writer,
  34.         buffer,
  35.         // Initially, we've written none of the bytes from buffer.
  36.         pos: 0,
  37.     }))
  38. }
  40. impl<W, T> Future for WriteAll<W, T>
  41. where
  42.     W: AsyncWrite,
  43.     T: AsRef<[u8]>,
  44. {
  45.     // When we've written out the entire buffer, we want to return
  46.     // both the buffer and the writer so that the user can re-use them.
  47.     type Item = (W, T);
  48.     type Error = io::Error;
  50.     fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
  51.         match self.0 {
  52.             Some(Writing {
  53.                 ref mut writer,
  54.                 ref buffer,
  55.                 ref mut pos,
  56.             }) => {
  57.                 let buffer = buffer.as_ref();
  58.                 // Check that we haven't finished
  59.                 while *pos < buffer.len() {
  60.                     // Try to write the remainder of the buffer into the writer.
  61.                     // Just like write in std::io::Write, poll_write *can* write
  62.                     // fewer bytes than the length of the buffer it is given,
  63.                     // and we need to handle that by looking at its return
  64.                     // value, which is the number of bytes actually written.
  65.                     //
  66.                     // We are using try_ready! here, just like in poll_read in
  67.                     // ReadExact, so that if poll_write returns NotReady (or an
  68.                     // error), we will do the same! We uphold the contract that
  69.                     // we have arranged to be notified later because poll_write
  70.                     // follows that same contract, and _it_ returned NotReady.
  71.                     let n = try_ready!(writer.poll_write(&buffer[*pos..]));
  72.                     *pos += n;
  74.                     // If no bytes were written, but there was no error, this
  75.                     // generally implies that something weird happened under us.
  76.                     // We make sure to turn this into an error for the caller to
  77.                     // deal with.
  78.                     if n == 0 {
  79.                         return Err(io::Error::new(
  80.                             io::ErrorKind::WriteZero,
  81.                             "zero-length write",
  82.                         ));
  83.                     }
  84.                 }
  85.             }
  86.             None => panic!("poll a WriteAll after it's done"),
  87.         }
  89.         // We use the same trick as in ReadExact to ensure that we can return
  90.         // the buffer and the writer once the entire buffer has been written out.
  91.         let writing = self.0.take().expect("must have seen Some above");
  92.         Ok(Async::Ready((writing.writer, writing.buffer)))
  93.     }
  94. }

Next up: Implementing Async Read/Write