实现 Master-Slave 架构的跨区域用户和配额同步功能,支持国内主节点与 海外从节点之间的用户创建推送、余额同步、预扣费批量回传。 主要变更: - 新增 region_sync 服务层(SyncManager/SyncClient/推送逻辑) - User 模型扩展 Source/RemoteUserId/SyncedQuota/LastSyncAt 字段 - BillingSession 支持 SyncedUserFunding 资金来源 - 新增 SyncAuth 中间件(X-Sync-API-Key + 常量时间比较) - 新增 PendingSyncRecord/QuotaSyncLog 数据模型 - 前端新增区域同步配置面板 - 支持 SESSION_NAME 环境变量避免多节点 cookie 冲突 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>feat/alipay-payment
| @@ -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 | |||||
| } | |||||
| @@ -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) | |||||
| } | |||||
| }) | |||||
| } | |||||
| } | |||||
| @@ -51,6 +51,7 @@ const ( | |||||
| ContextKeyUserGroup ContextKey = "user_group" | ContextKeyUserGroup ContextKey = "user_group" | ||||
| ContextKeyUsingGroup ContextKey = "group" | ContextKeyUsingGroup ContextKey = "group" | ||||
| ContextKeyUserName ContextKey = "username" | ContextKeyUserName ContextKey = "username" | ||||
| ContextKeyUserSource ContextKey = "user_source" | |||||
| ContextKeyLocalCountTokens ContextKey = "local_count_tokens" | ContextKeyLocalCountTokens ContextKey = "local_count_tokens" | ||||
| @@ -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()) | |||||
| } | |||||
| } | |||||
| @@ -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) | |||||
| } | |||||
| @@ -159,6 +159,8 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { | |||||
| if priceData.FreeModel { | if priceData.FreeModel { | ||||
| logger.LogInfo(c, fmt.Sprintf("模型 %s 免费,跳过预扣费", relayInfo.OriginModelName)) | logger.LogInfo(c, fmt.Sprintf("模型 %s 免费,跳过预扣费", relayInfo.OriginModelName)) | ||||
| } else { | } 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) | newAPIError = service.PreConsumeBilling(c, priceData.QuotaToPreConsume, relayInfo) | ||||
| if newAPIError != nil { | if newAPIError != nil { | ||||
| return | return | ||||
| @@ -16,6 +16,7 @@ import ( | |||||
| "github.com/QuantumNous/new-api/logger" | "github.com/QuantumNous/new-api/logger" | ||||
| "github.com/QuantumNous/new-api/model" | "github.com/QuantumNous/new-api/model" | ||||
| "github.com/QuantumNous/new-api/service" | "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/setting" | ||||
| "github.com/QuantumNous/new-api/constant" | "github.com/QuantumNous/new-api/constant" | ||||
| @@ -184,12 +185,8 @@ func Register(c *gin.Context) { | |||||
| return | 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 { | if constant.GenerateDefaultToken { | ||||
| key, err := common.GenerateKey() | key, err := common.GenerateKey() | ||||
| @@ -200,7 +197,7 @@ func Register(c *gin.Context) { | |||||
| } | } | ||||
| // 生成默认令牌 | // 生成默认令牌 | ||||
| token := model.Token{ | token := model.Token{ | ||||
| UserId: insertedUser.Id, // 使用插入后的用户ID | |||||
| UserId: cleanUser.Id, // GORM Create 后已填充 ID | |||||
| Name: cleanUser.Username + "的初始令牌", | Name: cleanUser.Username + "的初始令牌", | ||||
| Key: key, | Key: key, | ||||
| CreatedTime: common.GetTimestamp(), | CreatedTime: common.GetTimestamp(), | ||||
| @@ -383,6 +380,11 @@ func GetSelf(c *gin.Context) { | |||||
| userSetting := user.GetSetting() | userSetting := user.GetSetting() | ||||
| // 构建响应数据,包含用户信息和权限 | // 构建响应数据,包含用户信息和权限 | ||||
| // Slave 节点使用 SyncedQuota 作为用户额度 | |||||
| quota := user.Quota | |||||
| if user.IsSyncedUser() { | |||||
| quota = user.SyncedQuota | |||||
| } | |||||
| responseData := map[string]interface{}{ | responseData := map[string]interface{}{ | ||||
| "id": user.Id, | "id": user.Id, | ||||
| "username": user.Username, | "username": user.Username, | ||||
| @@ -396,7 +398,7 @@ func GetSelf(c *gin.Context) { | |||||
| "wechat_id": user.WeChatId, | "wechat_id": user.WeChatId, | ||||
| "telegram_id": user.TelegramId, | "telegram_id": user.TelegramId, | ||||
| "group": user.Group, | "group": user.Group, | ||||
| "quota": user.Quota, | |||||
| "quota": quota, | |||||
| "used_quota": user.UsedQuota, | "used_quota": user.UsedQuota, | ||||
| "request_count": user.RequestCount, | "request_count": user.RequestCount, | ||||
| "aff_code": user.AffCode, | "aff_code": user.AffCode, | ||||
| @@ -22,8 +22,10 @@ import ( | |||||
| "github.com/QuantumNous/new-api/relay" | "github.com/QuantumNous/new-api/relay" | ||||
| "github.com/QuantumNous/new-api/router" | "github.com/QuantumNous/new-api/router" | ||||
| "github.com/QuantumNous/new-api/service" | "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/performance_setting" | ||||
| "github.com/QuantumNous/new-api/setting/ratio_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/bytedance/gopkg/util/gopool" | ||||
| "github.com/gin-contrib/sessions" | "github.com/gin-contrib/sessions" | ||||
| @@ -174,7 +176,12 @@ func main() { | |||||
| Secure: false, | Secure: false, | ||||
| SameSite: http.SameSiteStrictMode, | 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() | InjectUmamiAnalytics() | ||||
| InjectGoogleAnalytics() | InjectGoogleAnalytics() | ||||
| @@ -309,5 +316,19 @@ func InitResources() error { | |||||
| // Don't return error, custom OAuth is not critical | // 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 | return nil | ||||
| } | } | ||||
| @@ -22,7 +22,7 @@ DOCKER_TAG := $(DOCKER_REGISTRY):$(BUILD_TIME)-$(BRANCH_NAME) | |||||
| docker-build: | docker-build: | ||||
| @echo "Building Docker image with tag: $(DOCKER_TAG)" | @echo "Building Docker image with tag: $(DOCKER_TAG)" | ||||
| docker build -t $(DOCKER_TAG) . | docker build -t $(DOCKER_TAG) . | ||||
| docker push $(DOCKER_TAG) | |||||
| #docker push $(DOCKER_TAG) | |||||
| docker-push: docker-build | docker-push: docker-build | ||||
| @echo "Pushing Docker image: $(DOCKER_TAG)" | @echo "Pushing Docker image: $(DOCKER_TAG)" | ||||
| @@ -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 "" | |||||
| } | |||||
| @@ -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 ")) | |||||
| } | |||||
| @@ -280,6 +280,8 @@ func migrateDB() error { | |||||
| &UserOAuthBinding{}, | &UserOAuthBinding{}, | ||||
| &ChannelPricing{}, | &ChannelPricing{}, | ||||
| &PricingTag{}, | &PricingTag{}, | ||||
| &PendingSyncRecord{}, | |||||
| &QuotaSyncLog{}, | |||||
| ) | ) | ||||
| if err != nil { | if err != nil { | ||||
| return err | return err | ||||
| @@ -332,6 +334,8 @@ func migrateDBFast() error { | |||||
| {&UserOAuthBinding{}, "UserOAuthBinding"}, | {&UserOAuthBinding{}, "UserOAuthBinding"}, | ||||
| {&ChannelPricing{}, "ChannelPricing"}, | {&ChannelPricing{}, "ChannelPricing"}, | ||||
| {&PricingTag{}, "PricingTag"}, | {&PricingTag{}, "PricingTag"}, | ||||
| {&PendingSyncRecord{}, "PendingSyncRecord"}, | |||||
| {&QuotaSyncLog{}, "QuotaSyncLog"}, | |||||
| } | } | ||||
| // 动态计算migration数量,确保errChan缓冲区足够大 | // 动态计算migration数量,确保errChan缓冲区足够大 | ||||
| errChan := make(chan error, len(migrations)) | errChan := make(chan error, len(migrations)) | ||||
| @@ -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 | |||||
| } | |||||
| @@ -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 | |||||
| } | |||||
| @@ -154,6 +154,21 @@ func Redeem(key string, userId int) (quota int, err error) { | |||||
| common.SysError("redemption failed: " + err.Error()) | common.SysError("redemption failed: " + err.Error()) | ||||
| return 0, ErrRedeemFailed | 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)) | RecordLog(userId, LogTypeTopup, fmt.Sprintf("通过兑换码充值 %s,兑换码ID %d", logger.LogQuota(redemption.Quota), redemption.Id)) | ||||
| return redemption.Quota, nil | return redemption.Quota, nil | ||||
| } | } | ||||
| @@ -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) | |||||
| } | |||||
| @@ -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) | |||||
| } | |||||
| @@ -6,6 +6,7 @@ import ( | |||||
| "fmt" | "fmt" | ||||
| "strconv" | "strconv" | ||||
| "strings" | "strings" | ||||
| "time" | |||||
| "github.com/QuantumNous/new-api/common" | "github.com/QuantumNous/new-api/common" | ||||
| "github.com/QuantumNous/new-api/dto" | "github.com/QuantumNous/new-api/dto" | ||||
| @@ -17,6 +18,16 @@ import ( | |||||
| const UserNameMaxLength = 20 | 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. | // 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! | // Otherwise, the sensitive information will be saved on local storage in plain text! | ||||
| type User struct { | type User struct { | ||||
| @@ -49,6 +60,21 @@ type User struct { | |||||
| Setting string `json:"setting" gorm:"type:text;column:setting"` | Setting string `json:"setting" gorm:"type:text;column:setting"` | ||||
| Remark string `json:"remark,omitempty" gorm:"type:varchar(255)" validate:"max=255"` | 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"` | 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 { | func (user *User) ToBaseUser() *UserBase { | ||||
| @@ -60,6 +86,7 @@ func (user *User) ToBaseUser() *UserBase { | |||||
| Username: user.Username, | Username: user.Username, | ||||
| Setting: user.Setting, | Setting: user.Setting, | ||||
| Email: user.Email, | Email: user.Email, | ||||
| Source: user.Source, | |||||
| } | } | ||||
| return cache | return cache | ||||
| } | } | ||||
| @@ -767,17 +794,8 @@ func ValidateAccessToken(token string) (user *User) { | |||||
| } | } | ||||
| // GetUserQuota gets quota from Redis first, falls back to DB if needed | // GetUserQuota gets quota from Redis first, falls back to DB if needed | ||||
| // 同步用户返回 SyncedQuota,本地用户返回 Quota | |||||
| func GetUserQuota(id int, fromDB bool) (quota int, err error) { | 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 { | if !fromDB && common.RedisEnabled { | ||||
| quota, err := getUserQuotaCache(id) | quota, err := getUserQuotaCache(id) | ||||
| if err == nil { | 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 | // Don't return error - fall through to DB | ||||
| } | } | ||||
| fromDB = true | 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 { | if err != nil { | ||||
| return 0, err | 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 | return quota, nil | ||||
| } | } | ||||
| @@ -877,7 +914,26 @@ func IncreaseUserQuota(id int, quota int, db bool) (err error) { | |||||
| addNewRecord(BatchUpdateTypeUserQuota, id, quota) | addNewRecord(BatchUpdateTypeUserQuota, id, quota) | ||||
| return nil | 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) { | 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) { | //func GetRootUserEmail() (email string) { | ||||
| // DB.Model(&User{}).Where("role = ?", common.RoleRootUser).Select("email").Find(&email) | // DB.Model(&User{}).Where("role = ?", common.RoleRootUser).Select("email").Find(&email) | ||||
| // return email | // return email | ||||
| @@ -15,6 +15,7 @@ import ( | |||||
| // UserBase struct remains the same as it represents the cached data structure | // UserBase struct remains the same as it represents the cached data structure | ||||
| type UserBase struct { | type UserBase struct { | ||||
| Source string `json:"source"` | |||||
| Id int `json:"id"` | Id int `json:"id"` | ||||
| Group string `json:"group"` | Group string `json:"group"` | ||||
| Email string `json:"email"` | 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.ContextKeyUserEmail, user.Email) | ||||
| common.SetContextKey(c, constant.ContextKeyUserName, user.Username) | common.SetContextKey(c, constant.ContextKeyUserName, user.Username) | ||||
| common.SetContextKey(c, constant.ContextKeyUserSetting, user.GetSetting()) | common.SetContextKey(c, constant.ContextKeyUserSetting, user.GetSetting()) | ||||
| common.SetContextKey(c, constant.ContextKeyUserSource, user.Source) | |||||
| } | } | ||||
| func (user *UserBase) GetSetting() dto.UserSetting { | func (user *UserBase) GetSetting() dto.UserSetting { | ||||
| @@ -101,6 +103,7 @@ func GetUserCache(userId int) (userCache *UserBase, err error) { | |||||
| // Create cache object from user data | // Create cache object from user data | ||||
| userCache = &UserBase{ | userCache = &UserBase{ | ||||
| Id: user.Id, | Id: user.Id, | ||||
| Source: user.Source, | |||||
| Group: user.Group, | Group: user.Group, | ||||
| Quota: user.Quota, | Quota: user.Quota, | ||||
| Status: user.Status, | Status: user.Status, | ||||
| @@ -87,9 +87,8 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens | |||||
| if meta.MaxTokens != 0 { | if meta.MaxTokens != 0 { | ||||
| preConsumedTokens += meta.MaxTokens | preConsumedTokens += meta.MaxTokens | ||||
| } | } | ||||
| // 如果没有找到渠道定价,需要获取全局 modelRatio 和 completionRatio | |||||
| // 但如果 ChannelMeta 为 nil(渠道选择前),跳过检查,延迟到渠道选择后 | |||||
| if channelMetaAvailable && !channelPricingFound { | |||||
| // 如果没有找到渠道定价,获取全局 modelRatio 和 completionRatio | |||||
| if !channelPricingFound { | |||||
| var success bool | var success bool | ||||
| var matchName string | var matchName string | ||||
| modelRatio, success, matchName = ratio_setting.GetModelRatio(info.OriginModelName) | modelRatio, success, matchName = ratio_setting.GetModelRatio(info.OriginModelName) | ||||
| @@ -388,5 +388,15 @@ func SetApiRouter(router *gin.Engine) { | |||||
| deploymentsRoute.POST("/:id/extend", controller.ExtendDeployment) | deploymentsRoute.POST("/:id/extend", controller.ExtendDeployment) | ||||
| deploymentsRoute.DELETE("/:id", controller.DeleteDeployment) | 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) | |||||
| } | |||||
| } | } | ||||
| } | } | ||||
| @@ -7,6 +7,7 @@ import ( | |||||
| "sync" | "sync" | ||||
| "github.com/QuantumNous/new-api/common" | "github.com/QuantumNous/new-api/common" | ||||
| "github.com/QuantumNous/new-api/constant" | |||||
| "github.com/QuantumNous/new-api/logger" | "github.com/QuantumNous/new-api/logger" | ||||
| "github.com/QuantumNous/new-api/model" | "github.com/QuantumNous/new-api/model" | ||||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | 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 { | if sub, ok := s.funding.(*SubscriptionFunding); ok && sub.preConsumed > 0 { | ||||
| return true | return true | ||||
| } | } | ||||
| // 同步用户可能在 tokenConsumed=0 时仍预扣了额度 | |||||
| if synced, ok := s.funding.(*SyncedUserFunding); ok && synced.consumed > 0 { | |||||
| return true | |||||
| } | |||||
| return false | 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()) | 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) | pref := common.NormalizeBillingPreference(relayInfo.UserSetting.BillingPreference) | ||||
| // 钱包路径需要先检查用户额度 | // 钱包路径需要先检查用户额度 | ||||
| @@ -345,3 +362,29 @@ func NewBillingSession(c *gin.Context, relayInfo *relaycommon.RelayInfo, preCons | |||||
| return session, nil | 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 | |||||
| } | |||||
| @@ -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()) | |||||
| } | |||||
| @@ -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 | |||||
| } | |||||
| @@ -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) | |||||
| } | |||||
| @@ -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) | |||||
| } | |||||
| @@ -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) | |||||
| } | |||||
| @@ -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)) | |||||
| } | |||||
| } | |||||
| @@ -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"}) | |||||
| } | |||||
| @@ -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"` | |||||
| } | |||||
| @@ -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) | |||||
| } | |||||
| @@ -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 | |||||
| } | |||||
| @@ -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 不匹配,优雅处理 | |||||
| } | |||||
| @@ -256,6 +256,20 @@ func updateConfigFromMap(config interface{}, configMap map[string]string) error | |||||
| // 复杂类型使用JSON反序列化 | // 复杂类型使用JSON反序列化 | ||||
| err := json.Unmarshal([]byte(strValue), field.Addr().Interface()) | err := json.Unmarshal([]byte(strValue), field.Addr().Interface()) | ||||
| if err != nil { | 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 | continue | ||||
| } | } | ||||
| } | } | ||||
| @@ -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 | |||||
| } | |||||
| @@ -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) | |||||
| } | |||||
| } | |||||
| @@ -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 "==========================================" | |||||
| @@ -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 "==========================================" | |||||
| @@ -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 | |||||
| @@ -27,6 +27,7 @@ import SettingsLog from '../../pages/Setting/Operation/SettingsLog'; | |||||
| import SettingsMonitoring from '../../pages/Setting/Operation/SettingsMonitoring'; | import SettingsMonitoring from '../../pages/Setting/Operation/SettingsMonitoring'; | ||||
| import SettingsCreditLimit from '../../pages/Setting/Operation/SettingsCreditLimit'; | import SettingsCreditLimit from '../../pages/Setting/Operation/SettingsCreditLimit'; | ||||
| import SettingsCheckin from '../../pages/Setting/Operation/SettingsCheckin'; | import SettingsCheckin from '../../pages/Setting/Operation/SettingsCheckin'; | ||||
| import SettingsRegionSync from '../../pages/Setting/Operation/SettingsRegionSync'; | |||||
| import { API, showError, toBoolean } from '../../helpers'; | import { API, showError, toBoolean } from '../../helpers'; | ||||
| const OperationSetting = () => { | const OperationSetting = () => { | ||||
| @@ -79,6 +80,20 @@ const OperationSetting = () => { | |||||
| 'checkin_setting.min_quota': 1000, | 'checkin_setting.min_quota': 1000, | ||||
| 'checkin_setting.max_quota': 10000, | '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, | 'token_setting.max_user_tokens': 1000, | ||||
| }); | }); | ||||
| @@ -89,16 +104,22 @@ const OperationSetting = () => { | |||||
| const res = await API.get('/api/option/'); | const res = await API.get('/api/option/'); | ||||
| const { success, message, data } = res.data; | const { success, message, data } = res.data; | ||||
| if (success) { | if (success) { | ||||
| let newInputs = {}; | |||||
| let apiInputs = {}; | |||||
| data.forEach((item) => { | 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 { | } else { | ||||
| showError(message); | showError(message); | ||||
| } | } | ||||
| @@ -154,6 +175,10 @@ const OperationSetting = () => { | |||||
| <Card style={{ marginTop: '10px' }}> | <Card style={{ marginTop: '10px' }}> | ||||
| <SettingsCheckin options={inputs} refresh={onRefresh} /> | <SettingsCheckin options={inputs} refresh={onRefresh} /> | ||||
| </Card> | </Card> | ||||
| {/* 跨区域同步设置 */} | |||||
| <Card style={{ marginTop: '10px' }}> | |||||
| <SettingsRegionSync options={inputs} refresh={onRefresh} /> | |||||
| </Card> | |||||
| </Spin> | </Spin> | ||||
| </> | </> | ||||
| ); | ); | ||||
| @@ -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 <https://www.gnu.org/licenses/>. | |||||
| 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 ( | |||||
| <> | |||||
| <Spin spinning={loading}> | |||||
| <Form | |||||
| values={inputs} | |||||
| getFormApi={(formAPI) => (refForm.current = formAPI)} | |||||
| style={{ marginBottom: 15 }} | |||||
| > | |||||
| <Form.Section text={t('跨区域同步设置')}> | |||||
| <Typography.Text | |||||
| type='tertiary' | |||||
| style={{ marginBottom: 16, display: 'block' }} | |||||
| > | |||||
| {t('配置国内与海外节点之间的用户数据同步,实现跨区域余额管理')} | |||||
| </Typography.Text> | |||||
| {/* 启用开关 - 始终显示 */} | |||||
| <Row gutter={16}> | |||||
| <Col xs={24} sm={12} md={8} lg={8} xl={8}> | |||||
| <Form.Switch | |||||
| field={'region_sync.enabled'} | |||||
| label={t('启用跨区域同步')} | |||||
| size='default' | |||||
| checkedText='|' | |||||
| uncheckedText='〇' | |||||
| onChange={handleFieldChange('region_sync.enabled')} | |||||
| /> | |||||
| </Col> | |||||
| </Row> | |||||
| {/* 保存按钮 - 未启用时也显示 */} | |||||
| {!enabled && ( | |||||
| <Row style={{ marginTop: 16 }}> | |||||
| <Button size='default' onClick={onSubmit}> | |||||
| {t('保存跨区域同步设置')} | |||||
| </Button> | |||||
| </Row> | |||||
| )} | |||||
| {/* 启用后才显示以下内容 */} | |||||
| {enabled && ( | |||||
| <> | |||||
| {/* 区域 ID */} | |||||
| <Row gutter={16} style={{ marginTop: 16 }}> | |||||
| <Col xs={24} sm={12} md={8} lg={8} xl={8}> | |||||
| <Form.Input | |||||
| field={'region_sync.region_id'} | |||||
| label={t('区域 ID')} | |||||
| placeholder={t('例如: cn-east, us-west')} | |||||
| onChange={handleFieldChange('region_sync.region_id')} | |||||
| /> | |||||
| </Col> | |||||
| </Row> | |||||
| {/* 角色选择 */} | |||||
| <Row gutter={16} style={{ marginTop: 16 }}> | |||||
| <Col xs={24} sm={12} md={8} lg={8} xl={8}> | |||||
| <Form.Switch | |||||
| field={'region_sync.is_master'} | |||||
| label={t('当前节点为主节点')} | |||||
| size='default' | |||||
| checkedText='|' | |||||
| uncheckedText='〇' | |||||
| onChange={handleFieldChange('region_sync.is_master')} | |||||
| /> | |||||
| </Col> | |||||
| </Row> | |||||
| {/* 共用设置:API 密钥 */} | |||||
| <Row gutter={16} style={{ marginTop: 16 }}> | |||||
| <Col xs={24} sm={24} md={12} lg={12} xl={12}> | |||||
| <Form.Input | |||||
| field={'region_sync.sync_api_key'} | |||||
| label={t('同步 API 密钥')} | |||||
| placeholder={t('用于节点间认证的密钥(已设置的密钥不会显示)')} | |||||
| onChange={handleFieldChange('region_sync.sync_api_key')} | |||||
| type='password' | |||||
| mode='password' | |||||
| /> | |||||
| </Col> | |||||
| </Row> | |||||
| {/* 主节点专属设置 */} | |||||
| {isMaster && ( | |||||
| <> | |||||
| <Row gutter={16} style={{ marginTop: 16 }}> | |||||
| <Col xs={24}> | |||||
| <div className='semi-form-field'> | |||||
| <label className='semi-form-field-label'> | |||||
| {t('从节点地址列表')} | |||||
| </label> | |||||
| <TagInput | |||||
| value={slaveEndpoints} | |||||
| onChange={handleTagInputChange( | |||||
| 'region_sync.slave_endpoints', | |||||
| setSlaveEndpoints | |||||
| )} | |||||
| placeholder={t('输入地址后按回车添加')} | |||||
| separator=',' | |||||
| addOnBlur | |||||
| style={{ marginTop: 4 }} | |||||
| /> | |||||
| <Typography.Text | |||||
| type='tertiary' | |||||
| size='small' | |||||
| style={{ marginTop: 4, display: 'block' }} | |||||
| > | |||||
| {t('例如: https://slave1.example.com')} | |||||
| </Typography.Text> | |||||
| </div> | |||||
| </Col> | |||||
| </Row> | |||||
| <Row gutter={16} style={{ marginTop: 16 }}> | |||||
| <Col xs={24} sm={12} md={6} lg={6} xl={6}> | |||||
| <Form.InputNumber | |||||
| field={'region_sync.min_balance_threshold'} | |||||
| label={t('最小余额阈值')} | |||||
| placeholder={t('同步触发阈值')} | |||||
| onChange={handleFieldChange( | |||||
| 'region_sync.min_balance_threshold' | |||||
| )} | |||||
| min={0} | |||||
| /> | |||||
| </Col> | |||||
| <Col xs={24} sm={12} md={6} lg={6} xl={6}> | |||||
| <Form.InputNumber | |||||
| field={'region_sync.sync_interval_seconds'} | |||||
| label={t('同步间隔 (秒)')} | |||||
| placeholder={t('批量同步间隔')} | |||||
| onChange={handleFieldChange( | |||||
| 'region_sync.sync_interval_seconds' | |||||
| )} | |||||
| min={10} | |||||
| /> | |||||
| </Col> | |||||
| <Col xs={24} sm={12} md={6} lg={6} xl={6}> | |||||
| <Form.InputNumber | |||||
| field={'region_sync.max_retry_count'} | |||||
| label={t('最大重试次数')} | |||||
| placeholder={t('失败重试次数')} | |||||
| onChange={handleFieldChange('region_sync.max_retry_count')} | |||||
| min={0} | |||||
| max={10} | |||||
| /> | |||||
| </Col> | |||||
| <Col xs={24} sm={12} md={6} lg={6} xl={6}> | |||||
| <Form.InputNumber | |||||
| field={'region_sync.sync_batch_size'} | |||||
| label={t('同步批次大小')} | |||||
| placeholder={t('每批处理记录数')} | |||||
| onChange={handleFieldChange('region_sync.sync_batch_size')} | |||||
| min={1} | |||||
| max={1000} | |||||
| /> | |||||
| </Col> | |||||
| </Row> | |||||
| <Row gutter={16} style={{ marginTop: 16 }}> | |||||
| <Col xs={24} sm={12} md={8} lg={8} xl={8}> | |||||
| <Form.Switch | |||||
| field={'region_sync.disable_cached_consume'} | |||||
| label={t('禁用缓存消费')} | |||||
| size='default' | |||||
| checkedText='|' | |||||
| uncheckedText='〇' | |||||
| onChange={handleFieldChange( | |||||
| 'region_sync.disable_cached_consume' | |||||
| )} | |||||
| /> | |||||
| </Col> | |||||
| </Row> | |||||
| </> | |||||
| )} | |||||
| {/* 从节点专属设置 */} | |||||
| {!isMaster && ( | |||||
| <> | |||||
| <Row gutter={16} style={{ marginTop: 16 }}> | |||||
| <Col xs={24} sm={24} md={12} lg={12} xl={12}> | |||||
| <div className='semi-form-field'> | |||||
| <label className='semi-form-field-label'> | |||||
| {t('主节点地址')} | |||||
| </label> | |||||
| <TagInput | |||||
| value={masterEndpoints} | |||||
| onChange={handleTagInputChange( | |||||
| 'region_sync.master_endpoint', | |||||
| setMasterEndpoints | |||||
| )} | |||||
| placeholder={t('输入地址后按回车添加')} | |||||
| separator=',' | |||||
| addOnBlur | |||||
| style={{ marginTop: 4 }} | |||||
| /> | |||||
| <Typography.Text | |||||
| type='tertiary' | |||||
| size='small' | |||||
| style={{ marginTop: 4, display: 'block' }} | |||||
| > | |||||
| {t('例如: https://master.example.com')} | |||||
| </Typography.Text> | |||||
| </div> | |||||
| </Col> | |||||
| </Row> | |||||
| <Row gutter={16} style={{ marginTop: 16 }}> | |||||
| <Col xs={24} sm={12} md={8} lg={8} xl={8}> | |||||
| <Form.InputNumber | |||||
| field={'region_sync.quota_sync_interval_seconds'} | |||||
| label={t('余额同步间隔 (秒)')} | |||||
| placeholder={t('从主节点拉取余额的间隔')} | |||||
| onChange={handleFieldChange( | |||||
| 'region_sync.quota_sync_interval_seconds' | |||||
| )} | |||||
| min={30} | |||||
| /> | |||||
| </Col> | |||||
| </Row> | |||||
| </> | |||||
| )} | |||||
| {/* 保存按钮 */} | |||||
| <Row style={{ marginTop: 16 }}> | |||||
| <Button size='default' onClick={onSubmit}> | |||||
| {t('保存跨区域同步设置')} | |||||
| </Button> | |||||
| </Row> | |||||
| </> | |||||
| )} | |||||
| </Form.Section> | |||||
| </Form> | |||||
| </Spin> | |||||
| </> | |||||
| ); | |||||
| } | |||||