【golang】根据key做处理速度rate限制+key过期功能

一、背景

公司iot设备某个固件版本异常，导致单机器上抛百万条数据（正常只有几百条），需要在业务端做速度限制。

二、技术选择

1.time/rate

go get	golang.org/x/time/rate

Golang 标准库 time/rate
限流器是后台服务中的非常重要的组件，可以用来限制请求速率，保护服务，以免服务过载。

限流器的实现方法有很多种，例如滑动窗口法、Token Bucket、Leaky Bucket 等。

其实 Golang 标准库中就自带了限流算法的实现，即 golang.org/x/time/rate。该限流器是基于 Token Bucket(令牌桶) 实现的。

简单来说，令牌桶就是想象有一个固定大小的桶，系统会以恒定速率向桶中放 Token，桶满则暂时不放。
而用户则从桶中取 Token，如果有剩余 Token 就可以一直取。如果没有剩余 Token，则需要等到系统中被放置了 Token 才行。

构造一个限流器

我们可以使用以下方法构造一个限流器对象：

limiter := NewLimiter(10, 1);

这里有两个参数：

1. 第一个参数是 r Limit。代表每秒可以向 Token 桶中产生多少 token。Limit 实际上是 float64 的别名。
2. 第二个参数是 b int。b 代表 Token 桶的容量大小。

那么，对于以上例子来说，其构造出的限流器含义为，其令牌桶大小为 1, 以每秒 10 个 Token 的速率向桶中放置 Token。

除了直接指定每秒产生的 Token 个数外，还可以用 Every 方法来指定向 Token 桶中放置 Token 的间隔，例如：

limit := Every(100 * time.Millisecond);
limiter := NewLimiter(limit, 1);

以上就表示每 100ms 往桶中放一个 Token。本质上也就是一秒钟产生 10 个。

Limiter 提供了三类方法供用户消费 Token，用户可以每次消费一个 Token，也可以一次性消费多个 Token。
而每种方法代表了当 Token 不足时，各自不同的对应手段。

Wait/WaitN

func (lim *Limiter) Wait(ctx context.Context) (err error)
func (lim *Limiter) WaitN(ctx context.Context, n int) (err error)

Wait 实际上就是 WaitN(ctx,1)。

当使用 Wait 方法消费 Token 时，如果此时桶内 Token 数组不足 (小于 N)，那么 Wait 方法将会阻塞一段时间，直至 Token 满足条件。如果充足则直接返回。

这里可以看到，Wait 方法有一个 context 参数。
我们可以设置 context 的 Deadline 或者 Timeout，来决定此次 Wait 的最长时间。

Allow/AllowN

func (lim *Limiter) Allow() bool
func (lim *Limiter) AllowN(now time.Time, n int) bool

Allow 实际上就是 AllowN(time.Now(),1)。

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

通常对应这样的线上场景，如果请求速率过快，就直接丢到某些请求。（这里符合我们的需求）

Reserve/ReserveN

func (lim *Limiter) Reserve() *Reservation
func (lim *Limiter) ReserveN(now time.Time, n int) *Reservation

Reserve 相当于 ReserveN(time.Now(), 1)。

ReserveN 的用法就相对来说复杂一些，当调用完成后，无论 Token 是否充足，都会返回一个 Reservation * 对象。

你可以调用该对象的 Delay() 方法，该方法返回了需要等待的时间。如果等待时间为 0，则说明不用等待。
必须等到等待时间之后，才能进行接下来的工作。

或者，如果不想等待，可以调用 Cancel() 方法，该方法会将 Token 归还。

2.cache

go get github.com/patrickmn/go-cache

go-cache [1]基于内存的 K/V 存储/缓存 : (类似于Memcached)，适用于单机应用程序

Example

package main
 
import (
 "fmt"
 "time"
 
 "github.com/patrickmn/go-cache"
)
 
type MyStruct struct {
 Name string
}
 
func main() {
 // 设置超时时间和清理时间
 c := cache.New(5*time.Minute, 10*time.Minute)
 
 // 设置缓存值并带上过期时间
 c.Set("foo", "bar", cache.DefaultExpiration)
 
 
 // 设置没有过期时间的KEY，这个KEY不会被自动清除，想清除使用：c.Delete("baz")
 c.Set("baz", 42, cache.NoExpiration)
 
 
 var foo interface{}
 var found bool
 // 获取值
 foo, found = c.Get("foo")
 if found {
  fmt.Println(foo)
 }
 
 var foos string
 // 获取值， 并断言
 if x, found := c.Get("foo"); found {
  foos = x.(string)
  fmt.Println(foos)
 }
 // 对结构体指针进行操作
 var my *MyStruct
 c.Set("foo", &MyStruct{Name: "NameName"}, cache.DefaultExpiration)
 if x, found := c.Get("foo"); found {
  my = x.(*MyStruct)
  // ...
 }
 fmt.Println(my)
}
 

三、示例代码

go get github.com/patrickmn/go-cache
go get	golang.org/x/time/rate

package main

import (
	"log"
	"sync"
	"time"

	"github.com/patrickmn/go-cache"
	"golang.org/x/time/rate"
)

type Limiter struct {
	tokenBucketsNoTTL   map[string]*rate.Limiter
	tokenBucketsWithTTL *cache.Cache
	lock                sync.Mutex
}

func (l *Limiter) limitReachedNoTokenBucketTTL(key string) bool {
	l.lock.Lock()
	defer l.lock.Unlock()

	if _, found := l.tokenBucketsNoTTL[key]; !found {
		//limit := rate.Every(20 * time.Millisecond)
		l.tokenBucketsNoTTL[key] = rate.NewLimiter(1, 1)
	}
	//log.Println(l.tokenBucketsNoTTL[key].Allow())
	return l.tokenBucketsNoTTL[key].Allow()
}
func (l *Limiter) limitReachedWithCustomTokenBucketTTL(key string, tokenBucketTTL time.Duration) bool {
	l.lock.Lock()
	defer l.lock.Unlock()

	if _, found := l.tokenBucketsWithTTL.Get(key); !found {
		limit := rate.Every(60 * time.Second)
		x := rate.NewLimiter(limit, 10)
		log.Println(key, "过期了，重新申请")
		l.tokenBucketsWithTTL.Set(
			key,
			x,
			tokenBucketTTL,
		)
	}

	expiringMap, found := l.tokenBucketsWithTTL.Get(key)
	if !found {
		return false
	}

	return expiringMap.(*rate.Limiter).Allow()
}

var limiter Limiter

func init() {
	limiter.tokenBucketsNoTTL = make(map[string]*rate.Limiter)
	c := cache.New(5*time.Minute, 10*time.Minute)

	limiter.tokenBucketsWithTTL = c
}
func main() {
	test1()
}
func test0() {
	b := map[string]int{}
	limiter.tokenBucketsNoTTL = make(map[string]*rate.Limiter)
	c := cache.New(5*time.Minute, 10*time.Minute)

	limiter.tokenBucketsWithTTL = c
	b["001"] = 0
	for i := 0; i < 50; i++ {
		time.Sleep(10 * time.Second)
		did := "001"
		if limiter.limitReachedNoTokenBucketTTL(did) {
			b[did]++

		}
	}
	log.Println(b["001"])
}
func test1() {
	b := map[string]int{}

	b["001"] = 0
	did := "001"
	// limiter.limitReachedWithCustomTokenBucketTTL(did, cache.DefaultExpiration)
	// time.Sleep(30 * time.Second)

	for i := 0; i < 50; i++ {
		time.Sleep(1 * time.Second)
		//cache.DefaultExpiration
		if limiter.limitReachedWithCustomTokenBucketTTL(did, cache.DefaultExpiration) {
			log.Println("i:", i)
			b[did]++

		}
	}
	log.Println(b["001"])
}