Keep every database a cancelled fan-out opened

`sessions_for` recorded what opened from `gather`'s results, which arrive only
when it runs to completion. Cancelling it discarded them, so a database that
opened while a sibling was still pending was never recorded and `aclose` never
closed it. `return_exceptions=True` covers a failing child, not a cancelled
parent.

`_open` now registers its own session, so what opened is reachable however the
fan-out ends.
This commit is contained in:
Yiorgis Gozadinos 2026-08-27 11:40:24 +03:00
parent b2c617f9f2
commit cfb880e8a0
No known key found for this signature in database
2 changed files with 46 additions and 17 deletions

View file

@ -288,25 +288,22 @@ class FederatedSession:
missing = [name for name in names if name not in self._sessions]
if missing:
opened = await asyncio.gather(
*(self._open(self._refs[name]) for name in missing),
*(self._open(name) for name in missing),
return_exceptions=True,
)
# Whatever opened is tracked before the failure is reported, so
# teardown closes it: `gather` does not cancel the siblings of the
# one that raised, and an untracked connection leaks.
failure: BaseException | None = None
for name, result in zip(missing, opened, strict=True):
for result in opened:
if isinstance(result, BaseException):
failure = failure or result
else:
self._sessions[name] = result
if failure is not None:
raise failure
raise result
return [self._sessions[name] for name in names]
async def _open(self, ref: DatabaseRef) -> SingleDatabaseSession:
async def _open(self, name: str) -> None:
"""Open and register one database before returning to the fan-out.
Registered here because a cancelled `gather` discards its results.
"""
ref = self._refs[name]
one, db_path = ref.connection(self._config)
return await SingleDatabaseSession(
self._sessions[name] = await SingleDatabaseSession(
db_path if db_path is not None else default_db_path(one),
one,
skip_validation=self._skip_validation,

View file

@ -35,11 +35,11 @@ class TestOpeningDatabases:
barrier = asyncio.Barrier(len(names))
open_one = rag._session._open
async def gated(ref):
async def gated(name):
# Every open has to be in flight before any of them finishes, so
# a serial loop cannot get past this and the wait times out.
await barrier.wait()
return await open_one(ref)
await open_one(name)
rag._session._open = gated
clients = await asyncio.wait_for(rag.clients_for(names), timeout=15)
@ -61,6 +61,38 @@ class TestOpeningDatabases:
assert isinstance(rag._session, FederatedSession)
assert set(rag._session._sessions) == {"alpha"}
@pytest.mark.asyncio
async def test_a_cancelled_open_does_not_leak_the_ones_that_worked(self, tmp_path):
"""Cancellation discards the fan-out's results rather than returning them,
so a database that opened while a sibling was still pending is reachable
only because the opener recorded it."""
config = _config(tmp_path, ["alpha", "beta"])
await _seed(config, "alpha", ["alpha document about cats"])
await _seed(config, "beta", ["beta document about cats"])
async with HaikuRAG(config=config) as rag:
assert isinstance(rag._session, FederatedSession)
open_one = rag._session._open
alpha_open = asyncio.Event()
async def staged(name):
if name == "beta":
await asyncio.sleep(60)
await open_one(name)
alpha_open.set()
rag._session._open = staged
fanout = asyncio.create_task(rag.clients_for(["alpha", "beta"]))
await asyncio.wait_for(alpha_open.wait(), timeout=15)
fanout.cancel()
with pytest.raises(asyncio.CancelledError):
await fanout
assert set(rag._session._sessions) == {"alpha"}
alpha = rag._session._sessions["alpha"]
assert not alpha.store.db.is_open()
@pytest.mark.asyncio
async def test_a_database_named_twice_is_opened_once(self, tmp_path):
"""Fusion would count a repeated database as two rank lists."""
@ -392,8 +424,8 @@ class TestDatabaseIndependentWork:
opened: list[str] = []
async def refuse(self, ref):
opened.append(ref.name)
async def refuse(self, name):
opened.append(name)
raise AssertionError("opened a database to chunk a document")
monkeypatch.setattr(FederatedSession, "_open", refuse)