From 6d012889752f65c1a51d0ad6e5970fc97d19c4ca Mon Sep 17 00:00:00 2001 From: Jeff Scudder Date: Fri, 25 Sep 2026 20:32:42 -0700 Subject: [PATCH] fix: reap aiohttp sessions and auth locks belonging to closed event loops `BaseApiClient` caches its aiohttp session and asyncio auth lock per event loop but never released an entry once that loop closed. Servers that run each request on a fresh `asyncio.run()` loop therefore retained one session, connector and socket set per request served. Entries whose loop is closed are now dropped on the next access, bounding both caches by the number of live loops. PiperOrigin-RevId: 988653992 --- google/genai/_api_client.py | 16 ++++ .../client/test_client_initialization.py | 83 ++++++++++++++++++- 2 files changed, 97 insertions(+), 2 deletions(-) diff --git a/google/genai/_api_client.py b/google/genai/_api_client.py index 27adcf813..172dd0fd4 100644 --- a/google/genai/_api_client.py +++ b/google/genai/_api_client.py @@ -953,6 +953,20 @@ def _is_session_closed(self, session: Any) -> bool: return bool(session._auth_request._closed) return False + @staticmethod + def _discard_closed_loop_entries(entries: dict[Any, Any]) -> None: + """Drops entries of a loop-keyed cache whose event loop has been closed. + + Servers that run every request on a fresh loop (ADK's sync `Runner.run()`, + used by Agent Engine, calls `asyncio.run()` per request) would otherwise + grow these caches without bound, pinning one aiohttp session, connector and + set of sockets per request served. A session on a closed loop cannot be + awaited shut, but dropping the last reference to it lets the garbage + collector release the connector and its sockets. + """ + for loop in [loop for loop in entries if loop.is_closed()]: + del entries[loop] + @property def _aiohttp_session( self, @@ -976,6 +990,7 @@ async def _get_aiohttp_session( loop = asyncio.get_running_loop() with self._sync_auth_lock: + self._discard_closed_loop_entries(self._aiohttp_sessions) session = self._aiohttp_sessions.get(loop) if session is not None and self._is_session_closed(session): session = None @@ -1319,6 +1334,7 @@ async def _get_async_auth_lock(self) -> asyncio.Lock: """ loop = asyncio.get_running_loop() with self._sync_auth_lock: + self._discard_closed_loop_entries(self._async_auth_locks) if loop not in self._async_auth_locks: self._async_auth_locks[loop] = asyncio.Lock() return self._async_auth_locks[loop] diff --git a/google/genai/tests/client/test_client_initialization.py b/google/genai/tests/client/test_client_initialization.py index 659d1ca4b..d353db2d0 100644 --- a/google/genai/tests/client/test_client_initialization.py +++ b/google/genai/tests/client/test_client_initialization.py @@ -22,6 +22,7 @@ import os import ssl import sys +import threading from unittest import mock import certifi @@ -1676,8 +1677,6 @@ async def mock_async_operation(op_id: int): @pytest.mark.asyncio async def test_get_async_auth_lock_creation_lock_lifecycle(): """Tests the creation lock lifecycle and cleanup.""" - import threading - client = Client( vertexai=True, project="fake_project_id", location="fake-location" ) @@ -1947,3 +1946,83 @@ async def test_async_mtls_uses_refreshable_credentials(monkeypatch): assert passed_creds.valid == True mock_creds.expired = True assert passed_creds.valid == False + + +def _run_on_fresh_loops(coro_fn, count): + """Runs coro_fn() count times, each on its own thread and its own loop. + + This is the shape ADK's sync `Runner.run()` gives Agent Engine: every request + gets a new thread and a new `asyncio.run()` loop that is closed on the way + out, while the genai client itself is a long-lived singleton. + + Args: + coro_fn: Zero-argument callable returning the coroutine to run. + count: How many loops to run it on, one after another. + """ + for _ in range(count): + thread = threading.Thread(target=lambda: asyncio.run(coro_fn())) + thread.start() + thread.join() + + +@requires_aiohttp +def test_aiohttp_sessions_not_retained_for_closed_event_loops(): + """Sessions belonging to finished loops must not accumulate. b/496663148.""" + client = Client( + vertexai=True, project="fake_project_id", location="fake-location" + ) + api_client.has_aiohttp = True + base_client = client._api_client + + _run_on_fresh_loops(base_client._get_aiohttp_session, 10) + + # Reaping happens on access, so the newest loop's entry survives until the + # next call. What matters is that the cache stays bounded by the number of + # live loops rather than growing once per request served. + assert len(base_client._aiohttp_sessions) <= 1 + + +def test_async_auth_locks_not_retained_for_closed_event_loops(): + """Auth locks belonging to finished loops must not accumulate. b/496663148.""" + client = Client( + vertexai=True, project="fake_project_id", location="fake-location" + ) + base_client = client._api_client + + _run_on_fresh_loops(base_client._get_async_auth_lock, 10) + + assert len(base_client._async_auth_locks) <= 1 + + +@requires_aiohttp +def test_aiohttp_session_kept_per_live_event_loop(): + """Each live loop keeps its own session; reaping must not steal it. + + Reusing one loop's session on another is what raised `RuntimeError: ... got + Future ... attached to a different loop` in b/496663148. + """ + client = Client( + vertexai=True, project="fake_project_id", location="fake-location" + ) + api_client.has_aiohttp = True + base_client = client._api_client + sessions = [] + barrier = threading.Barrier(3) + + def hold_loop_open(): + async def run(): + session = await base_client._get_aiohttp_session() + sessions.append(session) + # Keep this loop alive until every thread has its own session. + await asyncio.get_running_loop().run_in_executor(None, barrier.wait) + assert await base_client._get_aiohttp_session() is session + + asyncio.run(run()) + + threads = [threading.Thread(target=hold_loop_open) for _ in range(3)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert len({id(session) for session in sessions}) == 3