// Package quota 用量计量与扣减:先扣当日试用(180s/天),余下扣余额。 // Redis 实时计数(Lua 原子),会话结束异步落库 ledger / asr_sessions。 // 设计见 doc/backend-architecture.html 第五章。 package quota import ( "context" "time" "github.com/redis/go-redis/v9" "gorm.io/gorm" "dudu/server/internal/store" "dudu/server/pkg/protocol" ) type Manager struct { rdb *redis.Client db *gorm.DB } func New(rdb *redis.Client, db *gorm.DB) *Manager { return &Manager{rdb: rdb, db: db} } // consumeScript 原子扣减:先试用桶后余额桶。 // KEYS[1]=trial key, KEYS[2]=balance key // ARGV[1]=need 秒, ARGV[2]=trial 每日上限 // 返回 {trialPart, balancePart, remainBalance, trialUsed};不足部分允许扣穿到 0(识别中增量扣减, // 是否拒绝由 Precheck 与网关在句间判断——"扣穿当句不掐断")。 var consumeScript = redis.NewScript(` local trialKey, balKey = KEYS[1], KEYS[2] local need, trialLimit = tonumber(ARGV[1]), tonumber(ARGV[2]) local trialUsed = tonumber(redis.call('GET', trialKey) or '0') local bal = tonumber(redis.call('GET', balKey) or '0') local trialPart = math.min(need, math.max(0, trialLimit - trialUsed)) local balancePart = math.min(need - trialPart, math.max(0, bal)) if trialPart > 0 then redis.call('INCRBY', trialKey, trialPart) redis.call('EXPIRE', trialKey, 172800) end if balancePart > 0 then redis.call('DECRBY', balKey, balancePart) end return {trialPart, balancePart, bal - balancePart, trialUsed + trialPart} `) type Snapshot struct { BalanceSeconds int64 TrialUsedToday int } // ensureBalanceLoaded 余额键懒加载:缺失时从 DB 冗余列加载(SET NX 防并发覆盖)。 func (m *Manager) ensureBalanceLoaded(ctx context.Context, uid string) error { key := store.KeyQuotaBalance(uid) if exists, err := m.rdb.Exists(ctx, key).Result(); err != nil || exists == 1 { return err } var u store.User if err := m.db.WithContext(ctx).Select("balance_seconds").First(&u, "id = ?", uid).Error; err != nil { return err } return m.rdb.SetNX(ctx, key, u.BalanceSeconds, 0).Err() } // Get 当前余额与今日试用(驱动 /v1/me 与 usage 帧)。 func (m *Manager) Get(ctx context.Context, uid string) (Snapshot, error) { if err := m.ensureBalanceLoaded(ctx, uid); err != nil { return Snapshot{}, err } day := store.Day(time.Now()) bal, err := m.rdb.Get(ctx, store.KeyQuotaBalance(uid)).Int64() if err != nil && err != redis.Nil { return Snapshot{}, err } trial, err := m.rdb.Get(ctx, store.KeyQuotaTrial(uid, day)).Int() if err != nil && err != redis.Nil { return Snapshot{}, err } return Snapshot{BalanceSeconds: bal, TrialUsedToday: trial}, nil } // Precheck start 预检:试用剩余 + 余额 ≤ 0 → 拒绝(QUOTA_EXCEEDED)。 func (m *Manager) Precheck(ctx context.Context, uid string) (bool, Snapshot, error) { s, err := m.Get(ctx, uid) if err != nil { return false, s, err } trialRemain := protocol.TrialDailySeconds - s.TrialUsedToday return trialRemain > 0 || s.BalanceSeconds > 0, s, nil } type ConsumeResult struct { TrialPart int BalancePart int BalanceSeconds int64 // 扣减后余额 TrialUsedToday int } // Exhausted 本次扣减后余额与今日试用是否均已耗尽(扣穿)。 // 网关据此在句间下发 QUOTA_EXCEEDED 并优雅结束会话("扣穿当句不掐断")。 func (r ConsumeResult) Exhausted() bool { return r.BalanceSeconds <= 0 && r.TrialUsedToday >= protocol.TrialDailySeconds } // Consume 扣减 seconds 秒(识别中每 2s 增量调用 + 会话结束补差)。 func (m *Manager) Consume(ctx context.Context, uid string, seconds int) (ConsumeResult, error) { if err := m.ensureBalanceLoaded(ctx, uid); err != nil { return ConsumeResult{}, err } day := store.Day(time.Now()) vals, err := consumeScript.Run(ctx, m.rdb, []string{store.KeyQuotaTrial(uid, day), store.KeyQuotaBalance(uid)}, seconds, protocol.TrialDailySeconds).Int64Slice() if err != nil { return ConsumeResult{}, err } return ConsumeResult{ TrialPart: int(vals[0]), BalancePart: int(vals[1]), BalanceSeconds: vals[2], TrialUsedToday: int(vals[3]), }, nil } // Exhausted 判断是否已无可用时长(网关句间检查,决定何时下发 QUOTA_EXCEEDED)。 func (s Snapshot) Exhausted() bool { return s.BalanceSeconds <= 0 && s.TrialUsedToday >= protocol.TrialDailySeconds } // CreditPurchase 购买入账:事务内写 ledger + 更新冗余列,并同步 Redis 余额键。 func (m *Manager) CreditPurchase(ctx context.Context, uid string, seconds int64, orderID string) error { err := m.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { if err := tx.Create(&store.BalanceLedger{ UserID: uid, DeltaSeconds: seconds, Reason: store.LedgerPurchase, OrderID: &orderID, }).Error; err != nil { return err } return tx.Model(&store.User{}).Where("id = ?", uid). UpdateColumn("balance_seconds", gorm.Expr("balance_seconds + ?", seconds)).Error }) if err != nil { return err } // Redis 键存在则同步增加;不存在则下次懒加载 if exists, _ := m.rdb.Exists(ctx, store.KeyQuotaBalance(uid)).Result(); exists == 1 { return m.rdb.IncrBy(ctx, store.KeyQuotaBalance(uid), seconds).Err() } return nil } // SettleSession 会话结束落库:用量 ledger(负向)+ asr_sessions 明细 + 冗余列。 func (m *Manager) SettleSession(ctx context.Context, sess store.ASRSession) error { return m.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { if err := tx.Create(&sess).Error; err != nil { return err } if sess.BalancePart > 0 { sid := sess.ID if err := tx.Create(&store.BalanceLedger{ UserID: sess.UserID, DeltaSeconds: -int64(sess.BalancePart), Reason: store.LedgerUsage, SessionID: &sid, }).Error; err != nil { return err } if err := tx.Model(&store.User{}).Where("id = ?", sess.UserID). UpdateColumn("balance_seconds", gorm.Expr("balance_seconds - ?", sess.BalancePart)).Error; err != nil { return err } } if sess.TrialPart > 0 { day := store.Day(sess.CreatedAt) if err := tx.Exec(`INSERT INTO trial_usages (user_id, date, used_seconds) VALUES (?, ?, ?) ON CONFLICT (user_id, date) DO UPDATE SET used_seconds = trial_usages.used_seconds + ?`, sess.UserID, day, sess.TrialPart, sess.TrialPart).Error; err != nil { return err } } return nil }) }