Repository: encode/httpx

Issue:
Keep-Alive timeouts
We should automatically close any Keep-Alive connections after a timeout period.

Discussion:
tomchristie: Have made some cleanups based on the review.

I'd like for all of this to be a bit tighter, as part of #624, but it seems sensible to approach it incrementally, and at least address the issue for now.

Patch:
diff --git a/docs/compatibility.md b/docs/compatibility.md
index 8f4a19b..506340c 100644
--- a/docs/compatibility.md
+++ b/docs/compatibility.md
@@ -4,7 +4,7 @@ HTTPX aims to be compatible with the `requests` API wherever possible.
 
 This documentation outlines places where the API differs...
 
-## Request URLS
+## Request URLs
 
 Accessing `response.url` will return a `URL` instance, rather than a string.
 Use `str(response.url)` if you need a string instance.
diff --git a/httpx/concurrency/asyncio.py b/httpx/concurrency/asyncio.py
index c73a138..bb795b1 100644
--- a/httpx/concurrency/asyncio.py
+++ b/httpx/concurrency/asyncio.py
@@ -242,6 +242,10 @@ class AsyncioBackend(ConcurrencyBackend):
 
         return SocketStream(stream_reader=stream_reader, stream_writer=stream_writer)
 
+    def time(self) -> float:
+        loop = asyncio.get_event_loop()
+        return loop.time()
+
     async def run_in_threadpool(
         self, func: typing.Callable, *args: typing.Any, **kwargs: typing.Any
     ) -> typing.Any:
diff --git a/httpx/concurrency/auto.py b/httpx/concurrency/auto.py
index c90ce24..c11c064 100644
--- a/httpx/concurrency/auto.py
+++ b/httpx/concurrency/auto.py
@@ -41,6 +41,9 @@ class AutoBackend(ConcurrencyBackend):
     ) -> BaseSocketStream:
         return await self.backend.open_uds_stream(path, hostname, ssl_context, timeout)
 
+    def time(self) -> float:
+        return self.backend.time()
+
     def get_semaphore(self, max_value: int) -> BasePoolSemaphore:
         return self.backend.get_semaphore(max_value)
 
diff --git a/httpx/concurrency/base.py b/httpx/concurrency/base.py
index da6dd98..366e574 100644
--- a/httpx/concurrency/base.py
+++ b/httpx/concurrency/base.py
@@ -99,6 +99,9 @@ class ConcurrencyBackend:
     ) -> BaseSocketStream:
         raise NotImplementedError()  # pragma: no cover
 
+    def time(self) -> float:
+        raise NotImplementedError()  # pragma: no cover
+
     def get_semaphore(self, max_value: int) -> BasePoolSemaphore:
         raise NotImplementedError()  # pragma: no cover
 
diff --git a/httpx/concurrency/trio.py b/httpx/concurrency/trio.py
index 6f2ab2d..0c64988 100644
--- a/httpx/concurrency/trio.py
+++ b/httpx/concurrency/trio.py
@@ -156,6 +156,9 @@ class TrioBackend(ConcurrencyBackend):
             functools.partial(coroutine, **kwargs) if kwargs else coroutine, *args
         )
 
+    def time(self) -> float:
+        return trio.current_time()
+
     def get_semaphore(self, max_value: int) -> BasePoolSemaphore:
         return PoolSemaphore(max_value)
 
diff --git a/httpx/dispatch/connection.py b/httpx/dispatch/connection.py
index 009895a..15ed455 100644
--- a/httpx/dispatch/connection.py
+++ b/httpx/dispatch/connection.py
@@ -38,6 +38,7 @@ class HTTPConnection(Dispatcher):
         self.release_func = release_func
         self.uds = uds
         self.open_connection: typing.Optional[OpenConnection] = None
+        self.expires_at: typing.Optional[float] = None
 
     async def send(
         self,
diff --git a/httpx/dispatch/connection_pool.py b/httpx/dispatch/connection_pool.py
index 4fe3998..db23640 100644
--- a/httpx/dispatch/connection_pool.py
+++ b/httpx/dispatch/connection_pool.py
@@ -78,6 +78,8 @@ class ConnectionStore:
 
 
 class ConnectionPool(Dispatcher):
+    KEEP_ALIVE_EXPIRY = 5.0
+
     def __init__(
         self,
         *,
@@ -101,6 +103,7 @@ class ConnectionPool(Dispatcher):
         self.active_connections = ConnectionStore()
 
         self.backend = lookup_backend(backend)
+        self.next_keepalive_check = 0.0
 
     @property
     def max_connections(self) -> BasePoolSemaphore:
@@ -118,6 +121,21 @@ class ConnectionPool(Dispatcher):
     def num_connections(self) -> int:
         return len(self.keepalive_connections) + len(self.active_connections)
 
+    async def check_keepalive_expiry(self) -> None:
+        now = self.backend.time()
+        if now < self.next_keepalive_check:
+            return
+        self.next_keepalive_check = now + 1.0
+
+        # Iterate through all the keep alive connections.
+        # We create a list here to avoid any 'changed during iteration' errors.
+        keepalives = list(self.keepalive_connections.all.keys())
+        for connection in keepalives:
+            if connection.expires_at is not None and now > connection.expires_at:
+                self.keepalive_connections.remove(connection)
+                self.max_connections.release()
+                await connection.close()
+
     async def send(
         self,
         request: Request,
@@ -125,6 +143,7 @@ class ConnectionPool(Dispatcher):
         cert: CertTypes = None,
         timeout: Timeout = None,
     ) -> Response:
+        await self.check_keepalive_expiry()
         connection = await self.acquire_connection(
             origin=request.url.origin, timeout=timeout
         )
@@ -180,6 +199,8 @@ class ConnectionPool(Dispatcher):
             self.max_connections.release()
             await connection.close()
         else:
+            now = self.backend.time()
+            connection.expires_at = now + self.KEEP_ALIVE_EXPIRY
             self.active_connections.remove(connection)
             self.keepalive_connections.add(connection)
 
diff --git a/httpx/dispatch/http2.py b/httpx/dispatch/http2.py
index 3439b61..e0ddf23 100644
--- a/httpx/dispatch/http2.py
+++ b/httpx/dispatch/http2.py
@@ -64,7 +64,7 @@ class HTTP2Connection(OpenConnection):
             await self.init_complete.wait()
             stream_id = self.state.get_next_available_stream_id()
 
-        stream = HTTP2Stream(stream_id=stream_id, connection=self, state=self.state)
+        stream = HTTP2Stream(stream_id=stream_id, connection=self)
         self.streams[stream_id] = stream
         self.events[stream_id] = []
         return await stream.send(request, timeout)
@@ -158,10 +158,32 @@ class HTTP2Connection(OpenConnection):
         data_to_send = self.state.data_to_send()
         await self.socket.write(data_to_send, timeout)
 
-    async def send_outgoing_data(self, timeout: Timeout) -> None:
+    async def send_headers(
+        self,
+        stream_id: int,
+        headers: typing.List[typing.Tuple[bytes, bytes]],
+        timeout: Timeout,
+    ) -> None:
+        self.state.send_headers(stream_id, headers)
+        data_to_send = self.state.data_to_send()
+        await self.socket.write(data_to_send, timeout)
+
+    async def send_data(self, stream_id: int, chunk: bytes, timeout: Timeout) -> None:
+        self.state.send_data(stream_id, chunk)
+        data_to_send = self.state.data_to_send()
+        await self.socket.write(data_to_send, timeout)
+
+    async def end_stream(self, stream_id: int, timeout: Timeout) -> None:
+        self.state.end_stream(stream_id)
         data_to_send = self.state.data_to_send()
-        if data_to_send:
-            await self.socket.write(data_to_send, timeout)
+        await self.socket.write(data_to_send, timeout)
+
+    async def acknowledge_received_data(
+        self, stream_id: int, amount: int, timeout: Timeout
+    ) -> None:
+        self.state.acknowledge_received_data(amount, stream_id)
+        data_to_send = self.state.data_to_send()
+        await self.socket.write(data_to_send, timeout)
 
     async def close_stream(self, stream_id: int) -> None:
         del self.streams[stream_id]
@@ -172,15 +194,9 @@ class HTTP2Connection(OpenConnection):
 
 
 class HTTP2Stream:
-    def __init__(
-        self,
-        stream_id: int,
-        connection: HTTP2Connection,
-        state: h2.connection.H2Connection,
-    ) -> None:
+    def __init__(self, stream_id: int, connection: HTTP2Connection) -> None:
         self.stream_id = stream_id
         self.connection = connection
-        self.state = state
 
     async def send(self, request: Request, timeout: Timeout) -> Response:
         # Send the request.
@@ -214,8 +230,7 @@ class HTTP2Stream:
             f"target={request.url.full_path!r} "
             f"headers={headers!r}"
         )
-        self.state.send_headers(self.stream_id, headers)
-        await self.connection.send_outgoing_data(timeout)
+        await self.connection.send_headers(self.stream_id, headers, timeout)
 
     async def send_body(self, request: Request, timeout: Timeout) -> None:
         logger.trace(f"send_body stream_id={self.stream_id}")
@@ -226,11 +241,9 @@ class HTTP2Stream:
                 )
                 chunk_size = min(len(data), max_flow)
                 chunk, data = data[:chunk_size], data[chunk_size:]
-                self.state.send_data(self.stream_id, chunk)
-                await self.connection.send_outgoing_data(timeout)
+                await self.connection.send_data(self.stream_id, chunk, timeout)
 
-        self.state.end_stream(self.stream_id)
-        await self.connection.send_outgoing_data(timeout)
+        await self.connection.end_stream(self.stream_id, timeout)
 
     async def receive_response(
         self, timeout: Timeout
@@ -257,10 +270,10 @@ class HTTP2Stream:
         while True:
             event = await self.connection.wait_for_event(self.stream_id, timeout)
             if isinstance(event, h2.events.DataReceived):
-                self.state.acknowledge_received_data(
-                    event.flow_controlled_length, self.stream_id
+                amount = event.flow_controlled_length
+                await self.connection.acknowledge_received_data(
+                    self.stream_id, amount, timeout
                 )
-                await self.connection.send_outgoing_data(timeout)
                 yield event.data
             elif isinstance(event, (h2.events.StreamEnded, h2.events.StreamReset)):
                 break


Test patch:
diff --git a/tests/dispatch/test_connection_pools.py b/tests/dispatch/test_connection_pools.py
index f10b16a..447b8d6 100644
--- a/tests/dispatch/test_connection_pools.py
+++ b/tests/dispatch/test_connection_pools.py
@@ -18,6 +18,37 @@ async def test_keepalive_connections(server, backend):
         assert len(http.keepalive_connections) == 1
 
 
+async def test_keepalive_timeout(server, backend):
+    """
+    Keep-alive connections should timeout.
+    """
+    async with ConnectionPool() as http:
+        response = await http.request("GET", server.url)
+        await response.read()
+        assert len(http.active_connections) == 0
+        assert len(http.keepalive_connections) == 1
+
+        http.next_keepalive_check = 0.0
+        await http.check_keepalive_expiry()
+
+        assert len(http.active_connections) == 0
+        assert len(http.keepalive_connections) == 1
+
+    async with ConnectionPool() as http:
+        http.KEEP_ALIVE_EXPIRY = 0.0
+
+        response = await http.request("GET", server.url)
+        await response.read()
+        assert len(http.active_connections) == 0
+        assert len(http.keepalive_connections) == 1
+
+        http.next_keepalive_check = 0.0
+        await http.check_keepalive_expiry()
+
+        assert len(http.active_connections) == 0
+        assert len(http.keepalive_connections) == 0
+
+
 async def test_differing_connection_keys(server, backend):
     """
     Connections to differing connection keys should result in multiple connections.