Go实现限流器(计数器模式)

1.实现手段描述

定义一个结构:代表某秒,最大数,当前数
发起某个请求时,判断当前时间是不是这个结构所代表的秒,
代表:看看当前数达到了最大值没有,没有++,达到则拒绝请求
不代表某秒了:代表则更新
注意:赋值需要原子操作,atomic.Addint64
利用time.Now().Unix()得到当前的unix秒数
每次请求前就判断是否是同一个Unix秒数时间,
如果是:判断是否达到了最大值,没有就++,表示允许通过.
达到了就直接不允许,结合中间件使用

2.talk is cheap

package utils
import (
    "sync/atomic"
    "time"
)
//Limiter 限流器对象
type Limiter struct {
    value  int64
    max    int64
    ts     int64
}
    
//NewLimiter 产生一个限流器
func NewLimiter(count int64) *Limiter {
    return &Limiter{
        value:  0,
        max:   count,
        ts:     time.Now().Unix(),
    }
}

//Ok 是否可以通过
func (l *Limiter) Ok() bool {
    ts := time.Now().Unix()
    tsOld := atomic.LoadInt64(&l.ts)
   
   //不在同1s内,就重置时间戳和计数器为1
    if ts != tsOld {
        atomic.StoreInt64(&l.ts, ts)
        atomic.StoreInt64(&l.value, 1)
        return true
    }

   //如果还在同一s内,使用次数达最大值没有,没超过就+1
    return atomic.AddInt64(&(l.value), 1) < l.max
}

//SetMax 设置最大限制
func (l *Limiter) SetMax(m int64) {
    l.max = m
}


//MaxAllowed 限流器
func MaxAllowed(limitValue int64) func(c *gin.Context) {
    limiter := utils.NewLimiter(limitValue)
    log.Println("limiter.SetMax:", limitValue)
    // 返回限流逻辑
    return func(c *gin.Context) {
        if !limiter.Ok() {
            c.AbortWithStatus(http.StatusServiceUnavailable)  //超过每秒200，就返回503错误码
            return
        }
        c.Next()
    }
}
func main(){
    router := gin.New()
    router.Use(MaxAllowed(200))  //限制每秒最多允许200个请求
}