diff --git a/cloud/tests/test_usage_quota.py b/cloud/tests/test_usage_quota.py index 91eea31..1e4c8fa 100644 --- a/cloud/tests/test_usage_quota.py +++ b/cloud/tests/test_usage_quota.py @@ -288,6 +288,36 @@ class LedgerTestCase(unittest.TestCase): self.assertNotIn("token", cols) self.assertGreater(row_id, 0) + def test_concurrent_records_keep_chain_intact(self): + # round-12 修复轮(项 5,P11):R12 真机同款并发复测—— + # 4 线程 × 10 条,0 写错,但旧 record() 读-改-写无锁, + # verify_chain 返回 (False, first_bad=11),链 G、链 H + # 两次复现同一断点。修后两轮都须全真、40 条一条不丢。 + for _round in range(2): + with tempfile.TemporaryDirectory() as tmp: + db_path = _make_db(tmp) + ledger = UsageLedger(db_path) + errors = [] + + def _worker(): + try: + for _ in range(10): + ledger.record("u1", "message", 1) + except Exception as e: # noqa: BLE001 - 汇总后断言 + errors.append(e) + + threads = [threading.Thread(target=_worker) + for _ in range(4)] + for t in threads: + t.start() + for t in threads: + t.join() + self.assertEqual(errors, []) + self.assertEqual(ledger.sum_kind("u1", "message"), 40) + ok, bad = ledger.verify_chain("u1") + self.assertTrue(ok, bad) + self.assertIsNone(bad) + class WarningTestCase(unittest.TestCase): def test_once_per_window_at_80_percent(self): diff --git a/cloud/usage.py b/cloud/usage.py index 882a07e..6bf566a 100644 --- a/cloud/usage.py +++ b/cloud/usage.py @@ -32,6 +32,7 @@ import hashlib import json import os import sqlite3 +import threading import time from datetime import datetime, timezone @@ -227,6 +228,13 @@ class UsageLedger: self._conn = _db_connect(db_path) self._conn.execute(self._DDL) self._conn.commit() + # round-12 修复轮(项 5,P11):record() 读-改-写的 + # 串行锁。无锁时两线程读到同一 prev_hash、各写一条, + # 链分叉(R12 真机:4 线程 × 10 条,0 写错, + # verify_chain (False, first_bad=11),链 G、链 H 两次 + # 同断点)。锁只串行化同一实例;跨实例/进程的并发由 + # record() 里的 BEGIN IMMEDIATE 持写锁兜底。 + self._lock = threading.Lock() def _new_conn(self) -> sqlite3.Connection: """Fresh connection per write/read (mirrors CostStore). @@ -256,26 +264,34 @@ class UsageLedger: created_at = time.time() conn = self._new_conn() try: - prev = conn.execute( - "SELECT entry_hash FROM usage_ledger" - " WHERE user_id = ? ORDER BY id DESC LIMIT 1", - (user_id,)).fetchone() - prev_hash = prev[0] if prev is not None else None - fields = {"user_id": user_id, "kind": kind, - "amount": int(amount), "unit": unit, - "task_id": task_id or "", "model": model or "", - "source": source or "", "day": day, - "created_at": created_at} - digest = _entry_hash(prev_hash, fields) - cursor = conn.execute( - """INSERT INTO usage_ledger - (user_id, kind, amount, unit, task_id, model, - source, day, created_at, prev_hash, entry_hash) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", - (user_id, kind, int(amount), unit, task_id, model, - source, day, created_at, prev_hash, digest)) - conn.commit() - return int(cursor.lastrowid or 0) + # 读-改-写原子化:锁串行化同一实例;BEGIN IMMEDIATE + # 在 SELECT 前即持库级写锁(WAL + busy_timeout 下等 + # 待而非立即失败),使另一实例/进程读到的链头必是已 + # 提交的最新值——链不断、条不丢。锁与事务缺一不可: + # 光有事务、同一实例的两线程仍可能各持一个连接同时读 + # 到同一链头。 + with self._lock: + conn.execute("BEGIN IMMEDIATE") + prev = conn.execute( + "SELECT entry_hash FROM usage_ledger" + " WHERE user_id = ? ORDER BY id DESC LIMIT 1", + (user_id,)).fetchone() + prev_hash = prev[0] if prev is not None else None + fields = {"user_id": user_id, "kind": kind, + "amount": int(amount), "unit": unit, + "task_id": task_id or "", "model": model or "", + "source": source or "", "day": day, + "created_at": created_at} + digest = _entry_hash(prev_hash, fields) + cursor = conn.execute( + """INSERT INTO usage_ledger + (user_id, kind, amount, unit, task_id, model, + source, day, created_at, prev_hash, entry_hash) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", + (user_id, kind, int(amount), unit, task_id, model, + source, day, created_at, prev_hash, digest)) + conn.commit() + return int(cursor.lastrowid or 0) finally: if conn is not self._conn: conn.close()