Go语言使用漏桶算法和令牌桶算法来实现API限流_Golang

在现代 web 应用程序中，流量的突增是不可避免的。为防止服务器被过多的请求压垮，限流（rate limiting） 是一个至关重要的技术手段。

本文将通过 go 语言的 gin 框架，演示如何使用漏桶算法和令牌桶算法来实现 api 的限流。

限流的意义

限流的主要目的是保护系统资源，防止因请求量过大导致服务器崩溃。同时，它也能防止恶意用户对系统的攻击，确保服务的稳定性和可用性。

两种常见的限流算法

1.漏桶算法（leaky bucket）

漏桶算法将请求视为水滴，水滴先进入桶中，然后以固定的速率从桶中流出。如果请求的速率超过了桶的流出速率，多余的请求将会被丢弃。

这个算法的优点很明显，就是让请求非常稳定，但是缺点也很明显，因为请求非常稳定，就不适于一些秒杀等一些可能在某一段时间会有洪峰流量的场景。不太好适情况控制流量的进入。

2.令牌桶算法（token bucket）

令牌桶算法中，系统会以固定的速率向桶中加入令牌，每个请求需要获取一个令牌才能执行。如果桶中没有足够的令牌，请求将被拒绝。

代码实现

在这个示例中，我们将展示如何在 gin 框架中应用这两种算法来实现 api 的限流。

package main

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

	"github.com/gin-gonic/gin"
	ratelimit2 "github.com/juju/ratelimit" // 令牌桶算法
	ratelimit1 "go.uber.org/ratelimit"     // 漏桶算法
)

func pinghandler(c *gin.context) {
	c.json(200, gin.h{
		"message": "pong",
	})
}

func pinghandler2(c *gin.context) {
	c.json(200, gin.h{
		"message": "pong2",
	})
}

// ratelimit1 使用漏桶算法来限制请求速率
func ratelimit1() func(ctx *gin.context) {
	// 漏桶算法，第一个参数为两滴水滴之间的时间间隔。
	// 此时表示两滴水之间的时间间隔是 100 纳秒
	rl := ratelimit1.new(100)

	return func(ctx *gin.context) {
		// 尝试取出水滴
		if waittime := rl.take().sub(time.now()); waittime > 0 {
			fmt.printf("需要等待 %v 秒，下一滴水才会滴下来\n", waittime)
			// 这里我们可以让程序继续等待，也可以直接拒绝掉
			// time.sleep(waittime)
			ctx.string(http.statusok, "rate limit, try again later")
			ctx.abort()
			return
		}
		// 证明可以继续执行
		ctx.next()
	}
}

// ratelimit2 使用令牌桶算法来限制请求速率
func ratelimit2() func(ctx *gin.context) {
	// 令牌桶算法：第一个参数为每秒填充令牌的速率为多少
	// 第二个参数为令牌桶的容量
	// 这里表示每秒填充 10 个令牌
	rl := ratelimit2.newbucket(time.second, 10)

	return func(ctx *gin.context) {
		// 尝试取出令牌
		var num int64 = 1
		// 这里表示需要 num 个令牌和已经取出的令牌数是否相等
		// 不相等，则表示超过了限流
                // 比如，假设每一个请求过来消耗2个令牌，但是从桶中取出的令牌个数为 1 ，那么则认为超过了限流（一般而言是一个请求消耗一个令牌，这里仅为举例）
		if rl.takeavailable(num) != num {
			// 此次没有取到令牌，说明超过了限流
			ctx.string(http.statusok, "rate limit, try again later")
			ctx.abort()
			return
		}
		// 证明可以继续执行
		ctx.next()
	}
}

func main() {
	r := gin.default()

	// 漏桶算法限流
	r.get("/ping", ratelimit1(), pinghandler)

	// 令牌桶算法限流
	r.get("/ping2", ratelimit2(), pinghandler2)

	r.run()
}