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
}