openteam/server/internal/usage/usage.go

// Package usage 异步记账：请求完成后写入 usage_logs，批量落库（PLANNING §3.2）。
// 每个请求在 flush 时同步完成：写明细 + 扣余额 + 写流水 + 日聚合。
package usage

import (
	"log"
	"sync"
	"time"

	"github.com/openteam/server/internal/store"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

// Recorder 异步记账器：缓冲队列 + 批量事务落库。
type Recorder struct {
	db     *gorm.DB
	ch     chan *store.UsageLog
	wg     sync.WaitGroup
	closed chan struct{}
	once   sync.Once // 保证 Close 只执行一次（main 的 defer Shutdown 与显式 Close 双调用）

	// mu 保护 stopped：Close 之后 Record 不再进队（ch 永不 close，杜绝 send-on-closed panic），
	// 改走同步落库兜底。
	mu      sync.RWMutex
	stopped bool
}

const batchSize = 32

// flushRetries 批量落库的瞬时错误重试次数（SQLite busy / PG 连接抖动）。
const flushRetries = 3

func NewRecorder(db *gorm.DB) *Recorder {
	r := &Recorder{
		db:     db,
		ch:     make(chan *store.UsageLog, 512),
		closed: make(chan struct{}),
	}
	r.wg.Add(1)
	go r.run()
	return r
}

// Record 提交一条用量（非阻塞；队列满或已停止接收时同步写入，保证不丢账）。
func (r *Recorder) Record(l *store.UsageLog) {
	r.mu.RLock()
	if r.stopped {
		r.mu.RUnlock()
		r.flushRetry([]*store.UsageLog{l})
		return
	}
	select {
	case r.ch <- l:
		r.mu.RUnlock()
	default:
		r.mu.RUnlock()
		r.flushRetry([]*store.UsageLog{l})
	}
}

// Close 停止接收新记账并等待后台 goroutine 排空队列后返回。幂等。
func (r *Recorder) Close() {
	r.once.Do(func() {
		r.mu.Lock()
		r.stopped = true
		r.mu.Unlock()
		close(r.closed)
		r.wg.Wait()
		// 注意：r.ch 故意不 close——在途请求的 Record 会安全落到同步落库分支
	})
}

func (r *Recorder) run() {
	defer r.wg.Done()
	buf := make([]*store.UsageLog, 0, batchSize)
	flushBuf := func() {
		if len(buf) > 0 {
			r.flushRetry(buf)
			buf = buf[:0]
		}
	}
	tick := time.NewTicker(2 * time.Second)
	defer tick.Stop()
	for {
		select {
		case l := <-r.ch:
			buf = append(buf, l)
			if len(buf) >= batchSize {
				flushBuf()
			}
		case <-r.closed:
			// 排空残留：Shutdown 已等 HTTP 在途请求退出，但仍有极小窗口的
			// Record 已拿到 RLock 还没写进 channel，这里等它落地再收尾。
			for {
				select {
				case l := <-r.ch:
					buf = append(buf, l)
					if len(buf) >= batchSize {
						flushBuf()
					}
				default:
					flushBuf()
					return
				}
			}
		case <-tick.C:
			flushBuf()
		}
	}
}

// flushRetry 瞬时 DB 错误重试：整批丢弃等于丢账，这里宁可多试几次。
func (r *Recorder) flushRetry(logs []*store.UsageLog) {
	if len(logs) == 0 {
		return
	}
	var err error
	for i := 0; i < flushRetries; i++ {
		if err = r.flush(logs); err == nil {
			return
		}
		time.Sleep(time.Duration(i+1) * 200 * time.Millisecond)
	}
	log.Printf("usage: batch write failed after %d retries (%d records lost): %v", flushRetries, len(logs), err)
}

// flush 批量插入用量明细，并同步更新余额、余额流水与日聚合。
// 记账口径：单次成本 = in×in_price + out×out_price + cache_read×cache_read_price（每百万 token）。
func (r *Recorder) flush(logs []*store.UsageLog) error {
	if len(logs) == 0 {
		return nil
	}
	return r.db.Transaction(func(tx *gorm.DB) error {
		if err := tx.Create(logs).Error; err != nil {
			return err
		}
		for _, l := range logs {
			// 计费范围：success（正常完成）与 canceled（流式中断，按已生成部分收费）
			if (l.Status != store.UsageStatusSuccess && l.Status != store.UsageStatusCanceled) || l.Cost <= 0 {
				continue
			}
			if err := r.billUser(tx, l); err != nil {
				// 明细已插入但扣费失败：宁可吵也不能静默漏账
				log.Printf("usage: billing failed trace=%s user=%d cost=%.6f: %v", l.TraceID, l.UserID, l.Cost, err)
			}
		}
		return nil
	})
}

// billUser 对单条明细扣余额、写流水、更新日聚合。
// 扣费用 SQL 侧原子递减（balance = balance - cost）：SQLite 驱动不支持 SELECT ... FOR UPDATE，
// Go 侧读-改-写在并发下会互相覆盖丢扣费。
func (r *Recorder) billUser(tx *gorm.DB, l *store.UsageLog) error {
	var user store.User
	if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).First(&user, l.UserID).Error; err != nil {
		return err
	}
	// 余额允许为负（流式请求不中断；后续请求被 402 拒），所以 WHERE 只按 id 匹配
	// 同一 trace 已记过账（重试/重复 flush）则跳过：流水 ref_id 有唯一索引，
	// 但扣款本身没有，不挡住会重复扣钱
	var billed int64
	tx.Model(&store.BalanceLog{}).Where("ref_id = ?", usageRefID(l.TraceID)).Count(&billed)
	if billed > 0 {
		return nil
	}
	res := tx.Model(&store.User{}).Where("id = ?", l.UserID).Update("balance", gorm.Expr("balance - ?", l.Cost))
	if res.Error != nil {
		return res.Error
	}
	if res.RowsAffected == 0 {
		return gorm.ErrRecordNotFound
	}
	newBalance := user.Balance - l.Cost
	if err := tx.Create(&store.BalanceLog{
		UserID:       l.UserID,
		Change:       -l.Cost,
		BalanceAfter: newBalance,
		Type:         store.BalanceTypeUsage,
		RefID:        usageRefID(l.TraceID),
		Remark:       "usage: " + l.ModelName,
	}).Error; err != nil {
		return err
	}

	// 日聚合 upsert
	date := l.CreatedAt.UTC().Format("2006-01-02")
	if err := tx.Clauses(clause.OnConflict{
		Columns: []clause.Column{{Name: "user_id"}, {Name: "model_id"}, {Name: "date"}},
		DoUpdates: clause.Assignments(map[string]any{
			"requests":          gorm.Expr("requests + 1"),
			"input_tokens":      gorm.Expr("input_tokens + ?", l.InputTokens),
			"output_tokens":     gorm.Expr("output_tokens + ?", l.OutputTokens),
			"cache_read_tokens": gorm.Expr("cache_read_tokens + ?", l.CacheReadTokens),
			"cost":              gorm.Expr("cost + ?", l.Cost),
		}),
	}).Create(&store.UsageDaily{
		UserID:          l.UserID,
		ModelID:         l.ModelID,
		Date:            date,
		Requests:        1,
		InputTokens:     l.InputTokens,
		OutputTokens:    l.OutputTokens,
		CacheReadTokens: l.CacheReadTokens,
		Cost:            l.Cost,
	}).Error; err != nil {
		return err
	}
	return nil
}

func usageRefID(traceID string) string {
	if traceID == "" {
		traceID = "unknown"
	}
	return "usage:" + traceID
}