SpringCloudGatewayで同時に受け付けるリクエスト数を制限する

概要

• spring-cloud-gatewayには「RequestRateLimiterGatewayFilterFactory」が用意されていて、秒間のリクエスト数についてlimitがかけられる
• 一方、Javaアプリケーションでspring-webを利用している場合、秒間のリクエスト数というよりもスレッド数による制約がかかってくる
• →同時に受け付けるリクエスト数でlimitをかけたい！

作ってみた

ということでそんなGatewayFilterFactoryを作ってみました！

build.gradle

plugins {
    id 'org.springframework.boot' version '2.3.6.RELEASE'
    id 'io.spring.dependency-management' version '1.0.10.RELEASE'
    id 'java'
}

...

dependencies {
    implementation 'org.springframework.boot:spring-boot-starter-webflux'
    implementation 'org.springframework.cloud:spring-cloud-starter-gateway'
    implementation 'org.springframework.boot:spring-boot-starter-data-redis-reactive'
    // 組み込みのredisサーバ（任意）
    implementation 'it.ozimov:embedded-redis:0.7.2'

    compileOnly 'org.projectlombok:lombok'
    annotationProcessor 'org.projectlombok:lombok'
}

dependencyManagement {
    imports {
        mavenBom "org.springframework.cloud:spring-cloud-dependencies:Hoxton.SR9"
    }
}

CurrentSessionRateLimitGatewayFilterFactory

@Component
public class CurrentSessionRateLimitGatewayFilterFactory extends AbstractGatewayFilterFactory<CurrentSessionRateLimitGatewayFilterFactory.Config> {
    private final ReactiveRedisTemplate<String, String> redisTemplate;

    public CurrentSessionRateLimitGatewayFilterFactory(
            ReactiveRedisTemplate<String, String> redisTemplate
    ) {
        super(Config.class);
        this.redisTemplate = redisTemplate;
    }

    @Override
    public GatewayFilter apply(Config config) {
        return (exchange, chain) -> {
                var key = config.getKeyPrefix() + "." + UUID.randomUUID().toString();
                // セッションの一覧を取得
                return redisTemplate.keys(config.getKeyPrefix() + ".*")
                        // Flux → Mono変換
                        .collectList()
                        .flatMap(list -> {
                            // キーの個数がlimit未満であればバックポスト
                            if (list.size() < config.getLimit()) {
                                return redisTemplate.opsForValue().set(key, "", Duration.ofSeconds(config.getLifetimeSeconds()))
                                        .flatMap(bool -> chain.filter(exchange));
                            }
                            // それ以外は429を返却
                            setResponseStatus(exchange, HttpStatus.TOO_MANY_REQUESTS);
                            return exchange.getResponse().setComplete();
                        })
                        .then(Mono.defer(() -> redisTemplate.delete(key)))
                        .then();
        };
    }

    @Data
    public static class Config {
        // 同時に受け付けるリクエストの最大数
        private int limit;
        // limitをかける単位を指定（別のidでも同じkeyPrefixを指定すれば、まとめてlimitをかけられる）
        private String keyPrefix;
        private int lifetimeSeconds = 60;
    }
}

使い方はこんな感じ

application.yml

spring:
  redis:
    host: localhost
    port: 6379
  cloud:
    gateway:
      routes:
        - id: sample
          uri: "<バックポスト先>"
          predicates:
            - Path=/sample/**
          filters:
            - name: CurrentSessionRateLimit
              args:
                limit: 10
                keyPrefix: sample

受付中のリクエストの管理にredisを使用しているので、spring.redis周りの設定が必要です

補足

組み込みのredisサーバは以下のように起動させてます

RedisAutoConfiguration

@Component
public class RedisAutoConfiguration {
    private final RedisServer redisServer;

    public RedisAutoConfiguration(RedisProperties redisProperties) {
        redisServer = new RedisServer(redisProperties.getPort());
        // サーバーを起動
        redisServer.start();
    }

    @PreDestroy
    public void preDestroy() {
        // サーバーを停止
        redisServer.stop();
    }
}

普通はtestで使うのだろうけれども、redisを立てるのがめんどうでtest以外に使ってしまいました😇