Ви не можете вибрати більше 25 тем Теми мають розпочинатися з літери або цифри, можуть містити дефіси (-) і не повинні перевищувати 35 символів.
 
 
 

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