diff --git a/common/constants.go b/common/constants.go index 274c514f9146..69a746e88a5c 100644 --- a/common/constants.go +++ b/common/constants.go @@ -173,7 +173,7 @@ var ( CriticalRateLimitEnable bool CriticalRateLimitNum = 20 - CriticalRateLimitDuration int64 = 20 * 60 + CriticalRateLimitDuration int64 = 60 UploadRateLimitNum = 10 UploadRateLimitDuration int64 = 60 diff --git a/common/init.go b/common/init.go index 35b4c6be17ee..107301832678 100644 --- a/common/init.go +++ b/common/init.go @@ -120,7 +120,7 @@ func InitEnv() { CriticalRateLimitEnable = GetEnvOrDefaultBool("CRITICAL_RATE_LIMIT_ENABLE", true) CriticalRateLimitNum = GetEnvOrDefault("CRITICAL_RATE_LIMIT", 20) - CriticalRateLimitDuration = int64(GetEnvOrDefault("CRITICAL_RATE_LIMIT_DURATION", 20*60)) + CriticalRateLimitDuration = int64(GetEnvOrDefault("CRITICAL_RATE_LIMIT_DURATION", 60)) SearchRateLimitEnable = GetEnvOrDefaultBool("SEARCH_RATE_LIMIT_ENABLE", true) SearchRateLimitNum = GetEnvOrDefault("SEARCH_RATE_LIMIT", 10) diff --git a/middleware/rate-limit.go b/middleware/rate-limit.go index d8dd15d9c5d7..9c6ab6728eab 100644 --- a/middleware/rate-limit.go +++ b/middleware/rate-limit.go @@ -18,6 +18,14 @@ var defNext = func(c *gin.Context) { c.Next() } +func abortRateLimited(c *gin.Context) { + c.JSON(http.StatusTooManyRequests, gin.H{ + "success": false, + "message": "请求过于频繁,请稍后再试", + }) + c.Abort() +} + func redisRateLimiter(c *gin.Context, maxRequestNum int, duration int64, mark string) { ctx := context.Background() rdb := common.RDB @@ -53,8 +61,7 @@ func redisRateLimiter(c *gin.Context, maxRequestNum int, duration int64, mark st // See: https://stackoverflow.com/questions/50970900/why-is-time-since-returning-negative-durations-on-windows if int64(nowTime.Sub(oldTime).Seconds()) < duration { rdb.Expire(ctx, key, common.RateLimitKeyExpirationDuration) - c.Status(http.StatusTooManyRequests) - c.Abort() + abortRateLimited(c) return } else { rdb.LPush(ctx, key, time.Now().Format(timeFormat)) @@ -67,8 +74,7 @@ func redisRateLimiter(c *gin.Context, maxRequestNum int, duration int64, mark st func memoryRateLimiter(c *gin.Context, maxRequestNum int, duration int64, mark string) { key := mark + c.ClientIP() if !inMemoryRateLimiter.Request(key, maxRequestNum, duration) { - c.Status(http.StatusTooManyRequests) - c.Abort() + abortRateLimited(c) return } } @@ -143,8 +149,7 @@ func userRateLimitFactory(maxRequestNum int, duration int64, mark string) func(c } key := fmt.Sprintf("%s:user:%d", mark, userId) if !inMemoryRateLimiter.Request(key, maxRequestNum, duration) { - c.Status(http.StatusTooManyRequests) - c.Abort() + abortRateLimited(c) return } } @@ -184,8 +189,7 @@ func userRedisRateLimiter(c *gin.Context, maxRequestNum int, duration int64, key } if int64(nowTime.Sub(oldTime).Seconds()) < duration { rdb.Expire(ctx, key, common.RateLimitKeyExpirationDuration) - c.Status(http.StatusTooManyRequests) - c.Abort() + abortRateLimited(c) return } else { rdb.LPush(ctx, key, time.Now().Format(timeFormat)) diff --git a/model/log.go b/model/log.go index 3cc9b74b24b5..43ac690f24b2 100644 --- a/model/log.go +++ b/model/log.go @@ -37,7 +37,8 @@ type Log struct { Ip string `json:"ip" gorm:"index;default:''"` RequestId string `json:"request_id,omitempty" gorm:"type:varchar(64);index:idx_logs_request_id;default:''"` Other string `json:"other"` - Balance int `json:"balance" gorm:"-"` + Balance int `json:"balance" gorm:"default:0"` + BalanceValid bool `json:"-" gorm:"default:false"` } // don't use iota, avoid change log type value @@ -87,7 +88,7 @@ func AttachLogBalances(logs []*Log) { } balances := make(map[int]int) for _, log := range logs { - if log == nil || log.UserId == 0 { + if log == nil || log.UserId == 0 || log.BalanceValid { continue } if _, ok := balances[log.UserId]; ok { @@ -105,7 +106,7 @@ func AttachLogBalances(logs []*Log) { balances[log.UserId] = quota - newerDelta } for _, log := range logs { - if log == nil || log.UserId == 0 { + if log == nil || log.UserId == 0 || log.BalanceValid { continue } balance, ok := balances[log.UserId] @@ -117,6 +118,20 @@ func AttachLogBalances(logs []*Log) { } } +func attachCurrentBalanceSnapshot(log *Log) { + if log == nil || log.UserId == 0 { + return + } + quota, err := GetUserQuota(log.UserId, true) + if err != nil { + common.SysLog(fmt.Sprintf("failed to attach log balance snapshot for user %d: %v", log.UserId, err)) + return + } + quota += pendingBatchUpdateValue(BatchUpdateTypeUserQuota, log.UserId) + log.Balance = quota + log.BalanceValid = true +} + func GetLogByTokenId(tokenId int) (logs []*Log, err error) { err = LOG_DB.Model(&Log{}).Where("token_id = ?", tokenId).Order("id desc").Limit(common.MaxRecentItems).Find(&logs).Error formatUserLogs(logs, 0) @@ -135,6 +150,7 @@ func RecordLog(userId int, logType int, content string) { Type: logType, Content: content, } + attachCurrentBalanceSnapshot(log) err := LOG_DB.Create(log).Error if err != nil { common.SysLog("failed to record log: " + err.Error()) @@ -154,6 +170,7 @@ func RecordLogWithAdminInfo(userId int, logType int, content string, adminInfo m Type: logType, Content: content, } + attachCurrentBalanceSnapshot(log) if len(adminInfo) > 0 { other := map[string]interface{}{ "admin_info": adminInfo, @@ -187,6 +204,7 @@ func RecordTopupLog(userId int, content string, callerIp string, paymentMethod s Ip: callerIp, Other: common.MapToJsonStr(other), } + attachCurrentBalanceSnapshot(log) err := LOG_DB.Create(log).Error if err != nil { common.SysLog("failed to record topup log: " + err.Error()) @@ -231,6 +249,7 @@ func RecordErrorLog(c *gin.Context, userId int, channelId int, modelName string, RequestId: requestId, Other: otherStr, } + attachCurrentBalanceSnapshot(log) err := LOG_DB.Create(log).Error if err != nil { logger.LogError(c, "failed to record log: "+err.Error()) @@ -292,6 +311,7 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams) RequestId: requestId, Other: otherStr, } + attachCurrentBalanceSnapshot(log) err := LOG_DB.Create(log).Error if err != nil { logger.LogError(c, "failed to record log: "+err.Error()) @@ -340,6 +360,7 @@ func RecordTaskBillingLog(params RecordTaskBillingLogParams) { Group: params.Group, Other: common.MapToJsonStr(params.Other), } + attachCurrentBalanceSnapshot(log) err := LOG_DB.Create(log).Error if err != nil { common.SysLog("failed to record task billing log: " + err.Error()) diff --git a/model/utils.go b/model/utils.go index adfd8e139a05..ca53f47f4ed6 100644 --- a/model/utils.go +++ b/model/utils.go @@ -49,6 +49,15 @@ func addNewRecord(type_ int, id int, value int) { } } +func pendingBatchUpdateValue(type_ int, id int) int { + if !common.BatchUpdateEnabled { + return 0 + } + batchUpdateLocks[type_].Lock() + defer batchUpdateLocks[type_].Unlock() + return batchUpdateStores[type_][id] +} + func batchUpdate() { // check if there's any data to update hasData := false diff --git a/service/billing_session.go b/service/billing_session.go index 7ca158bd5255..4a5119299d22 100644 --- a/service/billing_session.go +++ b/service/billing_session.go @@ -119,10 +119,8 @@ func (s *BillingSession) Refund(c *gin.Context) { common.SysLog("error refunding token quota: " + err.Error()) } } - // 3) 退还用户已用额度(tokenConsumed已经是折扣后的值) - if tokenConsumed > 0 { - model.UpdateUserUsedQuotaAndRequestCount(s.relayInfo.UserId, -tokenConsumed) - } + // 提交阶段失败不会写消费日志,也不会增加 users.used_quota。 + // 这里只回滚预扣的资金和令牌额度;异步任务已生成消费日志后的失败退款由 RefundTaskQuota 回退 used_quota。 }) } diff --git a/service/task_billing.go b/service/task_billing.go index 54a1f3fc4266..aa3c54c6e301 100644 --- a/service/task_billing.go +++ b/service/task_billing.go @@ -165,9 +165,9 @@ func RefundTaskQuota(ctx context.Context, task *model.Task, reason string) { // 2. 退还令牌额度 taskAdjustTokenQuota(ctx, task, -quota) - //3. 更新用户和渠道的使用额度 - //model.UpdateUserUsedQuotaAndRequestCount(task.UserId, -quota) - //model.UpdateChannelUsedQuota(task.ChannelId, -quota) + // 3. 更新用户和渠道的使用额度 + model.UpdateUserUsedQuotaAndRequestCount(task.UserId, -quota) + model.UpdateChannelUsedQuota(task.ChannelId, -quota) // 4. 记录日志 other := taskBillingOther(task) @@ -255,8 +255,8 @@ func RecalculateTaskQuota(ctx context.Context, task *model.Task, actualQuota int } else { logType = model.LogTypeRefund logQuota = -quotaDelta - //model.UpdateUserUsedQuotaAndRequestCount(task.UserId, -quotaDelta) - //model.UpdateChannelUsedQuota(task.ChannelId, -quotaDelta) + model.UpdateUserUsedQuotaAndRequestCount(task.UserId, quotaDelta) + model.UpdateChannelUsedQuota(task.ChannelId, quotaDelta) } other := taskBillingOther(task) other["task_id"] = task.TaskID