对接口实现流量控制,使其在N秒内访问控制在M次
lua脚本实现令牌桶算法(来自网络):
--令牌秒杀 Lua脚本 tokenTimeLimiter.lua
local key = KEYS[1] --限流KEY(一秒一个)
local limit = tonumber(ARGV[1]) --限流大小
local exprie = ARGV[2] --过期时间
-- 获取当前计数值
local current = tonumber(redis.call('get', key) or "0")
if current + 1 > limit then --如果超出限流大小
return 0
else
current = tonumber(redis.call("INCRBY", key, "1")) --请求数+1
if current == 1 then --第一次访问需要设置过期时间
redis.call("expire", key,exprie) --设置过期时间
end
end
return 1 --返回1代表不限流
自定义注解:
package com.suning.shengyu.web.aop;
import java.lang.annotation.*;
@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
@Documented
public @interface RateLimit {
}
切面:
package com.suning.shengyu.web.aop;
import com.alibaba.fastjson.JSON;
import com.suning.framework.sedis.ShardedJedisAction;
import com.suning.framework.sedis.impl.ShardedJedisClientImpl;
import com.suning.shengyu.common.support.log.ShengyuLogger;
import com.suning.shengyu.common.support.outeropenapi.OpenApiRequest;
import com.suning.shengyu.common.support.outeropenapi.OpenResultCode;
import com.suning.shengyu.common.support.outeropenapi.ResultMap;
import com.suning.shengyu.util.AESCoder;
import com.suning.shengyu.web.util.SignPlugin;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import redis.clients.jedis.Jedis;
import redis.clients.jedis.ShardedJedis;
import java.util.ArrayList;
import java.util.List;
@Aspect
@Component
public class RateLimitAdvice {
private static final ShengyuLogger LOGGER = ShengyuLogger.getLogger(RateLimitAdvice.class);
private static final String FLOWCONTROL_CACHE_STRING_KEY = "SHENGYU:MULTI-STRING:RATELIMIT:";
//时间间隔
private static final String INTERVAL_SECOND = "600";
//方法调用次数
private static final String METHOD_CALL_TIMES = "120";
@Autowired
private SignPlugin signPlugin;
@Autowired
// private ShardedJedisClientImpl jedisClientForSys;
private ShardedJedisClientImpl jedisClient;
@Autowired
private String luaScript;
@Pointcut(value = "@annotation(com.suning.shengyu.web.aop.RateLimit)")
private void pointCut(){
System.out.println("我是切入点。。。");
};
@Around("pointCut()")
public ResultMap around(ProceedingJoinPoint proceedingJoinPoint) {
Object[] args = proceedingJoinPoint.getArgs();
OpenApiRequest openApiRequest = (OpenApiRequest)args[0];
String mallId = AESCoder.decrypt(openApiRequest.getAppKey(), signPlugin.md5(AESCoder.KEYSTRHEXSTR_APPKEY));
ResultMap resultMap;
String method = null;
try {
method = openApiRequest.getMethod();
LOGGER.info("流控切面获取的mallId={}, method={}", mallId, method);
List<String> values = new ArrayList<>();
values.add(METHOD_CALL_TIMES);
values.add(INTERVAL_SECOND);
String key = FLOWCONTROL_CACHE_STRING_KEY + mallId + method;
Long executeResult = luaEval(luaScript, key, values);
// Long executeResult = 0l;
LOGGER.info("流控令牌数量剩余{}", executeResult);
if (executeResult == 1) {
LOGGER.info("当前访问时间段内剩余{}次访问次数", executeResult);
Object obj = proceedingJoinPoint.proceed();
System.out.println(JSON.toJSONString(obj));
resultMap = (ResultMap)obj;
} else {
LOGGER.info("当前接口调用超过时间段内限流,key:{}", key);
return ResultMap.fail(OpenResultCode.REQUEST_LIMITED);
}
} catch (Throwable throwable) {
LOGGER.error("接口限流层异常,mallId={},method={}", mallId, method, throwable);
return ResultMap.fail(OpenResultCode.SYSTEM_ERROR);
}
return resultMap;
}
public Long luaEval(final String LUA_SCRIPT, final String key, final List<String> args) {
return jedisClient.execute(new ShardedJedisAction<Long>() {
@Override
public Long doAction(ShardedJedis shardedJedis) {
Jedis jedis = shardedJedis.getShard(key);
// 参数KEY
List<String> keys = new ArrayList<>();
keys.add(key);
// 脚本调用
return (Long) jedis.eval(LUA_SCRIPT, keys, args);
}
});
}
}
读取lua脚本:
@Component
public class RateLimitConfig {
private static final ShengyuLogger LOGGER = ShengyuLogger.getLogger(RateLimitConfig.class);
private static final String SCRIPT_FILE_NAME = "ratelimit.lua";
@Bean({ "luaScript" })
public String luaScript() {
String luaScript = "";
try {
InputStream in = RateLimitConfig.class.getClassLoader().getResourceAsStream(SCRIPT_FILE_NAME);
luaScript = new String(ByteStreams.toByteArray(in));
LOGGER.info("流控脚本:{}", luaScript);
} catch (IOException e) {
LOGGER.info("加载lua脚本时异常", e);
}
return luaScript;
}
}
使用举例:
@ResponseBody
@RequestMapping("/testApi")
@RateLimit
public ResultMap testApi(@RequestBody OpenApiRequest openApiRequest) {
//业务逻辑
return ResultMap.ok();
}