在Go语言中,实现一个支持限流的HTTP中间件是一个常见的需求,特别是在构建微服务或高并发的Web应用时。限流(Rate Limiting)是控制客户端请求频率的一种手段,以防止系统过载或被恶意攻击。
一个常见的限流算法是漏桶算法(Leaky Bucket)或令牌桶算法(Token Bucket)。在Go中,我们可以使用第三方库如golang.org/x/time/rate
来实现令牌桶算法。
以下是一个使用golang.org/x/time/rate
库实现支持限流的HTTP中间件的示例:
安装
golang.org/x/time/rate
库:go get -u golang.org/x/time/rate
编写限流中间件:
package main import ( "context" "net/http" "golang.org/x/time/rate" "time" ) // LimiterMiddleware 创建一个中间件,用于限制请求频率 func LimiterMiddleware(limiter *rate.Limiter) func(http.Handler) http.Handler { return func(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { // 尝试获取令牌 if limiter.Allow() { // 如果有令牌,则继续处理请求 next.ServeHTTP(w, r) } else { // 如果没有令牌,则返回HTTP 429 Too Many Requests http.Error(w, http.StatusText(http.StatusTooManyRequests), http.StatusTooManyRequests) } }) } } // 创建一个简单的HTTP处理器 func helloHandler(w http.ResponseWriter, r *http.Request) { w.Write([]byte("Hello, World!")) } func main() { // 创建一个每秒放入1个令牌,桶容量为5的令牌桶 limiter := rate.NewLimiter(1, 5) // 创建一个http.Handler handler := http.HandlerFunc(helloHandler) // 应用限流中间件 handlerWithLimiter := LimiterMiddleware(limiter)(handler) // 启动HTTP服务器 http.ListenAndServe(":8080", handlerWithLimiter) }
在这个示例中,LimiterMiddleware
函数接受一个*rate.Limiter
对象作为参数,并返回一个HTTP中间件函数。这个中间件函数检查每个进入的请求是否可以从令牌桶中获取到令牌。如果可以,则继续处理请求;如果不可以,则返回HTTP 429状态码,表示请求过多。
你可以根据需要调整rate.NewLimiter
函数的参数,以控制请求的速率和桶的容量。
这种中间件方式使得限流逻辑与具体的业务逻辑解耦,便于复用和维护。此外,由于golang.org/x/time/rate
库提供了灵活的令牌桶实现,你可以很容易地根据实际需求调整限流策略。