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/QuantumNous/new-api/setting/system_setting" "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 } // 查询扣费后的余额(同一事务内可读到自己的更新,MySQL REPEATABLE READ / SQLite 均安全) 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()) } } // GetSyncConfig Master 节点返回同步配置给 Slave func GetSyncConfig(c *gin.Context) { settings := system_setting.GetRegionSyncSettings() c.JSON(http.StatusOK, region_sync.SyncConfigResponse{ Success: true, MinBalanceThreshold: settings.MinBalanceThreshold, MaxRetryCount: settings.MaxRetryCount, SyncBatchSize: settings.SyncBatchSize, SyncIntervalSeconds: settings.SyncIntervalSeconds, }) }