06 - 限流系统设计

package ratelimit

import (
    "sync"
    "time"
)

// FixedWindowLimiter 固定窗口限流器
type FixedWindowLimiter struct {
    limit    int64         // 限流阈值
    window   time.Duration // 时间窗口
    counter  int64         // 当前计数
    lastTime time.Time     // 窗口开始时间
    mu       sync.Mutex    // 互斥锁
}

// NewFixedWindowLimiter 创建固定窗口限流器
func NewFixedWindowLimiter(limit int64, window time.Duration) *FixedWindowLimiter {
    return &FixedWindowLimiter{
        limit:    limit,
        window:   window,
        counter:  0,
        lastTime: time.Now(),
    }
}

// Allow 判断是否允许请求
func (l *FixedWindowLimiter) Allow() bool {
    l.mu.Lock()
    defer l.mu.Unlock()

    now := time.Now()
    // 窗口是否过期
    if now.Sub(l.lastTime) >= l.window {
        // 重置窗口
        l.counter = 0
        l.lastTime = now
    }

    // 判断是否超过限流阈值
    if l.counter < l.limit {
        l.counter++
        return true
    }

    return false
}

// 使用示例
func ExampleFixedWindow() {
    // 限制每秒100次请求
    limiter := NewFixedWindowLimiter(100, time.Second)

    // 模拟请求
    for i := 0; i < 150; i++ {
        if limiter.Allow() {
            // 处理请求
            println("Request", i, "allowed")
        } else {
            // 拒绝请求
            println("Request", i, "rejected")
        }
    }
}

package ratelimit

import (
    "sync"
    "time"
)

// SlidingWindowLimiter 滑动窗口限流器
type SlidingWindowLimiter struct {
    limit     int64              // 限流阈值
    window    time.Duration      // 时间窗口
    slots     int                // 窗口分片数
    counters  map[int64]int64    // 时间槽 -> 计数
    mu        sync.Mutex         // 互斥锁
}

// NewSlidingWindowLimiter 创建滑动窗口限流器
func NewSlidingWindowLimiter(limit int64, window time.Duration, slots int) *SlidingWindowLimiter {
    return &SlidingWindowLimiter{
        limit:    limit,
        window:   window,
        slots:    slots,
        counters: make(map[int64]int64),
    }
}

// Allow 判断是否允许请求
func (l *SlidingWindowLimiter) Allow() bool {
    l.mu.Lock()
    defer l.mu.Unlock()

    now := time.Now().UnixNano()
    slotDuration := int64(l.window) / int64(l.slots)

    // 计算当前时间槽
    currentSlot := now / slotDuration

    // 计算滑动窗口的起始时间槽
    windowStart := currentSlot - int64(l.slots) + 1

    // 清理过期的时间槽
    for slot := range l.counters {
        if slot < windowStart {
            delete(l.counters, slot)
        }
    }

    // 统计当前窗口内的请求总数
    var total int64
    for slot := windowStart; slot <= currentSlot; slot++ {
        total += l.counters[slot]
    }

    // 判断是否超过限流阈值
    if total < l.limit {
        l.counters[currentSlot]++
        return true
    }

    return false
}

// 使用示例
func ExampleSlidingWindow() {
    // 限制每秒100次请求，窗口分为10个槽
    limiter := NewSlidingWindowLimiter(100, time.Second, 10)

    // 模拟请求
    for i := 0; i < 150; i++ {
        if limiter.Allow() {
            println("Request", i, "allowed")
        } else {
            println("Request", i, "rejected")
        }
        time.Sleep(5 * time.Millisecond)
    }
}

package ratelimit

import (
    "sync"
    "time"
)

// TokenBucketLimiter 令牌桶限流器
type TokenBucketLimiter struct {
    capacity   int64         // 桶容量
    tokens     int64         // 当前令牌数
    rate       int64         // 令牌生成速率（个/秒）
    lastTime   time.Time     // 上次生成令牌时间
    mu         sync.Mutex    // 互斥锁
}

// NewTokenBucketLimiter 创建令牌桶限流器
func NewTokenBucketLimiter(capacity int64, rate int64) *TokenBucketLimiter {
    return &TokenBucketLimiter{
        capacity: capacity,
        tokens:   capacity,  // 初始令牌数等于容量
        rate:     rate,
        lastTime: time.Now(),
    }
}

// Allow 判断是否允许请求（消耗1个令牌）
func (l *TokenBucketLimiter) Allow() bool {
    return l.AllowN(1)
}

// AllowN 判断是否允许请求（消耗n个令牌）
func (l *TokenBucketLimiter) AllowN(n int64) bool {
    l.mu.Lock()
    defer l.mu.Unlock()

    // 计算当前应该有多少令牌
    now := time.Now()
    elapsed := now.Sub(l.lastTime)

    // 计算新生成的令牌数
    newTokens := int64(elapsed.Seconds() * float64(l.rate))

    // 更新令牌数（不超过容量）
    l.tokens = min(l.capacity, l.tokens+newTokens)
    l.lastTime = now

    // 判断令牌是否足够
    if l.tokens >= n {
        l.tokens -= n
        return true
    }

    return false
}

func min(a, b int64) int64 {
    if a < b {
        return a
    }
    return b
}

// 使用示例
func ExampleTokenBucket() {
    // 容量100，每秒生成50个令牌
    limiter := NewTokenBucketLimiter(100, 50)

    // 模拟突发请求（前100个直接通过）
    for i := 0; i < 150; i++ {
        if limiter.Allow() {
            println("Request", i, "allowed")
        } else {
            println("Request", i, "rejected")
        }
    }

    // 等待1秒，令牌恢复
    time.Sleep(time.Second)

    // 继续请求（恢复50个令牌）
    for i := 150; i < 200; i++ {
        if limiter.Allow() {
            println("Request", i, "allowed")
        } else {
            println("Request", i, "rejected")
        }
    }
}

package ratelimit

import (
    "context"
    "sync"
    "time"
)

// LeakyBucketLimiter 漏桶限流器
type LeakyBucketLimiter struct {
    capacity int           // 桶容量
    rate     time.Duration // 流出速率（每个请求的间隔时间）
    queue    chan struct{} // 请求队列
    mu       sync.Mutex    // 互斥锁
}

// NewLeakyBucketLimiter 创建漏桶限流器
func NewLeakyBucketLimiter(capacity int, rate time.Duration) *LeakyBucketLimiter {
    limiter := &LeakyBucketLimiter{
        capacity: capacity,
        rate:     rate,
        queue:    make(chan struct{}, capacity),
    }

    // 启动漏水协程
    go limiter.leak()

    return limiter
}

// leak 漏水（以恒定速率流出）
func (l *LeakyBucketLimiter) leak() {
    ticker := time.NewTicker(l.rate)
    defer ticker.Stop()

    for range ticker.C {
        // 从队列中取出一个请求（流出）
        select {
        case <-l.queue:
            // 请求流出
        default:
            // 队列为空，无需处理
        }
    }
}

// Allow 判断是否允许请求
func (l *LeakyBucketLimiter) Allow() bool {
    select {
    case l.queue <- struct{}{}:
        // 成功加入队列
        return true
    default:
        // 队列已满，拒绝请求
        return false
    }
}

// Wait 等待直到允许请求（带超时）
func (l *LeakyBucketLimiter) Wait(ctx context.Context) error {
    select {
    case l.queue <- struct{}{}:
        return nil
    case <-ctx.Done():
        return ctx.Err()
    }
}

// 使用示例
func ExampleLeakyBucket() {
    // 容量100，每10ms流出一个请求（100个/秒）
    limiter := NewLeakyBucketLimiter(100, 10*time.Millisecond)

    // 模拟突发请求
    for i := 0; i < 150; i++ {
        if limiter.Allow() {
            println("Request", i, "allowed (queued)")
        } else {
            println("Request", i, "rejected (bucket full)")
        }
    }

    // 等待队列处理
    time.Sleep(2 * time.Second)
}

// RateLimiter 限流器接口
type RateLimiter interface {
    // Allow 判断是否允许请求通过
    // 返回值: true-允许, false-拒绝
    Allow(ctx context.Context, resource string, identifier string) (bool, error)

    // AllowN 判断是否允许n个请求通过
    AllowN(ctx context.Context, resource string, identifier string, n int64) (bool, error)

    // Wait 等待直到允许请求通过（阻塞）
    // 返回值: 等待时间, error
    Wait(ctx context.Context, resource string, identifier string) (time.Duration, error)

    // Reserve 预约令牌（不阻塞，返回等待时间）
    Reserve(ctx context.Context, resource string, identifier string) (*Reservation, error)
}

// Reservation 预约信息
type Reservation struct {
    OK        bool          // 是否成功
    Delay     time.Duration // 需要等待的时间
    CancelFunc func()       // 取消预约
}

// RuleManager 规则管理器接口
type RuleManager interface {
    // AddRule 添加限流规则
    AddRule(rule *RateLimitRule) error

    // UpdateRule 更新限流规则
    UpdateRule(rule *RateLimitRule) error

    // DeleteRule 删除限流规则
    DeleteRule(resource string) error

    // GetRule 获取限流规则
    GetRule(resource string) (*RateLimitRule, error)

    // ListRules 列出所有规则
    ListRules() ([]*RateLimitRule, error)
}

package ratelimit

import (
    "context"
    "sync"
    "time"
)

// LocalRateLimiter 本地限流器（单机版）
type LocalRateLimiter struct {
    limiters map[string]*TokenBucketLimiter
    rules    map[string]*RateLimitRule
    mu       sync.RWMutex
}

// RateLimitRule 限流规则
type RateLimitRule struct {
    Resource  string        // 资源标识
    Algorithm string        // 算法类型
    Limit     int64         // 限流阈值
    Window    time.Duration // 时间窗口
    Burst     int64         // 突发容量
}

func NewLocalRateLimiter() *LocalRateLimiter {
    return &LocalRateLimiter{
        limiters: make(map[string]*TokenBucketLimiter),
        rules:    make(map[string]*RateLimitRule),
    }
}

// AddRule 添加限流规则
func (l *LocalRateLimiter) AddRule(rule *RateLimitRule) {
    l.mu.Lock()
    defer l.mu.Unlock()

    // 创建限流器
    limiter := NewTokenBucketLimiter(rule.Burst, rule.Limit)

    l.rules[rule.Resource] = rule
    l.limiters[rule.Resource] = limiter
}

// Allow 判断是否允许请求
func (l *LocalRateLimiter) Allow(ctx context.Context, resource string, identifier string) (bool, error) {
    l.mu.RLock()
    limiter, exists := l.limiters[resource]
    l.mu.RUnlock()

    if !exists {
        // 没有配置规则，默认允许
        return true, nil
    }

    // 组合key: resource + identifier
    // (简化实现，生产环境需要分别存储)
    return limiter.Allow(), nil
}

// 使用示例
func ExampleLocalRateLimiter() {
    limiter := NewLocalRateLimiter()

    // 添加规则：订单API每秒100次
    limiter.AddRule(&RateLimitRule{
        Resource:  "/api/order/create",
        Algorithm: "token_bucket",
        Limit:     100,
        Window:    time.Second,
        Burst:     150,
    })

    ctx := context.Background()

    // 模拟请求
    for i := 0; i < 200; i++ {
        allowed, _ := limiter.Allow(ctx, "/api/order/create", "user:12345")
        if allowed {
            println("Request", i, "allowed")
        } else {
            println("Request", i, "rejected")
        }
    }
}

package ratelimit

import (
    "context"
    "errors"
    "time"

    "github.com/go-redis/redis/v8"
)

// RedisRateLimiter Redis分布式限流器
type RedisRateLimiter struct {
    client *redis.Client
    script *redis.Script  // Lua脚本
}

// Lua脚本：令牌桶算法
const tokenBucketScript = `
-- KEYS[1]: 限流key
-- ARGV[1]: 令牌生成速率(个/秒)
-- ARGV[2]: 桶容量
-- ARGV[3]: 当前时间戳(秒)
-- ARGV[4]: 消耗令牌数

local rate = tonumber(ARGV[1])
local capacity = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local requested = tonumber(ARGV[4])

-- 获取上次更新时间和令牌数
local last_time = tonumber(redis.call('HGET', KEYS[1], 'last_time') or now)
local tokens = tonumber(redis.call('HGET', KEYS[1], 'tokens') or capacity)

-- 计算新生成的令牌数
local elapsed = math.max(0, now - last_time)
local new_tokens = math.min(capacity, tokens + elapsed * rate)

-- 判断令牌是否足够
if new_tokens >= requested then
    -- 扣减令牌
    new_tokens = new_tokens - requested

    -- 更新Redis
    redis.call('HSET', KEYS[1], 'tokens', new_tokens)
    redis.call('HSET', KEYS[1], 'last_time', now)
    redis.call('EXPIRE', KEYS[1], 3600)  -- 1小时过期

    return {1, new_tokens}  -- 允许，返回剩余令牌
else
    return {0, new_tokens}  -- 拒绝，返回当前令牌
end
`

func NewRedisRateLimiter(client *redis.Client) *RedisRateLimiter {
    return &RedisRateLimiter{
        client: client,
        script: redis.NewScript(tokenBucketScript),
    }
}

// Allow 判断是否允许请求
func (r *RedisRateLimiter) Allow(ctx context.Context, resource string, identifier string, rate int64, capacity int64) (bool, error) {
    // 构造限流key
    key := "ratelimit:token:" + resource + ":" + identifier

    now := time.Now().Unix()

    // 执行Lua脚本（原子操作）
    result, err := r.script.Run(ctx, r.client, []string{key}, rate, capacity, now, 1).Result()
    if err != nil {
        return false, err
    }

    // 解析结果
    values, ok := result.([]interface{})
    if !ok || len(values) != 2 {
        return false, errors.New("invalid script result")
    }

    allowed := values[0].(int64) == 1
    return allowed, nil
}

// AllowN 判断是否允许n个请求
func (r *RedisRateLimiter) AllowN(ctx context.Context, resource string, identifier string, rate int64, capacity int64, n int64) (bool, error) {
    key := "ratelimit:token:" + resource + ":" + identifier
    now := time.Now().Unix()

    result, err := r.script.Run(ctx, r.client, []string{key}, rate, capacity, now, n).Result()
    if err != nil {
        return false, err
    }

    values, ok := result.([]interface{})
    if !ok || len(values) != 2 {
        return false, errors.New("invalid script result")
    }

    allowed := values[0].(int64) == 1
    return allowed, nil
}

// 使用示例
func ExampleRedisRateLimiter() {
    // 连接Redis
    client := redis.NewClient(&redis.Options{
        Addr: "localhost:6379",
    })

    limiter := NewRedisRateLimiter(client)
    ctx := context.Background()

    // 限流参数：每秒100个令牌，桶容量150
    rate := int64(100)
    capacity := int64(150)

    // 模拟请求
    for i := 0; i < 200; i++ {
        allowed, err := limiter.Allow(ctx, "/api/order/create", "user:12345", rate, capacity)
        if err != nil {
            println("Error:", err.Error())
            continue
        }

        if allowed {
            println("Request", i, "allowed")
        } else {
            println("Request", i, "rejected")
        }
    }
}

package ratelimit

import (
    "context"
    "time"
)

// MultiLevelRateLimiter 多级限流器
type MultiLevelRateLimiter struct {
    local  *LocalRateLimiter
    redis  *RedisRateLimiter
    config *MultiLevelConfig
}

// MultiLevelConfig 多级限流配置
type MultiLevelConfig struct {
    // 本地限流配置
    LocalLimit    int64   // 本地限流阈值
    LocalCapacity int64   // 本地突发容量

    // Redis限流配置
    RedisLimit    int64   // 全局限流阈值
    RedisCapacity int64   // 全局突发容量

    // 策略
    LocalFirst    bool    // 优先本地限流
    RedisRequired bool    // Redis必须检查
}

func NewMultiLevelRateLimiter(local *LocalRateLimiter, redis *RedisRateLimiter, config *MultiLevelConfig) *MultiLevelRateLimiter {
    return &MultiLevelRateLimiter{
        local:  local,
        redis:  redis,
        config: config,
    }
}

// Allow 判断是否允许请求
func (m *MultiLevelRateLimiter) Allow(ctx context.Context, resource string, identifier string) (bool, error) {
    // 第一级：本地限流（快速拦截）
    localAllowed, err := m.local.Allow(ctx, resource, identifier)
    if err != nil {
        return false, err
    }

    if !localAllowed {
        // 本地限流拦截
        return false, nil
    }

    // 第二级：Redis限流（精准控制）
    if m.config.RedisRequired {
        redisAllowed, err := m.redis.Allow(ctx, resource, identifier, m.config.RedisLimit, m.config.RedisCapacity)
        if err != nil {
            // Redis故障，降级到仅本地限流
            return localAllowed, nil
        }
        return redisAllowed, nil
    }

    return true, nil
}

// 使用示例
func ExampleMultiLevelRateLimiter() {
    // 初始化本地限流器
    local := NewLocalRateLimiter()
    local.AddRule(&RateLimitRule{
        Resource:  "/api/order/create",
        Algorithm: "token_bucket",
        Limit:     15,   // 单机限15（假设10台机器，总限100，每台分配10+冗余5）
        Window:    time.Second,
        Burst:     20,
    })

    // 初始化Redis限流器
    redis := NewRedisRateLimiter(nil) // 假设已连接

    // 配置多级限流
    config := &MultiLevelConfig{
        LocalLimit:    15,
        LocalCapacity: 20,
        RedisLimit:    100,
        RedisCapacity: 150,
        LocalFirst:    true,
        RedisRequired: true,
    }

    limiter := NewMultiLevelRateLimiter(local, redis, config)

    ctx := context.Background()

    // 模拟请求
    for i := 0; i < 200; i++ {
        allowed, err := limiter.Allow(ctx, "/api/order/create", "user:12345")
        if err != nil {
            println("Error:", err.Error())
            continue
        }

        if allowed {
            println("Request", i, "allowed")
        } else {
            println("Request", i, "rejected")
        }
    }
}

package ratelimit

import (
    "context"
    "fmt"
    "sync"
    "time"

    "github.com/go-redis/redis/v8"
)

// 完整的限流系统实现
type RateLimitSystem struct {
    ruleManager  *RuleManager
    localLimiter *LocalRateLimiter
    redisLimiter *RedisRateLimiter
    config       *SystemConfig
    mu           sync.RWMutex
}

type SystemConfig struct {
    Mode          string  // "local", "redis", "hybrid"
    EnableMetrics bool    // 是否启用指标
    EnableLog     bool    // 是否启用日志
}

// Check 检查限流
func (s *RateLimitSystem) Check(ctx context.Context, resource string, identifier string) (*CheckResult, error) {
    // 获取规则
    rule, err := s.ruleManager.GetRule(resource)
    if err != nil || rule == nil {
        // 没有规则，默认允许
        return &CheckResult{Allowed: true}, nil
    }

    if !rule.Enabled {
        return &CheckResult{Allowed: true, Reason: "rule disabled"}, nil
    }

    start := time.Now()

    var allowed bool
    var remaining int64

    // 根据模式选择限流器
    switch s.config.Mode {
    case "local":
        allowed, err = s.localLimiter.Allow(ctx, resource, identifier)
    case "redis":
        allowed, err = s.redisLimiter.Allow(ctx, resource, identifier, rule.Limit, rule.Burst)
    case "hybrid":
        // 多级限流
        localAllowed, _ := s.local Limiter.Allow(ctx, resource, identifier)
        if !localAllowed {
            allowed = false
        } else {
            allowed, err = s.redisLimiter.Allow(ctx, resource, identifier, rule.Limit, rule.Burst)
        }
    }

    latency := time.Since(start)

    // 记录指标
    if s.config.EnableMetrics {
        s.recordMetrics(resource, allowed, latency)
    }

    result := &CheckResult{
        Allowed:   allowed,
        Remaining: remaining,
        ResetTime: time.Now().Add(rule.Window),
        Latency:   latency,
    }

    if !allowed {
        result.Reason = "rate limit exceeded"
        result.RetryAfter = rule.Window
    }

    return result, err
}

type CheckResult struct {
    Allowed    bool          // 是否允许
    Remaining  int64         // 剩余配额
    ResetTime  time.Time     // 重置时间
    RetryAfter time.Duration // 重试等待时间
    Reason     string        // 原因
    Latency    time.Duration // 延迟
}

// 指标记录
func (s *RateLimitSystem) recordMetrics(resource string, allowed bool, latency time.Duration) {
    // 使用Prometheus记录指标
    // rateLimitRequestsTotal.WithLabelValues(resource, fmt.Sprint(allowed)).Inc()
    // rateLimitLatency.WithLabelValues(resource).Observe(latency.Seconds())
}

package middleware

import (
    "net/http"
    "ratelimit"
)

// RateLimitMiddleware HTTP限流中间件
func RateLimitMiddleware(limiter *ratelimit.RateLimitSystem) func(http.Handler) http.Handler {
    return func(next http.Handler) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
            ctx := r.Context()

            // 获取资源标识（API路径）
            resource := r.URL.Path

            // 获取用户标识（从请求中提取）
            identifier := extractIdentifier(r) // 如: user ID 或 IP

            // 检查限流
            result, err := limiter.Check(ctx, resource, identifier)
            if err != nil {
                http.Error(w, "Internal Server Error", http.StatusInternalServerError)
                return
            }

            if !result.Allowed {
                // 限流触发，返回429
                w.Header().Set("X-RateLimit-Limit", "100")
                w.Header().Set("X-RateLimit-Remaining", "0")
                w.Header().Set("X-RateLimit-Reset", result.ResetTime.Format(time.RFC3339))
                w.Header().Set("Retry-After", fmt.Sprintf("%d", int(result.RetryAfter.Seconds())))

                http.Error(w, "Too Many Requests", http.StatusTooManyRequests)
                return
            }

            // 添加响应头
            w.Header().Set("X-RateLimit-Limit", "100")
            w.Header().Set("X-RateLimit-Remaining", fmt.Sprintf("%d", result.Remaining))

            // 继续处理请求
            next.ServeHTTP(w, r)
        })
    }
}

// extractIdentifier 提取用户标识
func extractIdentifier(r *http.Request) string {
    // 优先使用用户ID
    if userID := r.Header.Get("X-User-ID"); userID != "" {
        return "user:" + userID
    }

    // 其次使用IP
    ip := r.Header.Get("X-Real-IP")
    if ip == "" {
        ip = r.Header.Get("X-Forwarded-For")
    }
    if ip == "" {
        ip = r.RemoteAddr
    }

    return "ip:" + ip
}

// 使用示例
func main() {
    // 初始化限流系统
    limiter := ratelimit.NewRateLimitSystem(config)

    // 创建HTTP服务
    mux := http.NewServeMux()
    mux.HandleFunc("/api/order/create", handleOrderCreate)

    // 应用限流中间件
    handler := RateLimitMiddleware(limiter)(mux)

    http.ListenAndServe(":8080", handler)
}

package interceptor

import (
    "context"
    "google.golang.org/grpc"
    "google.golang.org/grpc/codes"
    "google.golang.org/grpc/metadata"
    "google.golang.org/grpc/status"
    "ratelimit"
)

// RateLimitInterceptor gRPC限流拦截器
func RateLimitInterceptor(limiter *ratelimit.RateLimitSystem) grpc.UnaryServerInterceptor {
    return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
        // 获取方法名作为资源标识
        resource := info.FullMethod

        // 从metadata获取用户标识
        identifier := extractUserFromMetadata(ctx)

        // 检查限流
        result, err := limiter.Check(ctx, resource, identifier)
        if err != nil {
            return nil, status.Error(codes.Internal, "rate limit check failed")
        }

        if !result.Allowed {
            // 设置响应元数据
            md := metadata.Pairs(
                "x-ratelimit-limit", "100",
                "x-ratelimit-remaining", "0",
                "x-ratelimit-reset", result.ResetTime.Format(time.RFC3339),
            )
            grpc.SetTrailer(ctx, md)

            return nil, status.Error(codes.ResourceExhausted, "rate limit exceeded")
        }

        // 继续处理请求
        return handler(ctx, req)
    }
}

func extractUserFromMetadata(ctx context.Context) string {
    md, ok := metadata.FromIncomingContext(ctx)
    if !ok {
        return "unknown"
    }

    if userIDs := md.Get("user-id"); len(userIDs) > 0 {
        return "user:" + userIDs[0]
    }

    return "unknown"
}

// HotKeyOptimizer 热点key优化器
type HotKeyOptimizer struct {
    localCache  *sync.Map  // 本地缓存热点数据
    hotKeyStats map[string]int64
    threshold   int64      // 热点阈值
    mu          sync.RWMutex
}

// 检测热点key
func (h *HotKeyOptimizer) isHotKey(key string) bool {
    h.mu.RLock()
    count := h.hotKeyStats[key]
    h.mu.RUnlock()

    return count > h.threshold
}

// 优化后的限流检查
func (s *RateLimitSystem) CheckWithHotKeyOptimization(ctx context.Context, resource string, identifier string) (*CheckResult, error) {
    key := resource + ":" + identifier

    // 检测是否为热点key
    if s.hotKeyOptimizer.isHotKey(key) {
        // 使用本地缓存
        if cachedResult, ok := s.hotKeyOptimizer.localCache.Load(key); ok {
            return cachedResult.(*CheckResult), nil
        }
    }

    // 正常限流检查
    result, err := s.Check(ctx, resource, identifier)

    // 缓存热点key结果
    if s.hotKeyOptimizer.isHotKey(key) {
        s.hotKeyOptimizer.localCache.Store(key, result)
    }

    return result, err
}

// BatchCheck 批量检查限流
func (s *RateLimitSystem) BatchCheck(ctx context.Context, checks []CheckRequest) ([]CheckResult, error) {
    results := make([]CheckResult, len(checks))

    // 使用Pipeline批量执行Redis命令
    pipe := s.redisLimiter.client.Pipeline()

    for i, check := range checks {
        key := "ratelimit:token:" + check.Resource + ":" + check.Identifier

        // 添加到Pipeline
        pipe.HGetAll(ctx, key)
    }

    // 执行Pipeline
    cmds, err := pipe.Exec(ctx)
    if err != nil {
        return nil, err
    }

    // 解析结果
    for i, cmd := range cmds {
        // 处理每个命令的结果
        // ...
        results[i] = CheckResult{Allowed: true}
    }

    return results, nil
}

// WarmupConfig 预热配置
type WarmupConfig struct {
    Duration     time.Duration // 预热时长
    InitialLimit int64         // 初始限流值
    TargetLimit  int64         // 目标限流值
}

// CalculateLimit 计算预热期间的限流值
func (w *WarmupConfig) CalculateLimit(elapsedTime time.Duration) int64 {
    if elapsedTime >= w.Duration {
        return w.TargetLimit
    }

    // 线性增长
    ratio := float64(elapsedTime) / float64(w.Duration)
    currentLimit := float64(w.InitialLimit) + ratio*float64(w.TargetLimit-w.InitialLimit)

    return int64(currentLimit)
}

// 使用示例
func ExampleWarmup() {
    config := &WarmupConfig{
        Duration:     5 * time.Minute, // 5分钟预热
        InitialLimit: 10,              // 初始10 QPS
        TargetLimit:  100,             // 目标100 QPS
    }

    startTime := time.Now()

    // 每秒计算当前限流值
    for i := 0; i < 300; i++ {
        elapsed := time.Duration(i) * time.Second
        currentLimit := config.CalculateLimit(elapsed)

        println("Time:", elapsed, "Limit:", currentLimit)

        time.Sleep(time.Second)
    }
}

// CircuitBreaker 熔断器
type CircuitBreaker struct {
    maxFailures  int           // 最大失败次数
    resetTimeout time.Duration // 重置超时
    failures     int           // 当前失败次数
    lastFailTime time.Time     // 最后失败时间
    state        string        // 状态: "closed", "open", "half-open"
    mu           sync.Mutex
}

func NewCircuitBreaker(maxFailures int, resetTimeout time.Duration) *CircuitBreaker {
    return &CircuitBreaker{
        maxFailures:  maxFailures,
        resetTimeout: resetTimeout,
        state:        "closed",
    }
}

// Call 执行调用（带熔断）
func (cb *CircuitBreaker) Call(fn func() error) error {
    cb.mu.Lock()
    defer cb.mu.Unlock()

    // 检查状态
    if cb.state == "open" {
        // 检查是否可以进入半开状态
        if time.Since(cb.lastFailTime) > cb.resetTimeout {
            cb.state = "half-open"
        } else {
            return fmt.Errorf("circuit breaker is open")
        }
    }

    // 执行调用
    err := fn()

    if err != nil {
        // 记录失败
        cb.failures++
        cb.lastFailTime = time.Now()

        if cb.failures >= cb.maxFailures {
            cb.state = "open"
        }

        return err
    }

    // 成功，重置计数器
    cb.failures = 0
    cb.state = "closed"

    return nil
}

// 集成到限流系统
func (s *RateLimitSystem) CheckWithCircuitBreaker(ctx context.Context, resource string, identifier string) (*CheckResult, error) {
    var result *CheckResult
    var err error

    // 使用熔断器保护Redis调用
    cbErr := s.circuitBreaker.Call(func() error {
        result, err = s.redisLimiter.Allow(ctx, resource, identifier, 100, 150)
        return err
    })

    if cbErr != nil {
        // 熔断器打开，降级到本地限流
        result, err = s.localLimiter.Allow(ctx, resource, identifier)
    }

    return result, err
}

package metrics

import (
    "github.com/prometheus/client_golang/prometheus"
    "github.com/prometheus/client_golang/prometheus/promauto"
)

var (
    // 限流请求总数
    RateLimitRequestsTotal = promauto.NewCounterVec(
        prometheus.CounterOpts{
            Name: "rate_limit_requests_total",
            Help: "Total number of rate limit checks",
        },
        []string{"resource", "allowed"},
    )

    // 限流延迟
    RateLimitLatency = promauto.NewHistogramVec(
        prometheus.HistogramOpts{
            Name:    "rate_limit_latency_seconds",
            Help:    "Latency of rate limit checks",
            Buckets: prometheus.DefBuckets,
        },
        []string{"resource"},
    )

    // Redis操作总数
    RedisOperationsTotal = promauto.NewCounterVec(
        prometheus.CounterOpts{
            Name: "redis_operations_total",
            Help: "Total number of Redis operations",
        },
        []string{"operation", "status"},
    )

    // 本地缓存命中率
    LocalCacheHits = promauto.NewCounter(
        prometheus.CounterOpts{
            Name: "local_cache_hits_total",
            Help: "Total number of local cache hits",
        },
    )

    LocalCacheMisses = promauto.NewCounter(
        prometheus.CounterOpts{
            Name: "local_cache_misses_total",
            Help: "Total number of local cache misses",
        },
    )

    // 熔断器状态
    CircuitBreakerState = promauto.NewGaugeVec(
        prometheus.GaugeOpts{
            Name: "circuit_breaker_state",
            Help: "Circuit breaker state (0=closed, 1=open, 2=half-open)",
        },
        []string{"resource"},
    )
)

// 记录限流检查
func RecordRateLimitCheck(resource string, allowed bool, latency time.Duration) {
    RateLimitRequestsTotal.WithLabelValues(resource, fmt.Sprint(allowed)).Inc()
    RateLimitLatency.WithLabelValues(resource).Observe(latency.Seconds())
}