From 9a34517454ed9a7f3403652f7941d40843a36b42 Mon Sep 17 00:00:00 2001 From: snowy-yang <102860110+snowy-yang@users.noreply.github.com> Date: Sun, 27 Sep 2026 11:32:03 +0800 Subject: [PATCH] =?UTF-8?q?perf:=20get=5Fconn=20=E6=AF=8F=E8=BF=9E?= =?UTF-8?q?=E6=8E=A5=E5=9B=BA=E5=AE=9A=E5=BC=80=E9=94=80=201.6ms=20?= =?UTF-8?q?=E2=86=92=200.5ms=EF=BC=88-70%=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 为什么：同步层每次数据库操作都新建连接并执行 mkdir + PRAGMA journal_mode=WAL。实测拆解（各 100 次中位数）： mkdir(exist_ok) 190us sqlite3.connect 435us PRAGMA WAL(新连接首语句) ~1000us ← 大头 get_conn 合计 1641us → 501us 登录路径走同步层 4 次查询（身份解析/验证/位置/日志），每次请求 固定省 ~4.6ms；下载任务的 ledger/日志调用同样受益。 改动（对齐 core/db.py 既有模式——异步层早已做过同款优化并记录过 「并发 PRAGMA 争排他锁报 database is locked」的实测坑，同步层漏同步）： - _WAL_READY 按路径缓存：journal_mode 持久化在库文件头，设一次即可； _WAL_SETUP_LOCK 串行化首连切换（读路径不持 _DB_LOCK） - mkdir 移出 happy path：connect 失败才兜底补建（用 DB_PATH.parent， monkeypatch 场景下与 DATA_DIR 可能不同） - exists() 检查在 connect 之前（connect 会创建空文件，之后恒真）， 覆盖「删库重建后新文件回退 delete 模式」的边界 新增 tests/test_db_conn.py 6 例：WAL 首连生效、后续连接零 PRAGMA （set_trace_callback 实证）、换路径重设、目录缺失自动建、删库重建 恢复 WAL、8 线程并发首连无报错。 门禁：ruff 全过 · pyright 0/0 · pytest 382 passed（375→382）。 --- auth_system/db.py | 30 +++++++++- tests/test_db_conn.py | 131 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 158 insertions(+), 3 deletions(-) create mode 100644 tests/test_db_conn.py diff --git a/auth_system/db.py b/auth_system/db.py index 2e7afe8..3976ae9 100644 --- a/auth_system/db.py +++ b/auth_system/db.py @@ -39,15 +39,39 @@ DEFAULT_REGISTER_CREDITS = 3 # 注册直接赠送 3 次下载次数（无 _DB_LOCK = threading.RLock() +# journal_mode=WAL 持久化在库文件头里，按路径设一次即可。每连接重复执行 +# 实测多付 ~1ms（新连接首条语句要建 WAL 索引读锁），且并发首连会争排他锁 +# 报 database is locked（core/db.py 同款教训）。专用锁 + 二次复查防护。 +_WAL_READY: set[str] = set() +_WAL_SETUP_LOCK = threading.Lock() + # ============ 数据库 ============ def get_conn() -> sqlite3.Connection: - DATA_DIR.mkdir(parents=True, exist_ok=True) - conn = sqlite3.connect(str(DB_PATH), timeout=15) + key = str(DB_PATH) + # exists() 检查（~10us）必须在 connect 之前：connect 会创建空文件， + # 之后 exists 恒为 True，删库重建的场景就检测不到了 + force_wal = not DB_PATH.exists() + # mkdir 移出 happy path（实测 190us/次）：connect 直接失败（目录缺失） + # 才补建重试。库文件被删的场景同样由此兜底重建。 + try: + conn = sqlite3.connect(str(DB_PATH), timeout=15) + except sqlite3.OperationalError: + # DB_PATH.parent 而非 DATA_DIR：生产两者相同，测试 monkeypatch 只换 + # DB_PATH 时兜底仍能命中正确目录 + DB_PATH.parent.mkdir(parents=True, exist_ok=True) + conn = sqlite3.connect(str(DB_PATH), timeout=15) conn.row_factory = sqlite3.Row - conn.execute('PRAGMA journal_mode=WAL') + if key not in _WAL_READY or force_wal: + # 读路径不持 _DB_LOCK，多线程可能同时首连：锁内串行执行 journal_mode + # 切换（防并发争排他锁报 database is locked）。锁内不复查 DB_PATH.exists()—— + # connect 已把文件建出来，exists 恒真；force_wal 是本线程的观测结论， + # 重复执行 PRAGMA 在已是 WAL 的库上幂等且廉价（~2us） + with _WAL_SETUP_LOCK: + conn.execute('PRAGMA journal_mode=WAL') + _WAL_READY.add(key) return conn diff --git a/tests/test_db_conn.py b/tests/test_db_conn.py new file mode 100644 index 0000000..2f99357 --- /dev/null +++ b/tests/test_db_conn.py @@ -0,0 +1,131 @@ +"""get_conn 连接策略契约：WAL 首连设置一次 + 目录自动创建。 + +背景：get_conn 原实现每个连接都执行 mkdir + PRAGMA journal_mode=WAL， +实测每次多付 ~1.2ms（mkdir 190us + 新连接上 PRAGMA ~1ms）。WAL 模式 +持久化在库文件头里，与 core/db.py 的 _WAL_READY_FOR 同款，按路径设一次即可。 +""" + +import sqlite3 +from pathlib import Path + +from auth_system import db + + +def _point_db_to(monkeypatch, tmp_path: Path, conn_counter: list[int] | None = None): + """把 DB_PATH 指到临时库，并清空 WAL 已就绪缓存。""" + new_path = tmp_path / 'site.db' + monkeypatch.setattr(db, 'DB_PATH', new_path) + monkeypatch.setattr(db, '_WAL_READY', set()) + if conn_counter is not None: + real_connect = sqlite3.connect + + def counting_connect(*a, **k): + conn_counter[0] += 1 + return real_connect(*a, **k) + + monkeypatch.setattr(db.sqlite3, 'connect', counting_connect) + return new_path + + +def test_wal_mode_active_on_fresh_db(monkeypatch, tmp_path): + """新路径首连：WAL 生效。""" + _point_db_to(monkeypatch, tmp_path) + conn = db.get_conn() + try: + mode = conn.execute('PRAGMA journal_mode').fetchone()[0] + assert mode.lower() == 'wal' + finally: + conn.close() + + +def test_wal_pragma_only_once_per_path(monkeypatch, tmp_path): + """同一路径第二次起不再执行 PRAGMA journal_mode=WAL（用 SQL 追踪器实证）。""" + _point_db_to(monkeypatch, tmp_path) + + conn1 = db.get_conn() + conn1.close() + conn2 = db.get_conn() + conn2.close() + + # 用 sqlite3 的 set_trace_callback 记录第三条连接上的语句（2 连接已建立缓存） + statements: list[str] = [] + conn3 = db.get_conn() + conn3.set_trace_callback(statements.append) + conn3.execute('SELECT 1').fetchall() + conn3.close() + assert statements == ['SELECT 1'], f'第三条连接不应再跑 PRAGMA，实际: {statements}' + + +def test_new_path_resets_wal_cache(monkeypatch, tmp_path): + """换库路径（测试隔离场景）：新路径重新设置 WAL。""" + _point_db_to(monkeypatch, tmp_path / 'a') + c = db.get_conn() + c.close() + + monkeypatch.setattr(db, 'DB_PATH', tmp_path / 'b' / 'site.db') + c = db.get_conn() + try: + mode = c.execute('PRAGMA journal_mode').fetchone()[0] + assert mode.lower() == 'wal' + finally: + c.close() + + +def test_missing_dir_auto_created(monkeypatch, tmp_path): + """目录不存在：connect 失败兜底 mkdir 重试，不经 init_db 也能用。""" + deep = tmp_path / 'x' / 'y' + _point_db_to(monkeypatch, deep) + assert not deep.exists() + + conn = db.get_conn() + try: + conn.execute('CREATE TABLE t (v)').close() + conn.commit() + finally: + conn.close() + assert (deep / 'site.db').exists() + + +def test_deleted_db_file_recovers(monkeypatch, tmp_path): + """库文件被删（目录还在）：重连自动重建且恢复 WAL 模式，不崩。""" + _point_db_to(monkeypatch, tmp_path) + c = db.get_conn() + c.close() + + (tmp_path / 'site.db').unlink() + conn = db.get_conn() + try: + conn.execute('SELECT 1').fetchall() + mode = conn.execute('PRAGMA journal_mode').fetchone()[0] + assert mode.lower() == 'wal', '重建的库文件必须恢复 WAL（journal_mode 不持久于新文件）' + finally: + conn.close() + + +def test_concurrent_first_connect_no_lease_conflict(monkeypatch, tmp_path): + """多线程同时首连（读路径不持 _DB_LOCK）：WAL 设置正确、无连接报错。 + + core/db.py 记录过并发 PRAGMA 争排他锁报 database is locked 的实测坑， + 同步层用专用锁 + 二次复查同款防护。 + """ + import threading + + _point_db_to(monkeypatch, tmp_path) + errors: list[Exception] = [] + + def worker(): + try: + for _ in range(5): + c = db.get_conn() + c.execute('SELECT 1').fetchall() + c.close() + except Exception as e: # 收集全部失败供断言 + errors.append(e) + + threads = [threading.Thread(target=worker) for _ in range(8)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert errors == []