Day 2：限流与熔断——生产环境第一道防线

今天做什么

网关没有限流，一个突发流量就能把所有 LLM API 额度打爆。今天加两道防线：令牌桶限流 + 熔断器。

代码

package main

import (
    "sync"
    "sync/atomic"
    "time"
)

// ===== 令牌桶限流器 =====
type TokenBucket struct {
    rate       float64    // 每秒生成令牌数
    capacity   float64    // 桶容量（允许突发）
    tokens     float64    // 当前令牌数
    lastRefill time.Time
    mu         sync.Mutex
}

func NewTokenBucket(rate, capacity float64) *TokenBucket {
    return &TokenBucket{
        rate:       rate,
        capacity:   capacity,
        tokens:     capacity,
        lastRefill: time.Now(),
    }
}

func (tb *TokenBucket) Allow() bool {
    tb.mu.Lock()
    defer tb.mu.Unlock()

    // 补充令牌
    now := time.Now()
    elapsed := now.Sub(tb.lastRefill).Seconds()
    tb.tokens += elapsed * tb.rate
    if tb.tokens > tb.capacity {
        tb.tokens = tb.capacity
    }
    tb.lastRefill = now

    // 尝试消耗一个令牌
    if tb.tokens >= 1 {
        tb.tokens--
        return true
    }
    return false
}

// ===== 熔断器 =====
type CircuitBreaker struct {
    failureCount atomic.Int64
    successCount atomic.Int64
    threshold    int64 // 连续失败N次触发熔断
    state        atomic.Int32 // 0=关闭 1=开启 2=半开
    lastFailTime atomic.Int64
    cooldown     time.Duration
}

const (
    StateClosed    = 0
    StateOpen      = 1
    StateHalfOpen  = 2
)

func NewCircuitBreaker(threshold int64, cooldown time.Duration) *CircuitBreaker {
    return &CircuitBreaker{
        threshold: threshold,
        cooldown:  cooldown,
    }
}

func (cb *CircuitBreaker) Allow() bool {
    state := cb.state.Load()
    if state == StateClosed {
        return true
    }
    if state == StateHalfOpen {
        return true
    }
    // 熔断中，检查是否过了冷却期
    if time.Now().Unix()-cb.lastFailTime.Load() > int64(cb.cooldown.Seconds()) {
        cb.state.Store(StateHalfOpen)
        return true
    }
    return false
}

func (cb *CircuitBreaker) RecordSuccess() {
    cb.failureCount.Store(0)
    cb.state.Store(StateClosed)
}

func (cb *CircuitBreaker) RecordFailure() {
    count := cb.failureCount.Add(1)
    cb.lastFailTime.Store(time.Now().Unix())
    if count >= cb.threshold {
        cb.state.Store(StateOpen)
        log.Printf("🔴 熔断触发！连续 %d 次失败", count)
    }
}

// ===== 网关集成 =====
type Gateway struct {
    router    *ModelRouter
    limiter   *TokenBucket
    breaker   *CircuitBreaker
}

func (g *Gateway) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    // 第一关：限流
    if !g.limiter.Allow() {
        http.Error(w, `{"error":"rate limited"}`, 429)
        return
    }

    // 第二关：熔断
    if !g.breaker.Allow() {
        http.Error(w, `{"error":"service unavailable"}`, 503)
        return
    }

    // 正常转发
    g.router.ProxyHandler(w, r)
}

func main() {
    gw := &Gateway{
        router:  NewRouter(),
        limiter: NewTokenBucket(100, 200), // 每秒100个请求，突发200
        breaker: NewCircuitBreaker(10, 30*time.Second), // 10次失败熔断30秒
    }

    http.Handle("/v1/chat/completions", http.HandlerFunc(gw.ServeHTTP))
    log.Fatal(http.ListenAndServe(":8080", nil))
}

运行

go run main.go

# 压测限流
ab -n 1000 -c 100 http://localhost:8080/v1/chat/completions
# 部分请求会返回 429

---

Day 2 完成。网关不怕突发流量了。