diff --git a/context_intelligence_server/config.py b/context_intelligence_server/config.py index 48b20e97..fd72680c 100644 --- a/context_intelligence_server/config.py +++ b/context_intelligence_server/config.py @@ -797,6 +797,26 @@ def resolve_neo4j_query(self) -> Neo4jClientConfig: access_mode="READ", ) + # Upper bound on concurrent bolt connections for a driver shared across many + # logical callers (the lifespan admin driver, the lifespan query driver, the + # registry's shared per-session driver). Well under the server's default + # bolt thread-pool size so a driver leak can no longer starve it. The neo4j + # driver's own default is 100; 50 is a deliberate reduction, since every + # session now shares one pool instead of holding a private one. + # + # No companion max_connection_lifetime knob: the driver already recycles + # pooled connections at 3600 s by default, so a setting whose default equals + # the library default would change nothing. + neo4j_max_connection_pool_size: int = 50 + + @field_validator("neo4j_max_connection_pool_size") + @classmethod + def _validate_neo4j_max_connection_pool_size(cls, v: int) -> int: + """Fail loud on a non-positive pool size.""" + if v <= 0: + raise ValueError(f"neo4j_max_connection_pool_size must be > 0, got {v}") + return v + # ------------------------------------------------------------------------- # Storage paths # ------------------------------------------------------------------------- diff --git a/context_intelligence_server/main.py b/context_intelligence_server/main.py index f172a845..dba618a6 100644 --- a/context_intelligence_server/main.py +++ b/context_intelligence_server/main.py @@ -16,7 +16,7 @@ from fastapi import Depends, FastAPI, HTTPException, Request from fastapi.responses import JSONResponse, Response -from neo4j import READ_ACCESS, WRITE_ACCESS, AsyncGraphDatabase +from neo4j import READ_ACCESS, WRITE_ACCESS from context_intelligence_server import __version__ from context_intelligence_server.auth import ( @@ -41,6 +41,7 @@ EventResponse, ) from context_intelligence_server.neo4j_store import ( + build_bounded_neo4j_driver, count_untagged_nodes, ensure_neo4j_schema, ) @@ -61,13 +62,18 @@ def _neo4j_access_const(mode: str) -> str: def build_neo4j_driver(config: Neo4jClientConfig) -> Any: - """Construct an AsyncGraphDatabase driver from a resolved Neo4j client config. + """Construct the pool-bounded admin AsyncGraphDatabase driver. Shared by ``lifespan()`` (the admin driver, on every server boot) and ``doctor.run_doctor()`` (the CLI), so the two entry points can never - construct the connection differently. + construct the connection differently. Delegates the actual driver + construction to ``build_bounded_neo4j_driver`` so the pool-bounding kwargs + have one source of truth, shared with ``SessionRegistry``'s driver. """ - return AsyncGraphDatabase.driver(config.url, auth=config.auth) + return build_bounded_neo4j_driver( + config, + max_connection_pool_size=_settings.neo4j_max_connection_pool_size, + ) # --------------------------------------------------------------------------- @@ -255,9 +261,11 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: # build_neo4j_driver() is the SAME helper doctor.run_doctor() uses, so the # server and the doctor CLI can never construct this connection differently. app.state.neo4j_driver = build_neo4j_driver(_admin) - # Cypher-query (read-intent): /cypher + dashboard reads. - app.state.neo4j_query_driver = AsyncGraphDatabase.driver( - _query.url, auth=_query.auth + # Cypher-query (read-intent): /cypher + dashboard reads. Bounded through the + # same helper as the admin driver so every process-wide pool shares one cap. + app.state.neo4j_query_driver = build_bounded_neo4j_driver( + _query, + max_connection_pool_size=_settings.neo4j_max_connection_pool_size, ) # Stash the resolved query access_mode so /cypher opens READ sessions without # re-resolving settings on every request. @@ -411,9 +419,25 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: _sweep_task.cancel() with suppress(asyncio.CancelledError): await _sweep_task + # ORDER IS LOAD-BEARING. Quiesce the drainers FIRST. Every session's + # graph store now shares ONE driver, so closing it under a live drainer + # is no longer a per-session concern: the drainer's batch fails, it + # spends its max_delivery_attempts budget in ~250 ms, and + # _handle_exhausted_batch dead-letters each line AND commits the offset + # past it -- discarding healthy events that merely happened to be + # queued at shutdown, with no replay on the next boot. Cancelling first + # routes each drainer through CancelledError -> _safe_close -> a final + # flush while the driver is still open. + logger.info("lifespan_shutdown: quiescing drain workers") + await registry.shutdown_workers() logger.info("lifespan_shutdown: closing Neo4j drivers") await app.state.neo4j_driver.close() await app.state.neo4j_query_driver.close() + # The registry's shared per-session driver is independent of the two + # above (its own pool, built from settings.resolve_neo4j_admin() the + # first time a session is created) -- close it here too so no bolt + # connection outlives the process. + await registry.close_neo4j_driver() app = FastAPI( diff --git a/context_intelligence_server/neo4j_store.py b/context_intelligence_server/neo4j_store.py index 52ddc7d9..821db6c5 100644 --- a/context_intelligence_server/neo4j_store.py +++ b/context_intelligence_server/neo4j_store.py @@ -22,8 +22,58 @@ from neo4j import unit_of_work as _unit_of_work from neo4j.exceptions import DriverError, Neo4jError +from context_intelligence_server.config import Neo4jClientConfig + _LOG = logging.getLogger(__name__) + +def build_bounded_neo4j_driver( + config: Neo4jClientConfig, + *, + max_connection_pool_size: int, + connection_acquisition_timeout: float | None = None, +) -> Any: + """Construct an AsyncGraphDatabase driver with a bounded connection pool. + + Single source of truth for the pool-bounding kwargs applied to any driver + meant to be shared across many logical callers (the lifespan admin driver, + the lifespan query driver, the registry's shared per-session driver). All + three construct through here so they can never diverge. + + Args: + max_connection_pool_size: + Hard cap on concurrent bolt connections for this driver. + connection_acquisition_timeout: + How long a caller waits for a free pooled connection before + failing. ``None`` (the lifespan drivers) leaves the driver default + in place, matching what those two did before they were routed + through this helper. ``SessionRegistry`` passes + ``settings.neo4j_lock_timeout`` so the shared per-session driver + keeps the acquisition budget the per-session drivers it replaced + carried -- and it matters more now, not less: one bounded pool is + shared by every session, so acquisition can genuinely queue. + + ``max_connection_lifetime`` is deliberately not set: the neo4j driver + already recycles pooled connections at 3600 s by default, so passing it + would be a knob that changes nothing. + """ + kwargs: dict[str, Any] = { + "max_connection_pool_size": max_connection_pool_size, + # Explicit auto-retry budget for transient errors (e.g. deadlocks) so + # the managed-transaction retry window is deliberate and reviewable + # rather than relying on the driver default implicitly. Carried over + # verbatim from the per-session driver construction this helper + # subsumed. + "max_transaction_retry_time": 30.0, + } + if ( + connection_acquisition_timeout is not None + and connection_acquisition_timeout > 0 + ): + kwargs["connection_acquisition_timeout"] = connection_acquisition_timeout + return AsyncGraphDatabase.driver(config.url, auth=config.auth, **kwargs) + + # --------------------------------------------------------------------------- # Cypher identifier validation # --------------------------------------------------------------------------- @@ -1179,12 +1229,15 @@ def __init__( flush_chunk_rows: int = 100, flush_chunk_bytes: int = 4_194_304, neo4j_lock_timeout: float | None = None, + driver: Any | None = None, ) -> None: - """Initialise the store and create the async Neo4j driver. + """Initialise the store, reusing or creating the async Neo4j driver. Args: uri: Bolt/neo4j URI, e.g. ``bolt://localhost:7687``. + Ignored when ``driver`` is provided. auth: ``(username, password)`` tuple, or ``None`` for no-auth. + Ignored when ``driver`` is provided. database: Target Neo4j database name (default: ``"neo4j"``). workspace: Workspace to scope writes to. ``None`` resolves to ``"default"`` via the ``workspace`` property. @@ -1195,18 +1248,29 @@ def __init__( so a blocked flush raises ``Neo4jError`` instead of parking forever. ``None`` disables the timeout (default: no per-transaction limit). - Also sets ``connection_acquisition_timeout`` on - the driver to the same value so pool-exhaustion - failures also surface quickly. + When the store builds its own driver (``driver`` + not provided), this also sets + ``connection_acquisition_timeout`` on it to the + same value so pool-exhaustion failures surface + quickly. + driver: A pre-built async driver to reuse instead of + constructing a new one. When provided, this store + does not own the driver's lifecycle: ``close()`` + flushes and no-ops on the driver itself, leaving + it open for other stores sharing it. """ - # Explicit auto-retry budget for transient errors (e.g. deadlocks) so the - # managed-transaction retry window is deliberate and reviewable rather than - # relying on the driver default implicitly. 30.0s is a working default; - # design Open Question #3 — verify driver 6.1.0 backoff constants before tuning. - driver_kwargs: dict[str, Any] = {"max_transaction_retry_time": 30.0} - if neo4j_lock_timeout is not None and neo4j_lock_timeout > 0: - driver_kwargs["connection_acquisition_timeout"] = neo4j_lock_timeout - self._driver = AsyncGraphDatabase.driver(uri, auth=auth, **driver_kwargs) + if driver is not None: + self._driver = driver + self._owns_driver = False + else: + # Explicit auto-retry budget for transient errors (e.g. deadlocks) so + # the managed-transaction retry window is deliberate and reviewable + # rather than relying on the driver default implicitly. + driver_kwargs: dict[str, Any] = {"max_transaction_retry_time": 30.0} + if neo4j_lock_timeout is not None and neo4j_lock_timeout > 0: + driver_kwargs["connection_acquisition_timeout"] = neo4j_lock_timeout + self._driver = AsyncGraphDatabase.driver(uri, auth=auth, **driver_kwargs) + self._owns_driver = True self._database = database self._workspace = workspace self._created_by: str | None = None @@ -1224,6 +1288,19 @@ def __init__( else None ) + # ------------------------------------------------------------------ + # owns_driver property + # ------------------------------------------------------------------ + + @property + def owns_driver(self) -> bool: + """True when this store built its own driver; False when injected. + + Governs ``close()``: a store that does not own its driver must never + close it, since other stores may still be using it. + """ + return self._owns_driver + # ------------------------------------------------------------------ # workspace property # ------------------------------------------------------------------ @@ -1633,10 +1710,15 @@ async def _ensure_schema(self) -> None: # once Neo4j is reachable / duplicates are cleared by the dedup pass). async def close(self) -> None: - """Flush pending writes, await any background task, and close the driver. + """Flush pending writes and close the driver, if this store owns it. Handles event-loop mismatch gracefully when closing the driver from a different loop context. Sets ``_closed`` on completion. + + When the driver was injected (``owns_driver`` is False), the driver is + left open: it is shared with other stores/callers and closing it here + would break them out from under their own in-flight work. The shared + driver's owner is responsible for closing it exactly once. """ # Final flush to persist remaining buffer contents try: @@ -1646,11 +1728,12 @@ async def close(self) -> None: "Final flush failed during close; buffered writes may be lost" ) - # Close the driver, ignoring event-loop mismatch errors - try: - await self._driver.close() - except RuntimeError: - pass + if self._owns_driver: + # Close the driver, ignoring event-loop mismatch errors + try: + await self._driver.close() + except RuntimeError: + pass self._closed = True diff --git a/context_intelligence_server/registry.py b/context_intelligence_server/registry.py index 59796802..d31143d2 100644 --- a/context_intelligence_server/registry.py +++ b/context_intelligence_server/registry.py @@ -12,11 +12,14 @@ from context_intelligence_server.blob_store import AsyncDiskBlobStore from context_intelligence_server.config import get_settings -from context_intelligence_server.status import EventRecord, ring_buffer -from context_intelligence_server.neo4j_store import Neo4jGraphStore +from context_intelligence_server.neo4j_store import ( + Neo4jGraphStore, + build_bounded_neo4j_driver, +) from context_intelligence_server.pipeline import process_event, setup_handlers from context_intelligence_server.queue_manager import Batch, QueueManager from context_intelligence_server.services import HookStateService +from context_intelligence_server.status import EventRecord, ring_buffer logger = logging.getLogger("context_intelligence_server") @@ -76,6 +79,11 @@ def __init__(self) -> None: self._queue_manager: QueueManager | None = None self._write_semaphore: asyncio.Semaphore | None = None self._max_delivery_attempts: int = 0 + # Shared, pool-bounded Neo4j driver for every per-session Neo4jGraphStore + # (see _ensure_neo4j_driver). Built lazily for the same reason as + # _queue_manager; kept separate from _ensure_infra so the two concerns + # can evolve independently. + self._neo4j_driver: Any | None = None # Live conservation counters surfaced via /status (accepted/written/ # replayed/write_retries) so silently-dropped events are observable. self._accepted_total: int = 0 @@ -106,6 +114,72 @@ def queue_manager(self) -> QueueManager: assert self._queue_manager is not None return self._queue_manager + def _ensure_neo4j_driver(self) -> Any: + """Build the shared, pool-bounded Neo4j driver on first use. + + Lazy for the same reason as ``_ensure_infra``. Kept as its own method + (not folded into ``_ensure_infra``) so the two constructions stay + independent edits. + """ + if self._neo4j_driver is None: + settings = get_settings() + admin = settings.resolve_neo4j_admin() + self._neo4j_driver = build_bounded_neo4j_driver( + admin, + max_connection_pool_size=settings.neo4j_max_connection_pool_size, + # Parity with the per-session driver this one replaces: a + # blocked acquisition must surface on the SAME budget as a + # blocked transaction. Load-bearing now in a way it was not + # before -- every session shares this one bounded pool, so + # acquisition can actually queue. + connection_acquisition_timeout=settings.neo4j_lock_timeout, + ) + return self._neo4j_driver + + @property + def neo4j_driver(self) -> Any: + """The single shared, pool-bounded driver used by every per-session + Neo4jGraphStore -- never closed by a per-session finalize.""" + return self._ensure_neo4j_driver() + + async def shutdown_workers(self) -> None: + """Quiesce every drain worker BEFORE the shared driver is closed. + + Ordering invariant (must run before ``close_neo4j_driver``): a live + drainer that meets a closed shared driver fails its batch, spends its + ``max_delivery_attempts`` budget in ~250 ms (5 attempts x the 50 ms + ``_DRAIN_POLL_INTERVAL`` backoff), and falls into + ``_handle_exhausted_batch`` -- which dead-letters each line AND commits + the offset past it. Those are healthy events that merely happened to be + queued at shutdown, and once dead-lettered they never replay. + + Cancelling instead routes each drainer through its ``CancelledError`` + handler -> ``_safe_close(worker)`` -> a final flush on a driver that is + still open. Anything left uncommitted stays in the durable queue and + replays on the next boot, which is the pre-shared-driver behaviour. + + Exceptions are collected, not raised: shutdown must not be derailed by + one failing worker. + """ + tasks = [w.task for w in self._workers.values() if w.task is not None] + if not tasks: + return + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + async def close_neo4j_driver(self) -> None: + """Close the shared driver exactly once, at process shutdown. + + Call ``shutdown_workers()`` first -- see its docstring for why closing + this driver under a live drainer dead-letters good events. + + No-op if the driver was never built (no session has run yet). + """ + if self._neo4j_driver is not None: + await self._neo4j_driver.close() + self._neo4j_driver = None + @property def write_semaphore(self) -> asyncio.Semaphore: """The single shared global cap on concurrent Neo4j-write flushes.""" @@ -742,6 +816,7 @@ def get_or_create( neo4j_store = Neo4jGraphStore( uri=_admin.url, auth=_admin.auth, + driver=self.neo4j_driver, flush_chunk_rows=settings.neo4j_flush_chunk_rows, flush_chunk_bytes=settings.neo4j_flush_chunk_bytes, neo4j_lock_timeout=settings.neo4j_lock_timeout, diff --git a/docs/architecture/README.md b/docs/architecture/README.md index 8c58426f..965a3ac2 100644 --- a/docs/architecture/README.md +++ b/docs/architecture/README.md @@ -259,22 +259,42 @@ and their typed relationships (`HAS_EVENT`, `EMITTED`, `REFERENCES_BLOB`). ## Neo4j client topology -The server connects to Neo4j through **two drivers** rather than one. The **admin -driver** (`neo4j_driver`) is opened in **WRITE** access mode and carries the ingest -path — the drainer's batch flushes and schema work. The **cypher_query driver** -(`neo4j_query_driver`) is opened in **READ** access mode and serves `POST /cypher` -reads. Both drivers are created together in the dual-driver lifespan -at startup and closed together at shutdown. With legacy flat config -(`neo4j_url` / `neo4j_user` / `neo4j_password`), both drivers fall back to the same -endpoint and shared credentials, differing only by access-mode hint; a structured -`neo4j:` block lets the read driver take a separate credential and/or URL (e.g. a -read replica). Their connection health is reported independently on `/status` as -`neo4j_connected` (admin) and `neo4j_query_connected` (cypher_query). +The server connects to Neo4j through **three drivers**, each with a distinct job. + +Two are owned by the lifespan. The **admin driver** (`neo4j_driver`) is opened in +**WRITE** access mode and handles boot-time and operational work: schema creation, +the untagged-node integrity check, and the `/status` connectivity probe. The +**cypher_query driver** (`neo4j_query_driver`) is opened in **READ** access mode and +serves `POST /cypher` reads. Both are created together at startup and closed +together at shutdown. With legacy flat config (`neo4j_url` / `neo4j_user` / +`neo4j_password`), both fall back to the same endpoint and shared credentials, +differing only by access-mode hint; a structured `neo4j:` block lets the read driver +take a separate credential and/or URL (e.g. a read replica). Their connection health +is reported independently on `/status` as `neo4j_connected` (admin) and +`neo4j_query_connected` (cypher_query). + +The third is the **shared session driver**, owned by `SessionRegistry`, and it is the +one that carries the **ingest path** — every per-session `Neo4jGraphStore`'s batch +flushes. It is built lazily on the first session (from the same resolved admin client +config) with a bounded pool (`neo4j_max_connection_pool_size`, default 50) and +injected into every store, so a per-session finalize can never close the driver other +live sessions are still using. Before it existed, each session built and held its own +unshared driver, which is how bolt connections accumulated until the server's thread +pool starved. It is closed exactly once, at shutdown, and only *after* +`SessionRegistry.shutdown_workers()` has quiesced the drain workers — closing it +under a live drainer makes that drainer exhaust its retry budget and dead-letter +healthy queued events. + +Because `/status` probes only the admin and cypher_query drivers, +`neo4j_connected: true` does **not** by itself mean the ingest path is healthy. +Ingest health is visible through the pipeline-conservation counters in the `metrics` +block (`written_total`, `residual`, `degraded`). > **Note:** the existing `.dot` diagrams (e.g. `05-durable-ingest-queue`) show only -> the **write path** through the admin driver. The two-driver split is not yet -> rendered in any diagram — a dedicated topology diagram is a known follow-up. This -> prose subsection is the interim reference; no new `.dot` is authored in this pass. +> the **write path**, and label it as the admin driver — that write path is now the +> shared session driver. The three-driver split is not yet rendered in any diagram — +> a dedicated topology diagram is a known follow-up. This prose subsection is the +> interim reference; no new `.dot` is authored in this pass. --- diff --git a/tests/conftest.py b/tests/conftest.py index ecfd94d9..2c3b7c9f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -122,6 +122,7 @@ class _SettingsProxy: neo4j_flush_chunk_rows: int = _real.neo4j_flush_chunk_rows neo4j_flush_chunk_bytes: int = _real.neo4j_flush_chunk_bytes neo4j_lock_timeout: float = _real.neo4j_lock_timeout + neo4j_max_connection_pool_size: int = _real.neo4j_max_connection_pool_size # Neo4j two-client split (doc 12): SessionRegistry.get_or_create() calls # settings.resolve_neo4j_admin() directly, so this proxy (which stands diff --git a/tests/neo4j/test_driver_leak_bounded.py b/tests/neo4j/test_driver_leak_bounded.py new file mode 100644 index 00000000..132ac3e0 --- /dev/null +++ b/tests/neo4j/test_driver_leak_bounded.py @@ -0,0 +1,145 @@ +"""Behavioral evidence that the per-session driver leak is gone. + +Drives many sessions through the real registry construction path +(``get_or_create`` -> shared ``Neo4jGraphStore``) against a live Neo4j, doing +a real write per session so bolt connections are actually opened, then queries +the server's own connection list to prove the open bolt connections stay +bounded by the pool (never scale with the session count) and are released on +driver close. +""" + +from __future__ import annotations + +import asyncio +from typing import Any + +import pytest +from context_intelligence_server.config import Settings +from context_intelligence_server.registry import SessionRegistry +from neo4j import AsyncGraphDatabase # type: ignore[attr-defined] + +pytestmark = pytest.mark.neo4j + +# Many more sessions than the pool can hold: if each session built its own +# driver (the leak), open bolt connections would scale with this number. +SESSION_COUNT = 30 +# Small, explicit pool so the bound is unmistakable in the observed count. +POOL_SIZE = 8 + +# The Python driver identifies itself with this user-agent prefix; the probe +# driver below uses a different one so it is excluded from the count. +_PY_DRIVER_UA_PREFIX = "neo4j-python" +_PROBE_UA = "leak-probe/1.0" + +_COUNT_QUERY = ( + "CALL dbms.listConnections() YIELD connector, userAgent " + f"WHERE connector = 'bolt' AND userAgent STARTS WITH '{_PY_DRIVER_UA_PREFIX}' " + "RETURN count(*) AS c" +) + + +async def _count_python_bolt_connections(probe_driver: Any) -> int: + """Return the number of open bolt connections opened by the Python driver. + + Excludes the probe driver itself (distinct user-agent) so the count + reflects only the registry's shared driver. + """ + async with probe_driver.session() as session: + result = await session.run(_COUNT_QUERY) + record = await result.single() + return int(record["c"]) + + +@pytest.mark.asyncio +async def test_bolt_connections_stay_bounded_and_release( + neo4j_container: dict[str, Any], + monkeypatch: pytest.MonkeyPatch, + tmp_path: Any, + capsys: Any, +) -> None: + bolt_url = neo4j_container["bolt_url"] + user = neo4j_container["user"] + password = neo4j_container["password"] + + # Real settings pointing at the container, with a small bounded pool and + # writable scratch paths. The registry resolves its shared driver from these. + settings = Settings( + neo4j_url=bolt_url, + neo4j_user=user, + neo4j_password=password, + neo4j_max_connection_pool_size=POOL_SIZE, + blob_path=str(tmp_path / "blobs"), + queues_path=str(tmp_path / "queues"), + ) + monkeypatch.setattr( + "context_intelligence_server.registry.get_settings", + lambda: settings, + ) + + probe_driver = AsyncGraphDatabase.driver( + bolt_url, auth=(user, password), user_agent=_PROBE_UA + ) + + reg = SessionRegistry() + + baseline = await _count_python_bolt_connections(probe_driver) + + async def run_session(i: int) -> None: + # Exactly the construction path the leak came from: the registry builds + # (or reuses) its one shared driver and injects it into this session's + # store. The write forces a real bolt connection through the pool. + worker = reg.get_or_create(f"leak-session-{i}", f"/workspace/{i}") + graph = worker.services.graph + await graph.upsert_node( + f"node-{i}", {"label": "Event", "session": f"leak-session-{i}"} + ) + await graph.flush() + + await asyncio.gather(*(run_session(i) for i in range(SESSION_COUNT))) + + during_load = await _count_python_bolt_connections(probe_driver) + + # Stop the idle drain workers before tearing the driver down. + for worker in list(reg._workers.values()): + if worker.task is not None: + worker.task.cancel() + await asyncio.gather( + *(w.task for w in reg._workers.values() if w.task is not None), + return_exceptions=True, + ) + + # Reclaim: closing the one shared driver must release every bolt connection. + await reg.close_neo4j_driver() + + # Poll briefly for the server to observe the closed connections. + after_close = during_load + for _ in range(20): + after_close = await _count_python_bolt_connections(probe_driver) + if after_close == 0: + break + await asyncio.sleep(0.25) + + await probe_driver.close() + + with capsys.disabled(): + print( + f"\n[driver-leak evidence] sessions={SESSION_COUNT} pool_size={POOL_SIZE} " + f"baseline={baseline} during_load={during_load} after_close={after_close}" + ) + + # One shared driver was built for all sessions, not one per session. + assert reg._neo4j_driver is None # closed above + # The core proof: open bolt connections are bounded by the pool and do NOT + # scale with the session count. + assert during_load <= POOL_SIZE + 2, ( + f"open bolt connections ({during_load}) exceeded the pool bound " + f"({POOL_SIZE}); a per-session driver would scale with {SESSION_COUNT}" + ) + assert during_load < SESSION_COUNT, ( + f"open bolt connections ({during_load}) scaled with session count " + f"({SESSION_COUNT}) -- the leak is not fixed" + ) + # Reclaim proof: the shared driver's pool is fully released on close. + assert after_close == 0, ( + f"bolt connections not released after driver close (still {after_close})" + ) diff --git a/tests/neo4j/test_shared_driver_adverse_state.py b/tests/neo4j/test_shared_driver_adverse_state.py new file mode 100644 index 00000000..c817ac12 --- /dev/null +++ b/tests/neo4j/test_shared_driver_adverse_state.py @@ -0,0 +1,59 @@ +"""Adverse-state test: shared driver survives one session's close while +another session is mid-drain. + +Two Neo4jGraphStore instances share one driver (mirrors SessionRegistry's +per-session construction). Closing session A must not disturb session B's +in-flight write; the shared driver is closed exactly once, by its owner, +after both sessions are done with it. +""" + +from __future__ import annotations + +from typing import Any + +import pytest +from context_intelligence_server.neo4j_store import Neo4jGraphStore +from neo4j import AsyncGraphDatabase # type: ignore[attr-defined] + +pytestmark = pytest.mark.neo4j + + +@pytest.mark.asyncio +async def test_session_a_close_does_not_disrupt_session_b( + neo4j_container: dict[str, Any], +) -> None: + shared_driver = AsyncGraphDatabase.driver( + neo4j_container["bolt_url"], + auth=(neo4j_container["user"], neo4j_container["password"]), + ) + + store_a = Neo4jGraphStore( + uri=neo4j_container["bolt_url"], driver=shared_driver, workspace="session-a" + ) + store_b = Neo4jGraphStore( + uri=neo4j_container["bolt_url"], driver=shared_driver, workspace="session-b" + ) + + await store_b.upsert_node("node-b", {"label": "Event"}) + + # Session A finalizes and closes its store while B still has unflushed + # work buffered -- this is the adverse state: A's close must not touch + # the driver B is still using. + await store_a.close() + + # B's write still lands: the shared driver was never closed under it. + await store_b.flush() + fetched = await store_b.get_node("node-b") + assert fetched is not None + + await store_b.close() + + # The shared driver is still open after both stores are done with it -- + # neither store owned it. Its owner closes it exactly once, at shutdown. + async with shared_driver.session() as session: + result = await session.run("RETURN 1 AS one") + record = await result.single() + assert record is not None + assert record["one"] == 1 + + await shared_driver.close() diff --git a/tests/neo4j/test_shutdown_no_deadletter.py b/tests/neo4j/test_shutdown_no_deadletter.py new file mode 100644 index 00000000..39f84097 --- /dev/null +++ b/tests/neo4j/test_shutdown_no_deadletter.py @@ -0,0 +1,102 @@ +"""Behavioral evidence that shutdown does not discard queued events. + +The shared Neo4j driver made driver lifetime a cross-session concern: closing +it while a drain worker is still running is no longer "that session's driver +going away", it is *every* session's driver going away mid-flight. + +A drainer that meets a closed driver fails its batch, spends its +``max_delivery_attempts`` budget in ~250 ms (5 attempts x the 50 ms +``_DRAIN_POLL_INTERVAL`` backoff), and lands in ``_handle_exhausted_batch`` -- +which dead-letters each line AND commits the offset past it. Those are healthy +events that merely happened to be queued when the process stopped, and once +dead-lettered they never replay. + +This drives real events through the real registry against a live Neo4j and +asserts the shutdown sequence used by ``lifespan`` (quiesce the drainers, then +close the shared driver) dead-letters nothing. +""" + +from __future__ import annotations + +import asyncio +import json +from typing import Any + +import pytest +from context_intelligence_server.config import Settings +from context_intelligence_server.registry import SessionRegistry + +pytestmark = pytest.mark.neo4j + +# Enough to guarantee undrained lines are still queued when shutdown starts -- +# _DRAIN_MAX_BATCH is 100, so this is several batches deep. +EVENT_COUNT = 400 + + +def _event_line(session_id: str, i: int) -> bytes: + return json.dumps( + { + "event": "tool:pre", + "workspace": "/ws", + "data": { + "session_id": session_id, + "timestamp": "2024-01-01T00:00:00+00:00", + "tool_name": f"tool-{i}", + }, + } + ).encode("utf-8") + + +@pytest.mark.asyncio +async def test_shutdown_quiesce_then_close_deadletters_nothing( + neo4j_container: dict[str, Any], + monkeypatch: pytest.MonkeyPatch, + tmp_path: Any, +) -> None: + settings = Settings( + neo4j_url=neo4j_container["bolt_url"], + neo4j_user=neo4j_container["user"], + neo4j_password=neo4j_container["password"], + blob_path=str(tmp_path / "blobs"), + queues_path=str(tmp_path / "queues"), + ) + monkeypatch.setattr( + "context_intelligence_server.registry.get_settings", lambda: settings + ) + + reg = SessionRegistry() + session_id = "shutdown-session" + qm = reg.queue_manager + + dead_lettered: list[str] = [] + real_dead_letter = qm.dead_letter + + async def spy_dead_letter(sid: str, raw: bytes, error: str) -> None: + dead_lettered.append(error) + await real_dead_letter(sid, raw, error) + + monkeypatch.setattr(qm, "dead_letter", spy_dead_letter) + + for i in range(EVENT_COUNT): + await qm.append(session_id, _event_line(session_id, i)) + + reg.get_or_create(session_id, "/ws") + # Let the drainer get into its loop with work still queued behind it. + await asyncio.sleep(0.4) + + # The lifespan shutdown sequence, in order. Reversing these two lines is the + # regression this test exists to catch. + await reg.shutdown_workers() + await reg.close_neo4j_driver() + + # Stay on the loop as the server would during its shutdown window: a + # still-live drainer would burn its retry budget and dead-letter here. + await asyncio.sleep(2.0) + + assert dead_lettered == [], ( + f"{len(dead_lettered)} healthy queued events were dead-lettered during " + "shutdown; the drain workers must be quiesced before the shared driver " + "closes. First error: " + f"{dead_lettered[0] if dead_lettered else ''}" + ) + assert reg._neo4j_driver is None diff --git a/tests/test_main.py b/tests/test_main.py index f8d71322..a5ab80a5 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -718,7 +718,7 @@ async def test_lifespan_creates_and_closes_driver( "context_intelligence_server.main.setup_logging", ) as mock_setup_logging, patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ) as mock_driver_factory, ): @@ -741,6 +741,50 @@ async def test_lifespan_creates_and_closes_driver( assert mock_driver.close.await_count == 2 +async def test_lifespan_quiesces_drain_workers_before_closing_shared_driver( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Shutdown MUST cancel the drain workers before closing the registry's + shared Neo4j driver. + + Every session's graph store shares one driver now, so closing it under a + live drainer is not a per-session concern: the drainer's batch fails, it + burns its max_delivery_attempts budget in ~250ms, and + _handle_exhausted_batch dead-letters each line AND commits the offset past + it -- silently discarding healthy events that merely happened to be queued + at shutdown, with no replay on the next boot. + """ + mock_driver = MagicMock() + mock_driver.close = AsyncMock() + + calls: list[str] = [] + + async def record_shutdown_workers() -> None: + calls.append("shutdown_workers") + + async def record_close_driver() -> None: + calls.append("close_neo4j_driver") + + monkeypatch.setattr( + main_module.registry, "shutdown_workers", record_shutdown_workers + ) + monkeypatch.setattr(main_module.registry, "close_neo4j_driver", record_close_driver) + + with ( + patch("context_intelligence_server.main.setup_logging"), + patch( + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", + return_value=mock_driver, + ), + ): + async with lifespan(main_module.app): + assert calls == [] + + assert calls == ["shutdown_workers", "close_neo4j_driver"], ( + f"drain workers must be quiesced BEFORE the shared driver closes; got {calls}" + ) + + # --------------------------------------------------------------------------- # Lifespan crash-recovery + workers==1 guard tests (Phase B2) # --------------------------------------------------------------------------- @@ -780,7 +824,7 @@ async def test_lifespan_recovers_and_respawns_drainers( with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch( @@ -824,7 +868,7 @@ async def test_lifespan_skips_recovery_for_empty_workspace( with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch( @@ -875,7 +919,7 @@ async def test_lifespan_default_respawns_all_recovered_sessions_unbounded( with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch("context_intelligence_server.main.ensure_neo4j_schema", new=AsyncMock()), @@ -908,7 +952,7 @@ async def test_lifespan_respawn_cap_defers_remainder_and_logs_warning( with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch("context_intelligence_server.main.ensure_neo4j_schema", new=AsyncMock()), @@ -954,7 +998,7 @@ async def test_lifespan_deferred_sessions_untouched_and_recoverable_next_boot( with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch("context_intelligence_server.main.ensure_neo4j_schema", new=AsyncMock()), @@ -995,7 +1039,7 @@ async def test_lifespan_respawn_cap_zero_defers_everything( with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch("context_intelligence_server.main.ensure_neo4j_schema", new=AsyncMock()), @@ -1065,7 +1109,7 @@ async def test_lifespan_enables_sweep_under_finite_limit( with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch("context_intelligence_server.main.ensure_neo4j_schema", new=AsyncMock()), @@ -1092,7 +1136,7 @@ async def test_lifespan_no_sweep_when_limit_unbounded( with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch("context_intelligence_server.main.ensure_neo4j_schema", new=AsyncMock()), @@ -1122,7 +1166,7 @@ async def test_lifespan_no_sweep_when_interval_zero( with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch("context_intelligence_server.main.ensure_neo4j_schema", new=AsyncMock()), @@ -1161,7 +1205,7 @@ async def test_lifespan_calls_ensure_schema_with_fail_on_data_conflict() -> None with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch( @@ -1193,7 +1237,7 @@ async def test_lifespan_raises_on_ensure_schema_data_conflict() -> None: with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch( @@ -1218,7 +1262,7 @@ async def test_lifespan_raises_on_untagged_nodes() -> None: with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch( @@ -1247,7 +1291,7 @@ async def test_lifespan_does_not_raise_on_clean_graph() -> None: with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch( @@ -1276,7 +1320,7 @@ async def test_lifespan_does_not_raise_when_health_check_itself_fails( with ( patch("context_intelligence_server.main.setup_logging"), patch( - "context_intelligence_server.main.AsyncGraphDatabase.driver", + "context_intelligence_server.neo4j_store.AsyncGraphDatabase.driver", return_value=mock_driver, ), patch( diff --git a/tests/test_neo4j_driver_sharing.py b/tests/test_neo4j_driver_sharing.py new file mode 100644 index 00000000..d9c49e7b --- /dev/null +++ b/tests/test_neo4j_driver_sharing.py @@ -0,0 +1,369 @@ +"""Tests for shared, pool-bounded Neo4j driver reuse across sessions. + +Covers: +- Neo4jGraphStore accepts an injected driver and never closes it (owns_driver). +- The self-built path (no injected driver) is unchanged: it owns and closes + its own driver. +- SessionRegistry hands the same driver instance to every per-session store + instead of building one per session. +- The shared driver keeps the pool cap AND the acquisition/retry budgets the + per-session drivers it replaced carried. +- shutdown_workers() quiesces every drainer before the shared driver closes. +""" + +from __future__ import annotations + +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from context_intelligence_server.config import Neo4jClientConfig, get_settings +from context_intelligence_server.neo4j_store import ( + Neo4jGraphStore, + build_bounded_neo4j_driver, +) +from context_intelligence_server.registry import SessionRegistry, SessionWorker + +# --------------------------------------------------------------------------- +# Injected-driver seam (owns_driver) +# --------------------------------------------------------------------------- + + +def test_injected_driver_reports_owns_driver_false() -> None: + """A store built with driver= must report owns_driver False.""" + shared_driver = AsyncMock() + + store_a = Neo4jGraphStore(uri="bolt://unused:7687", driver=shared_driver) + store_b = Neo4jGraphStore(uri="bolt://unused:7687", driver=shared_driver) + + assert store_a.owns_driver is False + assert store_b.owns_driver is False + assert store_a._driver is shared_driver + assert store_b._driver is shared_driver + + +@pytest.mark.asyncio +async def test_close_on_injected_driver_does_not_close_it() -> None: + """Closing one store sharing an injected driver must not close the driver + out from under a second store still using it -- the safety property the + whole shared-driver change rests on.""" + shared_driver = AsyncMock() + + store_a = Neo4jGraphStore(uri="bolt://unused:7687", driver=shared_driver) + store_b = Neo4jGraphStore(uri="bolt://unused:7687", driver=shared_driver) + + await store_a.close() + + shared_driver.close.assert_not_awaited() + # store_b's driver reference is untouched and still the live shared mock -- + # a real driver would still be open and usable by store_b at this point. + assert store_b._driver is shared_driver + + +@pytest.mark.asyncio +async def test_self_built_driver_still_owned_and_closed() -> None: + """With driver=None (default), behavior is unchanged: the store builds + and owns its driver, and close() closes it.""" + with patch( + "context_intelligence_server.neo4j_store.AsyncGraphDatabase" + ) as mock_adb: + mock_driver = AsyncMock() + mock_adb.driver.return_value = mock_driver + + store = Neo4jGraphStore(uri="bolt://localhost:7687", auth=("u", "p")) + assert store.owns_driver is True + + await store.close() + mock_driver.close.assert_awaited_once() + + +# --------------------------------------------------------------------------- +# Registry driver reuse +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_registry_shares_one_driver_across_sessions(monkeypatch) -> None: + """get_or_create must hand the SAME driver object to every per-session + Neo4jGraphStore -- reuse, not a new driver per session_id.""" + reg = SessionRegistry() + + built_drivers: list[object] = [] + + def fake_build_bounded_driver(config, **kwargs): + driver = AsyncMock() + built_drivers.append(driver) + return driver + + monkeypatch.setattr( + "context_intelligence_server.registry.build_bounded_neo4j_driver", + fake_build_bounded_driver, + ) + + worker_a = reg.get_or_create("session-a", "/workspace/a") + worker_b = reg.get_or_create("session-b", "/workspace/b") + + # Exactly one driver was ever built, and both sessions' stores share it. + assert len(built_drivers) == 1 + assert worker_a.services.graph._driver is built_drivers[0] + assert worker_b.services.graph._driver is built_drivers[0] + assert worker_a.services.graph.owns_driver is False + assert worker_b.services.graph.owns_driver is False + + for worker in (worker_a, worker_b): + if worker.task is not None: + worker.task.cancel() + + +def _spy_driver_factory(reg: SessionRegistry, monkeypatch) -> list[dict]: + """Patch the registry's driver factory to record every build call. + + Returns the list of recorded builds; each entry is + ``{"driver": , "kwargs": {...}}``. + """ + builds: list[dict] = [] + + def fake_build_bounded_driver(config, **kwargs): + driver = AsyncMock() + builds.append({"driver": driver, "kwargs": kwargs}) + return driver + + monkeypatch.setattr( + "context_intelligence_server.registry.build_bounded_neo4j_driver", + fake_build_bounded_driver, + ) + return builds + + +def _cancel_workers(reg: SessionRegistry) -> None: + """Cancel any real drain tasks started by get_or_create (test cleanup).""" + for worker in reg._workers.values(): + if worker.task is not None: + worker.task.cancel() + + +@pytest.mark.asyncio +async def test_n_sessions_build_exactly_one_driver(monkeypatch) -> None: + """The core leak-gone proof: N distinct sessions must build the driver + exactly ONCE, not once per session_id (which was the leak).""" + reg = SessionRegistry() + builds = _spy_driver_factory(reg, monkeypatch) + + n = 30 + workers = [reg.get_or_create(f"session-{i}", f"/workspace/{i}") for i in range(n)] + + assert len(builds) == 1, ( + f"expected exactly 1 driver build across {n} sessions, " + f"got {len(builds)} (a per-session build is the leak)" + ) + the_driver = builds[0]["driver"] + for worker in workers: + assert worker.services.graph._driver is the_driver + assert worker.services.graph.owns_driver is False + + _cancel_workers(reg) + + +@pytest.mark.asyncio +async def test_shared_driver_built_with_bounded_kwargs(monkeypatch) -> None: + """The single shared driver must be built WITH the bounded pool cap + (default 50) -- an unbounded build is the leak -- AND with the acquisition + budget the per-session drivers it replaced carried.""" + reg = SessionRegistry() + builds = _spy_driver_factory(reg, monkeypatch) + + reg.get_or_create("session-1", "/workspace/1") + + assert len(builds) == 1 + kwargs = builds[0]["kwargs"] + assert kwargs["max_connection_pool_size"] == 50 + # Parity with neo4j_lock_timeout (default 30.0): without this the driver + # silently falls back to its own 60.0s default, doubling how long a caller + # parks on an exhausted pool -- and the pool is now SHARED, so exhaustion + # is reachable in a way it never was with a private pool per session. + assert kwargs["connection_acquisition_timeout"] == get_settings().neo4j_lock_timeout + + _cancel_workers(reg) + + +# --------------------------------------------------------------------------- +# Driver-kwarg parity with the per-session driver this helper replaced +# --------------------------------------------------------------------------- + + +def test_build_bounded_driver_preserves_per_session_driver_kwargs() -> None: + """build_bounded_neo4j_driver must carry over BOTH kwargs the per-session + Neo4jGraphStore driver set, not just the new pool cap. + + Regression guard: extracting the construction into a shared helper silently + dropped ``connection_acquisition_timeout`` (30.0 -> the driver's 60.0 + default) and the explicit ``max_transaction_retry_time``. + """ + config = Neo4jClientConfig(url="bolt://unused:7687", username="u", password="p") + + with patch( + "context_intelligence_server.neo4j_store.AsyncGraphDatabase" + ) as mock_adb: + build_bounded_neo4j_driver( + config, + max_connection_pool_size=50, + connection_acquisition_timeout=30.0, + ) + + kwargs = mock_adb.driver.call_args.kwargs + assert kwargs["max_connection_pool_size"] == 50 + assert kwargs["connection_acquisition_timeout"] == 30.0 + assert kwargs["max_transaction_retry_time"] == 30.0 + # Deliberately absent: the driver already defaults to 3600 s, so setting it + # would be a knob that changes nothing. + assert "max_connection_lifetime" not in kwargs + + +def test_build_bounded_driver_omits_acquisition_timeout_when_not_given() -> None: + """The lifespan admin/query drivers pass no acquisition timeout, so the + helper must leave the driver default in place for them -- matching exactly + what those two did before they were routed through this helper.""" + config = Neo4jClientConfig(url="bolt://unused:7687", username="u", password="p") + + with patch( + "context_intelligence_server.neo4j_store.AsyncGraphDatabase" + ) as mock_adb: + build_bounded_neo4j_driver(config, max_connection_pool_size=50) + + kwargs = mock_adb.driver.call_args.kwargs + assert "connection_acquisition_timeout" not in kwargs + assert kwargs["max_transaction_retry_time"] == 30.0 + + +@pytest.mark.asyncio +async def test_concurrent_first_sessions_build_exactly_one_driver(monkeypatch) -> None: + """Racing many get_or_create calls as the FIRST sessions must still build + exactly one driver -- the lazy build is synchronous with no await between + the None-check and the assignment, so concurrent coroutines cannot double-build.""" + reg = SessionRegistry() + builds = _spy_driver_factory(reg, monkeypatch) + + async def make(i: int) -> None: + # Wrap the sync get_or_create so many run concurrently under gather. + reg.get_or_create(f"session-{i}", f"/workspace/{i}") + + await asyncio.gather(*(make(i) for i in range(40))) + + assert len(builds) == 1, ( + f"concurrent first-sessions raced into {len(builds)} driver builds; " + "the lazy build must be single-shot" + ) + + _cancel_workers(reg) + + +@pytest.mark.asyncio +async def test_close_neo4j_driver_reclaims_once_and_is_idempotent(monkeypatch) -> None: + """After sessions run, close_neo4j_driver() must close the shared driver + exactly once and clear it; a second call is a safe no-op.""" + reg = SessionRegistry() + builds = _spy_driver_factory(reg, monkeypatch) + + reg.get_or_create("session-1", "/workspace/1") + reg.get_or_create("session-2", "/workspace/2") + assert len(builds) == 1 + shared_driver = builds[0]["driver"] + + await reg.close_neo4j_driver() + shared_driver.close.assert_awaited_once() + assert reg._neo4j_driver is None + + # Second call: no driver left to close, must not raise or double-close. + await reg.close_neo4j_driver() + shared_driver.close.assert_awaited_once() + assert reg._neo4j_driver is None + + _cancel_workers(reg) + + +@pytest.mark.asyncio +async def test_close_neo4j_driver_none_safe_when_no_session_ran() -> None: + """close_neo4j_driver() must be safe when no session ever built a driver.""" + reg = SessionRegistry() + assert reg._neo4j_driver is None + # Must not raise. + await reg.close_neo4j_driver() + assert reg._neo4j_driver is None + + +# --------------------------------------------------------------------------- +# Shutdown quiesce (shutdown_workers) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_shutdown_workers_cancels_and_awaits_every_drainer() -> None: + """shutdown_workers() must cancel AND await every live drain worker. + + This is the ordering guard for the shared driver: a drainer still running + when the shared driver closes fails its batch, exhausts its retry budget, + and dead-letters healthy events (committing the offset past them). + """ + reg = SessionRegistry() + + started = asyncio.Event() + + async def never_ending() -> None: + started.set() + await asyncio.sleep(3600) + + workers = [] + for i in range(3): + worker = SessionWorker( + session_id=f"s{i}", + workspace=f"/ws/{i}", + services=MagicMock(), + ) + worker.task = asyncio.create_task(never_ending()) + reg._workers[worker.session_id] = worker + workers.append(worker) + + await started.wait() + + await reg.shutdown_workers() + + for worker in workers: + assert worker.task is not None + assert worker.task.done(), "shutdown_workers must AWAIT, not just cancel" + assert worker.task.cancelled() + + +@pytest.mark.asyncio +async def test_shutdown_workers_no_op_with_no_workers() -> None: + """shutdown_workers() must be safe when no session ever ran.""" + reg = SessionRegistry() + await reg.shutdown_workers() # must not raise + + +@pytest.mark.asyncio +async def test_shutdown_workers_survives_a_failing_drainer() -> None: + """One worker raising during teardown must not abort the shutdown of the + others -- shutdown is not derailable by a single bad drainer.""" + reg = SessionRegistry() + + async def raises_on_cancel() -> None: + try: + await asyncio.sleep(3600) + except asyncio.CancelledError: + raise RuntimeError("teardown blew up") from None + + async def clean() -> None: + await asyncio.sleep(3600) + + bad = SessionWorker(session_id="bad", workspace="/ws", services=MagicMock()) + bad.task = asyncio.create_task(raises_on_cancel()) + good = SessionWorker(session_id="good", workspace="/ws", services=MagicMock()) + good.task = asyncio.create_task(clean()) + reg._workers["bad"] = bad + reg._workers["good"] = good + await asyncio.sleep(0) + + await reg.shutdown_workers() # must not raise + + assert bad.task.done() + assert good.task.done()