使用gin封装一个web脚手架(八)：限流器

一位不愿意透露姓名的杨先生 · · 861 次点击 · · 开始浏览

这是一个创建于的文章，其中的信息可能已经有所发展或是发生改变。

日常开发中，对于某些接口有请求频率的限制。比如登录的接口、发送短信的接口、秒杀商品的接口等等。

官方的golang.org/x/time/rate包中实现了令牌桶的算法，我们看一下应用。

package main

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

func main() {

	//实例化一个限流器，桶的容量是1，每秒生成一个令牌
	limiter := rate.NewLimiter(rate.Every(1*time.Second), 1)

	//获取令牌
	limiter.Allow()

}
复制代码

上方的限流器，桶的容量是1，生成令牌的速率是每秒一个，也就是说该限流器能起到每秒限制一个请求的效果。但是实际开发中，还有一些额外判断。比如登录接口，我需要限制一个ip每秒登录一次，再比如发送验证码接口，我需要一个手机号一秒只能发一次。这就需要额外封装一下官方限流器了。

在项目中新建component/limiter文件夹，并创建limiter.go文件

package limiter

//component/limiter/limiter.go

import (
	"golang.org/x/time/rate"
	"sync"
	"time"
)


type Limiters struct {
	limiters map[string]*Limiter
	lock     sync.Mutex
}

type Limiter struct {
	limiter *rate.Limiter
	lastGet time.Time //上一次获取token的时间
	key     string    //限流器的key,用来识别对应的限流器
}

//全局限流器，存储所有限流器
var GlobalLimiters = &Limiters{
	limiters: make(map[string]*Limiter),
	lock:     sync.Mutex{},
}
复制代码

将ip、手机号这种的作为限流器组的标识

接下来就是实例化限流器和获取令牌函数的实现

package limiter

//component/limiter/limiter.go

import (
	"golang.org/x/time/rate"
	"sync"
	"time"
)

type Limiters struct {
	limiters map[string]*Limiter
	lock     sync.Mutex
}

type Limiter struct {
	limiter *rate.Limiter
	lastGet time.Time //上一次获取token的时间
	key     string
}

var GlobalLimiters = &Limiters{
	limiters: make(map[string]*Limiter),
	lock:     sync.Mutex{},
}



func NewLimiter(r rate.Limit, b int, key string) *Limiter {

	
	keyLimiter := GlobalLimiters.getLimiter(r, b, key)

	return keyLimiter

}

func (l *Limiter) Allow() bool {

	l.lastGet = time.Now()

	return l.limiter.Allow()

}

func (ls *Limiters) getLimiter(r rate.Limit, b int, key string) *Limiter {

	ls.lock.Lock()

	defer ls.lock.Unlock()

	limiter, ok := ls.limiters[key]

	if ok {

		return limiter
	}

	l := &Limiter{
		limiter: rate.NewLimiter(r, b),
		lastGet: time.Now(),
		key:     key,
	}

	ls.limiters[key] = l

	return l
}
复制代码

在控制器中调用

func Index(context *context.Context) *response.Response {

	l := limiter.NewLimiter(rate.Every(1*time.Second), 1, context.ClientIP())

	if !l.Allow() {

		return response.Resp().String("您的访问过于频繁")
	}

	return response.Resp().String("success")
}
复制代码

访问浏览器，当你的刷新的手速过快时

这段代码建议在中间件中实现

由于限流器是一个全局的map变量，当一些限流器长时间不用时，会一直存在于内存中，所以需要定时对过期的限流器做一个清理。

完整代码如下：

package limiter

//component/limiter/limiter.go

import (
	"golang.org/x/time/rate"
	"sync"
	"time"
)

type Limiters struct {
	limiters map[string]*Limiter
	lock     sync.Mutex
}

type Limiter struct {
	limiter *rate.Limiter
	lastGet time.Time //上一次获取token的时间
	key     string
}

var GlobalLimiters = &Limiters{
	limiters: make(map[string]*Limiter),
	lock:     sync.Mutex{},
}

var once = sync.Once{}

func NewLimiter(r rate.Limit, b int, key string) *Limiter {

	once.Do(func() {

		go GlobalLimiters.clearLimiter()
	})

	keyLimiter := GlobalLimiters.getLimiter(r, b, key)

	return keyLimiter

}

func (l *Limiter) Allow() bool {

	l.lastGet = time.Now()

	return l.limiter.Allow()

}

func (ls *Limiters) getLimiter(r rate.Limit, b int, key string) *Limiter {

	ls.lock.Lock()

	defer ls.lock.Unlock()

	limiter, ok := ls.limiters[key]

	if ok {

		return limiter
	}

	l := &Limiter{
		limiter: rate.NewLimiter(r, b),
		lastGet: time.Now(),
		key:     key,
	}

	ls.limiters[key] = l

	return l
}

//清除过期的限流器
func (ls *Limiters) clearLimiter() {

	for {

		time.Sleep(1 * time.Minute)
                ls.lock.Lock()
		for i, i2 := range ls.limiters {

			//超过1分钟
			if time.Now().Unix()-i2.lastGet.Unix() > 60 {
				
				delete(ls.limiters, i)
				
			}

		}
                ls.lock.Unlock()

	}

}
复制代码