内存限流处理器
(duration int64, totalMaxCount, successMaxCount int)
| 130 | |
| 131 | // 内存限流处理器 |
| 132 | func memoryRateLimitHandler(duration int64, totalMaxCount, successMaxCount int) gin.HandlerFunc { |
| 133 | inMemoryRateLimiter.Init(time.Duration(setting.ModelRequestRateLimitDurationMinutes) * time.Minute) |
| 134 | |
| 135 | return func(c *gin.Context) { |
| 136 | userId := strconv.Itoa(c.GetInt("id")) |
| 137 | totalKey := ModelRequestRateLimitCountMark + userId |
| 138 | successKey := ModelRequestRateLimitSuccessCountMark + userId |
| 139 | |
| 140 | // 1. 检查总请求数限制（当totalMaxCount为0时跳过） |
| 141 | if totalMaxCount > 0 && !inMemoryRateLimiter.Request(totalKey, totalMaxCount, duration) { |
| 142 | c.Status(http.StatusTooManyRequests) |
| 143 | c.Abort() |
| 144 | return |
| 145 | } |
| 146 | |
| 147 | // 2. 检查成功请求数限制 |
| 148 | // 使用一个临时key来检查限制，这样可以避免实际记录 |
| 149 | checkKey := successKey + "_check" |
| 150 | if !inMemoryRateLimiter.Request(checkKey, successMaxCount, duration) { |
| 151 | c.Status(http.StatusTooManyRequests) |
| 152 | c.Abort() |
| 153 | return |
| 154 | } |
| 155 | |
| 156 | // 3. 处理请求 |
| 157 | c.Next() |
| 158 | |
| 159 | // 4. 如果请求成功，记录到实际的成功请求计数中 |
| 160 | if c.Writer.Status() < 400 { |
| 161 | inMemoryRateLimiter.Request(successKey, successMaxCount, duration) |
| 162 | } |
| 163 | } |
| 164 | } |
| 165 | |
| 166 | // ModelRequestRateLimit 模型请求限流中间件 |
| 167 | func ModelRequestRateLimit() func(c *gin.Context) { |
no test coverage detected