通过滑动窗口实现接口调用的多种限制策略

前言

有个邮箱发送的限制发送次数需求,为了防止用户恶意请求发送邮件的接口,对用户的发送邮件次数进行限制,每个邮箱60s内只能接收一次邮件,每个小时只能接收五次邮件,24小时只能接收十次邮件,一共有三个条件的限制。

实现方案

单机方案

单机简单实现可以用Caffeine,在Caffeine里面Key为mail的标识,value是个存这个mail每次接收邮件的时间戳List,数据结构如下图所示:

  1. list小于5个:每一次有新元素入队,都要判断队列里最新的时间戳和当前时间戳是否超过60s,不超过返回60s限制。
  2. 大于等于5个,小于10,则当前队列size-5,即往前数第五个值,取对应的value时间戳,判断和当前时间超不超1h,超过就放入list,不超就返回超过一小时的限制。
  3. 如果数量等于10个,得先判断24小时超不超10个,拿List里面的第一个值,判断和当前的时间戳是否超过24小时,不超则返回24小时限制,超再判断1小数超不超,判断逻辑往前数五个,如果超过,则把第一个值剔除(即最老的那个元素),加入新的元素。

通过上面的数据结构,其实也能把剩余多少时间接触限制一并返回到前端,在达到限制的时候,对比时间戳时间的差距即可。

caffeine单机方案代码

java 复制代码
 public boolean isMailCanSend(String mail){
        // 先判断缓存是否存在 不存在 则创建
        ArrayList<Long> mailTimeStampList = caffeineTemplate.getMailTimeStampFromCache(CacheKeyConstant.TIME_STAMP_LIST_FOR_MAIL + mail);
        if (mailTimeStampList == null) {
            ArrayList<Long> timeList = new ArrayList<>();
            timeList.add(System.currentTimeMillis());
            caffeineTemplate.addToMailTimeStampCache(CacheKeyConstant.TIME_STAMP_LIST_FOR_MAIL + mail, timeList);
            return true;
        } else {
            // 缓存存在
            // 存在先查60s
            Long timeStamp = mailTimeStampList.get(mailTimeStampList.size() - 1);
            // 判断与当前时间相差是否超过60s
            if (System.currentTimeMillis() - timeStamp > 60000) {
                // 再查数量是否小于5,满足直接加入缓存
                if (mailTimeStampList.size() < 5) {
                    mailTimeStampList.add(System.currentTimeMillis());
                    caffeineTemplate.addToMailTimeStampCache(CacheKeyConstant.TIME_STAMP_LIST_FOR_MAIL + mail, mailTimeStampList);
                    return true;
                } else {
                    // 大于等于5数量小于10
                    if (mailTimeStampList.size() < 10) {
                        // 则判断前面第五个是否满足一个小时
                        if (System.currentTimeMillis() - mailTimeStampList.get(mailTimeStampList.size() - 5) < 3600000) {
                            // 不满足大于一个小时 则不可发送
                            throw new EmailException(ResultCodeEnum.MAIL_ONE_HOUR_REQUEST_FREQUENT_ERROR, 3600000L - (System.currentTimeMillis() - mailTimeStampList.get(mailTimeStampList.size() - 5)));
                        } else {
                            mailTimeStampList.add(System.currentTimeMillis());
                            caffeineTemplate.addToMailTimeStampCache(CacheKeyConstant.TIME_STAMP_LIST_FOR_MAIL + mail, mailTimeStampList);
                            return true;
                        }
                    } else {
                        // 数量为 10的时候
                        // 等于10 判断大于24小时是否满足
                        if (System.currentTimeMillis() - mailTimeStampList.get(0) > 86400000) {
                            // 则判断前面第五个是否满足一个小时
                            if (System.currentTimeMillis() - mailTimeStampList.get(mailTimeStampList.size() - 5) < 3600000) {
                                // 不满足一个小时 则不可发送
                                throw new EmailException(ResultCodeEnum.MAIL_ONE_HOUR_REQUEST_FREQUENT_ERROR, 3600000L - (System.currentTimeMillis() - mailTimeStampList.get(mailTimeStampList.size() - 5)));
                            } else {
                                // 移除第一个
                                mailTimeStampList.remove(0);
                                mailTimeStampList.add(System.currentTimeMillis());
                                caffeineTemplate.addToMailTimeStampCache(CacheKeyConstant.TIME_STAMP_LIST_FOR_MAIL + mail, mailTimeStampList);
                                return true;
                            }
                        } else {
                            throw new EmailException(ResultCodeEnum.MAIL_24_HOUR_REQUEST_FREQUENT_ERROR, 86400000L - (System.currentTimeMillis() - mailTimeStampList.get(0)));
                        }
                    }
                }
            } else {
                throw new EmailException(ResultCodeEnum.MAIL_ONE_MIN_REQUEST_FREQUENT_ERROR, 60000L - (System.currentTimeMillis() - timeStamp));
            }


        }

分布式方案

分布式方案可以使用redis的zset数据结构来实现,同样是维护一个set,score存放的是时间戳,窗口元素都是24小时以内。

  1. 每次有新请求,先将时间戳位于窗口外的元素清除掉。
  2. set大小大于等于10,不放行,返回超过24小时限制。
  3. 判断set排名最大的元素的时间戳和当前时间戳是否超过60s,超过则放行,不超过返回60s限制。
  4. 判断set大小是否小于5,小于5则放行,并放入新元素。
  5. set大小小于10,大于等于5,取当前的set排名往前数5,即ZRANGE key size-5 size-5,拿出排行倒数第五的元素,判断是否超过一个小时,超过一个小时则可以放行,不超过返回1小时限制。

上述的执行应该以原子形式进行,防止出现不准确情况,这里采用lua脚本

lua脚本

lua 复制代码
local key = KEYS[1]
local limit1 = tonumber(ARGV[1])
local limit2 = tonumber(ARGV[2])
local windowStart = tonumber(ARGV[3])
local currentTime = tonumber(ARGV[4])

-- 清除窗口外的元素
redis.call('zremrangebyscore', key, 0 , windowStart)

-- 获取当前集合大小
local currentSize = tonumber(redis.call('zcard', key))

if currentSize >= limit2 then
    -- 集合大小大于等于 limit2,不放行,返回超过24小时限制
    return 0
end

-- 判断集合中最大元素与当前时间间隔是否超过60秒
local oldestTimestamp = tonumber(redis.call('zrange', key, -1, -1, 'WITHSCORES')[2])
if (currentTime - oldestTimestamp) < 60000 then
    -- 未超过60秒限制,返回60秒限制
    return 0
end

if currentSize < limit1 then
    -- 集合大小小于 limit1,放行请求并添加新元素
    redis.call('zadd', key, currentTime, currentTime)
    return 1
else
    -- 集合大小小于 limit2 且大于等于 limit1,判断是否超过1小时限制
    local hourAgoTimestamp = currentTime - 3600000          
    local fifthTimestamp = tonumber(redis.call('zrange', key, currentSize - limit1, currentSize - limit1, 'WITHSCORES')[2])
    if fifthTimestamp < hourAgoTimestamp then
        -- 未超过1小时限制,放行请求并添加新元素
        redis.call('zadd', key, currentTime, currentTime)
        return 1
    else
        -- 已超过1小时限制,返回1小时限制
        return 0
    end
end

java代码

java 复制代码
import org.redisson.Redisson;
import org.redisson.api.RScript;
import org.redisson.api.RedissonClient;
import org.redisson.config.Config;

public class RedisLuaScriptExample {

    public static void main(String[] args) {
        Config config = new Config();
        config.useSingleServer().setAddress("redis://127.0.0.1:6379");
        RedissonClient redisson = Redisson.create(config);

        // Lua 脚本
        String luaScript =
            "local key = KEYS[1] " +
            "local limit1 = tonumber(ARGV[1]) " +
            "local limit2 = tonumber(ARGV[2]) " +
            "local windowStart = tonumber(ARGV[3]) " +
            "local currentTime = tonumber(ARGV[4]) " +
            "redis.call('zremrangebyscore', key, '-inf', windowStart) " +
            "local currentSize = tonumber(redis.call('zcard', key)) " +
            "if currentSize >= limit2 then " +
            "  return 0 " +
            "end " +
            "local oldestTimestamp = tonumber(redis.call('zrange', key, -1, -1, 'WITHSCORES')[2]) " +
            "if (currentTime - oldestTimestamp) < 60000 then " +
            "  return 0 " +
            "end " +
            "if currentSize < limit1 then " +
            "  redis.call('zadd', key, currentTime, currentTime) " +
            "  return 1 " +
            "else " +
            "  local hourAgoTimestamp = currentTime - 3600000 " +
            "  local fifthTimestamp = tonumber(redis.call('zrange', key, currentSize - limit1 , currentSize - limit1, 'WITHSCORES')[2]) " +
            "  if fifthTimestamp < hourAgoTimestamp then " +
            "    redis.call('zadd', key, currentTime, currentTime) " +
            "    return 1 " +
            "  else " +
            "    return 0 " +
            "  end " +
            "end";

        RScript script = redisson.getScript();
        // 执行 Lua 脚本
        Long result = script.eval(RScript.Mode.READ_WRITE, luaScript, RScript.ReturnType.INTEGER,
                                  "your_key", // 这里替换成你的键
                                  "5",        // 替换成 limit1 的值
                                  "10",       // 替换成 limit2 的值
                                  String.valueOf(System.currentTimeMillis() - 86400000), // 24小时前的时间戳
                                  String.valueOf(System.currentTimeMillis()));
        System.out.println("Result: " + result);

        redisson.shutdown();
    }
}