Close all outstanding responses on client close

httpx/dispatch/connection_pool.py

@@ -1,4 +1,5 @@
 import typing
+from itertools import chain
 
 from ..concurrency.asyncio import AsyncioBackend
 from ..concurrency.base import ConcurrencyBackend
@@ -62,10 +63,11 @@ def add(self, connection: HTTPConnection) -> None:
             self.by_origin[connection.origin] = {connection: 0.0}
 
     def remove(self, connection: HTTPConnection) -> None:
-        del self.all[connection]
-        del self.by_origin[connection.origin][connection]
-        if not self.by_origin[connection.origin]:
-            del self.by_origin[connection.origin]
+        if connection in self.all:
+            del self.all[connection]
+            del self.by_origin[connection.origin][connection]
+            if not self.by_origin[connection.origin]:
+                del self.by_origin[connection.origin]
 
     def clear(self) -> None:
         self.all.clear()
@@ -175,7 +177,9 @@ async def release_connection(self, connection: HTTPConnection) -> None:
 
     async def close(self) -> None:
         self.is_closed = True
-        connections = list(self.keepalive_connections)
-        self.keepalive_connections.clear()
-        for connection in connections:
+
+        for connection in chain(self.keepalive_connections, self.active_connections):
             await connection.close()
+
+        self.keepalive_connections.clear()
+        self.active_connections.clear()

tests/client/test_async_client.py

@@ -1,4 +1,5 @@
 import pytest
+from h11 import RemoteProtocolError
 
 import httpx
 
@@ -43,13 +44,24 @@ async def test_post_json(server, backend):
     assert response.status_code == 200
 
 
+async def test_stream_manger_release_resources(server, backend):
+    async with httpx.AsyncClient(backend=backend) as http:
+        response = await http.post(
+            server.url.copy_with(path="/echo_body"), stream=True, data=b"*" * 409600
+        )
+
+    assert response.status_code == 200
+    with pytest.raises(RemoteProtocolError):
+        await response.read()
+
+
 async def test_stream_response(server, backend):
     async with httpx.AsyncClient(backend=backend) as client:
         response = await client.request("GET", server.url, stream=True)
-    assert response.status_code == 200
-    body = await response.read()
-    assert body == b"Hello, world!"
-    assert response.content == b"Hello, world!"
+        assert response.status_code == 200
+        body = await response.read()
+        assert body == b"Hello, world!"
+        assert response.content == b"Hello, world!"
 
 
 async def test_access_content_stream_response(server, backend):

tests/client/test_client.py

@@ -1,4 +1,5 @@
 import pytest
+from h11 import RemoteProtocolError
 
 import httpx
 
@@ -33,33 +34,44 @@ def test_post_json(server):
     assert response.reason_phrase == "OK"
 
 
+def test_stream_manger_release_resources(server):
+    with httpx.Client() as http:
+        response = http.post(
+            server.url.copy_with(path="/echo_body"), stream=True, data=b"*" * 409600
+        )
+
+    assert response.status_code == 200
+    with pytest.raises(RemoteProtocolError):
+        response.read()
+
+
 def test_stream_response(server):
     with httpx.Client() as http:
         response = http.get(server.url, stream=True)
-    assert response.status_code == 200
-    content = response.read()
-    assert content == b"Hello, world!"
+        assert response.status_code == 200
+        content = response.read()
+        assert content == b"Hello, world!"
 
 
 def test_stream_iterator(server):
     with httpx.Client() as http:
         response = http.get(server.url, stream=True)
-    assert response.status_code == 200
-    body = b""
-    for chunk in response.stream():
-        body += chunk
-    assert body == b"Hello, world!"
+        assert response.status_code == 200
+        body = b""
+        for chunk in response.stream():
+            body += chunk
+        assert body == b"Hello, world!"
 
 
 def test_raw_iterator(server):
     with httpx.Client() as http:
         response = http.get(server.url, stream=True)
-    assert response.status_code == 200
-    body = b""
-    for chunk in response.raw():
-        body += chunk
-    assert body == b"Hello, world!"
-    response.close()  # TODO: should Response be available as context managers?
+        assert response.status_code == 200
+        body = b""
+        for chunk in response.raw():
+            body += chunk
+        assert body == b"Hello, world!"
+        response.close()  # TODO: should Response be available as context managers?
 
 
 def test_raise_for_status(server):

tests/dispatch/test_connection_pools.py

@@ -112,6 +112,17 @@ async def test_standard_response_close(server, backend):
         assert len(http.keepalive_connections) == 1
 
 
+async def test_close_connection_on_context_manager_exit(server, backend):
+    """
+    After block with all connections should be closed.
+    """
+    async with httpx.ConnectionPool(backend=backend) as http:
+        await http.request("GET", server.url)
+
+    assert len(http.active_connections) == 0
+    assert len(http.keepalive_connections) == 0
+
+
 async def test_premature_response_close(server, backend):
     """
     A premature close should close the connection.