SpringCloudGateway内存中基于令牌桶进行限流

Posted 2022-11-30 北亮bl

研究了一下网上的很多文章，基本都是使用 RedisRateLimiter，即根据Redis来进行限流操作。
这样有个好处，就是网关的集群可以使用同一套数据进行限流；
当然也有缺点，网关本来就是所有流量的集中出入口，如果每个请求都要往返一次Redis，无疑加重了网关的负担，性能有下降。

本文介绍了如何直接在内存中使用令牌桶算法进行限流，在内存中限流的缺点，当然就是对集群不友好了，比如有3个网关实例在运行，每个网关按每秒1个令牌，令牌桶容量为10，那么实际上就是每秒3个令牌，最大容量是10~30波动，不过基于令牌桶算法的控制，我觉得还是可以接受的。

---

本文基于 Java1.8、 spring-cloud.version:2020.0.2、springboot:2.4.4进行开发。 1、直接引入现成的令牌桶组件，就不要自己开发了： ```java com.github.vladimir-bukhtoyarov bucket4j-core 4.0.0 ``` 注：[令牌桶算法介绍](https://baike.baidu.com/item/%E4%BB%A4%E7%89%8C%E6%A1%B6%E7%AE%97%E6%B3%95/6597000)

2、新建一个限流类MyRateLimiter，并实现AbstractRateLimiter<MyRateLimiter.Config>。限流的关键代码就在这里了
注：MyRateLimiter.Config是限流类使用的配置数据类

import io.github.bucket4j.*;
import lombok.Data;
import org.springframework.cloud.gateway.filter.ratelimit.AbstractRateLimiter;
import org.springframework.cloud.gateway.support.ConfigurationService;
import reactor.core.publisher.Mono;

import java.time.Duration;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class MyRateLimiter extends AbstractRateLimiter<MyRateLimiter.Config> 
    public static final String CONFIGURATION_PROPERTY_NAME = "in-memory-rate-limiter";
    private final Config defaultConfig;
    private final Map<String, Bucket> ipBucketMap = new ConcurrentHashMap<>();

    /**
     * 构造函数
     * @param service 用于读取配置文件里的令牌桶配置
     */
    public MyRateLimiter(ConfigurationService service) 
        super(Config.class, CONFIGURATION_PROPERTY_NAME, service);
        defaultConfig = new Config();
        defaultConfig.setReplenishRate(10);
        defaultConfig.setBurstCapacity(100);
    

    @Override
    public Mono<Response> isAllowed(String routeId, String id) 
        Config routeConfig = getConfig().get(routeId);
        if (routeConfig == null) 
            if (defaultConfig == null) 
                throw new IllegalArgumentException("No Configuration found for route " + routeId);
            
            routeConfig = defaultConfig;
        

        // 每秒生成多少个令牌，就是当令牌桶为空时，每秒最多允许多少个用户进入，比如10
        int replenishRate = routeConfig.getReplenishRate();
        // 令牌桶的最大容量，就是突发涌入大量请求时，最多允许多少用户进入，比如100
        int burstCapacity = routeConfig.getBurstCapacity();
        // 初始化当前id的桶
        Bucket bucket = ipBucketMap.computeIfAbsent(id, k -> 
            Refill refill = Refill.of(replenishRate, Duration.ofSeconds(1));
            Bandwidth limit = Bandwidth.classic(burstCapacity, refill);
            return Bucket4j.builder().addLimit(limit).build();
        );
        Map<String, String> headers = new HashMap<>();
        // 尝试获取，同时得到剩余令牌数
        ConsumptionProbe probe = bucket.tryConsumeAndReturnRemaining(1);
        // 把剩余令牌数写入Header
        headers.put("remaining", String.valueOf(probe.getRemainingTokens()));
        if (probe.isConsumed()) 
            // 拿到令牌，允许进入
            return Mono.just(new Response(true, headers));
         else 
            // 没令牌了，返回429，不允许进入
            return Mono.just(new Response(false, headers));
        
    
    @Data
    public static class Config 
        private int replenishRate;
        private int burstCapacity;
    

3、定义限流使用的2个Bean：
KeyResolver：用于得到限流的唯一id，就是上面类里 isAllowed方法的id；
RateLimiter：用于得到限流类，就是上面的MyRateLimiter了。

import org.springframework.cloud.gateway.filter.ratelimit.KeyResolver;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
import org.springframework.cloud.gateway.support.ConfigurationService;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import reactor.core.publisher.Mono;

@Configuration
public class DemoAutoConfiguration 
    @Bean
    public KeyResolver myKeyResolver() 
            return exchange -> Mono.just(exchange.getRequest().getPath().value());
    
    @Bean
    public RateLimiter myRateLimiter(ConfigurationService service) 
        return new MyRateLimiter(service);
    

4、最后就是去修改配置文件，添加一些路由的限流配置了，参考：

spring:
  application:
    name: gateway-server
  cloud:
    gateway:
      routes:
        - id: mall-service            # 路由ID，必须唯一
          uri: lb://demo-mall         # proxy_pass转发到的目标地址
          predicates:
            - Path=/mall/**           # 匹配规则
          filters:
            - name: RequestRateLimiter
              args:
                key-resolver: "#@myKeyResolver"
                #rate-limiter: "#@myRateLimiter"
                in-memory-rate-limiter:
                  replenish-rate: 1   # 对应 MyRateLimiter.Config的属性，不能用 replenishRate
                  burst-capacity: 2   # 对应 MyRateLimiter.Config的属性，不能用 burstCapacity

5、好了，到这里，基于内存的限流开发和配置就完成了，可以跑起项目来看效果了。
完整的示例代码已经上传到github，你可以直接下载并启动它进行尝试：
去看代码