scope pooled clients per event loop

Author: purificant

src/bubble_data_api_client/pool.py

@@ -3,6 +3,7 @@
 import asyncio
 import atexit
 import threading
+import weakref
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 
@@ -12,47 +13,86 @@
 from bubble_data_api_client.exceptions import ConfigurationError
 from bubble_data_api_client.http_client import httpx_client_factory
 
-# global client pool keyed by config
-_clients: dict[tuple[str, str], httpx.AsyncClient] = {}
+# type aliases
+_ConfigKey = tuple[str, str]
+_LoopClientMap = weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, httpx.AsyncClient]
+
+# global client pool: config_key → { loop → client }
+# WeakKeyDictionary auto-removes entries when the loop is garbage collected
+_clients: dict[_ConfigKey, _LoopClientMap] = {}
 _lock = threading.Lock()
 
 
-def _make_client_key(config: BubbleConfig) -> tuple[str, str]:
+def _make_client_key(config: BubbleConfig) -> _ConfigKey:
     """Generate a unique key for client pooling based on config."""
     return (config["data_api_root_url"], config["api_key"])
 
 
+def _create_client_from_config(config: BubbleConfig) -> httpx.AsyncClient:
+    """Create a new httpx client from config."""
+    base_url = config["data_api_root_url"]
+    if not base_url:
+        raise ConfigurationError("data_api_root_url")
+    api_key = config["api_key"]
+    if not api_key:
+        raise ConfigurationError("api_key")
+    return httpx_client_factory(base_url=base_url, api_key=api_key)
+
+
 def get_client() -> httpx.AsyncClient:
-    """Get or create a client for the current config. Thread-safe."""
+    """Get or create a client for the current config and event loop. Thread-safe.
+
+    Each (config, event_loop) pair gets its own client. When an event loop is
+    garbage collected, its associated clients are automatically removed.
+
+    If called outside an async context (no running loop), returns a fresh
+    uncached client as a fallback.
+    """
     config = get_config()
-    key = _make_client_key(config)
 
-    # fast path: no lock if client exists
-    if key in _clients:
-        return _clients[key]
+    try:
+        current_loop = asyncio.get_running_loop()
+    except RuntimeError:
+        # no running loop - return uncached client
+        # it will bind to whatever loop it's eventually used in
+        return _create_client_from_config(config)
+
+    key = _make_client_key(config)
 
-    # slow path: acquire lock for creation
     with _lock:
-        # double-check after acquiring lock
         if key not in _clients:
-            base_url = config["data_api_root_url"]
-            if not base_url:
-                raise ConfigurationError("data_api_root_url")
-            api_key = config["api_key"]
-            if not api_key:
-                raise ConfigurationError("api_key")
-            _clients[key] = httpx_client_factory(base_url=base_url, api_key=api_key)
-        return _clients[key]
+            _clients[key] = weakref.WeakKeyDictionary()
+
+        loop_clients = _clients[key]
+
+        # return cached client if it exists and is still open
+        if current_loop in loop_clients and not loop_clients[current_loop].is_closed:
+            return loop_clients[current_loop]
+
+        # create new client (either missing or was closed externally)
+        loop_clients[current_loop] = _create_client_from_config(config)
+        return loop_clients[current_loop]
 
 
 async def close_clients() -> None:
-    """Close all clients in the pool. Thread-safe. Safe to call multiple times."""
+    """Close all clients for the current event loop. Thread-safe. Safe to call multiple times.
+
+    Only closes clients bound to the calling loop. Clients for other loops are
+    left alone (they should be closed by their respective loops, or will be
+    garbage collected when those loops die).
+    """
+    current_loop = asyncio.get_running_loop()
+
     with _lock:
-        clients_to_close = list(_clients.values())
-        _clients.clear()
+        clients_to_close: list[httpx.AsyncClient] = [
+            loop_clients.pop(current_loop) for loop_clients in _clients.values() if current_loop in loop_clients
+        ]
 
     for client in clients_to_close:
-        await client.aclose()
+        try:
+            await client.aclose()
+        except Exception:
+            pass  # best-effort cleanup, continue with remaining clients
 
 
 @asynccontextmanager
@@ -65,9 +105,12 @@ async def client_scope() -> AsyncIterator[None]:
 
 
 def _atexit_cleanup() -> None:
-    """Best-effort cleanup at interpreter exit."""
+    """Best-effort cleanup of all clients at interpreter exit."""
     with _lock:
-        clients_to_close = list(_clients.values())
+        # collect all clients from all config/loop combinations
+        clients_to_close: list[httpx.AsyncClient] = []
+        for loop_clients in _clients.values():
+            clients_to_close.extend(loop_clients.values())
         _clients.clear()
 
     if not clients_to_close:
@@ -92,7 +135,10 @@ def _atexit_cleanup() -> None:
             # no running loop, create one and close all clients
             async def _close_all() -> None:
                 for client in clients_to_close:
-                    await client.aclose()
+                    try:
+                        await client.aclose()
+                    except Exception:
+                        pass
 
             try:
                 asyncio.run(_close_all())