go学习笔记——gin限流

如果想在Gin Web服务中实现限流功能，可以使用ulule/limiter，uber-go/ratelimit

1.ulule/limiter

ulule/limiter是一款支持分布式限流的框架，其可以在Redis中存储和共享限流状态，从而在分布式环境中实现一致的限流逻辑。

ulule/limiter基于令牌桶（Token Bucket）算法，因为允许累积令牌，若桶中有令牌，短时间可处理大量请求，所以可能会短时超限。令牌桶算法允许突发流量的业务，适用场景如 API 请求、流媒体加载。

其他基于令牌桶算法的限流框架还有 golang.org/x/time/rate

1.引用依赖

go get github.com/ulule/limiter/v3@v3.11.2

ulule/limiter可以和gin web框架集成，实现服务接收外部请求场景下的限流，或者服务请求外部API场景下的限流

2.使用Gin限流中间件实现服务接收外部请求场景下的限流

参考：https://github.com/ulule/limiter-examples/blob/master/gin/main.go

package middleware

import (
	"github.com/gin-gonic/gin"
	"github.com/redis/go-redis/v9"
	"github.com/ulule/limiter/v3"
	mgin "github.com/ulule/limiter/v3/drivers/middleware/gin"
	redisstore "github.com/ulule/limiter/v3/drivers/store/redis"
	"time"
)

// 在中间件中对接收的请求进行限流
func RateLimitMiddleWare(rdb *redis.Client) gin.HandlerFunc {
	// 定义限流规则：1 请求/10秒
	rate := limiter.Rate{
		Period: 10 * time.Second,
		Limit:  1,
	}
	// 创建限流存储
	store, err := redisstore.NewStoreWithOptions(rdb, limiter.StoreOptions{
		Prefix:   "rate_limiter", // Redis 键的前缀
	})
	if err != nil {
		panic(err)
	}
	// 创建限流实例
	instance := limiter.New(store, rate)
	middleware := mgin.NewMiddleware(instance)
	return middleware
}

在router中使用限流中间件

	// ratelimit
	rateLimitMiddleWare := middleware.RateLimitMiddleWare(redisClient)
	authGroup.Use(rateLimitMiddleWare)

限流效果，如果限流的话，接口将会返回429的状态码

curl -i --request GET 'http://172.17.0.229:18080/api/v1/user/3' \
--header 'Content-Type: application/json'
HTTP/1.1 429 Too Many Requests
Content-Type: text/plain; charset=utf-8
X-Ratelimit-Limit: 1
X-Ratelimit-Remaining: 0
X-Ratelimit-Reset: 1734772666
Date: Sat, 21 Dec 2024 09:17:38 GMT
Content-Length: 14

Limit exceeded

在redis中会存储限流状态，如下

其中默认的key是prefix+ip，value是API的请求次数

3.使用limiter的Reached方法实现服务请求外部API场景下的限流

// 对外部的API进行请求的时候限流
func CallExternalAPIWithRateLimit(limiter *limiter.Limiter, restClient *client.RestClient) error {
	ctx := context.Background()
	key := "rate_limiter:external_api"

	// 检查是否允许请求
	res, err := limiter.Get(ctx, key)
	if err != nil {
		return fmt.Errorf("error checking rate limit: %v", err)
	}

	if res.Reached {
		return fmt.Errorf("rate limit exceeded")
	}

	// 调用外部 API
	var result any
	response, err := restClient.RestyClient.R().SetResult(&result).Get("https://www.baidu.com")
	if err != nil {
		return err
	}
	fmt.Println(response.RawResponse.StatusCode)

	return nil
}

func NewRateLimiter(redisClient *redis.Client) (*limiter.Limiter, error) {
	store, err := redisstore.NewStoreWithOptions(redisClient, limiter.StoreOptions{
		Prefix:   "rate_limiter",
		MaxRetry: 3,
	})
	if err != nil {
		return nil, err
	}

	// 定义限流规则
	rate := limiter.Rate{
		Period: 10 * time.Second, // 每秒
		Limit:  1,              // 最大10次请求
	}

	return limiter.New(store, rate), nil
}

service层中，在调用data层之前添加调用外部API接口的代码

func (s *UserService) FindByID(ctx context.Context, id int64) (*model.User, error) {
	err := middleware.CallExternalAPIWithRateLimit(s.limiter, s.restClient)
	if err != nil {
		return nil, err
	}
	u, err := s.uRepo.FindByID(ctx, id)
	if err != nil {
		return nil, err
	}
	return u, nil
}

如果触发限流条件，返回

{
    "code": 500,
    "msg": "find user by id fail",
    "data": "rate limit exceeded"
}

redis中存储的key

未触发限流，查询正常返回

{
    "code": 200,
    "msg": "find user by id success",
    "data": {
        "id": 1,
        "username": "test",
        "email": "test@test",
        "department_id": 1
    }
}

2.uber-go/ratelimit

uber-go/ratelimit是uber提供的Leaky Bucket（漏桶）限流算法的Golang实现，其不支持分布式限流。

Leaky Bucket（漏桶）限流算法 始终匀速处理请求，即使请求突增，速率也不会提高，所以不会超限。适用场景：需要严格限流的业务，如音视频播放、网络流控。

uber-go/ratelimit的使用例子如下

因为限制了每秒的请求数量最多是100个，令牌以1s/100=10ms的速度生成，take()的时候获得令牌，所以2次获得令牌的时间差是10ms，没有获取令牌的时候阻塞。

import (
	"fmt"
	"time"

	"go.uber.org/ratelimit"
)

func main() {
    rl := ratelimit.New(100) // per second

    prev := time.Now()
    for i := 0; i < 10; i++ {
        now := rl.Take()
        fmt.Println(i, now.Sub(prev))
        prev = now
    }

    // Output:
    // 0 0
    // 1 10ms
    // 2 10ms
    // 3 10ms
    // 4 10ms
    // 5 10ms
    // 6 10ms
    // 7 10ms
    // 8 10ms
    // 9 10ms
}