From b4c15eef2ce9302b948f44e01831549fde8a79d5 Mon Sep 17 00:00:00 2001
From: Jacob Hoffman-Andrews <github@hoffman-andrews.com>
Date: Fri, 12 Jun 2020 20:55:19 -0700
Subject: [PATCH] Check for server closed connections.

This builds on 753d61b. Before we send a request, we can do a 1-byte
nonblocking peek on the connection. If the server has closed the
connection, this will give us an EOF, and we can take the connection out
of the pool before sending any request on it. This will reduce the
likelihood that we send a non-retryable POST on an already-closed
connection.

The server could still potentially close the connection between when we
make this check and when we finish sending the request, but this should
handle the majority of cases.
---
 src/stream.rs | 27 +++++++++++++++++++++++++++
 src/unit.rs   | 10 ++++++++--
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/src/stream.rs b/src/stream.rs
index 7057a93..a3f91c5 100644
--- a/src/stream.rs
+++ b/src/stream.rs
@@ -46,6 +46,33 @@ impl ::std::fmt::Debug for Stream {
 }
 
 impl Stream {
+    // Check if the server has closed a stream by performing a one-byte
+    // non-blocking read. If this returns EOF, the server has closed the
+    // connection: return true. If this returns WouldBlock (aka EAGAIN),
+    // that means the connection is still open: return false. Otherwise
+    // return an error.
+    fn serverclosed_stream(stream: &std::net::TcpStream) -> IoResult<bool> {
+        let mut buf = [0; 1];
+        stream.set_nonblocking(true)?;
+
+        let result = match stream.peek(&mut buf) {
+            Ok(0) => Ok(true),
+            Ok(_) => Ok(false), // TODO: Maybe this should produce an "unexpected response" error
+            Err(e) if e.kind() == ErrorKind::WouldBlock => Ok(false),
+            Err(e) => Err(e),
+        };
+        stream.set_nonblocking(false)?;
+
+        result
+    }
+    // Return true if the server has closed this connection.
+    pub(crate) fn server_closed(&self) -> IoResult<bool> {
+        match self {
+            Stream::Http(tcpstream) => Stream::serverclosed_stream(tcpstream),
+            Stream::Https(rustls_stream) => Stream::serverclosed_stream(&rustls_stream.sock),
+            _ => Ok(false),
+        }
+    }
     pub fn is_poolable(&self) -> bool {
         match self {
             Stream::Http(_) => true,
diff --git a/src/unit.rs b/src/unit.rs
index b461caa..a3b4205 100644
--- a/src/unit.rs
+++ b/src/unit.rs
@@ -274,8 +274,14 @@ fn connect_socket(unit: &Unit, use_pooled: bool) -> Result<(Stream, bool), Error
     if use_pooled {
         let state = &mut unit.agent.lock().unwrap();
         if let Some(agent) = state.as_mut() {
-            if let Some(stream) = agent.pool.try_get_connection(&unit.url) {
-                return Ok((stream, true));
+            // The connection may have been closed by the server
+            // due to idle timeout while it was sitting in the pool.
+            // Loop until we find one that is still good or run out of connections.
+            while let Some(stream) = agent.pool.try_get_connection(&unit.url) {
+                let server_closed = stream.server_closed()?;
+                if !server_closed {
+                    return Ok((stream, true));
+                }
             }
         }
     }