xsl
2009aa938c
The product plan is the only source for quota numbers, but the 0001 seed used a second set (agent_chat 20/50, tryon hard 10) and omitted hairstyle, hd_export and report regen. Migration 0003 replaces the seed and renames points_reserved → quota_reserved and bonus_value → bonus_granted so the docs and the schema stop contradicting each other. Effective limit now clips at hard_cap. CI runs the integration slice against the Postgres/Redis services it already starts.
277 lines
9.8 KiB
Python
Executable File
277 lines
9.8 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 cast
|
|
|
|
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
|
|
|
|
|
|
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,
|
|
)
|
|
|
|
# -- 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
|