redis做接口限流(eladmin)
LimitAspect
package me.zhengjie.aspect;
import com.google.common.collect.ImmutableList;
import me.zhengjie.annotation.Limit;
import me.zhengjie.exception.BadRequestException;
import me.zhengjie.utils.RequestHolder;
import me.zhengjie.utils.StringUtils;
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.aspectj.lang.reflect.MethodSignature;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
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.stereotype.Component;
import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
/**
* @author /
*/
@Aspect
@Component
public class LimitAspect {
private final RedisTemplate<Object,Object> redisTemplate;
private static final Logger logger = LoggerFactory.getLogger(LimitAspect.class);
public LimitAspect(RedisTemplate<Object,Object> redisTemplate) {
this.redisTemplate = redisTemplate;
}
@Pointcut("@annotation(me.zhengjie.annotation.Limit)")
public void pointcut() {
}
@Around("pointcut()")
public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
HttpServletRequest request = RequestHolder.getHttpServletRequest();
MethodSignature signature = (MethodSignature) joinPoint.getSignature();
Method signatureMethod = signature.getMethod();
Limit limit = signatureMethod.getAnnotation(Limit.class);
LimitType limitType = limit.limitType();
String key = limit.key();
if (StringUtils.isEmpty(key)) {
if (limitType == LimitType.IP) {
key = StringUtils.getIp(request);
} else {
key = signatureMethod.getName();
}
}
ImmutableList<Object> keys = ImmutableList.of(StringUtils.join(limit.prefix(), "_", key, "_", request.getRequestURI().replaceAll("/","_")));
String luaScript = buildLuaScript();
RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
Number count = redisTemplate.execute(redisScript, keys, limit.count(), limit.period());
if (null != count && count.intValue() <= limit.count()) {
logger.info("第{}次访问key为 {},描述为 [{}] 的接口", count, keys, limit.name());
return joinPoint.proceed();
} else {
throw new BadRequestException("访问次数受限制");
}
}
/**
* 限流脚本
*/
private String buildLuaScript() {
return "local c" +
"\nc = redis.call('get',KEYS[1])" +
"\nif c and tonumber(c) > tonumber(ARGV[1]) then" +
"\nreturn c;" +
"\nend" +
"\nc = redis.call('incr',KEYS[1])" +
"\nif tonumber(c) == 1 then" +
"\nredis.call('expire',KEYS[1],ARGV[2])" +
"\nend" +
"\nreturn c;";
}
}
Limit注解
package me.zhengjie.annotation;
import me.zhengjie.aspect.LimitType;
import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;
/**
* @author jacky
*/
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface Limit {
// 资源名称,用于描述接口功能
String name() default "";
// 资源 key
String key() default "";
// key prefix
String prefix() default "";
// 时间的,单位秒
int period();
// 限制访问次数
int count();
// 限制类型
LimitType limitType() default LimitType.CUSTOMER;
}
LimitType
/**
* 限流枚举
* @author /
*/
public enum LimitType {
// 默认
CUSTOMER,
// by ip addr
IP
}
使用
@Limit(key = “test”, period = 60, count = 10, name = “testLimit”, prefix = “limit”)