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"
|
||||
|
||||
Reference in new issue
Block a user