Redis+注解实现限流机制（IP、自定义等）

简介

在项目的使用过程中，限流的场景是很多的，尤其是要提供接口给外部使用的时候，但是自己去封装的话，相对比较耗时。

本方式可以使用默认（方法），ip、自定义参数进行限流，根据时间和次数进行。

整合步骤

依赖

  <dependency><groupId>org.springframework.boot</groupId><artifactId>spring-boot-starter-data-redis</artifactId></dependency><dependency><groupId>org.springframework.boot</groupId><artifactId>spring-boot-starter-web</artifactId></dependency><dependency><groupId>org.springframework.boot</groupId><artifactId>spring-boot-starter-aop</artifactId></dependency><dependency><groupId>org.projectlombok</groupId><artifactId>lombok</artifactId></dependency><dependency><groupId>cn.hutool</groupId><artifactId>hutool-all</artifactId><version>5.8.27</version></dependency><dependency><groupId>com.googlecode.aviator</groupId><artifactId>aviator</artifactId><version>5.4.1</version><scope>compile</scope></dependency>

限流注解

package com.walker.ratelimiter.annotation;import com.walker.ratelimiter.enums.LimitType;import java.lang.annotation.*;@Documented
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimiter {/*** 限流key*/String key() default "rate_limit";/*** 限流时间,单位秒*/int time() default 60;/*** 限流次数*/int count() default 50;/*** 限流类型*/LimitType limitType() default LimitType.DEFAULT;/*** 自定义编码* 支持SPEL表达式* 如果使用多参数，则使用:分割**/String customerCode() default "";/*** 自定义编码分割符*/String customerCodeSplit() default ":";
}

限流配置：获取限流lua脚本

package com.walker.ratelimiter.config;import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;@Configuration
public class RateLimitConfig {@Beanpublic DefaultRedisScript<Long> limitScript() {DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("lua/limit.lua")));redisScript.setResultType(Long.class);return redisScript;}}

基础变量

package com.walker.ratelimiter.constants;public interface BaseConstants {String COLON = ":";
}

枚举类型

package com.walker.ratelimiter.enums;public enum LimitType {/*** 默认策略*/DEFAULT,/*** 根据IP进行限流*/IP,/*** 自定义*/CUSTOME,}

lua脚本

local key = KEYS[1]
local count = tonumber(ARGV[1])
local time = tonumber(ARGV[2])
local current = redis.call('get', key)
if current and tonumber(current) > count thenreturn tonumber(current)
end
current = redis.call('incr', key)
if tonumber(current) == 1 thenredis.call('expire', key, time)
end
return tonumber(current)

切面类

package com.walker.ratelimiter.aspect;import cn.hutool.core.util.StrUtil;
import com.walker.ratelimiter.annotation.RateLimiter;
import com.walker.ratelimiter.constants.BaseConstants;
import com.walker.ratelimiter.enums.LimitType;
import lombok.extern.slf4j.Slf4j;
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.context.expression.MethodBasedEvaluationContext;
import org.springframework.core.DefaultParameterNameDiscoverer;
import org.springframework.core.ParameterNameDiscoverer;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.expression.EvaluationContext;
import org.springframework.expression.Expression;
import org.springframework.expression.TypedValue;
import org.springframework.expression.spel.standard.SpelExpressionParser;
import org.springframework.stereotype.Component;import java.lang.reflect.Method;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.Collections;
import java.util.List;@Slf4j
@Aspect
@Component
public class RateLimiterAspect {private final RedisTemplate redisTemplate;private final RedisScript<Long> limitScript;private SpelExpressionParser spelExpressionParser = new SpelExpressionParser();public RateLimiterAspect(RedisTemplate redisTemplate, RedisScript<Long> limitScript) {this.redisTemplate = redisTemplate;this.limitScript = limitScript;}@Around("@annotation(com.walker.ratelimiter.annotation.RateLimiter)")public Object doBefore(ProceedingJoinPoint joinPoint) throws Throwable {MethodSignature methodSignature = (MethodSignature) joinPoint.getSignature();RateLimiter rateLimiter = methodSignature.getMethod().getAnnotation(RateLimiter.class);//判断该方法是否存在限流的注解if (null != rateLimiter) {//获得注解中的配置信息int count = rateLimiter.count();int time = rateLimiter.time();//调用getCombineKey()获得存入redis中的key   key -> 注解中配置的key前缀-ip地址-方法路径-方法名String combineKey = getCombineKey(rateLimiter, methodSignature, joinPoint);log.info("combineKey->,{}", combineKey);//将combineKey放入集合List<Object> keys = Collections.singletonList(combineKey);log.info("keys->", keys);try {//执行lua脚本获得返回值Long number = (Long) redisTemplate.execute(limitScript, keys, count, time);//如果返回null或者返回次数大于配置次数，则限制访问if (number == null || number.intValue() > count) {throw new RuntimeException("访问过于频繁，请稍候再试");}log.info("限制请求'{}',当前请求'{}',缓存key'{}'", count, number.intValue(), combineKey);} catch (RuntimeException e) {throw e;} catch (Exception e) {throw new RuntimeException("服务器限流异常，请稍候再试");}}return joinPoint.proceed();}/*** Gets combine key.** @param rateLimiter the rate limiter* @param signature   the signature* @param joinPoint* @return the combine key*/public String getCombineKey(RateLimiter rateLimiter, MethodSignature signature, ProceedingJoinPoint joinPoint) throws UnknownHostException {StringBuilder stringBuffer = new StringBuilder(rateLimiter.key());
//        ip限流if (rateLimiter.limitType() == LimitType.IP) {InetAddress ip = InetAddress.getLocalHost();log.info("获取ip地址为：{}", ip);String hostAddress = ip.getHostAddress();stringBuffer.append(hostAddress).append(BaseConstants.COLON);
//        自定义编码限流} else if (rateLimiter.limitType() == LimitType.CUSTOME) {if (StrUtil.isEmpty(rateLimiter.customerCode())) {throw new RuntimeException("自定义编码不能为空");}String customerCode = rateLimiter.customerCode();String split = rateLimiter.customerCodeSplit();String[] customerCodes = customerCode.split(split);for (String code : customerCodes) {ParameterNameDiscoverer parameterNameDiscoverer = new DefaultParameterNameDiscoverer();EvaluationContext evaluationContext = new MethodBasedEvaluationContext(TypedValue.NULL, signature.getMethod(), joinPoint.getArgs(), parameterNameDiscoverer);Expression expression = spelExpressionParser.parseExpression(code);String resolvedCustomerCode =  String.valueOf(expression.getValue(evaluationContext));if(StrUtil.isEmpty(resolvedCustomerCode)){throw new RuntimeException("自定义编码不能为空");}stringBuffer.append(BaseConstants.COLON).append(resolvedCustomerCode);}}Method method = signature.getMethod();Class<?> targetClass = method.getDeclaringClass();stringBuffer.append(BaseConstants.COLON).append(targetClass.getName()).append(BaseConstants.COLON).append(method.getName());return stringBuffer.toString();}}

使用

根据ip进行限流

limitType = LimitType.IP

默认

limitType = LimitType.DEFAULT

自定义参数限流

使用Spel表达式，从参数中获取自定义的code，然后60s限流5次

@RateLimiter(limitType = LimitType.CUSTOME,customerCode = "#form.appCode:#toUserInfo.userUid",count = 5,time = 60)
public Result<Boolean> message(ImSendMsgForm form, MissuUsers toUserInfo) {}