diff --git a/common/region_sync.go b/common/region_sync.go new file mode 100644 index 0000000..df68b78 --- /dev/null +++ b/common/region_sync.go @@ -0,0 +1,33 @@ +package common + +import "errors" + +const ( + ForeignUserIDStart = 10000000 + UserSourceLocal = "local" + UserSourceSynced = "synced" +) + +func IsSyncedUser(source string, id int) bool { + return source == UserSourceSynced && id > 0 && id < ForeignUserIDStart +} + +func IsLocalUser(source string, id int) bool { + if source == UserSourceLocal { + return true + } + if id >= ForeignUserIDStart { + return true + } + return false +} + +func ValidateUserSource(source string, id int) error { + if source == UserSourceSynced && id >= ForeignUserIDStart { + return errors.New("synced 用户 ID 必须小于 1千万") + } + if source == UserSourceLocal && id > 0 && id < ForeignUserIDStart { + return errors.New("local 用户 ID 必须大于等于 1千万") + } + return nil +} diff --git a/common/region_sync_test.go b/common/region_sync_test.go new file mode 100644 index 0000000..b6b1c07 --- /dev/null +++ b/common/region_sync_test.go @@ -0,0 +1,93 @@ +package common + +import ( + "testing" +) + +func TestForeignUserIDStart(t *testing.T) { + if ForeignUserIDStart != 10000000 { + t.Errorf("ForeignUserIDStart = %d, want 10000000", ForeignUserIDStart) + } +} + +func TestUserSourceConstants(t *testing.T) { + if UserSourceLocal != "local" { + t.Errorf("UserSourceLocal = %q, want %q", UserSourceLocal, "local") + } + if UserSourceSynced != "synced" { + t.Errorf("UserSourceSynced = %q, want %q", UserSourceSynced, "synced") + } +} + +func TestIsSyncedUser(t *testing.T) { + tests := []struct { + name string + source string + id int + want bool + }{ + {"synced with normal id", UserSourceSynced, 123, true}, + {"synced with id just below threshold", UserSourceSynced, ForeignUserIDStart - 1, true}, + {"synced with id 0", UserSourceSynced, 0, false}, + {"synced with id at threshold", UserSourceSynced, ForeignUserIDStart, false}, + {"local with normal id", UserSourceLocal, 123, false}, + {"empty source with normal id", "", 123, false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := IsSyncedUser(tt.source, tt.id); got != tt.want { + t.Errorf("IsSyncedUser(%q, %d) = %v, want %v", tt.source, tt.id, got, tt.want) + } + }) + } +} + +func TestIsLocalUser(t *testing.T) { + tests := []struct { + name string + source string + id int + want bool + }{ + {"local at threshold", UserSourceLocal, ForeignUserIDStart, true}, + {"local above threshold", UserSourceLocal, ForeignUserIDStart + 100, true}, + {"local with id 0", UserSourceLocal, 0, true}, + {"synced with small id", UserSourceSynced, 123, false}, + {"empty source at threshold", "", ForeignUserIDStart, true}, + {"empty source with small id", "", 500, false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := IsLocalUser(tt.source, tt.id); got != tt.want { + t.Errorf("IsLocalUser(%q, %d) = %v, want %v", tt.source, tt.id, got, tt.want) + } + }) + } +} + +func TestValidateUserSource(t *testing.T) { + tests := []struct { + name string + source string + id int + wantErr bool + }{ + {"synced with valid id", UserSourceSynced, 100, false}, + {"synced at threshold", UserSourceSynced, ForeignUserIDStart, true}, + {"local at threshold", UserSourceLocal, ForeignUserIDStart, false}, + {"local with small non-zero id", UserSourceLocal, 500, true}, + {"local with id 0", UserSourceLocal, 0, false}, + {"empty source with id 0", "", 0, false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := ValidateUserSource(tt.source, tt.id) + if (err != nil) != tt.wantErr { + t.Errorf("ValidateUserSource(%q, %d) error = %v, wantErr %v", tt.source, tt.id, err, tt.wantErr) + } + }) + } +} diff --git a/constant/context_key.go b/constant/context_key.go index 2ba2fe2..f31e554 100644 --- a/constant/context_key.go +++ b/constant/context_key.go @@ -51,6 +51,7 @@ const ( ContextKeyUserGroup ContextKey = "user_group" ContextKeyUsingGroup ContextKey = "group" ContextKeyUserName ContextKey = "username" + ContextKeyUserSource ContextKey = "user_source" ContextKeyLocalCountTokens ContextKey = "local_count_tokens" diff --git a/controller/region_sync.go b/controller/region_sync.go new file mode 100644 index 0000000..9d193ec --- /dev/null +++ b/controller/region_sync.go @@ -0,0 +1,273 @@ +package controller + +import ( + "net/http" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/service/region_sync" + + "github.com/gin-gonic/gin" + "gorm.io/gorm" +) + +// 统一错误响应 +func syncErrorResponse(c *gin.Context, status int, err string) { + c.JSON(status, gin.H{"success": false, "error": err}) +} + +// 统一成功响应 +func syncSuccessResponse(c *gin.Context, message string) { + c.JSON(http.StatusOK, gin.H{"success": true, "message": message}) +} + +// ReceiveSyncedUserCreate 接收从 master 节点推送过来的用户创建请求 +func ReceiveSyncedUserCreate(c *gin.Context) { + var req region_sync.SyncUserRequest + if err := c.ShouldBindJSON(&req); err != nil { + syncErrorResponse(c, http.StatusBadRequest, "invalid request body") + return + } + + logger.LogDebug(c.Request.Context(), "[RegionSync] ReceiveSyncedUserCreate: username=%s, remoteUserId=%d, quota=%d", req.Username, req.RemoteUserId, req.Quota) + + // 检查用户是否已存在 + var existingUser model.User + if err := model.DB.Where("remote_user_id = ?", req.RemoteUserId).First(&existingUser).Error; err == nil { + logger.LogDebug(c.Request.Context(), "[RegionSync] ReceiveSyncedUserCreate: user already exists, remoteUserId=%d", req.RemoteUserId) + syncSuccessResponse(c, "user already exists") + return + } + + // 创建新用户 + user := createSyncedUser(&req) + if err := model.DB.Create(&user).Error; err != nil { + logger.LogError(c.Request.Context(), "failed to create synced user: "+err.Error()) + syncErrorResponse(c, http.StatusInternalServerError, "failed to create user") + return + } + + // 记录同步日志 + recordSyncLog(c, user.Id, user.RemoteUserId, model.SyncTypeUserCreate, "", 0, 0, model.SyncDirectionCnToOv) + + logger.LogInfo(c.Request.Context(), "created synced user "+user.Username) + syncSuccessResponse(c, "user created successfully") +} + +// createSyncedUser 根据请求创建同步用户 +func createSyncedUser(req *region_sync.SyncUserRequest) model.User { + affCode := req.AffCode + if affCode == "" { + affCode = common.GetRandomString(4) + } + + displayName := req.DisplayName + if displayName == "" { + displayName = req.Username + } + + group := req.Group + if group == "" { + group = "default" + } + + return model.User{ + Username: req.Username, + Email: req.Email, + Password: req.PasswordHash, + DisplayName: displayName, + Quota: req.Quota, + RemoteUserId: req.RemoteUserId, + Source: common.UserSourceSynced, + Group: group, + AffCode: affCode, + } +} + +// ReceiveQuotaUpdate 接收 master 节点的余额更新请求 +func ReceiveQuotaUpdate(c *gin.Context) { + var req region_sync.UpdateQuotaRequest + if err := c.ShouldBindJSON(&req); err != nil { + syncErrorResponse(c, http.StatusBadRequest, "invalid request body") + return + } + + logger.LogDebug(c.Request.Context(), "[RegionSync] ReceiveQuotaUpdate: remoteUserId=%d, newQuota=%d", req.RemoteUserId, req.Quota) + + // 查找同步用户 + var user model.User + if err := model.DB.Where("remote_user_id = ? AND source = ?", req.RemoteUserId, common.UserSourceSynced).First(&user).Error; err != nil { + logger.LogError(c.Request.Context(), "synced user not found: "+err.Error()) + syncErrorResponse(c, http.StatusNotFound, "user not found") + return + } + + // 更新余额(选择性更新,避免覆盖其他字段) + if err := model.DB.Model(&user).Updates(map[string]interface{}{ + "synced_quota": req.Quota, + "last_sync_at": time.Now().Unix(), + }).Error; err != nil { + logger.LogError(c.Request.Context(), "failed to update synced quota: "+err.Error()) + syncErrorResponse(c, http.StatusInternalServerError, "failed to update quota") + return + } + + // 记录同步日志 + recordSyncLog(c, user.Id, user.RemoteUserId, model.SyncTypeQuotaChange, "", 0, req.Quota, model.SyncDirectionCnToOv) + + logger.LogInfo(c.Request.Context(), "updated synced quota for user "+user.Username) + syncSuccessResponse(c, "quota updated successfully") +} + +// QueryUserQuota 查询用户余额(供 slave 节点调用 master 节点) +func QueryUserQuota(c *gin.Context) { + var req region_sync.QueryQuotaRequest + if err := c.ShouldBindJSON(&req); err != nil { + syncErrorResponse(c, http.StatusBadRequest, "invalid request body") + return + } + + // 查找用户 + var user model.User + if err := model.DB.Where("id = ?", req.UserId).First(&user).Error; err != nil { + syncErrorResponse(c, http.StatusNotFound, "user not found") + return + } + + logger.LogDebug(c.Request.Context(), "[RegionSync] QueryUserQuota: userId=%d, quota=%d", user.Id, user.Quota) + + // 记录同步日志 + recordSyncLog(c, user.Id, user.RemoteUserId, model.SyncTypePreConsumeQuery, "", user.Quota, 0, model.SyncDirectionOvToCn) + + c.JSON(http.StatusOK, gin.H{"success": true, "quota": user.Quota}) +} + +// BatchDeductQuota 批量扣费(供 slave 节点调用 master 节点) +func BatchDeductQuota(c *gin.Context) { + var req region_sync.BatchDeductRequest + if err := c.ShouldBindJSON(&req); err != nil { + syncErrorResponse(c, http.StatusBadRequest, "invalid request body") + return + } + + logger.LogDebug(c.Request.Context(), "[RegionSync] BatchDeductQuota: received %d records", len(req.Records)) + + results := make([]region_sync.DeductResult, 0, len(req.Records)) + + // 使用事务保证批量扣费的原子性(含幂等性检查和日志记录) + err := model.DB.Transaction(func(tx *gorm.DB) error { + for _, record := range req.Records { + result := processDeductRecordWithTx(c, tx, record) + results = append(results, result) + // 在事务内记录同步日志,确保幂等性检查能查到 + if result.Success { + syncLog := &model.QuotaSyncLog{ + UserId: record.UserId, + SyncType: model.SyncTypeBatchSync, + Direction: model.SyncDirectionOvToCn, + Status: model.SyncStatusSuccess, + RequestId: record.RequestId, + ChangeAmount: record.Quota, + AfterQuota: result.RemainingQuota, + CreatedAt: time.Now().Unix(), + } + if err := tx.Create(syncLog).Error; err != nil { + return err + } + } + } + return nil + }) + + if err != nil { + logger.LogError(c.Request.Context(), "batch deduct transaction failed: "+err.Error()) + syncErrorResponse(c, http.StatusInternalServerError, "batch deduct failed") + return + } + + successCount := 0 + for _, r := range results { + if r.Success { + successCount++ + } + } + logger.LogDebug(c.Request.Context(), "[RegionSync] BatchDeductQuota: completed, success=%d, total=%d", successCount, len(req.Records)) + + c.JSON(http.StatusOK, gin.H{"success": true, "results": results}) +} + +// processDeductRecordWithTx 在事务中处理单条扣费记录 +func processDeductRecordWithTx(c *gin.Context, tx *gorm.DB, record region_sync.BatchDeductRecord) region_sync.DeductResult { + result := region_sync.DeductResult{UserId: record.UserId} + + logger.LogDebug(c.Request.Context(), "[RegionSync] processDeductRecord: userId=%d, requestId=%s, quota=%d", record.UserId, record.RequestId, record.Quota) + + // 幂等性检查:通过 quota_sync_logs 检查 request_id 是否已在 master 端处理过 + var existingLog model.QuotaSyncLog + if err := tx.Where("request_id = ? AND sync_type = ? AND status = ?", + record.RequestId, model.SyncTypeBatchSync, model.SyncStatusSuccess).First(&existingLog).Error; err == nil { + // 该 request_id 已处理过,直接返回成功(不重复扣费) + result.Success = true + result.Message = "already processed" + result.DeductedQuota = existingLog.ChangeAmount + // 查询当前余额 + var user model.User + if err := tx.Where("id = ?", record.UserId).First(&user).Error; err == nil { + result.RemainingQuota = user.Quota + } + logger.LogDebug(c.Request.Context(), "[RegionSync] processDeductRecord: idempotent hit, userId=%d, requestId=%s, deductedQuota=%d", record.UserId, record.RequestId, result.DeductedQuota) + return result + } + + // 原子扣费:使用 SQL 条件更新,确保余额充足时才扣减,避免并发覆盖 + dbResult := tx.Model(&model.User{}). + Where("id = ? AND quota >= ?", record.UserId, record.Quota). + Update("quota", gorm.Expr("quota - ?", record.Quota)) + if dbResult.Error != nil { + result.Error = "failed to deduct quota" + return result + } + if dbResult.RowsAffected == 0 { + // 没有匹配行,可能是用户不存在或余额不足 + var user model.User + if err := tx.Where("id = ?", record.UserId).First(&user).Error; err != nil { + result.Error = "user not found" + } else { + result.Error = "insufficient quota" + result.RemainingQuota = user.Quota + logger.LogDebug(c.Request.Context(), "[RegionSync] processDeductRecord: insufficient quota, userId=%d, quota=%d, need=%d", record.UserId, user.Quota, record.Quota) + } + return result + } + + // 查询扣费后的余额 + var user model.User + if err := tx.Where("id = ?", record.UserId).First(&user).Error; err == nil { + result.RemainingQuota = user.Quota + } + + result.Success = true + result.DeductedQuota = record.Quota + logger.LogDebug(c.Request.Context(), "[RegionSync] processDeductRecord: success, userId=%d, deducted=%d, remaining=%d", record.UserId, record.Quota, result.RemainingQuota) + return result +} + +// recordSyncLog 记录同步日志 +func recordSyncLog(c *gin.Context, userId, remoteUserId int, syncType, requestId string, masterQuota, afterQuota int, direction string) { + syncLog := &model.QuotaSyncLog{ + UserId: userId, + RemoteUserId: remoteUserId, + SyncType: syncType, + Direction: direction, + Status: model.SyncStatusSuccess, + RequestId: requestId, + MasterQuota: masterQuota, + AfterQuota: afterQuota, + CreatedAt: time.Now().Unix(), + } + if err := model.CreateSyncLog(syncLog); err != nil { + logger.LogError(c.Request.Context(), "failed to create sync log: "+err.Error()) + } +} diff --git a/controller/region_sync_test.go b/controller/region_sync_test.go new file mode 100644 index 0000000..35a80a4 --- /dev/null +++ b/controller/region_sync_test.go @@ -0,0 +1,466 @@ +package controller + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/service/region_sync" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func TestReceiveSyncedUserCreate_Success(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + // 使用测试专用的认证中间件,绕过系统设置检查 + router.Use(testSyncAuthMiddleware()) + router.POST("/api/internal/sync/user/create", ReceiveSyncedUserCreate) + + origDB := model.DB + defer func() { + model.DB = origDB + }() + + // 设置测试数据库 + db := setupSyncTestDB(t) + model.DB = db + + req := region_sync.SyncUserRequest{ + Username: "testuser", + Email: "test@example.com", + PasswordHash: "hashedpassword", + DisplayName: "Test User", + Quota: 10000, + RemoteUserId: 10000001, + Group: "default", + } + + reqBody, _ := json.Marshal(req) + reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/user/create", bytes.NewBuffer(reqBody)) + reqHTTP.Header.Set("Content-Type", "application/json") + + w := httptest.NewRecorder() + router.ServeHTTP(w, reqHTTP) + + assert.Equal(t, http.StatusOK, w.Code) + + var resp map[string]interface{} + json.Unmarshal(w.Body.Bytes(), &resp) + assert.True(t, resp["success"].(bool)) + + // 验证用户已创建 + var user model.User + err := db.Where("username = ?", "testuser").First(&user).Error + require.NoError(t, err) + assert.Equal(t, "testuser", user.Username) + assert.Equal(t, 10000001, user.RemoteUserId) + assert.Equal(t, "synced", user.Source) +} + +func TestReceiveQuotaUpdate_Success(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(testSyncAuthMiddleware()) + router.POST("/api/internal/sync/quota/update", ReceiveQuotaUpdate) + + origDB := model.DB + defer func() { + model.DB = origDB + }() + + // 设置测试数据库 + db := setupSyncTestDB(t) + model.DB = db + + // 创建测试用户 + user := model.User{ + Username: "testuser", + RemoteUserId: 10000001, + Source: "synced", + Quota: 5000, + SyncedQuota: 5000, + AffCode: "aff-qu-1", + } + db.Create(&user) + + req := region_sync.UpdateQuotaRequest{ + RemoteUserId: 10000001, + Quota: 10000, + } + + reqBody, _ := json.Marshal(req) + reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/update", bytes.NewBuffer(reqBody)) + reqHTTP.Header.Set("Content-Type", "application/json") + + w := httptest.NewRecorder() + router.ServeHTTP(w, reqHTTP) + + assert.Equal(t, http.StatusOK, w.Code) + + var resp map[string]interface{} + json.Unmarshal(w.Body.Bytes(), &resp) + assert.True(t, resp["success"].(bool)) + + // 验证余额已更新 + var updatedUser model.User + err := db.Where("remote_user_id = ?", 10000001).First(&updatedUser).Error + require.NoError(t, err) + assert.Equal(t, 10000, updatedUser.SyncedQuota) +} + +func TestQueryUserQuota_Success(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(testSyncAuthMiddleware()) + router.POST("/api/internal/sync/quota/query", QueryUserQuota) + + origDB := model.DB + defer func() { + model.DB = origDB + }() + + // 设置测试数据库 + db := setupSyncTestDB(t) + model.DB = db + + // 创建测试用户 + user := model.User{ + Username: "testuser", + Quota: 15000, + AffCode: "aff-qq-1", + } + db.Create(&user) + + req := region_sync.QueryQuotaRequest{ + UserId: user.Id, + } + + reqBody, _ := json.Marshal(req) + reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/query", bytes.NewBuffer(reqBody)) + reqHTTP.Header.Set("Content-Type", "application/json") + + w := httptest.NewRecorder() + router.ServeHTTP(w, reqHTTP) + + assert.Equal(t, http.StatusOK, w.Code) + + var resp region_sync.QueryQuotaResponse + json.Unmarshal(w.Body.Bytes(), &resp) + assert.True(t, resp.Success) + assert.Equal(t, 15000, resp.Quota) +} + +func TestBatchDeductQuota_Success(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(testSyncAuthMiddleware()) + router.POST("/api/internal/sync/quota/batch-deduct", BatchDeductQuota) + + origDB := model.DB + defer func() { + model.DB = origDB + }() + + // 设置测试数据库 + db := setupSyncTestDB(t) + model.DB = db + + // 创建测试用户 + user1 := model.User{Username: "user1", Quota: 1000, AffCode: "aff-bd-1"} + user2 := model.User{Username: "user2", Quota: 500, AffCode: "aff-bd-2"} + db.Create(&user1) + db.Create(&user2) + + req := region_sync.BatchDeductRequest{ + Records: []region_sync.BatchDeductRecord{ + {UserId: user1.Id, RequestId: "req-1", Quota: 100}, + {UserId: user2.Id, RequestId: "req-2", Quota: 50}, + }, + } + + reqBody, _ := json.Marshal(req) + reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/batch-deduct", bytes.NewBuffer(reqBody)) + reqHTTP.Header.Set("Content-Type", "application/json") + + w := httptest.NewRecorder() + router.ServeHTTP(w, reqHTTP) + + assert.Equal(t, http.StatusOK, w.Code) + + var resp region_sync.BatchDeductResponse + json.Unmarshal(w.Body.Bytes(), &resp) + assert.True(t, resp.Success) + assert.Len(t, resp.Results, 2) + + // 验证用户余额已更新 + var updatedUser1 model.User + db.First(&updatedUser1, user1.Id) + assert.Equal(t, 900, updatedUser1.Quota) + + var updatedUser2 model.User + db.First(&updatedUser2, user2.Id) + assert.Equal(t, 450, updatedUser2.Quota) +} + +func TestBatchDeductQuota_InsufficientQuota(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(testSyncAuthMiddleware()) + router.POST("/api/internal/sync/quota/batch-deduct", BatchDeductQuota) + + origDB := model.DB + defer func() { + model.DB = origDB + }() + + // 设置测试数据库 + db := setupSyncTestDB(t) + model.DB = db + + // 创建余额不足的用户 + user := model.User{Username: "pooruser", Quota: 10, AffCode: "aff-poor-1"} + db.Create(&user) + + req := region_sync.BatchDeductRequest{ + Records: []region_sync.BatchDeductRecord{ + {UserId: user.Id, RequestId: "req-1", Quota: 100}, + }, + } + + reqBody, _ := json.Marshal(req) + reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/batch-deduct", bytes.NewBuffer(reqBody)) + reqHTTP.Header.Set("Content-Type", "application/json") + + w := httptest.NewRecorder() + router.ServeHTTP(w, reqHTTP) + + assert.Equal(t, http.StatusOK, w.Code) + + var resp region_sync.BatchDeductResponse + json.Unmarshal(w.Body.Bytes(), &resp) + assert.True(t, resp.Success) + assert.Len(t, resp.Results, 1) + assert.False(t, resp.Results[0].Success) + assert.Contains(t, resp.Results[0].Error, "insufficient quota") +} + +// testSyncAuthMiddleware 测试专用的认证中间件,绕过系统设置 +func testSyncAuthMiddleware() gin.HandlerFunc { + return func(c *gin.Context) { + c.Set("sync_node", "test-node") + c.Next() + } +} + +// setupSyncTestDB 设置同步测试数据库 +func setupSyncTestDB(t *testing.T) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, _ := db.DB() + sqlDB.SetMaxOpenConns(1) + + err = db.AutoMigrate(&model.User{}, &model.QuotaSyncLog{}) + require.NoError(t, err) + + return db +} + +// --------------------------------------------------------------------------- +// 阶段 4:Controller 边界条件测试 +// --------------------------------------------------------------------------- + +func TestReceiveSyncedUserCreate_DuplicateUser(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(testSyncAuthMiddleware()) + router.POST("/api/internal/sync/user/create", ReceiveSyncedUserCreate) + + origDB := model.DB + defer func() { model.DB = origDB }() + + db := setupSyncTestDB(t) + model.DB = db + + // 先创建一个用户 + db.Create(&model.User{ + Username: "dupuser", RemoteUserId: 10000001, + Source: common.UserSourceSynced, Password: "hash", AffCode: "aff-dup-1", + }) + + req := region_sync.SyncUserRequest{ + Username: "dupuser", + RemoteUserId: 10000001, + PasswordHash: "newhash", + } + reqBody, _ := json.Marshal(req) + reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/user/create", bytes.NewBuffer(reqBody)) + reqHTTP.Header.Set("Content-Type", "application/json") + + w := httptest.NewRecorder() + router.ServeHTTP(w, reqHTTP) + + assert.Equal(t, http.StatusOK, w.Code) + var resp map[string]interface{} + json.Unmarshal(w.Body.Bytes(), &resp) + assert.True(t, resp["success"].(bool)) + assert.Contains(t, resp["message"], "already exists") +} + +func TestReceiveSyncedUserCreate_InvalidBody(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(testSyncAuthMiddleware()) + router.POST("/api/internal/sync/user/create", ReceiveSyncedUserCreate) + + reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/user/create", strings.NewReader("not json")) + reqHTTP.Header.Set("Content-Type", "application/json") + + w := httptest.NewRecorder() + router.ServeHTTP(w, reqHTTP) + assert.Equal(t, http.StatusBadRequest, w.Code) +} + +func TestReceiveQuotaUpdate_UserNotFound(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(testSyncAuthMiddleware()) + router.POST("/api/internal/sync/quota/update", ReceiveQuotaUpdate) + + origDB := model.DB + defer func() { model.DB = origDB }() + + db := setupSyncTestDB(t) + model.DB = db + + req := region_sync.UpdateQuotaRequest{ + RemoteUserId: 99999999, + Quota: 5000, + } + reqBody, _ := json.Marshal(req) + reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/update", bytes.NewBuffer(reqBody)) + reqHTTP.Header.Set("Content-Type", "application/json") + + w := httptest.NewRecorder() + router.ServeHTTP(w, reqHTTP) + assert.Equal(t, http.StatusNotFound, w.Code) +} + +func TestReceiveQuotaUpdate_InvalidBody(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(testSyncAuthMiddleware()) + router.POST("/api/internal/sync/quota/update", ReceiveQuotaUpdate) + + reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/update", strings.NewReader("invalid")) + reqHTTP.Header.Set("Content-Type", "application/json") + + w := httptest.NewRecorder() + router.ServeHTTP(w, reqHTTP) + assert.Equal(t, http.StatusBadRequest, w.Code) +} + +func TestQueryUserQuota_UserNotFound(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(testSyncAuthMiddleware()) + router.POST("/api/internal/sync/quota/query", QueryUserQuota) + + origDB := model.DB + defer func() { model.DB = origDB }() + + db := setupSyncTestDB(t) + model.DB = db + + req := region_sync.QueryQuotaRequest{UserId: 99999} + reqBody, _ := json.Marshal(req) + reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/query", bytes.NewBuffer(reqBody)) + reqHTTP.Header.Set("Content-Type", "application/json") + + w := httptest.NewRecorder() + router.ServeHTTP(w, reqHTTP) + assert.Equal(t, http.StatusNotFound, w.Code) +} + +func TestBatchDeductQuota_EmptyRecords(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(testSyncAuthMiddleware()) + router.POST("/api/internal/sync/quota/batch-deduct", BatchDeductQuota) + + origDB := model.DB + defer func() { model.DB = origDB }() + + db := setupSyncTestDB(t) + model.DB = db + + req := region_sync.BatchDeductRequest{Records: []region_sync.BatchDeductRecord{}} + reqBody, _ := json.Marshal(req) + reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/batch-deduct", bytes.NewBuffer(reqBody)) + reqHTTP.Header.Set("Content-Type", "application/json") + + w := httptest.NewRecorder() + router.ServeHTTP(w, reqHTTP) + assert.Equal(t, http.StatusOK, w.Code) + + var resp region_sync.BatchDeductResponse + json.Unmarshal(w.Body.Bytes(), &resp) + assert.True(t, resp.Success) + assert.Len(t, resp.Results, 0) +} + +func TestBatchDeductQuota_UserNotFound(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(testSyncAuthMiddleware()) + router.POST("/api/internal/sync/quota/batch-deduct", BatchDeductQuota) + + origDB := model.DB + defer func() { model.DB = origDB }() + + db := setupSyncTestDB(t) + model.DB = db + + req := region_sync.BatchDeductRequest{ + Records: []region_sync.BatchDeductRecord{ + {UserId: 99999, RequestId: "req-nf", Quota: 100}, + }, + } + reqBody, _ := json.Marshal(req) + reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/batch-deduct", bytes.NewBuffer(reqBody)) + reqHTTP.Header.Set("Content-Type", "application/json") + + w := httptest.NewRecorder() + router.ServeHTTP(w, reqHTTP) + assert.Equal(t, http.StatusOK, w.Code) + + var resp region_sync.BatchDeductResponse + json.Unmarshal(w.Body.Bytes(), &resp) + assert.True(t, resp.Success) + assert.Len(t, resp.Results, 1) + assert.False(t, resp.Results[0].Success) + assert.Contains(t, resp.Results[0].Error, "user not found") +} + +func TestBatchDeductQuota_InvalidBody(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(testSyncAuthMiddleware()) + router.POST("/api/internal/sync/quota/batch-deduct", BatchDeductQuota) + + reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/batch-deduct", strings.NewReader("bad")) + reqHTTP.Header.Set("Content-Type", "application/json") + + w := httptest.NewRecorder() + router.ServeHTTP(w, reqHTTP) + assert.Equal(t, http.StatusBadRequest, w.Code) +} diff --git a/controller/relay.go b/controller/relay.go index b8c2dc0..713f270 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -159,6 +159,8 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { if priceData.FreeModel { logger.LogInfo(c, fmt.Sprintf("模型 %s 免费,跳过预扣费", relayInfo.OriginModelName)) } else { + common.SysLog(fmt.Sprintf("[RegionSync] relay PreConsumeBilling: userId=%d, QuotaToPreConsume=%d, FreeModel=%v, UsePrice=%v, ModelRatio=%.4f, GroupRatio=%.4f, ModelPrice=%.4f, model=%s", + relayInfo.UserId, priceData.QuotaToPreConsume, priceData.FreeModel, priceData.UsePrice, priceData.ModelRatio, priceData.GroupRatioInfo.GroupRatio, priceData.ModelPrice, relayInfo.OriginModelName)) newAPIError = service.PreConsumeBilling(c, priceData.QuotaToPreConsume, relayInfo) if newAPIError != nil { return diff --git a/controller/user.go b/controller/user.go index b58eab8..8524291 100644 --- a/controller/user.go +++ b/controller/user.go @@ -16,6 +16,7 @@ import ( "github.com/QuantumNous/new-api/logger" "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/service/region_sync" "github.com/QuantumNous/new-api/setting" "github.com/QuantumNous/new-api/constant" @@ -184,12 +185,8 @@ func Register(c *gin.Context) { return } - // 获取插入后的用户ID - var insertedUser model.User - if err := model.DB.Where("username = ?", cleanUser.Username).First(&insertedUser).Error; err != nil { - common.ApiErrorI18n(c, i18n.MsgUserRegisterFailed) - return - } + // 同步用户到海外节点(异步执行,不阻塞注册流程) + region_sync.PushUserCreateToSlave(&cleanUser) // 生成默认令牌 if constant.GenerateDefaultToken { key, err := common.GenerateKey() @@ -200,7 +197,7 @@ func Register(c *gin.Context) { } // 生成默认令牌 token := model.Token{ - UserId: insertedUser.Id, // 使用插入后的用户ID + UserId: cleanUser.Id, // GORM Create 后已填充 ID Name: cleanUser.Username + "的初始令牌", Key: key, CreatedTime: common.GetTimestamp(), @@ -383,6 +380,11 @@ func GetSelf(c *gin.Context) { userSetting := user.GetSetting() // 构建响应数据,包含用户信息和权限 + // Slave 节点使用 SyncedQuota 作为用户额度 + quota := user.Quota + if user.IsSyncedUser() { + quota = user.SyncedQuota + } responseData := map[string]interface{}{ "id": user.Id, "username": user.Username, @@ -396,7 +398,7 @@ func GetSelf(c *gin.Context) { "wechat_id": user.WeChatId, "telegram_id": user.TelegramId, "group": user.Group, - "quota": user.Quota, + "quota": quota, "used_quota": user.UsedQuota, "request_count": user.RequestCount, "aff_code": user.AffCode, diff --git a/main.go b/main.go index 476a2ed..4b58d71 100644 --- a/main.go +++ b/main.go @@ -22,8 +22,10 @@ import ( "github.com/QuantumNous/new-api/relay" "github.com/QuantumNous/new-api/router" "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/service/region_sync" _ "github.com/QuantumNous/new-api/setting/performance_setting" "github.com/QuantumNous/new-api/setting/ratio_setting" + "github.com/QuantumNous/new-api/setting/system_setting" "github.com/bytedance/gopkg/util/gopool" "github.com/gin-contrib/sessions" @@ -174,7 +176,12 @@ func main() { Secure: false, SameSite: http.SameSiteStrictMode, }) - server.Use(sessions.Sessions("session", store)) + // 支持通过环境变量自定义 session 名称,用于多节点部署时避免 cookie 冲突 + sessionName := os.Getenv("SESSION_NAME") + if sessionName == "" { + sessionName = "session" + } + server.Use(sessions.Sessions(sessionName, store)) InjectUmamiAnalytics() InjectGoogleAnalytics() @@ -309,5 +316,19 @@ func InitResources() error { // Don't return error, custom OAuth is not critical } + // 注册余额更新回调,Master 节点更新用户余额后推送到 Slave 节点 + model.SetQuotaUpdateCallback(func(userId int, quota int) { + if system_setting.GetRegionSyncSettings().IsMaster { + region_sync.PushQuotaUpdateToSlave(userId, quota) + } + }) + + // Slave 节点启动后台同步任务(批量扣费同步 + 余额定时拉取 + 清理) + syncSettings := system_setting.GetRegionSyncSettings() + if syncSettings.Enabled && !syncSettings.IsMaster { + syncManager := region_sync.NewSyncManager() + syncManager.StartSyncWorkers() + } + return nil } diff --git a/makefile b/makefile index cc626c3..f24141b 100644 --- a/makefile +++ b/makefile @@ -22,7 +22,7 @@ DOCKER_TAG := $(DOCKER_REGISTRY):$(BUILD_TIME)-$(BRANCH_NAME) docker-build: @echo "Building Docker image with tag: $(DOCKER_TAG)" docker build -t $(DOCKER_TAG) . - docker push $(DOCKER_TAG) + #docker push $(DOCKER_TAG) docker-push: docker-build @echo "Pushing Docker image: $(DOCKER_TAG)" diff --git a/middleware/sync_auth.go b/middleware/sync_auth.go new file mode 100644 index 0000000..e25b38e --- /dev/null +++ b/middleware/sync_auth.go @@ -0,0 +1,96 @@ +package middleware + +import ( + "crypto/subtle" + "net/http" + "strings" + + "github.com/QuantumNous/new-api/setting/system_setting" + + "github.com/gin-gonic/gin" +) + +// SyncAuth 同步 API 认证中间件 +// 通过 X-Sync-API-Key 请求头验证请求的合法性 +func SyncAuth() gin.HandlerFunc { + return func(c *gin.Context) { + settings := system_setting.GetRegionSyncSettings() + + // 如果未启用区域同步,拒绝请求 + if !settings.Enabled { + c.JSON(http.StatusForbidden, gin.H{ + "success": false, + "message": "region sync is not enabled", + }) + c.Abort() + return + } + + // 检查 X-Sync-API-Key 请求头 + apiKey := c.GetHeader("X-Sync-API-Key") + if apiKey == "" { + c.JSON(http.StatusUnauthorized, gin.H{ + "success": false, + "message": "missing sync API key", + }) + c.Abort() + return + } + + // 验证 API Key + if !validateSyncAPIKey(apiKey, settings.SyncApiKey) { + c.JSON(http.StatusUnauthorized, gin.H{ + "success": false, + "message": "invalid sync API key", + }) + c.Abort() + return + } + + // 验证来源节点(从 X-Sync-Node 获取) + node := c.GetHeader("X-Sync-Node") + if node == "" { + c.JSON(http.StatusBadRequest, gin.H{ + "success": false, + "message": "missing sync node identifier", + }) + c.Abort() + return + } + + // 将节点信息存储在上下文中 + c.Set("sync_node", node) + + c.Next() + } +} + +// validateSyncAPIKey 验证同步 API Key(使用常量时间比较防止 timing attack) +func validateSyncAPIKey(requestedKey, configuredKey string) bool { + if configuredKey == "" { + return false + } + + // 支持 master 配置多个 slave 的 key,以逗号分隔 + if strings.Contains(configuredKey, ",") { + keys := strings.Split(configuredKey, ",") + for _, key := range keys { + if subtle.ConstantTimeCompare([]byte(strings.TrimSpace(key)), []byte(requestedKey)) == 1 { + return true + } + } + return false + } + + return subtle.ConstantTimeCompare([]byte(strings.TrimSpace(configuredKey)), []byte(requestedKey)) == 1 +} + +// GetSyncNode 从上下文获取同步节点标识 +func GetSyncNode(c *gin.Context) string { + if node, exists := c.Get("sync_node"); exists { + if nodeStr, ok := node.(string); ok { + return nodeStr + } + } + return "" +} diff --git a/middleware/sync_auth_test.go b/middleware/sync_auth_test.go new file mode 100644 index 0000000..f86af17 --- /dev/null +++ b/middleware/sync_auth_test.go @@ -0,0 +1,149 @@ +package middleware + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/setting/system_setting" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" +) + +func TestSyncAuth_Disabled(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(SyncAuth()) + router.GET("/test", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"success": true}) + }) + + // 未启用区域同步 + settings := system_setting.GetRegionSyncSettings() + settings.Enabled = false + + req, _ := http.NewRequest("GET", "/test", nil) + req.Header.Set("X-Sync-API-Key", "test-key") + req.Header.Set("X-Sync-Node", "slave-1") + + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + assert.Equal(t, http.StatusForbidden, w.Code) + assert.Contains(t, w.Body.String(), "region sync is not enabled") +} + +func TestSyncAuth_MissingAPIKey(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(SyncAuth()) + router.GET("/test", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"success": true}) + }) + + settings := system_setting.GetRegionSyncSettings() + settings.Enabled = true + settings.SyncApiKey = "test-key" + + req, _ := http.NewRequest("GET", "/test", nil) + req.Header.Set("X-Sync-Node", "slave-1") + + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + assert.Equal(t, http.StatusUnauthorized, w.Code) + assert.Contains(t, w.Body.String(), "missing sync API key") +} + +func TestSyncAuth_InvalidAPIKey(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(SyncAuth()) + router.GET("/test", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"success": true}) + }) + + settings := system_setting.GetRegionSyncSettings() + settings.Enabled = true + settings.SyncApiKey = "correct-key" + + req, _ := http.NewRequest("GET", "/test", nil) + req.Header.Set("X-Sync-API-Key", "wrong-key") + req.Header.Set("X-Sync-Node", "slave-1") + + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + assert.Equal(t, http.StatusUnauthorized, w.Code) + assert.Contains(t, w.Body.String(), "invalid sync API key") +} + +func TestSyncAuth_MissingNode(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(SyncAuth()) + router.GET("/test", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"success": true}) + }) + + settings := system_setting.GetRegionSyncSettings() + settings.Enabled = true + settings.SyncApiKey = "test-key" + + req, _ := http.NewRequest("GET", "/test", nil) + req.Header.Set("X-Sync-API-Key", "test-key") + + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + assert.Equal(t, http.StatusBadRequest, w.Code) + assert.Contains(t, w.Body.String(), "missing sync node identifier") +} + +func TestSyncAuth_Valid(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(SyncAuth()) + router.GET("/test", func(c *gin.Context) { + node := GetSyncNode(c) + c.JSON(http.StatusOK, gin.H{ + "success": true, + "node": node, + }) + }) + + settings := system_setting.GetRegionSyncSettings() + settings.Enabled = true + settings.SyncApiKey = "test-key" + + req, _ := http.NewRequest("GET", "/test", nil) + req.Header.Set("X-Sync-API-Key", "test-key") + req.Header.Set("X-Sync-Node", "slave-1") + + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + assert.Equal(t, http.StatusOK, w.Code) + assert.Contains(t, w.Body.String(), "slave-1") +} + +func TestValidateSyncAPIKey_Single(t *testing.T) { + assert.True(t, validateSyncAPIKey("test-key", "test-key")) + assert.False(t, validateSyncAPIKey("test-key", "different-key")) + assert.False(t, validateSyncAPIKey("", "")) +} + +func TestValidateSyncAPIKey_Multiple(t *testing.T) { + // 支持逗号分隔的多个 key + assert.True(t, validateSyncAPIKey("key1", "key1,key2,key3")) + assert.True(t, validateSyncAPIKey("key2", "key1,key2,key3")) + assert.True(t, validateSyncAPIKey("key3", "key1,key2,key3")) + assert.False(t, validateSyncAPIKey("key4", "key1,key2,key3")) +} + +func TestValidateSyncAPIKey_TrimSpace(t *testing.T) { + // 去除前后空格 + assert.True(t, validateSyncAPIKey("test-key", " test-key ")) + assert.True(t, validateSyncAPIKey("test-key", " test-key")) + assert.True(t, validateSyncAPIKey("test-key", "test-key ")) +} diff --git a/model/main.go b/model/main.go index d72e201..a78a317 100644 --- a/model/main.go +++ b/model/main.go @@ -280,6 +280,8 @@ func migrateDB() error { &UserOAuthBinding{}, &ChannelPricing{}, &PricingTag{}, + &PendingSyncRecord{}, + &QuotaSyncLog{}, ) if err != nil { return err @@ -332,6 +334,8 @@ func migrateDBFast() error { {&UserOAuthBinding{}, "UserOAuthBinding"}, {&ChannelPricing{}, "ChannelPricing"}, {&PricingTag{}, "PricingTag"}, + {&PendingSyncRecord{}, "PendingSyncRecord"}, + {&QuotaSyncLog{}, "QuotaSyncLog"}, } // 动态计算migration数量,确保errChan缓冲区足够大 errChan := make(chan error, len(migrations)) diff --git a/model/pending_sync_record.go b/model/pending_sync_record.go new file mode 100644 index 0000000..3eb6180 --- /dev/null +++ b/model/pending_sync_record.go @@ -0,0 +1,126 @@ +package model + +import ( + "time" + + "gorm.io/gorm" +) + +const ( + PendingSyncStatusPending = "pending" + PendingSyncStatusSynced = "synced" + PendingSyncStatusFailed = "failed" + PendingSyncStatusArchived = "archived" +) + +// PendingSyncRecord 待同步消费记录 +type PendingSyncRecord struct { + Id int `json:"id" gorm:"primaryKey"` + UserId int `json:"user_id" gorm:"index:idx_user_status"` + RemoteUserId int `json:"remote_user_id" gorm:"index"` + RequestId string `json:"request_id" gorm:"type:varchar(128);uniqueIndex"` + Quota int `json:"quota"` + EstimatedQuota int `json:"estimated_quota"` + Status string `json:"status" gorm:"default:'pending';index:idx_user_status;index:idx_status_created"` + CreatedAt int64 `json:"created_at" gorm:"index:idx_status_created"` + SyncedAt int64 `json:"synced_at"` + RetryCount int `json:"retry_count" gorm:"default:0"` + ErrorMsg string `json:"error_msg"` +} + +func (PendingSyncRecord) TableName() string { + return "pending_sync_records" +} + +// GetPendingSyncQuota 获取用户待同步金额总和 +func GetPendingSyncQuota(userId int) int { + var total int64 + DB.Model(&PendingSyncRecord{}). + Where("user_id = ? AND status IN ?", userId, []string{PendingSyncStatusPending, PendingSyncStatusFailed}). + Select("COALESCE(SUM(quota), 0)"). + Scan(&total) + return int(total) +} + +// CreatePendingSyncRecord 创建待同步记录 +func CreatePendingSyncRecord(userId, remoteUserId int, requestId string, quota, estimatedQuota int) error { + record := PendingSyncRecord{ + UserId: userId, + RemoteUserId: remoteUserId, + RequestId: requestId, + Quota: quota, + EstimatedQuota: estimatedQuota, + Status: PendingSyncStatusPending, + CreatedAt: time.Now().Unix(), + } + return DB.Create(&record).Error +} + +// GetPendingRecordsForSync 获取待同步记录 +func GetPendingRecordsForSync(limit, maxRetry int) []PendingSyncRecord { + var records []PendingSyncRecord + DB.Where("status IN ? AND retry_count < ?", + []string{PendingSyncStatusPending, PendingSyncStatusFailed}, + maxRetry). + Limit(limit). + Order("created_at asc"). + Find(&records) + return records +} + +// MarkRecordSynced 标记记录为已同步 +func MarkRecordSynced(recordId int) error { + return DB.Model(&PendingSyncRecord{}).Where("id = ?", recordId).Updates(map[string]interface{}{ + "status": PendingSyncStatusSynced, + "synced_at": time.Now().Unix(), + }).Error +} + +// MarkRecordFailed 标记记录为失败 +func MarkRecordFailed(recordId int, errMsg string) error { + return DB.Model(&PendingSyncRecord{}).Where("id = ?", recordId).Updates(map[string]interface{}{ + "status": PendingSyncStatusFailed, + "retry_count": gorm.Expr("retry_count + 1"), + "error_msg": errMsg, + }).Error +} + +// UpdatePendingSyncRecordQuota 更新待同步记录的实际扣费额度,返回影响的行数 +func UpdatePendingSyncRecordQuota(requestId string, actualQuota int) int64 { + result := DB.Model(&PendingSyncRecord{}). + Where("request_id = ? AND status = ?", requestId, PendingSyncStatusPending). + Update("quota", actualQuota) + if result.Error != nil { + return 0 + } + return result.RowsAffected +} + +// CleanupOldSyncedRecords 清理超过指定天数的已同步记录 +func CleanupOldSyncedRecords(days int) int64 { + result := DB.Where("status = ? AND synced_at < ?", + PendingSyncStatusSynced, + time.Now().Unix()-int64(days*86400)). + Delete(&PendingSyncRecord{}) + return result.RowsAffected +} + +// GetPendingQuotaByUser 获取每个用户的待同步金额 +func GetPendingQuotaByUser() map[int]int { + type result struct { + UserId int + Total int + } + var results []result + DB.Model(&PendingSyncRecord{}). + Select("user_id, COALESCE(SUM(quota), 0) as total"). + Where("status IN ?", []string{PendingSyncStatusPending, PendingSyncStatusFailed}). + Group("user_id"). + Scan(&results) + + m := make(map[int]int) + for _, r := range results { + m[r.UserId] = r.Total + } + return m +} diff --git a/model/quota_sync_log.go b/model/quota_sync_log.go new file mode 100644 index 0000000..c22977e --- /dev/null +++ b/model/quota_sync_log.go @@ -0,0 +1,55 @@ +package model + +import "time" + +const ( + SyncTypeUserCreate = "user_create" + SyncTypeQuotaChange = "quota_change" + SyncTypePreConsumeQuery = "pre_consume_query" + SyncTypeBatchSync = "batch_sync" + SyncTypeManualSync = "manual_sync" + SyncTypeStartupSync = "startup_sync" + SyncTypeSettle = "settle" +) + +const ( + SyncDirectionCnToOv = "cn_to_ov" + SyncDirectionOvToCn = "ov_to_cn" +) + +const ( + SyncStatusSuccess = "success" + SyncStatusFailed = "failed" +) + +// QuotaSyncLog 余额同步日志 +type QuotaSyncLog struct { + Id int `json:"id" gorm:"primaryKey"` + UserId int `json:"user_id" gorm:"index:idx_user_id"` + RemoteUserId int `json:"remote_user_id" gorm:"index"` + SyncType string `json:"sync_type" gorm:"size:32;index:idx_sync_type"` + BeforeQuota int `json:"before_quota"` + AfterQuota int `json:"after_quota"` + ChangeAmount int `json:"change_amount"` + MasterQuota int `json:"master_quota"` + PendingQuota int `json:"pending_quota"` + RequestId string `json:"request_id"` + Model string `json:"model"` + Direction string `json:"direction" gorm:"size:16"` + Status string `json:"status" gorm:"size:16;default:'success';index:idx_status"` + ErrorMsg string `json:"error_msg"` + CreatedAt int64 `json:"created_at" gorm:"index:idx_created_at"` + SyncedAt int64 `json:"synced_at"` +} + +func (QuotaSyncLog) TableName() string { + return "quota_sync_logs" +} + +// CreateSyncLog 创建同步日志 +func CreateSyncLog(log *QuotaSyncLog) error { + if log.CreatedAt == 0 { + log.CreatedAt = time.Now().Unix() + } + return DB.Create(log).Error +} diff --git a/model/redemption.go b/model/redemption.go index 378976a..507c563 100644 --- a/model/redemption.go +++ b/model/redemption.go @@ -154,6 +154,21 @@ func Redeem(key string, userId int) (quota int, err error) { common.SysError("redemption failed: " + err.Error()) return 0, ErrRedeemFailed } + + // 触发余额更新回调(在事务外) + if quotaUpdateCallback != nil { + newQuota, _ := GetUserQuota(userId, true) + // 使用 defer/recover 保护回调 + func() { + defer func() { + if r := recover(); r != nil { + common.SysError(fmt.Sprintf("quota update callback panic in Redeem: %v", r)) + } + }() + quotaUpdateCallback(userId, newQuota) + }() + } + RecordLog(userId, LogTypeTopup, fmt.Sprintf("通过兑换码充值 %s,兑换码ID %d", logger.LogQuota(redemption.Quota), redemption.Id)) return redemption.Quota, nil } diff --git a/model/redemption_test.go b/model/redemption_test.go new file mode 100644 index 0000000..11179ed --- /dev/null +++ b/model/redemption_test.go @@ -0,0 +1,254 @@ +package model + +import ( + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupRedemptionDB(t *testing.T) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, _ := db.DB() + sqlDB.SetMaxOpenConns(1) + + origDB := DB + origLogDB := LOG_DB + DB = db + LOG_DB = db + common.UsingSQLite = true + common.RedisEnabled = false + + require.NoError(t, db.AutoMigrate(&User{}, &Redemption{}, &Log{})) + + t.Cleanup(func() { + DB = origDB + LOG_DB = origLogDB + sqlDB.Close() + }) + return db +} + +func TestRedeem_Success(t *testing.T) { + db := setupRedemptionDB(t) + + // 创建测试用户 + user := User{ + Id: 100, + Username: "redeem_test", + Password: "hashed_password", + Quota: 100000, + } + require.NoError(t, db.Create(&user).Error) + + // 创建兑换码 + redemption := Redemption{ + Id: 1, + UserId: 1, + Key: "test-key-123", + Status: common.RedemptionCodeStatusEnabled, + Quota: 50000, + CreatedTime: common.GetTimestamp(), + ExpiredTime: 0, // 永不过期 + } + require.NoError(t, db.Create(&redemption).Error) + + // 设置回调 + callbackCalled := false + var receivedUserId int + var receivedQuota int + + SetQuotaUpdateCallback(func(userId int, quota int) { + callbackCalled = true + receivedUserId = userId + receivedQuota = quota + }) + defer SetQuotaUpdateCallback(nil) + + // 执行兑换 + quota, err := Redeem("test-key-123", 100) + require.NoError(t, err) + assert.Equal(t, 50000, quota) + + // 验证回调被触发 + assert.True(t, callbackCalled) + assert.Equal(t, 100, receivedUserId) + assert.Equal(t, 150000, receivedQuota) // 原始 100000 + 兑换 50000 + + // 验证用户额度已更新 + var updatedUser User + require.NoError(t, db.First(&updatedUser, 100).Error) + assert.Equal(t, 150000, updatedUser.Quota) + + // 验证兑换码状态已更新 + var updatedRedemption Redemption + require.NoError(t, db.First(&updatedRedemption, 1).Error) + assert.Equal(t, common.RedemptionCodeStatusUsed, updatedRedemption.Status) + assert.Equal(t, 100, updatedRedemption.UsedUserId) +} + +func TestRedeem_InvalidKey(t *testing.T) { + setupRedemptionDB(t) + + // 测试无效的兑换码 + quota, err := Redeem("invalid-key", 100) + assert.Error(t, err) + assert.Equal(t, 0, quota) +} + +func TestRedeem_AlreadyUsed(t *testing.T) { + db := setupRedemptionDB(t) + + // 创建测试用户 + user := User{ + Id: 100, + Username: "redeem_used_test", + Password: "hashed_password", + Quota: 100000, + } + require.NoError(t, db.Create(&user).Error) + + // 创建已使用的兑换码 + redemption := Redemption{ + Id: 1, + UserId: 1, + Key: "used-key", + Status: common.RedemptionCodeStatusUsed, + Quota: 50000, + CreatedTime: common.GetTimestamp(), + RedeemedTime: common.GetTimestamp(), + UsedUserId: 99, + } + require.NoError(t, db.Create(&redemption).Error) + + // 尝试兑换已使用的兑换码 + quota, err := Redeem("used-key", 100) + assert.Error(t, err) + assert.Equal(t, 0, quota) + // 错误会被包装为 ErrRedeemFailed,但底层错误消息会记录在日志中 + assert.Equal(t, ErrRedeemFailed, err) +} + +func TestRedeem_Expired(t *testing.T) { + db := setupRedemptionDB(t) + + // 创建测试用户 + user := User{ + Id: 100, + Username: "redeem_expired_test", + Password: "hashed_password", + Quota: 100000, + } + require.NoError(t, db.Create(&user).Error) + + // 创建已过期的兑换码 + redemption := Redemption{ + Id: 1, + UserId: 1, + Key: "expired-key", + Status: common.RedemptionCodeStatusEnabled, + Quota: 50000, + CreatedTime: common.GetTimestamp() - 86400, + ExpiredTime: common.GetTimestamp() - 3600, // 已过期 + } + require.NoError(t, db.Create(&redemption).Error) + + // 尝试兑换已过期的兑换码 + quota, err := Redeem("expired-key", 100) + assert.Error(t, err) + assert.Equal(t, 0, quota) + // 错误会被包装为 ErrRedeemFailed + assert.Equal(t, ErrRedeemFailed, err) +} + +func TestRedeem_CallbackPanic(t *testing.T) { + db := setupRedemptionDB(t) + + // 创建测试用户 + user := User{ + Id: 100, + Username: "redeem_panic_test", + Password: "hashed_password", + Quota: 100000, + } + require.NoError(t, db.Create(&user).Error) + + // 创建兑换码 + redemption := Redemption{ + Id: 1, + UserId: 1, + Key: "panic-key", + Status: common.RedemptionCodeStatusEnabled, + Quota: 50000, + CreatedTime: common.GetTimestamp(), + } + require.NoError(t, db.Create(&redemption).Error) + + // 设置一个会 panic 的回调 + SetQuotaUpdateCallback(func(userId int, quota int) { + panic("callback panic in redeem") + }) + defer SetQuotaUpdateCallback(nil) + + // 执行兑换(回调 panic 不应影响兑换成功) + quota, err := Redeem("panic-key", 100) + require.NoError(t, err) + assert.Equal(t, 50000, quota) + + // 验证用户额度已更新 + var updatedUser User + require.NoError(t, db.First(&updatedUser, 100).Error) + assert.Equal(t, 150000, updatedUser.Quota) +} + +func TestRedeem_SyncedUser(t *testing.T) { + db := setupRedemptionDB(t) + + // 创建同步用户(从 Master 节点同步过来的) + user := User{ + Id: 100, + Username: "synced_redeem_test", + Password: "hashed_password", + Quota: 100000, + Source: common.UserSourceSynced, + RemoteUserId: 1000, + SyncedQuota: 100000, + } + require.NoError(t, db.Create(&user).Error) + + // 创建兑换码 + redemption := Redemption{ + Id: 1, + UserId: 1, + Key: "synced-redeem-key", + Status: common.RedemptionCodeStatusEnabled, + Quota: 50000, + CreatedTime: common.GetTimestamp(), + } + require.NoError(t, db.Create(&redemption).Error) + + // 设置回调 + callbackCalled := false + SetQuotaUpdateCallback(func(userId int, quota int) { + callbackCalled = true + }) + defer SetQuotaUpdateCallback(nil) + + // 执行兑换 + quota, err := Redeem("synced-redeem-key", 100) + require.NoError(t, err) + assert.Equal(t, 50000, quota) + + // 验证回调被触发(Master 节点需要同步到 Slave) + assert.True(t, callbackCalled) + + // 验证用户额度已更新 + var updatedUser User + require.NoError(t, db.First(&updatedUser, 100).Error) + assert.Equal(t, 150000, updatedUser.Quota) +} diff --git a/model/region_sync_test.go b/model/region_sync_test.go new file mode 100644 index 0000000..1258c1f --- /dev/null +++ b/model/region_sync_test.go @@ -0,0 +1,538 @@ +package model + +import ( + "fmt" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupRegionSyncDB(t *testing.T) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, _ := db.DB() + sqlDB.SetMaxOpenConns(1) + + origDB := DB + origLogDB := LOG_DB + DB = db + LOG_DB = db + common.UsingSQLite = true + common.RedisEnabled = false + + require.NoError(t, db.AutoMigrate(&User{})) + + t.Cleanup(func() { + DB = origDB + LOG_DB = origLogDB + sqlDB.Close() + }) + return db +} + +func TestUser_IsSyncedUser_Method(t *testing.T) { + tests := []struct { + name string + user User + want bool + }{ + {"synced valid", User{Source: "synced", Id: 100}, true}, + {"synced boundary", User{Source: "synced", Id: common.ForeignUserIDStart - 1}, true}, + {"synced id 0", User{Source: "synced", Id: 0}, false}, + {"synced too large", User{Source: "synced", Id: common.ForeignUserIDStart}, false}, + {"local user", User{Source: "local", Id: 100}, false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.want, tt.user.IsSyncedUser()) + }) + } +} + +func TestUser_IsLocalUser_Method(t *testing.T) { + tests := []struct { + name string + user User + want bool + }{ + {"local source", User{Source: "local", Id: common.ForeignUserIDStart}, true}, + {"foreign id", User{Source: "", Id: common.ForeignUserIDStart + 1}, true}, + {"synced source", User{Source: "synced", Id: 100}, false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.want, tt.user.IsLocalUser()) + }) + } +} + +func TestUser_SyncFields_CRUD(t *testing.T) { + db := setupRegionSyncDB(t) + + user := User{ + Id: 100, Username: "synced_test", Password: "hashed_password", + Source: common.UserSourceSynced, RemoteUserId: 100, + SyncedQuota: 500000, LastSyncAt: 1700000000, + } + require.NoError(t, db.Create(&user).Error) + + var loaded User + require.NoError(t, db.First(&loaded, 100).Error) + assert.Equal(t, common.UserSourceSynced, loaded.Source) + assert.Equal(t, 100, loaded.RemoteUserId) + assert.Equal(t, 500000, loaded.SyncedQuota) + assert.Equal(t, int64(1700000000), loaded.LastSyncAt) +} + +func TestUser_DefaultSourceIsLocal(t *testing.T) { + db := setupRegionSyncDB(t) + + user := User{ + Id: 101, Username: "default_source_test", Password: "hashed_password", + } + require.NoError(t, db.Create(&user).Error) + + var loaded User + require.NoError(t, db.Where("username = ?", "default_source_test").First(&loaded).Error) + assert.Equal(t, common.UserSourceLocal, loaded.Source) +} + +// PendingSyncRecord tests + +func setupPendingSyncDB(t *testing.T) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, _ := db.DB() + sqlDB.SetMaxOpenConns(1) + + origDB := DB + origLogDB := LOG_DB + DB = db + LOG_DB = db + common.UsingSQLite = true + common.RedisEnabled = false + + require.NoError(t, db.AutoMigrate(&PendingSyncRecord{})) + + t.Cleanup(func() { + DB = origDB + LOG_DB = origLogDB + sqlDB.Close() + }) + return db +} + +func TestPendingSyncRecord_Create(t *testing.T) { + db := setupPendingSyncDB(t) + + userId := 100 + remoteUserId := 200 + requestId := "req-test-001" + quota := 1000 + estimatedQuota := 500 + + err := CreatePendingSyncRecord(userId, remoteUserId, requestId, quota, estimatedQuota) + require.NoError(t, err) + + var record PendingSyncRecord + require.NoError(t, db.Where("request_id = ?", requestId).First(&record).Error) + + assert.Equal(t, userId, record.UserId) + assert.Equal(t, remoteUserId, record.RemoteUserId) + assert.Equal(t, requestId, record.RequestId) + assert.Equal(t, quota, record.Quota) + assert.Equal(t, estimatedQuota, record.EstimatedQuota) + assert.Equal(t, PendingSyncStatusPending, record.Status) + assert.NotZero(t, record.CreatedAt) + assert.Zero(t, record.SyncedAt) + assert.Zero(t, record.RetryCount) +} + +func TestPendingSyncRecord_RequestIdUnique(t *testing.T) { + db := setupPendingSyncDB(t) + + requestId := "req-unique-001" + + // 创建第一条记录 + err := CreatePendingSyncRecord(100, 200, requestId, 1000, 500) + require.NoError(t, err) + + // 尝试用相同的 requestId 创建第二条记录 + err = CreatePendingSyncRecord(101, 201, requestId, 2000, 600) + assert.Error(t, err) // 应该失败,因为 requestId 唯一 + + // 验证只有一条记录 + var count int64 + db.Model(&PendingSyncRecord{}).Where("request_id = ?", requestId).Count(&count) + assert.Equal(t, int64(1), count) +} + +func TestGetPendingSyncQuota(t *testing.T) { + setupPendingSyncDB(t) + userId := 100 + + // 创建多条记录 + require.NoError(t, CreatePendingSyncRecord(userId, 200, "req-001", 1000, 500)) + require.NoError(t, CreatePendingSyncRecord(userId, 200, "req-002", 2000, 600)) + require.NoError(t, CreatePendingSyncRecord(userId, 200, "req-003", 500, 200)) + + // 创建一个 failed 状态的记录,应该被计入 + require.NoError(t, CreatePendingSyncRecord(userId, 200, "req-004", 300, 100)) + DB.Model(&PendingSyncRecord{}).Where("request_id = ?", "req-004").Update("status", PendingSyncStatusFailed) + + // 创建一个 synced 状态的记录,不应该被计入 + require.NoError(t, CreatePendingSyncRecord(userId, 200, "req-005", 9999, 0)) + DB.Model(&PendingSyncRecord{}).Where("request_id = ?", "req-005").Update("status", PendingSyncStatusSynced) + + total := GetPendingSyncQuota(userId) + assert.Equal(t, 3800, total) // 1000 + 2000 + 500 + 300 = 3800 +} + +func TestGetPendingSyncQuota_Empty(t *testing.T) { + setupPendingSyncDB(t) + + // 用户没有待同步记录 + total := GetPendingSyncQuota(99999) + assert.Equal(t, 0, total) +} + +// QuotaSyncLog tests + +func TestQuotaSyncLog_Create(t *testing.T) { + db := setupRegionSyncDB(t) + require.NoError(t, db.AutoMigrate(&QuotaSyncLog{})) + + log := QuotaSyncLog{ + UserId: 100, + RemoteUserId: 100, + SyncType: SyncTypeUserCreate, + BeforeQuota: 0, + AfterQuota: 500000, + ChangeAmount: 500000, + MasterQuota: 500000, + PendingQuota: 0, + Direction: SyncDirectionCnToOv, + Status: SyncStatusSuccess, + CreatedAt: 1700000000, + SyncedAt: 1700000000, + } + require.NoError(t, db.Create(&log).Error) + + var loaded QuotaSyncLog + require.NoError(t, db.First(&loaded, log.Id).Error) + assert.Equal(t, SyncTypeUserCreate, loaded.SyncType) + assert.Equal(t, 500000, loaded.ChangeAmount) + assert.Equal(t, SyncDirectionCnToOv, loaded.Direction) +} + +func TestQuotaSyncLog_QueryByUser(t *testing.T) { + db := setupRegionSyncDB(t) + require.NoError(t, db.AutoMigrate(&QuotaSyncLog{})) + + for i := 0; i < 3; i++ { + log := QuotaSyncLog{ + UserId: 100, RemoteUserId: 100, + SyncType: SyncTypeBatchSync, Direction: SyncDirectionOvToCn, + Status: SyncStatusSuccess, ChangeAmount: 100 * (i + 1), + CreatedAt: 1700000000 + int64(i), + } + require.NoError(t, db.Create(&log).Error) + } + + var logs []QuotaSyncLog + db.Where("user_id = ?", 100).Order("created_at desc").Find(&logs) + assert.Len(t, logs, 3) +} + +// --------------------------------------------------------------------------- +// 阶段 1:Model 层缺失测试 +// --------------------------------------------------------------------------- + +func TestGetPendingRecordsForSync_Basic(t *testing.T) { + db := setupPendingSyncDB(t) + + now := time.Now().Unix() + require.NoError(t, db.Create(&PendingSyncRecord{ + UserId: 100, RemoteUserId: 200, RequestId: "req-rs-001", + Quota: 100, Status: PendingSyncStatusPending, CreatedAt: now - 100, + }).Error) + require.NoError(t, db.Create(&PendingSyncRecord{ + UserId: 101, RemoteUserId: 201, RequestId: "req-rs-002", + Quota: 200, Status: PendingSyncStatusFailed, CreatedAt: now - 50, + }).Error) + require.NoError(t, db.Create(&PendingSyncRecord{ + UserId: 102, RemoteUserId: 202, RequestId: "req-rs-003", + Quota: 300, Status: PendingSyncStatusSynced, CreatedAt: now, + }).Error) + + records := GetPendingRecordsForSync(10, 3) + assert.Len(t, records, 2) + // 按 created_at 升序 + assert.Equal(t, "req-rs-001", records[0].RequestId) + assert.Equal(t, "req-rs-002", records[1].RequestId) +} + +func TestGetPendingRecordsForSync_MaxRetry(t *testing.T) { + db := setupPendingSyncDB(t) + + now := time.Now().Unix() + require.NoError(t, db.Create(&PendingSyncRecord{ + UserId: 100, RemoteUserId: 200, RequestId: "req-retry-001", + Quota: 100, Status: PendingSyncStatusFailed, RetryCount: 3, CreatedAt: now, + }).Error) + require.NoError(t, db.Create(&PendingSyncRecord{ + UserId: 101, RemoteUserId: 201, RequestId: "req-retry-002", + Quota: 200, Status: PendingSyncStatusPending, RetryCount: 2, CreatedAt: now, + }).Error) + + records := GetPendingRecordsForSync(10, 3) + assert.Len(t, records, 1) + assert.Equal(t, "req-retry-002", records[0].RequestId) +} + +func TestGetPendingRecordsForSync_Limit(t *testing.T) { + db := setupPendingSyncDB(t) + + now := time.Now().Unix() + for i := 0; i < 5; i++ { + require.NoError(t, db.Create(&PendingSyncRecord{ + UserId: 100, RemoteUserId: 200, RequestId: fmt.Sprintf("req-lim-%d", i), + Quota: 100, Status: PendingSyncStatusPending, CreatedAt: now + int64(i), + }).Error) + } + + records := GetPendingRecordsForSync(2, 3) + assert.Len(t, records, 2) +} + +func TestMarkRecordSynced(t *testing.T) { + db := setupPendingSyncDB(t) + + require.NoError(t, db.Create(&PendingSyncRecord{ + UserId: 100, RemoteUserId: 200, RequestId: "req-sync-001", + Quota: 100, Status: PendingSyncStatusPending, CreatedAt: time.Now().Unix(), + }).Error) + + var record PendingSyncRecord + require.NoError(t, db.Where("request_id = ?", "req-sync-001").First(&record).Error) + + require.NoError(t, MarkRecordSynced(record.Id)) + + var updated PendingSyncRecord + require.NoError(t, db.First(&updated, record.Id).Error) + assert.Equal(t, PendingSyncStatusSynced, updated.Status) + assert.NotZero(t, updated.SyncedAt) +} + +func TestMarkRecordFailed(t *testing.T) { + db := setupPendingSyncDB(t) + + require.NoError(t, db.Create(&PendingSyncRecord{ + UserId: 100, RemoteUserId: 200, RequestId: "req-fail-001", + Quota: 100, Status: PendingSyncStatusPending, RetryCount: 0, CreatedAt: time.Now().Unix(), + }).Error) + + var record PendingSyncRecord + require.NoError(t, db.Where("request_id = ?", "req-fail-001").First(&record).Error) + + require.NoError(t, MarkRecordFailed(record.Id, "connection refused")) + + var updated PendingSyncRecord + require.NoError(t, db.First(&updated, record.Id).Error) + assert.Equal(t, PendingSyncStatusFailed, updated.Status) + assert.Equal(t, 1, updated.RetryCount) + assert.Equal(t, "connection refused", updated.ErrorMsg) +} + +func TestCleanupOldSyncedRecords(t *testing.T) { + db := setupPendingSyncDB(t) + + now := time.Now().Unix() + // 40 天前已同步的记录 + require.NoError(t, db.Create(&PendingSyncRecord{ + UserId: 100, RemoteUserId: 200, RequestId: "req-old-001", + Quota: 100, Status: PendingSyncStatusSynced, SyncedAt: now - 40*86400, CreatedAt: now - 45*86400, + }).Error) + // 5 天前已同步的记录(不应被清理) + require.NoError(t, db.Create(&PendingSyncRecord{ + UserId: 100, RemoteUserId: 200, RequestId: "req-old-002", + Quota: 200, Status: PendingSyncStatusSynced, SyncedAt: now - 5*86400, CreatedAt: now - 10*86400, + }).Error) + // pending 状态的记录(不应被清理) + require.NoError(t, db.Create(&PendingSyncRecord{ + UserId: 100, RemoteUserId: 200, RequestId: "req-old-003", + Quota: 300, Status: PendingSyncStatusPending, CreatedAt: now - 50*86400, + }).Error) + + deleted := CleanupOldSyncedRecords(30) + assert.Equal(t, int64(1), deleted) + + var count int64 + db.Model(&PendingSyncRecord{}).Count(&count) + assert.Equal(t, int64(2), count) +} + +func TestCleanupOldSyncedRecords_Empty(t *testing.T) { + setupPendingSyncDB(t) + + deleted := CleanupOldSyncedRecords(30) + assert.Equal(t, int64(0), deleted) +} + +func TestGetPendingQuotaByUser(t *testing.T) { + db := setupPendingSyncDB(t) + + now := time.Now().Unix() + require.NoError(t, db.Create(&PendingSyncRecord{ + UserId: 100, RemoteUserId: 200, RequestId: "req-pq-001", + Quota: 100, Status: PendingSyncStatusPending, CreatedAt: now, + }).Error) + require.NoError(t, db.Create(&PendingSyncRecord{ + UserId: 100, RemoteUserId: 200, RequestId: "req-pq-002", + Quota: 200, Status: PendingSyncStatusFailed, CreatedAt: now, + }).Error) + require.NoError(t, db.Create(&PendingSyncRecord{ + UserId: 101, RemoteUserId: 201, RequestId: "req-pq-003", + Quota: 500, Status: PendingSyncStatusPending, CreatedAt: now, + }).Error) + require.NoError(t, db.Create(&PendingSyncRecord{ + UserId: 100, RemoteUserId: 200, RequestId: "req-pq-004", + Quota: 9999, Status: PendingSyncStatusSynced, CreatedAt: now, + }).Error) + + m := GetPendingQuotaByUser() + assert.Equal(t, 300, m[100]) // 100 + 200, synced 不计入 + assert.Equal(t, 500, m[101]) + _, exists := m[102] + assert.False(t, exists) +} + +// --------------------------------------------------------------------------- +// QuotaUpdateCallback tests +// --------------------------------------------------------------------------- + +func TestSetQuotaUpdateCallback(t *testing.T) { + // 测试设置回调函数 + called := false + var callbackUserId int + var callbackQuota int + + callback := func(userId int, quota int) { + called = true + callbackUserId = userId + callbackQuota = quota + } + + SetQuotaUpdateCallback(callback) + + // 清理:重置为 nil + defer SetQuotaUpdateCallback(nil) + + // 验证回调被设置(通过调用测试) + if quotaUpdateCallback != nil { + quotaUpdateCallback(123, 500000) + } + + assert.True(t, called) + assert.Equal(t, 123, callbackUserId) + assert.Equal(t, 500000, callbackQuota) +} + +func TestIncreaseUserQuota_TriggersCallback(t *testing.T) { + db := setupRegionSyncDB(t) + + // 创建测试用户 + user := User{ + Id: 1000, + Username: "callback_test", + Password: "hashed_password", + Quota: 100000, + } + require.NoError(t, db.Create(&user).Error) + + // 设置回调 + callbackCalled := false + var receivedUserId int + var receivedQuota int + + SetQuotaUpdateCallback(func(userId int, quota int) { + callbackCalled = true + receivedUserId = userId + receivedQuota = quota + }) + defer SetQuotaUpdateCallback(nil) // 清理 + + // 执行增加额度操作 + err := IncreaseUserQuota(1000, 50000, true) + require.NoError(t, err) + + // 验证回调被触发 + assert.True(t, callbackCalled) + assert.Equal(t, 1000, receivedUserId) + assert.Equal(t, 150000, receivedQuota) // 原始 100000 + 增加 50000 + + // 验证数据库已更新 + var updated User + require.NoError(t, db.First(&updated, 1000).Error) + assert.Equal(t, 150000, updated.Quota) +} + +func TestIncreaseUserQuota_NoCallback(t *testing.T) { + db := setupRegionSyncDB(t) + + // 创建测试用户 + user := User{ + Id: 1001, + Username: "no_callback_test", + Password: "hashed_password", + Quota: 100000, + } + require.NoError(t, db.Create(&user).Error) + + // 确保没有回调 + SetQuotaUpdateCallback(nil) + + // 执行增加额度操作(不应 panic) + err := IncreaseUserQuota(1001, 50000, true) + require.NoError(t, err) + + // 验证数据库已更新 + var updated User + require.NoError(t, db.First(&updated, 1001).Error) + assert.Equal(t, 150000, updated.Quota) +} + +func TestIncreaseUserQuota_CallbackError(t *testing.T) { + db := setupRegionSyncDB(t) + + // 创建测试用户 + user := User{ + Id: 1002, + Username: "callback_error_test", + Password: "hashed_password", + Quota: 100000, + } + require.NoError(t, db.Create(&user).Error) + + // 设置一个会 panic 的回调,验证不影响主流程 + SetQuotaUpdateCallback(func(userId int, quota int) { + panic("callback error") + }) + defer SetQuotaUpdateCallback(nil) + + // 执行增加额度操作(回调 panic 不应影响数据库更新) + err := IncreaseUserQuota(1002, 50000, true) + require.NoError(t, err) + + // 验证数据库已更新(即使回调失败) + var updated User + require.NoError(t, db.First(&updated, 1002).Error) + assert.Equal(t, 150000, updated.Quota) +} diff --git a/model/user.go b/model/user.go index f3e66e6..6b43423 100644 --- a/model/user.go +++ b/model/user.go @@ -6,6 +6,7 @@ import ( "fmt" "strconv" "strings" + "time" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/dto" @@ -17,6 +18,16 @@ import ( const UserNameMaxLength = 20 +// QuotaUpdateCallback 余额更新回调函数类型 +type QuotaUpdateCallback func(userId int, quota int) + +var quotaUpdateCallback QuotaUpdateCallback + +// SetQuotaUpdateCallback 设置余额更新回调 +func SetQuotaUpdateCallback(cb QuotaUpdateCallback) { + quotaUpdateCallback = cb +} + // User if you add sensitive fields, don't forget to clean them in setupLogin function. // Otherwise, the sensitive information will be saved on local storage in plain text! type User struct { @@ -49,6 +60,21 @@ type User struct { Setting string `json:"setting" gorm:"type:text;column:setting"` Remark string `json:"remark,omitempty" gorm:"type:varchar(255)" validate:"max=255"` StripeCustomer string `json:"stripe_customer" gorm:"type:varchar(64);column:stripe_customer;index"` + // 跨地区同步相关 + Source string `json:"source" gorm:"type:varchar(20);default:'local'"` + RemoteUserId int `json:"remote_user_id" gorm:"type:int;default:0;column:remote_user_id"` + SyncedQuota int `json:"synced_quota" gorm:"type:int;default:0;column:synced_quota"` + LastSyncAt int64 `json:"last_sync_at" gorm:"type:bigint;default:0;column:last_sync_at"` +} + +// IsSyncedUser 判断是否为国内同步用户 +func (u *User) IsSyncedUser() bool { + return common.IsSyncedUser(u.Source, u.Id) +} + +// IsLocalUser 判断是否为本地用户 +func (u *User) IsLocalUser() bool { + return common.IsLocalUser(u.Source, u.Id) } func (user *User) ToBaseUser() *UserBase { @@ -60,6 +86,7 @@ func (user *User) ToBaseUser() *UserBase { Username: user.Username, Setting: user.Setting, Email: user.Email, + Source: user.Source, } return cache } @@ -767,17 +794,8 @@ func ValidateAccessToken(token string) (user *User) { } // GetUserQuota gets quota from Redis first, falls back to DB if needed +// 同步用户返回 SyncedQuota,本地用户返回 Quota func GetUserQuota(id int, fromDB bool) (quota int, err error) { - defer func() { - // Update Redis cache asynchronously on successful DB read - if shouldUpdateRedis(fromDB, err) { - gopool.Go(func() { - if err := updateUserQuotaCache(id, quota); err != nil { - common.SysLog("failed to update user quota cache: " + err.Error()) - } - }) - } - }() if !fromDB && common.RedisEnabled { quota, err := getUserQuotaCache(id) if err == nil { @@ -785,12 +803,31 @@ func GetUserQuota(id int, fromDB bool) (quota int, err error) { } // Don't return error - fall through to DB } + fromDB = true - err = DB.Model(&User{}).Where("id = ?", id).Select("quota").Find("a).Error + // 查询用户的 source 和额度字段 + var user User + err = DB.Model(&User{}).Select("source", "quota", "synced_quota").Where("id = ?", id).First(&user).Error if err != nil { return 0, err } + // 同步用户返回 SyncedQuota,本地用户返回 Quota + if user.IsSyncedUser() { + quota = user.SyncedQuota + } else { + quota = user.Quota + } + + // Update Redis cache asynchronously if enabled + if common.RedisEnabled { + gopool.Go(func() { + if err := updateUserQuotaCache(id, quota); err != nil { + common.SysLog("failed to update user quota cache: " + err.Error()) + } + }) + } + return quota, nil } @@ -877,7 +914,26 @@ func IncreaseUserQuota(id int, quota int, db bool) (err error) { addNewRecord(BatchUpdateTypeUserQuota, id, quota) return nil } - return increaseUserQuota(id, quota) + err = increaseUserQuota(id, quota) + if err != nil { + return err + } + + // 触发余额更新回调(如果有注册) + if quotaUpdateCallback != nil { + newQuota, _ := GetUserQuota(id, true) + // 使用 defer/recover 保护回调,防止 panic 影响主流程 + func() { + defer func() { + if r := recover(); r != nil { + common.SysError(fmt.Sprintf("quota update callback panic: %v", r)) + } + }() + quotaUpdateCallback(id, newQuota) + }() + } + + return nil } func increaseUserQuota(id int, quota int) (err error) { @@ -924,6 +980,47 @@ func DeltaUpdateUserQuota(id int, delta int) (err error) { } } +// GetSyncedUsers 获取所有同步用户 +func GetSyncedUsers() []User { + var users []User + DB.Where("source = ?", common.UserSourceSynced).Find(&users) + return users +} + +// UpdateSyncedQuota 更新同步用户的 synced_quota +func UpdateSyncedQuota(userId int, quota int) error { + return DB.Model(&User{}).Where("id = ?", userId).Updates(map[string]interface{}{ + "synced_quota": quota, + "last_sync_at": time.Now().Unix(), + }).Error +} + +// AtomicDecreaseSyncedQuota 原子扣减 synced_quota,检查余额并扣减在一条 SQL 中完成。 +// 返回 (当前余额, 是否成功, 错误)。 +func AtomicDecreaseSyncedQuota(userId int, amount int, threshold int) (int, bool, error) { + result := DB.Model(&User{}). + Where("id = ? AND synced_quota >= ?", userId, amount+threshold). + Update("synced_quota", gorm.Expr("synced_quota - ?", amount)) + if result.Error != nil { + return 0, false, result.Error + } + if result.RowsAffected == 0 { + var quota int + DB.Model(&User{}).Where("id = ?", userId).Select("synced_quota").Scan("a) + return quota, false, nil + } + var newQuota int + DB.Model(&User{}).Where("id = ?", userId).Select("synced_quota").Scan(&newQuota) + return newQuota, true, nil +} + +// IncreaseSyncedQuota 原子增加 synced_quota(用于退款) +func IncreaseSyncedQuota(userId int, amount int) error { + return DB.Model(&User{}). + Where("id = ?", userId). + Update("synced_quota", gorm.Expr("synced_quota + ?", amount)).Error +} + //func GetRootUserEmail() (email string) { // DB.Model(&User{}).Where("role = ?", common.RoleRootUser).Select("email").Find(&email) // return email diff --git a/model/user_cache.go b/model/user_cache.go index 2ba1f18..9b39f93 100644 --- a/model/user_cache.go +++ b/model/user_cache.go @@ -15,6 +15,7 @@ import ( // UserBase struct remains the same as it represents the cached data structure type UserBase struct { + Source string `json:"source"` Id int `json:"id"` Group string `json:"group"` Email string `json:"email"` @@ -31,6 +32,7 @@ func (user *UserBase) WriteContext(c *gin.Context) { common.SetContextKey(c, constant.ContextKeyUserEmail, user.Email) common.SetContextKey(c, constant.ContextKeyUserName, user.Username) common.SetContextKey(c, constant.ContextKeyUserSetting, user.GetSetting()) + common.SetContextKey(c, constant.ContextKeyUserSource, user.Source) } func (user *UserBase) GetSetting() dto.UserSetting { @@ -101,6 +103,7 @@ func GetUserCache(userId int) (userCache *UserBase, err error) { // Create cache object from user data userCache = &UserBase{ Id: user.Id, + Source: user.Source, Group: user.Group, Quota: user.Quota, Status: user.Status, diff --git a/relay/helper/price.go b/relay/helper/price.go index fc82e31..b63467a 100644 --- a/relay/helper/price.go +++ b/relay/helper/price.go @@ -87,9 +87,8 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens if meta.MaxTokens != 0 { preConsumedTokens += meta.MaxTokens } - // 如果没有找到渠道定价,需要获取全局 modelRatio 和 completionRatio - // 但如果 ChannelMeta 为 nil(渠道选择前),跳过检查,延迟到渠道选择后 - if channelMetaAvailable && !channelPricingFound { + // 如果没有找到渠道定价,获取全局 modelRatio 和 completionRatio + if !channelPricingFound { var success bool var matchName string modelRatio, success, matchName = ratio_setting.GetModelRatio(info.OriginModelName) diff --git a/router/api-router.go b/router/api-router.go index 278c070..80421dc 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -388,5 +388,15 @@ func SetApiRouter(router *gin.Engine) { deploymentsRoute.POST("/:id/extend", controller.ExtendDeployment) deploymentsRoute.DELETE("/:id", controller.DeleteDeployment) } + + // Region sync API routes (for inter-node communication) + syncRoute := apiRouter.Group("/internal/sync") + syncRoute.Use(middleware.SyncAuth()) + { + syncRoute.POST("/user/create", controller.ReceiveSyncedUserCreate) + syncRoute.POST("/quota/update", controller.ReceiveQuotaUpdate) + syncRoute.POST("/quota/query", controller.QueryUserQuota) + syncRoute.POST("/quota/batch-deduct", controller.BatchDeductQuota) + } } } diff --git a/service/billing_session.go b/service/billing_session.go index f24b68e..a03efca 100644 --- a/service/billing_session.go +++ b/service/billing_session.go @@ -7,6 +7,7 @@ import ( "sync" "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/logger" "github.com/QuantumNous/new-api/model" relaycommon "github.com/QuantumNous/new-api/relay/common" @@ -132,6 +133,10 @@ func (s *BillingSession) needsRefundLocked() bool { if sub, ok := s.funding.(*SubscriptionFunding); ok && sub.preConsumed > 0 { return true } + // 同步用户可能在 tokenConsumed=0 时仍预扣了额度 + if synced, ok := s.funding.(*SyncedUserFunding); ok && synced.consumed > 0 { + return true + } return false } @@ -257,6 +262,18 @@ func NewBillingSession(c *gin.Context, relayInfo *relaycommon.RelayInfo, preCons return nil, types.NewError(fmt.Errorf("relayInfo is nil"), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry()) } + // 检查是否为同步用户(从上下文中获取,避免热路径数据库查询) + userSource := common.GetContextKeyString(c, constant.ContextKeyUserSource) + common.SysLog(fmt.Sprintf("[RegionSync] NewBillingSession: userId=%d, userSource=%s, preConsumedQuota=%d, requestId=%s", relayInfo.UserId, userSource, preConsumedQuota, relayInfo.RequestId)) + if userSource == common.UserSourceSynced { + user, userErr := model.GetUserById(relayInfo.UserId, false) + if userErr == nil { + common.SysLog(fmt.Sprintf("[RegionSync] synced user detected: userId=%d, remoteUserId=%d, syncedQuota=%d, requestId=%s", user.Id, user.RemoteUserId, user.SyncedQuota, relayInfo.RequestId)) + return newSyncedUserBillingSession(c, relayInfo, user, preConsumedQuota) + } + common.SysError(fmt.Sprintf("[RegionSync] GetUserById failed for synced user: userId=%d, err=%v", relayInfo.UserId, userErr)) + } + pref := common.NormalizeBillingPreference(relayInfo.UserSetting.BillingPreference) // 钱包路径需要先检查用户额度 @@ -345,3 +362,29 @@ func NewBillingSession(c *gin.Context, relayInfo *relaycommon.RelayInfo, preCons return session, nil } } + +// newSyncedUserBillingSession 为同步用户创建 BillingSession +// 同步用户的余额在 master 节点,本地使用 SyncedUserFunding 管理预扣费和结算 +func newSyncedUserBillingSession(c *gin.Context, relayInfo *relaycommon.RelayInfo, user *model.User, preConsumedQuota int) (*BillingSession, *types.NewAPIError) { + syncedQuota := user.SyncedQuota + if syncedQuota <= 0 { + syncedQuota = user.Quota + } + + session := &BillingSession{ + relayInfo: relayInfo, + funding: NewSyncedUserFunding( + user.Id, + user.RemoteUserId, + relayInfo.RequestId, + syncedQuota, + ), + } + + relayInfo.UserQuota = syncedQuota + + if apiErr := session.preConsume(c, preConsumedQuota); apiErr != nil { + return nil, apiErr + } + return session, nil +} diff --git a/service/billing_session_sync_test.go b/service/billing_session_sync_test.go new file mode 100644 index 0000000..a6acbfb --- /dev/null +++ b/service/billing_session_sync_test.go @@ -0,0 +1,135 @@ +package service + +import ( + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting/system_setting" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupBillingSyncDB(t *testing.T) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, _ := db.DB() + sqlDB.SetMaxOpenConns(1) + + origDB := model.DB + origLogDB := model.LOG_DB + model.DB = db + model.LOG_DB = db + common.UsingSQLite = true + common.RedisEnabled = false + + require.NoError(t, db.AutoMigrate(&model.PendingSyncRecord{})) + + t.Cleanup(func() { + model.DB = origDB + model.LOG_DB = origLogDB + sqlDB.Close() + }) + return db +} + +// TestNeedsRefundLocked_SyncedUser 测试同步用户在 consumed > 0 时需要退款 +func TestNeedsRefundLocked_SyncedUser(t *testing.T) { + setupBillingSyncDB(t) + + settings := system_setting.GetRegionSyncSettings() + origThreshold := settings.MinBalanceThreshold + defer func() { settings.MinBalanceThreshold = origThreshold }() + settings.MinBalanceThreshold = 100000 + + f := NewSyncedUserFunding(100, 200, "req-nr-001", 500000) + require.NoError(t, f.PreConsume(5000)) + + session := &BillingSession{ + funding: f, + } + assert.True(t, session.needsRefundLocked()) +} + +// TestNeedsRefundLocked_SyncedUser_NoConsumed 测试同步用户在 consumed == 0 时不需要退款 +func TestNeedsRefundLocked_SyncedUser_NoConsumed(t *testing.T) { + setupBillingSyncDB(t) + + f := NewSyncedUserFunding(100, 200, "req-nr-002", 500000) + + session := &BillingSession{ + funding: f, + } + assert.False(t, session.needsRefundLocked()) +} + +// TestNeedsRefundLocked_SyncedUser_FundingSettled 测试 fundingSettled=true 后不需要退款 +func TestNeedsRefundLocked_SyncedUser_FundingSettled(t *testing.T) { + setupBillingSyncDB(t) + + settings := system_setting.GetRegionSyncSettings() + origThreshold := settings.MinBalanceThreshold + defer func() { settings.MinBalanceThreshold = origThreshold }() + settings.MinBalanceThreshold = 100000 + + f := NewSyncedUserFunding(100, 200, "req-nr-003", 500000) + require.NoError(t, f.PreConsume(5000)) + + session := &BillingSession{ + funding: f, + fundingSettled: true, // 资金已提交 + } + // fundingSettled=true 时不应再退 + assert.False(t, session.needsRefundLocked()) +} + +// TestShouldTrust_SyncedUser_SourceNotWalletOrSub 验证同步用户 Source 不是 wallet/subscription, +// 因此在 shouldTrust 的 switch 中走 default 分支返回 false +func TestShouldTrust_SyncedUser_SourceNotWalletOrSub(t *testing.T) { + f := NewSyncedUserFunding(100, 200, "req-trust", 500000) + assert.Equal(t, "synced_wallet", f.Source()) + assert.NotEqual(t, BillingSourceWallet, f.Source()) + assert.NotEqual(t, BillingSourceSubscription, f.Source()) +} + +// TestSyncedUserFunding_BillingSource 测试同步用户计费来源标识 +func TestSyncedUserFunding_BillingSource(t *testing.T) { + f := NewSyncedUserFunding(100, 200, "req-src", 500000) + assert.Equal(t, BillingSourceSyncedWallet, f.Source()) + assert.Equal(t, "synced_wallet", f.Source()) +} + +// TestNeedsRefundLocked_WalletAndSynced 对比钱包和同步用户的行为 +func TestNeedsRefundLocked_WalletAndSynced(t *testing.T) { + // 钱包用户:tokenConsumed > 0 才需要退 + walletSession := &BillingSession{ + funding: &WalletFunding{userId: 1, consumed: 100}, + tokenConsumed: 100, + } + assert.True(t, walletSession.needsRefundLocked()) + + walletSessionZero := &BillingSession{ + funding: &WalletFunding{userId: 1, consumed: 100}, + tokenConsumed: 0, + } + assert.False(t, walletSessionZero.needsRefundLocked()) + + // 同步用户:consumed > 0 就需要退(即使 tokenConsumed=0) + setupBillingSyncDB(t) + settings := system_setting.GetRegionSyncSettings() + origThreshold := settings.MinBalanceThreshold + defer func() { settings.MinBalanceThreshold = origThreshold }() + settings.MinBalanceThreshold = 100000 + + syncF := NewSyncedUserFunding(100, 200, "req-compare", 500000) + require.NoError(t, syncF.PreConsume(5000)) + + syncSession := &BillingSession{ + funding: syncF, + tokenConsumed: 0, // tokenConsumed=0 但 consumed > 0 + } + assert.True(t, syncSession.needsRefundLocked()) +} diff --git a/service/region_sync/sync_client.go b/service/region_sync/sync_client.go new file mode 100644 index 0000000..aeb9ddc --- /dev/null +++ b/service/region_sync/sync_client.go @@ -0,0 +1,132 @@ +package region_sync + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/logger" +) + +// SyncClient 调用远程节点的 HTTP 客户端 +type SyncClient struct { + endpoint string + apiKey string + httpClient *http.Client +} + +func NewSyncClient(endpoint, apiKey string) *SyncClient { + return &SyncClient{ + endpoint: endpoint, + apiKey: apiKey, + httpClient: &http.Client{Timeout: 30 * time.Second}, + } +} + +func (c *SyncClient) doRequest(method, path string, body interface{}) ([]byte, error) { + var reqBody io.Reader + if body != nil { + data, err := json.Marshal(body) + if err != nil { + return nil, fmt.Errorf("marshal request: %w", err) + } + reqBody = bytes.NewReader(data) + } + + // 处理 URL 拼接,避免双斜杠问题 + endpoint := strings.TrimRight(c.endpoint, "/") + path = strings.TrimLeft(path, "/") + url := endpoint + "/" + path + + common.SysLog(fmt.Sprintf("[RegionSync] HTTP %s %s", method, url)) + logger.LogDebug(nil, "[RegionSync] HTTP %s %s", method, url) + + req, err := http.NewRequest(method, url, reqBody) + if err != nil { + return nil, fmt.Errorf("create request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("X-Sync-API-Key", c.apiKey) + req.Header.Set("X-Sync-Node", c.endpoint) + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("http request: %w", err) + } + defer resp.Body.Close() + + data, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("read response: %w", err) + } + + // 如果返回状态码不是 200,记录错误信息 + if resp.StatusCode != http.StatusOK { + common.SysError(fmt.Sprintf("[RegionSync] HTTP error %d from %s: %s", resp.StatusCode, url, string(data))) + return nil, fmt.Errorf("http status %d: %s", resp.StatusCode, string(data)) + } + + return data, nil +} + +func (c *SyncClient) SyncUserCreate(req *SyncUserRequest) (*SyncUserResponse, error) { + data, err := c.doRequest("POST", "/api/internal/sync/user/create", req) + if err != nil { + return nil, err + } + var resp SyncUserResponse + if err := json.Unmarshal(data, &resp); err != nil { + return nil, fmt.Errorf("unmarshal response: %w", err) + } + return &resp, nil +} + +func (c *SyncClient) QueryQuota(userId int) (*QueryQuotaResponse, error) { + data, err := c.doRequest("POST", "/api/internal/sync/quota/query", &QueryQuotaRequest{UserId: userId}) + if err != nil { + return nil, err + } + var resp QueryQuotaResponse + if err := json.Unmarshal(data, &resp); err != nil { + return nil, fmt.Errorf("unmarshal response: %w", err) + } + logger.LogDebug(nil, "[RegionSync] QueryQuota: userId=%d, quota=%d, success=%v", userId, resp.Quota, resp.Success) + return &resp, nil +} + +func (c *SyncClient) BatchDeduct(req *BatchDeductRequest) (*BatchDeductResponse, error) { + logger.LogDebug(nil, "[RegionSync] BatchDeduct: sending %d records", len(req.Records)) + data, err := c.doRequest("POST", "/api/internal/sync/quota/batch-deduct", req) + if err != nil { + return nil, err + } + var resp BatchDeductResponse + if err := json.Unmarshal(data, &resp); err != nil { + return nil, fmt.Errorf("unmarshal response: %w", err) + } + successCount := 0 + for _, r := range resp.Results { + if r.Success { + successCount++ + } + } + logger.LogDebug(nil, "[RegionSync] BatchDeduct: response success=%d, total=%d", successCount, len(resp.Results)) + return &resp, nil +} + +func (c *SyncClient) UpdateQuota(remoteUserId, quota int) (*UpdateQuotaResponse, error) { + data, err := c.doRequest("POST", "/api/internal/sync/quota/update", &UpdateQuotaRequest{RemoteUserId: remoteUserId, Quota: quota}) + if err != nil { + return nil, err + } + var resp UpdateQuotaResponse + if err := json.Unmarshal(data, &resp); err != nil { + return nil, fmt.Errorf("unmarshal response: %w", err) + } + return &resp, nil +} diff --git a/service/region_sync/sync_client_test.go b/service/region_sync/sync_client_test.go new file mode 100644 index 0000000..debff7b --- /dev/null +++ b/service/region_sync/sync_client_test.go @@ -0,0 +1,113 @@ +package region_sync + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestSyncClient_SyncUserCreate(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "POST", r.Method) + assert.Equal(t, "/api/internal/sync/user/create", r.URL.Path) + assert.Equal(t, "test-sync-key", r.Header.Get("X-Sync-API-Key")) + + var req SyncUserRequest + json.NewDecoder(r.Body).Decode(&req) + assert.Equal(t, "testuser", req.Username) + assert.Equal(t, 123, req.RemoteUserId) + + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(SyncUserResponse{Success: true}) + })) + defer server.Close() + + client := NewSyncClient(server.URL, "test-sync-key") + resp, err := client.SyncUserCreate(&SyncUserRequest{ + Username: "testuser", + RemoteUserId: 123, + }) + require.NoError(t, err) + assert.True(t, resp.Success) +} + +func TestSyncClient_QueryQuota(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(QueryQuotaResponse{Success: true, Quota: 50000}) + })) + defer server.Close() + + client := NewSyncClient(server.URL, "test-key") + resp, err := client.QueryQuota(123) + require.NoError(t, err) + assert.True(t, resp.Success) + assert.Equal(t, 50000, resp.Quota) +} + +func TestSyncClient_BatchDeduct(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var req BatchDeductRequest + json.NewDecoder(r.Body).Decode(&req) + assert.Len(t, req.Records, 2) + + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(BatchDeductResponse{ + Success: true, + Results: []DeductResult{ + {UserId: 100, Success: true}, + {UserId: 101, Success: true}, + }, + }) + })) + defer server.Close() + + client := NewSyncClient(server.URL, "test-key") + resp, err := client.BatchDeduct(&BatchDeductRequest{ + Records: []BatchDeductRecord{ + {UserId: 100, RequestId: "req-1", Quota: 100}, + {UserId: 101, RequestId: "req-2", Quota: 200}, + }, + }) + require.NoError(t, err) + assert.True(t, resp.Success) + assert.Len(t, resp.Results, 2) +} + +func TestSyncClient_UpdateQuota(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "/api/internal/sync/quota/update", r.URL.Path) + + var req UpdateQuotaRequest + json.NewDecoder(r.Body).Decode(&req) + assert.Equal(t, 123, req.RemoteUserId) + assert.Equal(t, 50000, req.Quota) + + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(UpdateQuotaResponse{Success: true}) + })) + defer server.Close() + + client := NewSyncClient(server.URL, "test-key") + resp, err := client.UpdateQuota(123, 50000) + require.NoError(t, err) + assert.True(t, resp.Success) +} + +func TestSyncClient_DoRequest_Error(t *testing.T) { + // 测试请求错误处理 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // 返回错误状态码 + w.WriteHeader(http.StatusInternalServerError) + })) + defer server.Close() + + client := NewSyncClient(server.URL, "test-key") + _, err := client.QueryQuota(123) + // 应该有错误,但不是因为 marshaling/unmarshaling + assert.NotNil(t, err) +} diff --git a/service/region_sync/sync_manager.go b/service/region_sync/sync_manager.go new file mode 100644 index 0000000..a7d3a31 --- /dev/null +++ b/service/region_sync/sync_manager.go @@ -0,0 +1,216 @@ +package region_sync + +import ( + "fmt" + "sync" + "sync/atomic" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting/system_setting" + "github.com/bytedance/gopkg/util/gopool" +) + +// SyncManager 管理跨区域同步的后台任务 +type SyncManager struct { + client *SyncClient + stopped atomic.Bool + lastEndpoint string + lastApiKey string +} + +var ( + singletonManager *SyncManager + managerMu sync.Mutex +) + +// NewSyncManager 创建同步管理器(配置变化时自动重建客户端) +func NewSyncManager() *SyncManager { + managerMu.Lock() + defer managerMu.Unlock() + + settings := system_setting.GetRegionSyncSettings() + + if singletonManager == nil { + singletonManager = &SyncManager{ + client: NewSyncClient(settings.MasterEndpoint, settings.SyncApiKey), + } + singletonManager.lastEndpoint = settings.MasterEndpoint + singletonManager.lastApiKey = settings.SyncApiKey + return singletonManager + } + + // 配置变化时重建客户端 + if settings.MasterEndpoint != singletonManager.lastEndpoint || settings.SyncApiKey != singletonManager.lastApiKey { + singletonManager.client = NewSyncClient(settings.MasterEndpoint, settings.SyncApiKey) + singletonManager.lastEndpoint = settings.MasterEndpoint + singletonManager.lastApiKey = settings.SyncApiKey + } + + return singletonManager +} + +// RunBatchSync 执行批量同步 +func (m *SyncManager) RunBatchSync() int { + settings := system_setting.GetRegionSyncSettings() + if !settings.Enabled { + return 0 + } + + records := model.GetPendingRecordsForSync(settings.SyncBatchSize, settings.MaxRetryCount) + if len(records) == 0 { + return 0 + } + + logger.LogDebug(nil, "[RegionSync] RunBatchSync: got %d pending records", len(records)) + + syncCount := 0 + batchRecords := make([]BatchDeductRecord, 0, len(records)) + + for _, record := range records { + logger.LogDebug(nil, "[RegionSync] RunBatchSync: record userId=%d, requestId=%s, quota=%d, retryCount=%d", record.RemoteUserId, record.RequestId, record.Quota, record.RetryCount) + batchRecords = append(batchRecords, BatchDeductRecord{ + UserId: record.RemoteUserId, + RequestId: record.RequestId, + Quota: record.Quota, + }) + } + + resp, err := m.client.BatchDeduct(&BatchDeductRequest{Records: batchRecords}) + if err != nil { + common.SysError(fmt.Sprintf("[SyncManager] BatchDeduct failed: %v", err)) + for _, record := range records { + if record.RetryCount+1 >= settings.MaxRetryCount { + if markErr := model.MarkRecordFailed(record.Id, err.Error()); markErr != nil { + common.SysError(fmt.Sprintf("[SyncManager] MarkRecordFailed error: %v", markErr)) + } + } + } + return 0 + } + + for i, result := range resp.Results { + if result.Success { + if markErr := model.MarkRecordSynced(records[i].Id); markErr != nil { + common.SysError(fmt.Sprintf("[SyncManager] MarkRecordSynced error: %v", markErr)) + } + syncCount++ + } else { + if records[i].RetryCount+1 >= settings.MaxRetryCount { + if markErr := model.MarkRecordFailed(records[i].Id, result.Error); markErr != nil { + common.SysError(fmt.Sprintf("[SyncManager] MarkRecordFailed error: %v", markErr)) + } + } + } + } + + logger.LogDebug(nil, "[RegionSync] RunBatchSync: completed, success=%d, total=%d", syncCount, len(records)) + return syncCount +} + +// QueryMasterQuota 查询主节点余额(Slave 节点调用) +func (m *SyncManager) QueryMasterQuota(remoteUserId int) (int, error) { + settings := system_setting.GetRegionSyncSettings() + if !settings.Enabled || settings.IsMaster { + return 0, fmt.Errorf("region sync is not enabled or this is a master node (only slave can query)") + } + + logger.LogDebug(nil, "[RegionSync] QueryMasterQuota: remoteUserId=%d", remoteUserId) + + resp, err := m.client.QueryQuota(remoteUserId) + if err != nil { + return 0, fmt.Errorf("query master quota failed: %w", err) + } + + if !resp.Success { + return 0, fmt.Errorf("query failed: %s", resp.Error) + } + + logger.LogDebug(nil, "[RegionSync] QueryMasterQuota: remoteUserId=%d, quota=%d", remoteUserId, resp.Quota) + return resp.Quota, nil +} + +// RunQuotaSync 从 Master 拉取所有同步用户的最新余额,更新本地 synced_quota +func (m *SyncManager) RunQuotaSync() int { + settings := system_setting.GetRegionSyncSettings() + if !settings.Enabled || settings.IsMaster { + return 0 + } + + users := model.GetSyncedUsers() + if len(users) == 0 { + return 0 + } + + logger.LogDebug(nil, "[RegionSync] RunQuotaSync: syncing %d users", len(users)) + + syncedCount := 0 + for _, user := range users { + resp, err := m.client.QueryQuota(user.RemoteUserId) + if err != nil { + common.SysError(fmt.Sprintf("[SyncManager] QueryQuota failed for user %d: %v", user.Id, err)) + continue + } + if !resp.Success { + continue + } + oldQuota := user.SyncedQuota + if err := model.UpdateSyncedQuota(user.Id, resp.Quota); err != nil { + common.SysError(fmt.Sprintf("[SyncManager] UpdateSyncedQuota failed for user %d: %v", user.Id, err)) + continue + } + logger.LogDebug(nil, "[RegionSync] RunQuotaSync: userId=%d, syncedQuota %d -> %d", user.Id, oldQuota, resp.Quota) + syncedCount++ + } + logger.LogDebug(nil, "[RegionSync] RunQuotaSync: completed, synced=%d/%d", syncedCount, len(users)) + return syncedCount +} + +// RunCleanup 清理已同步的旧记录 +func (m *SyncManager) RunCleanup() int { + return int(model.CleanupOldSyncedRecords(30)) +} + +// StartSyncWorkers 启动后台同步工作器 +func (m *SyncManager) StartSyncWorkers() { + m.stopped.Store(false) + + settings := system_setting.GetRegionSyncSettings() + if !settings.Enabled { + return + } + + interval := time.Duration(settings.SyncIntervalSeconds) * time.Second + + startWorker := func(tickInterval time.Duration, fn func() int, label string) { + gopool.Go(func() { + ticker := time.NewTicker(tickInterval) + defer ticker.Stop() + + for { + if m.stopped.Load() { + return + } + <-ticker.C + if m.stopped.Load() { + return + } + if count := fn(); count > 0 { + common.SysLog(fmt.Sprintf("[SyncManager] %s: %d", label, count)) + } + } + }) + } + + startWorker(interval, m.RunBatchSync, "batch sync completed") + quotaSyncInterval := time.Duration(settings.QuotaSyncIntervalSeconds) * time.Second + startWorker(quotaSyncInterval, m.RunQuotaSync, "quota sync completed") + startWorker(time.Hour, m.RunCleanup, "cleanup completed") +} + +// Stop 停止同步管理器 +func (m *SyncManager) Stop() { + m.stopped.Store(true) +} diff --git a/service/region_sync/sync_manager_test.go b/service/region_sync/sync_manager_test.go new file mode 100644 index 0000000..affb02f --- /dev/null +++ b/service/region_sync/sync_manager_test.go @@ -0,0 +1,324 @@ +package region_sync + +import ( + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/glebarez/sqlite" + "gorm.io/gorm" + + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting/system_setting" +) + +// setupTestDB 设置测试数据库 +func setupTestDB(t *testing.T) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, _ := db.DB() + sqlDB.SetMaxOpenConns(1) + + // 迁移模型 + err = db.AutoMigrate(&model.PendingSyncRecord{}) + require.NoError(t, err) + + return db +} + +func TestSyncManager_NewSyncManager(t *testing.T) { + manager := NewSyncManager() + assert.NotNil(t, manager.client) +} + +func TestSyncManager_RunBatchSync_Empty(t *testing.T) { + // 保存原始设置 + settings := system_setting.GetRegionSyncSettings() + settings.Enabled = false + + manager := NewSyncManager() + count := manager.RunBatchSync() + assert.Equal(t, 0, count) +} + +func TestSyncManager_RunBatchSync_WithRecords(t *testing.T) { + origDB := model.DB + defer func() { + model.DB = origDB + }() + + // 设置测试数据库 + db := setupTestDB(t) + model.DB = db + + // 插入一些待同步记录 + record1 := &model.PendingSyncRecord{ + Id: 1, + UserId: 100, + RemoteUserId: 100, + RequestId: "req-001", + Quota: 100, + Status: model.PendingSyncStatusPending, + CreatedAt: time.Now().Unix(), + } + record2 := &model.PendingSyncRecord{ + Id: 2, + UserId: 100, + RemoteUserId: 100, + RequestId: "req-002", + Quota: 200, + Status: model.PendingSyncStatusPending, + CreatedAt: time.Now().Unix(), + } + + require.NoError(t, db.Create(record1).Error) + require.NoError(t, db.Create(record2).Error) + + manager := NewSyncManager() + count := manager.RunBatchSync() + + // 因为没有真正的服务器,返回 0 + assert.Equal(t, 0, count) +} + +func TestSyncManager_QueryMasterQuota(t *testing.T) { + origDB := model.DB + defer func() { + model.DB = origDB + }() + + manager := NewSyncManager() + quota, err := manager.QueryMasterQuota(100) + + // 在测试环境中应该返回错误(因为没有启用配置) + require.Error(t, err) + assert.Equal(t, 0, quota) +} + +func TestSyncManager_RunCleanup(t *testing.T) { + origDB := model.DB + defer func() { + model.DB = origDB + }() + + // 设置测试数据库 + db := setupTestDB(t) + model.DB = db + + // 插入一些已同步的旧记录 + oldRecord := &model.PendingSyncRecord{ + Id: 1, + UserId: 100, + RemoteUserId: 100, + RequestId: "req-001", + Quota: 100, + Status: model.PendingSyncStatusSynced, + CreatedAt: time.Now().Add(-40 * 24 * time.Hour).Unix(), + } + require.NoError(t, db.Create(oldRecord).Error) + + manager := NewSyncManager() + deletedCount := manager.RunCleanup() + assert.GreaterOrEqual(t, deletedCount, 0) +} + +func TestSyncManager_RunBatchSync_MultipleRecords(t *testing.T) { + origDB := model.DB + defer func() { + model.DB = origDB + }() + + // 设置测试数据库 + db := setupTestDB(t) + model.DB = db + + // 插入待同步记录 + for i := 1; i <= 3; i++ { + record := &model.PendingSyncRecord{ + Id: i, + UserId: 100, + RemoteUserId: 100, + RequestId: fmt.Sprintf("req-batch-%d", i), + Quota: i * 100, + Status: model.PendingSyncStatusPending, + CreatedAt: time.Now().Unix(), + } + require.NoError(t, db.Create(record).Error) + } + + manager := NewSyncManager() + // 执行批量同步 + count := manager.RunBatchSync() + + // 因为没有真正的服务器,返回 0 + assert.Equal(t, 0, count) +} + +func TestSyncManager_StartSyncWorkers(t *testing.T) { + manager := NewSyncManager() + + // 这个测试只是验证不会 panic + // 实际的协程测试需要更复杂的设置 + manager.StartSyncWorkers() + + // 等待一小段时间让协程启动 + time.Sleep(100 * time.Millisecond) + + // 停止管理器 + manager.Stop() +} + +// --------------------------------------------------------------------------- +// 阶段 3 补充:SyncManager 端到端测试 +// --------------------------------------------------------------------------- + +func TestRunBatchSync_WithMockServer(t *testing.T) { + origDB := model.DB + defer func() { model.DB = origDB }() + + db := setupTestDB(t) + model.DB = db + + now := time.Now().Unix() + + settings := system_setting.GetRegionSyncSettings() + origEnabled := settings.Enabled + origBatchSize := settings.SyncBatchSize + origMaxRetry := settings.MaxRetryCount + defer func() { + settings.Enabled = origEnabled + settings.SyncBatchSize = origBatchSize + settings.MaxRetryCount = origMaxRetry + }() + + settings.Enabled = true + settings.SyncBatchSize = 100 + settings.MaxRetryCount = 3 + + t.Run("AllSuccess", func(t *testing.T) { + require.NoError(t, db.Exec("DELETE FROM pending_sync_records").Error) + require.NoError(t, db.Create(&model.PendingSyncRecord{ + UserId: 100, RemoteUserId: 200, RequestId: "req-batch-ok-1", + Quota: 100, Status: model.PendingSyncStatusPending, CreatedAt: now, + }).Error) + require.NoError(t, db.Create(&model.PendingSyncRecord{ + UserId: 101, RemoteUserId: 201, RequestId: "req-batch-ok-2", + Quota: 200, Status: model.PendingSyncStatusPending, CreatedAt: now + 1, + }).Error) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(BatchDeductResponse{ + Success: true, + Results: []DeductResult{ + {UserId: 200, Success: true, DeductedQuota: 100, RemainingQuota: 999}, + {UserId: 201, Success: true, DeductedQuota: 200, RemainingQuota: 888}, + }, + }) + })) + defer server.Close() + + manager := &SyncManager{client: NewSyncClient(server.URL, "test-key")} + count := manager.RunBatchSync() + assert.Equal(t, 2, count) + + var syncedCount int64 + db.Model(&model.PendingSyncRecord{}).Where("status = ?", model.PendingSyncStatusSynced).Count(&syncedCount) + assert.Equal(t, int64(2), syncedCount) + }) + + t.Run("PartialSuccess", func(t *testing.T) { + require.NoError(t, db.Exec("DELETE FROM pending_sync_records").Error) + require.NoError(t, db.Create(&model.PendingSyncRecord{ + UserId: 100, RemoteUserId: 200, RequestId: "req-part-1", + Quota: 100, Status: model.PendingSyncStatusPending, CreatedAt: now, + }).Error) + require.NoError(t, db.Create(&model.PendingSyncRecord{ + UserId: 101, RemoteUserId: 201, RequestId: "req-part-2", + Quota: 200, Status: model.PendingSyncStatusPending, CreatedAt: now + 1, + }).Error) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(BatchDeductResponse{ + Success: true, + Results: []DeductResult{ + {UserId: 200, Success: true, DeductedQuota: 100}, + {UserId: 201, Success: false, Error: "insufficient quota"}, + }, + }) + })) + defer server.Close() + + manager := &SyncManager{client: NewSyncClient(server.URL, "test-key")} + count := manager.RunBatchSync() + assert.Equal(t, 1, count) + + var syncedCount, pendingCount int64 + db.Model(&model.PendingSyncRecord{}).Where("status = ?", model.PendingSyncStatusSynced).Count(&syncedCount) + db.Model(&model.PendingSyncRecord{}).Where("status = ?", model.PendingSyncStatusPending).Count(&pendingCount) + assert.Equal(t, int64(1), syncedCount) + assert.Equal(t, int64(1), pendingCount) + }) + + t.Run("ServerError", func(t *testing.T) { + require.NoError(t, db.Exec("DELETE FROM pending_sync_records").Error) + require.NoError(t, db.Create(&model.PendingSyncRecord{ + UserId: 100, RemoteUserId: 200, RequestId: "req-srv-err", + Quota: 100, Status: model.PendingSyncStatusPending, CreatedAt: now, + }).Error) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(500) + json.NewEncoder(w).Encode(BatchDeductResponse{Success: false, Error: "internal error"}) + })) + defer server.Close() + + manager := &SyncManager{client: NewSyncClient(server.URL, "test-key")} + count := manager.RunBatchSync() + assert.Equal(t, 0, count) + }) +} + +func TestRunBatchSync_RetryExhaustion(t *testing.T) { + origDB := model.DB + defer func() { model.DB = origDB }() + + db := setupTestDB(t) + model.DB = db + + now := time.Now().Unix() + require.NoError(t, db.Create(&model.PendingSyncRecord{ + UserId: 100, RemoteUserId: 200, RequestId: "req-retry-exh", + Quota: 100, Status: model.PendingSyncStatusFailed, RetryCount: 2, CreatedAt: now, + }).Error) + + settings := system_setting.GetRegionSyncSettings() + origEnabled := settings.Enabled + origBatchSize := settings.SyncBatchSize + origMaxRetry := settings.MaxRetryCount + defer func() { + settings.Enabled = origEnabled + settings.SyncBatchSize = origBatchSize + settings.MaxRetryCount = origMaxRetry + }() + settings.Enabled = true + settings.SyncBatchSize = 100 + settings.MaxRetryCount = 3 + + manager := &SyncManager{client: NewSyncClient("http://127.0.0.1:0", "key")} + count := manager.RunBatchSync() + assert.Equal(t, 0, count) + + var record model.PendingSyncRecord + require.NoError(t, db.Where("request_id = ?", "req-retry-exh").First(&record).Error) + assert.Equal(t, model.PendingSyncStatusFailed, record.Status) + assert.Equal(t, 3, record.RetryCount) +} diff --git a/service/region_sync/sync_push.go b/service/region_sync/sync_push.go new file mode 100644 index 0000000..20b1732 --- /dev/null +++ b/service/region_sync/sync_push.go @@ -0,0 +1,116 @@ +package region_sync + +import ( + "fmt" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting/system_setting" + "github.com/bytedance/gopkg/util/gopool" +) + +// PushUserCreateToSlave 在 master 节点创建用户后,异步推送到 slave 节点 +func PushUserCreateToSlave(user *model.User) { + settings := system_setting.GetRegionSyncSettings() + + if !settings.Enabled || !settings.IsMaster { + return + } + if len(settings.SlaveEndpoints) == 0 || settings.SyncApiKey == "" { + return + } + + gopool.Go(func() { + req := &SyncUserRequest{ + Username: user.Username, + Email: user.Email, + PasswordHash: user.Password, + DisplayName: user.DisplayName, + Quota: user.Quota, + RemoteUserId: user.Id, + Group: user.Group, + AffCode: user.AffCode, + } + + for _, endpoint := range settings.SlaveEndpoints { + client := NewSyncClient(endpoint, settings.SyncApiKey) + resp, err := client.SyncUserCreate(req) + if err != nil { + logSyncError("user", user.Username, endpoint, "", err) + continue + } + if !resp.Success { + logSyncError("user", user.Username, endpoint, resp.Error, nil) + continue + } + logSyncSuccess("user", user.Username, user.Id, endpoint) + } + + recordSyncLog(user.Id, 0, model.SyncTypeUserCreate, "") + }) +} + +// PushQuotaUpdateToSlave 在 master 节点更新用户余额后,推送到 slave 节点 +func PushQuotaUpdateToSlave(userId int, quota int) { + settings := system_setting.GetRegionSyncSettings() + if !settings.Enabled || !settings.IsMaster { + return + } + if len(settings.SlaveEndpoints) == 0 || settings.SyncApiKey == "" { + return + } + + logger.LogDebug(nil, "[RegionSync] PushQuotaUpdateToSlave: userId=%d, quota=%d, endpoints=%v", userId, quota, settings.SlaveEndpoints) + + gopool.Go(func() { + for _, endpoint := range settings.SlaveEndpoints { + client := NewSyncClient(endpoint, settings.SyncApiKey) + resp, err := client.UpdateQuota(userId, quota) + if err != nil { + logSyncError("quota update", fmt.Sprintf("user %d", userId), endpoint, "", err) + continue + } + if !resp.Success { + logSyncError("quota update", fmt.Sprintf("user %d", userId), endpoint, resp.Error, nil) + continue + } + logSyncSuccess("quota update", fmt.Sprintf("user %d", userId), 0, endpoint) + } + }) +} + +// logSyncError 记录同步错误日志 +func logSyncError(operation, target, endpoint string, errMsg string, err error) { + if err != nil { + common.SysError(fmt.Sprintf("[RegionSync] Failed to push %s %s to slave %s: %v", operation, target, endpoint, err)) + } else { + common.SysError(fmt.Sprintf("[RegionSync] Slave %s rejected %s %s: %s", endpoint, operation, target, errMsg)) + } +} + +// logSyncSuccess 记录同步成功日志 +func logSyncSuccess(operation, target string, id int, endpoint string) { + if id > 0 { + common.SysLog(fmt.Sprintf("[RegionSync] Successfully pushed %s %s (id=%d) to slave %s", operation, target, id, endpoint)) + } else { + common.SysLog(fmt.Sprintf("[RegionSync] Successfully pushed %s %s to slave %s", operation, target, endpoint)) + } +} + +// recordSyncLog 创建同步日志记录 +func recordSyncLog(userId, remoteUserId int, syncType string, requestId string) { + syncLog := &model.QuotaSyncLog{ + UserId: userId, + RemoteUserId: remoteUserId, + SyncType: syncType, + Direction: model.SyncDirectionCnToOv, + Status: model.SyncStatusSuccess, + RequestId: requestId, + CreatedAt: time.Now().Unix(), + } + if err := model.CreateSyncLog(syncLog); err != nil { + common.SysError(fmt.Sprintf("[RegionSync] Failed to create sync log: %v", err)) + } +} diff --git a/service/region_sync/sync_push_test.go b/service/region_sync/sync_push_test.go new file mode 100644 index 0000000..0cc0aa2 --- /dev/null +++ b/service/region_sync/sync_push_test.go @@ -0,0 +1,171 @@ +package region_sync + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting/system_setting" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupPushDB(t *testing.T) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, _ := db.DB() + sqlDB.SetMaxOpenConns(1) + + origDB := model.DB + origLogDB := model.LOG_DB + model.DB = db + model.LOG_DB = db + + require.NoError(t, db.AutoMigrate(&model.QuotaSyncLog{})) + + t.Cleanup(func() { + model.DB = origDB + model.LOG_DB = origLogDB + sqlDB.Close() + }) + return db +} + +func setupPushSettings(t *testing.T, fn func(s *system_setting.RegionSyncSettings)) { + t.Helper() + s := system_setting.GetRegionSyncSettings() + orig := *s + fn(s) + t.Cleanup(func() { *s = orig }) +} + +func TestPushUserCreateToSlave_Disabled(t *testing.T) { + setupPushSettings(t, func(s *system_setting.RegionSyncSettings) { + s.Enabled = false + }) + + // 不应该 panic 或发请求 + PushUserCreateToSlave(&model.User{Id: 1, Username: "test"}) +} + +func TestPushUserCreateToSlave_NotMaster(t *testing.T) { + setupPushSettings(t, func(s *system_setting.RegionSyncSettings) { + s.Enabled = true + s.IsMaster = false + }) + + PushUserCreateToSlave(&model.User{Id: 1, Username: "test"}) +} + +func TestPushUserCreateToSlave_Success(t *testing.T) { + setupPushDB(t) + + var receivedBody []byte + var receivedKey string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedKey = r.Header.Get("X-Sync-API-Key") + receivedBody, _ = io.ReadAll(r.Body) + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(SyncUserResponse{Success: true, Message: "ok"}) + })) + defer server.Close() + + setupPushSettings(t, func(s *system_setting.RegionSyncSettings) { + s.Enabled = true + s.IsMaster = true + s.SlaveEndpoints = []string{server.URL} + s.SyncApiKey = "test-key-123" + }) + + PushUserCreateToSlave(&model.User{ + Id: 100, Username: "syncuser", Email: "a@b.com", + Password: "hashed", DisplayName: "Sync User", Quota: 500000, Group: "default", + }) + + // 等待异步 goroutine 完成 + time.Sleep(200 * time.Millisecond) + + assert.Equal(t, "test-key-123", receivedKey) + var req SyncUserRequest + require.NoError(t, json.Unmarshal(receivedBody, &req)) + assert.Equal(t, "syncuser", req.Username) + assert.Equal(t, "hashed", req.PasswordHash) + assert.Equal(t, 100, req.RemoteUserId) + assert.Equal(t, 500000, req.Quota) +} + +func TestPushUserCreateToSlave_PartialFailure(t *testing.T) { + setupPushDB(t) + + var successCount int32 + server1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt32(&successCount, 1) + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(SyncUserResponse{Success: true}) + })) + defer server1.Close() + + server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(SyncUserResponse{Success: false, Error: "rejected"}) + })) + defer server2.Close() + + setupPushSettings(t, func(s *system_setting.RegionSyncSettings) { + s.Enabled = true + s.IsMaster = true + s.SlaveEndpoints = []string{server1.URL, server2.URL} + s.SyncApiKey = "test-key" + }) + + PushUserCreateToSlave(&model.User{Id: 1, Username: "test"}) + time.Sleep(200 * time.Millisecond) + + assert.Equal(t, int32(1), atomic.LoadInt32(&successCount)) +} + +func TestPushQuotaUpdateToSlave_Success(t *testing.T) { + var receivedBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedBody, _ = io.ReadAll(r.Body) + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(UpdateQuotaResponse{Success: true}) + })) + defer server.Close() + + setupPushSettings(t, func(s *system_setting.RegionSyncSettings) { + s.Enabled = true + s.IsMaster = true + s.SlaveEndpoints = []string{server.URL} + s.SyncApiKey = "key" + }) + + PushQuotaUpdateToSlave(100, 500000) + time.Sleep(200 * time.Millisecond) + + var req UpdateQuotaRequest + require.NoError(t, json.Unmarshal(receivedBody, &req)) + assert.Equal(t, 100, req.RemoteUserId) + assert.Equal(t, 500000, req.Quota) +} + +func TestPushUserCreateToSlave_NoEndpoints(t *testing.T) { + setupPushDB(t) + setupPushSettings(t, func(s *system_setting.RegionSyncSettings) { + s.Enabled = true + s.IsMaster = true + s.SlaveEndpoints = nil + s.SyncApiKey = "key" + }) + + // 不应该 panic + PushUserCreateToSlave(&model.User{Id: 1, Username: "test"}) +} diff --git a/service/region_sync/sync_types.go b/service/region_sync/sync_types.go new file mode 100644 index 0000000..ba094b3 --- /dev/null +++ b/service/region_sync/sync_types.go @@ -0,0 +1,73 @@ +package region_sync + +// SyncUserRequest 创建同步用户的请求 +type SyncUserRequest struct { + Username string `json:"username"` + Email string `json:"email"` + PasswordHash string `json:"password_hash"` + DisplayName string `json:"display_name"` + Quota int `json:"quota"` + RemoteUserId int `json:"remote_user_id"` + Group string `json:"group"` + AffCode string `json:"aff_code"` +} + +// SyncUserResponse 创建同步用户的响应 +type SyncUserResponse struct { + Success bool `json:"success"` + Message string `json:"message,omitempty"` + Error string `json:"error,omitempty"` +} + +// QueryQuotaRequest 查询余额请求 +type QueryQuotaRequest struct { + UserId int `json:"user_id"` +} + +// QueryQuotaResponse 查询余额响应 +type QueryQuotaResponse struct { + Success bool `json:"success"` + Quota int `json:"quota"` + Error string `json:"error,omitempty"` +} + +// BatchDeductRecord 批量扣费记录 +type BatchDeductRecord struct { + UserId int `json:"user_id"` + RequestId string `json:"request_id"` + Quota int `json:"quota"` +} + +// BatchDeductRequest 批量扣费请求 +type BatchDeductRequest struct { + Records []BatchDeductRecord `json:"records"` +} + +// DeductResult 单条扣费结果 +type DeductResult struct { + UserId int `json:"user_id"` + Success bool `json:"success"` + DeductedQuota int `json:"deducted_quota,omitempty"` + RemainingQuota int `json:"remaining_quota,omitempty"` + Error string `json:"error,omitempty"` + Message string `json:"message,omitempty"` +} + +// BatchDeductResponse 批量扣费响应 +type BatchDeductResponse struct { + Success bool `json:"success"` + Results []DeductResult `json:"results"` + Error string `json:"error,omitempty"` +} + +// UpdateQuotaRequest 更新同步余额请求 +type UpdateQuotaRequest struct { + RemoteUserId int `json:"remote_user_id"` + Quota int `json:"quota"` +} + +// UpdateQuotaResponse 更新同步余额响应 +type UpdateQuotaResponse struct { + Success bool `json:"success"` + Error string `json:"error,omitempty"` +} diff --git a/service/region_sync/sync_types_test.go b/service/region_sync/sync_types_test.go new file mode 100644 index 0000000..e8192f4 --- /dev/null +++ b/service/region_sync/sync_types_test.go @@ -0,0 +1,37 @@ +package region_sync + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestSyncTypes_Structures(t *testing.T) { + // 验证结构体是否正确定义 + req := SyncUserRequest{} + assert.NotNil(t, req.Username) + assert.NotNil(t, req.Email) + assert.NotNil(t, req.PasswordHash) + assert.NotNil(t, req.Quota) + assert.NotNil(t, req.RemoteUserId) + assert.NotNil(t, req.Group) + + resp := SyncUserResponse{} + assert.NotNil(t, resp.Success) + assert.NotNil(t, resp.Message) + + queryReq := QueryQuotaRequest{} + assert.NotNil(t, queryReq.UserId) + + queryResp := QueryQuotaResponse{} + assert.NotNil(t, queryResp.Success) + + batchRecord := BatchDeductRecord{} + assert.NotNil(t, batchRecord.UserId) + assert.NotNil(t, batchRecord.RequestId) + assert.NotNil(t, batchRecord.Quota) + + deductResult := DeductResult{} + assert.NotNil(t, deductResult.UserId) + assert.NotNil(t, deductResult.Success) +} \ No newline at end of file diff --git a/service/synced_user_funding.go b/service/synced_user_funding.go new file mode 100644 index 0000000..8ea3a54 --- /dev/null +++ b/service/synced_user_funding.go @@ -0,0 +1,137 @@ +package service + +import ( + "fmt" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting/system_setting" +) + +const BillingSourceSyncedWallet = "synced_wallet" + +// SyncedUserFunding 同步用户的资金来源实现 +// 同步用户的余额存储在 master 节点,本地通过原子 SQL 管理 synced_quota 快照 +// 预扣时原子扣减 synced_quota 并记录到 PendingSyncRecord,结算时同步到 master +type SyncedUserFunding struct { + userId int + remoteUserId int + requestId string + consumed int // 实际预扣额度 + syncedQuota int // master 同步过来的余额快照(PreConsume 后为最新值) +} + +func NewSyncedUserFunding(userId, remoteUserId int, requestId string, syncedQuota int) *SyncedUserFunding { + return &SyncedUserFunding{ + userId: userId, + remoteUserId: remoteUserId, + requestId: requestId, + syncedQuota: syncedQuota, + } +} + +func (s *SyncedUserFunding) Source() string { return BillingSourceSyncedWallet } + +func (s *SyncedUserFunding) PreConsume(amount int) error { + if amount <= 0 { + common.SysLog(fmt.Sprintf("[RegionSync] PreConsume skipped: userId=%d, amount=%d (<=0)", s.userId, amount)) + return nil + } + + settings := system_setting.GetRegionSyncSettings() + common.SysLog(fmt.Sprintf("[RegionSync] PreConsume: userId=%d, remoteUserId=%d, amount=%d, syncedQuota=%d, threshold=%d, requestId=%s", + s.userId, s.remoteUserId, amount, s.syncedQuota, settings.MinBalanceThreshold, s.requestId)) + + // 原子 SQL:检查余额 + 扣减 synced_quota + newQuota, ok, err := model.AtomicDecreaseSyncedQuota(s.userId, amount, settings.MinBalanceThreshold) + if err != nil { + return fmt.Errorf("扣减同步额度失败: %w", err) + } + if !ok { + return fmt.Errorf("同步用户余额不足 (userId=%d, syncedQuota=%d, need=%d, threshold=%d)", + s.userId, newQuota, amount, settings.MinBalanceThreshold) + } + + // 创建待同步记录(预估扣费) + err = model.CreatePendingSyncRecord(s.userId, s.remoteUserId, s.requestId, amount, amount) + if err != nil { + // 回滚:退还已扣减的额度 + common.SysError(fmt.Sprintf("[RegionSync] PreConsume failed to create pending record, rolling back: userId=%d, requestId=%s, err=%v", s.userId, s.requestId, err)) + model.IncreaseSyncedQuota(s.userId, amount) + return fmt.Errorf("创建同步扣费记录失败: %w", err) + } + + s.consumed = amount + s.syncedQuota = newQuota + common.SysLog(fmt.Sprintf("[RegionSync] PreConsume success: userId=%d, remoteUserId=%d, consumed=%d, syncedQuota=%d, requestId=%s", s.userId, s.remoteUserId, s.consumed, s.syncedQuota, s.requestId)) + return nil +} + +func (s *SyncedUserFunding) Settle(delta int) error { + if delta == 0 { + return nil + } + + // 结算时更新待同步记录的实际扣费额度 + if delta > 0 { + // 需要额外扣减 + s.consumed += delta + } else { + // 需要退还部分额度 + s.consumed += delta // delta 是负数 + } + + // 更新本地同步余额快照 + s.syncedQuota -= delta + + common.SysLog(fmt.Sprintf("[RegionSync] Settle: userId=%d, delta=%d, consumed=%d, syncedQuota=%d, requestId=%s", s.userId, delta, s.consumed, s.syncedQuota, s.requestId)) + + // 尝试更新已有的 PendingSyncRecord;如果不存在(PreConsume 时 amount=0 跳过),则创建新记录 + updated := model.UpdatePendingSyncRecordQuota(s.requestId, s.consumed) + if updated == 0 { + // PreConsume 跳过时没有创建记录,此处用实际消耗量创建 + if err := model.CreatePendingSyncRecord(s.userId, s.remoteUserId, s.requestId, s.consumed, s.consumed); err != nil { + common.SysError(fmt.Sprintf("[RegionSync] Settle: failed to create pending record, requestId=%s, err=%v", s.requestId, err)) + } else { + common.SysLog(fmt.Sprintf("[RegionSync] Settle: created pending record (PreConsume skipped), userId=%d, quota=%d, requestId=%s", s.userId, s.consumed, s.requestId)) + } + } + + // 更新数据库中用户的 synced_quota 快照 + if err := model.UpdateSyncedQuota(s.userId, s.syncedQuota); err != nil { + common.SysError(fmt.Sprintf("[RegionSync] Settle: failed to update synced_quota, userId=%d, err=%v", s.userId, err)) + } + + return nil +} + +func (s *SyncedUserFunding) Refund() error { + if s.consumed <= 0 { + return nil + } + + common.SysLog(fmt.Sprintf("[RegionSync] Refund: userId=%d, consumed=%d, requestId=%s", s.userId, s.consumed, s.requestId)) + + // 取消待同步记录 + var record model.PendingSyncRecord + err := model.DB.Where("request_id = ? AND status = ?", s.requestId, model.PendingSyncStatusPending).First(&record).Error + if err != nil { + common.SysLog(fmt.Sprintf("[RegionSync] Refund: no pending record found for requestId=%s", s.requestId)) + return nil // 记录不存在,无需处理 + } + + // 标记为已同步(实际上已取消) + if err := model.MarkRecordSynced(record.Id); err != nil { + common.SysError(fmt.Sprintf("[RegionSync] Refund: failed to mark record synced, recordId=%d, err=%v", record.Id, err)) + return err + } + + // 原子退还 DB synced_quota + if err := model.IncreaseSyncedQuota(s.userId, s.consumed); err != nil { + common.SysError(fmt.Sprintf("[RegionSync] Refund: failed to increase synced_quota, userId=%d, err=%v", s.userId, err)) + return err + } + + common.SysLog(fmt.Sprintf("[RegionSync] Refund success: userId=%d, recordId=%d, refunded=%d, requestId=%s", s.userId, record.Id, s.consumed, s.requestId)) + return nil +} diff --git a/service/synced_user_funding_test.go b/service/synced_user_funding_test.go new file mode 100644 index 0000000..0406ef2 --- /dev/null +++ b/service/synced_user_funding_test.go @@ -0,0 +1,212 @@ +package service + +import ( + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting/system_setting" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupSyncedFundingDB(t *testing.T) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, _ := db.DB() + sqlDB.SetMaxOpenConns(1) + + origDB := model.DB + origLogDB := model.LOG_DB + model.DB = db + model.LOG_DB = db + common.UsingSQLite = true + common.RedisEnabled = false + + require.NoError(t, db.AutoMigrate(&model.PendingSyncRecord{})) + + t.Cleanup(func() { + model.DB = origDB + model.LOG_DB = origLogDB + sqlDB.Close() + }) + return db +} + +func TestNewSyncedUserFunding(t *testing.T) { + f := NewSyncedUserFunding(100, 200, "req-001", 500000) + assert.Equal(t, 100, f.userId) + assert.Equal(t, 200, f.remoteUserId) + assert.Equal(t, "req-001", f.requestId) + assert.Equal(t, 500000, f.syncedQuota) + assert.Equal(t, 0, f.consumed) +} + +func TestSyncedUserFunding_Source(t *testing.T) { + f := NewSyncedUserFunding(100, 200, "req-001", 500000) + assert.Equal(t, BillingSourceSyncedWallet, f.Source()) +} + +func TestPreConsume_Success(t *testing.T) { + setupSyncedFundingDB(t) + + settings := system_setting.GetRegionSyncSettings() + origThreshold := settings.MinBalanceThreshold + settings.MinBalanceThreshold = 100000 + defer func() { settings.MinBalanceThreshold = origThreshold }() + + f := NewSyncedUserFunding(100, 200, "req-pre-success", 1000000) + err := f.PreConsume(50000) + require.NoError(t, err) + assert.Equal(t, 50000, f.consumed) + + // 验证 PendingSyncRecord 已创建 + var record model.PendingSyncRecord + require.NoError(t, model.DB.Where("request_id = ?", "req-pre-success").First(&record).Error) + assert.Equal(t, 100, record.UserId) + assert.Equal(t, 200, record.RemoteUserId) + assert.Equal(t, 50000, record.Quota) + assert.Equal(t, model.PendingSyncStatusPending, record.Status) +} + +func TestPreConsume_ZeroAmount(t *testing.T) { + setupSyncedFundingDB(t) + + f := NewSyncedUserFunding(100, 200, "req-pre-zero", 1000000) + err := f.PreConsume(0) + assert.NoError(t, err) + assert.Equal(t, 0, f.consumed) +} + +func TestPreConsume_NegativeAmount(t *testing.T) { + setupSyncedFundingDB(t) + + f := NewSyncedUserFunding(100, 200, "req-pre-neg", 1000000) + err := f.PreConsume(-100) + assert.NoError(t, err) + assert.Equal(t, 0, f.consumed) +} + +func TestPreConsume_InsufficientQuota(t *testing.T) { + setupSyncedFundingDB(t) + + settings := system_setting.GetRegionSyncSettings() + origThreshold := settings.MinBalanceThreshold + settings.MinBalanceThreshold = 100000 + defer func() { settings.MinBalanceThreshold = origThreshold }() + + // syncedQuota=150000, amount=50000, threshold=100000 + // 150000 < 50000 + 100000 = 150000 → 不满足(严格小于) + f := NewSyncedUserFunding(100, 200, "req-pre-insuf", 149999) + err := f.PreConsume(50000) + assert.Error(t, err) + assert.Contains(t, err.Error(), "余额不足") + assert.Equal(t, 0, f.consumed) +} + +func TestPreConsume_ExactlyAtThreshold(t *testing.T) { + setupSyncedFundingDB(t) + + settings := system_setting.GetRegionSyncSettings() + origThreshold := settings.MinBalanceThreshold + settings.MinBalanceThreshold = 100000 + defer func() { settings.MinBalanceThreshold = origThreshold }() + + // syncedQuota=150000, amount=50000, threshold=100000 + // 150000 >= 50000 + 100000 = 150000 → 刚好满足 + f := NewSyncedUserFunding(100, 200, "req-pre-exact", 150000) + err := f.PreConsume(50000) + assert.NoError(t, err) + assert.Equal(t, 50000, f.consumed) +} + +func TestSettle_PositiveDelta(t *testing.T) { + f := NewSyncedUserFunding(100, 200, "req-settle-pos", 500000) + f.consumed = 1000 + + err := f.Settle(500) + assert.NoError(t, err) + assert.Equal(t, 1500, f.consumed) + assert.Equal(t, 499500, f.syncedQuota) // 500000 - 500 +} + +func TestSettle_NegativeDelta(t *testing.T) { + f := NewSyncedUserFunding(100, 200, "req-settle-neg", 500000) + f.consumed = 1000 + + err := f.Settle(-300) + assert.NoError(t, err) + assert.Equal(t, 700, f.consumed) + assert.Equal(t, 500300, f.syncedQuota) // 500000 - (-300) = 500000 + 300 +} + +func TestSettle_ZeroDelta(t *testing.T) { + f := NewSyncedUserFunding(100, 200, "req-settle-zero", 500000) + f.consumed = 1000 + + err := f.Settle(0) + assert.NoError(t, err) + assert.Equal(t, 1000, f.consumed) // 不变 + assert.Equal(t, 500000, f.syncedQuota) // 不变 +} + +func TestRefund_Success(t *testing.T) { + db := setupSyncedFundingDB(t) + + // 先创建一条 pending 记录 + require.NoError(t, model.CreatePendingSyncRecord(100, 200, "req-refund-ok", 5000, 5000)) + + f := NewSyncedUserFunding(100, 200, "req-refund-ok", 500000) + f.consumed = 5000 + + err := f.Refund() + assert.NoError(t, err) + + // 验证记录已被标记为 synced + var record model.PendingSyncRecord + require.NoError(t, db.Where("request_id = ?", "req-refund-ok").First(&record).Error) + assert.Equal(t, model.PendingSyncStatusSynced, record.Status) + assert.NotZero(t, record.SyncedAt) +} + +func TestRefund_NoConsumed(t *testing.T) { + setupSyncedFundingDB(t) + + f := NewSyncedUserFunding(100, 200, "req-refund-nocons", 500000) + f.consumed = 0 + + err := f.Refund() + assert.NoError(t, err) +} + +func TestRefund_RecordNotFound(t *testing.T) { + setupSyncedFundingDB(t) + + // 没有 PendingSyncRecord,但 consumed > 0 + f := NewSyncedUserFunding(100, 200, "req-refund-notfound", 500000) + f.consumed = 5000 + + err := f.Refund() + assert.NoError(t, err) // 应优雅处理,不报错 +} + +func TestRefund_RecordAlreadySynced(t *testing.T) { + db := setupSyncedFundingDB(t) + + // 创建已 synced 的记录 + require.NoError(t, db.Create(&model.PendingSyncRecord{ + UserId: 100, RemoteUserId: 200, RequestId: "req-refund-synced", + Quota: 5000, Status: model.PendingSyncStatusSynced, + SyncedAt: time.Now().Unix(), CreatedAt: time.Now().Unix(), + }).Error) + + f := NewSyncedUserFunding(100, 200, "req-refund-synced", 500000) + f.consumed = 5000 + + err := f.Refund() + assert.NoError(t, err) // status=pending 不匹配,优雅处理 +} diff --git a/setting/config/config.go b/setting/config/config.go index 8b3d051..3555196 100644 --- a/setting/config/config.go +++ b/setting/config/config.go @@ -256,6 +256,20 @@ func updateConfigFromMap(config interface{}, configMap map[string]string) error // 复杂类型使用JSON反序列化 err := json.Unmarshal([]byte(strValue), field.Addr().Interface()) if err != nil { + // 对于字符串切片,尝试兼容换行分隔的格式 + if field.Kind() == reflect.Slice && field.Type().Elem().Kind() == reflect.String { + lines := strings.Split(strValue, "\n") + var result []string + for _, line := range lines { + line = strings.TrimSpace(line) + if line != "" { + result = append(result, line) + } + } + if len(result) > 0 { + field.Set(reflect.ValueOf(result)) + } + } continue } } diff --git a/setting/system_setting/region_sync.go b/setting/system_setting/region_sync.go new file mode 100644 index 0000000..2e5fd51 --- /dev/null +++ b/setting/system_setting/region_sync.go @@ -0,0 +1,38 @@ +package system_setting + +import "github.com/QuantumNous/new-api/setting/config" + +// RegionSyncSettings 跨地区同步配置 +type RegionSyncSettings struct { + Enabled bool `json:"enabled"` + RegionId string `json:"region_id"` + IsMaster bool `json:"is_master"` + MasterEndpoint string `json:"master_endpoint"` + SlaveEndpoints []string `json:"slave_endpoints"` + SyncApiKey string `json:"sync_api_key"` + + MinBalanceThreshold int `json:"min_balance_threshold"` + SyncIntervalSeconds int `json:"sync_interval_seconds"` + MaxRetryCount int `json:"max_retry_count"` + SyncBatchSize int `json:"sync_batch_size"` + + QuotaSyncIntervalSeconds int `json:"quota_sync_interval_seconds"` // 余额同步间隔(秒) + DisableCachedConsume bool `json:"disable_cached_consume"` +} + +var defaultRegionSyncSettings = RegionSyncSettings{ + Enabled: false, + MinBalanceThreshold: 100000, + SyncIntervalSeconds: 60, + MaxRetryCount: 3, + SyncBatchSize: 100, + QuotaSyncIntervalSeconds: 300, +} + +func init() { + config.GlobalConfig.Register("region_sync", &defaultRegionSyncSettings) +} + +func GetRegionSyncSettings() *RegionSyncSettings { + return &defaultRegionSyncSettings +} diff --git a/setting/system_setting/region_sync_test.go b/setting/system_setting/region_sync_test.go new file mode 100644 index 0000000..252cd40 --- /dev/null +++ b/setting/system_setting/region_sync_test.go @@ -0,0 +1,27 @@ +package system_setting + +import ( + "testing" +) + +func TestGetRegionSyncSettings_Defaults(t *testing.T) { + settings := GetRegionSyncSettings() + if settings == nil { + t.Fatal("GetRegionSyncSettings() returned nil") + } + if settings.Enabled { + t.Error("default Enabled should be false") + } + if settings.MinBalanceThreshold != 100000 { + t.Errorf("MinBalanceThreshold = %d, want 100000", settings.MinBalanceThreshold) + } + if settings.SyncIntervalSeconds != 60 { + t.Errorf("SyncIntervalSeconds = %d, want 60", settings.SyncIntervalSeconds) + } + if settings.MaxRetryCount != 3 { + t.Errorf("MaxRetryCount = %d, want 3", settings.MaxRetryCount) + } + if settings.SyncBatchSize != 100 { + t.Errorf("SyncBatchSize = %d, want 100", settings.SyncBatchSize) + } +} diff --git a/test_region_sync.sh b/test_region_sync.sh new file mode 100644 index 0000000..511c711 --- /dev/null +++ b/test_region_sync.sh @@ -0,0 +1,153 @@ +#!/bin/bash + +# Master-Slave 余额同步测试脚本 +# 用法: ./test_region_sync.sh [USER_ID] +# 示例: ./test_region_sync.sh 1 + +MASTER="http://172.22.57.85:23000" +SLAVE="http://172.22.57.85:13000" +API_KEY="123456789" +USER_ID=${1:-1} # 默认用户ID为1,可通过参数传入 + +echo "==========================================" +echo " Master-Slave 余额同步测试" +echo "==========================================" +echo "Master: $MASTER" +echo "Slave: $SLAVE" +echo "API Key: $API_KEY" +echo "测试用户ID: $USER_ID" +echo "==========================================" +echo "" + +# 测试7:完整消费流程测试 +echo "=== 测试7: 完整消费流程测试 ===" +echo "" + +echo "[步骤1] 查询Master当前余额..." +RESULT=$(curl -s -X POST "$MASTER/api/internal/sync/quota/query" \ + -H "Content-Type: application/json" \ + -H "X-Sync-API-Key: $API_KEY" \ + -H "X-Sync-Node: $SLAVE" \ + -d "{\"user_id\": $USER_ID}") +echo "结果: $RESULT" +BEFORE_QUOTA=$(echo $RESULT | grep -o '"quota":[0-9]*' | grep -o '[0-9]*') +echo "当前余额: $BEFORE_QUOTA" +echo "" + +echo "[步骤2] 执行扣费500..." +RESULT=$(curl -s -X POST "$MASTER/api/internal/sync/quota/batch-deduct" \ + -H "Content-Type: application/json" \ + -H "X-Sync-API-Key: $API_KEY" \ + -H "X-Sync-Node: $SLAVE" \ + -d "{\"records\": [{\"user_id\": $USER_ID, \"request_id\": \"req-test-$(date +%s)\", \"quota\": 500}]}") +echo "结果: $RESULT" +echo "" + +echo "[步骤3] 再次查询Master余额验证扣费成功..." +RESULT=$(curl -s -X POST "$MASTER/api/internal/sync/quota/query" \ + -H "Content-Type: application/json" \ + -H "X-Sync-API-Key: $API_KEY" \ + -H "X-Sync-Node: $SLAVE" \ + -d "{\"user_id\": $USER_ID}") +echo "结果: $RESULT" +AFTER_QUOTA=$(echo $RESULT | grep -o '"quota":[0-9]*' | grep -o '[0-9]*') +echo "扣费后余额: $AFTER_QUOTA" + +if [ "$BEFORE_QUOTA" != "" ] && [ "$AFTER_QUOTA" != "" ]; then + DIFF=$((BEFORE_QUOTA - AFTER_QUOTA)) + if [ "$DIFF" -eq 500 ]; then + echo "✅ 测试7通过: 余额正确扣除了500" + else + echo "❌ 测试7失败: 余额变化为$DIFF,期望为500" + fi +fi +echo "" + +# 测试8:并发扣费测试 +echo "=== 测试8: 并发扣费测试 ===" +echo "" +echo "[步骤1] 查询当前余额..." +RESULT=$(curl -s -X POST "$MASTER/api/internal/sync/quota/query" \ + -H "Content-Type: application/json" \ + -H "X-Sync-API-Key: $API_KEY" \ + -H "X-Sync-Node: $SLAVE" \ + -d "{\"user_id\": $USER_ID}") +BEFORE_QUOTA=$(echo $RESULT | grep -o '"quota":[0-9]*' | grep -o '[0-9]*') +echo "当前余额: $BEFORE_QUOTA" +echo "" + +echo "[步骤2] 并发发送10个扣费请求(每个100)..." +for i in {1..10}; do + curl -s -X POST "$MASTER/api/internal/sync/quota/batch-deduct" \ + -H "Content-Type: application/json" \ + -H "X-Sync-API-Key: $API_KEY" \ + -H "X-Sync-Node: $SLAVE" \ + -d "{\"records\": [{\"user_id\": $USER_ID, \"request_id\": \"req-concurrent-$i-$(date +%s%N)\", \"quota\": 100}]}" > /dev/null & +done +wait +echo "并发请求发送完成" +echo "" + +echo "[步骤3] 查询最终余额..." +sleep 1 +RESULT=$(curl -s -X POST "$MASTER/api/internal/sync/quota/query" \ + -H "Content-Type: application/json" \ + -H "X-Sync-API-Key: $API_KEY" \ + -H "X-Sync-Node: $SLAVE" \ + -d "{\"user_id\": $USER_ID}") +AFTER_QUOTA=$(echo $RESULT | grep -o '"quota":[0-9]*' | grep -o '[0-9]*') +echo "最终余额: $AFTER_QUOTA" + +if [ "$AFTER_QUOTA" != "" ] && [ "$AFTER_QUOTA" -ge 0 ]; then + echo "✅ 测试8通过: 余额不为负数" +else + echo "❌ 测试8失败: 余额为负数或查询失败" +fi +echo "" + +# 测试10:幂等性测试 +echo "=== 测试10: 重复请求幂等性测试 ===" +echo "" +REQUEST_ID="req-idempotent-$(date +%s)" + +echo "[步骤1] 发送扣费请求 request_id=$REQUEST_ID..." +RESULT=$(curl -s -X POST "$MASTER/api/internal/sync/quota/batch-deduct" \ + -H "Content-Type: application/json" \ + -H "X-Sync-API-Key: $API_KEY" \ + -H "X-Sync-Node: $SLAVE" \ + -d "{\"records\": [{\"user_id\": $USER_ID, \"request_id\": \"$REQUEST_ID\", \"quota\": 100}]}") +echo "结果: $RESULT" +AFTER_FIRST=$(curl -s -X POST "$MASTER/api/internal/sync/quota/query" \ + -H "Content-Type: application/json" \ + -H "X-Sync-API-Key: $API_KEY" \ + -H "X-Sync-Node: $SLAVE" \ + -d "{\"user_id\": $USER_ID}" | grep -o '"quota":[0-9]*' | grep -o '[0-9]*') +echo "第一次扣费后余额: $AFTER_FIRST" +echo "" + +echo "[步骤2] 再次发送相同request_id的请求..." +RESULT=$(curl -s -X POST "$MASTER/api/internal/sync/quota/batch-deduct" \ + -H "Content-Type: application/json" \ + -H "X-Sync-API-Key: $API_KEY" \ + -H "X-Sync-Node: $SLAVE" \ + -d "{\"records\": [{\"user_id\": $USER_ID, \"request_id\": \"$REQUEST_ID\", \"quota\": 100}]}") +echo "结果: $RESULT" +AFTER_SECOND=$(curl -s -X POST "$MASTER/api/internal/sync/quota/query" \ + -H "Content-Type: application/json" \ + -H "X-Sync-API-Key: $API_KEY" \ + -H "X-Sync-Node: $SLAVE" \ + -d "{\"user_id\": $USER_ID}" | grep -o '"quota":[0-9]*' | grep -o '[0-9]*') +echo "第二次扣费后余额: $AFTER_SECOND" +echo "" + +if [ "$AFTER_FIRST" == "$AFTER_SECOND" ]; then + echo "✅ 测试10通过: 重复请求没有重复扣费" +else + echo "⚠️ 测试10警告: 余额发生了变化,当前实现可能不支持幂等性" + echo " (第一次: $AFTER_FIRST, 第二次: $AFTER_SECOND)" +fi +echo "" + +echo "==========================================" +echo " 测试完成" +echo "==========================================" diff --git a/test_region_sync_concurrent.sh b/test_region_sync_concurrent.sh new file mode 100644 index 0000000..ef63c85 --- /dev/null +++ b/test_region_sync_concurrent.sh @@ -0,0 +1,156 @@ +#!/bin/bash + +# Master + Slave 并发扣费测试脚本 +# 测试场景:Master本地扣费 和 Slave同步扣费 同时对同一用户扣费 +# 用法: ./test_region_sync_concurrent.sh [USER_ID] + +MASTER="http://172.22.57.85:23000" +SLAVE="http://172.22.57.85:13000" +API_KEY="123456789" +USER_ID=${1:-1} + +echo "==========================================" +echo " Master/Slave 并发扣费冲突测试" +echo "==========================================" +echo "Master: $MASTER" +echo "Slave: $SLAVE" +echo "测试用户ID: $USER_ID" +echo "==========================================" +echo "" + +# 先查询当前余额 +echo "[1] 查询初始余额..." +INITIAL=$(curl -s -X POST "$MASTER/api/internal/sync/quota/query" \ + -H "Content-Type: application/json" \ + -H "X-Sync-API-Key: $API_KEY" \ + -H "X-Sync-Node: $SLAVE" \ + -d "{\"user_id\": $USER_ID}") +INITIAL_QUOTA=$(echo $INITIAL | grep -o '"quota":[0-9]*' | grep -o '[0-9]*') +echo "初始余额: $INITIAL_QUOTA" +echo "" + +# 定义扣费金额 +MASTER_DEDUCT=1000 # Master端扣费 1000 +SLAVE_DEDUCT=1000 # Slave端扣费 1000 +EXPECTED_FINAL=$((INITIAL_QUOTA - MASTER_DEDUCT - SLAVE_DEDUCT)) + +echo "[2] 并发执行 Master本地扣费($MASTER_DEDUCT) + Slave同步扣费($SLAVE_DEDUCT)..." +echo " 期望最终余额: $INITIAL_QUOTA - $MASTER_DEDUCT - $SLAVE_DEDUCT = $EXPECTED_FINAL" +echo "" + +# 并发发送两个请求 +# 请求1: 模拟 Slave 同步扣费(调用 Master 的 batch-deduct 接口) +curl -s -X POST "$MASTER/api/internal/sync/quota/batch-deduct" \ + -H "Content-Type: application/json" \ + -H "X-Sync-API-Key: $API_KEY" \ + -H "X-Sync-Node: $SLAVE" \ + -d "{\"records\": [{\"user_id\": $USER_ID, \"request_id\": \"req-slave-concurrent-$(date +%s)\", \"quota\": $SLAVE_DEDUCT}]}" > /tmp/slave_result.json & + +# 请求2: 模拟 Master 本地扣费(通过正常API调用,这里用 DecreaseUserQuota 对应的管理接口) +# 直接用 batch-deduct 的不同 request_id 模拟第二个并发扣费 +curl -s -X POST "$MASTER/api/internal/sync/quota/batch-deduct" \ + -H "Content-Type: application/json" \ + -H "X-Sync-API-Key: $API_KEY" \ + -H "X-Sync-Node: $SLAVE" \ + -d "{\"records\": [{\"user_id\": $USER_ID, \"request_id\": \"req-master-concurrent-$(date +%s)\", \"quota\": $MASTER_DEDUCT}]}" > /tmp/master_result.json & + +wait +echo "Slave扣费结果: $(cat /tmp/slave_result.json)" +echo "Master扣费结果: $(cat /tmp/master_result.json)" +echo "" + +# 查询最终余额 +echo "[3] 查询最终余额..." +FINAL=$(curl -s -X POST "$MASTER/api/internal/sync/quota/query" \ + -H "Content-Type: application/json" \ + -H "X-Sync-API-Key: $API_KEY" \ + -H "X-Sync-Node: $SLAVE" \ + -d "{\"user_id\": $USER_ID}") +FINAL_QUOTA=$(echo $FINAL | grep -o '"quota":[0-9]*' | grep -o '[0-9]*') +echo "最终余额: $FINAL_QUOTA" +echo "" + +# 分析结果 +echo "==========================================" +echo " 结果分析" +echo "==========================================" +echo "初始余额: $INITIAL_QUOTA" +echo "Master扣费: -$MASTER_DEDUCT" +echo "Slave扣费: -$SLAVE_DEDUCT" +echo "期望最终余额: $EXPECTED_FINAL" +echo "实际最终余额: $FINAL_QUOTA" +echo "" + +ACTUAL_DEDUCTED=$((INITIAL_QUOTA - FINAL_QUOTA)) +echo "实际总共扣除: $ACTUAL_DEDUCTED" +echo "期望总共扣除: $((MASTER_DEDUCT + SLAVE_DEDUCT))" +echo "" + +if [ "$FINAL_QUOTA" -eq "$EXPECTED_FINAL" ]; then + echo "✅ 测试通过: 余额变化正确,并发扣费无丢失" +elif [ "$ACTUAL_DEDUCTED" -lt "$((MASTER_DEDUCT + SLAVE_DEDUCT))" ]; then + LOST=$((MASTER_DEDUCT + SLAVE_DEDUCT - ACTUAL_DEDUCTED)) + echo "❌ 测试失败: 扣费丢失 $LOST,并发扣费存在覆盖问题" + echo " 可能原因: BatchDeductQuota 使用 FOR UPDATE + 设定新值," + echo " 而 DecreaseUserQuota 使用 gorm.Expr 原子减,两者混用会互相覆盖" +else + echo "❌ 测试失败: 余额异常(扣多了或变成负数)" +fi +echo "" + +# 高并发版本: 多次并发测试 +echo "==========================================" +echo " 高并发压力测试 (10轮)" +echo "==========================================" +echo "" + +echo "[4] 查询当前余额..." +BEFORE=$(curl -s -X POST "$MASTER/api/internal/sync/quota/query" \ + -H "Content-Type: application/json" \ + -H "X-Sync-API-Key: $API_KEY" \ + -H "X-Sync-Node: $SLAVE" \ + -d "{\"user_id\": $USER_ID}") +BEFORE_QUOTA=$(echo $BEFORE | grep -o '"quota":[0-9]*' | grep -o '[0-9]*') +echo "当前余额: $BEFORE_QUOTA" +echo "" + +echo "[5] 并发发送20个扣费请求(每个100)..." +for i in {1..20}; do + curl -s -X POST "$MASTER/api/internal/sync/quota/batch-deduct" \ + -H "Content-Type: application/json" \ + -H "X-Sync-API-Key: $API_KEY" \ + -H "X-Sync-Node: $SLAVE" \ + -d "{\"records\": [{\"user_id\": $USER_ID, \"request_id\": \"req-stress-$i-$(date +%s%N)\", \"quota\": 100}]}" > /dev/null & +done +wait +echo "全部请求发送完成" +echo "" + +echo "[6] 查询最终余额..." +AFTER=$(curl -s -X POST "$MASTER/api/internal/sync/quota/query" \ + -H "Content-Type: application/json" \ + -H "X-Sync-API-Key: $API_KEY" \ + -H "X-Sync-Node: $SLAVE" \ + -d "{\"user_id\": $USER_ID}") +AFTER_QUOTA=$(echo $AFTER | grep -o '"quota":[0-9]*' | grep -o '[0-9]*') +echo "最终余额: $AFTER_QUOTA" + +ACTUAL_TOTAL_DEDUCTED=$((BEFORE_QUOTA - AFTER_QUOTA)) +EXPECTED_TOTAL_DEDUCTED=$((20 * 100)) + +echo "" +echo "实际总共扣除: $ACTUAL_TOTAL_DEDUCTED" +echo "期望总共扣除: $EXPECTED_TOTAL_DEDUCTED" + +if [ "$ACTUAL_TOTAL_DEDUCTED" -eq "$EXPECTED_TOTAL_DEDUCTED" ]; then + echo "✅ 高并发测试通过: 所有扣费都正确" +elif [ "$ACTUAL_TOTAL_DEDUCTED" -lt "$EXPECTED_TOTAL_DEDUCTED" ]; then + echo "❌ 高并发测试失败: 少扣了 $((EXPECTED_TOTAL_DEDUCTED - ACTUAL_TOTAL_DEDUCTED))" +else + echo "❌ 高并发测试失败: 多扣了 $((ACTUAL_TOTAL_DEDUCTED - EXPECTED_TOTAL_DEDUCTED))" +fi + +echo "" +echo "==========================================" +echo " 测试完成" +echo "==========================================" diff --git a/test_region_sync_e2e.sh b/test_region_sync_e2e.sh new file mode 100644 index 0000000..769eb28 --- /dev/null +++ b/test_region_sync_e2e.sh @@ -0,0 +1,533 @@ +#!/bin/bash + +# ============================================ +# 余额同步端到端测试脚本 +# 仿照真实用户使用场景测试 Master-Slave 余额同步 +# +# 用法: ./test_region_sync_e2e.sh +# ============================================ + +set -euo pipefail + +# ==================== 配置 ==================== +MASTER="http://172.22.57.85:23000" +SLAVE="http://172.22.57.85:13000" +SYNC_API_KEY="123456789" +MASTER_USER_KEY="sk-LDEPGBf5QJYgralE0cIJ47NGsaN8odTqn3fs5IGGJtOdLKSE:1" +SLAVE_USER_KEY="sk-5zVnpGLzKxN4xeiWTBh8nuIdXajhEcBPlIZN29L0gpLdAwsu:1" +USER_ID=7 +MODEL="deepseek-ai/DeepSeek-V3-0324" +SYNC_WAIT_SECONDS=70 # 等待同步的时间(秒),需大于 SyncIntervalSeconds +SYNC_POLL_INTERVAL=5 # 轮询间隔(秒) +CONCURRENT_COUNT=5 # 并发请求数量 +DEDUCT_AMOUNT=500 # 直接同步 API 测试的扣费金额 +TIMEOUT=60 # curl 超时时间(秒) + +# ==================== 颜色定义 ==================== +GREEN='\033[0;32m' +RED='\033[0;31m' +YELLOW='\033[1;33m' +BLUE='\033[0;34m' +CYAN='\033[0;36m' +BOLD='\033[1m' +NC='\033[0m' + +# ==================== 计数器 ==================== +PASS_COUNT=0 +FAIL_COUNT=0 +SKIP_COUNT=0 +TOTAL_TESTS=0 + +# ==================== 工具函数 ==================== +log_pass() { + echo -e " ${GREEN}[PASS]${NC} $1" + ((PASS_COUNT++)) || true + ((TOTAL_TESTS++)) || true +} + +log_fail() { + echo -e " ${RED}[FAIL]${NC} $1" + ((FAIL_COUNT++)) || true + ((TOTAL_TESTS++)) || true +} + +log_skip() { + echo -e " ${YELLOW}[SKIP]${NC} $1" + ((SKIP_COUNT++)) || true + ((TOTAL_TESTS++)) || true +} + +log_info() { + echo -e " ${BLUE}[INFO]${NC} $1" +} + +log_warn() { + echo -e " ${YELLOW}[WARN]${NC} $1" +} + +print_header() { + echo "" + echo -e "${CYAN}============================================${NC}" + echo -e "${BOLD}$1${NC}" + echo -e "${CYAN}============================================${NC}" + echo "" +} + +print_section() { + echo "" + echo -e "${BOLD}--- $1 ---${NC}" +} + +# 查询 Master 余额 +query_master_balance() { + local user_id=${1:-$USER_ID} + local result=$(curl -s -m $TIMEOUT -X POST "$MASTER/api/internal/sync/quota/query" \ + -H "Content-Type: application/json" \ + -H "X-Sync-API-Key: $SYNC_API_KEY" \ + -H "X-Sync-Node: test-node" \ + -d "{\"user_id\": $user_id}") + echo "$result" | grep -o '"quota":[0-9]*' | grep -o '[0-9]*' +} + +# 等待同步完成(轮询 Master 余额,直到变化或超时) +wait_for_sync() { + local balance_before=$1 + local description=$2 + local waited=0 + + echo -e " ${BLUE}[SYNC]${NC} 等待余额同步到 Master($description)..." >&2 + + while [ $waited -lt $SYNC_WAIT_SECONDS ]; do + local current=$(query_master_balance) + if [ -n "$current" ] && [ "$current" != "$balance_before" ]; then + echo -e " ${GREEN}[SYNC]${NC} 检测到余额变化(等待了 ${waited}s):$balance_before → $current" >&2 + # 只向 stdout 输出纯数字 + echo "$current" + return 0 + fi + sleep $SYNC_POLL_INTERVAL + waited=$((waited + SYNC_POLL_INTERVAL)) + echo -e " ${BLUE}[SYNC]${NC} 已等待 ${waited}s / ${SYNC_WAIT_SECONDS}s" >&2 + done + + # 超时后再查一次 + local final=$(query_master_balance) + echo -e " ${YELLOW}[SYNC]${NC} 等待超时(${SYNC_WAIT_SECONDS}s),最终余额: $final" >&2 + echo "${final:-$balance_before}" +} + +# 发送聊天请求 +send_chat() { + local stream=${1:-false} + local max_tokens=${2:-10} + local prompt=${3:-"Say hello in one word"} + + curl -s -m $TIMEOUT --no-buffer -X POST "$SLAVE/v1/chat/completions" \ + -H "Authorization: Bearer $SLAVE_USER_KEY" \ + -H "Content-Type: application/json" \ + -d "{ + \"model\": \"$MODEL\", + \"messages\": [{\"role\": \"user\", \"content\": \"$prompt\"}], + \"max_tokens\": $max_tokens, + \"stream\": $stream + }" +} + +# ==================== 脚本开始 ==================== +print_header "余额同步端到端测试 (真实用户场景)" + +echo -e " Master: $MASTER" +echo -e " Slave: $SLAVE" +echo -e " 测试模型: $MODEL" +echo -e " 用户 ID: $USER_ID" +echo -e " 同步等待: ${SYNC_WAIT_SECONDS}s" +echo -e " 并发数量: $CONCURRENT_COUNT" + +# ==================== 阶段 0: 环境验证 ==================== +print_header "阶段 0: 环境验证" + +# T0.1: Master 连通性 +print_section "T0.1 Master 同步 API 连通性" +MASTER_RESULT=$(curl -s -m $TIMEOUT -w "\n%{http_code}" -X POST "$MASTER/api/internal/sync/quota/query" \ + -H "Content-Type: application/json" \ + -H "X-Sync-API-Key: $SYNC_API_KEY" \ + -H "X-Sync-Node: test-node" \ + -d "{\"user_id\": $USER_ID}") +MASTER_HTTP_CODE=$(echo "$MASTER_RESULT" | tail -1) +MASTER_BODY=$(echo "$MASTER_RESULT" | sed '$d') + +if [ "$MASTER_HTTP_CODE" = "200" ]; then + log_pass "Master 同步 API 可访问 (HTTP $MASTER_HTTP_CODE)" + log_info "响应: $MASTER_BODY" +else + log_fail "Master 同步 API 不可访问 (HTTP $MASTER_HTTP_CODE)" + log_info "响应: $MASTER_BODY" + echo -e " ${RED}Master 不可访问,无法继续测试${NC}" + exit 1 +fi + +# T0.2: Slave 连通性 +print_section "T0.2 Slave API 连通性" +SLAVE_RESULT=$(curl -s -m $TIMEOUT -w "\n%{http_code}" "$SLAVE/v1/models" \ + -H "Authorization: Bearer $SLAVE_USER_KEY") +SLAVE_HTTP_CODE=$(echo "$SLAVE_RESULT" | tail -1) + +if [ "$SLAVE_HTTP_CODE" = "200" ]; then + log_pass "Slave API 可访问 (HTTP $SLAVE_HTTP_CODE)" +else + log_fail "Slave API 不可访问 (HTTP $SLAVE_HTTP_CODE)" + echo -e " ${RED}Slave 不可访问,无法继续测试${NC}" + exit 1 +fi + +# ==================== 阶段 1: 基线余额 ==================== +print_header "阶段 1: 基线余额" + +print_section "T1.1 查询初始余额" +BALANCE_BEFORE=$(query_master_balance) + +if [ -n "$BALANCE_BEFORE" ]; then + log_pass "Master 初始余额: $BALANCE_BEFORE" +else + log_fail "无法查询 Master 余额" + exit 1 +fi + +if [ "$BALANCE_BEFORE" -lt 1000 ]; then + log_warn "余额较低 ($BALANCE_BEFORE),部分测试可能因余额不足而失败" +fi + +# ==================== 阶段 2: 非流式聊天调用 ==================== +print_header "阶段 2: 非流式聊天调用 (Slave → 上游 → 同步 Master)" + +print_section "T2.1 通过 Slave 发起非流式聊天请求" +log_info "模型: $MODEL, max_tokens: 10" + +CHAT_RESULT=$(send_chat false 10 "Say hello in one word") +CHAT_HTTP_OK=true + +# 检查是否有错误 +if echo "$CHAT_RESULT" | grep -q '"error"'; then + ERROR_MSG=$(echo "$CHAT_RESULT" | grep -o '"message":"[^"]*"' | head -1) + log_fail "非流式请求失败: $ERROR_MSG" + CHAT_HTTP_OK=false +elif echo "$CHAT_RESULT" | grep -q '"choices"'; then + log_pass "非流式请求成功" + # 提取使用量 + USAGE=$(echo "$CHAT_RESULT" | grep -o '"total_tokens":[0-9]*' | grep -o '[0-9]*') + if [ -n "$USAGE" ]; then + log_info "Token 使用量: $USAGE" + fi + # 提取回复内容(截断显示) + CONTENT=$(echo "$CHAT_RESULT" | grep -o '"content":"[^"]*"' | head -1 | sed 's/"content":"//;s/"$//') + if [ -n "$CONTENT" ]; then + log_info "模型回复: ${CONTENT:0:80}" + fi +else + log_warn "响应格式未知: ${CHAT_RESULT:0:100}" + CHAT_HTTP_OK=false +fi + +if [ "$CHAT_HTTP_OK" = true ]; then + print_section "T2.2 等待同步到 Master" + BALANCE_AFTER_NON_STREAM=$(wait_for_sync "$BALANCE_BEFORE" "非流式调用") + + print_section "T2.3 验证 Master 余额变化" + if [ "$BALANCE_AFTER_NON_STREAM" != "$BALANCE_BEFORE" ]; then + DELTA=$((BALANCE_BEFORE - BALANCE_AFTER_NON_STREAM)) + if [ "$DELTA" -gt 0 ]; then + log_pass "Master 余额已扣减: $BALANCE_BEFORE → $BALANCE_AFTER_NON_STREAM (扣除 $DELTA)" + else + log_fail "Master 余额变化异常: $BALANCE_BEFORE → $BALANCE_AFTER_NON_STREAM (变化 $DELTA)" + fi + else + log_fail "Master 余额未变化,同步可能失败" + fi +else + log_skip "T2.2-T2.3(非流式请求失败)" + BALANCE_AFTER_NON_STREAM=$BALANCE_BEFORE +fi + +# ==================== 阶段 3: 流式聊天调用 ==================== +print_header "阶段 3: 流式聊天调用 (SSE)" + +print_section "T3.1 通过 Slave 发起流式聊天请求" +log_info "模型: $MODEL, stream: true, max_tokens: 20" + +STREAM_RESULT=$(curl -s -m $TIMEOUT --no-buffer -X POST "$SLAVE/v1/chat/completions" \ + -H "Authorization: Bearer $SLAVE_USER_KEY" \ + -H "Content-Type: application/json" \ + -d "{ + \"model\": \"$MODEL\", + \"messages\": [{\"role\": \"user\", \"content\": \"Count from 1 to 3\"}], + \"max_tokens\": 20, + \"stream\": true + }" 2>&1) + +STREAM_OK=true + +if echo "$STREAM_RESULT" | grep -q '"error"'; then + ERROR_MSG=$(echo "$STREAM_RESULT" | grep -o '"message":"[^"]*"' | head -1) + log_fail "流式请求失败: $ERROR_MSG" + STREAM_OK=false +elif echo "$STREAM_RESULT" | grep -q 'data:'; then + log_pass "流式请求成功,收到 SSE 数据" + + # 验证 SSE 格式 + SSE_DATA_LINES=$(echo "$STREAM_RESULT" | grep -c "data:" || echo "0") + HAS_DONE=$(echo "$STREAM_RESULT" | grep -c "data: \[DONE\]" || echo "0") + + if [ "$SSE_DATA_LINES" -gt 0 ]; then + log_info "收到 $SSE_DATA_LINES 行 SSE 数据" + fi + if [ "$HAS_DONE" -gt 0 ]; then + log_pass "流式响应包含 [DONE] 标记" + else + log_warn "流式响应未包含 [DONE] 标记" + fi +else + log_warn "响应中无 SSE 数据: ${STREAM_RESULT:0:100}" + STREAM_OK=false +fi + +if [ "$STREAM_OK" = true ]; then + print_section "T3.2 等待同步到 Master" + BALANCE_BEFORE_STREAM=$BALANCE_AFTER_NON_STREAM + BALANCE_AFTER_STREAM=$(wait_for_sync "$BALANCE_BEFORE_STREAM" "流式调用") + + print_section "T3.3 验证 Master 余额变化" + if [ "$BALANCE_AFTER_STREAM" != "$BALANCE_BEFORE_STREAM" ]; then + DELTA=$((BALANCE_BEFORE_STREAM - BALANCE_AFTER_STREAM)) + if [ "$DELTA" -gt 0 ]; then + log_pass "流式调用同步成功: $BALANCE_BEFORE_STREAM → $BALANCE_AFTER_STREAM (扣除 $DELTA)" + else + log_fail "流式调用余额变化异常: $DELTA" + fi + else + log_warn "流式调用后 Master 余额未变化(可能 token 消耗为 0)" + fi +else + log_skip "T3.2-T3.3(流式请求失败)" + BALANCE_AFTER_STREAM=$BALANCE_AFTER_NON_STREAM +fi + +# ==================== 阶段 4: 并发请求 ==================== +print_header "阶段 4: 并发请求 ($CONCURRENT_COUNT 个)" + +print_section "T4.1 并发发送 $CONCURRENT_COUNT 个聊天请求" +BALANCE_BEFORE_CONCURRENT=$BALANCE_AFTER_STREAM + +# 准备不同的 prompt 以避免缓存 +PROMPTS=("What is 1+1?" "Tell me a color" "Say yes or no" "What day is it?" "Name a fruit" + "What is 2+2?" "Say goodbye" "What is water?" "Name an animal" "Say OK") + +PIDS=() +for i in $(seq 1 $CONCURRENT_COUNT); do + PROMPT="${PROMPTS[$((i - 1))]}" + send_chat false 10 "$PROMPT" > "/tmp/sync_test_concurrent_$i.json" 2>&1 & + PIDS+=($!) +done + +log_info "已发送 $CONCURRENT_COUNT 个并发请求,等待完成..." + +SUCCESS_COUNT=0 +FAIL_COUNT_CONCURRENT=0 +for pid in "${PIDS[@]}"; do + if wait "$pid"; then + ((SUCCESS_COUNT++)) || true + else + ((FAIL_COUNT_CONCURRENT++)) || true + fi +done + +if [ "$FAIL_COUNT_CONCURRENT" -eq 0 ]; then + log_pass "全部 $CONCURRENT_COUNT 个请求完成" +else + log_warn "$SUCCESS_COUNT 成功, $FAIL_COUNT_CONCURRENT 失败" +fi + +# 检查响应 +print_section "T4.2 验证并发请求响应" +VALID_RESPONSES=0 +for i in $(seq 1 $CONCURRENT_COUNT); do + if [ -f "/tmp/sync_test_concurrent_$i.json" ]; then + if grep -q '"choices"' "/tmp/sync_test_concurrent_$i.json"; then + ((VALID_RESPONSES++)) || true + fi + fi +done +log_info "有效响应: $VALID_RESPONSES / $CONCURRENT_COUNT" + +if [ "$VALID_RESPONSES" -gt 0 ]; then + print_section "T4.3 等待同步到 Master" + BALANCE_AFTER_CONCURRENT=$(wait_for_sync "$BALANCE_BEFORE_CONCURRENT" "并发调用") + + print_section "T4.4 验证 Master 余额变化" + if [ "$BALANCE_AFTER_CONCURRENT" != "$BALANCE_BEFORE_CONCURRENT" ]; then + DELTA=$((BALANCE_BEFORE_CONCURRENT - BALANCE_AFTER_CONCURRENT)) + if [ "$DELTA" -gt 0 ]; then + log_pass "并发调用同步成功: $BALANCE_BEFORE_CONCURRENT → $BALANCE_AFTER_CONCURRENT (共扣除 $DELTA)" + else + log_fail "并发调用余额变化异常: $DELTA" + fi + else + log_warn "并发调用后 Master 余额未变化" + fi +else + log_skip "T4.3-T4.4(无有效并发响应)" + BALANCE_AFTER_CONCURRENT=$BALANCE_BEFORE_CONCURRENT +fi + +# 清理临时文件 +rm -f /tmp/sync_test_concurrent_*.json + +# ==================== 阶段 5: 错误场景 ==================== +print_header "阶段 5: 错误场景测试" + +# T5.1: 错误模型名 +print_section "T5.1 错误模型名" +ERR_RESULT=$(curl -s -m $TIMEOUT -w "\n%{http_code}" -X POST "$SLAVE/v1/chat/completions" \ + -H "Authorization: Bearer $SLAVE_USER_KEY" \ + -H "Content-Type: application/json" \ + -d '{"model":"nonexistent-model-xyz-12345","messages":[{"role":"user","content":"hi"}],"max_tokens":5}') +ERR_HTTP=$(echo "$ERR_RESULT" | tail -1) +ERR_BODY=$(echo "$ERR_RESULT" | sed '$d') + +if [ "$ERR_HTTP" != "200" ]; then + log_pass "错误模型名被拒绝 (HTTP $ERR_HTTP)" +else + log_fail "错误模型名未被拒绝 (HTTP $ERR_HTTP)" +fi +log_info "响应: ${ERR_BODY:0:100}" + +# T5.2: 无认证 +print_section "T5.2 无认证请求" +NO_AUTH_RESULT=$(curl -s -m $TIMEOUT -w "\n%{http_code}" -X POST "$SLAVE/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d '{"model":"gpt-3.5-turbo","messages":[{"role":"user","content":"hi"}],"max_tokens":5}') +NO_AUTH_HTTP=$(echo "$NO_AUTH_RESULT" | tail -1) + +if [ "$NO_AUTH_HTTP" != "200" ]; then + log_pass "无认证请求被拒绝 (HTTP $NO_AUTH_HTTP)" +else + log_fail "无认证请求未被拒绝 (HTTP $NO_AUTH_HTTP)" +fi + +# T5.3: 错误 API Key +print_section "T5.3 错误 API Key" +BAD_KEY_RESULT=$(curl -s -m $TIMEOUT -w "\n%{http_code}" -X POST "$SLAVE/v1/chat/completions" \ + -H "Authorization: Bearer sk-invalid-key-12345" \ + -H "Content-Type: application/json" \ + -d '{"model":"gpt-3.5-turbo","messages":[{"role":"user","content":"hi"}],"max_tokens":5}') +BAD_KEY_HTTP=$(echo "$BAD_KEY_RESULT" | tail -1) + +if [ "$BAD_KEY_HTTP" != "200" ]; then + log_pass "错误 API Key 被拒绝 (HTTP $BAD_KEY_HTTP)" +else + log_fail "错误 API Key 未被拒绝 (HTTP $BAD_KEY_HTTP)" +fi + +# T5.4: 空消息 +print_section "T5.4 空消息列表" +EMPTY_MSG_RESULT=$(curl -s -m $TIMEOUT -w "\n%{http_code}" -X POST "$SLAVE/v1/chat/completions" \ + -H "Authorization: Bearer $SLAVE_USER_KEY" \ + -H "Content-Type: application/json" \ + -d '{"model":"gpt-3.5-turbo","messages":[],"max_tokens":5}') +EMPTY_MSG_HTTP=$(echo "$EMPTY_MSG_RESULT" | tail -1) + +if [ "$EMPTY_MSG_HTTP" != "200" ]; then + log_pass "空消息被拒绝 (HTTP $EMPTY_MSG_HTTP)" +else + log_warn "空消息未被拒绝 (HTTP $EMPTY_MSG_HTTP) — 某些实现可能允许" +fi + +# ==================== 阶段 6: 直接同步 API 测试 ==================== +print_header "阶段 6: 直接同步 API 测试" + +BALANCE_BEFORE_DIRECT=$(query_master_balance) +log_info "直接测试前 Master 余额: $BALANCE_BEFORE_DIRECT" + +# T6.1: 批量扣费 +print_section "T6.1 批量扣费 ($DEDUCT_AMOUNT)" +REQUEST_ID="req-e2e-$(date +%s%N)" + +DEDUCT_RESULT=$(curl -s -m $TIMEOUT -X POST "$MASTER/api/internal/sync/quota/batch-deduct" \ + -H "Content-Type: application/json" \ + -H "X-Sync-API-Key: $SYNC_API_KEY" \ + -H "X-Sync-Node: test-node" \ + -d "{\"records\": [{\"user_id\": $USER_ID, \"request_id\": \"$REQUEST_ID\", \"quota\": $DEDUCT_AMOUNT}]}") +log_info "响应: $DEDUCT_RESULT" + +if echo "$DEDUCT_RESULT" | grep -qi "success\|true\|200"; then + log_pass "批量扣费请求成功" +else + log_fail "批量扣费请求失败: $DEDUCT_RESULT" +fi + +# T6.2: 幂等性测试(相同 request_id) +print_section "T6.2 幂等性测试(相同 request_id)" +IDEMPOTENT_RESULT=$(curl -s -m $TIMEOUT -X POST "$MASTER/api/internal/sync/quota/batch-deduct" \ + -H "Content-Type: application/json" \ + -H "X-Sync-API-Key: $SYNC_API_KEY" \ + -H "X-Sync-Node: test-node" \ + -d "{\"records\": [{\"user_id\": $USER_ID, \"request_id\": \"$REQUEST_ID\", \"quota\": $DEDUCT_AMOUNT}]}") +log_info "响应: $IDEMPOTENT_RESULT" + +BALANCE_AFTER_FIRST=$(query_master_balance) +log_info "幂等测试后余额: $BALANCE_AFTER_FIRST" + +# 再查一次确认没变 +sleep 1 +BALANCE_AFTER_IDEMPOTENT=$(query_master_balance) + +if [ "$BALANCE_AFTER_IDEMPOTENT" = "$BALANCE_AFTER_FIRST" ]; then + log_pass "幂等性验证通过:重复请求未重复扣费" +else + log_fail "幂等性验证失败:余额发生了变化 ($BALANCE_AFTER_FIRST → $BALANCE_AFTER_IDEMPOTENT)" +fi + +# T6.3: 精确金额验证 +print_section "T6.3 精确扣费金额验证" +log_info "扣费前: $BALANCE_BEFORE_DIRECT" +log_info "扣费后: $BALANCE_AFTER_IDEMPOTENT" + +ACTUAL_DEDUCT=$((BALANCE_BEFORE_DIRECT - BALANCE_AFTER_IDEMPOTENT)) +if [ "$ACTUAL_DEDUCT" -eq $DEDUCT_AMOUNT ]; then + log_pass "精确金额验证通过:扣除了 $ACTUAL_DEDUCT (期望 $DEDUCT_AMOUNT)" +else + log_fail "精确金额验证失败:扣除了 $ACTUAL_DEDUCT (期望 $DEDUCT_AMOUNT)" +fi + +# ==================== 阶段 7: 总结报告 ==================== +print_header "测试总结" + +FINAL_BALANCE=$(query_master_balance) + +echo -e " ${BOLD}测试统计:${NC}" +echo -e " ${GREEN}PASS${NC}: $PASS_COUNT" +echo -e " ${RED}FAIL${NC}: $FAIL_COUNT" +echo -e " ${YELLOW}SKIP${NC}: $SKIP_COUNT" +echo -e " 总计: $TOTAL_TESTS" +echo "" +echo -e " ${BOLD}余额变化:${NC}" +echo -e " 初始余额: $BALANCE_BEFORE" +echo -e " 最终余额: $FINAL_BALANCE" +echo -e " 总扣除: $((BALANCE_BEFORE - FINAL_BALANCE))" +echo "" +echo -e " 测试模型: $MODEL" + +echo "" +if [ "$FAIL_COUNT" -eq 0 ]; then + echo -e " ${GREEN}${BOLD}全部测试通过!${NC}" +else + echo -e " ${RED}${BOLD}存在 $FAIL_COUNT 个失败${NC}" +fi + +echo "" +echo -e "${CYAN}============================================${NC}" + +# 返回退出码 +if [ "$FAIL_COUNT" -gt 0 ]; then + exit 1 +fi +exit 0 diff --git a/web/public/logo.png b/web/public/logo.png index aa247ab..000d76d 100644 Binary files a/web/public/logo.png and b/web/public/logo.png differ diff --git a/web/src/components/settings/OperationSetting.jsx b/web/src/components/settings/OperationSetting.jsx index 8585a3e..7bf5454 100644 --- a/web/src/components/settings/OperationSetting.jsx +++ b/web/src/components/settings/OperationSetting.jsx @@ -27,6 +27,7 @@ import SettingsLog from '../../pages/Setting/Operation/SettingsLog'; import SettingsMonitoring from '../../pages/Setting/Operation/SettingsMonitoring'; import SettingsCreditLimit from '../../pages/Setting/Operation/SettingsCreditLimit'; import SettingsCheckin from '../../pages/Setting/Operation/SettingsCheckin'; +import SettingsRegionSync from '../../pages/Setting/Operation/SettingsRegionSync'; import { API, showError, toBoolean } from '../../helpers'; const OperationSetting = () => { @@ -79,6 +80,20 @@ const OperationSetting = () => { 'checkin_setting.min_quota': 1000, 'checkin_setting.max_quota': 10000, + /* 跨区域同步设置 */ + 'region_sync.enabled': false, + 'region_sync.region_id': '', + 'region_sync.is_master': false, + 'region_sync.master_endpoint': '', + 'region_sync.slave_endpoints': '', + 'region_sync.sync_api_key': '', + 'region_sync.min_balance_threshold': 100000, + 'region_sync.sync_interval_seconds': 60, + 'region_sync.max_retry_count': 3, + 'region_sync.sync_batch_size': 100, + 'region_sync.quota_sync_interval_seconds': 300, + 'region_sync.disable_cached_consume': false, + /* 令牌设置 */ 'token_setting.max_user_tokens': 1000, }); @@ -89,16 +104,22 @@ const OperationSetting = () => { const res = await API.get('/api/option/'); const { success, message, data } = res.data; if (success) { - let newInputs = {}; + let apiInputs = {}; data.forEach((item) => { - if (typeof inputs[item.key] === 'boolean') { - newInputs[item.key] = toBoolean(item.value); - } else { - newInputs[item.key] = item.value; + // 只处理已知字段,避免添加未知字段 + if (inputs.hasOwnProperty(item.key)) { + if (typeof inputs[item.key] === 'boolean') { + apiInputs[item.key] = toBoolean(item.value); + } else if (typeof inputs[item.key] === 'number') { + apiInputs[item.key] = parseFloat(item.value) || 0; + } else { + apiInputs[item.key] = item.value; + } } }); - setInputs(newInputs); + // 合并 API 数据和初始默认值 + setInputs((prevInputs) => ({ ...prevInputs, ...apiInputs })); } else { showError(message); } @@ -154,6 +175,10 @@ const OperationSetting = () => { + {/* 跨区域同步设置 */} + + + > ); diff --git a/web/src/pages/Setting/Operation/SettingsRegionSync.jsx b/web/src/pages/Setting/Operation/SettingsRegionSync.jsx new file mode 100644 index 0000000..74a2c64 --- /dev/null +++ b/web/src/pages/Setting/Operation/SettingsRegionSync.jsx @@ -0,0 +1,425 @@ +/* +Copyright (C) 2025 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ + +import React, { useEffect, useState, useRef } from 'react'; +import { Button, Col, Form, Row, Spin, Typography, TagInput } from '@douyinfe/semi-ui'; +import { + compareObjects, + API, + showError, + showSuccess, + showWarning, + toBoolean, +} from '../../../helpers'; +import { useTranslation } from 'react-i18next'; + +// 将字符串转换为数组(用于 TagInput 显示) +function stringToArray(str) { + if (!str || typeof str !== 'string') return []; + return str + .split('\n') + .map((s) => s.trim()) + .filter((s) => s.length > 0); +} + +// 将数组转换为 JSON 字符串(用于保存) +function arrayToString(arr) { + if (!arr || !Array.isArray(arr)) return '[]'; + const filtered = arr + .map((s) => s.trim()) + .filter((s) => s.length > 0); + return JSON.stringify(filtered); +} + +// 字段类型定义,用于从 props.options 加载时正确转换类型 +const FIELD_TYPES = { + 'region_sync.enabled': 'boolean', + 'region_sync.region_id': 'string', + 'region_sync.is_master': 'boolean', + 'region_sync.master_endpoint': 'string', + 'region_sync.slave_endpoints': 'string', + 'region_sync.sync_api_key': 'string', + 'region_sync.min_balance_threshold': 'number', + 'region_sync.sync_interval_seconds': 'number', + 'region_sync.max_retry_count': 'number', + 'region_sync.sync_batch_size': 'number', + 'region_sync.quota_sync_interval_seconds': 'number', + 'region_sync.disable_cached_consume': 'boolean', +}; + +// 默认值常量,确保所有字段始终存在(即使 API 不返回某些敏感字段) +const DEFAULT_INPUTS = { + 'region_sync.enabled': false, + 'region_sync.region_id': '', + 'region_sync.is_master': false, + 'region_sync.master_endpoint': '', + 'region_sync.slave_endpoints': '', + 'region_sync.sync_api_key': '', + 'region_sync.min_balance_threshold': 100000, + 'region_sync.sync_interval_seconds': 60, + 'region_sync.max_retry_count': 3, + 'region_sync.sync_batch_size': 100, + 'region_sync.quota_sync_interval_seconds': 300, + 'region_sync.disable_cached_consume': false, +}; + +export default function SettingsRegionSync(props) { + const { t } = useTranslation(); + const [loading, setLoading] = useState(false); + // 原始数据(字符串格式,用于保存) + const [inputs, setInputs] = useState({ ...DEFAULT_INPUTS }); + // TagInput 显示用的数组格式 + const [masterEndpoints, setMasterEndpoints] = useState([]); + const [slaveEndpoints, setSlaveEndpoints] = useState([]); + + const refForm = useRef(); + const [inputsRow, setInputsRow] = useState({ ...DEFAULT_INPUTS }); + + function handleFieldChange(fieldName) { + return (value) => { + setInputs((inputs) => ({ ...inputs, [fieldName]: value })); + }; + } + + // 处理 TagInput 变化 + function handleTagInputChange(fieldName, setter) { + return (value) => { + setter(value); + const strValue = arrayToString(value); + setInputs((inputs) => ({ ...inputs, [fieldName]: strValue })); + }; + } + + function onSubmit() { + const updateArray = compareObjects(inputs, inputsRow); + if (!updateArray.length) return showWarning(t('你似乎并没有修改什么')); + const requestQueue = updateArray + .filter((item) => { + // 跳过敏感字段的空值提交,避免覆盖已保存的密钥 + if ( + item.key === 'region_sync.sync_api_key' && + (!inputs[item.key] || inputs[item.key] === '') + ) { + return false; + } + return true; + }) + .map((item) => { + let value = ''; + if (typeof inputs[item.key] === 'boolean') { + value = String(inputs[item.key]); + } else { + value = String(inputs[item.key]); + } + return API.put('/api/option/', { + key: item.key, + value, + }); + }); + if (!requestQueue.length) return showWarning(t('你似乎并没有修改什么')); + setLoading(true); + Promise.all(requestQueue) + .then((res) => { + if (requestQueue.length === 1) { + if (res.includes(undefined)) return; + } else if (requestQueue.length > 1) { + if (res.includes(undefined)) + return showError(t('部分保存失败,请重试')); + } + showSuccess(t('保存成功')); + props.refresh(); + }) + .catch(() => { + showError(t('保存失败,请重试')); + }) + .finally(() => { + setLoading(false); + }); + } + + useEffect(() => { + // 始终从默认值开始,确保所有字段(包括 API 不返回的敏感字段)都有值 + const currentInputs = { ...DEFAULT_INPUTS }; + for (let key in props.options) { + if (FIELD_TYPES[key]) { + let value = props.options[key]; + const fieldType = FIELD_TYPES[key]; + // 确保类型正确:API 返回的都是字符串 + if (fieldType === 'number' && typeof value === 'string') { + value = parseInt(value, 10); + if (isNaN(value)) value = 0; + } else if (fieldType === 'boolean' && typeof value === 'string') { + value = toBoolean(value); + } + currentInputs[key] = value; + } + } + setInputs(currentInputs); + setInputsRow(structuredClone(currentInputs)); + + // 转换为数组用于 TagInput + setMasterEndpoints(stringToArray(currentInputs['region_sync.master_endpoint'])); + setSlaveEndpoints(stringToArray(currentInputs['region_sync.slave_endpoints'])); + }, [props.options]); + + // 当条件渲染的组件挂载后,重新设置表单值 + // 原因:setValues 在 enabled/isMaster 变为 true 之前调用, + // 此时 InputNumber 等组件尚未挂载,无法接收值 + useEffect(() => { + if (refForm.current) { + refForm.current.setValues(inputs); + } + }, [inputs['region_sync.enabled'], inputs['region_sync.is_master']]); + + const enabled = inputs['region_sync.enabled']; + const isMaster = inputs['region_sync.is_master']; + + return ( + <> + + (refForm.current = formAPI)} + style={{ marginBottom: 15 }} + > + + + {t('配置国内与海外节点之间的用户数据同步,实现跨区域余额管理')} + + + {/* 启用开关 - 始终显示 */} + + + + + + + {/* 保存按钮 - 未启用时也显示 */} + {!enabled && ( + + + {t('保存跨区域同步设置')} + + + )} + + {/* 启用后才显示以下内容 */} + {enabled && ( + <> + {/* 区域 ID */} + + + + + + + {/* 角色选择 */} + + + + + + + {/* 共用设置:API 密钥 */} + + + + + + + {/* 主节点专属设置 */} + {isMaster && ( + <> + + + + + {t('从节点地址列表')} + + + + {t('例如: https://slave1.example.com')} + + + + + + + + + + + + + + + + + + + + + + + + > + )} + + {/* 从节点专属设置 */} + {!isMaster && ( + <> + + + + + {t('主节点地址')} + + + + {t('例如: https://master.example.com')} + + + + + + + + + + > + )} + + {/* 保存按钮 */} + + + {t('保存跨区域同步设置')} + + + > + )} + + + + > + ); +}