From d68ca9774ff43840141bb804e9da79aa44ee3d2e Mon Sep 17 00:00:00 2001 From: ThearchyHelios Date: Mon, 19 Oct 2026 13:01:21 +0800 Subject: [PATCH] =?UTF-8?q?=E5=8A=A0=E7=9B=B4=E8=BF=9E=E5=9C=B0=E5=9D=80?= =?UTF-8?q?=20AGENTLAB=5FDATABASE=5FDIRECT=5FURL=20=E5=92=8C=20PgBouncer?= =?UTF-8?q?=20=E4=BA=8B=E5=8A=A1=E6=A8=A1=E5=BC=8F=E5=BC=80=E5=85=B3=20AGE?= =?UTF-8?q?NTLAB=5FDB=5FPGBOUNCER=EF=BC=88=E7=94=A8=E6=88=B7=E6=8B=8D?= =?UTF-8?q?=E6=9D=BF=EF=BC=89=EF=BC=9A=E4=BE=9D=E8=B5=96=E4=BC=9A=E8=AF=9D?= =?UTF-8?q?=E7=8A=B6=E6=80=81=E7=9A=84=E8=BF=9E=E6=8E=A5=E2=80=94=E2=80=94?= =?UTF-8?q?=E6=94=B6=E5=8F=91=E8=BF=9B=E7=A8=8B=E9=97=B4=E9=80=9A=E7=9F=A5?= =?UTF-8?q?=E7=9A=84=E9=82=A3=E4=B8=80=E6=9D=A1=EF=BC=88LISTEN=EF=BC=89?= =?UTF-8?q?=E3=80=81=E8=B7=A8=E8=BF=9B=E7=A8=8B=E7=9A=84=E9=94=81=EF=BC=88?= =?UTF-8?q?=E4=BC=9A=E8=AF=9D=E7=BA=A7=E5=92=A8=E8=AF=A2=E9=94=81=EF=BC=8C?= =?UTF-8?q?=E5=8F=A6=E5=BB=BA=E4=B8=80=E4=B8=AA=E7=9B=B4=E8=BF=9E=E7=9A=84?= =?UTF-8?q?=E5=B0=8F=E6=B1=A0=E5=AD=90=EF=BC=89=E3=80=81=E6=96=AD=E7=82=B9?= =?UTF-8?q?=E5=BA=93=EF=BC=88=E7=BC=BA=E7=9C=81=EF=BC=89=E2=80=94=E2=80=94?= =?UTF-8?q?=E8=B5=B0=E7=9B=B4=E8=BF=9E=E5=9C=B0=E5=9D=80=EF=BC=8C=E6=99=AE?= =?UTF-8?q?=E9=80=9A=E6=9F=A5=E8=AF=A2=E8=B5=B0=20AGENTLAB=5FDATABASE=5FUR?= =?UTF-8?q?L=EF=BC=8C=E4=B8=8D=E8=AE=BE=E5=B0=B1=E5=85=A8=E7=94=A8?= =?UTF-8?q?=E5=90=8C=E4=B8=80=E4=B8=AA=EF=BC=9B=E5=BC=80=E5=85=B3=E6=89=93?= =?UTF-8?q?=E5=BC=80=E6=97=B6=E6=99=AE=E9=80=9A=E5=9C=B0=E5=9D=80=E4=B8=8D?= =?UTF-8?q?=E7=94=A8=E5=AE=A2=E6=88=B7=E7=AB=AF=E8=BF=9E=E6=8E=A5=E6=B1=A0?= =?UTF-8?q?=E3=80=81=E5=85=B3=E6=8E=89=20SQLAlchemy=20=E5=92=8C=20asyncpg?= =?UTF-8?q?=20=E7=9A=84=E9=A2=84=E5=A4=87=E8=AF=AD=E5=8F=A5=E7=BC=93?= =?UTF-8?q?=E5=AD=98=E3=80=81=E9=A2=84=E5=A4=87=E8=AF=AD=E5=8F=A5=E7=94=A8?= =?UTF-8?q?=E9=9A=8F=E6=9C=BA=E5=90=8D=E5=AD=97=EF=BC=88=E7=85=A7=20SQLAlc?= =?UTF-8?q?hemy=20=E6=96=87=E6=A1=A3=EF=BC=89=EF=BC=8C=E4=BA=8B=E5=8A=A1?= =?UTF-8?q?=E7=BA=A7=E5=92=A8=E8=AF=A2=E9=94=81=E7=85=A7=E5=B8=B8=E6=88=90?= =?UTF-8?q?=E7=AB=8B=EF=BC=9B=E5=BC=80=E5=85=B3=E6=89=93=E5=BC=80=E5=8D=B4?= =?UTF-8?q?=E6=B2=A1=E7=BB=99=E7=9B=B4=E8=BF=9E=E5=9C=B0=E5=9D=80=E6=97=B6?= =?UTF-8?q?=E6=8B=92=E7=BB=9D=E5=90=AF=E5=8A=A8=E3=80=82README=20=E5=86=99?= =?UTF-8?q?=E4=B8=89=E7=A7=8D=E9=83=A8=E7=BD=B2=EF=BC=88=E7=9B=B4=E8=BF=9E?= =?UTF-8?q?=E3=80=81PgBouncer=20=E4=BC=9A=E8=AF=9D=E6=A8=A1=E5=BC=8F?= =?UTF-8?q?=E3=80=81=E4=BA=8B=E5=8A=A1=E6=A8=A1=E5=BC=8F=E5=8A=A0=E7=9B=B4?= =?UTF-8?q?=E8=BF=9E=E5=9C=B0=E5=9D=80=EF=BC=89=E5=92=8C=E8=A6=81=E9=85=8D?= =?UTF-8?q?=E7=9A=84=20server=5Freset=5Fquery=EF=BC=8C=E7=85=A7=E5=AE=9E?= =?UTF-8?q?=E8=AF=B4=E6=B2=A1=E5=9C=A8=E7=9C=9F=E7=9A=84=20PgBouncer=20?= =?UTF-8?q?=E4=B8=8A=E8=B7=91=E8=BF=87=E3=80=82=E8=A1=A5=E6=B5=8B=E8=AF=95?= =?UTF-8?q?=EF=BC=9A=E5=9C=B0=E5=9D=80=E8=A7=A3=E6=9E=90=E5=92=8C=E8=A7=84?= =?UTF-8?q?=E8=8C=83=E5=8C=96=E3=80=81=E6=96=AD=E7=82=B9=E5=BA=93=E7=BC=BA?= =?UTF-8?q?=E7=9C=81=E8=B5=B0=E7=9B=B4=E8=BF=9E=E3=80=81=E4=BA=8B=E5=8A=A1?= =?UTF-8?q?=E6=A8=A1=E5=BC=8F=E5=85=B3=E6=8E=89=E7=9A=84=E5=87=A0=E6=A0=B7?= =?UTF-8?q?=E3=80=81LISTEN=20=E5=92=8C=E8=B7=A8=E8=BF=9B=E7=A8=8B=E9=94=81?= =?UTF-8?q?=E8=B5=B0=E7=9B=B4=E8=BF=9E=E3=80=81=E4=BA=8B=E5=8A=A1=E6=A8=A1?= =?UTF-8?q?=E5=BC=8F=E6=B2=A1=E7=BB=99=E7=9B=B4=E8=BF=9E=E5=9C=B0=E5=9D=80?= =?UTF-8?q?=E6=8B=92=E7=BB=9D=E5=90=AF=E5=8A=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 25 ++++++++--- backend/app/cluster/mutex.py | 10 +++-- backend/app/cluster/signal.py | 11 ++++- backend/app/core/config.py | 21 +++++++-- backend/app/db/base.py | 76 ++++++++++++++++++++++++++++----- backend/app/main.py | 7 ++- backend/tests/test_db_direct.py | 69 ++++++++++++++++++++++++++++++ 7 files changed, 193 insertions(+), 26 deletions(-) create mode 100644 backend/tests/test_db_direct.py diff --git a/README.md b/README.md index c31a10c3..8bec1a43 100644 --- a/README.md +++ b/README.md @@ -1050,9 +1050,8 @@ AGENTLAB_DATABASE_URL=postgresql+asyncpg://agentlab:<密码>@127.0.0.1:5432/agen - 断点（LangGraph 的 checkpoints、checkpoint_blobs、checkpoint_writes、checkpoint_migrations 四张表，不和应用的表重名） 缺省放在同一个库；要放别处设 `AGENTLAB_CHECKPOINT_URL=postgresql://…`（这一项是 libpq 的写法，`sslrootcert` 等可以直接写）。 - 连接池：每个进程常驻 `AGENTLAB_DB_POOL_SIZE`（缺省 10）条、高峰再多开 `AGENTLAB_DB_MAX_OVERFLOW`（缺省 20）条，取连接前探活， - 30 分钟换一次；断点库另有一个小池子。前面有 PgBouncer（事务模式）时：`AGENTLAB_DB_POOL_SIZE` 设成 0（不用连接池）、 - 地址后加 `?prepared_statement_cache_size=0`；断点库用预备语句，要 PgBouncer 1.21 以上开 `max_prepared_statements`， - 或者用 `AGENTLAB_CHECKPOINT_URL` 让它直连 PostgreSQL。 + 30 分钟换一次；断点库另有一个小池子。前面有 PgBouncer 时照下面「直连与 PgBouncer」的三种部署写（事务模式要开 + `AGENTLAB_DB_PGBOUNCER=1` 并另给直连地址）。 - 连接数的账：每个进程最多 `AGENTLAB_DB_POOL_SIZE + AGENTLAB_DB_MAX_OVERFLOW`（缺省 10 + 20 = 30）条应用库连接，断点库的池子 最多 `min(AGENTLAB_MAX_CONCURRENT_RUNS, 20)` 条（至少 4，缺省 20），合起来一个进程最多 50 条；导入命令另占几条。PostgreSQL 缺省 `max_connections=100`，还要留几条给 `psql`、备份和超级用户（`superuser_reserved_connections` 缺省 3）：一个进程够用， @@ -1127,8 +1126,24 @@ AGENTLAB_WORKERS=4 … | `AGENTLAB_SCHEDULER_TICK_S` | 15 | 调度循环一拍的间隔 | 连接数照上一节「连接数的账」按进程数乘，另外每个进程多一条专门收发进程间通知（LISTEN / NOTIFY）的连接，不占连接池；跨进程的锁 -（咨询锁）持锁期间占连接池里的一条。前面放 PgBouncer（事务模式）时，通知那一条要直连 PostgreSQL（事务模式下 LISTEN 收不到通知）， -会话级的咨询锁也要求直连或会话模式——多进程部署建议应用库直连。 +（会话级咨询锁）持锁期间占一条。 + +**直连与 PgBouncer。** 有几样连接依赖会话状态：收发进程间通知的那一条（LISTEN）、跨进程的锁（会话级咨询锁，取锁和放锁要在同一条 +会话上）、断点库（预备语句，建表时取会话锁）。普通查询走 `AGENTLAB_DATABASE_URL`；设了 `AGENTLAB_DATABASE_DIRECT_URL`（写法同它） +时，上面这几样走直连地址（断点库缺省也走它，`AGENTLAB_CHECKPOINT_URL` 另给的除外）；不设就全用同一个地址。三种部署： + +| 部署 | 怎么设 | +| --- | --- | +| 直连 PostgreSQL | 只设 `AGENTLAB_DATABASE_URL`。最简单，多进程建议这样 | +| PgBouncer 会话模式（`pool_mode=session`） | `AGENTLAB_DATABASE_URL` 指 PgBouncer 就行：会话模式下一个客户端连接独占一条服务端连接，会话状态都在 | +| PgBouncer 事务模式（`pool_mode=transaction`） | `AGENTLAB_DATABASE_URL` 指 PgBouncer，另设 `AGENTLAB_DB_PGBOUNCER=1` 和 `AGENTLAB_DATABASE_DIRECT_URL`（直连 PostgreSQL，或指 PgBouncer 的会话模式那个库）。没给直连地址时拒绝启动 | + +`AGENTLAB_DB_PGBOUNCER=1` 做的事（只对普通地址）：不用客户端的连接池（交给 PgBouncer）；关掉 SQLAlchemy 和 asyncpg 的预备语句缓存 +（事务模式下下一条语句可能落在另一条服务端连接上，缓存着的预备语句在那边不存在）；asyncpg 每条语句仍先预备再执行（SQLAlchemy 的 +asyncpg 方言一律这样做，同一个事务里、在同一条服务端连接上），名字用随机的，不和别的客户端留下的撞名——PgBouncer 那边要配 +`server_reset_query = DISCARD ALL` 加 `server_reset_query_always = 1` 收掉留下的预备语句，或用 1.21 以上开 `max_prepared_statements`。 +事务级的咨询锁（发名额、同源队、建表时用）在事务模式下照常成立：取锁、放锁在同一个事务里。这几条照 SQLAlchemy、asyncpg、psycopg +的文档和代码核过，没有在真的 PgBouncer 上跑过（本机没有它的镜像）。 **认领和租约。** 每个进程在库里登记一行（主机、pid、启动时的随机串），隔一拍心跳续租约。运行、问数据的一轮、后台任务（探查、起草、 剖析、评测、影子起草）、处理中的文档、执行中的发布单都记着认领它的进程；租约按进程算，进程没了它认领的行一起到期。平稳关停的进程 diff --git a/backend/app/cluster/mutex.py b/backend/app/cluster/mutex.py index 2de45376..96214ba3 100644 --- a/backend/app/cluster/mutex.py +++ b/backend/app/cluster/mutex.py @@ -55,9 +55,9 @@ class ClusterLock: try: from sqlalchemy import text - from app.db.base import engine + from app.db.base import session_engine - conn = await engine.connect() + conn = await session_engine().connect() try: conn = await conn.execution_options(isolation_level="AUTOCOMMIT") got = bool(await conn.scalar(text("SELECT pg_try_advisory_lock(:k)"), {"k": self.key})) @@ -86,9 +86,11 @@ class ClusterLock: try: from sqlalchemy import text - from app.db.base import engine + from app.db.base import session_engine - conn = await engine.connect() + # 会话级的锁：走直连（AGENTLAB_DATABASE_DIRECT_URL，没设就是应用库）——PgBouncer 事务模式下取锁、放锁可能落在 + # 两条不同的服务端连接上 + conn = await session_engine().connect() try: conn = await conn.execution_options(isolation_level="AUTOCOMMIT") await conn.execute(text("SELECT pg_advisory_lock(:k)"), {"k": self.key}) diff --git a/backend/app/cluster/signal.py b/backend/app/cluster/signal.py index a2592d97..46785494 100644 --- a/backend/app/cluster/signal.py +++ b/backend/app/cluster/signal.py @@ -37,6 +37,14 @@ RETRY_S = 2.0 Handler = Callable[[list[str]], Any] +def listen_dsn() -> str: + """收发通知那条专用连接的地址：LISTEN 要一条一直开着的会话，设了直连地址（AGENTLAB_DATABASE_DIRECT_URL）就走它——普通地址 + 可能是 PgBouncer 的事务模式，那里收不到通知。""" + from app.db.urls import libpq_conninfo + + return libpq_conninfo(settings.session_db_url) + + def _tag() -> str: return identity.owner().rsplit(":", 1)[-1] @@ -126,9 +134,8 @@ class Signals: import asyncpg from app.db.base import APPLICATION_NAME - from app.db.urls import libpq_conninfo - dsn = libpq_conninfo(settings.database_url) + dsn = listen_dsn() first = True while not self._stopping: lost = asyncio.Event() diff --git a/backend/app/core/config.py b/backend/app/core/config.py index bdb311f7..bd91811d 100644 --- a/backend/app/core/config.py +++ b/backend/app/core/config.py @@ -56,6 +56,13 @@ class Settings(BaseSettings): # LangGraph 的断点也放进同一个库；checkpoint_url 另给就放那里。数据目录照旧要有：工件、上传表格、沙箱等仍是文件 database_url: str = "" checkpoint_url: str = "" + # 直连地址（可选，写法同 database_url）：普通查询走 database_url（可以是 PgBouncer），依赖会话状态的连接走这里——进程间通知的 + # LISTEN 那一条、跨进程锁（会话级咨询锁）、断点库（预备语句、建表时的会话锁）。不设就全用 database_url + database_direct_url: str = "" + # database_url 是 PgBouncer 的事务模式（AGENTLAB_DB_PGBOUNCER=1）：应用库不用客户端连接池（交给 PgBouncer）、关掉 asyncpg 和 + # SQLAlchemy 的预备语句缓存、预备语句用不会撞名的名字。这时必须另给 database_direct_url（事务模式下 LISTEN、会话级咨询锁不成立）， + # 启动时检查 + db_pgbouncer: bool = False # PostgreSQL 的连接池（SQLite 不用这几项）。pool_size=0 表示不用连接池、每次新建连接（前面有 PgBouncer 时用， # 测试里也这么用：每个测试一个事件循环，asyncpg 的连接不能跨循环） db_pool_size: int = 10 @@ -221,7 +228,7 @@ class Settings(BaseSettings): return True return value - @field_validator("database_url", "checkpoint_url", mode="before") + @field_validator("database_url", "checkpoint_url", "database_direct_url", mode="before") @classmethod def _valid_postgres_url(cls, value: object, info: Any) -> object: # 写错了启动时就报出来，别等到第一次查库。只认 PostgreSQL：SQLite 就是不设 @@ -232,7 +239,7 @@ class Settings(BaseSettings): from app.db.urls import normalize_database_url, normalize_checkpoint_url name = f"AGENTLAB_{(info.field_name or '').upper()}" - normalize = normalize_database_url if info.field_name == "database_url" else normalize_checkpoint_url + normalize = normalize_checkpoint_url if info.field_name == "checkpoint_url" else normalize_database_url try: return normalize(value.strip()) except ValueError as e: @@ -293,15 +300,21 @@ class Settings(BaseSettings): """LangGraph checkpointer 用独立的库文件，避免和业务表抢写锁。只在 SQLite 时用。""" return self.data_dir / "checkpoints.db" + @property + def session_db_url(self) -> str: + """依赖会话状态的连接（LISTEN、会话级咨询锁、断点库）用的应用库地址：设了直连地址用它，否则同 database_url。""" + return self.database_direct_url or self.database_url + @property def checkpoint_conninfo(self) -> str | None: - """PostgreSQL 断点库的 libpq 连接串（psycopg 用）：checkpoint_url 优先，否则由 database_url 换算。SQLite 时 None。""" + """PostgreSQL 断点库的 libpq 连接串（psycopg 用）：checkpoint_url 优先，否则由直连地址（没设就是 database_url）换算。 + SQLite 时 None。断点库用预备语句、建表时取会话级的锁，所以缺省走直连。""" if self.checkpoint_url: return self.checkpoint_url if self.database_url: from app.db.urls import libpq_conninfo - return libpq_conninfo(self.database_url) + return libpq_conninfo(self.session_db_url) return None @property diff --git a/backend/app/db/base.py b/backend/app/db/base.py index ee3edaca..83d0b117 100644 --- a/backend/app/db/base.py +++ b/backend/app/db/base.py @@ -188,12 +188,45 @@ def pg_json_dumps(value: Any) -> str: APPLICATION_NAME = "agentlab" -def make_engine(url: str | None = None, *, application_name: str = APPLICATION_NAME) -> AsyncEngine: +def pg_engine_kwargs(*, application_name: str = APPLICATION_NAME, pgbouncer: bool = False, + pool_size: int | None = None) -> dict[str, Any]: + """PostgreSQL 引擎的连接池和连接参数。pgbouncer：这个地址是 PgBouncer 的事务模式（AGENTLAB_DB_PGBOUNCER）—— + + - 不用客户端的连接池（NullPool，连接池交给 PgBouncer；SQLAlchemy 的说明里同样这么建议）； + - SQLAlchemy 的预备语句缓存（prepared_statement_cache_size）和 asyncpg 自己的（statement_cache_size）都关掉：事务模式下同一个 + 客户端连接的下一条语句可能落在另一条服务端连接上，缓存着的预备语句在那边不存在； + - asyncpg 每条语句仍要先预备再执行（SQLAlchemy 的 asyncpg 方言一律这么做，同一个事务里，落在同一条服务端连接上），名字用随机的 + （prepared_statement_name_func），不和别的客户端在同一条服务端连接上留下的撞名。PgBouncer 要配 server_reset_query=DISCARD ALL + （且 server_reset_query_always=1）收掉留下的预备语句，或用 1.21 以上开 max_prepared_statements。 + 事务级的咨询锁（pg_advisory_xact_lock，发名额、同源队、建表时用）在事务模式下照常成立：取锁和放锁都在同一个事务里。""" + connect_args: dict[str, Any] = { + # 会话时区定成 UTC：时间列是不带时区的 UTC（UTCDateTime），库里 now() 之类的也按 UTC 算 + "server_settings": {"application_name": application_name, "timezone": "UTC"}} + size = settings.db_pool_size if pool_size is None else pool_size + if pgbouncer: + connect_args.update(statement_cache_size=0, prepared_statement_cache_size=0, + prepared_statement_name_func=_unique_statement_name) + return {"poolclass": NullPool, "connect_args": connect_args} + if size <= 0: + return {"poolclass": NullPool, "connect_args": connect_args} + return {"pool_size": size, "max_overflow": max(settings.db_max_overflow, 0), "pool_timeout": settings.db_pool_timeout, + "pool_recycle": settings.db_pool_recycle, "connect_args": connect_args} + + +def _unique_statement_name() -> str: + return f"__agentlab_{uuid.uuid4().hex}__" + + +def make_engine(url: str | None = None, *, application_name: str = APPLICATION_NAME, pgbouncer: bool | None = None, + pool_size: int | None = None) -> AsyncEngine: """应用库的引擎。SQLite：WAL、外键、忙等，单个文件；PostgreSQL：连接池、方言调整（db/postgres.py）。 copy-to-postgres 建目标库的引擎也走这里，两边的列类型、排序规则是同一套。 """ url = url or settings.db_url + if pgbouncer is None: + # 只有应用库的普通地址按 AGENTLAB_DB_PGBOUNCER 处理：直连地址、copy-to-postgres 给的目标库照常 + pgbouncer = bool(settings.db_pgbouncer) and url == settings.database_url if url.startswith("sqlite"): # pool_pre_ping：取连接前先探一下活。取消异常打进 SQLAlchemy 的清理流程时 # （_close_connection 记完日志会把 CancelledError 原样抛出，跳过了把连接记录 @@ -206,15 +239,9 @@ def make_engine(url: str | None = None, *, application_name: str = APPLICATION_N from app.db import postgres _require_postgres_driver("asyncpg") - pool: dict[str, Any] = ( - {"poolclass": NullPool} if settings.db_pool_size <= 0 else { - "pool_size": settings.db_pool_size, "max_overflow": max(settings.db_max_overflow, 0), - "pool_timeout": settings.db_pool_timeout, "pool_recycle": settings.db_pool_recycle}) - made = create_async_engine( - url, echo=False, future=True, pool_pre_ping=True, json_serializer=pg_json_dumps, - # 会话时区定成 UTC：时间列是不带时区的 UTC（UTCDateTime），库里 now() 之类的也按 UTC 算 - connect_args={"server_settings": {"application_name": application_name, "timezone": "UTC"}}, - **pool) + made = create_async_engine(url, echo=False, future=True, pool_pre_ping=True, json_serializer=pg_json_dumps, + **pg_engine_kwargs(application_name=application_name, pgbouncer=pgbouncer, + pool_size=pool_size)) postgres.install(made.sync_engine.dialect) event.listen(made.sync_engine, "before_cursor_execute", postgres.strip_nul, retval=True) return made @@ -250,6 +277,35 @@ engine = make_engine() SessionLocal = async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession) +_session_engine: AsyncEngine | None = None + + +def session_engine() -> AsyncEngine: + """依赖会话状态的连接用的引擎（跨进程锁的会话级咨询锁，app/cluster/mutex）：设了 AGENTLAB_DATABASE_DIRECT_URL 时是直连那个库的 + 一个小池子，否则就是应用库的引擎。""" + global _session_engine + if not settings.database_direct_url or settings.database_direct_url == settings.database_url: + return engine + if _session_engine is None: + _session_engine = make_engine(settings.database_direct_url, pgbouncer=False, pool_size=max(2, min(settings.db_pool_size or 5, 10))) + return _session_engine + + +async def dispose_session_engine() -> None: + """关停时关掉直连那个小池子（没另建就什么都不做）。""" + global _session_engine + if _session_engine is not None: + await _session_engine.dispose() + _session_engine = None + + +def check_pgbouncer() -> None: + """AGENTLAB_DB_PGBOUNCER=1 却没给直连地址：拒绝启动（事务模式下 LISTEN、会话级咨询锁不成立，跨进程的锁会锁不住、放不掉）。""" + if settings.db_pgbouncer and settings.db_dialect == "postgresql" and not settings.database_direct_url: + raise RuntimeError("AGENTLAB_DB_PGBOUNCER=1（应用库地址是 PgBouncer 的事务模式）时要另给 AGENTLAB_DATABASE_DIRECT_URL：" + "进程间通知（LISTEN）、跨进程的锁（会话级咨询锁）、断点库要直连 PostgreSQL（或 PgBouncer 的会话模式）。" + "见 README「用 PostgreSQL」里的三种部署") + async def get_session() -> AsyncIterator[AsyncSession]: """FastAPI 依赖注入用的会话。""" diff --git a/backend/app/main.py b/backend/app/main.py index b926852f..c558ae52 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -58,6 +58,10 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]: # SQLite 只支持单进程：同一个数据目录已经有一个进程在用（uvicorn --workers、两个服务指着同一个目录）就拒绝启动、照实说 # （app/cluster/guard.py）。PostgreSQL 可以多进程，什么都不做 guard.acquire() + # PgBouncer 事务模式（AGENTLAB_DB_PGBOUNCER=1）要另给直连地址，没给就拒绝启动（db/base.check_pgbouncer） + from app.db.base import check_pgbouncer + + check_pgbouncer() await init_db() # 多进程协调（app/cluster）：登记这个工作进程、开始心跳和进程间通知、争调度当班。要在下面的启动整理之前：整理只碰属主不在了的 # 行，要先知道别的进程谁还活着 @@ -150,9 +154,10 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]: await cluster.stop() # 应用库的连接全关掉：SQLite 最后一个连接关掉时才把 -wal 并回、删掉（copy-to-postgres 按它认服务停没停）， # PostgreSQL 的连接池也正常关 - from app.db.base import engine as db_engine + from app.db.base import dispose_session_engine, engine as db_engine await db_engine.dispose() + await dispose_session_engine() logger.info("已停止") diff --git a/backend/tests/test_db_direct.py b/backend/tests/test_db_direct.py new file mode 100644 index 00000000..552a0c40 --- /dev/null +++ b/backend/tests/test_db_direct.py @@ -0,0 +1,69 @@ +"""直连地址和 PgBouncer 事务模式（AGENTLAB_DATABASE_DIRECT_URL、AGENTLAB_DB_PGBOUNCER）：地址怎么解析、哪些连接走直连、事务模式下 +关掉了哪些东西。只建引擎、不连库，两种库都跑。""" +from __future__ import annotations + +import pytest +from sqlalchemy.pool import NullPool + +from app.core.config import Settings, settings + +APP = "postgres://agentlab:secret@pgbouncer.internal:6432/agentlab" +DIRECT = "postgresql://agentlab:secret@db.internal:5432/agentlab" + + +def test_the_direct_url_is_normalized_and_used_for_session_bound_connections(): + s = Settings(database_url=APP, database_direct_url=DIRECT, _env_file=None) + assert s.database_url.startswith("postgresql+asyncpg://") and "pgbouncer.internal:6432" in s.database_url + assert s.database_direct_url.startswith("postgresql+asyncpg://") and "db.internal:5432" in s.database_direct_url + assert s.session_db_url == s.database_direct_url + assert "db.internal" in s.checkpoint_conninfo and s.checkpoint_conninfo.startswith("postgresql://"), "断点库缺省走直连" + plain = Settings(database_url=APP, _env_file=None) + assert plain.session_db_url == plain.database_url and "pgbouncer.internal" in plain.checkpoint_conninfo, "不设就全用同一个" + with pytest.raises(ValueError, match="AGENTLAB_DATABASE_DIRECT_URL"): + Settings(database_url=APP, database_direct_url="mysql://x@y/z", _env_file=None) + + +def test_transaction_mode_turns_off_client_pooling_and_statement_caches(): + from app.db.base import pg_engine_kwargs + + kw = pg_engine_kwargs(pgbouncer=True) + assert kw["poolclass"] is NullPool + args = kw["connect_args"] + assert args["statement_cache_size"] == 0 and args["prepared_statement_cache_size"] == 0 + names = {args["prepared_statement_name_func"]() for _ in range(3)} + assert len(names) == 3, "预备语句的名字每次都不一样，不和别的客户端撞名" + direct = pg_engine_kwargs(pgbouncer=False, pool_size=5) + assert direct["pool_size"] == 5 and "statement_cache_size" not in direct["connect_args"] + + +def test_listen_and_cluster_locks_go_direct(monkeypatch): + from app.cluster.signal import listen_dsn + from app.db import base + + monkeypatch.setattr(settings, "database_url", Settings(database_url=APP, _env_file=None).database_url) + monkeypatch.setattr(settings, "database_direct_url", Settings(database_url=APP, database_direct_url=DIRECT, + _env_file=None).database_direct_url) + monkeypatch.setattr(base, "_session_engine", None) + assert "db.internal:5432" in listen_dsn() and listen_dsn().startswith("postgresql://") + made = base.session_engine() + try: + assert made is not base.engine and made.url.host == "db.internal" and made.url.port == 5432 + assert base.session_engine() is made, "直连那个池子只建一个" + finally: + monkeypatch.setattr(base, "_session_engine", None) + made.sync_engine.dispose() + monkeypatch.setattr(settings, "database_direct_url", "") + assert base.session_engine() is base.engine and "pgbouncer.internal" in listen_dsn() + + +def test_transaction_mode_without_a_direct_url_refuses_to_start(monkeypatch): + from app.db import base + + monkeypatch.setattr(settings, "database_url", Settings(database_url=APP, _env_file=None).database_url) + monkeypatch.setattr(settings, "db_pgbouncer", True) + monkeypatch.setattr(settings, "database_direct_url", "") + with pytest.raises(RuntimeError, match="AGENTLAB_DATABASE_DIRECT_URL"): + base.check_pgbouncer() + monkeypatch.setattr(settings, "database_direct_url", Settings(database_url=APP, database_direct_url=DIRECT, + _env_file=None).database_direct_url) + base.check_pgbouncer()