diff --git a/cloud/db.py b/cloud/db.py index 2d94e74..c9b4e76 100644 --- a/cloud/db.py +++ b/cloud/db.py @@ -1165,6 +1165,49 @@ POSTGRES_SCHEMA = ( ) +# 启动期 schema 初始化 advisory lock key（Postgres）： +# 并发启动（cloud-server 与 worker 同时 init_db）时，CREATE TABLE / +# CREATE INDEX 的 IF NOT EXISTS 检查与建对象不是原子的——两个会话 +# 都查到“不存在”后同时建，会撞 pg_type_typname_nsp_index（建表） +# 或 pg_class_relname_nsp_index（建索引）的 UniqueViolation，启动 +# 即崩、靠 systemd 重启才自愈（37ba687 与 7983567 两次部署均撞上； +# 7983567 新增的 idx_tasks_parent 是当时的新对象）。 +# 持锁串行化整个初始化：init_postgres_schema 内持事务级锁（随本 +# 函数末尾的 commit 释放），server.init_db 另持会话级锁覆盖其后 +# 的 _ensure_*（task_continue_messages / user_gateway_keys / +# task_attempts / task_audit_loop_state 不在 POSTGRES_SCHEMA 里，同样 +# 会撞）。key 须全项目唯一：限流清理另有 +# RATE_LIMIT_SWEEP_ADVISORY_KEY = 20261027（cloud/server.py）， +# 两处绝不能撞 key。 +SCHEMA_INIT_ADVISORY_KEY = 20261108 + + +def acquire_schema_init_lock(conn): + """取 schema 初始化会话级 advisory 锁（阻塞等，不 try）。 + + 调用方须在初始化全部完成后调 release_schema_init_lock 释放； + 连接关闭时也会自动释放。初始化必须做完不能跳过，所以这里用 + 阻塞取锁，不用 pg_try_advisory_lock。 + """ + conn.execute( + "SELECT pg_advisory_lock(?)", + (SCHEMA_INIT_ADVISORY_KEY,)) + + +def release_schema_init_lock(conn): + """放 schema 初始化会话级 advisory 锁（best-effort）。 + + 事务已 abort 时 SELECT 本身也会失败——吞掉，连接关闭时锁 + 同样释放，绝不能让放锁失败再挡住启动。 + """ + try: + conn.execute( + "SELECT pg_advisory_unlock(?)", + (SCHEMA_INIT_ADVISORY_KEY,)) + except Exception: + pass + + def init_postgres_schema(conn): """在给定适配连接上建出全部 Postgres 表（幂等），并提交。 @@ -1178,6 +1221,15 @@ def init_postgres_schema(conn): 库是 no-op。server.init_db 的 Postgres 分支还会再跑一次 migrate_projects_table，两次补列互不冲突（第二次是 no-op）。 """ + # 事务级 advisory 锁：并发 init_postgres_schema 串行执行， + # 随本函数末尾的 commit 释放；独立调用方（如 migrations） + # 同样受保护。server.init_db 另持会话级锁覆盖 _ensure_*。 + # 不取结果行：pg_advisory_xact_lock 返回 void，且假连接 + # （test_p2_migrations / test_audit_batch1 的 FakeCursor） + # 没有 fetchone——execute 本身已阻塞到取锁成功。 + conn.execute( + "SELECT pg_advisory_xact_lock(?)", + (SCHEMA_INIT_ADVISORY_KEY,)) for statement in POSTGRES_SCHEMA: conn.execute(statement) cols = {row[1] for row in conn.execute( @@ -1379,6 +1431,9 @@ __all__ = [ "PgCursor", "PgRow", "POSTGRES_SCHEMA", + "SCHEMA_INIT_ADVISORY_KEY", + "acquire_schema_init_lock", + "release_schema_init_lock", "init_postgres_schema", "create_database", "database_from_env", diff --git a/cloud/server.py b/cloud/server.py index 79692c2..100d673 100644 --- a/cloud/server.py +++ b/cloud/server.py @@ -88,7 +88,9 @@ from cloud.timeouts import validate_server_timeout as _validate_timeout from cloud.auth import audit_on_conn as _audit_on_conn from cloud.db import connect as _db_connect from cloud.db import create_database as _create_database +from cloud.db import acquire_schema_init_lock as _acquire_schema_init_lock from cloud.db import init_postgres_schema as _init_postgres_schema +from cloud.db import release_schema_init_lock as _release_schema_init_lock from cloud.db import PostgresIntegrityError as _PgIntegrityError from cloud.db import PostgresOperationalError as _PgOperationalError from cloud.goal_command import parse_goal_command as _parse_goal_command @@ -1667,22 +1669,32 @@ def init_db(db_path): conn = _connect(db_path) try: if _is_postgres_backend(): - _init_postgres_schema(conn) - migrate_tasks_table(conn) - migrate_projects_table(conn) - _ensure_followups_table(conn) - _ensure_continue_messages_table(conn) - _ensure_environments_table(conn) - _ensure_quotas_table(conn) - _backfill_default_quotas(conn) - _ensure_gateway_keys(conn) - _ensure_uploads_tables(conn) - _ensure_task_checkpoints_table(conn) - _ensure_task_steps_table(conn) - _ensure_attempts(conn) - _ensure_audit_loop(conn) - conn.commit() - return + # 会话级 advisory 锁覆盖整个 PG 初始化：_init_postgres_schema + # 内还有事务级锁，但其后的 _ensure_* 建的对象不全在 + # POSTGRES_SCHEMA 里（task_continue_messages / + # user_gateway_keys / task_attempts / task_audit_loop_state）， + # 只锁前者仍会在后者上撞 UniqueViolation；锁在最终 commit + # 后才放。连接在外层 finally 里关闭，锁至迟随连接释放。 + _acquire_schema_init_lock(conn) + try: + _init_postgres_schema(conn) + migrate_tasks_table(conn) + migrate_projects_table(conn) + _ensure_followups_table(conn) + _ensure_continue_messages_table(conn) + _ensure_environments_table(conn) + _ensure_quotas_table(conn) + _backfill_default_quotas(conn) + _ensure_gateway_keys(conn) + _ensure_uploads_tables(conn) + _ensure_task_checkpoints_table(conn) + _ensure_task_steps_table(conn) + _ensure_attempts(conn) + _ensure_audit_loop(conn) + conn.commit() + return + finally: + _release_schema_init_lock(conn) conn.execute( "CREATE TABLE IF NOT EXISTS projects(" "id TEXT PRIMARY KEY, " diff --git a/cloud/tests/test_pg_init_fix.py b/cloud/tests/test_pg_init_fix.py index 27157cf..0f8d2f1 100644 --- a/cloud/tests/test_pg_init_fix.py +++ b/cloud/tests/test_pg_init_fix.py @@ -28,7 +28,9 @@ import sqlite3 import sys import tempfile import time +import traceback import unittest +from multiprocessing import Barrier, Process, Queue from pathlib import Path from types import SimpleNamespace from unittest import mock @@ -209,6 +211,18 @@ class SqliteUpsertTestCase(unittest.TestCase): conn.close() +def _concurrent_init_db_worker(barrier, queue, wid, env): + """并发启动子进程：与其它进程同时跑 server.init_db。""" + try: + os.environ.update(env) + barrier.wait(timeout=60) + from cloud.server import init_db + init_db(":memory:") + queue.put((wid, "ok", "")) + except Exception: + queue.put((wid, "FAIL", traceback.format_exc())) + + class PgInitFixTestCase(unittest.TestCase): """真实 PG 测试：无 DSN/驱动时全部 skip。""" @@ -331,6 +345,67 @@ class PgInitFixTestCase(unittest.TestCase): finally: conn.close() + def test_concurrent_init_db_no_unique_violation(self): + """并发启动回归：4 进程同时 init_db，须全部成功。 + + 覆盖整条启动路径（init_postgres_schema + _ensure_*）： + _ensure_* 建的 task_continue_messages / user_gateway_keys / + task_attempts / task_audit_loop_state 不在 POSTGRES_SCHEMA + 里，只锁 schema 函数不够；init_db 持会话级 advisory 锁 + 覆盖全程。空库上跑，未修复时第 1 轮即 3/4 进程报 + UniqueViolation（pg_type_typname_nsp_index）。 + """ + env = _pg_env() + if env is None: + self.skipTest("未设置可达的 CLOUD_TEST_PG_DSN 或未装驱动，" + "跳过真实 PG 测试") + from cloud.db import PostgresDatabase + db = PostgresDatabase( + host=env["CLOUD_DB_HOST"], port=int(env["CLOUD_DB_PORT"]), + dbname=env["CLOUD_DB_NAME"], user=env["CLOUD_DB_USER"], + password=env["CLOUD_DB_PASSWORD"]) + conn = db.connect() + try: + conn.execute("DROP SCHEMA public CASCADE") + conn.execute("CREATE SCHEMA public") + conn.commit() + finally: + conn.close() + nworkers = 4 + barrier = Barrier(nworkers) + queue = Queue() + procs = [Process(target=_concurrent_init_db_worker, + args=(barrier, queue, i, env)) + for i in range(nworkers)] + try: + for proc in procs: + proc.start() + for proc in procs: + proc.join(timeout=180) + finally: + for proc in procs: + if proc.is_alive(): + proc.terminate() + results = [queue.get() for _ in range(nworkers)] + failures = [r for r in results if r[1] == "FAIL"] + self.assertEqual([], failures, + "并发 init_db 有进程失败：%s" % (failures,)) + # 初始化完整：schema 表与 _ensure_* 表都在。 + conn = db.connect() + try: + rows = db.read_all( + conn, + "SELECT table_name FROM information_schema.tables" + " WHERE table_schema = 'public'", + ()) + names = {r["table_name"] for r in rows} + for table in ("users", "projects", "tasks", + "task_continue_messages", "user_gateway_keys", + "task_attempts", "task_audit_loop_state"): + self.assertIn(table, names) + finally: + conn.close() + if __name__ == "__main__": unittest.main() diff --git a/cloud/tests/test_postgres_backend.py b/cloud/tests/test_postgres_backend.py index 816cd6e..40a4239 100644 --- a/cloud/tests/test_postgres_backend.py +++ b/cloud/tests/test_postgres_backend.py @@ -17,7 +17,9 @@ import os import socket import sys +import traceback import unittest +from multiprocessing import Barrier, Process, Queue from pathlib import Path from unittest import mock from urllib.parse import unquote, urlparse @@ -386,3 +388,81 @@ class RealPostgresSchemaTestCase(unittest.TestCase): except Exception: pass conn.close() + + +def _concurrent_schema_worker(barrier, queue, wid, params): + """并发初始化子进程：与其它进程同时跑 init_postgres_schema。""" + try: + host, port, dbname, user, password = params + barrier.wait(timeout=60) + db = PostgresDatabase(host=host, port=port, dbname=dbname, + user=user, password=password) + conn = db.connect() + try: + init_postgres_schema(conn) + finally: + conn.close() + queue.put((wid, "ok", "")) + except Exception: + queue.put((wid, "FAIL", traceback.format_exc())) + + +class RealPostgresConcurrentInitTestCase(unittest.TestCase): + """真实 PG 并发初始化：无 DSN 或不可达时 skip。 + + 回归：启动竞态（37ba687 与 7983567 两次部署均撞上）—— + cloud-server 与 worker 并发 init 时，CREATE TABLE / CREATE + INDEX 的 IF NOT EXISTS 检查与建对象不是原子的，并发建同名 + 对象会撞 pg_type_typname_nsp_index（建表）或 + pg_class_relname_nsp_index（建索引）的 UniqueViolation。 + 修法是 schema 初始化持 advisory 锁串行化；本用例在空库上 + 4 进程同时初始化，必须全部成功、表齐全、无一例异常外泄。 + """ + + def test_concurrent_init_postgres_schema_no_unique_violation(self): + params = _real_pg_params() + if params is None: + self.skipTest("未设置可达的 CLOUD_TEST_PG_DSN，跳过真实 PG 测试") + if not _has_real_driver(): + self.skipTest("未装 psycopg 驱动，跳过真实 PG 测试") + host, port, dbname, user, password = params + db = PostgresDatabase(host=host, port=port, dbname=dbname, + user=user, password=password) + conn = db.connect() + try: + conn.execute("DROP SCHEMA public CASCADE") + conn.execute("CREATE SCHEMA public") + conn.commit() + finally: + conn.close() + nworkers = 4 + barrier = Barrier(nworkers) + queue = Queue() + procs = [Process(target=_concurrent_schema_worker, + args=(barrier, queue, i, params)) + for i in range(nworkers)] + try: + for proc in procs: + proc.start() + for proc in procs: + proc.join(timeout=120) + finally: + for proc in procs: + if proc.is_alive(): + proc.terminate() + results = [queue.get() for _ in range(nworkers)] + failures = [r for r in results if r[1] == "FAIL"] + self.assertEqual([], failures, + "并发初始化有进程失败：%s" % (failures,)) + conn = db.connect() + try: + rows = db.read_all( + conn, + "SELECT table_name FROM information_schema.tables" + " WHERE table_schema = 'public'", + ()) + names = {r["table_name"] for r in rows} + for table in SchemaStatementTestCase.EXPECTED_TABLES: + self.assertIn(table, names) + finally: + conn.close()