限流算法—go rate 包

• • 前言
  • 代码解释
  • • WaitN方法
  • 代码示例
  • • AllowN方法
  • 实战限流
  • 整合到Go-kit微服务

一直是微服务中比较重要的一点，常见的有令牌桶和漏桶算法两种。

Go内置的rate包，提供了令牌桶的实现，以此来实现限流

下面给出使用rate包实现限流的代码示例及解释。

r := rate.NewLimiter(1, 5) 

第一个参数是每秒生成的令牌个数，第二个参数是初始的令牌个数

err := r.WaitN(ctx, 2)

waitN方法，第二个参数表示执行一次业务逻辑需要消耗的令牌个数。

import (
	"context"
	"fmt"
	"golang.org/x/time/rate"
	"log"
	"time"
)

func main() {
	r := rate.NewLimiter(1, 5) //1表示每次放进筒内的数量，桶内的令牌数是5，最大令牌数也是5，这个筒子是自动补充的，你只要取了令牌不管你取多少个，这里都会在每次取完后自动加1个进来，因为我们设置的是1
	ctx := context.Background()

	for {
		err := r.WaitN(ctx, 2)//每次消耗2个，放入一个，消耗完了还会放进去，因为初始是5个，所以这段代码再执行到第4次的时候筒里面就空了，如果当前不够取两个了，本次就不取，再放一个进去，然后返回false
		if err != nil {
			log.Fatal(err)
		}
		fmt.Println(time.Now().Format("2021-11-02 15:04:05"))
		time.Sleep(time.Second)
	}

}

AllowN 方法表示，截止到某一时刻，目前桶中数目是否至少为 n 个，满足则返回 true，同时从桶中消费 n 个 token。反之不消费桶中的Token，返回false。

此外还有Allow方法，含义和作用等同于Allow(time.Now(),1)

import (
	"golang.org/x/time/rate"
	"net/http"
)

//1表示每次放进筒内的数量，桶内的令牌数是5，最大令牌数也是5，这个筒子是自动补充的，你只要取了令牌不管你取多少个，这里都会在每次取完后自动加1个进来，因为我们设置的是1
var r = rate.NewLimiter(1, 5) 
func Limiter(next http.Handler) http.Handler {
	//这里使用http.HandlerFunc进行类型转换把匿名函数转换成了type http.HandlerFunc，因为HandlerFunc实现了ServeHttp方法所以是http.Handler的实例
	return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { 
		if !r.Allow() { 
			http.Error(writer, "too many requests", http.StatusTooManyRequests)
		} else {
			// 处理完了之后，再调用next.ServerHTTP继续完成请求
			next.ServeHTTP(writer, request)
		}
	})
}

限流器参数和返回都是 http.Handler，实际上这就是对Handler进行了一层“过滤”，经典的go中间件写法和Java web中的拦截器写法。

整合到Go-kit微服务主要是通过中间件的形式实现

//加入限流功能中间件
func RateLimitMiddleWare(limit *rate.Limiter) endpoint.Middleware { //Middleware type Middleware func(Endpoint) Endpoint
	return func(next endpoint.Endpoint) endpoint.Endpoint { //Endpoint type Endpoint func(ctx context.Context, request interface{}) (response interface{}, err error)
		return func(ctx context.Context, request interface{}) (response interface{}, err error) {
			if !limit.Allow(){
				return nil,errors.New("too many request")
			}
			return next(ctx,request)
		}
	}
}

除此之外，go-kit中也内置了令牌桶实现