Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

274 строки
9.6 KiB

  1. package controller
  2. import (
  3. "net/http"
  4. "time"
  5. "github.com/QuantumNous/new-api/common"
  6. "github.com/QuantumNous/new-api/logger"
  7. "github.com/QuantumNous/new-api/model"
  8. "github.com/QuantumNous/new-api/service/region_sync"
  9. "github.com/gin-gonic/gin"
  10. "gorm.io/gorm"
  11. )
  12. // 统一错误响应
  13. func syncErrorResponse(c *gin.Context, status int, err string) {
  14. c.JSON(status, gin.H{"success": false, "error": err})
  15. }
  16. // 统一成功响应
  17. func syncSuccessResponse(c *gin.Context, message string) {
  18. c.JSON(http.StatusOK, gin.H{"success": true, "message": message})
  19. }
  20. // ReceiveSyncedUserCreate 接收从 master 节点推送过来的用户创建请求
  21. func ReceiveSyncedUserCreate(c *gin.Context) {
  22. var req region_sync.SyncUserRequest
  23. if err := c.ShouldBindJSON(&req); err != nil {
  24. syncErrorResponse(c, http.StatusBadRequest, "invalid request body")
  25. return
  26. }
  27. logger.LogDebug(c.Request.Context(), "[RegionSync] ReceiveSyncedUserCreate: username=%s, remoteUserId=%d, quota=%d", req.Username, req.RemoteUserId, req.Quota)
  28. // 检查用户是否已存在
  29. var existingUser model.User
  30. if err := model.DB.Where("remote_user_id = ?", req.RemoteUserId).First(&existingUser).Error; err == nil {
  31. logger.LogDebug(c.Request.Context(), "[RegionSync] ReceiveSyncedUserCreate: user already exists, remoteUserId=%d", req.RemoteUserId)
  32. syncSuccessResponse(c, "user already exists")
  33. return
  34. }
  35. // 创建新用户
  36. user := createSyncedUser(&req)
  37. if err := model.DB.Create(&user).Error; err != nil {
  38. logger.LogError(c.Request.Context(), "failed to create synced user: "+err.Error())
  39. syncErrorResponse(c, http.StatusInternalServerError, "failed to create user")
  40. return
  41. }
  42. // 记录同步日志
  43. recordSyncLog(c, user.Id, user.RemoteUserId, model.SyncTypeUserCreate, "", 0, 0, model.SyncDirectionCnToOv)
  44. logger.LogInfo(c.Request.Context(), "created synced user "+user.Username)
  45. syncSuccessResponse(c, "user created successfully")
  46. }
  47. // createSyncedUser 根据请求创建同步用户
  48. func createSyncedUser(req *region_sync.SyncUserRequest) model.User {
  49. affCode := req.AffCode
  50. if affCode == "" {
  51. affCode = common.GetRandomString(4)
  52. }
  53. displayName := req.DisplayName
  54. if displayName == "" {
  55. displayName = req.Username
  56. }
  57. group := req.Group
  58. if group == "" {
  59. group = "default"
  60. }
  61. return model.User{
  62. Username: req.Username,
  63. Email: req.Email,
  64. Password: req.PasswordHash,
  65. DisplayName: displayName,
  66. Quota: req.Quota,
  67. RemoteUserId: req.RemoteUserId,
  68. Source: common.UserSourceSynced,
  69. Group: group,
  70. AffCode: affCode,
  71. }
  72. }
  73. // ReceiveQuotaUpdate 接收 master 节点的余额更新请求
  74. func ReceiveQuotaUpdate(c *gin.Context) {
  75. var req region_sync.UpdateQuotaRequest
  76. if err := c.ShouldBindJSON(&req); err != nil {
  77. syncErrorResponse(c, http.StatusBadRequest, "invalid request body")
  78. return
  79. }
  80. logger.LogDebug(c.Request.Context(), "[RegionSync] ReceiveQuotaUpdate: remoteUserId=%d, newQuota=%d", req.RemoteUserId, req.Quota)
  81. // 查找同步用户
  82. var user model.User
  83. if err := model.DB.Where("remote_user_id = ? AND source = ?", req.RemoteUserId, common.UserSourceSynced).First(&user).Error; err != nil {
  84. logger.LogError(c.Request.Context(), "synced user not found: "+err.Error())
  85. syncErrorResponse(c, http.StatusNotFound, "user not found")
  86. return
  87. }
  88. // 更新余额(选择性更新,避免覆盖其他字段)
  89. if err := model.DB.Model(&user).Updates(map[string]interface{}{
  90. "synced_quota": req.Quota,
  91. "last_sync_at": time.Now().Unix(),
  92. }).Error; err != nil {
  93. logger.LogError(c.Request.Context(), "failed to update synced quota: "+err.Error())
  94. syncErrorResponse(c, http.StatusInternalServerError, "failed to update quota")
  95. return
  96. }
  97. // 记录同步日志
  98. recordSyncLog(c, user.Id, user.RemoteUserId, model.SyncTypeQuotaChange, "", 0, req.Quota, model.SyncDirectionCnToOv)
  99. logger.LogInfo(c.Request.Context(), "updated synced quota for user "+user.Username)
  100. syncSuccessResponse(c, "quota updated successfully")
  101. }
  102. // QueryUserQuota 查询用户余额(供 slave 节点调用 master 节点)
  103. func QueryUserQuota(c *gin.Context) {
  104. var req region_sync.QueryQuotaRequest
  105. if err := c.ShouldBindJSON(&req); err != nil {
  106. syncErrorResponse(c, http.StatusBadRequest, "invalid request body")
  107. return
  108. }
  109. // 查找用户
  110. var user model.User
  111. if err := model.DB.Where("id = ?", req.UserId).First(&user).Error; err != nil {
  112. syncErrorResponse(c, http.StatusNotFound, "user not found")
  113. return
  114. }
  115. logger.LogDebug(c.Request.Context(), "[RegionSync] QueryUserQuota: userId=%d, quota=%d", user.Id, user.Quota)
  116. // 记录同步日志
  117. recordSyncLog(c, user.Id, user.RemoteUserId, model.SyncTypePreConsumeQuery, "", user.Quota, 0, model.SyncDirectionOvToCn)
  118. c.JSON(http.StatusOK, gin.H{"success": true, "quota": user.Quota})
  119. }
  120. // BatchDeductQuota 批量扣费(供 slave 节点调用 master 节点)
  121. func BatchDeductQuota(c *gin.Context) {
  122. var req region_sync.BatchDeductRequest
  123. if err := c.ShouldBindJSON(&req); err != nil {
  124. syncErrorResponse(c, http.StatusBadRequest, "invalid request body")
  125. return
  126. }
  127. logger.LogDebug(c.Request.Context(), "[RegionSync] BatchDeductQuota: received %d records", len(req.Records))
  128. results := make([]region_sync.DeductResult, 0, len(req.Records))
  129. // 使用事务保证批量扣费的原子性(含幂等性检查和日志记录)
  130. err := model.DB.Transaction(func(tx *gorm.DB) error {
  131. for _, record := range req.Records {
  132. result := processDeductRecordWithTx(c, tx, record)
  133. results = append(results, result)
  134. // 在事务内记录同步日志,确保幂等性检查能查到
  135. if result.Success {
  136. syncLog := &model.QuotaSyncLog{
  137. UserId: record.UserId,
  138. SyncType: model.SyncTypeBatchSync,
  139. Direction: model.SyncDirectionOvToCn,
  140. Status: model.SyncStatusSuccess,
  141. RequestId: record.RequestId,
  142. ChangeAmount: record.Quota,
  143. AfterQuota: result.RemainingQuota,
  144. CreatedAt: time.Now().Unix(),
  145. }
  146. if err := tx.Create(syncLog).Error; err != nil {
  147. return err
  148. }
  149. }
  150. }
  151. return nil
  152. })
  153. if err != nil {
  154. logger.LogError(c.Request.Context(), "batch deduct transaction failed: "+err.Error())
  155. syncErrorResponse(c, http.StatusInternalServerError, "batch deduct failed")
  156. return
  157. }
  158. successCount := 0
  159. for _, r := range results {
  160. if r.Success {
  161. successCount++
  162. }
  163. }
  164. logger.LogDebug(c.Request.Context(), "[RegionSync] BatchDeductQuota: completed, success=%d, total=%d", successCount, len(req.Records))
  165. c.JSON(http.StatusOK, gin.H{"success": true, "results": results})
  166. }
  167. // processDeductRecordWithTx 在事务中处理单条扣费记录
  168. func processDeductRecordWithTx(c *gin.Context, tx *gorm.DB, record region_sync.BatchDeductRecord) region_sync.DeductResult {
  169. result := region_sync.DeductResult{UserId: record.UserId}
  170. logger.LogDebug(c.Request.Context(), "[RegionSync] processDeductRecord: userId=%d, requestId=%s, quota=%d", record.UserId, record.RequestId, record.Quota)
  171. // 幂等性检查:通过 quota_sync_logs 检查 request_id 是否已在 master 端处理过
  172. var existingLog model.QuotaSyncLog
  173. if err := tx.Where("request_id = ? AND sync_type = ? AND status = ?",
  174. record.RequestId, model.SyncTypeBatchSync, model.SyncStatusSuccess).First(&existingLog).Error; err == nil {
  175. // 该 request_id 已处理过,直接返回成功(不重复扣费)
  176. result.Success = true
  177. result.Message = "already processed"
  178. result.DeductedQuota = existingLog.ChangeAmount
  179. // 查询当前余额
  180. var user model.User
  181. if err := tx.Where("id = ?", record.UserId).First(&user).Error; err == nil {
  182. result.RemainingQuota = user.Quota
  183. }
  184. logger.LogDebug(c.Request.Context(), "[RegionSync] processDeductRecord: idempotent hit, userId=%d, requestId=%s, deductedQuota=%d", record.UserId, record.RequestId, result.DeductedQuota)
  185. return result
  186. }
  187. // 原子扣费:使用 SQL 条件更新,确保余额充足时才扣减,避免并发覆盖
  188. dbResult := tx.Model(&model.User{}).
  189. Where("id = ? AND quota >= ?", record.UserId, record.Quota).
  190. Update("quota", gorm.Expr("quota - ?", record.Quota))
  191. if dbResult.Error != nil {
  192. result.Error = "failed to deduct quota"
  193. return result
  194. }
  195. if dbResult.RowsAffected == 0 {
  196. // 没有匹配行,可能是用户不存在或余额不足
  197. var user model.User
  198. if err := tx.Where("id = ?", record.UserId).First(&user).Error; err != nil {
  199. result.Error = "user not found"
  200. } else {
  201. result.Error = "insufficient quota"
  202. result.RemainingQuota = user.Quota
  203. logger.LogDebug(c.Request.Context(), "[RegionSync] processDeductRecord: insufficient quota, userId=%d, quota=%d, need=%d", record.UserId, user.Quota, record.Quota)
  204. }
  205. return result
  206. }
  207. // 查询扣费后的余额(同一事务内可读到自己的更新,MySQL REPEATABLE READ / SQLite 均安全)
  208. var user model.User
  209. if err := tx.Where("id = ?", record.UserId).First(&user).Error; err == nil {
  210. result.RemainingQuota = user.Quota
  211. }
  212. result.Success = true
  213. result.DeductedQuota = record.Quota
  214. logger.LogDebug(c.Request.Context(), "[RegionSync] processDeductRecord: success, userId=%d, deducted=%d, remaining=%d", record.UserId, record.Quota, result.RemainingQuota)
  215. return result
  216. }
  217. // recordSyncLog 记录同步日志
  218. func recordSyncLog(c *gin.Context, userId, remoteUserId int, syncType, requestId string, masterQuota, afterQuota int, direction string) {
  219. syncLog := &model.QuotaSyncLog{
  220. UserId: userId,
  221. RemoteUserId: remoteUserId,
  222. SyncType: syncType,
  223. Direction: direction,
  224. Status: model.SyncStatusSuccess,
  225. RequestId: requestId,
  226. MasterQuota: masterQuota,
  227. AfterQuota: afterQuota,
  228. CreatedAt: time.Now().Unix(),
  229. }
  230. if err := model.CreateSyncLog(syncLog); err != nil {
  231. logger.LogError(c.Request.Context(), "failed to create sync log: "+err.Error())
  232. }
  233. }