SpringBoot - 流量控制(二)：基于REDIS+LUA脚本的限流
写在前面
公司近期在做对外接口的限流，由于是轻量级的处理，所以采用“基于REDIS+LUA脚本”的方式进行访问控制。
当然也可以使用RequestRateLimiterGatewayFilterFactory或者阿里巴巴的sentinel。
SpringCloud - 流量控制(一)：基于RequestRateLimiterGatewayFilterFactory的限流
如何编写？
① 定义注解
/**
* 限流注解
* @author ROCKY
*/
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface Limiter {
// 限流KEY,用于REDIS的存储
public String key()；
// 限流时间(秒)
public int time() default 60;
// 限流次数
public int count() default 100;
// 限流类型
public LimitType limitType() default LimitType.DEFAULT;
}
② 切面处理
/**
* 限流AOP
* @author ROCKY
*/
@Aspect
@Component
@Slf4j
public class LimiterAspect {
@Autowired
private RedisTemplate<Object, Object> redisTemplate;
@Autowired
private RedisScript<Long> limitScript;
@Before("@annotation(limiter)")
public void doBefore(JoinPoint point, Limiter limiter) throws Throwable {
// 可以使用指定的KEY，也可以自己封装
// final String limiterKey = limiter.key();
final String limiterKey = getLimiterKey(limiter, point);
int time = limiter.time();
int count = limiter.count();
List<Object> keys = Collections.singletonList(limiterKey);
try {
Long number = redisTemplate.execute(limitScript, keys, count, time);
if (StringUtils.isEmpty(number) || number.intValue() > count) {
log.error("限流请求总次数为: '{}', 当前为第'{}'请求, 缓存KEY为: '{}'", count, number.intValue(), limiterKey);
throw new MyException(MyException.LIMITER_IN_THE_CONTROL);
}
} catch (MyException me) {
throw me;
} catch (Exception e) {
throw new RuntimeException("限流控制内部异常，请稍候再试");
}
}
// 获取限流数据的KEY(小写)
public String getLimiterKey(Limiter limiter, JoinPoint point) {
StringBuffer stringBuffer = new StringBuffer(Constants.LIMITER_IN_REDIS_PREFIX_KEY);
if (limiter.limitType() == LimitType.IP) {
stringBuffer.append(IpUtils.getIpAddr(ServletUtils.getRequest())).append("-");
}
if (limiter.limitType() == LimitType.USER) {
// 从RequestContextHolder中获取用户名
...
}
if (limiter.limitType() == LimitType.TOKEN) {
// 从RequestContextHolder中获取TOKEN
...
}
MethodSignature signature = (MethodSignature) point.getSignature();
Method method = signature.getMethod();
Class<?> targetClass = method.getDeclaringClass();
stringBuffer.append(targetClass.getName()).append("-").append(method.getName());
return stringBuffer.toString().toLowerCase();
}
}
③ 定义脚本
/**
* LUA限流脚本
* @author ROCKY
*/
@Configuration
@EnableCaching
public class RedisScriptConfig extends CachingConfigurerSupport {
@Bean
public DefaultRedisScript<Long> limitScript() {
DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
redisScript.setScriptText(limitScriptText());
redisScript.setResultType(Long.class);
return redisScript;
}
/**
* 限流脚本
*/
private String limitScriptText() {
return "local key = KEYS[1]\n" +
"local count = tonumber(ARGV[1])\n" +
"local time = tonumber(ARGV[2])\n" +
"local current = redis.call('get', key);\n" +
"if current and tonumber(current) > count then\n" +
" return tonumber(current);\n" +
"end\n" +
"current = redis.call('incr', key)\n" +
"if tonumber(current) == 1 then\n" +
" redis.call('expire', key, time)\n" +
"end\n" +
"return tonumber(current);";
}
}
④ 其他定义
/**
* 限流类型
* @author ROCKY
*/
public enum LimitType {
// 默认策略全局限流
DEFAULT,
// 根据请求者IP进行限流
IP,
// 根据请求者进行限流
USER,
// 根据TOKEN进行限流
TOKEN
}
/**
* SERVLET工具类
* @author ROCKY
*/
public class ServletUtils {
// 获取String参数
public static String getParameter(String name) {
return getRequest().getParameter(name);
}
// 获取request
public static HttpServletRequest getRequest() {
try {
return getRequestAttributes().getRequest();
} catch (Exception e) {
return null;
}
}
// 获取response
public static HttpServletResponse getResponse() {
try {
return getRequestAttributes().getResponse();
} catch (Exception e) {
return null;
}
}
// 获取session
public static HttpSession getSession() {
return getRequest().getSession();
}
public static ServletRequestAttributes getRequestAttributes() {
try {
RequestAttributes attributes = RequestContextHolder.getRequestAttributes();
return (ServletRequestAttributes) attributes;
} catch (Exception e) {
return null;
}
}
public static Map<String, String> getHeaders(HttpServletRequest request) {
Map<String, String> map = new LinkedHashMap<>();
Enumeration<String> enumeration = request.getHeaderNames();
if (enumeration != null) {
while (enumeration.hasMoreElements()) {
String key = enumeration.nextElement();
String value = request.getHeader(key);
map.put(key, value);
}
}
return map;
}
/**
* 将字符串渲染到客户端
*/
public static String renderString(HttpServletResponse response, String string) {
try {
response.setStatus(200);
response.setContentType("application/json");
response.setCharacterEncoding("utf-8");
response.getWriter().print(string);
} catch (IOException e) {
e.printStackTrace();
}
return null;
}
/**
* 是否包含字符串
*/
private static boolean inStringIgnoreCase(String str, String... strs) {
if (str != null && strs != null) {
for (String s : strs) {
if (str.equalsIgnoreCase(StringUtils.trim(s))) {
return true;
}
}
}
return false;
}
}
/**
* 获取IP工具类
* @author ROCKY
*/
public class IpUtils {
public static String getIpAddr(HttpServletRequest request) {
if (request == null) {
return null;
}
String ip = null;
// X-Forwarded-For：Squid 服务代理
String ipAddresses = request.getHeader("X-Forwarded-For");
if (ipAddresses == null || ipAddresses.length() == 0 || "unknown".equalsIgnoreCase(ipAddresses)) {
// Proxy-Client-IP：apache 服务代理
ipAddresses = request.getHeader("Proxy-Client-IP");
}
if (ipAddresses == null || ipAddresses.length() == 0 || "unknown".equalsIgnoreCase(ipAddresses)) {
// WL-Proxy-Client-IP：weblogic 服务代理
ipAddresses = request.getHeader("WL-Proxy-Client-IP");
}
if (ipAddresses == null || ipAddresses.length() == 0 || "unknown".equalsIgnoreCase(ipAddresses)) {
// HTTP_CLIENT_IP：有些代理服务器
ipAddresses = request.getHeader("HTTP_CLIENT_IP");
}
if (ipAddresses == null || ipAddresses.length() == 0 || "unknown".equalsIgnoreCase(ipAddresses)) {
// X-Real-IP：nginx服务代理
ipAddresses = request.getHeader("X-Real-IP");
}
// 有些网络通过多层代理，那么获取到的ip就会有多个，一般都是通过逗号（,）分割开来，并且第一个ip为客户端的真实IP
if (ipAddresses != null && ipAddresses.length() != 0) {
ip = ipAddresses.split(",")[0];
}
// 还是不能获取到，最后再通过request.getRemoteAddr();获取
if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ipAddresses)) {
ip = request.getRemoteAddr();
}
return ip.equals("0:0:0:0:0:0:0:1") ? "127.0.0.1" : ip;
}
public static boolean internalIp(String ip) {
byte[] addr = textToNumericFormatV4(ip);
return internalIp(addr) || "127.0.0.1".equals(ip);
}
private static boolean internalIp(byte[] addr) {
if (null == addr || addr.length < 2) {
return true;
}
final byte b0 = addr[0];
final byte b1 = addr[1];
// 10.x.x.x/8
final byte SECTION_1 = 0x0A;
// 172.16.x.x/12
final byte SECTION_2 = (byte) 0xAC;
final byte SECTION_3 = (byte) 0x10;
final byte SECTION_4 = (byte) 0x1F;
// 192.168.x.x/16
final byte SECTION_5 = (byte) 0xC0;
final byte SECTION_6 = (byte) 0xA8;
switch (b0) {
case SECTION_1:
return true;
case SECTION_2:
if (b1 >= SECTION_3 && b1 <= SECTION_4) {
return true;
}
case SECTION_5:
switch (b1) {
case SECTION_6:
return true;
}
default:
return false;
}
}
// 将IPv4地址转换成字节
public static byte[] textToNumericFormatV4(String text) {
if (text.length() == 0) {
return null;
}
byte[] bytes = new byte[4];
String[] elements = text.split("\\.", -1);
try {
long l;
int i;
switch (elements.length) {
case 1:
l = Long.parseLong(elements[0]);
if ((l < 0L) || (l > 4294967295L)) {
return null;
}
bytes[0] = (byte) (int) (l >> 24 & 0xFF);
bytes[1] = (byte) (int) ((l & 0xFFFFFF) >> 16 & 0xFF);
bytes[2] = (byte) (int) ((l & 0xFFFF) >> 8 & 0xFF);
bytes[3] = (byte) (int) (l & 0xFF);
break;
case 2:
l = Integer.parseInt(elements[0]);
if ((l < 0L) || (l > 255L)) {
return null;
}
bytes[0] = (byte) (int) (l & 0xFF);
l = Integer.parseInt(elements[1]);
if ((l < 0L) || (l > 16777215L)) {
return null;
}
bytes[1] = (byte) (int) (l >> 16 & 0xFF);
bytes[2] = (byte) (int) ((l & 0xFFFF) >> 8 & 0xFF);
bytes[3] = (byte) (int) (l & 0xFF);
break;
case 3:
for (i = 0; i < 2; ++i) {
l = Integer.parseInt(elements[i]);
if ((l < 0L) || (l > 255L)) {
return null;
}
bytes[i] = (byte) (int) (l & 0xFF);
}
l = Integer.parseInt(elements[2]);
if ((l < 0L) || (l > 65535L)) {
return null;
}
bytes[2] = (byte) (int) (l >> 8 & 0xFF);
bytes[3] = (byte) (int) (l & 0xFF);
break;
case 4:
for (i = 0; i < 4; ++i) {
l = Integer.parseInt(elements[i]);
if ((l < 0L) || (l > 255L)) {
return null;
}
bytes[i] = (byte) (int) (l & 0xFF);
}
break;
default:
return null;
}
} catch (NumberFormatException e) {
return null;
}
return bytes;
}
public static String getHostIp() {
try {
return InetAddress.getLocalHost().getHostAddress();
} catch (UnknownHostException e) {
}
return "127.0.0.1";
}
public static String getHostPort(HttpServletRequest request) {
return String.valueOf(request.getLocalPort());
}
public static String getHostName() {
try {
return InetAddress.getLocalHost().getHostName();
} catch (UnknownHostException e) {
}
return "未知";
}
}
如何调试？
LUA调试器