SpringBoot--使用redis实现分布式限流

1、引入依赖

        <!-- 默认就内嵌了Tomcat 容器，如需要更换容器也极其简单-->
        <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>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-test</artifactId>
        </dependency>

2、在application配置文件中添加redis配置

spring:
  redis:
      host: *****
      password:****
      port: 6379
  # 连接超时时间（毫秒）
      timeout: 1000
  # Redis默认情况下有16个分片，这里配置具体使用的分片，默认是0
      database: 0
  # 连接池配置
      lettuce:
        pool:
  # 连接池最大连接数（使用负值表示没有限制） 默认 8
          max-active: 8
  # 连接池最大阻塞等待时间（使用负值表示没有限制） 默认 -1
          max-wait: -1
  # 连接池中的最大空闲连接 默认 8
          max-idle: 8
  # 连接池中的最小空闲连接 默认 0
          min-idle: 0

3、自定义redisTemplate

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

package com.example.demo.utils;

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;
    }
}

4、增加限定类型枚举类

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

package com.example.demo.entity;

public enum LimitType {
    //自定义key
    CUSTOMER,
    //根据请求者IP
    IP;
}

5、添加Limit注解

package com.example.demo.utils;

import com.example.demo.entity.LimitType;

import java.lang.annotation.*;
import java.util.concurrent.TimeUnit;

@Target({ElementType.METHOD,ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
@Inherited
@Documented
public @interface Limit {

    //资源名称
    String name() default "";
    //资源key
    String key() default "";
    //前缀
    String prefix() default "";
    //时间
    int period();//最多访问次数
    int count();
    //类型
    LimitType limintType() default LimitType.CUSTOMER;
}

6、增加Limit注解AOP实现类

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

package com.example.demo.utils;

import com.example.demo.entity.LimitType;
import com.google.common.collect.ImmutableList;
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.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;

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

    private final RedisTemplate<String, Serializable> limitRedisTemplate;

    public LimitInterceptor(RedisTemplate redisTemplate, RedisTemplate<String, Serializable> limitRedisTemplate) {
        this.limitRedisTemplate = limitRedisTemplate;
    }

    @Around("execution(public * *(..)) && @annotation(com.example.demo.utils.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 name = limitAnnotation.name();
        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());
        }
        ImmutableList<String> keys = ImmutableList.of(StringUtils.join(limitAnnotation.prefix(),key));
        try{
            String luaScript = buildLuaScript();
            RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
            Number count = limitRedisTemplate.execute(redisScript, keys, limitCount, limitPeriod);
            logger.info("Access try count is {} for name={} and key = {}", count, name, 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;
    }
}

6、增加访问控制类

　　在控制层添加Limit注解，返回访问次数。

@ResponseBody
    @GetMapping(value = "limit")
    @Limit(key = "test",period = 100, count = 5)
    public String testLimit(){
        return "第"+ATOMIC_INTEGER.incrementAndGet()+"次访问";
    }

7、测试

当访问超过次数后，抛出异常信息（此处无权限是由于添加了shiro集成的原因）