基于redis+lua脚本实现限流

 private void waitForToken(String tenantId) throws InterruptedException {
        String key = RATE_LIMIT_KEY_PREFIX + tenantId;
        long maxTokens = 5; // 最大令牌数
        long refillInterval = 1000; // 令牌补充间隔(ms)
        
        while (true) {
            // 使用 Lua 脚本保证原子性
            String luaScript = 
                "local current_tokens = redis.call('GET', KEYS[1])\n" +
                "local last_refill_time = redis.call('GET', KEYS[2])\n" +
                "local now = tonumber(ARGV[1])\n" +
                "local max_tokens = tonumber(ARGV[2])\n" +
                "local refill_interval = tonumber(ARGV[3])\n" +
                "\n" +
                "if current_tokens == false then\n" +
                "    redis.call('SET', KEYS[1], max_tokens - 1)\n" +
                "    redis.call('SET', KEYS[2], now)\n" +
                "    return 1\n" +
                "end\n" +
                "\n" +
                "current_tokens = tonumber(current_tokens)\n" +
                "last_refill_time = tonumber(last_refill_time)\n" +
                "\n" +
                "local tokens_to_add = math.floor((now - last_refill_time) / refill_interval)\n" +
                "if tokens_to_add > 0 then\n" +
                "    current_tokens = math.min(max_tokens, current_tokens + tokens_to_add)\n" +
                "    last_refill_time = now\n" +
                "end\n" +
                "\n" +
                "if current_tokens > 0 then\n" +
                "    redis.call('SET', KEYS[1], current_tokens - 1)\n" +
                "    redis.call('SET', KEYS[2], last_refill_time)\n" +
                "    return 1\n" +
                "else\n" +
                "    return 0\n" +
                "end";
            
            List<String> keys = Arrays.asList(key, key + ":last_refill");
            List<String> args = Arrays.asList(
                String.valueOf(System.currentTimeMillis()),
                String.valueOf(maxTokens),
                String.valueOf(refillInterval)
            );
            
            Long result = stringRedisTemplate.execute(
                new DefaultRedisScript<>(luaScript, Long.class),
                keys,
                args.toArray(new String[0])
            );
            
            if (result != null && result == 1) {
                // 获取到令牌
                break;
            } else {
                // 没有获取到令牌,等待一段时间后重试
                Thread.sleep(100);
            }
        }
    }
    
    /**
     * 归还令牌(可选,根据业务需求决定是否需要)
     * @param tenantId 租户ID
     */
    private void returnToken(String tenantId) {
        try {
            String key = RATE_LIMIT_KEY_PREFIX + tenantId;
            String luaScript = 
                "local current_tokens = redis.call('GET', KEYS[1])\n" +
                "local max_tokens = tonumber(ARGV[1])\n" +
                "\n" +
                "if current_tokens ~= false then\n" +
                "    current_tokens = tonumber(current_tokens)\n" +
                "    if current_tokens < max_tokens then\n" +
                "        redis.call('SET', KEYS[1], current_tokens + 1)\n" +
                "    end\n" +
                "end";
            
            stringRedisTemplate.execute(
                new DefaultRedisScript<>(luaScript),
                Arrays.asList(key),
                String.valueOf(5)
            );
        } catch (Exception e) {
            log.warn("归还令牌失败,tenantId={}", tenantId, e);
        }
    }
}

### 使用 RedisLua 实现限流算法 #### 固定窗口限流 固定窗口限流是一种简单的限流策略,在指定的时间窗口内允许一定数量的请求通过。一旦超过这个数量,则拒绝后续请求直到下一个时间窗口开始。 ```lua local key = KEYS[1] local limit = tonumber(ARGV[1]) local current_time = tonumber(ARGV[2]) -- 获取当前计数器值,如果不存在则初始化为0 local counter = redis.call("GET", key) if not counter then counter = 0 else counter = tonumber(counter) end -- 如果计数值小于限制,则增加计数并设置过期时间 if counter < limit then local new_counter = counter + 1 redis.call("SET", key, new_counter) redis.call("EXPIRE", key, 60) -- 设置60秒过期时间 return true else return false endif ``` 此代码片段展示了如何利用 Redis 的 `GET` 和 `SET` 命令以及 Lua 脚本实现基于固定窗口的限流逻辑[^1]。 #### 令牌桶算法 令牌桶算法更加灵活,它模拟了一个可以容纳有限数量令牌的桶。系统按照固定的速率向桶里添加令牌,当有新请求到来时会尝试从中取出一个令牌;如果没有足够的令牌可用就拒绝服务。 ```lua local function get_current_timestamp() return tonumber(redis.call('TIME')[1]) end local bucket_key = KEYS[1] local rate_per_second = tonumber(ARGV[1]) -- 每秒钟产生的令牌数目 local capacity = tonumber(ARGV[2]) -- 桶的最大容量 local now = get_current_timestamp() -- 尝试获取存储的状态信息 local state = cjson.decode(redis.call("GET", bucket_key)) if not state or type(state) ~= 'table' then state = {tokens=capacity, last_update=now} end -- 计算自上次更新以来应该补充多少个token local elapsed_seconds = now - state.last_update local tokens_to_add = math.min(capacity, elapsed_seconds * rate_per_second) state.tokens = math.min(capacity, state.tokens + tokens_to_add) state.last_update = now -- 处理本次请求所需的token量 local required_tokens = tonumber(ARGV[3]) if state.tokens >= required_tokens then state.tokens = state.tokens - required_tokens redis.call("SETEX", bucket_key, 86400, cjson.encode(state)) -- 存储一天有效期的数据 return "ALLOWED" else return "DENIED" end ``` 这段 Lua 脚本实现了完整的令牌桶机制,并且借助于 Redis 提供的持久化功能来保存桶的状态信息[^2]。 为了确保高并发环境下的安全性,上述两种方法都依赖于 Redis Lua 脚本提供的原子性执行特性,从而避免了竞争条件的发生。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值