You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

323 lines
11 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. // BatchQueryUserQuota 批量查询用户余额(供 slave 节点调用)
  122. func BatchQueryUserQuota(c *gin.Context) {
  123. var req region_sync.BatchQueryQuotaRequest
  124. if err := c.ShouldBindJSON(&req); err != nil {
  125. syncErrorResponse(c, http.StatusBadRequest, "invalid request body")
  126. return
  127. }
  128. if len(req.UserIds) == 0 {
  129. c.JSON(http.StatusOK, region_sync.BatchQueryQuotaResponse{Success: true, Quotas: []region_sync.QuotaEntry{}})
  130. return
  131. }
  132. if len(req.UserIds) > 500 {
  133. syncErrorResponse(c, http.StatusBadRequest, "too many user ids (max 500)")
  134. return
  135. }
  136. var users []model.User
  137. if err := model.DB.Where("id IN ?", req.UserIds).Find(&users).Error; err != nil {
  138. syncErrorResponse(c, http.StatusInternalServerError, "database error")
  139. return
  140. }
  141. quotas := make([]region_sync.QuotaEntry, 0, len(users))
  142. for _, user := range users {
  143. quotas = append(quotas, region_sync.QuotaEntry{UserId: user.Id, Quota: user.Quota})
  144. }
  145. c.JSON(http.StatusOK, region_sync.BatchQueryQuotaResponse{
  146. Success: true,
  147. Quotas: quotas,
  148. })
  149. }
  150. // BatchDeductQuota 批量扣费(供 slave 节点调用 master 节点)
  151. func BatchDeductQuota(c *gin.Context) {
  152. var req region_sync.BatchDeductRequest
  153. if err := c.ShouldBindJSON(&req); err != nil {
  154. syncErrorResponse(c, http.StatusBadRequest, "invalid request body")
  155. return
  156. }
  157. logger.LogDebug(c.Request.Context(), "[RegionSync] BatchDeductQuota: received %d records", len(req.Records))
  158. results := make([]region_sync.DeductResult, 0, len(req.Records))
  159. // 使用事务保证批量扣费的原子性(含幂等性检查和日志记录)
  160. err := model.DB.Transaction(func(tx *gorm.DB) error {
  161. for _, record := range req.Records {
  162. result := processDeductRecordWithTx(c, tx, record)
  163. results = append(results, result)
  164. // 在事务内记录同步日志,确保幂等性检查能查到
  165. if result.Success {
  166. syncLog := &model.QuotaSyncLog{
  167. UserId: record.UserId,
  168. SyncType: model.SyncTypeBatchSync,
  169. Direction: model.SyncDirectionOvToCn,
  170. Status: model.SyncStatusSuccess,
  171. RequestId: record.RequestId,
  172. ChangeAmount: record.Quota,
  173. AfterQuota: result.RemainingQuota,
  174. CreatedAt: time.Now().Unix(),
  175. }
  176. if err := tx.Create(syncLog).Error; err != nil {
  177. return err
  178. }
  179. }
  180. }
  181. return nil
  182. })
  183. if err != nil {
  184. logger.LogError(c.Request.Context(), "batch deduct transaction failed: "+err.Error())
  185. syncErrorResponse(c, http.StatusInternalServerError, "batch deduct failed")
  186. return
  187. }
  188. successCount := 0
  189. for _, r := range results {
  190. if r.Success {
  191. successCount++
  192. }
  193. }
  194. logger.LogDebug(c.Request.Context(), "[RegionSync] BatchDeductQuota: completed, success=%d, total=%d", successCount, len(req.Records))
  195. c.JSON(http.StatusOK, gin.H{"success": true, "results": results})
  196. }
  197. // processDeductRecordWithTx 在事务中处理单条扣费记录
  198. func processDeductRecordWithTx(c *gin.Context, tx *gorm.DB, record region_sync.BatchDeductRecord) region_sync.DeductResult {
  199. result := region_sync.DeductResult{UserId: record.UserId}
  200. logger.LogDebug(c.Request.Context(), "[RegionSync] processDeductRecord: userId=%d, requestId=%s, quota=%d", record.UserId, record.RequestId, record.Quota)
  201. // 幂等性检查:通过 quota_sync_logs 检查 request_id 是否已在 master 端处理过
  202. var existingLog model.QuotaSyncLog
  203. if err := tx.Where("request_id = ? AND sync_type = ? AND status = ?",
  204. record.RequestId, model.SyncTypeBatchSync, model.SyncStatusSuccess).First(&existingLog).Error; err == nil {
  205. // 该 request_id 已处理过,直接返回成功(不重复扣费)
  206. result.Success = true
  207. result.Message = "already processed"
  208. result.DeductedQuota = existingLog.ChangeAmount
  209. // 查询当前余额
  210. var user model.User
  211. if err := tx.Where("id = ?", record.UserId).First(&user).Error; err == nil {
  212. result.RemainingQuota = user.Quota
  213. }
  214. logger.LogDebug(c.Request.Context(), "[RegionSync] processDeductRecord: idempotent hit, userId=%d, requestId=%s, deductedQuota=%d", record.UserId, record.RequestId, result.DeductedQuota)
  215. return result
  216. }
  217. // 原子扣费:使用 SQL 条件更新,确保余额充足时才扣减,避免并发覆盖
  218. dbResult := tx.Model(&model.User{}).
  219. Where("id = ? AND quota >= ?", record.UserId, record.Quota).
  220. Update("quota", gorm.Expr("quota - ?", record.Quota))
  221. if dbResult.Error != nil {
  222. result.Error = "failed to deduct quota"
  223. return result
  224. }
  225. if dbResult.RowsAffected == 0 {
  226. // 没有匹配行,可能是用户不存在或余额不足
  227. var user model.User
  228. if err := tx.Where("id = ?", record.UserId).First(&user).Error; err != nil {
  229. result.Error = "user not found"
  230. } else {
  231. result.Error = "insufficient quota"
  232. result.RemainingQuota = user.Quota
  233. logger.LogDebug(c.Request.Context(), "[RegionSync] processDeductRecord: insufficient quota, userId=%d, quota=%d, need=%d", record.UserId, user.Quota, record.Quota)
  234. }
  235. return result
  236. }
  237. // 查询扣费后的余额(同一事务内可读到自己的更新,MySQL REPEATABLE READ / SQLite 均安全)
  238. var user model.User
  239. if err := tx.Where("id = ?", record.UserId).First(&user).Error; err == nil {
  240. result.RemainingQuota = user.Quota
  241. }
  242. result.Success = true
  243. result.DeductedQuota = record.Quota
  244. logger.LogDebug(c.Request.Context(), "[RegionSync] processDeductRecord: success, userId=%d, deducted=%d, remaining=%d", record.UserId, record.Quota, result.RemainingQuota)
  245. return result
  246. }
  247. // recordSyncLog 记录同步日志
  248. func recordSyncLog(c *gin.Context, userId, remoteUserId int, syncType, requestId string, masterQuota, afterQuota int, direction string) {
  249. syncLog := &model.QuotaSyncLog{
  250. UserId: userId,
  251. RemoteUserId: remoteUserId,
  252. SyncType: syncType,
  253. Direction: direction,
  254. Status: model.SyncStatusSuccess,
  255. RequestId: requestId,
  256. MasterQuota: masterQuota,
  257. AfterQuota: afterQuota,
  258. CreatedAt: time.Now().Unix(),
  259. }
  260. if err := model.CreateSyncLog(syncLog); err != nil {
  261. logger.LogError(c.Request.Context(), "failed to create sync log: "+err.Error())
  262. }
  263. }
  264. // GetSyncConfig Master 节点返回同步配置给 Slave
  265. func GetSyncConfig(c *gin.Context) {
  266. settings := system_setting.GetRegionSyncSettings()
  267. c.JSON(http.StatusOK, region_sync.SyncConfigResponse{
  268. Success: true,
  269. MinBalanceThreshold: settings.MinBalanceThreshold,
  270. MaxRetryCount: settings.MaxRetryCount,
  271. SyncBatchSize: settings.SyncBatchSize,
  272. SyncIntervalSeconds: settings.SyncIntervalSeconds,
  273. })
  274. }