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 异步记账器：缓冲队列 + 批量事务落库。
//
// 多个 worker 按用户分片：同一用户的记录始终交给同一个 worker，落库顺序与扣款
// 流水顺序保持一致；不同用户并行落库，PostgreSQL 上不再被单 goroutine 串行化。
type Recorder struct {
	db     *gorm.DB
	shards []chan *store.UsageLog
	// workers 本实例的落库并发度（SQLite 为 1，PG 为 flushWorkersPG）
	workers int
	wg      sync.WaitGroup
	closed  chan struct{}
	once    sync.Once // 保证 Close 只执行一次（main 的 defer Shutdown 与显式 Close 双调用）

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

const (
	batchSize = 32
	// flushWorkersPG 落库 worker 数（按用户分片）。SQLite 只有一个写者，
	// 多 worker 并发写只会互相抢锁、把 p99 拖长（实测 4 worker 时 p99 19–34ms），
	// 所以按驱动自适应：PG 用多个，SQLite 只用一个。
	flushWorkersPG = 4
)

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

func NewRecorder(db *gorm.DB) *Recorder {
	return newRecorder(db, workersFor(db))
}

// workersFor 按驱动决定落库并发度：SQLite 只有一个写者，多 worker 只会抢锁。
func workersFor(db *gorm.DB) int {
	if db != nil && db.Dialector != nil && db.Dialector.Name() == "postgres" {
		return flushWorkersPG
	}
	return 1
}

// newRecorder 指定并发度构造（测试用：在 SQLite 上也能验证多分片语义）。
func newRecorder(db *gorm.DB, workers int) *Recorder {
	if workers < 1 {
		workers = 1
	}
	r := &Recorder{
		db:      db,
		closed:  make(chan struct{}),
		workers: workers,
		shards:  make([]chan *store.UsageLog, workers),
	}
	for i := range r.shards {
		// 每片容量与原来的单队列等效（512 总量按片均分，至少 64）
		size := 512 / workers
		if size < 64 {
			size = 64
		}
		r.shards[i] = make(chan *store.UsageLog, size)
		r.wg.Add(1)
		go r.run(i)
	}
	return r
}

// shardOf 用户 → 队列分片：同一用户恒定落在同一片，保证其记账顺序。
func (r *Recorder) shardOf(userID uint64) chan *store.UsageLog {
	return r.shards[userID%uint64(len(r.shards))]
}

// Record 提交一条用量（非阻塞；队列满或已停止接收时同步写入，保证不丢账）。
func (r *Recorder) Record(l *store.UsageLog) {
	r.mu.RLock()
	if r.stopped {
		r.mu.RUnlock()
		r.flushRetry([]*store.UsageLog{l})
		return
	}
	ch := r.shardOf(l.UserID)
	select {
	case 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(shard int) {
	defer r.wg.Done()
	ch := r.shards[shard]
	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 := <-ch:
			buf = append(buf, l)
			if len(buf) >= batchSize {
				flushBuf()
			}
		case <-r.closed:
			// 排空残留：Shutdown 已等 HTTP 在途请求退出，但仍有极小窗口的
			// Record 已拿到 RLock 还没写进 channel，这里等它落地再收尾。
			for {
				select {
				case l := <-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）。
//
// 按批聚合而不是逐行：每行原本 5 条语句（查用户、查流水幂等、扣余额、写流水、
// upsert 日聚合），32 条一批就是 160 条；同一用户的并发请求还会在 users 行锁上排队。
// 现在按用户聚合一次扣款、批量写流水、按 (user, model, 日期) 聚合一次日表。
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
		}
		return r.bill(tx, logs)
	})
}

// bill 对一批明细计费：按用户聚合扣款，批量写流水，按 (user, model, 日期) 聚合日表。
func (r *Recorder) bill(tx *gorm.DB, logs []*store.UsageLog) error {
	// 1. 可计费行：success/canceled 且有成本
	billable := make([]*store.UsageLog, 0, len(logs))
	refs := make([]string, 0, len(logs))
	for _, l := range logs {
		if (l.Status != store.UsageStatusSuccess && l.Status != store.UsageStatusCanceled) || l.Cost <= 0 {
			continue
		}
		billable = append(billable, l)
		refs = append(refs, usageRefID(l.TraceID))
	}
	if len(billable) == 0 {
		return nil
	}

	// 2. 已记过账的 ref（重试/重复 flush）：一条查询取代每行一次 COUNT
	already := map[string]bool{}
	var billed []string
	if err := tx.Model(&store.BalanceLog{}).Where("ref_id IN ?", refs).Pluck("ref_id", &billed).Error; err != nil {
		return err
	}
	for _, ref := range billed {
		already[ref] = true
	}

	// 3. 过滤重复（含同批内同 trace），按用户分组
	byUser := map[uint64][]*store.UsageLog{}
	order := make([]uint64, 0, len(byUser))
	seen := map[string]bool{}
	for _, l := range billable {
		ref := usageRefID(l.TraceID)
		if already[ref] || seen[ref] {
			continue
		}
		seen[ref] = true
		if _, ok := byUser[l.UserID]; !ok {
			order = append(order, l.UserID)
		}
		byUser[l.UserID] = append(byUser[l.UserID], l)
	}
	if len(byUser) == 0 {
		return nil
	}

	// 4. 每个用户一次扣款，流水批量插入
	ledger := make([]store.BalanceLog, 0, len(billable))
	for _, uid := range order {
		rows := byUser[uid]
		total := 0.0
		for _, l := range rows {
			total += l.Cost
		}
		var user store.User
		if err := tx.First(&user, uid).Error; err != nil {
			// 用户被删了：明细已插入，扣费无处可扣，吵一声但不能拖垮整批
			log.Printf("usage: user %d not found, %d records unbilled: %v", uid, len(rows), err)
			continue
		}
		// 扣费用 SQL 侧原子递减（balance = balance - ?）：SQLite 驱动不支持
		// SELECT ... FOR UPDATE，Go 侧读-改-写在并发下会互相覆盖丢扣费。
		// 余额允许为负（流式请求不中断；后续请求被 402 拒）。
		if err := tx.Model(&store.User{}).Where("id = ?", uid).
			Update("balance", gorm.Expr("balance - ?", total)).Error; err != nil {
			return err
		}
		// 流水里的 balance_after 是批内累计值，与逐行扣款时的语义一致
		running := user.Balance
		for _, l := range rows {
			running -= l.Cost
			ledger = append(ledger, store.BalanceLog{
				UserID:       uid,
				Change:       -l.Cost,
				BalanceAfter: running,
				Type:         store.BalanceTypeUsage,
				RefID:        usageRefID(l.TraceID),
				Remark:       "usage: " + l.ModelName,
			})
		}
	}
	if len(ledger) > 0 {
		// ref_id 唯一索引兜底：并发下别处已写入同 ref 时跳过，不整批失败也不重复扣款
		if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&ledger).Error; err != nil {
			return err
		}
	}

	// 5. 日聚合：按 (user, model, 日期) 汇总后一次 upsert
	type dailyKey struct {
		user  uint64
		model uint64
		date  string
	}
	agg := map[dailyKey]*store.UsageDaily{}
	dailyOrder := make([]dailyKey, 0, len(billable))
	for _, l := range billable {
		ref := usageRefID(l.TraceID)
		if !seen[ref] {
			continue // 重复 trace，不计入日聚合
		}
		k := dailyKey{user: l.UserID, model: l.ModelID, date: l.CreatedAt.UTC().Format("2006-01-02")}
		d, ok := agg[k]
		if !ok {
			d = &store.UsageDaily{UserID: k.user, ModelID: k.model, Date: k.date}
			agg[k] = d
			dailyOrder = append(dailyOrder, k)
		}
		d.Requests++
		d.InputTokens += l.InputTokens
		d.OutputTokens += l.OutputTokens
		d.CacheReadTokens += l.CacheReadTokens
		d.Cost += l.Cost
	}
	for _, k := range dailyOrder {
		d := agg[k]
		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 + ?", d.Requests),
				"input_tokens":      gorm.Expr("input_tokens + ?", d.InputTokens),
				"output_tokens":     gorm.Expr("output_tokens + ?", d.OutputTokens),
				"cache_read_tokens": gorm.Expr("cache_read_tokens + ?", d.CacheReadTokens),
				"cost":              gorm.Expr("cost + ?", d.Cost),
			}),
		}).Create(d).Error; err != nil {
			return err
		}
	}
	return nil
}

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