Drain connection after PoolManager redirect (#1817)

5 files changed, 44 insertions, 28 deletions

diff --git a/CONTRIBUTORS.txt b/CONTRIBUTORS.txt
index b632da40..ef41e32b 100644
--- a/CONTRIBUTORS.txt
+++ b/CONTRIBUTORS.txt
@@ -297,5 +297,8 @@ In chronological order:
 * Hod Bin Noon <bin.noon.hod@gmail.com>
   * Test improvements
 
+* Chris Olufson <tycarac@gmail.com>
+  * Fix for connection not being released on HTTP redirect and response not preloaded
+
 * [Your name or handle] <[email or website]>
   * [Brief summary of your changes]
diff --git a/src/urllib3/connectionpool.py b/src/urllib3/connectionpool.py
index 7d267b68..5f044dbd 100644
--- a/src/urllib3/connectionpool.py
+++ b/src/urllib3/connectionpool.py
@@ -765,21 +765,6 @@ class HTTPConnectionPool(ConnectionPool, RequestMethods):
                 **response_kw
             )
 
-        def drain_and_release_conn(response):
-            try:
-                # discard any remaining response body, the connection will be
-                # released back to the pool once the entire response is read
-                response.read()
-            except (
-                TimeoutError,
-                HTTPException,
-                SocketError,
-                ProtocolError,
-                BaseSSLError,
-                SSLError,
-            ):
-                pass
-
         # Handle redirect?
         redirect_location = redirect and response.get_redirect_location()
         if redirect_location:
@@ -790,15 +775,11 @@ class HTTPConnectionPool(ConnectionPool, RequestMethods):
                 retries = retries.increment(method, url, response=response, _pool=self)
             except MaxRetryError:
                 if retries.raise_on_redirect:
-                    # Drain and release the connection for this response, since
-                    # we're not returning it to be released manually.
-                    drain_and_release_conn(response)
+                    response.drain_conn()
                     raise
                 return response
 
-            # drain and return the connection to the pool before recursing
-            drain_and_release_conn(response)
-
+            response.drain_conn()
             retries.sleep_for_retry(response)
             log.debug("Redirecting %s -> %s", url, redirect_location)
             return self.urlopen(
@@ -824,15 +805,11 @@ class HTTPConnectionPool(ConnectionPool, RequestMethods):
                 retries = retries.increment(method, url, response=response, _pool=self)
             except MaxRetryError:
                 if retries.raise_on_status:
-                    # Drain and release the connection for this response, since
-                    # we're not returning it to be released manually.
-                    drain_and_release_conn(response)
+                    response.drain_conn()
                     raise
                 return response
 
-            # drain and return the connection to the pool before recursing
-            drain_and_release_conn(response)
-
+            response.drain_conn()
             retries.sleep(response)
             log.debug("Retry: %s", url)
             return self.urlopen(
diff --git a/src/urllib3/poolmanager.py b/src/urllib3/poolmanager.py
index 242a2f82..c77784ce 100644
--- a/src/urllib3/poolmanager.py
+++ b/src/urllib3/poolmanager.py
@@ -6,7 +6,11 @@ import logging
 from ._collections import RecentlyUsedContainer
 from .connectionpool import HTTPConnectionPool, HTTPSConnectionPool
 from .connectionpool import port_by_scheme
-from .exceptions import LocationValueError, MaxRetryError, ProxySchemeUnknown
+from .exceptions import (
+    LocationValueError,
+    MaxRetryError,
+    ProxySchemeUnknown,
+)
 from .packages import six
 from .packages.six.moves.urllib.parse import urljoin
 from .request import RequestMethods
@@ -359,6 +363,7 @@ class PoolManager(RequestMethods):
             retries = retries.increment(method, url, response=response, _pool=conn)
         except MaxRetryError:
             if retries.raise_on_redirect:
+                response.drain_conn()
                 raise
             return response
 
@@ -366,6 +371,8 @@ class PoolManager(RequestMethods):
         kw["redirect"] = redirect
 
         log.info("Redirecting %s -> %s", url, redirect_location)
+
+        response.drain_conn()
         return self.urlopen(method, redirect_location, **kw)
 
 
diff --git a/src/urllib3/response.py b/src/urllib3/response.py
index 6090a735..7dc9b93c 100644
--- a/src/urllib3/response.py
+++ b/src/urllib3/response.py
@@ -20,6 +20,7 @@ from .exceptions import (
     ResponseNotChunked,
     IncompleteRead,
     InvalidHeader,
+    HTTPError,
 )
 from .packages.six import string_types as basestring, PY3
 from .packages.six.moves import http_client as httplib
@@ -277,6 +278,17 @@ class HTTPResponse(io.IOBase):
         self._pool._put_conn(self._connection)
         self._connection = None
 
+    def drain_conn(self):
+        """
+        Read and discard any remaining HTTP response data in the response connection.
+
+        Unread data in the HTTPResponse connection blocks the connection from being released back to the pool.
+        """
+        try:
+            self.read()
+        except (HTTPError, SocketError, BaseSSLError, HTTPException):
+            pass
+
     @property
     def data(self):
         # For backwords-compat with earlier urllib3 0.4 and earlier.
diff --git a/test/with_dummyserver/test_poolmanager.py b/test/with_dummyserver/test_poolmanager.py
index 4a47f627..1160a2b6 100644
--- a/test/with_dummyserver/test_poolmanager.py
+++ b/test/with_dummyserver/test_poolmanager.py
@@ -115,6 +115,7 @@ class TestPoolManager(HTTPDummyServerTestCase):
                         % (self.base_url, self.base_url)
                     },
                     retries=1,
+                    preload_content=False,
                 )
 
             with pytest.raises(MaxRetryError):
@@ -126,8 +127,15 @@ class TestPoolManager(HTTPDummyServerTestCase):
                         % (self.base_url, self.base_url)
                     },
                     retries=Retry(total=None, redirect=1),
+                    preload_content=False,
                 )
 
+            # Even with preload_content=False and raise on redirects, we reused the same
+            # connection
+            assert len(http.pools) == 1
+            pool = http.connection_from_host(self.host, self.port)
+            assert pool.num_connections == 1
+
     def test_redirect_cross_host_remove_headers(self):
         with PoolManager() as http:
             r = http.request(
@@ -206,6 +214,15 @@ class TestPoolManager(HTTPDummyServerTestCase):
             assert "X-API-Secret" not in data
             assert data["Authorization"] == "bar"
 
+    def test_redirect_without_preload_releases_connection(self):
+        with PoolManager(block=True, maxsize=2) as http:
+            r = http.request(
+                "GET", "%s/redirect" % self.base_url, preload_content=False
+            )
+            assert r._pool.num_requests == 2
+            assert r._pool.num_connections == 1
+            assert len(http.pools) == 1
+
     def test_raise_on_redirect(self):
         with PoolManager() as http:
             r = http.request(