springboot使用redis&lua脚本实现分布式限流starter

一，引入依赖

<dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-aop</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-data-redis</artifactId>
        </dependency>
        <dependency>
            <groupId>com.google.guava</groupId>
            <artifactId>guava</artifactId>
            <version>21.0</version>
        </dependency>
        <dependency>
            <groupId>org.apache.commons</groupId>
            <artifactId>commons-lang3</artifactId>
        </dependency>

二，在application配置文件中添加redis配置

spring:
  redis:
      host: *****
      password:****
      port: 6379
  # 连接超时时间（毫秒）
      timeout: 1000
  # Redis默认情况下有16个分片，这里配置具体使用的分片，默认是0
      database: 0

三，自定义RedisTemplate

由于后续要使用lua脚本来做权限控制，所以必须自定义一个redisTemplate，此处如果不自定义redisTemplate，则执行lua脚本时会报错。

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.redis.connection.lettuce.LettuceConnectionFactory;
import org.springframework.data.redis.core.RedisTemplate;
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(LettuceConnectionFactory redisConnectionFactory) {
        RedisTemplate<String, Serializable> template = new RedisTemplate<>();
        template.setKeySerializer(new StringRedisSerializer());
        template.setValueSerializer(new GenericJackson2JsonRedisSerializer());
        template.setConnectionFactory(redisConnectionFactory);
        return template;
    }
}

四，增加限定类型枚举类

自定义一个限定类型枚举类，后续根据类型判断，是根据ip、或是根据类型、或是根据方法名进行限流

public enum LimitType {
    //自定义key
    CUSTOMER,
    //默认，取方法名作为key
    OTHERS,
    //根据请求者IP作为key
    IP;
}

五，增加Limit注解

import java.lang.annotation.*;
import com.eoi.fly.redislimiter.constant.LimitType;
@Target({ElementType.METHOD,ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
@Inherited
@Documented
public @interface Limit {

    //资源key
    String key() default "";
    //时间
    int period();
    //最多访问次数
    int count();
    //类型
    LimitType limintType() default LimitType.OTHERS;
}

六，增加Limit注解AOP实现类

增加Limit注解的AOP切面，根据注解中的类型，使用lua脚本去redis获取访问次数

import com.eoi.fly.redislimiter.annotaion.Limit;
import com.eoi.fly.redislimiter.constant.LimitType;
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.slf4j.Logger;
import org.slf4j.LoggerFactory;
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.reflect.Method;
import java.util.Collections;

@Aspect
@Configuration
public class LimitInterceptor {
    private static final Logger logger = LoggerFactory.getLogger(LimitInterceptor.class);

    @Autowired
    private RedisTemplate<String, Serializable> limitRedisTemplate;


    @Around("execution(public * *(..)) && @annotation(com.eoi.fly.redislimiter.annotaion.Limit)")
    public Object interceptor(ProceedingJoinPoint joinPoint){
        //获取连接点的方法签名对象
        MethodSignature methodSignature = (MethodSignature) joinPoint.getSignature();
        //获取方法实例
        Method method = methodSignature.getMethod();
        //获取注解实例
        Limit limitAnnotation = method.getAnnotation(Limit.class);
        //注解中的类型
        LimitType limitType = limitAnnotation.limintType();
        //获取key名称
        String key;
        //获取限制时间范围
        int limitPeriod = limitAnnotation.period();
        //获取限制访问次数
        int limitCount = limitAnnotation.count();
        switch (limitType){
            //如果类型是IP，则根据IP限制访问次数，key取IP地址
            case IP:
                key = getIPAdress();
                break;
            //如果类型是customer，则根据key限制访问次数
            case CUSTOMER:
                key = limitAnnotation.key();
                break;
            //否则按照方法名称限制访问次数
            default:
                key = StringUtils.upperCase(method.getName());
        }
        try{
            String luaScript = buildLuaScript();
            RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
            /**
             * redisScript：构建的lua脚本对象
             *
             * List<K> keys是key的集合
             * Object... args是val的集合
             *
             * key的集合，在lua中可以使用KEYS[1]、KEYS[2]……获取，注意KEYS必须大写不能拼错
             * val的集合，在lua中可以使用ARGV[1]、ARGV[2]……获取，注意ARGV必须大写不能拼错
             * 说白了，使用redisTemplate操作lua，也就是传key的集合和val的集合，这一串lua脚本可以保证其原子性的
             *
             * lua中的redis.call命令就是操作redis的命令，第一个参数就是redis的原始命令，后面的参数就是redis命令的参数
             *
             * 这里构建的Lua脚本如下
             * ====================================================
             *  local c
             *  c = redis.call('get',KEYS[1])
             *  if c and tonumber(c) > tonumber(ARGV[1]) then
             *      return c;
             *  end
             *  c = redis.call('incr',KEYS[1])
             *  if tonumber(c) == 1 then
             *      redis.call('expire',KEYS[1],ARGV[2])
             *  end
             *  return c;
             *  ====================================================
             */
            Number count = limitRedisTemplate.execute(redisScript, Collections.singletonList(key), limitCount, limitPeriod);
            logger.info("Access try count is {} for key = {}", count, key);
            if(count !=null && count.intValue() <= limitCount){
                return joinPoint.proceed();
            }else{
                throw new RuntimeException("访问超限");
            }
        }catch(Throwable e){
            if(e instanceof RuntimeException){
                throw new RuntimeException(e.getLocalizedMessage());
            }
            throw new RuntimeException("服务异常");
        }
    }

    /**
     * lua限流脚本
     * @return
     */
    public String buildLuaScript(){
        StringBuilder sb = new StringBuilder();
        //定义c
        sb.append("local c");
        //获取redis中的值
        sb.append("\nc = redis.call('get',KEYS[1])");
        //如果调用不超过最大值
        sb.append("\nif c and tonumber(c) > tonumber(ARGV[1]) then");
        //直接返回
        sb.append("\n return c;");
        //结束
        sb.append("\nend");
        //访问次数加一
        sb.append("\nc = redis.call('incr',KEYS[1])");
        //如果是第一次调用
        sb.append("\nif tonumber(c) == 1 then");
        //设置对应值的过期设置
        sb.append("\nredis.call('expire',KEYS[1],ARGV[2])");
        //结束
        sb.append("\nend");
        //返回
        sb.append("\nreturn c;");

        return sb.toString();
    }

    private static final String UNKONW = "unknown";

    /**
     * 获取访问IP
     * @return
     */
    public String getIPAdress(){
        HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
        String ip = request.getHeader("x-forword-for");
        if(ip == null || ip.length() ==0 || UNKONW.equalsIgnoreCase(ip)){
            ip = request.getHeader("Proxy-Clent-IP");
        }
        if(ip == null || ip.length() ==0 || UNKONW.equalsIgnoreCase(ip)){
            ip = request.getHeader("WL-Clent-IP");
        }
        if(ip == null || ip.length() ==0 || UNKONW.equalsIgnoreCase(ip)){
            ip = request.getRemoteAddr();
        }
        return ip;
    }
}

七，效果预览

提供一个测试接口，加上限流注解
参数说明：在5秒内，此接口访问不能超过2次（没有指定key，默认方法名作为key，意味着此方法只能在单位时间内共被请求两次；如果指定key为IP，则每个IP在单位时间内请求两次）

@RequestMapping("/test")
    @Limit(period = 5,count = 2)
    public String test(){
        return "Test";
    }

前两次请求

超过两次请求后

过了五秒之后，时间刷新，再次请求后