dudu/server/internal/quota/quota.go

// 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
	})
}