go限流器的简单通用实现

原创
2022/04/11 22:57
阅读数 153

基于官网的令牌桶限流器封装

package rate

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

type Limiters struct {
  limiters *sync.Map
}

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

var GlobalLimiters = &Limiters{
  limiters: &sync.Map{},
}
var once = sync.Once{}

// 令牌桶的限流算法
/*
t 单位时间生成的令牌数,多少时间周期内产生一个令牌,比如:每2秒产生一个令牌: time.Second * 2
b 桶的大小,最多只能有多少个令牌
限流的标识
*/
func New(t time.Duration, b int, key string) *Limiter {
  once.Do(func() {
    go GlobalLimiters.clearLimiter()
  })
  keyLimiter := GlobalLimiters.getLimiter(rate.Every(t), b, key)
  return keyLimiter
}

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

  limiter, ok := ls.limiters.Load(key)
  if ok {
    return limiter.(*Limiter)
  }

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

  ls.limiters.Store(key, l)
  return l
}

//清除过期的限流器
func (ls *Limiters) clearLimiter() {
  for {
    time.Sleep(1 * time.Minute)
    ls.limiters.Range(func(key, value interface{}) bool {
      //超过1分钟
      if time.Now().Unix()-value.(*Limiter).lastGet.Unix() > 60 {
        ls.limiters.Delete(key)
      }
      return true
    })
  }
}

func (l *Limiter) Allow() bool {
  l.lastGet = time.Now()
  return l.limiter.Allow()
}

应用

#实例化一个限流器,每2秒一个令牌,最多只有一个,根据ip地址进行限流.
# 这里可以封装成中间件,并结合路由地址,请求方式来生成.
l := LimitRate.New(time.Second * 2, 1, "111.111.111.111")
if !l.Allow() {
  return //不允许,就直接返回
}
展开阅读全文
加载中
点击引领话题📣 发布并加入讨论🔥
打赏
0 评论
0 收藏
0
分享
返回顶部
顶部