Redis+Lua限制发送量及遇到的坑

业务中需要限制每个账号每天发送短信数量，如果没有超过设置的发送量，则正常发送，否则返回失败。

解决思路：

• 将账号ID+yyyyMMdd组成redis的key，value为当天的发送量。
• 在发送前获取账号ID+yyyyMMdd的值，如果没有超过发送量，则账号ID+yyyyMMdd对应的值+1，发送短信，如果账号ID+yyyyMMdd对应的值大于发送量，则返回失败。
• 设置账号ID+yyyyMMdd的过期时间为1天。
• 考虑并发问题，“查询账号ID+yyyyMMdd的值”和“账号ID+yyyyMMdd对应的值+1”需要同步执行。

综上考虑，决定采用redis+lua来解决上述问题。代码片段如下：

lua脚本 smslimit.lua

local key = KEYS[1]
local keyseconds = tonumber(ARGV[1])
local limitcount = tonumber(ARGV[2])

local val = redis.call('INCRBY', key, 1)

if val == 1 then
    redis.call('EXPIRE', key, keyseconds)
end

if val > limitcount then
    return '1'
end

return '0';

频率限制接口 DateLimitService.java

/**
 * @author conquer on 2021/7/5 5:04 下午
 */
public interface DateLimitService {

    /**
     * 是否超过发送限制
     *
     * @param accountId
     * @return
     */
    Boolean isSendExceedLimit(Integer accountId);
}

频率限制服务 DateLimitServiceImpl.java

@Service
@Slf4j
public class DateLimitServiceImpl implements DateLimitService {

    @Resource
    private StringRedisTemplate stringRedisTemplate;

    private static final DefaultRedisScript<String> SCRIPT = new DefaultRedisScript<>();

    public static final String LUA_RELATIVE_PATH = "lua/smslimit.lua";

    static {
        SCRIPT.setResultType(String.class);
        SCRIPT.setLocation(new ClassPathResource(LUA_RELATIVE_PATH));
    }


    /**
     * 账号是否超过今天发送量
     * 
     * @param accountId
     * @return
     */
    @Override
    public Boolean isSendExceedLimit(Integer accountId) {
        /**
         * List设置lua的KEYS
         */
        List<String> keyList = new ArrayList();
        String ttl = "";
        String limitCount = "";

        
        keyList.add("DAY_SEND_COUNT" + DateUtil.format(new Date(), "yyyyMMdd") + ":" + accountId);
        ttl = String.valueOf(24 * 60 * 60);

        // 一天最多发送1000条
        limitCount = String.valueOf(1000);
        

        String res = stringRedisTemplate.execute(SCRIPT, keyList, ttl, 1000);

        log.info("res: {}", res);

        // 比对结果为1表示，超出限制
        if ("1".equals(res)) {
            log.info("send limit, accountId: {}", accountId);
            return true;
        }
        return false;
    }
}

说一下踩到的坑：

刚开始我把ttl和limitCount两个字段直接放在keyList中，在lua脚本中使用KEYS[1]、KEYS[2]、KEYS[3] 取出ttl和limitCount字段，在测试环境（单机模式）运行正常，但是发布到线上（分片集群模式）后发现报错，提示key找不到。猜测可能是因为把参数放在keyList的原因，改完后测试发现没问题。

所以需要严格规范地把redis中的key放在keyList，参数放在stringRedisTemplate.execute()方法后面。