fix: 记账可靠性——关停排空队列、失败重试、扣费不再静默；余额改 SQL 原子递减
- Recorder.Close 改为停止接收+排空残留，channel 不再 close（杜绝 send-on-closed panic），Close 后的 Record 走同步落库 - flush 瞬时 DB 错误重试 3 次，不再整批丢弃；事务内扣费/流水/日聚合失败一律打日志，明细与扣费不再静默脱节 - 扣费与管理员调余额改 SQL 侧原子加减（SQLite 驱动不支持 FOR UPDATE，读-改-写并发会丢扣费）；调余额 ref_id 改纳秒不再撞唯一索引 - 最后一个启用的管理员不可被降级/禁用（管理面锁死后 reset-admin-password 救不回）
This commit is contained in:
Sakurasan committed
1 parent
7c597a6aa0
commit
984a31cbc5
3 files changed
+292
-65
No files matched your search
+116
-61
@@ -19,10 +19,18 @@ type Recorder struct {
|
||||
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,
|
||||
@@ -34,61 +42,89 @@ func NewRecorder(db *gorm.DB) *Recorder {
|
||||
return r
|
||||
}
|
||||
|
||||
// Record 提交一条用量（非阻塞；队列满时同步写入，保证不丢账）。
|
||||
// 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:
|
||||
if err := r.flush([]*store.UsageLog{l}); err != nil {
|
||||
log.Printf("usage: sync write failed: %v", err)
|
||||
}
|
||||
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()
|
||||
close(r.ch)
|
||||
// 注意：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, ok := <-r.ch:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
case l := <-r.ch:
|
||||
buf = append(buf, l)
|
||||
if len(buf) >= batchSize {
|
||||
if err := r.flush(buf); err != nil {
|
||||
log.Printf("usage: batch write failed: %v", err)
|
||||
}
|
||||
buf = buf[:0]
|
||||
flushBuf()
|
||||
}
|
||||
case <-r.closed:
|
||||
if len(buf) > 0 {
|
||||
if err := r.flush(buf); err != nil {
|
||||
log.Printf("usage: final batch write failed: %v", err)
|
||||
// 排空残留：Shutdown 已等 HTTP 在途请求退出，但仍有极小窗口的
|
||||
// Record 已拿到 RLock 还没写进 channel，这里等它落地再收尾。
|
||||
for {
|
||||
select {
|
||||
case l := <-r.ch:
|
||||
buf = append(buf, l)
|
||||
if len(buf) >= batchSize {
|
||||
flushBuf()
|
||||
}
|
||||
default:
|
||||
flushBuf()
|
||||
return
|
||||
}
|
||||
}
|
||||
return
|
||||
case <-tick.C:
|
||||
if len(buf) > 0 {
|
||||
if err := r.flush(buf); err != nil {
|
||||
log.Printf("usage: batch write failed: %v", err)
|
||||
}
|
||||
buf = buf[:0]
|
||||
}
|
||||
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 {
|
||||
@@ -104,50 +140,69 @@ func (r *Recorder) flush(logs []*store.UsageLog) error {
|
||||
if (l.Status != store.UsageStatusSuccess && l.Status != store.UsageStatusCanceled) || l.Cost <= 0 {
|
||||
continue
|
||||
}
|
||||
// 扣余额（余额可为负：流式请求不中断；后续请求被拒）
|
||||
var user store.User
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).First(&user, l.UserID).Error; err != nil {
|
||||
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)
|
||||
}
|
||||
newBalance := user.Balance - l.Cost
|
||||
if err := tx.Model(&store.User{}).Where("id = ?", l.UserID).Update("balance", newBalance).Error; err != nil {
|
||||
continue
|
||||
}
|
||||
tx.Create(&store.BalanceLog{
|
||||
UserID: l.UserID,
|
||||
Change: -l.Cost,
|
||||
BalanceAfter: newBalance,
|
||||
Type: store.BalanceTypeUsage,
|
||||
RefID: usageRefID(l.TraceID),
|
||||
Remark: "usage: " + l.ModelName,
|
||||
})
|
||||
|
||||
// 日聚合 upsert
|
||||
date := l.CreatedAt.UTC().Format("2006-01-02")
|
||||
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,
|
||||
})
|
||||
}
|
||||
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 匹配
|
||||
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"
|
||||
|
||||
@@ -0,0 +1,143 @@
|
||||
package usage
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/openteam/server/internal/store"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func TestCloseDrainsQueue(t *testing.T) {
|
||||
db := openTestDB(t)
|
||||
u := seedUser(t, db, 100)
|
||||
|
||||
rec := NewRecorder(db)
|
||||
// 快速灌入超过一批的记录（batchSize=32），大部分会留在 channel 里
|
||||
for i := 0; i < 100; i++ {
|
||||
rec.Record(&store.UsageLog{
|
||||
TraceID: fmt.Sprintf("t%d", i), UserID: u, ModelName: "m1",
|
||||
Status: store.UsageStatusError, ErrorCode: strPtr("test"), // error 状态不触发扣费分支
|
||||
CreatedAt: time.Now().UTC(),
|
||||
})
|
||||
}
|
||||
rec.Close()
|
||||
|
||||
var n int64
|
||||
db.Table("usage_logs").Count(&n)
|
||||
if n != 100 {
|
||||
t.Fatalf("Close 应排空队列全部落库，实际 %d/100", n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecordAfterCloseDoesNotPanic(t *testing.T) {
|
||||
db := openTestDB(t)
|
||||
u := seedUser(t, db, 100)
|
||||
|
||||
rec := NewRecorder(db)
|
||||
rec.Close()
|
||||
// Shutdown 超时后被强断的在途请求仍可能调 Record：不允许 send-on-closed panic
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
rec.Record(&store.UsageLog{
|
||||
TraceID: "late", UserID: u, ModelName: "m1",
|
||||
Status: store.UsageStatusError, ErrorCode: strPtr("late"),
|
||||
CreatedAt: time.Now().UTC(),
|
||||
})
|
||||
}()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("Record 在 Close 之后调用发生阻塞（疑似 panic 被吞）")
|
||||
}
|
||||
|
||||
var n int64
|
||||
db.Table("usage_logs").Where("trace_id = ?", "late").Count(&n)
|
||||
if n != 1 {
|
||||
t.Fatalf("Close 后的 Record 应同步落库，实际 %d 条", n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCloseIsIdempotent(t *testing.T) {
|
||||
db := openTestDB(t)
|
||||
seedUser(t, db, 0)
|
||||
rec := NewRecorder(db)
|
||||
rec.Close()
|
||||
rec.Close() // serve.go 的 defer Shutdown 与显式 Close 双调用
|
||||
}
|
||||
|
||||
func TestBillingDeductsBalanceAtomically(t *testing.T) {
|
||||
db := openTestDB(t)
|
||||
u := seedUser(t, db, 1)
|
||||
m := seedModel(t, db)
|
||||
|
||||
rec := NewRecorder(db)
|
||||
rec.Record(&store.UsageLog{
|
||||
TraceID: "bill", UserID: u, ModelID: m, ModelName: "m1",
|
||||
InputTokens: 1000, OutputTokens: 1000,
|
||||
InputPrice: 0.15, OutputPrice: 0.60,
|
||||
Cost: 0.15*1000/1e6 + 0.60*1000/1e6,
|
||||
Status: store.UsageStatusSuccess,
|
||||
LatencyMS: 5,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
})
|
||||
rec.Close()
|
||||
|
||||
var user store.User
|
||||
if err := db.First(&user, u).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := 1 - (0.15*1000/1e6 + 0.60*1000/1e6)
|
||||
if diff := user.Balance - want; diff > 1e-9 || diff < -1e-9 {
|
||||
t.Fatalf("余额应扣到 %.10f，实际 %.10f", want, user.Balance)
|
||||
}
|
||||
var bl store.BalanceLog
|
||||
if err := db.Where("user_id = ?", u).First(&bl).Error; err != nil {
|
||||
t.Fatalf("余额流水应落库: %v", err)
|
||||
}
|
||||
if bl.Change >= 0 {
|
||||
t.Fatalf("流水应为负数扣费，实际 %v", bl.Change)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// helpers
|
||||
|
||||
func openTestDB(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
dsn := "/tmp/usage_test_" + t.Name() + "_" + time.Now().Format("150405.000000000") + ".db"
|
||||
db, err := store.Open("sqlite", dsn)
|
||||
if err != nil {
|
||||
t.Fatalf("open db: %v", err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
sqlDB, _ := db.DB()
|
||||
sqlDB.Close()
|
||||
os.Remove(dsn)
|
||||
})
|
||||
return db
|
||||
}
|
||||
|
||||
func seedUser(t *testing.T, db *gorm.DB, balance float64) uint64 {
|
||||
t.Helper()
|
||||
u := store.User{Username: "u" + fmt.Sprint(time.Now().UnixNano()), Email: "u@x.com",
|
||||
PasswordHash: "x", Role: store.RoleUser, Balance: balance, Status: store.UserStatusActive}
|
||||
if err := db.Create(&u).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return u.ID
|
||||
}
|
||||
|
||||
func seedModel(t *testing.T, db *gorm.DB) uint64 {
|
||||
t.Helper()
|
||||
m := store.Model{Name: "m1", InputPrice: 0.15, OutputPrice: 0.60, Enabled: true}
|
||||
if err := db.Create(&m).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return m.ID
|
||||
}
|
||||
|
||||
func strPtr(s string) *string { return &s }
|
||||
Reference in new issue
Block a user