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集成的原因）