使用AOP+redis+lua做方法限流的实现
时间:2022-06-02 10:04:06|栏目:Redis|点击: 次
需求
公司里使用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 的方式来拦截请求,控制访问频率
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"; }