Golang API限流实现 令牌桶算法应用

golang.org/x/time/rate

rate.Limiter

Limiter

limiter := rate.NewLimiter(rate.Every(time.Second/10), 100)

package main

import (
    "fmt"
    "log"
    "net/http"
    "time"

    "golang.org/x/time/rate"
)

// 创建一个全局限流器，每秒允许10个请求，桶容量为20
var globalLimiter = rate.NewLimiter(rate.Every(time.Second/10), 20)

func rateLimitMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        // 尝试获取令牌，如果无法获取，则返回429
        if !globalLimiter.Allow() {
            http.Error(w, "Too Many Requests", http.StatusTooManyRequests)
            return
        }
        next.ServeHTTP(w, r)
    })
}

func helloHandler(w http.ResponseWriter, r *http.Request) {
    fmt.Fprintf(w, "Hello, you are not rate limited!")
}

func main() {
    http.Handle("/hello", rateLimitMiddleware(http.HandlerFunc(helloHandler)))
    log.Fatal(http.ListenAndServe(":8080", nil))
}

rate.NewLimiter

map

*rate.Limiter

package main

import (
    "fmt"
    "log"
    "net/http"
    "sync"
    "time"

    "golang.org/x/time/rate"
)

// clientLimiterStore 存储每个客户端的限流器
type clientLimiterStore struct {
    limiters map[string]*rate.Limiter
    mu       sync.Mutex
    // 默认限流参数
    r rate.Limit
    b int
}

func newClientLimiterStore(r rate.Limit, b int) *clientLimiterStore {
    return &clientLimiterStore{
        limiters: make(map[string]*rate.Limiter),
        r:        r,
        b:        b,
    }
}

// GetLimiter 为指定客户端获取或创建一个限流器
func (s *clientLimiterStore) GetLimiter(clientID string) *rate.Limiter {
    s.mu.Lock()
    defer s.mu.Unlock()

    limiter, exists := s.limiters[clientID]
    if !exists {
        limiter = rate.NewLimiter(s.r, s.b)
        s.limiters[clientID] = limiter
    }
    return limiter
}

// 清理不活跃的限流器，防止内存泄漏
func (s *clientLimiterStore) CleanupOldLimiters() {
    ticker := time.NewTicker(5 * time.Minute) // 每5分钟清理一次
    defer ticker.Stop()

    for range ticker.C {
        s.mu.Lock()
        for clientID, limiter := range s.limiters {
            // 这是一个简化的清理逻辑，实际可能需要更复杂的判断
            // 比如：如果某个limiter在N分钟内没有被访问过，则删除
            // 这里只是演示，实际生产不建议直接删除，需要记录上次访问时间
            _ = limiter // 占位符，避免unused变量警告
            // 假设我们有一个机制来判断哪些limiter不再需要
            // delete(s.limiters, clientID)
        }
        s.mu.Unlock()
        log.Println("Cleaned up old limiters (simplified logic)")
    }
}

// 初始化客户端限流存储，每秒5个请求，桶容量10
var clientLimiters = newClientLimiterStore(rate.Every(time.Second/5), 10)

func init() {
    go clientLimiters.CleanupOldLimiters() // 启动后台清理goroutine
}

func perClientRateLimitMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        // 实际应用中，clientID可以是用户ID、API Key或客户端IP
        // 这里简化为从查询参数获取，实际应从认证信息或请求头获取
        clientID := r.URL.Query().Get("client_id")
        if clientID == "" {
            clientID = r.RemoteAddr // 如果没有client_id，则使用IP作为标识
        }

        limiter := clientLimiters.GetLimiter(clientID)

        // 阻塞等待令牌，最多等待1秒
        ctx, cancel := context.WithTimeout(r.Context(), time.Second)
        defer cancel()

        if err := limiter.Wait(ctx); err != nil {
            // 如果等待超时或者context被取消
            http.Error(w, "Too Many Requests or Request Timeout", http.StatusTooManyRequests)
            return
        }
        next.ServeHTTP(w, r)
    })
}

func main() {
    http.Handle("/api/data", perClientRateLimitMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        fmt.Fprintf(w, "Data for client %s, not rate limited!", r.URL.Query().Get("client_id"))
    })))
    log.Fatal(http.ListenAndServe(":8081", nil))
}

map

map

rate.Limiter

-- KEYS[1]: key for last_refill_time
-- KEYS[2]: key for current_tokens
-- ARGV[1]: bucket_capacity
-- ARGV[2]: refill_rate_per_second
-- ARGV[3]: current_timestamp_in_ms
-- ARGV[4]: tokens_to_consume (usually 1)

local last_refill_time = tonumber(redis.call('GET', KEYS[1]) or "0")
local current_tokens = tonumber(redis.call('GET', KEYS[2]) or "0")
local bucket_capacity = tonumber(ARGV[1])
local refill_rate_per_second = tonumber(ARGV[2])
local current_timestamp_in_ms = tonumber(ARGV[3])
local tokens_to_consume = tonumber(ARGV[4])

-- 计算需要补充的令牌
local time_passed_ms = current_timestamp_in_ms - last_refill_time
local tokens_to_add = (time_passed_ms / 1000) * refill_rate_per_second

-- 更新令牌数量，不能超过桶容量
current_tokens = math.min(bucket_capacity, current_tokens + tokens_to_add)

-- 检查是否可以消费令牌
if current_tokens >= tokens_to_consume then
    current_tokens = current_tokens - tokens_to_consume
    redis.call('SET', KEYS[1], current_timestamp_in_ms)
    redis.call('SET', KEYS[2], current_tokens)
    return 1 -- 允许
else
    return 0 -- 拒绝
end