Ensure that remote streams are shutdown appropriately (#29)

Previously, when one end of the proxy connection stopped writing (so that
read() returns 0), the shutdown was not propagated to the other connection.
This is because the AsyncWrite::shutdown() api does nothing to actually
shutdown the connection.

We introduce Socket::tcp_shutdown to provide a consistent api for triggering a
TcpStream::shutdown.  ProxyStream uses both the AsyncWrite::shutdown and
Socket::tcp_shutdown APIs to signal when a write is complete.

This has been confirmed to fix the issue with manual testing.

Author: olix0r

src/lb/duplex.rs

@@ -34,7 +34,7 @@ impl Duplex {
         let dst = Rc::new(RefCell::new(dst));
         let tx_bytes_stat = tx_metrics.scope().stat("bytes".into());
         let rx_byte_stat = rx_metrics.scope().stat("bytes".into());
-        let metrics = rx_metrics.clone(); // doesn't matter which one,
+        let metrics = rx_metrics.clone(); // doesn't matter which one.
         let tx = ProxyStream::new(src.clone(), dst.clone(), buf.clone(), tx_metrics);
         let rx = ProxyStream::new(dst, src, buf, rx_metrics);
         Duplex {
@@ -70,6 +70,7 @@ impl Future for Duplex {
                     self.tx_bytes += sz;
                 }
                 Async::NotReady => {
+                    trace!("dstward not ready");
                     self.tx = Some(tx);
                 }
             }
@@ -87,17 +88,20 @@ impl Future for Duplex {
                     self.rx_bytes += sz;
                 }
                 Async::NotReady => {
+                    trace!("srcward not ready");
                     self.rx = Some(rx);
                 }
             }
         }
 
         if self.tx.is_none() && self.rx.is_none() {
+            trace!("complete");
             let mut rec = self.metrics.recorder();
             rec.add(&self.tx_bytes_stat, self.tx_bytes);
             rec.add(&self.rx_bytes_stat, self.rx_bytes);
             Ok(Async::Ready(()))
         } else {
+            trace!("not ready");
             Ok(Async::NotReady)
         }
     }

src/lb/proxy_stream.rs

@@ -5,6 +5,7 @@ use futures::{Async, Future, Poll};
 use lb::Socket;
 use std::cell::RefCell;
 use std::io::{self, Read, Write};
+use std::net::Shutdown;
 use std::rc::Rc;
 use tacho;
 use tokio_io::AsyncWrite;
@@ -69,14 +70,17 @@ impl Future for ProxyStream {
     /// all pending data before reading any more.
     fn poll(&mut self) -> Poll<u64, io::Error> {
         trace!("poll");
-        if self.completed {
-            return Ok(self.bytes_total.into());
-        }
-
         let mut rec = self.metrics.recorder();
         let mut writer = self.writer.borrow_mut();
         let mut reader = self.reader.borrow_mut();
         loop {
+            if self.completed {
+                try_nb!(writer.shutdown());
+                writer.tcp_shutdown(Shutdown::Write)?;
+                trace!("completed");
+                return Ok(self.bytes_total.into());
+            }
+
             // Try to flush pending bytes to the writer.
             if let Some(mut pending) = self.pending.take() {
                 let psz = pending.len();
@@ -106,16 +110,12 @@ impl Future for ProxyStream {
             let rsz = try_nb!(reader.read(&mut buf));
             if rsz == 0 {
                 // Nothing left to read, return the total number of bytes transferred.
-                trace!("completed: {}B", self.bytes_total);
+                trace!("completing: {}B", self.bytes_total);
                 self.completed = true;
-                match writer.shutdown()? {
-                    Async::NotReady => {
-                        return Ok(Async::NotReady);
-                    }
-                    Async::Ready(_) => {
-                        return Ok(self.bytes_total.into());
-                    }
-                }
+                try_nb!(writer.shutdown());
+                writer.tcp_shutdown(Shutdown::Write)?;
+                trace!("completed: {}B", self.bytes_total);
+                return Ok(self.bytes_total.into());
             }
             trace!("read {} bytes", rsz);