xsl
ecdfbb64c3
POST /quota/tasks/{task_key}/claim：配置驱动任务注册表、
grant_bonus 纯域函数、ledger 幂等键防重、task_bonus_max 封顶（FR-BILL-03/04）。
实现：packages/domain/billing/quota.py 新增 BonusDecision/grant_bonus（mypy strict，
恰好到顶允许、超一格整体拒绝）、grant_idempotency_key、PERIOD_TYPES；
packages/billing/service.py 新增 claim_task_bonus（幂等快捷路径 → 行锁 → 封顶判定
→ 同事务更新 bonus_granted 并写 system_grant 流水，delta=0 且 balance_after 按
remaining 口径；IntegrityError 并发兜底回滚为 40910）；响应模型 TaskBonusClaimResponse
放新建的 packages/billing/schemas.py；路由挂在 apps/api/routers/quota.py，
rules 由 settings.task_bonus_rules 注入（amount 为占位值，产品定稿后改 env）。
测试：新增单测 11 个（封顶边界、amount<=0、不动 used、hard_cap 独立生效、
默认规则形状与 env 覆盖通路）；新增集成 4 个（领取+重复 40910+纯 DB 流水断言、
五类任务全部到账、未知任务 40410、amount=11>max=10 → 40310 且无新流水）。
门禁全绿：ruff check / ruff format --check / mypy 两段 / lint-imports；
单测 180 passed，集成 7 passed。
偏差说明：1) 任务文档提到的 BillingService._load_rule 实际名为 _require_rule，
直接复用现有助手；2) packages/billing/repo.py 无需新增函数（现有
get_ledger_by_idempotency_key/lock_or_create_quota/add_ledger 已覆盖本任务）；
3) 文档设想的「TASK_BONUS_RULES_JSON 空值容错」经 env 源实测不成立——
pydantic-settings 对 dict 字段先做 JSON 解码，空白字符串在 Settings 构造时即
报错（共享校验器/config.py 属禁改，未动）；单测改为覆盖受支持的空配置形态 "{}"。
428 lines
16 KiB
Python
Executable File
428 lines
16 KiB
Python
Executable File
"""Quota three-phase protocol: reserve → commit → refund (架构文档 §7.4).
|
||
|
||
为什么预扣而不是直接扣: 长耗时任务可能失败，直接扣会产生「扣了额度没拿到
|
||
结果」的客诉。V1 没有客服退款兜底手段，只能靠机制保证 (FR-BILL-05).
|
||
|
||
Nothing here decides *how much* a feature costs — that comes from
|
||
``bl_quota_rule`` so it can change without a release (FR-BILL-04).
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass
|
||
from datetime import UTC, datetime
|
||
from typing import Any, cast
|
||
|
||
from sqlalchemy.exc import IntegrityError
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from packages.billing import repo
|
||
from packages.domain.billing import quota as quota_rules
|
||
from packages.domain.errors import QuotaExceededError
|
||
from packages.domain.ids import next_id
|
||
from packages.shared.errors import AppError, ErrorCode, QuotaExhaustedError
|
||
from packages.shared.logging import get_logger
|
||
from packages.shared.models import Ledger, QuotaRule
|
||
from packages.shared.observability import QUOTA_EVENTS
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class Reservation:
|
||
tx_id: int
|
||
user_id: int
|
||
quota_key: str
|
||
period_key: str
|
||
amount: int
|
||
remaining: int
|
||
idempotency_key: str
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class QuotaSnapshot:
|
||
quota_key: str
|
||
period_key: str
|
||
used: int
|
||
limit_value: int
|
||
bonus_granted: int
|
||
remaining: int
|
||
reset_at: datetime | None
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class TaskBonusClaim:
|
||
"""一次成功领取的结果（路由层据此组装 TaskBonusClaimResponse）."""
|
||
|
||
task_key: str
|
||
quota_key: str
|
||
amount: int
|
||
bonus_granted: int
|
||
remaining: int
|
||
|
||
|
||
def _parse_task_rule(
|
||
task_key: str, conf: dict[str, Any]
|
||
) -> tuple[str, int, quota_rules.PeriodType]:
|
||
"""解析 ``settings.task_bonus_rules`` 的一条配置.
|
||
|
||
坏配置是运维错误（5xxxx），不是用户错误：提前在这里拦下，比让
|
||
``period_key`` 在请求路径上抛 ValueError 变成 500 清晰得多。
|
||
"""
|
||
quota_key = conf.get("quota_key")
|
||
amount = conf.get("amount")
|
||
claim_period = conf.get("claim_period")
|
||
if not isinstance(quota_key, str) or not quota_key:
|
||
raise AppError(
|
||
ErrorCode.CONFIG_ERROR, message="任务奖励配置不合法", data={"task_key": task_key}
|
||
)
|
||
if not isinstance(amount, int) or isinstance(amount, bool) or amount <= 0:
|
||
raise AppError(
|
||
ErrorCode.CONFIG_ERROR, message="任务奖励配置不合法", data={"task_key": task_key}
|
||
)
|
||
if not isinstance(claim_period, str) or claim_period not in quota_rules.PERIOD_TYPES:
|
||
raise AppError(
|
||
ErrorCode.CONFIG_ERROR, message="任务奖励配置不合法", data={"task_key": task_key}
|
||
)
|
||
return quota_key, amount, cast(quota_rules.PeriodType, claim_period)
|
||
|
||
|
||
class BillingService:
|
||
def __init__(self, session: AsyncSession, *, tz_offset_hours: int = 8) -> None:
|
||
self._session = session
|
||
self._tz_offset = tz_offset_hours
|
||
|
||
# -- read ---------------------------------------------------------------
|
||
|
||
async def get_quota(self, *, user_id: int, quota_key: str) -> QuotaSnapshot:
|
||
rule = await self._require_rule(quota_key)
|
||
period_type = cast(quota_rules.PeriodType, rule.period_type)
|
||
now = datetime.now(UTC)
|
||
period_key = quota_rules.period_key(now, period_type, self._tz_offset)
|
||
row = await repo.lock_or_create_quota(
|
||
self._session,
|
||
user_id=user_id,
|
||
quota_key=quota_key,
|
||
period_key=period_key,
|
||
limit_value=rule.base_limit,
|
||
)
|
||
state = quota_rules.QuotaState(
|
||
used=row.used,
|
||
base_limit=row.limit_value,
|
||
bonus=row.bonus_granted,
|
||
hard_cap=rule.hard_cap,
|
||
)
|
||
return QuotaSnapshot(
|
||
quota_key=quota_key,
|
||
period_key=period_key,
|
||
used=row.used,
|
||
limit_value=row.limit_value,
|
||
bonus_granted=row.bonus_granted,
|
||
remaining=state.remaining,
|
||
reset_at=quota_rules.period_reset_at(now, period_type, self._tz_offset),
|
||
)
|
||
|
||
async def list_ledger(self, *, user_id: int, limit: int = 50) -> list[Ledger]:
|
||
return await repo.list_ledger(self._session, user_id, limit)
|
||
|
||
# -- phase 1: reserve ---------------------------------------------------
|
||
|
||
async def reserve(
|
||
self,
|
||
*,
|
||
user_id: int,
|
||
quota_key: str,
|
||
request_key: str,
|
||
amount: int = 1,
|
||
reason: str | None = None,
|
||
ref_type: str | None = None,
|
||
ref_id: str | None = None,
|
||
) -> Reservation:
|
||
"""Freeze ``amount`` units. Idempotent on ``request_key``."""
|
||
idem_key = quota_rules.reserve_idempotency_key(
|
||
user_id=user_id, quota_key=quota_key, request_key=request_key
|
||
)
|
||
existing = await repo.get_ledger_by_idempotency_key(self._session, idem_key)
|
||
if existing is not None:
|
||
QUOTA_EVENTS.labels(quota_key, "reserve", "idempotent_hit").inc()
|
||
return Reservation(
|
||
tx_id=existing.tx_id,
|
||
user_id=user_id,
|
||
quota_key=quota_key,
|
||
period_key=existing.period_key or "",
|
||
amount=abs(existing.delta),
|
||
remaining=existing.balance_after,
|
||
idempotency_key=idem_key,
|
||
)
|
||
|
||
rule = await self._require_rule(quota_key)
|
||
period_type = cast(quota_rules.PeriodType, rule.period_type)
|
||
now = datetime.now(UTC)
|
||
period_key = quota_rules.period_key(now, period_type, self._tz_offset)
|
||
row = await repo.lock_or_create_quota(
|
||
self._session,
|
||
user_id=user_id,
|
||
quota_key=quota_key,
|
||
period_key=period_key,
|
||
limit_value=rule.base_limit,
|
||
)
|
||
state = quota_rules.QuotaState(
|
||
used=row.used,
|
||
base_limit=row.limit_value,
|
||
bonus=row.bonus_granted,
|
||
hard_cap=rule.hard_cap,
|
||
)
|
||
|
||
try:
|
||
new_state = quota_rules.reserve(state, amount)
|
||
except QuotaExceededError as exc:
|
||
QUOTA_EVENTS.labels(quota_key, "reserve", "exhausted").inc()
|
||
reset_at = quota_rules.period_reset_at(now, period_type, self._tz_offset)
|
||
raise QuotaExhaustedError(
|
||
required=exc.required,
|
||
remaining=exc.remaining,
|
||
reset_at=reset_at.isoformat() if reset_at else None,
|
||
) from exc
|
||
|
||
row.used = new_state.used
|
||
tx_id = next_id()
|
||
repo.add_ledger(
|
||
self._session,
|
||
Ledger(
|
||
tx_id=tx_id,
|
||
user_id=user_id,
|
||
delta=-amount,
|
||
tx_type="consume",
|
||
status="reserved",
|
||
reason=reason,
|
||
ref_type=ref_type,
|
||
ref_id=ref_id,
|
||
quota_key=quota_key,
|
||
period_key=period_key,
|
||
balance_after=quota_rules.balance_after(new_state),
|
||
idempotency_key=idem_key,
|
||
),
|
||
)
|
||
await self._session.flush()
|
||
QUOTA_EVENTS.labels(quota_key, "reserve", "ok").inc()
|
||
logger.info(
|
||
"quota_reserved",
|
||
tx_id=tx_id,
|
||
quota_key=quota_key,
|
||
amount=amount,
|
||
remaining=new_state.remaining,
|
||
)
|
||
return Reservation(
|
||
tx_id=tx_id,
|
||
user_id=user_id,
|
||
quota_key=quota_key,
|
||
period_key=period_key,
|
||
amount=amount,
|
||
remaining=new_state.remaining,
|
||
idempotency_key=idem_key,
|
||
)
|
||
|
||
# -- phase 2: commit ----------------------------------------------------
|
||
|
||
async def commit(self, reservation: Reservation) -> None:
|
||
"""Confirm the deduction. Safe to call twice."""
|
||
ledger = await repo.get_ledger(self._session, reservation.tx_id)
|
||
if ledger is None:
|
||
raise AppError(ErrorCode.STATE_CONFLICT, message="配额流水不存在")
|
||
if ledger.status == "reserved":
|
||
ledger.status = "committed"
|
||
await self._session.flush()
|
||
QUOTA_EVENTS.labels(reservation.quota_key, "commit", "ok").inc()
|
||
logger.info("quota_committed", tx_id=ledger.tx_id)
|
||
else:
|
||
QUOTA_EVENTS.labels(reservation.quota_key, "commit", "noop").inc()
|
||
|
||
# -- phase 3: refund ----------------------------------------------------
|
||
|
||
async def refund(self, reservation: Reservation, *, reason: str = "task_failed") -> None:
|
||
"""Release a reservation and write the reverse ledger entry.
|
||
|
||
FR-BILL-05 requires refunds to be idempotent; the refund ledger key is
|
||
derived from the reserve key so a retry can never double-refund.
|
||
"""
|
||
ledger = await repo.get_ledger(self._session, reservation.tx_id)
|
||
if ledger is None:
|
||
raise AppError(ErrorCode.STATE_CONFLICT, message="配额流水不存在")
|
||
if ledger.status != "reserved":
|
||
QUOTA_EVENTS.labels(reservation.quota_key, "refund", "noop").inc()
|
||
return
|
||
|
||
refund_key = quota_rules.refund_idempotency_key(reservation.idempotency_key)
|
||
if await repo.get_ledger_by_idempotency_key(self._session, refund_key) is not None:
|
||
QUOTA_EVENTS.labels(reservation.quota_key, "refund", "noop").inc()
|
||
return
|
||
|
||
rule = await self._require_rule(reservation.quota_key)
|
||
row = await repo.lock_or_create_quota(
|
||
self._session,
|
||
user_id=reservation.user_id,
|
||
quota_key=reservation.quota_key,
|
||
period_key=reservation.period_key,
|
||
limit_value=rule.base_limit,
|
||
)
|
||
state = quota_rules.QuotaState(
|
||
used=row.used,
|
||
base_limit=row.limit_value,
|
||
bonus=row.bonus_granted,
|
||
hard_cap=rule.hard_cap,
|
||
)
|
||
new_state = quota_rules.refund(state, reservation.amount)
|
||
row.used = new_state.used
|
||
ledger.status = "refunded"
|
||
|
||
repo.add_ledger(
|
||
self._session,
|
||
Ledger(
|
||
tx_id=next_id(),
|
||
user_id=reservation.user_id,
|
||
delta=reservation.amount,
|
||
tx_type="refund",
|
||
status="committed",
|
||
reason=reason,
|
||
ref_type=ledger.ref_type,
|
||
ref_id=ledger.ref_id,
|
||
quota_key=reservation.quota_key,
|
||
period_key=reservation.period_key,
|
||
balance_after=quota_rules.balance_after(new_state),
|
||
idempotency_key=refund_key,
|
||
),
|
||
)
|
||
await self._session.flush()
|
||
QUOTA_EVENTS.labels(reservation.quota_key, "refund", "ok").inc()
|
||
logger.info(
|
||
"quota_refunded",
|
||
tx_id=ledger.tx_id,
|
||
amount=reservation.amount,
|
||
remaining=new_state.remaining,
|
||
)
|
||
|
||
# -- task bonus (FR-BILL-03) --------------------------------------------
|
||
|
||
async def claim_task_bonus(
|
||
self,
|
||
*,
|
||
user_id: int,
|
||
task_key: str,
|
||
rules: dict[str, dict[str, Any]],
|
||
) -> TaskBonusClaim:
|
||
"""领取任务奖励：发放即到账、幂等、封顶（FR-BILL-03）.
|
||
|
||
``rules`` 来自 ``settings.task_bonus_rules``（FR-BILL-04 配置化）：
|
||
amount 目前是占位值，产品在产品方案 §3.2 定稿后改 env 即生效，
|
||
不需要改代码。
|
||
"""
|
||
conf = rules.get(task_key)
|
||
if conf is None:
|
||
raise AppError(ErrorCode.TASK_BONUS_UNKNOWN, data={"task_key": task_key})
|
||
quota_key, amount, claim_period = _parse_task_rule(task_key, conf)
|
||
|
||
rule = await self._require_rule(quota_key)
|
||
quota_period_type = cast(quota_rules.PeriodType, rule.period_type)
|
||
now = datetime.now(UTC)
|
||
# 幂等粒度跟任务的 claim_period 走，额度行周期跟 bl_quota_rule 走：
|
||
# 「每天可领」的任务可以给「按月计量」的额度加成，两套周期不能混。
|
||
quota_pk = quota_rules.period_key(now, quota_period_type, self._tz_offset)
|
||
claim_pk = quota_rules.period_key(now, claim_period, self._tz_offset)
|
||
idem_key = quota_rules.grant_idempotency_key(
|
||
user_id=user_id, task_key=task_key, claim_period_key=claim_pk
|
||
)
|
||
|
||
# 幂等快捷路径：本周期已发放过就直接拒绝，不进任何写路径。
|
||
if await repo.get_ledger_by_idempotency_key(self._session, idem_key) is not None:
|
||
QUOTA_EVENTS.labels(quota_key, "grant", "idempotent_hit").inc()
|
||
raise AppError(ErrorCode.TASK_BONUS_ALREADY_CLAIMED, data={"task_key": task_key})
|
||
|
||
row = await repo.lock_or_create_quota(
|
||
self._session,
|
||
user_id=user_id,
|
||
quota_key=quota_key,
|
||
period_key=quota_pk,
|
||
limit_value=rule.base_limit,
|
||
)
|
||
state = quota_rules.QuotaState(
|
||
used=row.used,
|
||
base_limit=row.limit_value,
|
||
bonus=row.bonus_granted,
|
||
hard_cap=rule.hard_cap,
|
||
)
|
||
decision = quota_rules.grant_bonus(state, amount, rule.task_bonus_max)
|
||
if not decision.allowed:
|
||
# 封顶拒绝：不写任何流水，bonus_granted 保持原值（验收 3）。
|
||
QUOTA_EVENTS.labels(quota_key, "grant", "capped").inc()
|
||
raise AppError(
|
||
ErrorCode.TASK_BONUS_CAP_REACHED,
|
||
data={
|
||
"task_key": task_key,
|
||
"bonus_granted": state.bonus,
|
||
"task_bonus_max": rule.task_bonus_max,
|
||
},
|
||
)
|
||
|
||
row.bonus_granted = decision.bonus
|
||
tx_id = next_id()
|
||
try:
|
||
repo.add_ledger(
|
||
self._session,
|
||
Ledger(
|
||
tx_id=tx_id,
|
||
user_id=user_id,
|
||
# delta=0：system_grant 不改变 used，只提高有效额度；balance_after
|
||
# 仍按 remaining 语义记录 —— 流水可重算的是 used，不是余额本身，
|
||
# 与 reserve/commit 的记账口径一致。
|
||
delta=0,
|
||
tx_type="system_grant",
|
||
status="committed",
|
||
reason=task_key,
|
||
ref_type="task_bonus",
|
||
ref_id=task_key,
|
||
quota_key=quota_key,
|
||
period_key=quota_pk,
|
||
balance_after=decision.remaining,
|
||
idempotency_key=idem_key,
|
||
),
|
||
)
|
||
await self._session.flush()
|
||
except IntegrityError as exc:
|
||
# 并发兜底：两个请求同时通过上面的快捷检查，idempotency_key 唯一约束
|
||
# 让后提交者失败。对用户来说结果与快捷路径等价 —— 已领取，而非 500。
|
||
await self._session.rollback()
|
||
QUOTA_EVENTS.labels(quota_key, "grant", "idempotent_hit").inc()
|
||
raise AppError(
|
||
ErrorCode.TASK_BONUS_ALREADY_CLAIMED, data={"task_key": task_key}
|
||
) from exc
|
||
|
||
QUOTA_EVENTS.labels(quota_key, "grant", "ok").inc()
|
||
logger.info(
|
||
"task_bonus_granted",
|
||
tx_id=tx_id,
|
||
task_key=task_key,
|
||
quota_key=quota_key,
|
||
amount=amount,
|
||
bonus=decision.bonus,
|
||
remaining=decision.remaining,
|
||
)
|
||
return TaskBonusClaim(
|
||
task_key=task_key,
|
||
quota_key=quota_key,
|
||
amount=amount,
|
||
bonus_granted=decision.bonus,
|
||
remaining=decision.remaining,
|
||
)
|
||
|
||
# -- helpers ------------------------------------------------------------
|
||
|
||
async def _require_rule(self, quota_key: str) -> QuotaRule:
|
||
rule = await repo.get_rule(self._session, quota_key)
|
||
if rule is None:
|
||
raise AppError(
|
||
ErrorCode.CONFIG_ERROR,
|
||
message="配额规则未配置",
|
||
data={"quota_key": quota_key},
|
||
)
|
||
return rule
|