From 690f55c6f16d481ecc60d9247b98174abf8bbee4 Mon Sep 17 00:00:00 2001 From: JianFeeeee Date: Mon, 24 Aug 2026 02:00:25 +0800 Subject: [PATCH] feat: proactive rate limiting (RPM) + 429 short cooldown for sources MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - config.go: Source add RPM field (requests-per-minute cap, 0=unlimited) - provider.go: RecordRateLimit() — 429 uses fixed 30s cooldown, not exponential - provider.go: Throttle() — token-bucket proactive rate limiter, spaces requests at 60s/RPM interval, respects context cancellation - provider.go: ReportStatus() — 429 -> RecordRateLimit, 5xx -> RecordFailure - provider.go: Chat/ChatStream — wire Throttle after TryAcquire - api.go: sourcePayload + RPM, buildSource passes RPM through - ui/index.html: add RPM input field in source editor, bilingual i18n labels - deploy.sh: backup old binary + rollback on healthcheck failure - provider_test.go: TestModelStateRateLimitShortCooldown, TestThrottleSpacingAndCancel - config.yaml: sensenova rpm: 12 --- deploy.sh | 274 +++++++++++++++++++++++++++++ internal/config/config.go | 1 + internal/gateway/api.go | 2 + internal/gateway/ui/index.html | 7 + internal/provider/provider.go | 78 +++++++- internal/provider/provider_test.go | 71 ++++++++ 6 files changed, 429 insertions(+), 4 deletions(-) create mode 100755 deploy.sh diff --git a/deploy.sh b/deploy.sh new file mode 100755 index 0000000..ae09242 --- /dev/null +++ b/deploy.sh @@ -0,0 +1,274 @@ +#!/usr/bin/env bash +# ============================================================ +# llmsproxy 原子化滚动部署脚本 +# +# 策略： +# 1. /tmp (tmpfs) 构建，/usr/local/bin (ext4) 部署 +# 2. 构建后先 copy 到目标文件系统（.llmsproxy.tmp）， +# 再 rename(2) 做原子替换（同一文件系统内 mv 是原子的） +# 3. runtime.json 不替换（保留运行时密钥、审计状态） +# 4. systemctl restart → 服务中断仅在 Go 进程重启瞬间 (<500ms) +# 5. 健康检查确认新进程正常 +# +# 注意事项： +# - 本服务依赖自身（推理依赖），restart 前后会短暂中断， +# 但重试机制（systemd RestartSec=5 + 客户端 retry）可自愈 +# - 适配器文件替换后需重启才能生效（VM 启动时批量加载） +# ============================================================ +set -euo pipefail + +# ---------- 配置 ---------- +readonly TARGET_BIN="/usr/local/bin/llmsproxy" +readonly TARGET_ADAPTERS="/etc/llmsproxy/adapters" +readonly REPO_DIR="/home/program/llmsproxy" +readonly HEALTH_URL="http://127.0.0.1:8081/v1/models" +readonly BIN_NAME="llmsproxy" + +# 构建目录在 /tmp（tmpfs 更快） +readonly TMP_DIR="/tmp/llmsproxy-deploy" +# 同文件系统暂存文件（用于原子 rename） +readonly STAGING_BIN="${TARGET_BIN}.staging" +# 回滚备份：健康检查失败时从这里恢复旧二进制 +readonly BACKUP_BIN="${TARGET_BIN}.bak" + +# ---------- 日志 ---------- +log() { printf '\033[1;34m[%s]\033[0m %s\n' "$(date +%H:%M:%S)" "$*"; } +warn() { printf '\033[1;33m[WARN]\033[0m %s\n' "$*"; } +err() { printf '\033[1;31m[ERR]\033[0m %s\n' "$*" >&2; } +fail() { err "$@"; exit 1; } + +# ---------- 清理 ---------- +cleanup() { + log "清理临时文件" + rm -rf "$TMP_DIR" "$STAGING_BIN" +} +trap cleanup EXIT + +# ---------- 前置检查 ---------- +precheck() { + log "前置检查" + [[ -f "$REPO_DIR/go.mod" ]] || fail "仓库目录 $REPO_DIR 不存在" + command -v go &>/dev/null || fail "未找到 go" + command -v systemctl &>/dev/null || fail "未找到 systemctl" + [[ $(id -u) -eq 0 ]] || fail "需要 root 权限" + + OLD_SIZE=$(stat -c%s "$TARGET_BIN" 2>/dev/null || echo 0) + OLD_SHA256=$(sha256sum "$TARGET_BIN" 2>/dev/null | awk '{print $1}') + + log "旧二进制: $TARGET_BIN (${OLD_SIZE} bytes, sha256=${OLD_SHA256:0:16})" + log "go: $(go version)" +} + +# ---------- 备份旧二进制（用于回滚）---------- +backup_old_binary() { + log "备份旧二进制到 $BACKUP_BIN" + if [[ -f "$TARGET_BIN" ]]; then + cp -f "$TARGET_BIN" "$BACKUP_BIN" + log "✓ 旧二进制已备份 (sha256=$(sha256sum "$BACKUP_BIN" | awk '{print $1}' | cut -c1-16))" + else + warn "旧二进制不存在，跳过备份（无回滚点）" + fi +} + +# ---------- 构建 ---------- +build() { + log "构建新二进制" + rm -rf "$TMP_DIR" + mkdir -p "$TMP_DIR/bin" + + cd "$REPO_DIR" + CGO_ENABLED=1 go build -tags luajit \ + -trimpath \ + -ldflags="-s -w" \ + -o "$TMP_DIR/bin/$BIN_NAME" \ + ./cmd/llmsproxy + + NEW_BIN="$TMP_DIR/bin/$BIN_NAME" + NEW_SIZE=$(stat -c%s "$NEW_BIN") + + if [[ "$NEW_SIZE" -eq "$OLD_SIZE" ]]; then + NEW_SHA256=$(sha256sum "$NEW_BIN" | awk '{print $1}') + if [[ "$NEW_SHA256" == "$OLD_SHA256" ]]; then + warn "新二进制与旧版本完全一致（sha256 相同），跳过部署" + exit 0 + fi + fi + + log "新二进制: $NEW_BIN (${NEW_SIZE} bytes)" + + # 验证：确认包含关键修复 + local checks=( + "mergeUsage" + "prompt_tokens" + "completion_tokens" + "message_start" + ) + for c in "${checks[@]}"; do + if strings "$NEW_BIN" | grep -qF "$c"; then + log " ✓ 新二进制包含: $c" + else + warn " ⚠ 新二进制不包含: $c（非预期但非致命）" + fi + done +} + +# ---------- 原子替换二进制 ---------- +# 原理： +# /tmp 是 tmpfs，/usr/local/bin 是 ext4——不同文件系统 rename(2) 不原子。 +# 因此先 copy 到目标文件系统上的 .staging 文件，再 rename(2) 到目标路径。 +# rename(2) 在同一文件系统内是元数据级操作，对正在读写的进程安全： +# - 旧 inode 的 fd 继续有效（进程持有旧文件，不会读到截断内容） +# - 新进程打开 path 看到新 inode +# - 旧 inode 在最后一个 fd 关闭后释放 +atomic_replace_binary() { + log "原子替换二进制" + + local new_bin_tmp="$TMP_DIR/bin/$BIN_NAME" + [[ -f "$new_bin_tmp" ]] || fail "新二进制文件不存在: $new_bin_tmp" + + # Step 1: copy 到目标文件系统（非原子，但目标文件唯一） + log " copy → $STAGING_BIN (同文件系统暂存)" + cp -f "$new_bin_tmp" "$STAGING_BIN" + chmod 0755 "$STAGING_BIN" + + # Step 2: 原子 rename（同一文件系统，元数据级操作） + log " rename → $TARGET_BIN (原子)" + mv -f "$STAGING_BIN" "$TARGET_BIN" + + # 验证 + NEW_SHA256=$(sha256sum "$TARGET_BIN" | awk '{print $1}') + log "✓ 二进制已原子替换 (sha256=${NEW_SHA256:0:16})" +} + +# ---------- 同步适配器 ---------- +sync_adapters() { + log "同步适配器" + + # 只从内嵌适配器目录同步（完整集 10 个适配器，含最新 usage 透传修复）。 + # 注意：仓库根的 adapters/ 是运行时覆盖目录，可能不全（缺 opencode.lua 等）， + # 不可作为部署源。 + SRC_ADAPTERS="$REPO_DIR/internal/lua/adapters" + [[ -d "$SRC_ADAPTERS" ]] || fail "内嵌适配器目录不存在: $SRC_ADAPTERS" + + # 备份旧的适配器文件（含任何运行时覆盖版本，如 opencode.lua） + BACKUP_DIR="/etc/llmsproxy/adapters.bak.$(date +%Y%m%d%H%M%S)" + if [[ -d "$TARGET_ADAPTERS" ]]; then + mkdir -p "$BACKUP_DIR" + cp -rf "$TARGET_ADAPTERS/"*.lua "$BACKUP_DIR/" 2>/dev/null || true + log " 旧适配器已备份到 $BACKUP_DIR" + fi + + # 清空目标目录中的 .lua 文件（保留 .bak.* 归档），再复制全部新适配器。 + # 不依赖 rsync（部署机可能未安装），纯 shell 保证可移植。 + find "$TARGET_ADAPTERS" -maxdepth 1 -name '*.lua' -not -name '*.bak.*' -delete + cp -f "$SRC_ADAPTERS/"*.lua "$TARGET_ADAPTERS/" + chmod 0644 "$TARGET_ADAPTERS/"*.lua + + # 校验：每个适配器必须包含 usage 透传修复（opencode/openai/deepseek 等应有 'uses'） + local required_files=(openai deepseek anthropic gemini ollama opencode github groq kimicode mistral) + for f in "${required_files[@]}"; do + [[ -f "$TARGET_ADAPTERS/$f.lua" ]] || warn " 缺少适配器: $f.lua" + done + local usage_fixed=0 + for f in openai deepseek github groq kimicode mistral opencode; do + grep -qE '\buses\b' "$TARGET_ADAPTERS/$f.lua" 2>/dev/null && ((usage_fixed++)) || true + done + log " usage 透传修复适配器: ${usage_fixed}/7 (openai/deepseek/github/groq/kimicode/mistral/opencode)" + log "✓ 适配器已同步到 $TARGET_ADAPTERS/" +} + +# ---------- 重启服务 ---------- +restart_service() { + log "重启服务（中断 <500ms）" + systemctl restart llmsproxy.service || fail "systemd restart 失败" + log "✓ llmsproxy.service restarted" +} + +# ---------- 健康检查 ---------- +healthcheck() { + log "健康检查" + local attempts=20 + local wait_sec=1 + local attempt=1 + + while (( attempt <= attempts )); do + # 用 curl 检查（200=有 key 未鉴权但服务正常，401=需要认证，都说明服务在运行） + local code + code=$(curl -s -o /dev/null -w "%{http_code}" --max-time 3 "$HEALTH_URL" 2>/dev/null || echo "000") + if [[ "$code" == "200" || "$code" == "401" ]]; then + log "✓ 服务健康 (status=$code, ${attempt}/${attempts})" + return 0 + fi + log " 等待中 ($attempt/$attempts, status=$code)" + sleep "$wait_sec" + (( attempt++ )) + done + + warn "健康检查未通过，请检查服务状态: systemctl status llmsproxy.service" + return 1 +} + +# ---------- 回滚 ---------- +# 健康检查失败时恢复旧二进制并重启，保证推理服务可用。 +# 这是最后一道保险：即使新版本启动失败（编译错误、配置不兼容、 +# panic on start 等），也能让旧版本继续服务，避免推理能力丢失。 +rollback() { + err "健康检查失败，启动回滚流程" + if [[ ! -f "$BACKUP_BIN" ]]; then + err "无回滚备份 ($BACKUP_BIN 不存在)，无法自动恢复" + err "请手动检查: systemctl status llmsproxy.service && journalctl -u llmsproxy -n 50" + return 1 + fi + + log "恢复旧二进制" + cp -f "$BACKUP_BIN" "$TARGET_BIN" + chmod 0755 "$TARGET_BIN" + + log "重启服务（旧版本）" + systemctl restart llmsproxy.service || { + err "重启失败，服务可能完全宕机" + err "请手动检查: systemctl status llmsproxy.service" + return 1 + } + + # 回滚后再健康检查（只查 1 次，失败就认了） + sleep 2 + local code + code=$(curl -s -o /dev/null -w "%{http_code}" --max-time 5 "$HEALTH_URL" 2>/dev/null || echo "000") + if [[ "$code" == "200" || "$code" == "401" ]]; then + log "✓ 回滚成功，旧版本已恢复服务 (status=$code)" + log " 新版本部署失败，请检查构建/配置后重试" + return 0 + else + err "回滚后服务仍未恢复 (status=$code)" + err "服务可能已完全宕机，请手动介入" + return 1 + fi +} + +# ---------- 打印部署摘要 ---------- +print_summary() { + log "═════════════════════════════════════════════" + log "部署完成" + log " 二进制: $TARGET_BIN" + log " 适配器: $TARGET_ADAPTERS" + log "═════════════════════════════════════════════" +} + +# ---------- 主流程 ---------- +main() { + log "═══ llmsproxy 原子部署开始 ═══" + precheck + build + sync_adapters + backup_old_binary + atomic_replace_binary + restart_service + if ! healthcheck; then + rollback || true + exit 1 + fi + print_summary +} + +main "$@" \ No newline at end of file diff --git a/internal/config/config.go b/internal/config/config.go index 8c57e18..4820f48 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -62,6 +62,7 @@ type Source struct { Timeout time.Duration `yaml:"timeout" json:"-"` MaxConcurrent int `yaml:"max_concurrent" json:"max_concurrent"` QueueTimeout time.Duration `yaml:"queue_timeout" json:"-"` + RPM int `yaml:"rpm" json:"rpm,omitempty"` // optional requests-per-minute cap (0 = unlimited) } // Load reads and validates a config file. When the file does not exist yet a diff --git a/internal/gateway/api.go b/internal/gateway/api.go index 76c1e1f..c8da5ed 100644 --- a/internal/gateway/api.go +++ b/internal/gateway/api.go @@ -72,6 +72,7 @@ type sourcePayload struct { Temperature float64 `json:"temperature"` MaxTokens int `json:"max_tokens"` MaxConcurrent int `json:"max_concurrent"` + RPM int `json:"rpm"` // optional requests-per-minute cap, 0 = unlimited } func (g *Gateway) handleSourcesAPI(w http.ResponseWriter, r *http.Request) { @@ -105,6 +106,7 @@ func (g *Gateway) handleSourcesAPI(w http.ResponseWriter, r *http.Request) { Temperature: p.Temperature, MaxTokens: p.MaxTokens, MaxConcurrent: p.MaxConcurrent, + RPM: p.RPM, } if err := g.core.AddSource(src); err != nil { writeError(w, http.StatusBadRequest, "source_error", err.Error()) diff --git a/internal/gateway/ui/index.html b/internal/gateway/ui/index.html index bc40811..5846e04 100644 --- a/internal/gateway/ui/index.html +++ b/internal/gateway/ui/index.html @@ -847,6 +847,7 @@ mEp: "聊天端点", mImgEp: "生图端点", mConc: "并发上限", + mRPM: "RPM 限速 (0=不限)", mTemp: "温度", mModels: "模型列表", mAddModel: "+ 模型", @@ -1067,6 +1068,7 @@ mEp: "Chat endpoint", mImgEp: "Image endpoint", mConc: "Max concurrency", + mRPM: "RPM limit (0 = unlimited)", mTemp: "Temperature", mModels: "Models", mAddModel: "+ model", @@ -2479,6 +2481,10 @@
+
+
+
+
@@ -2537,6 +2543,7 @@ endpoint: $("#s-ep").value.trim(), image_endpoint: $("#s-img").value.trim(), max_concurrent: parseInt($("#s-conc").value) || 8, + rpm: parseInt($("#s-rpm").value) || 0, temperature: parseFloat($("#s-temp").value) || 0, models, meta, diff --git a/internal/provider/provider.go b/internal/provider/provider.go index dd16e3d..f4e2d94 100644 --- a/internal/provider/provider.go +++ b/internal/provider/provider.go @@ -47,6 +47,13 @@ const ( backoffCapN = 10 ) +// rateLimitCooldown is the fixed cooldown applied on HTTP 429. A rate-limit +// rejection means "too fast", not "broken": unlike transport/5xx failures it +// must NOT escalate exponentially to backoffCap, or a quota-rich source gets +// locked out for up to 30 minutes while its allowance is untouched. A short +// fixed window lets the slot re-enter the rotation quickly. +const rateLimitCooldown = 30 * time.Second + func clampPref(a *atomic.Int64, lo, hi int64) { for { cur := a.Load() @@ -100,6 +107,17 @@ func (s *ModelState) RecordFailure(auth bool) { clampPref(&s.pref, prefMin, prefMax) } +// RecordRateLimit records an HTTP 429: a short fixed cooldown instead of the +// exponential schedule, so a merely-throttled source recovers in seconds and +// its remaining quota stays usable. The preference penalty still applies so +// other slots are preferred while cooling. +func (s *ModelState) RecordRateLimit() { + s.failCount.Add(1) + s.cooldownUntil.Store(time.Now().Add(rateLimitCooldown).Unix()) + s.pref.Add(-prefFailStep) + clampPref(&s.pref, prefMin, prefMax) +} + // RecordSuccess resets the failure counter and cooldown and bumps the // preference score by one. func (s *ModelState) RecordSuccess() { @@ -139,6 +157,12 @@ type Provider struct { mu sync.Mutex sem chan struct{} states map[string]*ModelState // key = model id + // proactive rate limiting: nextOK is the earliest unix-nano time the next + // request may leave; throttle() spaces requests 60s/RPM apart. Zero when + // cfg.RPM <= 0 (unlimited). + rateMu sync.Mutex + nextOK int64 + rpmGap time.Duration lastProbe struct { ok bool err string @@ -178,6 +202,9 @@ func New(cfg config.Source, vm *lua.VM) *Provider { if cfg.MaxConcurrent <= 0 { p.sem = nil } + if cfg.RPM > 0 { + p.rpmGap = time.Minute / time.Duration(cfg.RPM) + } for _, m := range cfg.Models { p.states[m.ID] = &ModelState{} } @@ -491,15 +518,20 @@ func (p *Provider) RecordSuccess(model string) { // ReportStatus records an upstream HTTP status for the given model and drives // the (source, model) backoff state. 401/403 → capped self-healing cooldown -// with doubled penalty; 429 and 5xx → normal exponential backoff. Other codes -// (400 client schema errors, 402 billing errors) are not penalized here — -// they surface via the status page / audit instead. +// with doubled penalty; 5xx → normal exponential backoff; 429 → short fixed +// cooldown (rateLimitCooldown): a rate rejection means "too fast", not +// "broken", so the slot must re-enter rotation quickly instead of escalating +// to a 30-minute lockout that wastes remaining quota. func (p *Provider) ReportStatus(model string, code int) { if code == 401 || code == 403 { p.RecordFailure(model, code) return } - if code >= 500 || code == 429 { + if code == 429 { + p.state(model).RecordRateLimit() + return + } + if code >= 500 { p.RecordFailure(model, code) } } @@ -586,6 +618,37 @@ func (p *Provider) Release() { <-p.sem } +// ---- proactive rate limiting ---- + +// Throttle blocks until this source's rate-limit window allows the next +// request, spacing outbound requests at least rpmGap apart (60s/RPM). It is a +// no-op when no rpm cap is configured. The wait happens AFTER the concurrency +// slot is taken, so a queued request still counts against max_concurrent — +// bounded by the caller's ctx (queue timeout / client disconnect). +func (p *Provider) Throttle(ctx context.Context) error { + if p.rpmGap <= 0 { + return nil + } + p.rateMu.Lock() + now := time.Now().UnixNano() + wait := p.nextOK - now + if wait > 0 { + p.nextOK += int64(p.rpmGap) // reserve the next window for the follower + p.rateMu.Unlock() + t := time.NewTimer(time.Duration(wait)) + defer t.Stop() + select { + case <-t.C: + return nil + case <-ctx.Done(): + return ctx.Err() + } + } + p.nextOK = now + int64(p.rpmGap) + p.rateMu.Unlock() + return nil +} + // ---- request construction ---- func (p *Provider) buildHeaders(body, url string) (http.Header, error) { @@ -630,6 +693,9 @@ func (p *Provider) Chat(ctx context.Context, req *types.ChatRequest) (*types.Uni return nil, err } defer p.Release() + if err := p.Throttle(ctx); err != nil { + return nil, err + } model := p.ModelFor(req.Model) body, err := marshalTransform(p.vm, p.adapter, "transform_request", req) @@ -684,6 +750,10 @@ func (p *Provider) ChatStream(ctx context.Context, req *types.ChatRequest) (<-ch if err := p.TryAcquire(ctx); err != nil { return nil, err } + if err := p.Throttle(ctx); err != nil { + p.Release() + return nil, err + } model := p.ModelFor(req.Model) req.Stream = true body, err := marshalTransform(p.vm, p.adapter, "transform_request", req) diff --git a/internal/provider/provider_test.go b/internal/provider/provider_test.go index b0bb725..f84dad0 100644 --- a/internal/provider/provider_test.go +++ b/internal/provider/provider_test.go @@ -245,6 +245,77 @@ func TestModelStateCooldownAndRecovery(t *testing.T) { } } +func TestModelStateRateLimitShortCooldown(t *testing.T) { + // 429 must NOT use the exponential backoff schedule: a rate-limited but + // quota-rich source has to re-enter rotation after the short fixed window. + p := newTestProvider(t, src("mock", "http://127.0.0.1:1", "openai", "m")) + st := p.state("m") + + // repeated 429s must stay at the fixed short cooldown, never escalate + for i := 0; i < 15; i++ { + p.ReportStatus("m", 429) + } + until := st.CooldownUntil() + want := time.Now().Add(rateLimitCooldown).Unix() + if until < want-2 || until > want+2 { + t.Fatalf("429 cooldown = %d, want ~%d (fixed %v, not exponential)", until, want, rateLimitCooldown) + } + if st.FailCount() != 15 { + t.Fatalf("fail count = %d, want 15 (counted but not escalating)", st.FailCount()) + } + if st.Pref() != -15*int64(prefFailStep) && st.Pref() > int64(prefMin) { + t.Fatalf("pref = %d", st.Pref()) + } + if p.ModelAvailable("m") { + t.Fatal("model must be cooling right after a 429") + } +} + +func TestThrottleSpacingAndCancel(t *testing.T) { + p := newTestProvider(t, src("mock", "http://127.0.0.1:1", "openai", "m")) + p.cfg.RPM = 120 // gap = 500ms + p.rpmGap = time.Minute / time.Duration(p.cfg.RPM) + + start := time.Now() + for i := 0; i < 3; i++ { + if err := p.Throttle(context.Background()); err != nil { + t.Fatalf("throttle %d: %v", i, err) + } + } + elapsed := time.Since(start) + // first call passes immediately, the next two wait one gap each + want := 2 * time.Minute / time.Duration(p.cfg.RPM) + if elapsed < want-time.Duration(100*time.Millisecond) || elapsed > want+time.Second { + t.Fatalf("3 throttled calls took %v, want ~%v", elapsed, want) + } + + // unlimited source: Throttle is a no-op + p2 := newTestProvider(t, src("free", "http://127.0.0.1:1", "openai", "m")) + if err := p2.Throttle(context.Background()); err != nil { + t.Fatalf("unlimited throttle: %v", err) + } + + // cancelled context aborts a pending window reservation + p3 := newTestProvider(t, src("slow", "http://127.0.0.1:1", "openai", "m")) + p3.cfg.RPM = 6 // 10s gap + p3.rpmGap = time.Minute / time.Duration(p3.cfg.RPM) + // first call passes immediately (fresh provider) and reserves the next window + if err := p3.Throttle(context.Background()); err != nil { + t.Fatalf("first throttle: %v", err) + } + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + // second call must wait ~10s for the reserved window; the short deadline aborts it + start = time.Now() + err := p3.Throttle(ctx) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("cancelled throttle = %v, want DeadlineExceeded", err) + } + if time.Since(start) > time.Second { + t.Fatalf("cancel took %v, want fast abort", time.Since(start)) + } +} + func TestTryAcquire(t *testing.T) { p := newTestProvider(t, src("mock", "http://127.0.0.1:1", "openai", "m")) p.cfg.MaxConcurrent = 1