"""计费两阶段（D-P11）：预扣-结算-退款，作为 Ledger 的 mixin。 拆分为独立模块的工程原因：ledger.py 承载 DDL/CRUD/查询；金额敏感的 两阶段语义集中在此，便于评审与测试聚焦。 并发正确性：全局锁内 check-then-act（单写者模型）+ 余额原子防线 （语句条件 balance_milli >= ?，受影响行数为 0 即不足）。 """ from __future__ import annotations import threading import time from typing import Any, Dict, List, Optional from gateway.proxy.ledgerutil import _today # 预扣：余额原子防线；日上限判定在锁内 Python 侧（单写者模型下等价安全） # 占位顺序：est, spent_total, today, sid, est _SQL_HOLD = """UPDATE students SET balance_milli = balance_milli - ?, spent_today_milli = ?, spent_date = ? WHERE id = ? AND balance_milli >= ?""" _SQL_HOLD_MARK = """INSERT INTO usage_ledger (request_id, ts, key_id, model, bucket, charged_milli, status) VALUES (?, ?, ?, ?, ?, ?, 'holding')""" _SQL_INSUFFICIENT = """INSERT OR REPLACE INTO usage_ledger (request_id, ts, key_id, model, bucket, charged_milli, status) VALUES (?, ?, ?, ?, ?, 0, 'insufficient')""" _SQL_SETTLE_REFUND = """UPDATE students SET balance_milli = balance_milli + ?, spent_today_milli = MAX(0, spent_today_milli - ?) WHERE id = ?""" _SQL_SETTLE_USAGE = """UPDATE usage_ledger SET in_miss_tok = ?, in_hit_tok = ?, out_tok = ?, gateway_cached = ?, upstream_cost_milli = ?, charged_milli = ?, margin_milli = ? - ?, ttfb_ms = ?, total_ms = ?, status = ? WHERE request_id = ?""" class BillingMixin: """try_hold / settle / void（由 Ledger 继承；依赖 _lock/_connect）。""" _lock: threading.Lock def _connect(self): # 由 Ledger 提供 raise NotImplementedError def try_hold(self, request_id: str, key_id: int, student_id: int, model: str, bucket: str, est_milli: int, ts: float) -> bool: """预扣：锁内读现状 -> 日上限判定 -> 余额原子扣减 -> 写 holding 流水。 受影响行数为 0 即不足（余额不够）-> False（402）； request_id 幂等：重复 hold 直接返回 True（不重复扣）。 """ today = _today(ts) with self._lock, self._connect() as conn: dup = conn.execute( "SELECT 1 FROM usage_ledger WHERE request_id = ?", (request_id,)).fetchone() if dup: return True stu = conn.execute( "SELECT balance_milli, daily_cap_milli, spent_today_milli, spent_date" " FROM students WHERE id = ?", (student_id,)).fetchone() if stu is None: return False spent_base = stu["spent_today_milli"] if stu["spent_date"] == today else 0 if stu["daily_cap_milli"] > 0 and spent_base + est_milli > stu["daily_cap_milli"]: conn.execute(_SQL_INSUFFICIENT, (request_id, int(ts), key_id, model, bucket)) return False cur = conn.execute(_SQL_HOLD, ( est_milli, spent_base + est_milli, today, student_id, est_milli)) if cur.rowcount == 0: conn.execute(_SQL_INSUFFICIENT, (request_id, int(ts), key_id, model, bucket)) return False conn.execute(_SQL_HOLD_MARK, (request_id, int(ts), key_id, model, bucket, est_milli)) return True def settle(self, request_id: str, actual_milli: int, *, in_miss_tok: int = 0, in_hit_tok: int = 0, out_tok: int = 0, gateway_cached: int = 0, upstream_cost_milli: int = 0, ttfb_ms: Optional[int] = None, total_ms: Optional[int] = None, status: str = "ok") -> bool: """结算：按真实值更新流水并回补（预扣额 − 实际额）差额。""" with self._lock, self._connect() as conn: row = conn.execute( "SELECT key_id, charged_milli FROM usage_ledger WHERE request_id = ?", (request_id,)).fetchone() if row is None: return False key_id = row["key_id"] est = row["charged_milli"] student = conn.execute( "SELECT student_id FROM proxy_keys WHERE id = ?", (key_id,)).fetchone() if student is None: return False refund = est - actual_milli if refund != 0: conn.execute(_SQL_SETTLE_REFUND, (refund, refund, student["student_id"])) conn.execute(_SQL_SETTLE_USAGE, ( in_miss_tok, in_hit_tok, out_tok, gateway_cached, upstream_cost_milli, actual_milli, actual_milli, upstream_cost_milli, ttfb_ms, total_ms, status, request_id)) return True def void(self, request_id: str, status: str = "aborted") -> bool: """全额退款（上游失败）；流水保留审计。""" with self._lock, self._connect() as conn: row = conn.execute( "SELECT key_id, charged_milli, status FROM usage_ledger" " WHERE request_id = ?", (request_id,)).fetchone() if row is None or row["status"] not in ("holding", "ok"): return False est = row["charged_milli"] student = conn.execute( "SELECT student_id FROM proxy_keys WHERE id = ?", (row["key_id"],)).fetchone() if student is not None and est > 0: conn.execute(_SQL_SETTLE_REFUND, (est, est, student["student_id"])) conn.execute( "UPDATE usage_ledger SET charged_milli = 0, status = ?" " WHERE request_id = ?", (status, request_id)) return True # ---------- 流水查询 ---------- def get_usage(self, request_id: str) -> Optional[Dict[str, Any]]: with self._lock, self._connect() as conn: row = conn.execute( "SELECT * FROM usage_ledger WHERE request_id = ?", (request_id,)).fetchone() return dict(row) if row else None def list_usage(self, student_id: Optional[int] = None, limit: int = 50, offset: int = 0) -> List[Dict[str, Any]]: """流水分页（可按学生过滤，经其名下 key）；按学生过滤走联表常量语句。""" if student_id is None: with self._lock, self._connect() as conn: rows = conn.execute( "SELECT * FROM usage_ledger ORDER BY ts DESC LIMIT ? OFFSET ?", (limit, offset)).fetchall() return [dict(r) for r in rows] with self._lock, self._connect() as conn: rows = conn.execute( "SELECT u.* FROM usage_ledger u JOIN proxy_keys k ON k.id = u.key_id" " WHERE k.student_id = ? ORDER BY u.ts DESC LIMIT ? OFFSET ?", (student_id, limit, offset)).fetchall() return [dict(r) for r in rows]