redis 做分布式限流

发布时间 2023-04-23 19:25:19作者: 变换

参看来源:

https://blog.51cto.com/u_15708799/5703440

 

测试有效代码:

 

 @Test
    public void testLimitWait() throws InterruptedException {
        ExecutorService pool = Executors.newCachedThreadPool();
        log.info("--------{}", redisTemplate.opsForValue().get("abc"));
        for (int j = 1; j <= 5; j++) {
            int i=j;
            pool.execute(() -> {
                Thread.currentThread().setName( Thread.currentThread().getName().replace("-","_"));
                limitWait("abc", 3, 1);
                log.info(i + ":" + true + " ttl:" + redisTemplate.getExpire("abc", TimeUnit.MILLISECONDS));
                try {
                    // 线程等待,模拟执行业务逻辑
                    Thread.sleep(new Random().nextInt(100));
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            });
        }
        pool.shutdown();
        pool.awaitTermination(2,TimeUnit.SECONDS);
    }


    /**
     * 达到限流时,则等待,直到新的间隔。
     *
     * @param key 可以是ip + 当前秒  或者 是 uid +  当前秒  或者 固定的一个入口 key
     * @param limitCount   一定时间内最多访问次数
     * @param limitSecond  给定的时间范围 单位(秒)
     */
    public void limitWait(String key, int limitCount, int limitSecond) {
        boolean ok;//放行标志
        do {
            ok = limit(key, limitCount, limitSecond);
            log.info("放行标志={}", ok);
            if (!ok) {
                Long ttl = redisTemplate.getExpire(key, TimeUnit.MILLISECONDS);
                if (null != ttl && ttl > 0) {
                    try {
                        Thread.sleep(ttl);
                        log.info("sleeped:{}", ttl);
                    } catch (InterruptedException e) {
                        e.printStackTrace();
                    }
                }
            }
        } while (!ok);
    }


    /**
     * 限流方法    true-放行;false-限流
     *
     * @param key
     * @param limitCount
     * @param limitSecond
     * @return
     */
    public boolean limit(String key, int limitCount, int limitSecond) {
        List<String> keys = Collections.singletonList(key);
        String luaScript = buildLuaScript();
        RedisScript<Long> redisScript = new DefaultRedisScript<>(luaScript, Long.class);
        Long count = redisTemplate.execute(redisScript, keys, limitCount, limitSecond);
        log.info("Access try count is {} for key = {}", count, key);
        if (count != null && count.intValue() <= limitCount) {
            return true;//放行
        } else {
            return false;//限流
//            throw new RuntimeException("You have been dragged into the blacklist");
        }
    }


    /**
     * 编写 redis Lua 限流脚本
     */
    public String buildLuaScript() {
        StringBuilder lua = new StringBuilder();
        lua.append("local c");
        lua.append("\nc = redis.call('get',KEYS[1])");
        // 实际调用次数超过阈值,则直接返回
        lua.append("\nif c and tonumber(c) > tonumber(ARGV[1]) then");
        lua.append("\nreturn c;");
        lua.append("\nend");
        // 执行计算器自加
        lua.append("\nc = redis.call('incr',KEYS[1])");
        lua.append("\nif tonumber(c) == 1 then");
        // 从第一次调用开始限流,设置对应键值的过期
        lua.append("\nredis.call('expire',KEYS[1],ARGV[2])");
        lua.append("\nend");
        lua.append("\nreturn c;");
        return lua.toString();
    }