golang利用redis实现限速器

利用redis的String数据结构可以实现计数器的功能，利用incr或者incrby来实现计数器的增加或者减少。

同时redis对可以对key实现过期功能，这样我们就可以利用计数器和过期功能来实现限速器。

例如很多时候我们需要对单个用户进行限速，每分钟或每秒钟限制limit次，可以使用user_id作为key，访问次数作为value，到redis进行存储，当用户访问一次，我们就incr一下；同时我们设置key的过期时间为1秒或1分钟，当用户访问超出限制后，我们禁止用户的访问，当key过期后，我们重新设置key，这样用户在下一分钟或者下一秒就又有了limit次访问机会。

具体实现如下：

const (
	timeout = 60 //60 seconds
)

func RateLimitIncr(key string) (err error) {
	con := pool.Get()
	defer con.Close()
	count, err := redis.Int64(con.Do("incr", key))
	if err != nil {
		return
	}
	if count == 1 {
		_, err = con.Do("expire", timeout)
	}
	return
}

func GetRateLimitCount(key string) (count int64, err error) {
	con := pool.Get()
	defer con.Close()
	count, err = redis.Int64(con.Do("get", key))
	if err != nil {
		if err == redis.ErrNil {
			err = nil
		}
		return
	}
	return
}

一个简单的限速器就这样实现了，但是有两个问题：

• 不精确，只能适应于不需要很精确限速的场景

        如果我们设置每分钟能访问10次，想下当用户在第一分钟的开始访问了一次，这时候redis创建了一个key，在第一分钟的50秒时访问了9次，然后第60秒时key过期，用户在第61秒访问，redis会重新创建key，用户可以在61秒时访问10次，这样在第50秒到61秒的时间内，用户访问了19次。所以说这个限速器只适应于简单的不需要那么精确的限速场景。如果想精确的限速，可以使用令牌桶算法或者漏桶算法等。

• incr和expire并不是原子操作

当在创建key时，我们直接使用incr来创建，然后判断其返回值==1，再执行expire操作。这样有个风险点，就是当程序刚执行完incr时，我们的程序被重启或者崩溃了，那么这个key就没有设置过期时间，这样可能导致用户一直无法访问接口。

 

针对第二点，现在的做法一般是通过lua脚本去处理，在lua脚本中将两条命令放在一起，这样redis执行时就会是原子操作；

还有一种是我自己想出来的方法，我们可以在获取计数器值时用ttl命令判断一下key的过期时间，然后补偿给key，ttl的返回值如下：

当ttl返回-1时，我们知道这是异常数据，可以立即给他设置上过期时间，这样就能避免出现上面的问题。

实现如下：

func GetRateLimitCount(key string) (count int64, err error) {
	con := pool.Get()
	defer con.Close()
	count, err = redis.Int64(con.Do("get", key))
	if err != nil {
		if err == redis.ErrNil {
			err = nil
		}
		return
	}
	ttl, err := redis.Int64(con.Do("ttl", key))
	if err != nil {
		return
	}
	if ttl == -1 {
		_, err = con.Do("expire", timeout)
	}
	return
}