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