golang官方限流器rate

包名：golang.org/x/time/rate

实现原理：令牌桶

package main

import (
	"context"
	"fmt"
	"testing"
	"time"

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

func TestLimiter(t *testing.T) {
	// 第一个参数代表速率：Every代表每xx时间产生一个，比如这里的每0.1秒生产一个，换算一下就是每秒10个，
	// 第二个参数代表桶的大小：即，如果没有人来消费，那么桶里最多存着10个令牌，
	// 桶的初始状态是满的，后面才会按照速率来提供令牌，这一点很重要，
	l := rate.NewLimiter(rate.Every(time.Second/10), 10)

	for i := 0; i < 10; i++ {
		go func(i int) {
			for {
				if l.Allow() {
					fmt.Printf("allow %d\n", i)
				}
				// 每0.5秒请求一次
				time.Sleep(time.Second / 2)
			}
		}(i)
	}
	time.Sleep(time.Second * 10)
}

func TestLimiter2(t *testing.T) {
	l := rate.NewLimiter(rate.Every(time.Second), 10)

	for i := 0; i < 10; i++ {
		go func(i int) {
			for {
				ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(1*time.Second))

				// 等待获取令牌，可以传入 deadline context 设置最大等待时间
				// 这样在某些场景下我们可以让请求等待一会，而不是直接失败。
				if err := l.Wait(ctx); err != nil {
					fmt.Printf("timwout: %v\n", err)
				} else {
					fmt.Printf("allow %d\n", i)
				}
				cancel()
				time.Sleep(time.Second / 2)
			}
		}(i)
	}
	time.Sleep(time.Second * 10)
}

func TestLimiter3(t *testing.T) {
	l := rate.NewLimiter(rate.Every(time.Second), 10)

	for i := 0; i < 10; i++ {
		go func(i int) {
			for {
				// 先预留令牌，到指定时间不再需要去获取令牌，直接执行操作
				// 如果预留的令牌不想使用了，也可以使用r.Cancel()归还已预留的令牌
				if r := l.Reserve(); r.OK() {
					// 休眠直到令牌生效
					time.Sleep(r.Delay())
					fmt.Printf("allow %d\n", i)
				}
				time.Sleep(time.Second / 2)
			}
		}(i)
	}
	time.Sleep(time.Second * 10)
}

封装到 gin 中间件

func Limiter(l *rate.Limiter) gin.HandlerFunc {
	return func(c *gin.Context) {
		if !l.Allow() {
			c.AbortWithStatus(http.StatusTooManyRequests)
		}

		c.Next()
	}
}