Redis分布式限流依赖分享及使用

 1.基于tap包(com.github.taptap)

   <dependency>
            <groupId>com.github.taptap</groupId>
            <artifactId>ratelimiter-spring-boot-starter</artifactId>
            <version>1.3</version>
        </dependency>

lua脚本写法，此方法不需要额外引用redssionJar包，对依赖有严格控制的项目比较友好，当然不方便引入依赖的话可以直接用下面lua脚本。

-- https://gist.github.com/ptarjan/e38f45f2dfe601419ca3af937fff574d#file-1-check_request_rate_limiter-rb-L11-L34
redis.replicate_commands()

local tokens_key = KEYS[1]
local timestamp_key = KEYS[2]

local rate = tonumber(ARGV[1])
local capacity = tonumber(ARGV[2])
local requested = tonumber(ARGV[3])
local now = redis.call('TIME')[1]

local fill_time = capacity/rate
local ttl = math.floor(fill_time*2)

local last_tokens = tonumber(redis.call("get", tokens_key))
if last_tokens == nil then
  last_tokens = capacity
end

local last_refreshed = tonumber(redis.call("get", timestamp_key))
if last_refreshed == nil then
  last_refreshed = 0
end

local delta = math.max(0, now-last_refreshed)
local filled_tokens = math.min(capacity, last_tokens+(delta*rate))
local allowed = filled_tokens >= requested
local new_tokens = filled_tokens
local allowed_num = 0
if allowed then
  new_tokens = filled_tokens - requested
  allowed_num = 1
end

if ttl > 0 then
  redis.call("setex", tokens_key, ttl, new_tokens)
  redis.call("setex", timestamp_key, ttl, now)
end

return { allowed_num, new_tokens }

 Java端代码

//
// Source code recreated from a .class file by IntelliJ IDEA
// (powered by FernFlower decompiler)
//

package com.taptap.ratelimiter.core;

import com.taptap.ratelimiter.model.LuaScript;
import com.taptap.ratelimiter.model.Result;
import com.taptap.ratelimiter.model.Rule;
import java.util.Arrays;
import java.util.List;
import org.redisson.api.RScript;
import org.redisson.api.RedissonClient;
import org.redisson.api.RScript.Mode;
import org.redisson.api.RScript.ReturnType;
import org.redisson.client.codec.LongCodec;
import org.springframework.stereotype.Component;

@Component
public class TokenBucketRateLimiter implements RateLimiter {
    private final RScript rScript;

    public TokenBucketRateLimiter(RedissonClient client) {
        this.rScript = client.getScript(LongCodec.INSTANCE);
    }

    public Result isAllowed(Rule rule) {
        List<Object> keys = getKeys(rule.getKey());
        String script = LuaScript.getTokenBucketRateLimiterScript();
        List<Long> results = (List)this.rScript.eval(Mode.READ_WRITE, script, ReturnType.MULTI, keys, new Object[]{rule.getRate(), rule.getBucketCapacity(), rule.getRequestedTokens()});
        boolean isAllowed = (Long)results.get(0) == 1L;
        long newTokens = (Long)results.get(1);
        return new Result(isAllowed, newTokens);
    }

    static List<Object> getKeys(String key) {
        String prefix = "request_rate_limiter.{" + key;
        String tokenKey = prefix + "}.tokens";
        String timestampKey = prefix + "}.timestamp";
        return Arrays.asList(tokenKey, timestampKey);
    }
}

redis是可以执行lua脚本的，所以也不一定要用demo中提供的执行类rScript。

2.基于(io.github.forezp)

 <dependency>
            <groupId>io.github.forezp</groupId>
            <artifactId>distributed-limit-core</artifactId>
            <version>1.0.4</version>
        </dependency>

主题实现代码

//
// Source code recreated from a .class file by IntelliJ IDEA
// (powered by FernFlower decompiler)
//

package io.github.forezp.distributedlimitcore.limit;

import io.github.forezp.distributedlimitcore.entity.LimitEntity;
import io.github.forezp.distributedlimitcore.entity.LimitResult;
import io.github.forezp.distributedlimitcore.entity.LimitResult.ResultType;
import io.github.forezp.distributedlimitcore.util.KeyUtil;
import java.util.ArrayList;
import java.util.List;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.util.StringUtils;

public class RedisLimitExcutor implements LimitExcutor {
    private StringRedisTemplate stringRedisTemplate;
    Logger log = LoggerFactory.getLogger(RedisLimitExcutor.class);

    public RedisLimitExcutor(StringRedisTemplate stringRedisTemplate) {
        this.stringRedisTemplate = stringRedisTemplate;
    }

    public LimitResult tryAccess(LimitEntity limitEntity) {
        String identifier = limitEntity.getIdentifier();
        String key = KeyUtil.getKey(limitEntity);
        if (StringUtils.isEmpty(key)) {
            return null;
        } else {
            int seconds = limitEntity.getSeconds();
            int limitCount = limitEntity.getLimtNum();
            List<String> keys = new ArrayList();
            keys.add(key);
            String luaScript = this.buildLuaScript();
            RedisScript<Long> redisScript = new DefaultRedisScript(luaScript, Long.class);
            Long count = (Long)this.stringRedisTemplate.execute(redisScript, keys, new Object[]{"" + limitCount, "" + seconds});
            this.log.info("Access try count is {} for key={}", count, key);
            LimitResult result = new LimitResult();
            result.setUrl(key);
            result.setIdenfier(identifier);
            if (count != 0L) {
                result.setResultType(ResultType.SUCCESS);
            } else {
                result.setResultType(ResultType.FAIL);
            }

            return result;
        }
    }

    private String buildLuaScript() {
        StringBuilder lua = new StringBuilder();
        lua.append(" local key = KEYS[1]");
        lua.append("\nlocal limit = tonumber(ARGV[1])");
        lua.append("\nlocal curentLimit = tonumber(redis.call('get', key) or \"0\")");
        lua.append("\nif curentLimit + 1 > limit then");
        lua.append("\nreturn 0");
        lua.append("\nelse");
        lua.append("\n redis.call(\"INCRBY\", key, 1)");
        lua.append("\nredis.call(\"EXPIRE\", key, ARGV[2])");
        lua.append("\nreturn curentLimit + 1");
        lua.append("\nend");
        return lua.toString();
    }
}

这个的lua脚本更加简化。

其中identifier为识别身份的，key为限流的key,limtNum为限制的次数，seconds为多少秒，后2个配置的作用是在多少秒最大的请求次数 。其中identifier和key支持Spel表达式。如果仅API纬度，则identifier 为空即可；如果仅用户纬度，key为空即可。

3.redis集群lua脚本

local count
count = redis.call('get',KEYS[1])
--不超过最大值，则直接返回
if count and tonumber(count) > tonumber(ARGV[1]) then
    return count;
end
    --执行计算器自加
    count = redis.call('incr',KEYS[1])
if tonumber(count) == 1 then
    --从第一次调用开始限流，设置对应key的过期时间
    redis.call('expire',KEYS[1],ARGV[2])
end
return count;

或者

--KEYS[1]: 限流 key
--ARGV[1]: 时间戳 - 时间窗口
--ARGV[2]: 当前时间戳（作为score）
--ARGV[3]: 阈值
--ARGV[4]: score 对应的唯一value
-- 1. 移除时间窗口之前的数据
redis.call('zremrangeByScore', KEYS[1], 0, ARGV[1])
-- 2. 统计当前元素数量
local res = redis.call('zcard', KEYS[1])
-- 3. 是否超过阈值
if (res == nil) or (res < tonumber(ARGV[3])) then
    redis.call('zadd', KEYS[1], ARGV[2], ARGV[4])
    return 1
else
    return 0
end