|
- 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,
- })
- }
|