基于redis+lua进行限流的方法

package com.sport.sportcloudmarathonh5.config;

import java.lang.annotation.*;

/**
* @author zdj
* @version 1.0.0
* @description 自定义注解实现分布式限流
*/
@Target(value = ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface RedisLimitStream {
/**
* 请求限制，一秒内可以允许好多个进入(默认一秒可以支持100个)
* @return
*/
int reqLimit() default 1000;

/**
* 模块名称
* @return
*/
String reqName() default “”;
}

/**
* 压测接口
* @return
*/
@Login(isLogin = false)
@RedisLimitStream(reqName = “名额秒杀”, reqLimit = 1000)
@ApiOperation(value = “压测接口”, notes = “压测接口”, httpMethod = “GET”)
@RequestMapping(value = “/pressure”, method = RequestMethod.GET)
public ResultVO<Object> pressure(){
return ResultVO.success(“抢购成功！”);
}

package com.sport.sportcloudmarathonh5.config;

import com.alibaba.fastjson.JSONObject;
import com.sport.sportcloudmarathonh5.service.impl.RedisService;
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.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;
import org.springframework.util.ObjectUtils;

import javax.servlet.http.HttpServletResponse;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.List;

/**
* @author zdj
* @version 1.0.0
* @description MyRedisLimiter注解的切面类
*/
@Aspect
@Component
public class RedisLimiterAspect {
private final Logger logger = LoggerFactory.getLogger(RedisLimitStream.class);
/**
* 当前响应请求
*/
@Autowired
private HttpServletResponse response;

/**
* redis服务
*/
@Autowired
private RedisService redisService;

/**
* 执行redis的脚本文件
*/
@Autowired
private RedisScript<Boolean> rateLimitLua;

/**
* 对所有接口进行拦截
*/
@Pointcut(“execution(public * com.sport.sportcloudmarathonh5.controller.*.*(..))”)
public void pointcut(){}

/**
* 对切点进行继续处理
*/
@Around(“pointcut()”)
public Object process(ProceedingJoinPoint proceedingJoinPoint) throws Throwable{
//使用反射获取RedisLimitStream注解
MethodSignature signature = (MethodSignature) proceedingJoinPoint.getSignature();
//没有添加限流注解的方法直接放行
RedisLimitStream redisLimitStream = signature.getMethod().getDeclaredAnnotation(RedisLimitStream.class);
if(ObjectUtils.isEmpty(redisLimitStream)){
return proceedingJoinPoint.proceed();
}

//List设置Lua的KEYS[1]
List<String> keyList = new ArrayList<>();
keyList.add(“ip:” + (System.currentTimeMillis() / 1000));

//获取注解上的参数，获取配置的速率
//List设置Lua的ARGV[1]
int value = redisLimitStream.reqLimit();

// 调用Redis执行lua脚本，未拿到令牌的，直接返回提示
boolean acquired = redisService.execute(rateLimitLua, keyList, value);
logger.info(“执行lua结果：” + acquired);
if(!acquired){
this.limitStreamBackMsg();
return null;
}

//获取到令牌，继续向下执行
return proceedingJoinPoint.proceed();
}

/**
* 被拦截的人，提示消息
*/
private void limitStreamBackMsg() {
response.setHeader(“Content-Type”, “text/html;charset=UTF8”);
PrintWriter writer = null;
try {
writer = response.getWriter();
writer.println(“{\”code\”:503,\”message\”:\”当前排队人较多，请稍后再试！\”,\”data\”:\”null\”}”);
writer.flush();
} catch (Exception e) {
e.printStackTrace();
} finally {
if (writer != null) {
writer.close();
}
}
}
}

package com.sport.sportcloudmarathonh5.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;

/**
* @author zdj
* @version 1.0.0
* @description 实现redis的编码方式
*/
@Configuration
public class RedisConfiguration {

/**
* 初始化将lua脚本加载到redis脚本中
* @return
*/
@Bean
public DefaultRedisScript loadRedisScript() {
DefaultRedisScript redisScript = new DefaultRedisScript();
redisScript.setLocation(new ClassPathResource(“limit.lua”));
redisScript.setResultType(Boolean.class);
return redisScript;
}
}

/**
* 执行lua脚本
* @param redisScript lua源代码脚本
* @param keyList
* @param value
* @return
*/
public boolean execute(RedisScript<Boolean> redisScript, List<String> keyList, int value) {
return redisTemplate.execute(redisScript, keyList, String.valueOf(value));
}

local key = KEYS[1] –限流KEY（一秒一个）
local limit = tonumber(ARGV[1]) –限流大小
local current = tonumber(redis.call(‘get’, key) or “0”)
if current + 1 > limit then –如果超出限流大小
return false
else –请求数+1，并设置2秒过期
redis.call(“INCRBY”, key, “1”)
redis.call(“expire”, key, “2”)
end
return true