首页 > Golang > Go教程

Golang令牌桶算法实现API限流详解

时间：2025-08-15 23:13:50 409浏览收藏

哈喽！大家好，很高兴又见面了，我是golang学习网的一名作者，今天由我给大家带来一篇《Golang实现API限流：令牌桶算法详解》，本文主要会讲到等等知识点，希望大家一起学习进步，也欢迎大家关注、点赞、收藏、转发! 下面就一起来看看吧！

使用Golang实现API限流的推荐方法是采用golang.org/x/time/rate包实现令牌桶算法，该方法每秒生成10个令牌、桶容量为20，通过限流中间件对每个IP进行并发控制，结合sync.RWMutex保证并发安全，并可扩展支持动态配置与分布式限流，最终构建稳定的API防护体系。

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

在高并发场景下，API限流是保护后端服务稳定的重要手段。Golang 因其高并发支持和简洁的并发模型，非常适合实现高效的限流机制。其中，令牌桶算法（Token Bucket） 是一种常用且灵活的限流策略。本文将介绍如何使用 Golang 实现基于令牌桶算法的 API 限流。

什么是令牌桶算法？

令牌桶算法的核心思想是：

系统以固定速率向桶中添加令牌；
每次请求需要从桶中获取一个令牌；
如果桶中有足够的令牌，请求被放行；
如果令牌不足，请求被拒绝或等待。

相比漏桶算法，令牌桶允许一定程度的突发流量（burst），只要桶中还有令牌就可以通过，更加实用。

使用 Golang 实现令牌桶

我们可以使用标准库 time 和 sync 来实现一个简单的令牌桶，也可以借助 golang.org/x/time/rate 包（推荐）。

方法一：使用 `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)))
}

说明：

rate.NewLimiter(10, 20)：每秒生成10个令牌，最多积压20个。
limiter.Allow()：判断是否允许请求通过，返回布尔值。
使用 sync.RWMutex 保证并发安全。
可以按 IP、用户ID 或 API Key 做维度限流。

方法二：手动实现一个简单的令牌桶

如果你希望理解底层原理，可以自己实现一个基础版本。

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
}

使用方式与上面类似，可集成到 HTTP 中间件中。