-
-
Notifications
You must be signed in to change notification settings - Fork 2k
feat: 渠道级别限制每分钟最大请求数 #1711
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: alpha
Are you sure you want to change the base?
feat: 渠道级别限制每分钟最大请求数 #1711
Conversation
WalkthroughAdds per-channel rate limiting across stack: new Channel.RateLimit field and getter; in-memory and Redis-based limit checks integrated into channel selection; new context key for channel rate limit set in middleware and read into RelayInfo; UI adds rate_limit input on channel edit. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Client
participant API as API Gateway
participant MW as Middleware (distributor)
participant Sel as Channel Selector
participant Ch as Channel
participant Relay as Relay Layer
Client->>API: Request
API->>MW: SetupContextForSelectedChannel
MW->>Sel: Select channel
Sel->>Sel: Apply weight + rate-limit checks (in-mem/Redis)
Sel-->>MW: Selected Channel or not found
MW->>MW: Set context keys (status map, channel_rate_limit, next key)
MW-->>API: Context prepared
API->>Relay: GenRelayInfo (reads channel_rate_limit)
Relay-->>API: RelayInfo
API-->>Client: Response
sequenceDiagram
autonumber
participant Sel as Selector
participant RL1 as In-memory RL
participant RL2 as Redis RL
Sel->>Sel: Build candidate channels
alt Cache path uses Redis
Sel->>RL2: isRedisLimited(channelId)
RL2-->>Sel: allow/deny
else Ability path uses in-memory
Sel->>RL1: isRateLimited(channelId)
RL1-->>Sel: allow/deny
end
Sel->>Sel: Weighted pick among allowed
Sel-->>Sel: Return channel or "not found"
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related issues
Possibly related PRs
Poem
✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
model/channel_cache.go (1)
3-17
: Imports: drop unused redis/v8 and strconv; add context for INCR/EXPIRESwitching to atomic Redis ops removes the need for redis.Nil/strconv; add context for RDB calls.
import ( + "context" "errors" "fmt" "math/rand" "one-api/common" "one-api/setting" "sort" - "strconv" "strings" "sync" "time" "github.com/gin-gonic/gin" - "github.com/go-redis/redis/v8" )
🧹 Nitpick comments (9)
constant/context_key.go (1)
34-34
: Add a brief unit/semantics comment for the new keyClarify that 0 means “unlimited” and the unit is requests per minute.
Apply this diff:
- ContextKeyChannelRateLimit ContextKey = "channel_rate_limit" + // per-minute max requests for the selected channel; 0 means unlimited + ContextKeyChannelRateLimit ContextKey = "channel_rate_limit"relay/common/relay_info.go (1)
65-65
: Document the new field’s meaningAdd a short comment to avoid ambiguity downstream (unit + 0 semantics).
Apply this diff:
- ChannelRateLimit int + ChannelRateLimit int // per-minute cap; 0 => unlimitedmodel/channel.go (1)
401-406
: Getter behavior is correct; consider documenting 0 => unlimitedDefensive nil/<=0 handling is good. Add a comment for callers.
Apply this diff:
-func (channel *Channel) GetRateLimit() int { +// GetRateLimit returns the per-minute cap; 0 means unlimited. +func (channel *Channel) GetRateLimit() int {web/src/pages/Channel/EditChannel.js (1)
1375-1384
: Add user guidance: 0 means unlimited; coerce undefined/null to 0 on submitImprove UX and avoid sending nulls.
Apply this diff to annotate the control:
<Form.InputNumber field='rate_limit' label={t('每分钟最大请求数')} placeholder={t('每分钟最大请求数')} min={0} onNumberChange={(value) => handleInputChange('rate_limit', value)} style={{ width: '100%' }} + extraText={t('0 表示不限制')} />
And normalize before submit (outside this hunk; place near other normalizations in submit()):
// ensure non-negative integer; 0 => unlimited localInputs.rate_limit = Number.isFinite(localInputs.rate_limit) && localInputs.rate_limit > 0 ? Math.trunc(localInputs.rate_limit) : 0;model/channel_cache.go (2)
137-146
: Avoid Redis I/O while holding RLockYou’re calling Redis inside the read lock. Snapshot channel pointers under RLock, then unlock and run isRedisLimited outside to reduce contention and writer starvation.
200-207
: Minor: isRedisLimited naming/return semanticsisRedisLimited returns true when limited while checkRedisLimit returns true when allowed. Consider aligning names or invert check to avoid double negations in callers.
model/ability.go (3)
158-179
: Remove now-unused updateRateLimitStatus to avoid confusionWith single-writer semantics in checkRateLimit, this helper is redundant.
-func updateRateLimitStatus(channelId int) { - now := time.Now() - key := ChannelModelKey{ChannelID: channelId} - - rateLimitMutex.Lock() - defer rateLimitMutex.Unlock() - - val, _ := channelRateLimitStatus.Load(key) - if val == nil { - return - } - - rl := val.(*ChannelRateLimit) - if now.After(rl.ResetTime) { - rl.Count = 1 - rl.ResetTime = now.Add(time.Minute) - } else { - rl.Count++ - } - - channelRateLimitStatus.Store(key, rl) -}
106-117
: Consider per-channel mutex or atomic countersGlobal rateLimitMutex serializes all channels; per-channel locks or atomic.Int64 per entry would reduce contention at scale.
181-228
: Align semantics with Redis limiterAfter fixing Redis to fixed-window (TTL set once), this in-memory limiter matches. Document “fixed 60-second window” to avoid ambiguity with sliding windows.
Would you prefer fixed windows (current) or sliding windows? If sliding is desired, we should change both implementations accordingly.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (7)
constant/context_key.go
(1 hunks)middleware/distributor.go
(1 hunks)model/ability.go
(3 hunks)model/channel.go
(2 hunks)model/channel_cache.go
(4 hunks)relay/common/relay_info.go
(3 hunks)web/src/pages/Channel/EditChannel.js
(2 hunks)
🧰 Additional context used
🧬 Code graph analysis (5)
middleware/distributor.go (2)
common/gin.go (1)
SetContextKey
(49-51)constant/context_key.go (1)
ContextKeyChannelRateLimit
(34-34)
web/src/pages/Channel/EditChannel.js (2)
web/src/pages/Detail/index.js (1)
handleInputChange
(624-630)web/src/pages/Channel/EditTagModal.js (1)
handleInputChange
(57-109)
relay/common/relay_info.go (3)
model/ability.go (1)
ChannelRateLimit
(109-112)common/gin.go (1)
GetContextKeyInt
(61-63)constant/context_key.go (1)
ContextKeyChannelRateLimit
(34-34)
model/channel_cache.go (3)
model/channel.go (1)
Channel
(19-50)common/redis.go (2)
RedisGet
(72-79)RedisSet
(64-70)common/logger.go (1)
SysLog
(50-53)
model/ability.go (2)
model/channel.go (2)
Channel
(19-50)GetChannelById
(314-329)common/utils.go (1)
GetRandomInt
(187-190)
🔇 Additional comments (9)
relay/common/relay_info.go (1)
219-241
: LGTM: context propagation into RelayInfoReads the int from context and plumbs it into RelayInfo correctly; default 0 fallback via GetInt is appropriate.
middleware/distributor.go (1)
270-271
: LGTM: set channel rate limit in contextPlacement is fine and keeps semantics consistent with other channel keys.
web/src/pages/Channel/EditChannel.js (1)
105-105
: LGTM: default field addedDefaulting rate_limit to 0 matches backend semantics (unlimited).
model/channel_cache.go (3)
148-154
: Early return on single valid channel is goodThis avoids unnecessary weighting work and extra Redis calls.
157-159
: Priorities derived from filtered validChannelsCorrect to compute unique priorities post rate-limit filtering.
173-176
: Target channel selection uses filtered setUsing validChannels when building targetChannels is correct and prevents picking rate-limited or missing channels.
model/ability.go (3)
129-156
: checkRateLimit logic OK after above fixInitial Count=1 and boundary uses “>” ensuring exactly RateLimit requests per minute.
230-249
: Weighted selection helper looks goodHandles zero-sum fallback and weighted choice properly.
197-201
: Graceful empty-abilities handlingReturning “channel not found” is consistent with callers.
func isRateLimited(channel Channel, channelId int) bool { | ||
if (channel.RateLimit != nil && *channel.RateLimit > 0) { | ||
if _, ok := checkRateLimit(&channel, channelId); !ok { | ||
return true | ||
} | ||
updateRateLimitStatus(channelId) | ||
} | ||
return false | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🛠️ Refactor suggestion
Double-count bug: remove updateRateLimitStatus after successful check
checkRateLimit already increments Count. Calling updateRateLimitStatus increments again, reducing the effective limit.
func isRateLimited(channel Channel, channelId int) bool {
if (channel.RateLimit != nil && *channel.RateLimit > 0) {
if _, ok := checkRateLimit(&channel, channelId); !ok {
return true
}
- updateRateLimitStatus(channelId)
}
return false
}
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
func isRateLimited(channel Channel, channelId int) bool { | |
if (channel.RateLimit != nil && *channel.RateLimit > 0) { | |
if _, ok := checkRateLimit(&channel, channelId); !ok { | |
return true | |
} | |
updateRateLimitStatus(channelId) | |
} | |
return false | |
} | |
func isRateLimited(channel Channel, channelId int) bool { | |
if (channel.RateLimit != nil && *channel.RateLimit > 0) { | |
if _, ok := checkRateLimit(&channel, channelId); !ok { | |
return true | |
} | |
} | |
return false | |
} |
🤖 Prompt for AI Agents
In model/ability.go around lines 118 to 126, the function is double-incrementing
the rate counter because checkRateLimit already increments Count; remove the
post-check call to updateRateLimitStatus(channelId) so that on a successful
check the count is not incremented again, leaving updateRateLimitStatus only
where increments are actually needed (e.g., on explicit failures or outside this
successful path).
channelPtr, err := GetChannelById(selectedAbility.ChannelId, true) | ||
if err != nil { | ||
if err.Error() != "channel not found" { | ||
return nil, err | ||
} | ||
abilities = removeAbility(abilities, selectedIndex) | ||
continue | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🛠️ Refactor suggestion
Use errors.Is for gorm.ErrRecordNotFound
String compare on error is brittle; treat only “record not found” as skippable.
- channelPtr, err := GetChannelById(selectedAbility.ChannelId, true)
- if err != nil {
- if err.Error() != "channel not found" {
- return nil, err
- }
- abilities = removeAbility(abilities, selectedIndex)
- continue
- }
+ channelPtr, err := GetChannelById(selectedAbility.ChannelId, true)
+ if err != nil {
+ if !errors.Is(err, gorm.ErrRecordNotFound) {
+ return nil, err
+ }
+ abilities = removeAbility(abilities, selectedIndex)
+ continue
+ }
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
channelPtr, err := GetChannelById(selectedAbility.ChannelId, true) | |
if err != nil { | |
if err.Error() != "channel not found" { | |
return nil, err | |
} | |
abilities = removeAbility(abilities, selectedIndex) | |
continue | |
} | |
channelPtr, err := GetChannelById(selectedAbility.ChannelId, true) | |
if err != nil { | |
if !errors.Is(err, gorm.ErrRecordNotFound) { | |
return nil, err | |
} | |
abilities = removeAbility(abilities, selectedIndex) | |
continue | |
} |
🤖 Prompt for AI Agents
In model/ability.go around lines 209 to 216, the code compares err.Error() to a
string to detect a missing channel which is brittle; replace that string
comparison with errors.Is(err, gorm.ErrRecordNotFound) to correctly detect a
not-found error even if wrapped, return the error for other cases, and when the
record is not found remove the ability and continue; also ensure the file
imports the "errors" package and "gorm.io/gorm" (or references the project’s
gorm error variable) so the errors.Is check compiles.
func isRedisLimited(channel Channel, channelId int) bool { | ||
if channel.RateLimit != nil && *channel.RateLimit > 0 { | ||
if !checkRedisLimit(channel, channelId) { | ||
return true | ||
} | ||
} | ||
return false | ||
} | ||
|
||
func checkRedisLimit(channel Channel, channelId int) bool { | ||
key := fmt.Sprintf("rate_limit:%d", channelId) | ||
|
||
countStr, err := common.RedisGet(key) | ||
if err == redis.Nil { | ||
// Key doesn't exist, set it with expiration | ||
err = common.RedisSet(key, "1", time.Minute) | ||
if err != nil { | ||
common.SysLog(fmt.Sprintf("Error setting rate limit: %v", err)) | ||
return false | ||
} | ||
return true | ||
} else if err != nil { | ||
common.SysLog(fmt.Sprintf("Error checking rate limit: %v", err)) | ||
return false | ||
} | ||
|
||
count, err := strconv.ParseInt(countStr, 10, 64) | ||
if err != nil { | ||
common.SysLog(fmt.Sprintf("Error parsing rate limit count: %v", err)) | ||
return false | ||
} | ||
|
||
if count > int64(*channel.RateLimit) { | ||
return false | ||
} | ||
|
||
// 增加计数 | ||
newCount := strconv.FormatInt(count+1, 10) | ||
err = common.RedisSet(key, newCount, time.Minute) | ||
if err != nil { | ||
common.SysLog(fmt.Sprintf("Error incrementing rate limit: %v", err)) | ||
return false | ||
} | ||
|
||
return true | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🛠️ Refactor suggestion
Fix Redis rate limiter: use atomic INCR + one-time TTL set; correct boundary check
Current GET/SET is racy and allows limit+1 per window; also resets TTL on every hit (sliding), diverging from in-memory (fixed). Use INCR and set EXPIRE only when creating the key; compare with <= limit.
-func checkRedisLimit(channel Channel, channelId int) bool {
- key := fmt.Sprintf("rate_limit:%d", channelId)
-
- countStr, err := common.RedisGet(key)
- if err == redis.Nil {
- // Key doesn't exist, set it with expiration
- err = common.RedisSet(key, "1", time.Minute)
- if err != nil {
- common.SysLog(fmt.Sprintf("Error setting rate limit: %v", err))
- return false
- }
- return true
- } else if err != nil {
- common.SysLog(fmt.Sprintf("Error checking rate limit: %v", err))
- return false
- }
-
- count, err := strconv.ParseInt(countStr, 10, 64)
- if err != nil {
- common.SysLog(fmt.Sprintf("Error parsing rate limit count: %v", err))
- return false
- }
-
- if count > int64(*channel.RateLimit) {
- return false
- }
-
- // 增加计数
- newCount := strconv.FormatInt(count+1, 10)
- err = common.RedisSet(key, newCount, time.Minute)
- if err != nil {
- common.SysLog(fmt.Sprintf("Error incrementing rate limit: %v", err))
- return false
- }
-
- return true
-}
+func checkRedisLimit(channel Channel, channelId int) bool {
+ key := fmt.Sprintf("rate_limit:%d", channelId)
+ ctx := context.Background()
+
+ // Atomic increment
+ count, err := common.RDB.Incr(ctx, key).Result()
+ if err != nil {
+ common.SysLog(fmt.Sprintf("Error incrementing rate limit: %v", err))
+ return false
+ }
+ // First hit: set TTL once to create a fixed 1-minute window
+ if count == 1 {
+ if err := common.RDB.Expire(ctx, key, time.Minute).Err(); err != nil {
+ common.SysLog(fmt.Sprintf("Error setting rate limit TTL: %v", err))
+ return false
+ }
+ }
+ // Allow if within limit
+ return count <= int64(*channel.RateLimit)
+}
RateLimit *int `json:"rate_limit" gorm:"default:0"` | ||
// add after v0.8.5 | ||
ChannelInfo ChannelInfo `json:"channel_info" gorm:"type:json"` | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🛠️ Refactor suggestion
Add DB migration and constraints for rate_limit
Schema change requires migration. Backfill nulls to 0 and add a non-negative constraint to prevent invalid values.
Apply one of these migrations (adjust table name if needed):
- MySQL:
ALTER TABLE channels ADD COLUMN rate_limit INT NOT NULL DEFAULT 0;
ALTER TABLE channels ADD CONSTRAINT chk_rate_limit_nonneg CHECK (rate_limit >= 0);
- PostgreSQL:
ALTER TABLE channels ADD COLUMN rate_limit INT NOT NULL DEFAULT 0;
ALTER TABLE channels ADD CONSTRAINT chk_rate_limit_nonneg CHECK (rate_limit >= 0);
- SQLite:
ALTER TABLE channels ADD COLUMN rate_limit INTEGER DEFAULT 0;
UPDATE channels SET rate_limit = 0 WHERE rate_limit IS NULL;
-- (SQLite lacks ADD CONSTRAINT; enforce non-negativity in app code)
Also consider returning 0 instead of null in read APIs to simplify clients.
🤖 Prompt for AI Agents
In model/channel.go around lines 47 to 50, the new RateLimit field introduces a
schema change that needs a DB migration: add the rate_limit column with NOT NULL
DEFAULT 0, backfill existing NULLs to 0, and enforce non-negativity with a CHECK
constraint where supported (MySQL/Postgres); for SQLite add the column, run an
UPDATE to set NULLs to 0 and enforce non-negativity in application code. Create
and run the appropriate migration for your database (adjust table name if
needed), and update read APIs to return 0 instead of null to simplify clients.
PR 类型
PR 是否包含破坏性更新?
PR 描述
在渠道的高级设置中添加“每分钟最大请求数”设置
感觉报 429 合适一些,但没找到哪里调,现在会报“分组xxx没有可以用xxx模型的渠道”
Summary by CodeRabbit
New Features
Chores