jianf/ModelRouter
1
0
Fork 0
mirror of https://gitcode.com/JianFeeeee/ModelRouter.git synced 2026-10-05 23:17:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf(gateway): 拒绝路径只判定一次 + 补配额交互判据

Browse Source 
复查后修掉一个自己引入的缺陷，并补上此前缺失的交叉场景验证。

## 修复：拒绝路径重复判定

4 个入口原本先 checkModelScope（判是否为空）再 writeScopeReject
（内部又 checkQuota 一次）。即每个【被拒】的请求要跑两遍配额统计，
且两次之间用量可能变化 —— 判定与响应存在理论竞态。

改为 checkQuota 一次判定直接把 *quotaRejection 交给 writeReject，
消息与 Retry-After 都来自同一次读，不再有二次求值。
checkModelScope 保留（只需知道放行与否的调用方仍可用）。

## 补判据：此前完全没验证过的交叉场景

1. TestKeyQuotaWinsOverSlotQuota —— key 配额与 AUTO 槽位配额是两种
   不同作用域的限额（槽位是全网关共享的上游预算，key 配额属于单个
   调用方）。两者同时耗尽时必须报【key 配额】：报槽位配额会被表述成
   「无可用容量」，读起来像上游故障，而调用方能处理的恰恰是 key 配额。
2. TestUncappedKeyNeverBlockedByEmptyScope —— 只配模型范围、不配配额的
   key（生产上 5 把 user key 全是这种）绝不能被槽位检查误伤。

## 复查补测的实测数据

配额检查的真实开销（每请求一次，走完整 checkQuota 路径）：
  配了配额    149 ns  0 allocs
  未配配额     42.6 ns 0 allocs   <- 生产上 5/7 把 key 是这种
  admin key    37 ns  0 allocs

未配配额的 key 只付 FindKey 的开销、根本不碰桶。相对一次 LLM 请求
（秒级）可忽略。

生产配置副本（7 key / 16 源 / 真加密凭据 / 真上游）实测：
- 100 并发 -> 50 成功 / 50 容量拒绝，RSS 19.9 -> 25.8 MB
- 生产形态桶内存（7 key x 8 model x 2 源 x 40 天满 retention）
  = 3.73 MB，占 ~32MB 预算的 11%
- 配额记账与 stats 一致：配 63000 配额后报 64062/63000
- **跨重启存活**：重启后从审计日志回放，仍报 64062/63000 并拦截；
  未配配额的 key 仍 200

(cherry picked from commit 9811654b3e)
This commit is contained in:
JianFeeeee
2026-09-27 18:42:03 +08:00
parent d072a03c9a
commit a21ae84cbe
2 changed files with 102 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
internal/gateway/chat.go
View File

@ -197,6 +197,8 @@ func intersectModels(models []string, allow []config.ModelScope) []string {
// A scope entry with model "AUTO" only allows requests where the effective
// model is AUTO (the routing mode). It does NOT grant access to specific model
// ids — that requires an explicit scope entry for the model.
// checkModelScope keeps its string-returning signature for callers that only
// need to know whether the request may proceed.
func (g *Gateway) checkModelScope(ctx context.Context, model string) string {
if q := g.checkQuota(ctx, model); q != nil {
return q.msg
@ -329,26 +331,20 @@ func (g *Gateway) hasScopeModel(list []config.ModelScope, s string) bool {
// 429 (rate_limit_exceeded) with Retry-After, so a client waits and resumes
// after the reset; a model the key may not use stays 403 (model_not_allowed),
// because retrying cannot help.
func (g *Gateway) writeScopeReject(w http.ResponseWriter, r *http.Request, model string) {
if q := g.checkQuota(r.Context(), model); q != nil {
if q.retry > 0 {
w.Header().Set("Retry-After", strconv.FormatInt(q.retry, 10))
writeError(w, http.StatusTooManyRequests, "rate_limit_exceeded", q.msg)
return
}
// no window to wait for: the cap is either permanent or key-wide
// with no period. "rate_limit_exceeded" still says "come back
// after the operator raises the cap", which 403 would not.
if strings.Contains(q.msg, "quota exceeded") {
writeError(w, http.StatusTooManyRequests, "rate_limit_exceeded", q.msg)
return
}
writeError(w, http.StatusForbidden, "model_not_allowed", q.msg)
func (g *Gateway) writeReject(w http.ResponseWriter, q *quotaRejection) {
if q.retry > 0 {
w.Header().Set("Retry-After", strconv.FormatInt(q.retry, 10))
writeError(w, http.StatusTooManyRequests, "rate_limit_exceeded", q.msg)
return
}
// The scope check already passed; reaching here means the state changed
// between the two calls. Fall back to the pre-existing behaviour.
writeError(w, http.StatusForbidden, "model_not_allowed", fmt.Sprintf("model %q is not allowed for this key", model))
// no window to wait for: the cap is either permanent or key-wide with no
// period. "rate_limit_exceeded" still says "come back after the operator
// raises the cap", which 403 would not.
if strings.Contains(q.msg, "quota exceeded") {
writeError(w, http.StatusTooManyRequests, "rate_limit_exceeded", q.msg)
return
}
writeError(w, http.StatusForbidden, "model_not_allowed", q.msg)
}
func (g *Gateway) resolveByModel(model string) ([]*provider.Provider, string) {
@ -413,8 +409,8 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
writeError(w, http.StatusServiceUnavailable, "no_provider", "no auto slot configured")
return
}
if msg := g.checkModelScope(r.Context(), "AUTO"); msg != "" {
g.writeScopeReject(w, r, "AUTO")
if q := g.checkQuota(r.Context(), "AUTO"); q != nil {
g.writeReject(w, q)
return
}
ctx := r.Context()
@ -467,8 +463,8 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
if effective == "" {
effective = firstModel(cands[0])
}
if msg := g.checkModelScope(r.Context(), effective); msg != "" {
g.writeScopeReject(w, r, effective)
if q := g.checkQuota(r.Context(), effective); q != nil {
g.writeReject(w, q)
return
}
ctx := r.Context()
@ -1124,8 +1120,8 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
}
if isAuto(model) {
if chain := g.core.AutoImageChain(); chain != nil && len(chain.Tiers) > 0 {
if msg := g.checkModelScope(r.Context(), "AUTO"); msg != "" {
g.writeScopeReject(w, r, "AUTO")
if q := g.checkQuota(r.Context(), "AUTO"); q != nil {
g.writeReject(w, q)
return
}
done := g.stats.Begin()
@ -1171,8 +1167,8 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
writeError(w, http.StatusServiceUnavailable, "no_provider", "no image source configured")
return
}
if msg := g.checkModelScope(r.Context(), effectiveImageModel(model, cands)); msg != "" {
g.writeScopeReject(w, r, effectiveImageModel(model, cands))
if q := g.checkQuota(r.Context(), effectiveImageModel(model, cands)); q != nil {
g.writeReject(w, q)
return
}
done := g.stats.Begin()