diff --git a/src/kimi_cli/utils/aiohttp.py b/src/kimi_cli/utils/aiohttp.py index bff2690767..971e97c3b0 100644 --- a/src/kimi_cli/utils/aiohttp.py +++ b/src/kimi_cli/utils/aiohttp.py @@ -1,6 +1,12 @@ from __future__ import annotations +import asyncio +import atexit +import contextlib import ssl +import threading +import warnings +from typing import Final import aiohttp import certifi @@ -13,12 +19,156 @@ sock_connect=15, ) +_DEFAULT_POOL_LIMIT: Final = 100 +_DEFAULT_POOL_LIMIT_PER_HOST: Final = 30 +_DEFAULT_DNS_CACHE_TTL: Final = 300 + + +class _ConnectionPool: + """Thread-safe lazy-initialised connection pool for aiohttp. + + Each event loop gets its own :class:`aiohttp.TCPConnector` so that + concurrent loops (e.g. from different threads or sequential + ``asyncio.run()`` calls) never invalidate one another. + + Connectors belonging to closed event loops are reaped automatically + when a new connector is created. HTTP keep-alive therefore works + across multiple ``async with`` blocks as long as the same loop is + used. + """ + + def __init__( + self, + *, + limit: int = _DEFAULT_POOL_LIMIT, + limit_per_host: int = _DEFAULT_POOL_LIMIT_PER_HOST, + ttl_dns_cache: int = _DEFAULT_DNS_CACHE_TTL, + register_atexit: bool = True, + ) -> None: + self._limit = limit + self._limit_per_host = limit_per_host + self._ttl_dns_cache = ttl_dns_cache + # Map id(loop) -> (loop, connector). We store a strong reference + # to the loop alongside the connector so that id(loop) stays + # stable as a dict key until we explicitly reap the entry. + self._connectors: dict[int, tuple[asyncio.AbstractEventLoop, aiohttp.TCPConnector]] = {} + self._lock = threading.Lock() + self._atexit_registered: bool = False + self._register_atexit = register_atexit + + def get_connector(self) -> aiohttp.TCPConnector: + """Return a connector bound to the current event loop, creating one if necessary.""" + loop = asyncio.get_running_loop() + loop_id = id(loop) + + with self._lock: + # Opportunistically reap connectors from closed loops first. + # If a long-lived loop keeps taking the fast path, stale + # connectors from dead loops would otherwise never be cleaned. + stale_ids = [ + lid for lid, (sl, c) in self._connectors.items() if c.closed or sl.is_closed() + ] + for lid in stale_ids: + _, old = self._connectors.pop(lid, (None, None)) + if old is not None and not old.closed: + with contextlib.suppress(Exception): + # Best-effort synchronous close — same path aiohttp + # uses internally in BaseConnector.__del__. On a + # dead loop this may raise, in which case GC will + # reclaim the object. + old._close() # type: ignore[reportPrivateUsage] + + # Fast path: already have a live connector for this loop. + entry = self._connectors.get(loop_id) + if entry is not None: + stored_loop, connector = entry + if not connector.closed and not stored_loop.is_closed(): + return connector + + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", category=DeprecationWarning, module="aiohttp.connector" + ) + connector = aiohttp.TCPConnector( + ssl=_ssl_context, + limit=self._limit, + limit_per_host=self._limit_per_host, + ttl_dns_cache=self._ttl_dns_cache, + # On affected CPython versions aiohttp needs a periodic + # callback to reap SSL transports closed by the remote + # peer. aiohttp ignores the flag on versions that do + # not need it, so we always pass it. + enable_cleanup_closed=True, + ) + self._connectors[loop_id] = (loop, connector) + + if self._register_atexit and not self._atexit_registered: + atexit.register(self._sync_close) + self._atexit_registered = True + + return connector + + def new_session( + self, + *, + timeout: aiohttp.ClientTimeout | None = None, + ) -> aiohttp.ClientSession: + """Create a new client session backed by this pool.""" + return aiohttp.ClientSession( + connector=self.get_connector(), + connector_owner=False, + timeout=timeout or _DEFAULT_TIMEOUT, + ) + + async def close(self) -> None: + """Close all connectors. Safe to call multiple times.""" + connectors: list[tuple[asyncio.AbstractEventLoop, aiohttp.TCPConnector]] = [] + with self._lock: + connectors.extend(self._connectors.values()) + self._connectors.clear() + if self._atexit_registered: + atexit.unregister(self._sync_close) + self._atexit_registered = False + current_loop = asyncio.get_running_loop() + for loop, c in connectors: + if c.closed: + continue + if loop is current_loop and not loop.is_closed(): + await c.close() + else: + # Close connectors from dead/other loops synchronously to + # avoid cross-loop task/future errors. + with contextlib.suppress(Exception): + c._close() # type: ignore[reportPrivateUsage] + + def _sync_close(self) -> None: + """Synchronous best-effort cleanup for atexit. + + On process exit the OS reclaims any remaining file descriptors, + so we simply drop the references rather than risk interacting + with a partially-torn-down event loop. + """ + with self._lock: + self._connectors.clear() + + +# Module-level default pool used by all callers. +_default_pool = _ConnectionPool() + def new_client_session( *, timeout: aiohttp.ClientTimeout | None = None, ) -> aiohttp.ClientSession: - return aiohttp.ClientSession( - connector=aiohttp.TCPConnector(ssl=_ssl_context), - timeout=timeout or _DEFAULT_TIMEOUT, - ) + """Create a new client session backed by the default connection pool. + + The returned session must be used as an async context manager + (``async with``). The underlying TCPConnector is shared (per event + loop) and kept open across sessions so that HTTP keep-alive works. + """ + return _default_pool.new_session(timeout=timeout) + + +async def close_connector() -> None: + """Close the default shared connectors. Safe to call multiple times.""" + await _default_pool.close() diff --git a/tests/utils/test_aiohttp_timeout.py b/tests/utils/test_aiohttp_timeout.py index cfaf429093..b537d21edf 100644 --- a/tests/utils/test_aiohttp_timeout.py +++ b/tests/utils/test_aiohttp_timeout.py @@ -3,14 +3,22 @@ from __future__ import annotations import asyncio +import contextlib import pytest -from kimi_cli.utils.aiohttp import new_client_session +from kimi_cli.utils.aiohttp import _ConnectionPool, close_connector, new_client_session + + +@pytest.fixture(autouse=True) +async def _reset_pool(): + """Ensure a fresh default pool for every test.""" + await close_connector() + yield + await close_connector() async def test_default_session_has_timeout(): - """new_client_session() should create a session with non-None timeout values.""" async with new_client_session() as session: assert session.timeout.total == 120 assert session.timeout.sock_read == 60 @@ -18,7 +26,6 @@ async def test_default_session_has_timeout(): async def test_custom_timeout_override(): - """Callers can override the default timeout.""" import aiohttp custom = aiohttp.ClientTimeout(total=30, sock_read=10) @@ -28,7 +35,6 @@ async def test_custom_timeout_override(): async def test_slow_server_is_interrupted(): - """A server that accepts but never responds should be interrupted by timeout.""" hang_forever = asyncio.Event() async def _slow_handler(reader, writer): @@ -50,3 +56,196 @@ async def _slow_handler(reader, writer): hang_forever.set() server.close() await server.wait_closed() + + +async def test_sessions_share_connector_per_loop(): + async with new_client_session() as session1: + connector1 = session1.connector + + async with new_client_session() as session2: + connector2 = session2.connector + + assert connector1 is not None + assert connector1 is connector2 + assert connector1.limit == 100 + assert connector1.limit_per_host == 30 + assert not connector1.closed + + +async def test_connector_recreated_after_close(): + async with new_client_session() as session1: + connector1 = session1.connector + + await close_connector() + + async with new_client_session() as session2: + connector2 = session2.connector + + assert connector1 is not None + assert connector2 is not None + assert connector1 is not connector2 + assert connector1.closed + assert not connector2.closed + + +async def test_concurrent_creation_is_safe(): + await close_connector() + + async def _open(): + async with new_client_session() as session: + return session.connector + + connectors = await asyncio.gather(*(_open() for _ in range(20))) + + unique = {id(c) for c in connectors} + assert len(unique) == 1, f"Expected 1 unique connector, got {len(unique)}" + + connector = connectors[0] + assert connector is not None + assert connector.limit == 100 + assert connector.limit_per_host == 30 + assert not connector.closed + + +async def test_keepalive_pools_connections(): + seen_ports: set[int] = set() + + async def _handler(reader, writer): + peer = writer.get_extra_info("peername") + if peer is not None: + seen_ports.add(peer[1]) + while True: + try: + request = await asyncio.wait_for(reader.readuntil(b"\r\n\r\n"), timeout=1.0) + if not request: + break + except (TimeoutError, asyncio.LimitOverrunError, ConnectionResetError): + break + writer.write( + b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: keep-alive\r\n\r\nok" + ) + await writer.drain() + writer.close() + + server = await asyncio.start_server(_handler, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + + try: + async with new_client_session() as s1, s1.get(f"http://127.0.0.1:{port}/") as resp: + await resp.text() + + async with new_client_session() as s2, s2.get(f"http://127.0.0.1:{port}/") as resp: + await resp.text() + + assert len(seen_ports) == 1, f"Expected connection reuse (1 port), got {seen_ports}" + finally: + server.close() + await server.wait_closed() + + +async def test_stress_pool_under_load(): + await close_connector() + + request_count = 0 + + async def _handler(reader, writer): + nonlocal request_count + request_count += 1 + writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: keep-alive\r\n\r\nok") + await writer.drain() + with contextlib.suppress(TimeoutError): + await asyncio.wait_for(reader.read(), timeout=1.0) + writer.close() + + server = await asyncio.start_server(_handler, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + + try: + + async def _fetch(): + async with new_client_session() as s, s.get(f"http://127.0.0.1:{port}/") as r: + await r.text() + + await asyncio.gather(*(_fetch() for _ in range(20))) + + assert request_count == 20, f"Expected 20 requests, got {request_count}" + + from kimi_cli.utils.aiohttp import _default_pool + + connector = _default_pool.get_connector() + assert connector is not None + assert not connector.closed + finally: + server.close() + await server.wait_closed() + + +async def test_isolated_pool_does_not_bleed(): + pool = _ConnectionPool(limit=5, limit_per_host=2, register_atexit=False) + + async with pool.new_session() as s: + assert s.connector is not None + assert s.connector.limit == 5 + assert s.connector.limit_per_host == 2 + + await pool.close() + assert not pool._connectors + + +def test_separate_loops_get_separate_connectors(): + async def _get_connector_id(): + async with new_client_session() as s: + return id(s.connector) + + id1 = asyncio.run(_get_connector_id()) + id2 = asyncio.run(_get_connector_id()) + + assert id1 != id2, "Expected separate connectors for separate event loops" + + +def test_reaps_closed_loop_connectors(): + async def _get_connector(): + from kimi_cli.utils.aiohttp import _default_pool + + return _default_pool.get_connector() + + c1 = asyncio.run(_get_connector()) + c1_id = id(c1) + + c2 = asyncio.run(_get_connector()) + c2_id = id(c2) + + assert c1_id != c2_id, "Expected a new connector for the new loop" + from kimi_cli.utils.aiohttp import _default_pool + + assert c1_id not in {id(c) for _, c in _default_pool._connectors.values()} + + +async def _get_connector_in_thread() -> int: + async with new_client_session() as s: + return id(s.connector) + + +def test_multi_threaded_access(): + import threading + + results: list[int] = [] + errors: list[Exception] = [] + + def _run(): + try: + results.append(asyncio.run(_get_connector_in_thread())) + except Exception as exc: + errors.append(exc) + + threads = [threading.Thread(target=_run) for _ in range(8)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not errors, f"Threads raised exceptions: {errors}" + assert len(results) == 8 + # Connectors are keyed by id(loop); sequential threads may reuse the + # same memory address for a new loop after the old one is GC'd, so we + # only assert that every thread succeeded, not that all IDs are unique.