diff --git a/src/informix_db/aio.py b/src/informix_db/aio.py index b9405f5..3b2e3d9 100644 --- a/src/informix_db/aio.py +++ b/src/informix_db/aio.py @@ -63,7 +63,9 @@ import asyncio import contextlib import functools import threading +import weakref from collections.abc import AsyncIterator, Awaitable, Callable +from concurrent.futures import ThreadPoolExecutor from typing import Any, TypeVar from . import connect as _sync_connect @@ -80,6 +82,61 @@ def _to_thread(fn: Callable[..., T], *args: Any, **kwargs: Any) -> Awaitable[T]: return asyncio.to_thread(fn, *args, **kwargs) +def _run_on( + executor: ThreadPoolExecutor, + fn: Callable[..., T], + *args: Any, + **kwargs: Any, +) -> Awaitable[T]: + """Await ``fn`` on a specific executor rather than the loop's default.""" + loop = asyncio.get_running_loop() + return loop.run_in_executor(executor, functools.partial(fn, *args, **kwargs)) + + +def _connection_executor(conn: _SyncConnection) -> ThreadPoolExecutor: + """The dedicated worker thread for one connection, created on demand. + + ``asyncio.to_thread`` runs on the event loop's default executor, which + the whole process shares and which is sized from the CPU count — + ``min(32, cpu_count + 4)``, so six threads on a two-CPU container. + Two consequences, both silent. + + A cancelled await does not stop its worker. ``asyncio.to_thread`` + cannot interrupt the thread, so it keeps running the wire call until + the read timeout expires. Cancellation is ordinary in a web app — a + client disconnect cancels the request task — so a handful of them + pins every thread in the shared pool and unrelated ``to_thread`` work + anywhere else in the process stops dead. Measured: six cancelled + calls against a six-worker default executor starve an unrelated + ``to_thread`` indefinitely. + + And pool concurrency was capped by the same number without saying so. + A pool with ``max_size=20`` on a two-CPU box ran six queries at once. + + One thread per connection is the right size rather than a compromise: + the sync connection serializes every wire operation on its own lock, + so a second thread could do nothing but wait for the first. It also + rules out a deadlock that a shared pool-sized executor invites — with + N threads and N connections, N tasks blocked in ``acquire`` occupy + every thread while the connection they are waiting for is held by a + task that now needs a thread of its own to finish and release it. + + The executor lives on the sync connection so it survives being + returned to the pool and handed out again. The finalizer is the + backstop for a connection dropped without ``close()``: a + ``ThreadPoolExecutor`` that is never shut down leaves its worker + parked on the work queue for the life of the process. + """ + executor = getattr(conn, "_async_executor", None) + if executor is None: + executor = ThreadPoolExecutor( + max_workers=1, thread_name_prefix="informix-conn" + ) + conn._async_executor = executor + weakref.finalize(conn, executor.shutdown, wait=False) + return executor + + class AsyncCursor: """Async wrapper over a sync :class:`Cursor`. Each I/O call awaits a thread-offloaded version of the sync operation. @@ -89,10 +146,14 @@ class AsyncCursor: paying the thread-hop cost. """ - __slots__ = ("_cur",) + __slots__ = ("_cur", "_run") - def __init__(self, cur: _SyncCursor): + def __init__(self, cur: _SyncCursor, run: Callable[..., Awaitable[Any]]): self._cur = cur + # The owning connection's runner, so cursor I/O lands on that + # connection's dedicated thread rather than the shared default + # executor. See _connection_executor. + self._run = run # -- Pass-through synchronous attributes (no I/O) --------------------- @@ -121,32 +182,32 @@ class AsyncCursor: async def execute( self, operation: str, parameters: Any = None ) -> None: - await _to_thread(self._cur.execute, operation, parameters) + await self._run(self._cur.execute, operation, parameters) async def executemany( self, operation: str, seq_of_parameters: Any ) -> None: - await _to_thread( + await self._run( self._cur.executemany, operation, list(seq_of_parameters) ) async def fetchone(self) -> tuple | None: - return await _to_thread(self._cur.fetchone) + return await self._run(self._cur.fetchone) async def fetchmany(self, size: int | None = None) -> list[tuple]: - return await _to_thread(self._cur.fetchmany, size) + return await self._run(self._cur.fetchmany, size) async def fetchall(self) -> list[tuple]: - return await _to_thread(self._cur.fetchall) + return await self._run(self._cur.fetchall) async def close(self) -> None: - await _to_thread(self._cur.close) + await self._run(self._cur.close) # Phase 10/11 BLOB helpers (preserve the sync API surface) async def read_blob_column( self, sql: str, params: tuple = () ) -> bytes | None: - return await _to_thread(self._cur.read_blob_column, sql, params) + return await self._run(self._cur.read_blob_column, sql, params) async def write_blob_column( self, @@ -156,7 +217,7 @@ class AsyncCursor: *, clob: bool = False, ) -> None: - await _to_thread( + await self._run( functools.partial( self._cur.write_blob_column, sql, blob_data, params, clob=clob, @@ -178,14 +239,23 @@ class AsyncCursor: class AsyncConnection: """Async wrapper over a sync :class:`Connection`.""" - __slots__ = ("_conn",) + __slots__ = ("_conn", "_executor") def __init__(self, conn: _SyncConnection): self._conn = conn + self._executor = _connection_executor(conn) + + def _run(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Awaitable[T]: + """Run a blocking connection call on this connection's own thread.""" + return _run_on(self._executor, fn, *args, **kwargs) @classmethod async def connect(cls, *args: Any, **kwargs: Any) -> AsyncConnection: - """Open a connection. Same parameters as :func:`informix_db.connect`.""" + """Open a connection. Same parameters as :func:`informix_db.connect`. + + The connect itself still goes to the default executor -- there is + no connection yet to own a thread, and it is one bounded call. + """ sync_conn = await _to_thread( functools.partial(_sync_connect, *args, **kwargs) ) @@ -196,22 +266,46 @@ class AsyncConnection: return self._conn.closed async def cursor(self) -> AsyncCursor: - sync_cur = await _to_thread(self._conn.cursor) - return AsyncCursor(sync_cur) + sync_cur = await self._run(self._conn.cursor) + return AsyncCursor(sync_cur, self._run) async def commit(self) -> None: - await _to_thread(self._conn.commit) + await self._run(self._conn.commit) async def rollback(self) -> None: - await _to_thread(self._conn.rollback) + await self._run(self._conn.rollback) async def close(self) -> None: - await _to_thread(self._conn.close) + """Close the connection and stop the thread that served it. + + One thread per connection is only affordable if the thread goes + away with the connection. ``ThreadPoolExecutor`` workers park on + the work queue rather than exiting when idle, so an executor that + is never shut down leaks its thread for the life of the process — + the ``weakref.finalize`` in ``_connection_executor`` is a backstop + for connections dropped without ``close()``, not a substitute for + closing here. + + ``wait=False`` because we are on the event loop: the worker has + just finished the close and needs no waiting, and blocking the + loop to confirm that would be the one thing this module exists to + avoid. + """ + try: + await self._run(self._conn.close) + finally: + self._executor.shutdown(wait=False) + # Drop the reference too. A shut-down executor rejects new + # work, so leaving it attached would turn any later wrap of + # this sync connection into a RuntimeError rather than a + # fresh thread. + with contextlib.suppress(AttributeError): + del self._conn._async_executor async def fast_path_call( self, signature: str, *params: object ) -> list[object]: - return await _to_thread(self._conn.fast_path_call, signature, *params) + return await self._run(self._conn.fast_path_call, signature, *params) # Async context-manager support async def __aenter__(self) -> AsyncConnection: @@ -304,8 +398,15 @@ class AsyncConnectionPool: async def release( self, conn: AsyncConnection, *, broken: bool = False ) -> None: - await _to_thread( - functools.partial(self._pool.release, conn._conn, broken=broken) + # On the connection's own thread, not the default executor. That + # thread is idle by definition — the caller is done with the + # connection — and it keeps the release off a shared pool that + # tasks blocked in ``acquire`` may have filled. Release has to + # win that race: it is what frees the connection they are + # waiting for. + await _run_on( + conn._executor, + functools.partial(self._pool.release, conn._conn, broken=broken), ) @contextlib.asynccontextmanager diff --git a/tests/test_async_threads.py b/tests/test_async_threads.py new file mode 100644 index 0000000..1819997 --- /dev/null +++ b/tests/test_async_threads.py @@ -0,0 +1,216 @@ +"""The async layer used a thread pool it shared with the whole process. + +``asyncio.to_thread`` runs on the event loop's default executor. That +executor belongs to the process, not to us, and it is sized from the CPU +count — ``min(32, cpu_count + 4)``, which is six threads on a two-CPU +container. Every blocking call in ``informix_db.aio`` went through it. + +Two consequences, both silent. + +**Cancelled calls hold their threads.** ``asyncio.to_thread`` cannot +interrupt a worker, so a cancelled await leaves the thread running the +wire call until the read timeout expires. Cancellation is ordinary in a +web app — a client disconnect cancels the request task — so a handful of +them pins every thread in the shared pool. Unrelated ``to_thread`` work +anywhere else in the process then stops dead, and so does the driver. + +**Pool concurrency was capped by an unrelated number.** A pool with +``max_size=20`` on a two-CPU box ran six queries at a time, and nothing +said so. + +Each connection now owns one thread. That is the right size rather than +a compromise: the sync connection serializes every wire operation on its +own lock, so a second thread could do nothing but wait for the first. It +also avoids a deadlock that a shared pool-sized executor invites — with +N threads and N connections, N tasks blocked in ``acquire`` occupy every +thread while the connection they wait for is held by a task that now +needs a thread to finish and release it. +""" + +from __future__ import annotations + +import asyncio +import threading +import time +from concurrent.futures import ThreadPoolExecutor + +import pytest + +from informix_db import aio +from tests.conftest import ConnParams + +pytestmark = pytest.mark.integration + + +def _kw(conn_params: ConnParams) -> dict: + return { + "host": conn_params.host, + "port": conn_params.port, + "user": conn_params.user, + "password": conn_params.password, + "database": conn_params.database, + "server": conn_params.server, + "connect_timeout": 10.0, + "read_timeout": 20.0, + "autocommit": True, + } + + +@pytest.mark.asyncio +async def test_queries_run_with_the_default_executor_fully_held( + conn_params: ConnParams, +) -> None: + """The starvation case, with real queries. Four unrelated jobs occupy + every thread of a four-worker default executor; the driver must not + be waiting behind them.""" + loop = asyncio.get_running_loop() + previous = loop._default_executor + loop.set_default_executor(ThreadPoolExecutor(max_workers=4)) + conns = [await aio.connect(**_kw(conn_params)) for _ in range(4)] + cursors = [await c.cursor() for c in conns] + + gate = threading.Event() + hogs = [ + asyncio.create_task(asyncio.to_thread(gate.wait, 30)) for _ in range(4) + ] + await asyncio.sleep(0.3) + + async def query(cursor, i: int): + await cursor.execute(f"SELECT FIRST 1 tabid + {i} FROM systables") + return await cursor.fetchone() + + try: + rows = await asyncio.wait_for( + asyncio.gather(*(query(c, i) for i, c in enumerate(cursors))), + timeout=10.0, + ) + assert len(rows) == 4 + assert all(r is not None for r in rows) + finally: + gate.set() + await asyncio.gather(*hogs, return_exceptions=True) + for c in conns: + await c.close() + # set_default_executor rejects None, which is what the loop + # starts with before anything has used to_thread. + loop.set_default_executor(previous or ThreadPoolExecutor()) + + +@pytest.mark.asyncio +async def test_each_connection_gets_its_own_thread( + conn_params: ConnParams, +) -> None: + a = await aio.connect(**_kw(conn_params)) + b = await aio.connect(**_kw(conn_params)) + try: + assert a._executor is not b._executor + assert a._executor._max_workers == 1, ( + "more than one thread per connection cannot help — the wire " + "lock serializes them anyway" + ) + finally: + await a.close() + await b.close() + + +@pytest.mark.asyncio +async def test_cursor_runs_on_its_connections_thread( + conn_params: ConnParams, +) -> None: + """A cursor must not fall back to the default executor, or half the + work goes back to being shared.""" + conn = await aio.connect(**_kw(conn_params)) + try: + cur = await conn.cursor() + await cur.execute("SELECT FIRST 1 tabid FROM systables") + names: list[str] = [] + await conn._run(lambda: names.append(threading.current_thread().name)) + assert names[0].startswith("informix-conn") + assert cur._run.__self__ is conn + finally: + await conn.close() + + +@pytest.mark.asyncio +async def test_pool_reuses_one_thread_per_connection( + conn_params: ConnParams, +) -> None: + """The executor lives on the sync connection, so a connection handed + out, returned, and handed out again keeps the same thread instead of + spawning one per acquire.""" + pool = await aio.create_pool(**_kw(conn_params), min_size=1, max_size=2) + try: + seen = [] + for _ in range(4): + async with pool.connection() as conn: + cur = await conn.cursor() + await cur.execute("SELECT FIRST 1 tabid FROM systables") + assert await cur.fetchone() is not None + seen.append(id(conn._executor)) + assert len(set(seen)) <= 2, ( + f"expected at most one executor per pooled connection, saw " + f"{len(set(seen))}" + ) + finally: + await pool.close() + + +@pytest.mark.asyncio +async def test_pool_concurrency_is_not_capped_by_the_default_executor( + conn_params: ConnParams, +) -> None: + """Six concurrent pooled queries against a two-worker default + executor. Under the old arrangement at most two could run.""" + loop = asyncio.get_running_loop() + previous = loop._default_executor + pool = await aio.create_pool(**_kw(conn_params), min_size=6, max_size=6) + loop.set_default_executor(ThreadPoolExecutor(max_workers=2)) + try: + + async def one(i: int): + async with pool.connection(timeout=15.0) as conn: + cur = await conn.cursor() + # Long enough that serialized execution would be obvious. + await cur.execute( + "SELECT FIRST 200 a.tabid FROM systables a, systables b" + ) + return len(await cur.fetchall()) + + started = time.monotonic() + results = await asyncio.wait_for( + asyncio.gather(*(one(i) for i in range(6))), timeout=25.0 + ) + assert all(r > 0 for r in results) + assert len(results) == 6 + elapsed = time.monotonic() - started + assert elapsed < 20.0, f"queries appear serialized ({elapsed:.1f}s)" + finally: + await pool.close() + # set_default_executor rejects None, which is what the loop + # starts with before anything has used to_thread. + loop.set_default_executor(previous or ThreadPoolExecutor()) + + +@pytest.mark.asyncio +async def test_closing_a_connection_stops_its_thread( + conn_params: ConnParams, +) -> None: + """One thread per connection is only affordable if the thread goes + away with the connection.""" + before = threading.active_count() + conns = [await aio.connect(**_kw(conn_params)) for _ in range(5)] + for c in conns: + cur = await c.cursor() + await cur.execute("SELECT FIRST 1 tabid FROM systables") + for c in conns: + await c.close() + del conns + # Threads exit asynchronously after shutdown; give them a moment. + for _ in range(50): + if threading.active_count() <= before + 1: + break + await asyncio.sleep(0.1) + assert threading.active_count() <= before + 1, ( + f"connection threads outlived their connections " + f"({before} -> {threading.active_count()})" + )