如何用Golang实现API限流 令牌桶算法实践

time

sync

golang.org/x/time/rate

package main

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

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

// 为每个用户或IP维护一个限流器
var visitors = make(map[string]*rate.Limiter)
var mtx sync.RWMutex

// 获取对应IP的限流器，每秒允许10个请求，桶容量为20
func getVisitorLimiter(ip string) *rate.Limiter {
    mtx.Lock()
    defer mtx.Unlock()

    limiter, exists := visitors[ip]
    if !exists {
        // 每秒生成10个令牌，桶最多存20个
        limiter = rate.NewLimiter(10, 20)
        visitors[ip] = limiter
    }
    return limiter
}

// 限流中间件
func rateLimit(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        ip := r.RemoteAddr // 实际使用中建议提取真实IP
        limiter := getVisitorLimiter(ip)

        if !limiter.Allow() {
            http.StatusText(http.StatusTooManyRequests)
            w.WriteHeader(http.StatusTooManyRequests)
            w.Write([]byte("Too many requests"))
            return
        }

        next.ServeHTTP(w, r)
    })
}

func main() {
    mux := http.NewServeMux()
    mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
        fmt.Fprintf(w, "Hello, your request is processed at %v", time.Now())
    })

    log.Println("Server starting on :8080")
    log.Fatal(http.ListenAndServe(":8080", rateLimit(mux)))
}

package main

import (
    "sync"
    "time"
)

type TokenBucket struct {
    capacity  int           // 桶容量
    tokens    int           // 当前令牌数
    rate      time.Duration // 添加一个令牌的时间间隔
    lastToken time.Time     // 上次添加令牌的时间
    mutex     sync.Mutex
}

// 创建令牌桶：每 `rate` 时间添加一个令牌，最多 `capacity` 个
func NewTokenBucket(capacity int, rate time.Duration) *TokenBucket {
    return &TokenBucket{
        capacity:  capacity,
        tokens:    capacity,
        rate:      rate,
        lastToken: time.Now(),
    }
}

// 获取令牌，成功返回 true
func (tb *TokenBucket) Allow() bool {
    tb.mutex.Lock()
    defer tb.mutex.Unlock()

    now := time.Now()
    // 计算应该补充多少令牌
    elapsed := now.Sub(tb.lastToken)
    newTokens := int(elapsed / tb.rate)
    if newTokens > 0 {
        tb.tokens = min(tb.capacity, tb.tokens+newTokens)
        tb.lastToken = now
    }

    if tb.tokens > 0 {
        tb.tokens--
        return true
    }
    return false
}

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

golang.org/x/time/rate