redis+lua实现限流

柔光的暖阳◎ 2024-02-20 10:22 52阅读 0赞

## 1、需要引入Redis的maven坐标 ##

<dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-data-redis</artifactId>
        <version>2.3.0.RELEASE</version>
    </dependency>

## 2、redis配置 ##

spring:
      # Redis数据库索引
      redis:
        database: 0
      # Redis服务器地址
        host: 127.0.0.1
      # Redis服务器连接端口
        port: 6379
      # Redis服务器连接密码（默认为空）
        password:
      # 连接池最大连接数（使用负值表示没有限制）
        jedis:
          pool:
            max-active: 8
      # 连接池最大阻塞等待时间（使用负值表示没有限制）
            max-wait: -1
      # 连接池中的最大空闲连接
            max-idle: 8
      # 连接池中的最小空闲连接
            min-idle: 0
      # 连接超时时间（毫秒）
        timeout: 10000

## 3、新建脚本放在该项目的 resources 目录下，新建 limit.lua ##

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 0 else redis.call("INCRBY", key,"1") redis.call("expire", key,"2") return current + 1 end

## 4、自定义限流注解 ##

import java.lang.annotation.*;
    
    @Target(value = ElementType.METHOD)
    @Retention(RetentionPolicy.RUNTIME)
    @Documented
    public @interface RedisRateLimiter {
    
       //往令牌桶放入令牌的速率
        double value() default  Double.MAX_VALUE;
        //获取令牌的超时时间
        double limit() default  Double.MAX_VALUE;
    }

## 5、自定义切面类 RedisLimiterAspect 类 ，修改扫描自己controller类 ##

import com.imooc.annotation.RedisRateLimiter;
    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.annotation.Pointcut;
    import org.aspectj.lang.reflect.MethodSignature;
    import org.assertj.core.util.Lists;
    import org.json.JSONObject;
    import org.slf4j.Logger;
    import org.slf4j.LoggerFactory;
    import org.springframework.beans.factory.annotation.Autowired;
    import org.springframework.core.io.ClassPathResource;
    import org.springframework.data.redis.core.StringRedisTemplate;
    import org.springframework.data.redis.core.script.DefaultRedisScript;
    import org.springframework.scripting.support.ResourceScriptSource;
    import org.springframework.stereotype.Component;
    import javax.annotation.PostConstruct;
    import javax.servlet.http.HttpServletResponse;
    import java.io.PrintWriter;
    import java.util.List;
    
    @Aspect
    @Component
    public class RedisLimiterAspect {
        @Autowired
        private HttpServletResponse response;
    
        /**
         * 注入redis操作类
         */
        @Autowired
        private StringRedisTemplate stringRedisTemplate;
    
         private DefaultRedisScript<List> redisScript;
    
        /**
         * 初始化 redisScript 类
         * 返回值为 List
         */
        @PostConstruct
        public void init(){
            redisScript = new DefaultRedisScript<List>();
            redisScript.setResultType(List.class);
            redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("limit.lua")));
        }
    
        public final static Logger log = LoggerFactory.getLogger(RedisLimiterAspect.class);
    
        @Pointcut("execution( public * com.zz.controller.*.*(..))")
        public void pointcut(){
    
        }
        @Around("pointcut()")
        public Object process(ProceedingJoinPoint proceedingJoinPoint) throws  Throwable {
            MethodSignature  signature = (MethodSignature)proceedingJoinPoint.getSignature();
            //使用Java 反射技术获取方法上是否有@RedisRateLimiter 注解类
            RedisRateLimiter redisRateLimiter = signature.getMethod().getDeclaredAnnotation(RedisRateLimiter.class);
            if(redisRateLimiter == null){
                //正常执行方法,执行正常业务逻辑
                return proceedingJoinPoint.proceed();
            }
            //获取注解上的参数，获取配置的速率
            double value = redisRateLimiter.value();
            double time = redisRateLimiter.limit();
    
    
            //list设置lua的keys[1]
            //取当前时间戳到单位秒
            String key = "ip:"+ System.currentTimeMillis() / 1000;
    
            List<String> keyList = Lists.newArrayList(key);
    
            //用户Mpa设置Lua 的ARGV[1]
            //List<String> argList = Lists.newArrayList(String.valueOf(value));
    
            //调用脚本并执行
            List result = stringRedisTemplate.execute(redisScript, keyList, String.valueOf(value),String.valueOf(time));
    
            log.info("限流时间段内访问第：{} 次", result.toString());
    
            //lua 脚本返回 "0" 表示超出流量大小，返回1表示没有超出流量大小
            if(StringUtils.equals(result.get(0).toString(),"0")){
                //服务降级
                fullback();
                return null;
            }
    
            // 没有限流，直接放行
            return proceedingJoinPoint.proceed();
        }
    
        /**
         * 服务降级方法
         */
        private  void  fullback(){
            response.setCharacterEncoding("UTF-8");
            response.setContentType("application/json; charset=utf-8");
            PrintWriter writer = null;
            try {
                writer= response.getWriter();
                JSONObject o = new JSONObject();
                o.put("status",500);
                o.put("msg","Redis限流：请求太频繁，请稍后重试！");
                o.put("data",null);
                writer.printf(o.toString()
                );
    
            }catch (Exception e){
                e.printStackTrace();
            }finally {
                if(writer != null){
                    writer.close();
                }
            }
        }
    }

## 6、在需要限流的类添加注解 ##

import com.imooc.annotation.RedisRateLimiter;
    import io.swagger.annotations.Api;
    import io.swagger.annotations.ApiOperation;
    import org.springframework.web.bind.annotation.GetMapping;
    import org.springframework.web.bind.annotation.RequestMapping;
    import org.springframework.web.bind.annotation.RestController;
    
    import java.util.concurrent.TimeUnit;
    
    @RestController
    @Api(value = "限流", tags = {"限流测试接口"})
    @RequestMapping("limiter")
    public class LimiterController {
    
        @ApiOperation(value = "Redis限流注解测试接口",notes = "Redis限流注解测试接口", httpMethod = "GET")
        @RedisRateLimiter(value = 10, limit = 1)
        @GetMapping("/redislimit")
        public IMOOCJSONResult redislimit(){
    
            System.out.println("Redis限流注解测试接口");
            return IMOOCJSONResult.ok();
        }
    
    
    }