refactor: 抽象统一计费会话 BillingSession - #2877
Conversation
将散落在多个文件中的预扣费/结算/退款逻辑抽象为统一的 BillingSession 生命周期管理： - 新增 BillingSettler 接口 (relay/common/billing.go) 避免循环引用 - 新增 FundingSource 接口 + WalletFunding / SubscriptionFunding 实现 (service/funding_source.go) - 新增 BillingSession 封装预扣/结算/退款原子操作 (service/billing_session.go) - 新增 SettleBilling 统一结算辅助函数，替换各 handler 中的 quotaDelta 模式 - 重写 PreConsumeBilling 为 BillingSession 工厂入口 - controller/relay.go 退款守卫改用 BillingSession.Refund() 修复的 Bug： - 令牌额度泄漏：PreConsumeTokenQuota 成功但 DecreaseUserQuota 失败时未回滚 - 订阅退款遗漏：FinalPreConsumedQuota=0 但 SubscriptionPreConsumed>0 时跳过退款 - 订阅多扣费：subConsume 强制为 1 但 FinalPreConsumedQuota 不同步 - 退款路径不统一：钱包/订阅退款逻辑现统一由 FundingSource.Refund 分派
WalkthroughAdds a billing lifecycle abstraction: a Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant RelayHandler as Relay Handler
participant BillingSession as BillingSession
participant FundingSource as FundingSource
participant QuotaService as Quota Service
Client->>RelayHandler: Request
RelayHandler->>BillingSession: NewBillingSession(pre-consume)
activate BillingSession
BillingSession->>BillingSession: shouldTrust / token checks
alt Pre-consume required
BillingSession->>FundingSource: PreConsume(amount)
activate FundingSource
FundingSource->>QuotaService: Deduct/Validate quota
FundingSource-->>BillingSession: Pre-consume result
deactivate FundingSource
end
BillingSession-->>RelayHandler: Session ready
deactivate BillingSession
RelayHandler->>RelayHandler: Process request -> actualQuota
alt Success
RelayHandler->>BillingSession: Settle(actualQuota)
activate BillingSession
BillingSession->>BillingSession: calc delta
BillingSession->>FundingSource: Settle(delta)
FundingSource->>QuotaService: Adjust quota
FundingSource-->>BillingSession: Settled
BillingSession-->>RelayHandler: Settlement done
deactivate BillingSession
else Error / Abort
RelayHandler->>BillingSession: Refund()
activate BillingSession
BillingSession->>FundingSource: Refund()
FundingSource->>QuotaService: Restore quota (async)
FundingSource-->>BillingSession: Refunded
BillingSession-->>RelayHandler: Refund done
deactivate BillingSession
end
RelayHandler-->>Client: Response
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Fix all issues with AI agents
In `@service/billing_session.go`:
- Around line 36-69: The Settle method can partially apply the funding delta
then fail adjusting token quota, so modify Settle to mark the funding part as
completed immediately after s.funding.Settle(delta) succeeds: either set
s.settled = true right after funding.Settle or add a new boolean (e.g.,
fundingSettled) on BillingSession and set it there; then proceed to try token
quota adjustments and if those fail, log the error and return (but do not
re-apply funding on retry). Also update Refund (and any callers that rely on
s.settled) to check the new fundingSettled flag (or the early-set s.settled) so
they won't refund an already-applied funding delta; ensure
relayInfo.SubscriptionPostDelta and the final s.settled are only updated after
all durable adjustments are known or handled appropriately.
- Around line 275-293: The bug is a mismatch between SubscriptionFunding.amount
and the value passed into preConsume: trySubscription sets amount = max(1,
preConsumedQuota) but calls session.preConsume with the original
preConsumedQuota, so when preConsumedQuota==0 the subscription deducts 1 while
the session records 0. Fix by computing subConsume := preConsumedQuota; if
subConsume <= 0 { subConsume = 1 }, assign that value to
SubscriptionFunding.amount, and pass that same subConsume into
BillingSession.preConsume (i.e., call session.preConsume(c, subConsume)) so
SubscriptionFunding.amount, the preConsume parameter, and any
FinalPreConsumedQuota remain consistent; references: trySubscription,
SubscriptionFunding.amount, BillingSession.preConsume, PreConsume.
- Around line 159-174: When PreConsume fails and you attempt to rollback token
quota, do not discard the error from model.IncreaseTokenQuota; capture its
return value, log the failure (including context: s.relayInfo.TokenId,
tokenConsumed, and original PreConsume error) and, if appropriate,
combine/annotate it with the original error before returning so quota
reconciliation is visible; update the block around s.funding.PreConsume to check
and handle errFromRollback := model.IncreaseTokenQuota(...) instead of ignoring
it. Also replace fragile substring matching on errMsg (the strings.Contains
checks) with structured error handling: have FundingSource implementations
return sentinel errors or typed errors and use errors.Is / errors.As to branch
(update the conditionals that call types.NewErrorWithStatusCode accordingly,
e.g., check for ErrNoActiveSubscription, ErrSubscriptionQuotaInsufficient,
ErrUserQuotaInsufficient).
In `@service/billing.go`:
- Around line 34-73: The fallback branch in SettleBilling skips quota
notifications when quotaDelta == 0; ensure the notification still fires by
invoking the same notification path as the non-fallback branch. Modify
SettleBilling so that when relayInfo.Billing == nil and quotaDelta == 0 you
either call PostConsumeQuota(relayInfo, 0, relayInfo.FinalPreConsumedQuota,
true) (to preserve old behavior) or explicitly call
checkAndSendQuotaNotify(relayInfo, 0, relayInfo.FinalPreConsumedQuota) so that
notifications run even when no quota change occurred.
🧹 Nitpick comments (7)
service/funding_source.go (4)
57-62:WalletFunding.Refunddoes not resetconsumed, unlikeSubscriptionFundingwhich guards onpreConsumed.If
Refund()is called twice (e.g., due to a bug in a future caller), the wallet gets credited twice. TheBillingSessionlayer should prevent this via its own idempotency guard, but defensive coding withinWalletFundingwould be safer.🛡️ Suggested defensive reset
func (w *WalletFunding) Refund() error { if w.consumed <= 0 { return nil } - return model.IncreaseUserQuota(w.userId, w.consumed, false) + amount := w.consumed + w.consumed = 0 + return model.IncreaseUserQuota(w.userId, amount, false) }
109-116: Consider adding the same defensive reset pattern toSubscriptionFunding.Refund.Similar to
WalletFunding,preConsumedisn't reset after a successful refund. WhileBillingSessionmanages idempotency externally, resettingpreConsumed = 0afterrefundWithRetrysucceeds would make the guard on line 110 effective for any direct re-call.🛡️ Suggested defensive reset
func (s *SubscriptionFunding) Refund() error { if s.preConsumed <= 0 { return nil } - return refundWithRetry(func() error { + err := refundWithRetry(func() error { return model.RefundSubscriptionPreConsume(s.requestId) }) + if err == nil { + s.preConsumed = 0 + } + return err }
84-100:PreConsumeignores itsamountparameter — document at the interface level or consider a design adjustment.The
FundingSource.PreConsume(amount int)contract implies the caller controls the amount, butSubscriptionFundingsilently ignores it in favor ofs.amount. This is documented locally (line 85) but can surprise future implementors or callers. If subscription always uses an internal amount, consider either:
- Making the interface's
PreConsumeparameterless and havingWalletFundingreceive its amount at construction time too (for symmetry), or- Documenting at the interface level that
amountis advisory.This is a minor API design nit — current behavior is correct.
120-137: Retry backoff is reasonable; consider logging failed attempts for observability.The linear backoff (200ms, 400ms) with 3 attempts is appropriate for transient DB failures. However, failed intermediate attempts are silently discarded. Adding a log on retry could help diagnose recurring transient issues in production.
service/billing_session.go (3)
71-106: Refund goroutine drops observability context.The
gin.Contextis correctly not captured in the goroutine (it's recycled), but the async refund now runs with no request ID or trace context, making it hard to correlate refund failures in logs with the original request.Consider capturing
relayInfo.RequestIdalongside the other locals and including it in thecommon.SysLogmessages inside the goroutine.
115-127: Type assertion breaksFundingSourceabstraction.Line 123 reaches into
*SubscriptionFundinginternals to checksub.preConsumed > 0. This couplesBillingSessionto a concrete funding type, which undermines the interface-based design.Consider adding a
HasPreConsumed() boolmethod to theFundingSourceinterface so this logic stays polymorphic:Sketch
- // 订阅可能在 tokenConsumed=0 时仍预扣了额度 - if sub, ok := s.funding.(*SubscriptionFunding); ok && sub.preConsumed > 0 { - return true - } + if s.funding.HasPreConsumed() { + return true + }
WalletFunding.HasPreConsumed()would simply returnfalse(or derive from its own state).
213-231: Same type-assertion pattern asneedsRefundLocked.If the
FundingSourceinterface is extended per the earlier suggestion,syncRelayInfocould similarly use an interface method (e.g.,SubscriptionInfo()) to avoid reaching into*SubscriptionFundinginternals. This would keep both call sites consistent.
| func SettleBilling(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, actualQuota int) error { | ||
| if relayInfo.Billing != nil { | ||
| preConsumed := relayInfo.Billing.GetPreConsumedQuota() | ||
| delta := actualQuota - preConsumed | ||
|
|
||
| res, err := model.PreConsumeUserSubscription(relayInfo.RequestId, relayInfo.UserId, relayInfo.OriginModelName, quotaType, subConsume) | ||
| if err != nil { | ||
| // revert token pre-consume when subscription fails | ||
| if preConsumedQuota > 0 && !relayInfo.IsPlayground { | ||
| _ = model.IncreaseTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, preConsumedQuota) | ||
| } | ||
| errMsg := err.Error() | ||
| if strings.Contains(errMsg, "no active subscription") || strings.Contains(errMsg, "subscription quota insufficient") { | ||
| return types.NewErrorWithStatusCode(fmt.Errorf("订阅额度不足或未配置订阅: %s", errMsg), types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog()) | ||
| } | ||
| return types.NewErrorWithStatusCode(fmt.Errorf("订阅预扣失败: %s", errMsg), types.ErrorCodeQueryDataError, http.StatusInternalServerError) | ||
| if delta > 0 { | ||
| logger.LogInfo(ctx, fmt.Sprintf("预扣费后补扣费：%s（实际消耗：%s，预扣费：%s）", | ||
| logger.FormatQuota(delta), | ||
| logger.FormatQuota(actualQuota), | ||
| logger.FormatQuota(preConsumed), | ||
| )) | ||
| } else if delta < 0 { | ||
| logger.LogInfo(ctx, fmt.Sprintf("预扣费后返还扣费：%s（实际消耗：%s，预扣费：%s）", | ||
| logger.FormatQuota(-delta), | ||
| logger.FormatQuota(actualQuota), | ||
| logger.FormatQuota(preConsumed), | ||
| )) | ||
| } else { | ||
| logger.LogInfo(ctx, fmt.Sprintf("预扣费与实际消耗一致，无需调整：%s（按次计费）", | ||
| logger.FormatQuota(actualQuota), | ||
| )) | ||
| } | ||
|
|
||
| relayInfo.BillingSource = BillingSourceSubscription | ||
| relayInfo.SubscriptionId = res.UserSubscriptionId | ||
| relayInfo.SubscriptionPreConsumed = res.PreConsumed | ||
| relayInfo.SubscriptionPostDelta = 0 | ||
| relayInfo.SubscriptionAmountTotal = res.AmountTotal | ||
| relayInfo.SubscriptionAmountUsedAfterPreConsume = res.AmountUsedAfter | ||
| if planInfo, err := model.GetSubscriptionPlanInfoByUserSubscriptionId(res.UserSubscriptionId); err == nil && planInfo != nil { | ||
| relayInfo.SubscriptionPlanId = planInfo.PlanId | ||
| relayInfo.SubscriptionPlanTitle = planInfo.PlanTitle | ||
| if err := relayInfo.Billing.Settle(actualQuota); err != nil { | ||
| return err | ||
| } | ||
| relayInfo.FinalPreConsumedQuota = preConsumedQuota | ||
|
|
||
| logger.LogInfo(c, fmt.Sprintf("用户 %d 使用订阅计费预扣：订阅=%d，token_quota=%d", relayInfo.UserId, res.PreConsumed, preConsumedQuota)) | ||
| // 发送额度通知 | ||
| if actualQuota != 0 { | ||
| checkAndSendQuotaNotify(relayInfo, actualQuota-preConsumed, preConsumed) | ||
| } | ||
| return nil | ||
| } | ||
|
|
||
| tryWallet := func() *types.NewAPIError { | ||
| relayInfo.BillingSource = BillingSourceWallet | ||
| relayInfo.SubscriptionId = 0 | ||
| relayInfo.SubscriptionPreConsumed = 0 | ||
| return PreConsumeQuota(c, preConsumedQuota, relayInfo) | ||
| } | ||
|
|
||
| switch pref { | ||
| case "subscription_only": | ||
| return trySubscription() | ||
| case "wallet_only": | ||
| return tryWallet() | ||
| case "wallet_first": | ||
| if err := tryWallet(); err != nil { | ||
| // only fallback for insufficient wallet quota | ||
| if err.GetErrorCode() == types.ErrorCodeInsufficientUserQuota { | ||
| return trySubscription() | ||
| } | ||
| return err | ||
| } | ||
| return nil | ||
| case "subscription_first": | ||
| fallthrough | ||
| default: | ||
| if err := trySubscription(); err != nil { | ||
| // fallback only when subscription not available/insufficient | ||
| if err.GetErrorCode() == types.ErrorCodeInsufficientUserQuota { | ||
| return tryWallet() | ||
| } | ||
| return err | ||
| } | ||
| return nil | ||
| // 回退：无 BillingSession 时使用旧路径 | ||
| quotaDelta := actualQuota - relayInfo.FinalPreConsumedQuota | ||
| if quotaDelta != 0 { | ||
| return PostConsumeQuota(relayInfo, quotaDelta, relayInfo.FinalPreConsumedQuota, true) | ||
| } | ||
| return nil |
There was a problem hiding this comment.
Fallback path skips quota notification when delta is zero.
In the Billing != nil path (line 62), quota notification fires whenever actualQuota != 0. In the fallback path (lines 68–72), when quotaDelta == 0 (actual equals pre-consumed), PostConsumeQuota is never called — so its internal notification logic is also skipped, even though FinalPreConsumedQuota may be non-zero. The old code would have called PostConsumeQuota(relayInfo, 0, preConsumedQuota, true) which still triggered the notification check.
If this fallback is only temporary and will be removed once all paths adopt BillingSession, this is likely acceptable. Otherwise, consider adding the notification call for the quotaDelta == 0 case.
🤖 Prompt for AI Agents
In `@service/billing.go` around lines 34 - 73, The fallback branch in
SettleBilling skips quota notifications when quotaDelta == 0; ensure the
notification still fires by invoking the same notification path as the
non-fallback branch. Modify SettleBilling so that when relayInfo.Billing == nil
and quotaDelta == 0 you either call PostConsumeQuota(relayInfo, 0,
relayInfo.FinalPreConsumedQuota, true) (to preserve old behavior) or explicitly
call checkAndSendQuotaNotify(relayInfo, 0, relayInfo.FinalPreConsumedQuota) so
that notifications run even when no quota change occurred.
- Settle 部分失败保护：新增 fundingSettled 标记，资金来源提交后 令牌调整失败不再导致 Refund 误退已结算的资金 - 订阅多扣费修复：trySubscription 传 subConsume 而非 preConsumedQuota 给 preConsume，保证三者（amount/preConsume/FinalPreConsumedQuota）一致 - 令牌回滚错误记录：preConsume 中 funding 失败时令牌回滚错误不再丢弃 - 移除钱包路径死代码：用户额度不足的 strings.Contains 匹配不可能命中 - WalletFunding.Refund 不重试：IncreaseUserQuota 非幂等，重试会多退
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@service/billing_session.go`:
- Around line 80-114: Settle must also guard against an in-flight refund: update
the Settle method to check s.refunded (in addition to s.settled and
needsRefundLocked()) under the same mutex so Settle returns early if a refund
has been started, preventing funding.Settle(delta) from running after Refund
began; additionally, in Refund keep the async behavior but re-check the session
state inside the goroutine under s.mu (or re-acquire the mutex) before calling
funding.Refund() — for example, verify that !s.settled && s.refunded (or that
funding hasn't been settled) while holding s.mu before invoking
funding.Refund(), and only then perform the refund and token quota increase to
avoid the race between funding.Settle(delta) and funding.Refund().
🧹 Nitpick comments (2)
service/billing_session.go (2)
123-136:needsRefundLockeduses a concrete type assertion, coupling it toSubscriptionFunding.Line 132 uses
s.funding.(*SubscriptionFunding)to check whether the subscription has pre-consumed quota even whentokenConsumed == 0. This breaks theFundingSourceinterface boundary and will silently miss any futureFundingSourceimplementation that has a similar pre-consume-without-token pattern.Consider adding a method like
HasPendingPreConsume() boolto theFundingSourceinterface so the check can be polymorphic.
258-285:tryWalletmutatesrelayInfo.UserQuotaeven on failure — benign but worth noting.Line 275 sets
relayInfo.UserQuota = userQuotabeforepreConsume. IfpreConsumefails and the factory falls back totrySubscription(in thewallet_firstpath),relayInfo.UserQuotaretains the stale snapshot. This doesn't cause bugs today sincetrySubscriptiondoesn't rely onUserQuota, but it's a subtle coupling that could surprise future maintainers.
| func (s *BillingSession) Refund(c *gin.Context) { | ||
| s.mu.Lock() | ||
| if s.settled || s.refunded || !s.needsRefundLocked() { | ||
| s.mu.Unlock() | ||
| return | ||
| } | ||
| s.refunded = true | ||
| s.mu.Unlock() | ||
|
|
||
| logger.LogInfo(c, fmt.Sprintf("用户 %d 请求失败, 返还预扣费（token_quota=%s, funding=%s）", | ||
| s.relayInfo.UserId, | ||
| logger.FormatQuota(s.tokenConsumed), | ||
| s.funding.Source(), | ||
| )) | ||
|
|
||
| // 复制需要的值到闭包中 | ||
| tokenId := s.relayInfo.TokenId | ||
| tokenKey := s.relayInfo.TokenKey | ||
| isPlayground := s.relayInfo.IsPlayground | ||
| tokenConsumed := s.tokenConsumed | ||
| funding := s.funding | ||
|
|
||
| gopool.Go(func() { | ||
| // 1) 退还资金来源 | ||
| if err := funding.Refund(); err != nil { | ||
| common.SysLog("error refunding billing source: " + err.Error()) | ||
| } | ||
| // 2) 退还令牌额度 | ||
| if tokenConsumed > 0 && !isPlayground { | ||
| if err := model.IncreaseTokenQuota(tokenId, tokenKey, tokenConsumed); err != nil { | ||
| common.SysLog("error refunding token quota: " + err.Error()) | ||
| } | ||
| } | ||
| }) | ||
| } |
There was a problem hiding this comment.
Settle does not check s.refunded, allowing a funding double-operation if both paths overlap.
Refund sets s.refunded = true and releases the lock before the async goroutine runs funding.Refund(). If Settle is called between the lock release (line 87) and the goroutine execution, Settle will proceed because it only checks s.settled (line 42), not s.refunded. This results in both funding.Settle(delta) and funding.Refund() executing against the same funding source.
Even if the current call-sites guarantee mutual exclusion, the session abstraction itself should be safe against this since it already uses a mutex for exactly this purpose.
Proposed fix — guard Settle against in-flight refund
func (s *BillingSession) Settle(actualQuota int) error {
s.mu.Lock()
defer s.mu.Unlock()
- if s.settled {
+ if s.settled || s.refunded {
return nil
}Additionally, in Refund, the async goroutine calls funding.Refund() without holding the lock. If Settle runs concurrently (before the goroutine starts), fundingSettled is set to true, but the goroutine still proceeds with funding.Refund(). Consider moving the fundingSettled check inside the goroutine:
Proposed fix — re-check fundingSettled inside goroutine
gopool.Go(func() {
+ s.mu.Lock()
+ if s.fundingSettled {
+ s.mu.Unlock()
+ // Funding was settled between Refund dispatch and goroutine execution;
+ // only refund token quota.
+ } else {
+ s.mu.Unlock()
+ if err := funding.Refund(); err != nil {
+ common.SysLog("error refunding billing source: " + err.Error())
+ }
+ }
- // 1) 退还资金来源
- if err := funding.Refund(); err != nil {
- common.SysLog("error refunding billing source: " + err.Error())
- }
// 2) 退还令牌额度
if tokenConsumed > 0 && !isPlayground {🤖 Prompt for AI Agents
In `@service/billing_session.go` around lines 80 - 114, Settle must also guard
against an in-flight refund: update the Settle method to check s.refunded (in
addition to s.settled and needsRefundLocked()) under the same mutex so Settle
returns early if a refund has been started, preventing funding.Settle(delta)
from running after Refund began; additionally, in Refund keep the async behavior
but re-check the session state inside the goroutine under s.mu (or re-acquire
the mutex) before calling funding.Refund() — for example, verify that !s.settled
&& s.refunded (or that funding hasn't been settled) while holding s.mu before
invoking funding.Refund(), and only then perform the refund and token quota
increase to avoid the race between funding.Settle(delta) and funding.Refund().
…-session refactor: 抽象统一计费会话 BillingSession
1 participant
将散落在多个文件中的预扣费/结算/退款逻辑抽象为统一的 BillingSession 生命周期管理：
修复的 Bug：
Summary by CodeRabbit
Refactor
Bug Fixes