前言
通過使用註解(提供介面形式與引數形式2種方式)針對性的對介面進行限流，底層使用redis配合lua指令碼實現令牌桶。
Redis基礎實現
public class RedisRateLimiter {
private static String REDIS_KEY = "rate:limiter:%s";
private RedisTemplate redisTemplate;
public RedisRateLimiter(RedisTemplate redisTemplate) {
this.redisTemplate = redisTemplate;
}
public boolean tryAcquire(String flag, int maxPermits, int addRate,
long expireSeconds) {
String luaScript = buildLuaScript();
String key = String.format(REDIS_KEY, flag);
Object execute = redisTemplate.execute(new RedisCallback<Long>() {
@Override
public Long doInRedis(RedisConnection connection)
throws DataAccessException {
return connection.eval(luaScript.getBytes(), ReturnType.INTEGER,
1, key.getBytes(),
String.valueOf(maxPermits).getBytes(),
String.valueOf(expireSeconds).getBytes(),
String.valueOf(addRate).getBytes(),
String.valueOf(System.currentTimeMillis()).getBytes());
}
});
Long count = (Long) execute;
return count != 0;
}
private String buildLuaScript() {
StringBuilder lua = new StringBuilder();
lua.append(" local key = KEYS[1]");
lua.append("\nlocal limit = tonumber(ARGV[1])");
lua.append("\nlocal expireSeconds = tonumber(ARGV[2])");
lua.append("\nlocal rate = tonumber(ARGV[3])");
lua.append("\nlocal now = tonumber(ARGV[4])");
lua.append(
"\nlocal rateLimiterInfo = redis.call('HMGET', key, \"count\", \"time\")");
lua.append(
"\nlocal currentPermits =tonumber(rateLimiterInfo[1] or \"0\")");
lua.append(
"\nlocal lastRecoverTime =tonumber(rateLimiterInfo[2] or \"0\")");
lua.append("\nif lastRecoverTime > 0 then");
lua.append(
"\n local recoverPermits = math.floor(((now - lastRecoverTime) / 1000) * rate)");
lua.append("\nif recoverPermits > 0 then");
lua.append(
"\ncurrentPermits = math.max(0, currentPermits - recoverPermits);");
lua.append("\nend");
lua.append("\nend");
lua.append("\nif currentPermits + 1 > limit then");
lua.append("\nreturn 0");
lua.append("\nelse");
lua.append(
"\n redis.call(\"HMSET\", key, \"count\", currentPermits + 1, \"time\", now)");
lua.append("\nredis.call(\"EXPIRE\", key, expireSeconds)");
lua.append("\nreturn currentPermits + 1");
lua.append("\nend");
return lua.toString();
}
}
複製程式碼為什麼使用Hash: 使用叢集版Redis時，必須使所有key在1個slot上。
註解類
@Target({ ElementType.METHOD, ElementType.TYPE })
@Retention(RetentionPolicy.RUNTIME)
public @interface EnableRedisRateLimiter {
/**
* 通過key尋找配置檔案
*
* @return
*/
String key() default "";
}
複製程式碼新建配置檔案
demo.flag=
demo.params=demo
demo.maxPermits=10
demo.addRate=1
demo.expireSeconds=10
複製程式碼此配置檔案供下文RedisRateLimiterConfigManager類使用
配置類
@Configuration
public class RedisRateLimiterConfiguration implements ApplicationContextAware {
private ApplicationContext ctx;
@Bean("redisRateLimiter")
public RedisRateLimiter redisRateLimiter() {
String beanName = "redisTemplate-rl";
boolean containsBean = ctx.containsBean(beanName);
if (containsBean) {
RedisTemplate redisTemplate = (RedisTemplate) ctx.getBean(beanName);
return new RedisRateLimiter(redisTemplate);
} else {
return null;
}
}
@Override
public void setApplicationContext(ApplicationContext applicationContext)
throws BeansException {
this.ctx = applicationContext;
}
}
複製程式碼RedisTemplate類注入程式碼省略，整合jedis。
Aspect
@Aspect
@Component
public class RedisRateLimiterAspect {
private final static String FLAG_SUFFIX = ".flag";//根據flag指定限流
private final static String PARAMS_SUFFIX = ".params";//根據引數指定限流，如果根據引數指定了則不必填flag
private final static String MAX_PERMITS_SUFFIX = ".maxPermits";//最大令牌數
private final static String ADD_RATE_SUFFIX = ".addRate";//每秒增加的令牌數
private final static String EXPIRE_SECONDS_SUFFIX = ".expireSeconds";//過期時間
@Autowired
private RedisRateLimiter redisRateLimiter;
@Before("execution(* 此處根據自己環境填寫路徑")
public void beforeMethod(JoinPoint joinPoint) {
Signature signature = joinPoint.getSignature();
MethodSignature methodSignature = (MethodSignature) signature;
EnableRedisRateLimiter annotation = methodSignature.getMethod()
.getAnnotation(EnableRedisRateLimiter.class);
if (annotation != null) {
String key = annotation.key();
if (!StringUtil.isNullOrEmpty(key)) {
String flag = RedisRateLimiterConfigManager
.getProperty(combineKey(key, FLAG_SUFFIX));
if (StringUtil.isNullOrEmpty(flag)) {
String paramName = RedisRateLimiterConfigManager
.getProperty(combineKey(key, PARAMS_SUFFIX));
if (StringUtil.isNullOrEmpty(paramName)) {
return;
}
String[] paramNames = paramName.split(",");
String[] argNames = methodSignature.getParameterNames(); // 引數名
Object[] args = joinPoint.getArgs();// 引數值
flag = Arrays.stream(paramNames).map(p -> {
for (int i = 0; i < argNames.length; i++) {
if (argNames[i].equals(p)) {
try {
return args[i].toString();
} catch (Exception e) {
return null;
}
}
}
return null;
}).filter(Objects::nonNull)
.collect(Collectors.joining(":"));
}
if (StringUtil.isNullOrEmpty(flag)) {
return;
}
String strMaxPermits = RedisRateLimiterConfigManager
.getProperty(combineKey(key, MAX_PERMITS_SUFFIX));
int maxPermits = StringUtil.isNullOrEmpty(strMaxPermits) ? 0
: Integer.parseInt(strMaxPermits);
String strAddRate = RedisRateLimiterConfigManager
.getProperty(combineKey(key, ADD_RATE_SUFFIX));
int addRate = StringUtil.isNullOrEmpty(strAddRate) ? 0
: Integer.parseInt(strAddRate);
String strExpireSeconds = RedisRateLimiterConfigManager
.getProperty(combineKey(key, EXPIRE_SECONDS_SUFFIX));
long expireSeconds = StringUtil.isNullOrEmpty(strExpireSeconds)
? 0
: Integer.parseInt(strExpireSeconds);
String finalFlag = key + ":" + flag;
boolean success = redisRateLimiter.tryAcquire(finalFlag,
maxPermits, addRate, expireSeconds);
if (!success) {
throw error;
}
}
}
}
public String combineKey(String key, String suffix) {
return key + suffix;
}
}
複製程式碼上述程式碼中RedisRateLimiterConfigManager程式碼省略，內部實現為通過combineKey方法後的值去獲取對應的配置檔案中value。
結尾
上述程式碼都完成後，只需在需要限流的介面上註上@EnableRedisRateLimiter(key = "demo")即可，如果有更好的實現方式可以留言探討。