SpringBoot中如何編寫一個優雅的限流組件？

很早以前，我曾寫過兩篇介紹如何在SpringBoot中使用Guava和Redis實現接口限流的文章。具體包括：

1. 使用Guava實現單機令牌桶限流
2. 使用Redis實現分佈式限流

現在，一個問題擺在我們面前：如何將這兩種限流機制整合到同一個組件中，以便用戶隨時切換呢？

顯然，我們需要定義一個通用的限流組件，將其引入到業務中，並支持通過配置文件自由切換不同的限流機制。舉例而言，當使用limit.type=redis時，啓用Redis分佈式限流組件，當使用limit.type=local時，啓用Guava限流組件。這種自由切換機制能夠爲用戶提供更大的靈活性和可維護性。

接下來，讓我們開始動手實現吧！

第一步，創建通用模塊cloud-limiter-starter

首先在父項目下創建一個模塊

然後在pom文件中引入相關依賴

<dependencies>
  <dependency>
    <groupId>com.google.guava</groupId>
    <artifactId>guava</artifactId>
  </dependency>
  <!--SpringFramework-->
  <dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-web</artifactId>
    <scope>provided</scope>
  </dependency>
  <dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-data-redis</artifactId>
    <scope>provided</scope>
  </dependency>
  <dependency>
    <groupId>org.projectlombok</groupId>
    <artifactId>lombok</artifactId>
    <scope>provided</scope>
  </dependency>
  <dependency>
    <groupId>org.springframework</groupId>
    <artifactId>spring-aspects</artifactId>
    <scope>provided</scope>
  </dependency>

</dependencies>

小提示：通用模塊命名最好遵照規則以starter命名結束，同時通用模塊引入的依賴最好設置<scope>provided</scope>屬性。

第二步，實現限流功能

1. 創建限流接口

既然有兩種限流機制，按照套路肯定得先創建一個限流接口，就叫LimiterManager吧。

public interface LimiterManager {
    boolean tryAccess(Limiter limiter);
}

2. 分別實現Redis的限流功能和Guava的限流功能，這裏只給出核心代碼。

Guava限流的核心實現GuavaLimiter

@Slf4j
public class GuavaLimiter implements LimiterManager{
    private final Map<String, RateLimiter> limiterMap = Maps.newConcurrentMap();

    @Override
    public boolean tryAccess(Limiter limiter) {
        RateLimiter rateLimiter = getRateLimiter(limiter);
        if (rateLimiter == null) {
            return false;
        }

        boolean access = rateLimiter.tryAcquire(1,100, TimeUnit.MILLISECONDS);

        log.info("{} access :{}",limiter.getKey() , access);

        return access;
    }
}    

Redis限流的核心實現RedisLimiter

@Slf4j
public class RedisLimiter implements LimiterManager{

    private final StringRedisTemplate stringRedisTemplate;

    public RedisLimiter(StringRedisTemplate stringRedisTemplate) {
        this.stringRedisTemplate = stringRedisTemplate;
    }

    @Override
    public boolean tryAccess(Limiter limiter) {

        String key = limiter.getKey();
        if (StringUtils.isEmpty(key)) {
            throw new LimiterException( "redis limiter key cannot be null" );
        }

        List<String> keys = new ArrayList<>();
        keys.add( key );

        int seconds = limiter.getSeconds();
        int limitCount = limiter.getLimitNum();

        String luaScript = buildLuaScript();

        RedisScript<Long> redisScript = new DefaultRedisScript<>(luaScript, Long.class);

        Long count = stringRedisTemplate.execute( redisScript, keys, "" + limitCount, "" + seconds );

        log.info( "Access try count is {} for key={}", count, key );

        return count != null && count != 0;
    }
}    

第三步，創建配置類

編寫配置類根據配置文件注入限流實現類，當配置文件中屬性 limit.type=local 時啓用Guava限流機制，當limit.type=redis 時啓用Redis限流機制。

@Configuration
public class LimiterConfigure {

    @Bean
    @ConditionalOnProperty(name = "limit.type",havingValue = "local")
    public LimiterManager guavaLimiter(){
        return new GuavaLimiter();
    }


    @Bean
    @ConditionalOnProperty(name = "limit.type",havingValue = "redis")
    public LimiterManager redisLimiter(StringRedisTemplate stringRedisTemplate){
        return new RedisLimiter(stringRedisTemplate);
    }
}

第四步，創建AOP

根據前面的兩篇文章可知，避免限流功能污染業務邏輯的最好方式是藉助Spring AOP，所以很顯然還得需要創建一個AOP。

@Aspect
@EnableAspectJAutoProxy(proxyTargetClass = true) //使用CGLIB代理
@Conditional(LimitAspectCondition.class)
public class LimitAspect {

    @Setter(onMethod_ = @Autowired)
    private LimiterManager limiterManager;

    @Pointcut("@annotation(com.jianzh5.limit.aop.Limit)")
    private void check() {

    }

    @Before("check()")
    public void before(JoinPoint joinPoint){
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();

        Limit limit = method.getAnnotation(Limit.class);
        if(limit != null){

            Limiter limiter = Limiter.builder().limitNum(limit.limitNum())
                    .seconds(limit.seconds())
                    .key(limit.key()).build();

            if(!limiterManager.tryAccess(limiter)){
                throw new LimiterException( "There are currently many people , please try again later!" );
            }
        }
    }
}

注意到類上我加了一行@Conditional(LimitAspectCondition.class)，使用了自定義條件選擇器，意思是隻有當配置類中出現了limit.type屬性時纔會加載這個AOP。

public class LimitAspectCondition implements Condition {
    @Override
    public boolean matches(ConditionContext conditionContext, AnnotatedTypeMetadata annotatedTypeMetadata) {
        //檢查配置文件是否包含limit.type屬性
        return conditionContext.getEnvironment().containsProperty(ConfigConstant.LIMIT_TYPE);
    }
}

第四步，創建spring.factories文件，引導SpringBoot加載配置類

## AutoConfiguration
org.springframework.boot.autoconfigure.EnableAutoConfiguration=\
  com.jianzh5.limit.config.LimiterConfigure,\
  com.jianzh5.limit.aop.LimitAspect

完整目錄結構如下：

第五步，在項目中引用限流組件

1. 引入依賴

<dependency>
    <groupId>com.jianzh5</groupId>
    <artifactId>cloud-limit-starter</artifactId>
</dependency>

2. 在application.properties中設置加載的限流組件

limit.type = redis

如果不配置此屬性則不加載對應限流功能。

3. 在需要限流的接口上加上註解

@Limit(key = "Limiter:test",limitNum = 3,seconds = 1)

小結

通過上述步驟，我們已經成功實現了一個通用限流組件。在實際應用中，只需要根據場景需求選擇對應的限流機制，即可非常方便的進行限流操作。這種靈活性和便捷性，也是SpringBoot中定義Starter的一般套路。

如果你想詳細瞭解這兩種限流機制的原理，可以參考之前的文章中所介紹的內容。

老鳥系列源碼已經上傳至GitHub，需要的在公號【JAVA日知錄】回覆關鍵字 0923 獲取源碼地址。