diff --git a/backend/app/db/base.py b/backend/app/db/base.py index b8c1592c..fe02274c 100644 --- a/backend/app/db/base.py +++ b/backend/app/db/base.py @@ -40,7 +40,45 @@ class TimestampMixin: ) -engine = create_async_engine(settings.db_url, echo=False, future=True) +def _stop_aiosqlite_once() -> None: + """aiosqlite 连接的工作线程只停一次;之后再要求停的,拿到的是同一个结果。 + + aiosqlite 0.22 的 stop() 每调一次都往工作线程队列里放一个"关闭并退出线程" + 的请求,而线程处理完第一个就退出了——之后放进去的请求没人处理,等它的一方 + 永远挂着。close() 的 finally 里会调 stop(),SQLAlchemy 的强制 terminate 也会 + 直接调;两个 close 并发、或者 close 撞上强制 terminate,就会卡死。实测同一 + 连接并发 close 两次,10 次卡 10 次。 + + 在这里触发它的路径:并行的两条支路一条失败,LangGraph 取消另一条,而那一条 + 正在写工件。取消打进进行中的数据库操作,SQLAlchemy 在两条路径上各 terminate + 一次同一条连接——节点任务取消不完,整个运行卡在 running,占着的并发名额也不 + 还。大约五次里有一次。连接池和 checkpointer 都走 aiosqlite,补在这一层都管得到。 + """ + import weakref + + import aiosqlite + + original = aiosqlite.Connection.stop + if getattr(original, "_agentlab_once", False): + return + stopping: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary() + + def stop(self): # type: ignore[no-untyped-def] + if self not in stopping: + stopping[self] = original(self) + return stopping[self] + + stop._agentlab_once = True # type: ignore[attr-defined] + aiosqlite.Connection.stop = stop # type: ignore[method-assign] + + +_stop_aiosqlite_once() +# pool_pre_ping:取连接前先探一下活。取消异常打进 SQLAlchemy 的清理流程时 +# (_close_connection 记完日志会把 CancelledError 原样抛出,跳过了把连接记录 +# 清空的那一行),已经关掉的连接会留在池里,下一个取到它的人拿到的是 +# "no active connection"。探活失败会被当成断开,换一条新的——SQLite 本地一次 +# SELECT 1 的代价,换池子不被污染。 +engine = create_async_engine(settings.db_url, echo=False, future=True, pool_pre_ping=True) @event.listens_for(engine.sync_engine, "connect") diff --git a/backend/tests/test_aiosqlite_close.py b/backend/tests/test_aiosqlite_close.py new file mode 100644 index 00000000..3ca1b9ca --- /dev/null +++ b/backend/tests/test_aiosqlite_close.py @@ -0,0 +1,48 @@ +"""aiosqlite 连接的关闭不能把调用方永远挂住。 + +aiosqlite 0.22 的 stop() 每调一次都往工作线程队列里放一个"关闭并退出线程" +的请求:线程处理完第一个就退出,之后的请求没人处理,等它的一方永远挂着。 +close() 的 finally 会调 stop(),SQLAlchemy 的强制 terminate 也会直接调—— +不打补丁时同一连接并发 close 两次,10 次卡 10 次。 + +它在这个项目里的真实后果是运行卡死:并行支路一条失败,LangGraph 取消另一条, +那一条正在写工件,SQLAlchemy 就在两条路径上各 terminate 一次同一条连接—— +节点任务取消不完,运行永远停在 running(test_approval 里并行驳回那例约五次 +卡一次,就是这么被发现的)。补丁在 app.db.base。 +""" +from __future__ import annotations + +import asyncio + +import aiosqlite + +import app.db.base # noqa: F401 补丁在这里打上 + + +async def test_concurrent_close_of_one_connection_finishes(tmp_path): + conn = await aiosqlite.connect(tmp_path / "x.db") + await conn.execute("CREATE TABLE t (x INTEGER)") + await asyncio.wait_for(asyncio.gather(conn.close(), conn.close(), conn.close()), timeout=5) + + +async def test_close_racing_a_forced_stop_finishes(tmp_path): + """SQLAlchemy 的优雅关闭被取消时会退回强制 terminate,直接调 stop()。""" + conn = await aiosqlite.connect(tmp_path / "w.db") + closing = asyncio.ensure_future(conn.close()) + await asyncio.sleep(0) + conn.stop() + await asyncio.wait_for(closing, timeout=5) + + +async def test_close_after_a_cancelled_close_finishes(tmp_path): + conn = await aiosqlite.connect(tmp_path / "y.db") + first = asyncio.ensure_future(conn.close()) + await asyncio.sleep(0) + first.cancel() + await asyncio.wait_for(conn.close(), timeout=5) + + +async def test_closing_an_already_closed_connection_is_a_no_op(tmp_path): + conn = await aiosqlite.connect(tmp_path / "z.db") + await conn.close() + await asyncio.wait_for(conn.close(), timeout=5)