Resolve race condition around exceptions during streaming a response. (#491)

* Use 'socks5://' for the protocol in SOCKS proxy docs. Not 'socks://'

* Defensive programming around removing request status from the connection pool

* Tighten up behaviour around exceptions during streaming responses

* Call response_closed on async cancellations, KeyboardInterupts

* Include number of requests in-flight in exception if pool is closed early

* Update CHANGELOG.md

Author: lovelydinosaur

CHANGELOG.md

@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ## Master branch
 
 - Fix SOCKS support for `http://` URLs. (#492)
+- Resolve race condition around exceptions during streaming a response. (#491)
 
 ## 0.14.5 (January 18th, 2022)
 

httpcore/_async/connection_pool.py

@@ -239,7 +239,7 @@ async def handle_async_request(self, request: Request) -> Response:
                     # status so that the request becomes queued again.
                     status.unset_connection()
                     await self._attempt_to_acquire_connection(status)
-            except Exception as exc:
+            except BaseException as exc:
                 await self.response_closed(status)
                 raise exc
             else:
@@ -267,7 +267,8 @@ async def response_closed(self, status: RequestStatus) -> None:
 
         async with self._pool_lock:
             # Update the state of the connection pool.
-            self._requests.remove(status)
+            if status in self._requests:
+                self._requests.remove(status)
 
             if connection.is_closed() and connection in self._pool:
                 self._pool.remove(connection)
@@ -291,11 +292,19 @@ async def aclose(self) -> None:
         Close any connections in the pool.
         """
         async with self._pool_lock:
+            requests_still_in_flight = len(self._requests)
+
             for connection in self._pool:
                 await connection.aclose()
             self._pool = []
             self._requests = []
 
+            if requests_still_in_flight:
+                raise RuntimeError(
+                    f"The connection pool was closed while {requests_still_in_flight} "
+                    f"HTTP requests/responses were still in-flight."
+                )
+
     async def __aenter__(self) -> "AsyncConnectionPool":
         return self
 

httpcore/_async/http11.py

@@ -279,13 +279,23 @@ class HTTP11ConnectionByteStream:
     def __init__(self, connection: AsyncHTTP11Connection, request: Request) -> None:
         self._connection = connection
         self._request = request
+        self._closed = False
 
     async def __aiter__(self) -> AsyncIterator[bytes]:
         kwargs = {"request": self._request}
-        async with Trace("http11.receive_response_body", self._request, kwargs):
-            async for chunk in self._connection._receive_response_body(**kwargs):
-                yield chunk
+        try:
+            async with Trace("http11.receive_response_body", self._request, kwargs):
+                async for chunk in self._connection._receive_response_body(**kwargs):
+                    yield chunk
+        except BaseException as exc:
+            # If we get an exception while streaming the response,
+            # we want to close the response (and possibly the connection)
+            # before raising that exception.
+            await self.aclose()
+            raise exc
 
     async def aclose(self) -> None:
-        async with Trace("http11.response_closed", self._request):
-            await self._connection._response_closed()
+        if not self._closed:
+            self._closed = True
+            async with Trace("http11.response_closed", self._request):
+                await self._connection._response_closed()

httpcore/_async/http2.py

@@ -305,7 +305,6 @@ async def _response_closed(self, stream_id: int) -> None:
     async def aclose(self) -> None:
         # Note that this method unilaterally closes the connection, and does
         # not have any kind of locking in place around it.
-        # For task-safe/thread-safe operations call into 'attempt_close' instead.
         self._h2_state.close_connection()
         self._state = HTTPConnectionState.CLOSED
         await self._network_stream.aclose()
@@ -446,16 +445,26 @@ def __init__(
         self._connection = connection
         self._request = request
         self._stream_id = stream_id
+        self._closed = False
 
     async def __aiter__(self) -> typing.AsyncIterator[bytes]:
         kwargs = {"request": self._request, "stream_id": self._stream_id}
-        async with Trace("http2.receive_response_body", self._request, kwargs):
-            async for chunk in self._connection._receive_response_body(
-                request=self._request, stream_id=self._stream_id
-            ):
-                yield chunk
+        try:
+            async with Trace("http2.receive_response_body", self._request, kwargs):
+                async for chunk in self._connection._receive_response_body(
+                    request=self._request, stream_id=self._stream_id
+                ):
+                    yield chunk
+        except BaseException as exc:
+            # If we get an exception while streaming the response,
+            # we want to close the response (and possibly the connection)
+            # before raising that exception.
+            await self.aclose()
+            raise exc
 
     async def aclose(self) -> None:
-        kwargs = {"stream_id": self._stream_id}
-        async with Trace("http2.response_closed", self._request, kwargs):
-            await self._connection._response_closed(stream_id=self._stream_id)
+        if not self._closed:
+            self._closed = True
+            kwargs = {"stream_id": self._stream_id}
+            async with Trace("http2.response_closed", self._request, kwargs):
+                await self._connection._response_closed(stream_id=self._stream_id)

httpcore/_sync/connection_pool.py

@@ -239,7 +239,7 @@ def handle_request(self, request: Request) -> Response:
                     # status so that the request becomes queued again.
                     status.unset_connection()
                     self._attempt_to_acquire_connection(status)
-            except Exception as exc:
+            except BaseException as exc:
                 self.response_closed(status)
                 raise exc
             else:
@@ -267,7 +267,8 @@ def response_closed(self, status: RequestStatus) -> None:
 
         with self._pool_lock:
             # Update the state of the connection pool.
-            self._requests.remove(status)
+            if status in self._requests:
+                self._requests.remove(status)
 
             if connection.is_closed() and connection in self._pool:
                 self._pool.remove(connection)
@@ -291,11 +292,19 @@ def close(self) -> None:
         Close any connections in the pool.
         """
         with self._pool_lock:
+            requests_still_in_flight = len(self._requests)
+
             for connection in self._pool:
                 connection.close()
             self._pool = []
             self._requests = []
 
+            if requests_still_in_flight:
+                raise RuntimeError(
+                    f"The connection pool was closed while {requests_still_in_flight} "
+                    f"HTTP requests/responses were still in-flight."
+                )
+
     def __enter__(self) -> "ConnectionPool":
         return self
 

httpcore/_sync/http11.py

@@ -279,13 +279,23 @@ class HTTP11ConnectionByteStream:
     def __init__(self, connection: HTTP11Connection, request: Request) -> None:
         self._connection = connection
         self._request = request
+        self._closed = False
 
     def __iter__(self) -> Iterator[bytes]:
         kwargs = {"request": self._request}
-        with Trace("http11.receive_response_body", self._request, kwargs):
-            for chunk in self._connection._receive_response_body(**kwargs):
-                yield chunk
+        try:
+            with Trace("http11.receive_response_body", self._request, kwargs):
+                for chunk in self._connection._receive_response_body(**kwargs):
+                    yield chunk
+        except BaseException as exc:
+            # If we get an exception while streaming the response,
+            # we want to close the response (and possibly the connection)
+            # before raising that exception.
+            self.close()
+            raise exc
 
     def close(self) -> None:
-        with Trace("http11.response_closed", self._request):
-            self._connection._response_closed()
+        if not self._closed:
+            self._closed = True
+            with Trace("http11.response_closed", self._request):
+                self._connection._response_closed()

httpcore/_sync/http2.py

@@ -305,7 +305,6 @@ def _response_closed(self, stream_id: int) -> None:
     def close(self) -> None:
         # Note that this method unilaterally closes the connection, and does
         # not have any kind of locking in place around it.
-        # For task-safe/thread-safe operations call into 'attempt_close' instead.
         self._h2_state.close_connection()
         self._state = HTTPConnectionState.CLOSED
         self._network_stream.close()
@@ -446,16 +445,26 @@ def __init__(
         self._connection = connection
         self._request = request
         self._stream_id = stream_id
+        self._closed = False
 
     def __iter__(self) -> typing.Iterator[bytes]:
         kwargs = {"request": self._request, "stream_id": self._stream_id}
-        with Trace("http2.receive_response_body", self._request, kwargs):
-            for chunk in self._connection._receive_response_body(
-                request=self._request, stream_id=self._stream_id
-            ):
-                yield chunk
+        try:
+            with Trace("http2.receive_response_body", self._request, kwargs):
+                for chunk in self._connection._receive_response_body(
+                    request=self._request, stream_id=self._stream_id
+                ):
+                    yield chunk
+        except BaseException as exc:
+            # If we get an exception while streaming the response,
+            # we want to close the response (and possibly the connection)
+            # before raising that exception.
+            self.close()
+            raise exc
 
     def close(self) -> None:
-        kwargs = {"stream_id": self._stream_id}
-        with Trace("http2.response_closed", self._request, kwargs):
-            self._connection._response_closed(stream_id=self._stream_id)
+        if not self._closed:
+            self._closed = True
+            kwargs = {"stream_id": self._stream_id}
+            with Trace("http2.response_closed", self._request, kwargs):
+                self._connection._response_closed(stream_id=self._stream_id)

tests/_async/test_connection_pool.py

@@ -435,3 +435,29 @@ async def test_unsupported_protocol():
 
         with pytest.raises(UnsupportedProtocol):
             await pool.request("GET", "://www.example.com/")
+
+
+@pytest.mark.anyio
+async def test_connection_pool_closed_while_request_in_flight():
+    """
+    Closing a connection pool while a request/response is still in-flight
+    should raise an error.
+    """
+    network_backend = AsyncMockBackend(
+        [
+            b"HTTP/1.1 200 OK\r\n",
+            b"Content-Type: plain/text\r\n",
+            b"Content-Length: 13\r\n",
+            b"\r\n",
+            b"Hello, world!",
+        ]
+    )
+
+    async with AsyncConnectionPool(
+        network_backend=network_backend,
+    ) as pool:
+        # Send a request, and then close the connection pool while the
+        # response has not yet been streamed.
+        async with pool.stream("GET", "https://example.com/"):
+            with pytest.raises(RuntimeError):
+                await pool.aclose()

tests/_async/test_http11.py

@@ -91,6 +91,35 @@ async def test_http11_connection_with_remote_protocol_error():
         )
 
 
+@pytest.mark.anyio
+async def test_http11_connection_with_incomplete_response():
+    """
+    We should be gracefully handling the case where the connection ends prematurely.
+    """
+    origin = Origin(b"https", b"example.com", 443)
+    stream = AsyncMockStream(
+        [
+            b"HTTP/1.1 200 OK\r\n",
+            b"Content-Type: plain/text\r\n",
+            b"Content-Length: 13\r\n",
+            b"\r\n",
+            b"Hello, wor",
+        ]
+    )
+    async with AsyncHTTP11Connection(origin=origin, stream=stream) as conn:
+        with pytest.raises(RemoteProtocolError):
+            await conn.request("GET", "https://example.com/")
+
+        assert not conn.is_idle()
+        assert conn.is_closed()
+        assert not conn.is_available()
+        assert not conn.has_expired()
+        assert (
+            repr(conn)
+            == "<AsyncHTTP11Connection ['https://example.com:443', CLOSED, Request Count: 1]>"
+        )
+
+
 @pytest.mark.anyio
 async def test_http11_connection_with_local_protocol_error():
     """

tests/_sync/test_connection_pool.py

@@ -435,3 +435,29 @@ def test_unsupported_protocol():
 
         with pytest.raises(UnsupportedProtocol):
             pool.request("GET", "://www.example.com/")
+
+
+
+def test_connection_pool_closed_while_request_in_flight():
+    """
+    Closing a connection pool while a request/response is still in-flight
+    should raise an error.
+    """
+    network_backend = MockBackend(
+        [
+            b"HTTP/1.1 200 OK\r\n",
+            b"Content-Type: plain/text\r\n",
+            b"Content-Length: 13\r\n",
+            b"\r\n",
+            b"Hello, world!",
+        ]
+    )
+
+    with ConnectionPool(
+        network_backend=network_backend,
+    ) as pool:
+        # Send a request, and then close the connection pool while the
+        # response has not yet been streamed.
+        with pool.stream("GET", "https://example.com/"):
+            with pytest.raises(RuntimeError):
+                pool.close()