实现 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" | |||
| ContextKeyUsingGroup ContextKey = "group" | |||
| ContextKeyUserName ContextKey = "username" | |||
| ContextKeyUserSource ContextKey = "user_source" | |||
| 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 { | |||
| logger.LogInfo(c, fmt.Sprintf("模型 %s 免费,跳过预扣费", relayInfo.OriginModelName)) | |||
| } else { | |||
| common.SysLog(fmt.Sprintf("[RegionSync] relay PreConsumeBilling: userId=%d, QuotaToPreConsume=%d, FreeModel=%v, UsePrice=%v, ModelRatio=%.4f, GroupRatio=%.4f, ModelPrice=%.4f, model=%s", | |||
| relayInfo.UserId, priceData.QuotaToPreConsume, priceData.FreeModel, priceData.UsePrice, priceData.ModelRatio, priceData.GroupRatioInfo.GroupRatio, priceData.ModelPrice, relayInfo.OriginModelName)) | |||
| newAPIError = service.PreConsumeBilling(c, priceData.QuotaToPreConsume, relayInfo) | |||
| if newAPIError != nil { | |||
| return | |||
| @@ -16,6 +16,7 @@ import ( | |||
| "github.com/QuantumNous/new-api/logger" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/QuantumNous/new-api/service" | |||
| "github.com/QuantumNous/new-api/service/region_sync" | |||
| "github.com/QuantumNous/new-api/setting" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| @@ -184,12 +185,8 @@ func Register(c *gin.Context) { | |||
| return | |||
| } | |||
| // 获取插入后的用户ID | |||
| var insertedUser model.User | |||
| if err := model.DB.Where("username = ?", cleanUser.Username).First(&insertedUser).Error; err != nil { | |||
| common.ApiErrorI18n(c, i18n.MsgUserRegisterFailed) | |||
| return | |||
| } | |||
| // 同步用户到海外节点(异步执行,不阻塞注册流程) | |||
| region_sync.PushUserCreateToSlave(&cleanUser) | |||
| // 生成默认令牌 | |||
| if constant.GenerateDefaultToken { | |||
| key, err := common.GenerateKey() | |||
| @@ -200,7 +197,7 @@ func Register(c *gin.Context) { | |||
| } | |||
| // 生成默认令牌 | |||
| token := model.Token{ | |||
| UserId: insertedUser.Id, // 使用插入后的用户ID | |||
| UserId: cleanUser.Id, // GORM Create 后已填充 ID | |||
| Name: cleanUser.Username + "的初始令牌", | |||
| Key: key, | |||
| CreatedTime: common.GetTimestamp(), | |||
| @@ -383,6 +380,11 @@ func GetSelf(c *gin.Context) { | |||
| userSetting := user.GetSetting() | |||
| // 构建响应数据,包含用户信息和权限 | |||
| // Slave 节点使用 SyncedQuota 作为用户额度 | |||
| quota := user.Quota | |||
| if user.IsSyncedUser() { | |||
| quota = user.SyncedQuota | |||
| } | |||
| responseData := map[string]interface{}{ | |||
| "id": user.Id, | |||
| "username": user.Username, | |||
| @@ -396,7 +398,7 @@ func GetSelf(c *gin.Context) { | |||
| "wechat_id": user.WeChatId, | |||
| "telegram_id": user.TelegramId, | |||
| "group": user.Group, | |||
| "quota": user.Quota, | |||
| "quota": quota, | |||
| "used_quota": user.UsedQuota, | |||
| "request_count": user.RequestCount, | |||
| "aff_code": user.AffCode, | |||
| @@ -22,8 +22,10 @@ import ( | |||
| "github.com/QuantumNous/new-api/relay" | |||
| "github.com/QuantumNous/new-api/router" | |||
| "github.com/QuantumNous/new-api/service" | |||
| "github.com/QuantumNous/new-api/service/region_sync" | |||
| _ "github.com/QuantumNous/new-api/setting/performance_setting" | |||
| "github.com/QuantumNous/new-api/setting/ratio_setting" | |||
| "github.com/QuantumNous/new-api/setting/system_setting" | |||
| "github.com/bytedance/gopkg/util/gopool" | |||
| "github.com/gin-contrib/sessions" | |||
| @@ -174,7 +176,12 @@ func main() { | |||
| Secure: false, | |||
| SameSite: http.SameSiteStrictMode, | |||
| }) | |||
| server.Use(sessions.Sessions("session", store)) | |||
| // 支持通过环境变量自定义 session 名称,用于多节点部署时避免 cookie 冲突 | |||
| sessionName := os.Getenv("SESSION_NAME") | |||
| if sessionName == "" { | |||
| sessionName = "session" | |||
| } | |||
| server.Use(sessions.Sessions(sessionName, store)) | |||
| InjectUmamiAnalytics() | |||
| InjectGoogleAnalytics() | |||
| @@ -309,5 +316,19 @@ func InitResources() error { | |||
| // Don't return error, custom OAuth is not critical | |||
| } | |||
| // 注册余额更新回调,Master 节点更新用户余额后推送到 Slave 节点 | |||
| model.SetQuotaUpdateCallback(func(userId int, quota int) { | |||
| if system_setting.GetRegionSyncSettings().IsMaster { | |||
| region_sync.PushQuotaUpdateToSlave(userId, quota) | |||
| } | |||
| }) | |||
| // Slave 节点启动后台同步任务(批量扣费同步 + 余额定时拉取 + 清理) | |||
| syncSettings := system_setting.GetRegionSyncSettings() | |||
| if syncSettings.Enabled && !syncSettings.IsMaster { | |||
| syncManager := region_sync.NewSyncManager() | |||
| syncManager.StartSyncWorkers() | |||
| } | |||
| return nil | |||
| } | |||
| @@ -22,7 +22,7 @@ DOCKER_TAG := $(DOCKER_REGISTRY):$(BUILD_TIME)-$(BRANCH_NAME) | |||
| docker-build: | |||
| @echo "Building Docker image with tag: $(DOCKER_TAG)" | |||
| docker build -t $(DOCKER_TAG) . | |||
| docker push $(DOCKER_TAG) | |||
| #docker push $(DOCKER_TAG) | |||
| docker-push: docker-build | |||
| @echo "Pushing Docker image: $(DOCKER_TAG)" | |||
| @@ -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{}, | |||
| &ChannelPricing{}, | |||
| &PricingTag{}, | |||
| &PendingSyncRecord{}, | |||
| &QuotaSyncLog{}, | |||
| ) | |||
| if err != nil { | |||
| return err | |||
| @@ -332,6 +334,8 @@ func migrateDBFast() error { | |||
| {&UserOAuthBinding{}, "UserOAuthBinding"}, | |||
| {&ChannelPricing{}, "ChannelPricing"}, | |||
| {&PricingTag{}, "PricingTag"}, | |||
| {&PendingSyncRecord{}, "PendingSyncRecord"}, | |||
| {&QuotaSyncLog{}, "QuotaSyncLog"}, | |||
| } | |||
| // 动态计算migration数量,确保errChan缓冲区足够大 | |||
| errChan := make(chan error, len(migrations)) | |||
| @@ -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()) | |||
| return 0, ErrRedeemFailed | |||
| } | |||
| // 触发余额更新回调(在事务外) | |||
| if quotaUpdateCallback != nil { | |||
| newQuota, _ := GetUserQuota(userId, true) | |||
| // 使用 defer/recover 保护回调 | |||
| func() { | |||
| defer func() { | |||
| if r := recover(); r != nil { | |||
| common.SysError(fmt.Sprintf("quota update callback panic in Redeem: %v", r)) | |||
| } | |||
| }() | |||
| quotaUpdateCallback(userId, newQuota) | |||
| }() | |||
| } | |||
| RecordLog(userId, LogTypeTopup, fmt.Sprintf("通过兑换码充值 %s,兑换码ID %d", logger.LogQuota(redemption.Quota), redemption.Id)) | |||
| return redemption.Quota, nil | |||
| } | |||
| @@ -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" | |||
| "strconv" | |||
| "strings" | |||
| "time" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/dto" | |||
| @@ -17,6 +18,16 @@ import ( | |||
| const UserNameMaxLength = 20 | |||
| // QuotaUpdateCallback 余额更新回调函数类型 | |||
| type QuotaUpdateCallback func(userId int, quota int) | |||
| var quotaUpdateCallback QuotaUpdateCallback | |||
| // SetQuotaUpdateCallback 设置余额更新回调 | |||
| func SetQuotaUpdateCallback(cb QuotaUpdateCallback) { | |||
| quotaUpdateCallback = cb | |||
| } | |||
| // User if you add sensitive fields, don't forget to clean them in setupLogin function. | |||
| // Otherwise, the sensitive information will be saved on local storage in plain text! | |||
| type User struct { | |||
| @@ -49,6 +60,21 @@ type User struct { | |||
| Setting string `json:"setting" gorm:"type:text;column:setting"` | |||
| Remark string `json:"remark,omitempty" gorm:"type:varchar(255)" validate:"max=255"` | |||
| StripeCustomer string `json:"stripe_customer" gorm:"type:varchar(64);column:stripe_customer;index"` | |||
| // 跨地区同步相关 | |||
| Source string `json:"source" gorm:"type:varchar(20);default:'local'"` | |||
| RemoteUserId int `json:"remote_user_id" gorm:"type:int;default:0;column:remote_user_id"` | |||
| SyncedQuota int `json:"synced_quota" gorm:"type:int;default:0;column:synced_quota"` | |||
| LastSyncAt int64 `json:"last_sync_at" gorm:"type:bigint;default:0;column:last_sync_at"` | |||
| } | |||
| // IsSyncedUser 判断是否为国内同步用户 | |||
| func (u *User) IsSyncedUser() bool { | |||
| return common.IsSyncedUser(u.Source, u.Id) | |||
| } | |||
| // IsLocalUser 判断是否为本地用户 | |||
| func (u *User) IsLocalUser() bool { | |||
| return common.IsLocalUser(u.Source, u.Id) | |||
| } | |||
| func (user *User) ToBaseUser() *UserBase { | |||
| @@ -60,6 +86,7 @@ func (user *User) ToBaseUser() *UserBase { | |||
| Username: user.Username, | |||
| Setting: user.Setting, | |||
| Email: user.Email, | |||
| Source: user.Source, | |||
| } | |||
| return cache | |||
| } | |||
| @@ -767,17 +794,8 @@ func ValidateAccessToken(token string) (user *User) { | |||
| } | |||
| // GetUserQuota gets quota from Redis first, falls back to DB if needed | |||
| // 同步用户返回 SyncedQuota,本地用户返回 Quota | |||
| func GetUserQuota(id int, fromDB bool) (quota int, err error) { | |||
| defer func() { | |||
| // Update Redis cache asynchronously on successful DB read | |||
| if shouldUpdateRedis(fromDB, err) { | |||
| gopool.Go(func() { | |||
| if err := updateUserQuotaCache(id, quota); err != nil { | |||
| common.SysLog("failed to update user quota cache: " + err.Error()) | |||
| } | |||
| }) | |||
| } | |||
| }() | |||
| if !fromDB && common.RedisEnabled { | |||
| quota, err := getUserQuotaCache(id) | |||
| if err == nil { | |||
| @@ -785,12 +803,31 @@ func GetUserQuota(id int, fromDB bool) (quota int, err error) { | |||
| } | |||
| // Don't return error - fall through to DB | |||
| } | |||
| fromDB = true | |||
| err = DB.Model(&User{}).Where("id = ?", id).Select("quota").Find("a).Error | |||
| // 查询用户的 source 和额度字段 | |||
| var user User | |||
| err = DB.Model(&User{}).Select("source", "quota", "synced_quota").Where("id = ?", id).First(&user).Error | |||
| if err != nil { | |||
| return 0, err | |||
| } | |||
| // 同步用户返回 SyncedQuota,本地用户返回 Quota | |||
| if user.IsSyncedUser() { | |||
| quota = user.SyncedQuota | |||
| } else { | |||
| quota = user.Quota | |||
| } | |||
| // Update Redis cache asynchronously if enabled | |||
| if common.RedisEnabled { | |||
| gopool.Go(func() { | |||
| if err := updateUserQuotaCache(id, quota); err != nil { | |||
| common.SysLog("failed to update user quota cache: " + err.Error()) | |||
| } | |||
| }) | |||
| } | |||
| return quota, nil | |||
| } | |||
| @@ -877,7 +914,26 @@ func IncreaseUserQuota(id int, quota int, db bool) (err error) { | |||
| addNewRecord(BatchUpdateTypeUserQuota, id, quota) | |||
| return nil | |||
| } | |||
| return increaseUserQuota(id, quota) | |||
| err = increaseUserQuota(id, quota) | |||
| if err != nil { | |||
| return err | |||
| } | |||
| // 触发余额更新回调(如果有注册) | |||
| if quotaUpdateCallback != nil { | |||
| newQuota, _ := GetUserQuota(id, true) | |||
| // 使用 defer/recover 保护回调,防止 panic 影响主流程 | |||
| func() { | |||
| defer func() { | |||
| if r := recover(); r != nil { | |||
| common.SysError(fmt.Sprintf("quota update callback panic: %v", r)) | |||
| } | |||
| }() | |||
| quotaUpdateCallback(id, newQuota) | |||
| }() | |||
| } | |||
| return nil | |||
| } | |||
| func increaseUserQuota(id int, quota int) (err error) { | |||
| @@ -924,6 +980,47 @@ func DeltaUpdateUserQuota(id int, delta int) (err error) { | |||
| } | |||
| } | |||
| // GetSyncedUsers 获取所有同步用户 | |||
| func GetSyncedUsers() []User { | |||
| var users []User | |||
| DB.Where("source = ?", common.UserSourceSynced).Find(&users) | |||
| return users | |||
| } | |||
| // UpdateSyncedQuota 更新同步用户的 synced_quota | |||
| func UpdateSyncedQuota(userId int, quota int) error { | |||
| return DB.Model(&User{}).Where("id = ?", userId).Updates(map[string]interface{}{ | |||
| "synced_quota": quota, | |||
| "last_sync_at": time.Now().Unix(), | |||
| }).Error | |||
| } | |||
| // AtomicDecreaseSyncedQuota 原子扣减 synced_quota,检查余额并扣减在一条 SQL 中完成。 | |||
| // 返回 (当前余额, 是否成功, 错误)。 | |||
| func AtomicDecreaseSyncedQuota(userId int, amount int, threshold int) (int, bool, error) { | |||
| result := DB.Model(&User{}). | |||
| Where("id = ? AND synced_quota >= ?", userId, amount+threshold). | |||
| Update("synced_quota", gorm.Expr("synced_quota - ?", amount)) | |||
| if result.Error != nil { | |||
| return 0, false, result.Error | |||
| } | |||
| if result.RowsAffected == 0 { | |||
| var quota int | |||
| DB.Model(&User{}).Where("id = ?", userId).Select("synced_quota").Scan("a) | |||
| return quota, false, nil | |||
| } | |||
| var newQuota int | |||
| DB.Model(&User{}).Where("id = ?", userId).Select("synced_quota").Scan(&newQuota) | |||
| return newQuota, true, nil | |||
| } | |||
| // IncreaseSyncedQuota 原子增加 synced_quota(用于退款) | |||
| func IncreaseSyncedQuota(userId int, amount int) error { | |||
| return DB.Model(&User{}). | |||
| Where("id = ?", userId). | |||
| Update("synced_quota", gorm.Expr("synced_quota + ?", amount)).Error | |||
| } | |||
| //func GetRootUserEmail() (email string) { | |||
| // DB.Model(&User{}).Where("role = ?", common.RoleRootUser).Select("email").Find(&email) | |||
| // return email | |||
| @@ -15,6 +15,7 @@ import ( | |||
| // UserBase struct remains the same as it represents the cached data structure | |||
| type UserBase struct { | |||
| Source string `json:"source"` | |||
| Id int `json:"id"` | |||
| Group string `json:"group"` | |||
| Email string `json:"email"` | |||
| @@ -31,6 +32,7 @@ func (user *UserBase) WriteContext(c *gin.Context) { | |||
| common.SetContextKey(c, constant.ContextKeyUserEmail, user.Email) | |||
| common.SetContextKey(c, constant.ContextKeyUserName, user.Username) | |||
| common.SetContextKey(c, constant.ContextKeyUserSetting, user.GetSetting()) | |||
| common.SetContextKey(c, constant.ContextKeyUserSource, user.Source) | |||
| } | |||
| func (user *UserBase) GetSetting() dto.UserSetting { | |||
| @@ -101,6 +103,7 @@ func GetUserCache(userId int) (userCache *UserBase, err error) { | |||
| // Create cache object from user data | |||
| userCache = &UserBase{ | |||
| Id: user.Id, | |||
| Source: user.Source, | |||
| Group: user.Group, | |||
| Quota: user.Quota, | |||
| Status: user.Status, | |||
| @@ -87,9 +87,8 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens | |||
| if meta.MaxTokens != 0 { | |||
| preConsumedTokens += meta.MaxTokens | |||
| } | |||
| // 如果没有找到渠道定价,需要获取全局 modelRatio 和 completionRatio | |||
| // 但如果 ChannelMeta 为 nil(渠道选择前),跳过检查,延迟到渠道选择后 | |||
| if channelMetaAvailable && !channelPricingFound { | |||
| // 如果没有找到渠道定价,获取全局 modelRatio 和 completionRatio | |||
| if !channelPricingFound { | |||
| var success bool | |||
| var matchName string | |||
| modelRatio, success, matchName = ratio_setting.GetModelRatio(info.OriginModelName) | |||
| @@ -388,5 +388,15 @@ func SetApiRouter(router *gin.Engine) { | |||
| deploymentsRoute.POST("/:id/extend", controller.ExtendDeployment) | |||
| deploymentsRoute.DELETE("/:id", controller.DeleteDeployment) | |||
| } | |||
| // Region sync API routes (for inter-node communication) | |||
| syncRoute := apiRouter.Group("/internal/sync") | |||
| syncRoute.Use(middleware.SyncAuth()) | |||
| { | |||
| syncRoute.POST("/user/create", controller.ReceiveSyncedUserCreate) | |||
| syncRoute.POST("/quota/update", controller.ReceiveQuotaUpdate) | |||
| syncRoute.POST("/quota/query", controller.QueryUserQuota) | |||
| syncRoute.POST("/quota/batch-deduct", controller.BatchDeductQuota) | |||
| } | |||
| } | |||
| } | |||
| @@ -7,6 +7,7 @@ import ( | |||
| "sync" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/logger" | |||
| "github.com/QuantumNous/new-api/model" | |||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | |||
| @@ -132,6 +133,10 @@ func (s *BillingSession) needsRefundLocked() bool { | |||
| if sub, ok := s.funding.(*SubscriptionFunding); ok && sub.preConsumed > 0 { | |||
| return true | |||
| } | |||
| // 同步用户可能在 tokenConsumed=0 时仍预扣了额度 | |||
| if synced, ok := s.funding.(*SyncedUserFunding); ok && synced.consumed > 0 { | |||
| return true | |||
| } | |||
| return false | |||
| } | |||
| @@ -257,6 +262,18 @@ func NewBillingSession(c *gin.Context, relayInfo *relaycommon.RelayInfo, preCons | |||
| return nil, types.NewError(fmt.Errorf("relayInfo is nil"), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry()) | |||
| } | |||
| // 检查是否为同步用户(从上下文中获取,避免热路径数据库查询) | |||
| userSource := common.GetContextKeyString(c, constant.ContextKeyUserSource) | |||
| common.SysLog(fmt.Sprintf("[RegionSync] NewBillingSession: userId=%d, userSource=%s, preConsumedQuota=%d, requestId=%s", relayInfo.UserId, userSource, preConsumedQuota, relayInfo.RequestId)) | |||
| if userSource == common.UserSourceSynced { | |||
| user, userErr := model.GetUserById(relayInfo.UserId, false) | |||
| if userErr == nil { | |||
| common.SysLog(fmt.Sprintf("[RegionSync] synced user detected: userId=%d, remoteUserId=%d, syncedQuota=%d, requestId=%s", user.Id, user.RemoteUserId, user.SyncedQuota, relayInfo.RequestId)) | |||
| return newSyncedUserBillingSession(c, relayInfo, user, preConsumedQuota) | |||
| } | |||
| common.SysError(fmt.Sprintf("[RegionSync] GetUserById failed for synced user: userId=%d, err=%v", relayInfo.UserId, userErr)) | |||
| } | |||
| pref := common.NormalizeBillingPreference(relayInfo.UserSetting.BillingPreference) | |||
| // 钱包路径需要先检查用户额度 | |||
| @@ -345,3 +362,29 @@ func NewBillingSession(c *gin.Context, relayInfo *relaycommon.RelayInfo, preCons | |||
| return session, nil | |||
| } | |||
| } | |||
| // newSyncedUserBillingSession 为同步用户创建 BillingSession | |||
| // 同步用户的余额在 master 节点,本地使用 SyncedUserFunding 管理预扣费和结算 | |||
| func newSyncedUserBillingSession(c *gin.Context, relayInfo *relaycommon.RelayInfo, user *model.User, preConsumedQuota int) (*BillingSession, *types.NewAPIError) { | |||
| syncedQuota := user.SyncedQuota | |||
| if syncedQuota <= 0 { | |||
| syncedQuota = user.Quota | |||
| } | |||
| session := &BillingSession{ | |||
| relayInfo: relayInfo, | |||
| funding: NewSyncedUserFunding( | |||
| user.Id, | |||
| user.RemoteUserId, | |||
| relayInfo.RequestId, | |||
| syncedQuota, | |||
| ), | |||
| } | |||
| relayInfo.UserQuota = syncedQuota | |||
| if apiErr := session.preConsume(c, preConsumedQuota); apiErr != nil { | |||
| return nil, apiErr | |||
| } | |||
| return session, nil | |||
| } | |||
| @@ -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反序列化 | |||
| err := json.Unmarshal([]byte(strValue), field.Addr().Interface()) | |||
| if err != nil { | |||
| // 对于字符串切片,尝试兼容换行分隔的格式 | |||
| if field.Kind() == reflect.Slice && field.Type().Elem().Kind() == reflect.String { | |||
| lines := strings.Split(strValue, "\n") | |||
| var result []string | |||
| for _, line := range lines { | |||
| line = strings.TrimSpace(line) | |||
| if line != "" { | |||
| result = append(result, line) | |||
| } | |||
| } | |||
| if len(result) > 0 { | |||
| field.Set(reflect.ValueOf(result)) | |||
| } | |||
| } | |||
| continue | |||
| } | |||
| } | |||
| @@ -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 SettingsCreditLimit from '../../pages/Setting/Operation/SettingsCreditLimit'; | |||
| import SettingsCheckin from '../../pages/Setting/Operation/SettingsCheckin'; | |||
| import SettingsRegionSync from '../../pages/Setting/Operation/SettingsRegionSync'; | |||
| import { API, showError, toBoolean } from '../../helpers'; | |||
| const OperationSetting = () => { | |||
| @@ -79,6 +80,20 @@ const OperationSetting = () => { | |||
| 'checkin_setting.min_quota': 1000, | |||
| 'checkin_setting.max_quota': 10000, | |||
| /* 跨区域同步设置 */ | |||
| 'region_sync.enabled': false, | |||
| 'region_sync.region_id': '', | |||
| 'region_sync.is_master': false, | |||
| 'region_sync.master_endpoint': '', | |||
| 'region_sync.slave_endpoints': '', | |||
| 'region_sync.sync_api_key': '', | |||
| 'region_sync.min_balance_threshold': 100000, | |||
| 'region_sync.sync_interval_seconds': 60, | |||
| 'region_sync.max_retry_count': 3, | |||
| 'region_sync.sync_batch_size': 100, | |||
| 'region_sync.quota_sync_interval_seconds': 300, | |||
| 'region_sync.disable_cached_consume': false, | |||
| /* 令牌设置 */ | |||
| 'token_setting.max_user_tokens': 1000, | |||
| }); | |||
| @@ -89,16 +104,22 @@ const OperationSetting = () => { | |||
| const res = await API.get('/api/option/'); | |||
| const { success, message, data } = res.data; | |||
| if (success) { | |||
| let newInputs = {}; | |||
| let apiInputs = {}; | |||
| data.forEach((item) => { | |||
| if (typeof inputs[item.key] === 'boolean') { | |||
| newInputs[item.key] = toBoolean(item.value); | |||
| } else { | |||
| newInputs[item.key] = item.value; | |||
| // 只处理已知字段,避免添加未知字段 | |||
| if (inputs.hasOwnProperty(item.key)) { | |||
| if (typeof inputs[item.key] === 'boolean') { | |||
| apiInputs[item.key] = toBoolean(item.value); | |||
| } else if (typeof inputs[item.key] === 'number') { | |||
| apiInputs[item.key] = parseFloat(item.value) || 0; | |||
| } else { | |||
| apiInputs[item.key] = item.value; | |||
| } | |||
| } | |||
| }); | |||
| setInputs(newInputs); | |||
| // 合并 API 数据和初始默认值 | |||
| setInputs((prevInputs) => ({ ...prevInputs, ...apiInputs })); | |||
| } else { | |||
| showError(message); | |||
| } | |||
| @@ -154,6 +175,10 @@ const OperationSetting = () => { | |||
| <Card style={{ marginTop: '10px' }}> | |||
| <SettingsCheckin options={inputs} refresh={onRefresh} /> | |||
| </Card> | |||
| {/* 跨区域同步设置 */} | |||
| <Card style={{ marginTop: '10px' }}> | |||
| <SettingsRegionSync options={inputs} refresh={onRefresh} /> | |||
| </Card> | |||
| </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> | |||
| </> | |||
| ); | |||
| } | |||