tzt
a9f3405970
- ledger.py：students/keys CRUD + 日限额 check_and_count（跨日重置，注入时钟）+ 流水分页；计费两阶段拆分至 billing.py（BillingMixin，评审聚焦）： try_hold 锁内读现状->日上限判定->余额原子防线->holding 流水（request_id 幂等）、 settle 按真实值回补（预扣-实际）差额、void 全额退款（上游失败） - auth.py：issue_key（sk-campus- 前缀，明文只返回一次，库存 sha256+前缀）、 authenticate 校验链（形态/注销/停用/日额/rpm，热路径 LRU TTL30s D-P10）、 RateLimiter（令牌桶 rpm + per-key 并发信号量，D-P9 单进程）、 verify_admin（hmac.compare_digest；未配置仅 loopback） - 测试 +17：签发/错key/注销/停用/403/rpm 429/日额 429/并发槽/admin 策略 + 预扣-结算-回补一致/双向差额/余额不足/日上限/幂等/跨日重置/void/并发10路不超扣 - 全量 339 passed
158 lines
7.0 KiB
Python
158 lines
7.0 KiB
Python
"""计费两阶段（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]
|