工作中我们经常利用redis来实现限速, 比如限制一个手机号60秒最多发送3条短信.

如果不考虑原子性, 伪代码如下:

long count = incr('手机号') ; 

if count==1  expire('手机号',60) ; 

if count>3   return "发送频率超限" ; 

上面代码在执行时, 前后可能调用redis两次, 网络耗时的问题,忽略不说, 但另外一个问题应该被重视, 就是整个执行的原子性, 假如incr执行成功,而expire执行失败, 那么reids里面就会存在不主动失效的脏数据.

为了保证整个执行的原子性, 我们来利用Lua.

首先把incr和expire的调用逻辑一起封装在Lua里, 然后再通过redis调用. 

定义脚本文件rate_limit.lua,内容如下:

-- incr
local count = redis.call('incr',KEYS[1])

-- 第一次时,设置expire
if count == 1 then
  redis.call('expire',KEYS[1],ARGV[1])
end

-- 返回目前次数
return count

通过redis-cli调用lua测试:

redis-cli --eval rate_limit.lua 手机号 , 60
(integer) 1
redis-cli --eval rate_limit.lua 手机号 , 60
(integer) 2

如果项目是Java编写,且使用了spring-data-redis, 可参考代码如下:

String lua = "";
lua += "local count = redis.call('incr',KEYS[1]) ";
lua += "if count == 1 then ";
lua += "redis.call('expire',KEYS[1],ARGV[1]) ";
lua += "end ";
lua += "return count";
RedisScript redisScript=new DefaultRedisScript(lua,Long.TYPE);
Long count=(Long)redisTemplate.execute(redisScript, Arrays.asList("手机号"), 60);
if(count>3){
   return "发送频率超限";
}