sinksmell/rateLimit

// models/pool.go
// 连接池
var(
	POOL *redis.Pool
)


func init(){
	initPool()
}

// 初始化redis连接池
func initPool(){
	POOL=&redis.Pool{
		MaxIdle:16, // 最初连接数量
		MaxActive:0, // 最大连接数量 0表示按需创建
		IdleTimeout:300, // 连接关闭时间
		Dial: func() (redis.Conn, error) {
			return redis.Dial("tcp","127.0.0.1:6379")
		},
	}
}


// models/limit.go
// 判断请求是否达到上限
func RateLimit(window,maxCnt int) bool{
	var(
		conn redis.Conn
		res interface{}
		timeStamp int64
	)
	// 获取连接
	conn=POOL.Get()
	defer conn.Close()

	// 判断指定key是否存在
	res,_=conn.Do("EXISTS",KEY)
	isExits:=res.(int64)
	timeStamp=time.Now().Unix()
	// 如果key 不存在
	if isExits==0{
		// 向列表中添加时间戳 设置过期时间为时间窗口大小
		conn.Do("LPUSH",KEY,timeStamp)
		conn.Do("EXPIRE",KEY,window)
		// 可以继续放行请求
		return true
	}

	// key存在 获取当前list长度
	res,_=conn.Do("LLEN",KEY)
	// 转为int
	lens :=res.(int64)
	end:=0
	// 获取list
	list,_:=redis.Values(conn.Do("lrange",KEY,0, lens))
	// 遍历list找到最后一个没过期的记录
	for i := int(lens -1);i>=0 ;i--  {
		str:=string(list[i].([]byte))
		oldStamp,_:=strconv.ParseInt(str,10,64)
		if timeStamp-oldStamp<int64(window){
			end=i
			break
		}
	}
	

	// 删除过期时间戳记录
	conn.Do("LTRIM",KEY,0,end)

	// 判断记录是否达到限制
	if end+1<maxCnt{
		// beego.BeeLogger.Debug("添加记录")
		// 向列表中添加时间戳
		// 向列表中添加时间戳 设置过期时间为时间窗口大小
		conn.Do("LPUSH",KEY,timeStamp)
		conn.Do("EXPIRE",KEY,window)
		// 放行请求
		return true
	}
	// 不放行
	return false
}

// filters/limit.go
// 过滤器实现请求限流
var CheckFilter= func(ctx *context.Context) {
	// 请求被限制
	if pass:=models.RateLimit(models.TIME_WND,models.MAX_CNT);!pass{
		ctx.Redirect(302,"/err")
	}
} 

// main.go
// 对指定的url配置过滤器
beego.InsertFilter("/check",beego.BeforeRouter,filters.CheckFilter)