diff --git a/server/internal/api/admin.go b/server/internal/api/admin.go index 98ee6d4..b54b87f 100644 --- a/server/internal/api/admin.go +++ b/server/internal/api/admin.go @@ -115,6 +115,34 @@ func (h *Handler) AdminPatchUser(c *gin.Context) { } updates["status"] = *req.Status } + // 最后一个启用的管理员不可被降级/禁用：否则管理面被锁死， + // 且 reset-admin-password CLI 只在 admin 角色里找目标，救不回来。 + if req.Role != nil || req.Status != nil { + var target store.User + if err := h.a.DB.First(&target, id).Error; err != nil { + resp.Fail(c, http.StatusNotFound, "user not found") + return + } + if target.Role == store.RoleAdmin && target.Status == store.UserStatusActive { + newRole, newStatus := target.Role, target.Status + if req.Role != nil { + newRole = *req.Role + } + if req.Status != nil { + newStatus = *req.Status + } + if newRole != store.RoleAdmin || newStatus != store.UserStatusActive { + var admins int64 + h.a.DB.Model(&store.User{}). + Where("role = ? AND status = ? AND id != ?", store.RoleAdmin, store.UserStatusActive, id). + Count(&admins) + if admins == 0 { + resp.Fail(c, http.StatusBadRequest, "cannot demote or disable the last active admin") + return + } + } + } + } // 模型限制(jsonb)：手动序列化 if req.AllowedModels != nil { raw, _ := json.Marshal(*req.AllowedModels) @@ -166,16 +194,17 @@ func (h *Handler) AdminAdjustBalance(c *gin.Context) { if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).First(&u, id).Error; err != nil { return err } - newBalance := u.Balance + req.Amount - if err := tx.Model(&store.User{}).Where("id = ?", id).Update("balance", newBalance).Error; err != nil { + // SQL 侧原子加减：SQLite 驱动不支持行锁，Go 侧读-改-写会与记账扣费互相覆盖 + if err := tx.Model(&store.User{}).Where("id = ?", id).Update("balance", gorm.Expr("balance + ?", req.Amount)).Error; err != nil { return err } - ref := "admin:" + strconv.FormatUint(u.ID, 10) + ":" + time.Now().Format("20060102150405") + // ref_id 有单列唯一索引，秒级时间戳同一秒第二次调整必撞；纳秒粒度避免 + ref := "admin:" + strconv.FormatUint(u.ID, 10) + ":" + strconv.FormatInt(time.Now().UnixNano(), 10) _ = admin.ID // 流水里不冗余管理员 ID；需要时再加 return tx.Create(&store.BalanceLog{ UserID: u.ID, Change: req.Amount, - BalanceAfter: newBalance, + BalanceAfter: u.Balance + req.Amount, Type: store.BalanceTypeAdminAdjust, RefID: ref, Remark: req.Remark, diff --git a/server/internal/usage/usage.go b/server/internal/usage/usage.go index 6dbeb50..4f80dc3 100644 --- a/server/internal/usage/usage.go +++ b/server/internal/usage/usage.go @@ -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" diff --git a/server/internal/usage/usage_test.go b/server/internal/usage/usage_test.go new file mode 100644 index 0000000..8926673 --- /dev/null +++ b/server/internal/usage/usage_test.go @@ -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 }