Fix pool timeout edge-case. (#688)

Author: lovelydinosaur

CHANGELOG.md

@@ -10,6 +10,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - Improve correctness of tracebacks on network exceptions, by raising properly chained exceptions. (#678)
 - Prevent connection-hanging behaviour when HTTP/2 connections are closed by a server-sent 'GoAway' frame. (#679)
 - Fix edge-case exception when removing requests from the connection pool. (#680)
+- Fix pool timeout edge-case. (#688)
 
 ## 0.17.0 (March 16th, 2023)
 

httpcore/_async/connection_pool.py

@@ -31,7 +31,8 @@ def unset_connection(self) -> None:
     async def wait_for_connection(
         self, timeout: Optional[float] = None
     ) -> AsyncConnectionInterface:
-        await self._connection_acquired.wait(timeout=timeout)
+        if self.connection is None:
+            await self._connection_acquired.wait(timeout=timeout)
         assert self.connection is not None
         return self.connection
 

httpcore/_sync/connection_pool.py

@@ -31,7 +31,8 @@ def unset_connection(self) -> None:
     def wait_for_connection(
         self, timeout: Optional[float] = None
     ) -> ConnectionInterface:
-        self._connection_acquired.wait(timeout=timeout)
+        if self.connection is None:
+            self._connection_acquired.wait(timeout=timeout)
         assert self.connection is not None
         return self.connection
 

tests/_async/test_connection_pool.py

@@ -571,6 +571,68 @@ async def test_connection_pool_timeout():
                 await pool.request("GET", "https://example.com/", extensions=extensions)
 
 
+@pytest.mark.anyio
+async def test_connection_pool_timeout_zero():
+    """
+    A pool timeout of 0 shouldn't raise a PoolTimeout if there's
+    no need to wait on a new connection.
+    """
+    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!",
+            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!",
+        ]
+    )
+
+    # Use a pool timeout of zero.
+    extensions = {"timeout": {"pool": 0}}
+
+    # A connection pool configured to allow only one connection at a time.
+    async with AsyncConnectionPool(
+        network_backend=network_backend, max_connections=1
+    ) as pool:
+        # Two consecutive requests with a pool timeout of zero.
+        # Both succeed without raising a timeout.
+        response = await pool.request(
+            "GET", "https://example.com/", extensions=extensions
+        )
+        assert response.status == 200
+        assert response.content == b"Hello, world!"
+
+        response = await pool.request(
+            "GET", "https://example.com/", extensions=extensions
+        )
+        assert response.status == 200
+        assert response.content == b"Hello, world!"
+
+    # A connection pool configured to allow only one connection at a time.
+    async with AsyncConnectionPool(
+        network_backend=network_backend, max_connections=1
+    ) as pool:
+        # Two concurrent requests with a pool timeout of zero.
+        # Only the first will succeed without raising a timeout.
+        async with pool.stream(
+            "GET", "https://example.com/", extensions=extensions
+        ) as response:
+            # The first response hasn't yet completed.
+            with pytest.raises(PoolTimeout):
+                # So a pool timeout occurs.
+                await pool.request("GET", "https://example.com/", extensions=extensions)
+            # The first response now completes.
+            await response.aread()
+
+        assert response.status == 200
+        assert response.content == b"Hello, world!"
+
+
 @pytest.mark.anyio
 async def test_http11_upgrade_connection():
     """

tests/_sync/test_connection_pool.py

@@ -572,6 +572,68 @@ def test_connection_pool_timeout():
 
 
 
+def test_connection_pool_timeout_zero():
+    """
+    A pool timeout of 0 shouldn't raise a PoolTimeout if there's
+    no need to wait on a new connection.
+    """
+    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!",
+            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!",
+        ]
+    )
+
+    # Use a pool timeout of zero.
+    extensions = {"timeout": {"pool": 0}}
+
+    # A connection pool configured to allow only one connection at a time.
+    with ConnectionPool(
+        network_backend=network_backend, max_connections=1
+    ) as pool:
+        # Two consecutive requests with a pool timeout of zero.
+        # Both succeed without raising a timeout.
+        response = pool.request(
+            "GET", "https://example.com/", extensions=extensions
+        )
+        assert response.status == 200
+        assert response.content == b"Hello, world!"
+
+        response = pool.request(
+            "GET", "https://example.com/", extensions=extensions
+        )
+        assert response.status == 200
+        assert response.content == b"Hello, world!"
+
+    # A connection pool configured to allow only one connection at a time.
+    with ConnectionPool(
+        network_backend=network_backend, max_connections=1
+    ) as pool:
+        # Two concurrent requests with a pool timeout of zero.
+        # Only the first will succeed without raising a timeout.
+        with pool.stream(
+            "GET", "https://example.com/", extensions=extensions
+        ) as response:
+            # The first response hasn't yet completed.
+            with pytest.raises(PoolTimeout):
+                # So a pool timeout occurs.
+                pool.request("GET", "https://example.com/", extensions=extensions)
+            # The first response now completes.
+            response.read()
+
+        assert response.status == 200
+        assert response.content == b"Hello, world!"
+
+
+
 def test_http11_upgrade_connection():
     """
     HTTP "101 Switching Protocols" indicates an upgraded connection.