duplex: Ensure written data is flushed (#944)

In testing Postgres dumps through Linkerd, we see some situations where
data is written to rustls but does not become visible to the client.
Rustls documents that written data may need to be flushed.

This change updates the `HalfDuplex::copy_into` function to flush data
when all available data has been read. This change also increases the
copy buffer from 4KB to 64KB to reduce processing overhead.

Author: olix0r

linkerd/duplex/src/lib.rs

@@ -2,12 +2,11 @@
 
 use bytes::{Buf, BufMut};
 use futures::ready;
-use io::{AsyncRead, AsyncWrite};
-use linkerd_io as io;
+use linkerd_io::{self as io, AsyncRead, AsyncWrite};
 use pin_project::pin_project;
 use std::task::{Context, Poll};
 use std::{future::Future, pin::Pin};
-use tracing::trace;
+use tracing::{error, trace};
 
 /// A future piping data bi-directionally to In and Out.
 #[pin_project]
@@ -24,6 +23,7 @@ struct HalfDuplex<T> {
     #[pin]
     io: T,
     direction: &'static str,
+    flushing: bool,
 }
 
 /// A buffer used to copy bytes from one IO to another.
@@ -39,6 +39,18 @@ struct CopyBuf {
     write_pos: usize,
 }
 
+enum Buffered {
+    NotEmpty,
+    Read(usize),
+    Eof,
+}
+
+enum Drained {
+    BufferEmpty,
+    Partial(usize),
+    All(usize),
+}
+
 impl<In, Out> Duplex<In, Out>
 where
     In: AsyncRead + AsyncWrite + Unpin,
@@ -57,9 +69,9 @@ where
     In: AsyncRead + AsyncWrite + Unpin,
     Out: AsyncRead + AsyncWrite + Unpin,
 {
-    type Output = Result<(), io::Error>;
+    type Output = io::Result<()>;
 
-    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> io::Poll<()> {
         let mut this = self.project();
         // This purposefully ignores the Async part, since we don't want to
         // return early if the first half isn't ready, but the other half
@@ -85,17 +97,19 @@ where
             is_shutdown: false,
             io,
             direction,
+            flushing: false,
         }
     }
 
-    fn copy_into<U>(
+    /// Reads data from `self`, buffering it, and writing it to `dst.
+    ///
+    /// Returns ready when the stream has shutdown such that no more data may be
+    /// proxied.
+    fn copy_into<U: AsyncWrite + Unpin>(
         &mut self,
         dst: &mut HalfDuplex<U>,
         cx: &mut Context<'_>,
-    ) -> Poll<Result<(), io::Error>>
-    where
-        U: AsyncWrite + Unpin,
-    {
+    ) -> io::Poll<()> {
         // Since Duplex::poll() intentionally ignores the Async part of our
         // return value, we may be polled again after returning Ready, if the
         // other half isn't ready. In that case, if the destination has
@@ -105,61 +119,156 @@ where
             trace!(direction = %self.direction, "already shutdown");
             return Poll::Ready(Ok(()));
         }
+
+        // If the last invocation returned pending while flushing, resume
+        // flushing and only proceed when the flush is complete.
+        if self.flushing {
+            ready!(self.poll_flush(dst, cx))?;
+        }
+
+        // `needs_flush` is set to true if the buffer is written so that, if a
+        // read returns pending, that data may be flushed.
+        let mut needs_flush = false;
+
         loop {
-            ready!(self.poll_read(cx))?;
-            ready!(self.poll_write_into(dst, cx))?;
-            if self.buf.is_none() {
-                trace!(direction = %self.direction, "shutting down");
-                debug_assert!(!dst.is_shutdown, "attempted to shut down destination twice");
-                ready!(Pin::new(&mut dst.io).poll_shutdown(cx))?;
-                dst.is_shutdown = true;
-
-                return Poll::Ready(Ok(()));
+            // As long as the underlying socket is alive, ensure we've read data
+            // from it into the local buffer.
+            match self.poll_buffer(cx)? {
+                Poll::Pending => {
+                    // If there's no data to be read and we've written data, try
+                    // flushing before returning pending.
+                    if needs_flush {
+                        // The poll status of the flush isn't relevant, as we
+                        // have registered interest in the read (and maybe the
+                        // write as well). If the flush did not complete
+                        // `self.flushing` is true so that it may be resumed on
+                        // the next poll.
+                        let _ = self.poll_flush(dst, cx)?;
+                    }
+                    return Poll::Pending;
+                }
+
+                Poll::Ready(Buffered::NotEmpty) | Poll::Ready(Buffered::Read(_)) => {
+                    // Write buffered data to the destination.
+                    match self.drain_into(dst, cx)? {
+                        // All of the buffered data was written, so continue reading more.
+                        Drained::All(sz) => {
+                            debug_assert!(sz > 0);
+                            needs_flush = true;
+                        }
+                        // Only some of the buffered data could be written
+                        // before the destination became pending. Try to flush
+                        // the written data to get capacity.
+                        Drained::Partial(_) => {
+                            ready!(self.poll_flush(dst, cx))?;
+                            // If the flush completed, try writing again to
+                            // ensure that we have a notification registered. If
+                            // all of the buffered data still cannot be written,
+                            // return pending. Otherwise, continue.
+                            if let Drained::Partial(_) = self.drain_into(dst, cx)? {
+                                return Poll::Pending;
+                            }
+                            needs_flush = false;
+                        }
+                        Drained::BufferEmpty => {
+                            error!(
+                                direction = self.direction,
+                                "Invalid state: attempted to write from an empty buffer"
+                            );
+                            debug_assert!(false, "The write buffer should never be empty");
+                            return Poll::Ready(Ok(()));
+                        }
+                    }
+                }
+
+                // The socket closed, so initiate shutdown on the destination.
+                Poll::Ready(Buffered::Eof) => {
+                    trace!(direction = %self.direction, "shutting down");
+                    debug_assert!(!dst.is_shutdown, "attempted to shut down destination twice");
+                    ready!(Pin::new(&mut dst.io).poll_shutdown(cx))?;
+                    dst.is_shutdown = true;
+                    return Poll::Ready(Ok(()));
+                }
             }
         }
     }
 
-    fn poll_read(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
-        let mut is_eof = false;
-        if let Some(ref mut buf) = self.buf {
-            if !buf.has_remaining() {
-                buf.reset();
+    /// Attempts to read and buffer data from the underlying stream, returning
+    /// the number of bytes read. If the buffer already has data, no new data
+    /// will be read.
+    fn poll_buffer(&mut self, cx: &mut Context<'_>) -> io::Poll<Buffered> {
+        // Buffer data only if no data is buffered.
+        //
+        // TODO should we read more data as long as there's buffer capacity?
+        // To do this, we'd have to get more complex about handling EOF.
+        if let Some(buf) = self.buf.as_mut() {
+            if buf.has_remaining() {
+                // Data was already buffered, so just return immediately.
+                trace!(direction = %self.direction, remaining = buf.remaining(), "skipping read");
+                return Poll::Ready(Ok(Buffered::NotEmpty));
+            }
 
-                trace!(direction = %self.direction, "reading");
-                let n = ready!(io::poll_read_buf(Pin::new(&mut self.io), cx, buf))?;
-                trace!(direction = %self.direction, "read {}B", n);
+            buf.reset();
+            trace!(direction = %self.direction, "reading");
+            let sz = ready!(io::poll_read_buf(Pin::new(&mut self.io), cx, buf))?;
+            trace!(direction = %self.direction, "read {}B", sz);
 
-                is_eof = n == 0;
+            // If data was read, return the number of bytes read.
+            if sz > 0 {
+                return Poll::Ready(Ok(Buffered::Read(sz)));
             }
         }
-        if is_eof {
-            trace!("eof");
-            self.buf = None;
-        }
 
-        Poll::Ready(Ok(()))
+        // No more data can be read.
+        trace!("eof");
+        self.buf = None;
+        Poll::Ready(Ok(Buffered::Eof))
+    }
+
+    /// Attempts to flush the destination. `self.flushing` is set to true iff the
+    /// flush operation did not complete.
+    fn poll_flush<U: AsyncWrite + Unpin>(
+        &mut self,
+        dst: &mut HalfDuplex<U>,
+        cx: &mut Context<'_>,
+    ) -> io::Poll<()> {
+        trace!(direction = %self.direction, "flushing");
+        let poll = Pin::new(&mut dst.io).poll_flush(cx);
+        self.flushing = poll.is_pending();
+        if poll.is_ready() {
+            trace!(direction = %self.direction, "flushed");
+        }
+        poll
     }
 
-    fn poll_write_into<U>(
+    /// Writes as much buffered data as possible, returning the number of bytes written.
+    fn drain_into<U: AsyncWrite + Unpin>(
         &mut self,
         dst: &mut HalfDuplex<U>,
         cx: &mut Context<'_>,
-    ) -> Poll<Result<(), io::Error>>
-    where
-        U: AsyncWrite + Unpin,
-    {
-        if let Some(ref mut buf) = self.buf {
+    ) -> io::Result<Drained> {
+        let mut sz = 0;
+
+        if let Some(buf) = self.buf.as_mut() {
             while buf.has_remaining() {
                 trace!(direction = %self.direction, "writing {}B", buf.remaining());
-                let n = ready!(io::poll_write_buf(Pin::new(&mut dst.io), cx, buf))?;
+                let n = match io::poll_write_buf(Pin::new(&mut dst.io), cx, buf)? {
+                    Poll::Pending => return Ok(Drained::Partial(sz)),
+                    Poll::Ready(n) => n,
+                };
                 trace!(direction = %self.direction, "wrote {}B", n);
                 if n == 0 {
-                    return Poll::Ready(Err(write_zero()));
+                    return Err(write_zero());
                 }
+                sz += n;
             }
         }
 
-        Poll::Ready(Ok(()))
+        if sz == 0 {
+            Ok(Drained::BufferEmpty)
+        } else {
+            Ok(Drained::All(sz))
+        }
     }
 
     fn is_done(&self) -> bool {
@@ -174,7 +283,7 @@ fn write_zero() -> io::Error {
 impl CopyBuf {
     fn new() -> Self {
         CopyBuf {
-            buf: Box::new([0; 4096]),
+            buf: Box::new([0; 64 * 1024]),
             read_pos: 0,
             write_pos: 0,
         }
@@ -208,11 +317,9 @@ unsafe impl BufMut for CopyBuf {
     }
 
     fn chunk_mut(&mut self) -> &mut bytes::buf::UninitSlice {
+        // Safety: The memory is initialized. This is the only way to turn a
+        // `&[T]` into a `&[MaybeUninit<T>]` without ptr casting.
         unsafe {
-            // this is, in fact, _totally fine and safe_: all the memory is
-            // initialized.
-            // there's just no way to turn a `&[T]` into a `&[MaybeUninit<T>]`
-            // without ptr casting.
             bytes::buf::UninitSlice::from_raw_parts_mut(
                 &mut self.buf[self.write_pos] as *mut _,
                 self.buf.len() - self.write_pos,
@@ -225,61 +332,3 @@ unsafe impl BufMut for CopyBuf {
         self.write_pos += cnt;
     }
 }
-
-// #[cfg(test)]
-// mod tests {
-//     use std::io::{Error, Read, Result, Write};
-//     use std::sync::atomic::{AtomicBool, Ordering};
-
-//     use super::*;
-//     use tokio::io::{AsyncRead, AsyncWrite};
-
-//     #[derive(Debug)]
-//     struct DoneIo(AtomicBool);
-
-//     impl<'a> Read for &'a DoneIo {
-//         fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
-//             if self.0.swap(false, Ordering::Relaxed) {
-//                 Ok(buf.len())
-//             } else {
-//                 Ok(0)
-//             }
-//         }
-//     }
-
-//     impl<'a> AsyncRead for &'a DoneIo {
-//         unsafe fn prepare_uninitialized_buffer(&self, _buf: &mut [u8]) -> bool {
-//             true
-//         }
-//     }
-
-//     impl<'a> Write for &'a DoneIo {
-//         fn write(&mut self, buf: &[u8]) -> Result<usize> {
-//             Ok(buf.len())
-//         }
-//         fn flush(&mut self) -> Result<()> {
-//             Ok(())
-//         }
-//     }
-//     impl<'a> AsyncWrite for &'a DoneIo {
-//         fn shutdown(&mut self) -> Poll<(), Error> {
-//             if self.0.swap(false, Ordering::Relaxed) {
-//                 Ok(Async::NotReady)
-//             } else {
-//                 Ok(Async::Ready(()))
-//             }
-//         }
-//     }
-
-//     #[test]
-//     fn duplex_doesnt_hang_when_one_half_finishes() {
-//         // Test reproducing an infinite loop in Duplex that caused issue #519,
-//         // where a Duplex would enter an infinite loop when one half finishes.
-//         let io_1 = DoneIo(AtomicBool::new(true));
-//         let io_2 = DoneIo(AtomicBool::new(true));
-//         let mut duplex = Duplex::new(&io_1, &io_2);
-
-//         assert_eq!(duplex.poll().unwrap(), Async::NotReady);
-//         assert_eq!(duplex.poll().unwrap(), Async::Ready(()));
-//     }
-// }