使用AOP+redis+lua做方法限流的实现_F11 - 专业站长和开发者的学习网站

需求

公司里使用OneByOne的方式删除数据，为了防止一段时间内删除数据过多，让我这边做一个接口限流，超过一定阈值后报异常，终止删除操作。

实现方式

创建自定义注解 @limit 让使用者在需要的地方配置 count(一定时间内最多访问次数)、 period(给定的时间范围)，也就是访问频率。然后通过LimitInterceptor拦截方法的请求，通过 redis+lua 脚本的方式，控制访问频率。

源码

Limit 注解

用于配置方法的访问频率count、period

import javax.validation.constraints.Min;

import java.lang.annotation.*;

@Target({ElementType.METHOD, ElementType.TYPE})

@Retention(RetentionPolicy.RUNTIME)

@Inherited

@Documented

public @interface Limit {

/**

* key

String key() default "";

/**

* Key的前缀

String prefix() default "";

/**

* 一定时间内最多访问次数

@Min(1)

int count();

/**

* 给定的时间范围单位(秒)

@Min(1)

int period();

/**

* 限流的类型(用户自定义key或者请求ip)

LimitType limitType() default LimitType.CUSTOMER;

}

LimitKey

用于标记参数，作为redis key值的一部分

import java.lang.annotation.ElementType;

import java.lang.annotation.Retention;

import java.lang.annotation.RetentionPolicy;

import java.lang.annotation.Target;

@Target(ElementType.PARAMETER)

@Retention(RetentionPolicy.RUNTIME)

public @interface LimitKey {

}

LimitType

枚举，redis key值的类型，支持自定义key和ip、methodName中获取key

public enum LimitType {

/**

* 自定义key

CUSTOMER,

/**

* 请求者IP

IP,

/**

* 方法名称

METHOD_NAME;

}

RedisLimiterHelper

初始化一个限流用到的redisTemplate Bean

import org.springframework.beans.factory.annotation.Qualifier;

import org.springframework.context.annotation.Bean;

import org.springframework.context.annotation.Configuration;

import org.springframework.data.redis.core.RedisTemplate;

import org.springframework.data.redis.core.StringRedisTemplate;

import org.springframework.data.redis.serializer.GenericJackson2JsonRedisSerializer;

import org.springframework.data.redis.serializer.StringRedisSerializer;

import java.io.Serializable;

@Configuration

public class RedisLimiterHelper {

@Bean

public RedisTemplate<String, Serializable> limitRedisTemplate(@Qualifier("defaultStringRedisTemplate") StringRedisTemplate redisTemplate) {

RedisTemplate<String, Serializable> template = new RedisTemplate<String, Serializable>();

template.setKeySerializer(new StringRedisSerializer());

template.setValueSerializer(new GenericJackson2JsonRedisSerializer());

template.setConnectionFactory(redisTemplate.getConnectionFactory());

return template;

}

LimitInterceptor

使用 aop 的方式来拦截请求，控制访问频率

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

import com.google.common.collect.ImmutableList;

import com.yxt.qida.api.bean.service.xxv2.openapi.anno.Limit;

import com.yxt.qida.api.bean.service.xxv2.openapi.anno.LimitKey;

import com.yxt.qida.api.bean.service.xxv2.openapi.anno.LimitType;

import lombok.extern.slf4j.Slf4j;

import org.apache.commons.lang3.ArrayUtils;

import org.apache.commons.lang3.StringUtils;

import org.aspectj.lang.ProceedingJoinPoint;

import org.aspectj.lang.annotation.Around;

import org.aspectj.lang.annotation.Aspect;

import org.aspectj.lang.reflect.MethodSignature;

import org.springframework.beans.factory.annotation.Autowired;

import org.springframework.context.annotation.Configuration;

import org.springframework.data.redis.core.RedisTemplate;

import org.springframework.data.redis.core.script.DefaultRedisScript;

import org.springframework.data.redis.core.script.RedisScript;

import org.springframework.web.context.request.RequestContextHolder;

import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;

import java.io.Serializable;

import java.lang.annotation.Annotation;

import java.lang.reflect.Method;

@Slf4j

@Aspect

@Configuration

public class LimitInterceptor {

private static final String UNKNOWN = "unknown";

private final RedisTemplate<String, Serializable> limitRedisTemplate;

@Autowired

public LimitInterceptor(RedisTemplate<String, Serializable> limitRedisTemplate) {

this.limitRedisTemplate = limitRedisTemplate;

}

@Around("execution(public * *(..)) && @annotation(com.yxt.qida.api.bean.service.xxv2.openapi.anno.Limit)")

public Object interceptor(ProceedingJoinPoint pjp) {

MethodSignature signature = (MethodSignature) pjp.getSignature();

Method method = signature.getMethod();

Limit limitAnnotation = method.getAnnotation(Limit.class);

LimitType limitType = limitAnnotation.limitType();

int limitPeriod = limitAnnotation.period();

int limitCount = limitAnnotation.count();

/**

* 根据限流类型获取不同的key ,如果不传我们会以方法名作为key

String key;

switch (limitType) {

case IP:

key = getIpAddress();

break;

case CUSTOMER:

key = limitAnnotation.key();

break;

case METHOD_NAME:

String methodName = method.getName();

key = StringUtils.upperCase(methodName);

break;

default:

throw new RuntimeException("limitInterceptor - 无效的枚举值");

}

/**

* 获取注解标注的 key，这个是优先级最高的，会覆盖前面的 key 值

Object[] args = pjp.getArgs();

Annotation[][] paramAnnoAry = method.getParameterAnnotations();

for (Annotation[] item : paramAnnoAry) {

int paramIndex = ArrayUtils.indexOf(paramAnnoAry, item);

for (Annotation anno : item) {

if (anno instanceof LimitKey) {

Object arg = args[paramIndex];

if (arg instanceof String && StringUtils.isNotBlank((String) arg)) {

key = (String) arg;

break;

}

if (StringUtils.isBlank(key)) {

throw new RuntimeException("limitInterceptor - key值不能为空");

}

String prefix = limitAnnotation.prefix();

String[] keyAry = StringUtils.isBlank(prefix) ? new String[]{"limit", key} : new String[]{"limit", prefix, key};

ImmutableList<String> keys = ImmutableList.of(StringUtils.join(keyAry, "-"));

try {

String luaScript = buildLuaScript();

RedisScript<Number> redisScript = new DefaultRedisScript<Number>(luaScript, Number.class);

Number count = limitRedisTemplate.execute(redisScript, keys, limitCount, limitPeriod);

if (count != null && count.intValue() <= limitCount) {

return pjp.proceed();

} else {

String classPath = method.getDeclaringClass().getName() + "." + method.getName();

throw new RuntimeException("limitInterceptor - 限流被触发："

+ "class:" + classPath

+ ", keys:" + keys

+ ", limitcount:" + limitCount

+ ", limitPeriod:" + limitPeriod + "s");

}

} catch (Throwable e) {

if (e instanceof RuntimeException) {

throw new RuntimeException(e.getLocalizedMessage());

}

throw new RuntimeException("limitInterceptor - 限流服务异常");

}

/**

* lua 脚本，为了保证执行 redis 命令的原子性

public String buildLuaScript() {

StringBuilder lua = new StringBuilder();

lua.append("local c");

lua.append("\nc = redis.call('get',KEYS[1])");

// 调用不超过最大值，则直接返回

lua.append("\nif c and tonumber(c) > tonumber(ARGV[1]) then");

lua.append("\nreturn c;");

lua.append("\nend");

// 执行计算器自加

lua.append("\nc = redis.call('incr',KEYS[1])");

lua.append("\nif tonumber(c) == 1 then");

// 从第一次调用开始限流，设置对应键值的过期

lua.append("\nredis.call('expire',KEYS[1],ARGV[2])");

lua.append("\nend");

lua.append("\nreturn c;");

return lua.toString();

}

public String getIpAddress() {

HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();

String ip = request.getHeader("x-forwarded-for");

if (ip == null || ip.length() == 0 || UNKNOWN.equalsIgnoreCase(ip)) {

ip = request.getHeader("Proxy-Client-IP");

}

if (ip == null || ip.length() == 0 || UNKNOWN.equalsIgnoreCase(ip)) {

ip = request.getHeader("WL-Proxy-Client-IP");

}

if (ip == null || ip.length() == 0 || UNKNOWN.equalsIgnoreCase(ip)) {

ip = request.getRemoteAddr();

}

return ip;

}

TestService

使用方式示例

@Limit(period = 10, count = 10)

public String delUserByUrlTest(@LimitKey String token, String thirdId, String url) throws IOException {

return "success";

}