springcloud gateway限流实现

本文基于gateway自带的限流策略进行实现（通过接口、IP、用户进行限流）

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

/**
 * gateway自带的限流策略 缺点不能进行多级限流
 */
@Configuration
public class RedisRateLimiterConfig {

    /**
     * 接口限流
     * 获取请求地址的uri作为限流key
     *
     * @return
     */
    @Bean("pathResolver")
    public KeyResolver apiKeyResolver() {
        return exchange -> Mono.just(exchange.getRequest().getPath().value());
    }


    /**
     * ip限流
     *
     * @return
     */
    @Bean("ipResolver")
    @Primary
    public KeyResolver ipKeyResolver() {
        return exchange -> Mono.just(exchange.getRequest().getRemoteAddress().getHostName());
    }
    /**
     * 用户限流
     * 使用这种方式限流，请求路径中必须携带userId参数。
     * @return
     */
    @Bean
    public KeyResolver userKeyResolver() {
        return exchange -> Mono.just(exchange.getRequest().getQueryParams().getFirst("userId"));
    }
}

yml配置

server:
  port: 9003
spring:
  application:
    name: gateway
  cloud:
    gateway:
      globalcors:
        cors-configurations:
          '[/**]': # ??????
            allowCredentials: true
            allowedOriginsPatterns: "*" # 跨域问题
            allowedMethods: "*"
            allowedHeaders: "*"
      discovery:
        locator:
          enabled: true
      routes:
        - id: CLIENT01
          uri: lb://itsoft-task
          predicates:
            - Path=/systemTask/**
          filters:
            - StripPrefix=1
            - name: GatewayRequestRateLimiter
              args:
                redis-rate-limiter.replenishRate: 1
                redis-rate-limiter.burstCapacity: 1
                # 限流策略
                key-resolver: "#{@ipResolver}"

策略默认返回状态值429，这里自定义异常返回值

import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.RequestRateLimiterGatewayFilterFactory;
import org.springframework.cloud.gateway.filter.ratelimit.KeyResolver;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
import org.springframework.cloud.gateway.route.Route;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.util.Map;

/**
 * 默认限流返回状态码429   修改自定义限流返回值
 */
@Component
public class GatewayRequestRateLimiterGatewayFilterFactory extends RequestRateLimiterGatewayFilterFactory {

  private final RateLimiter defaultRateLimiter;

  private final KeyResolver defaultKeyResolver;

  public GatewayRequestRateLimiterGatewayFilterFactory(RateLimiter defaultRateLimiter, KeyResolver defaultKeyResolver) {
    super(defaultRateLimiter, defaultKeyResolver);
    this.defaultRateLimiter = defaultRateLimiter;
    this.defaultKeyResolver = defaultKeyResolver;
  }

  @Override
  public GatewayFilter apply(Config config) {
    KeyResolver resolver = getOrDefault(config.getKeyResolver(), defaultKeyResolver);
    RateLimiter<Object> limiter = getOrDefault(config.getRateLimiter(), defaultRateLimiter);
    return (exchange, chain) -> resolver.resolve(exchange).flatMap(key -> {
      String routeId = config.getRouteId();
      if (routeId == null) {
        Route route = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);
        routeId = route.getId();
      }
      String finalRouteId = routeId;
      return limiter.isAllowed(routeId, key).flatMap(response -> {
        for (Map.Entry<String, String> header : response.getHeaders().entrySet()) {
          exchange.getResponse().getHeaders().add(header.getKey(), header.getValue());
        }
        if (response.isAllowed()) {
          return chain.filter(exchange);
        }
        System.out.println("已限流: "+finalRouteId);
        ServerHttpResponse httpResponse = exchange.getResponse();
        //修改code为500
        httpResponse.setStatusCode(HttpStatus.INTERNAL_SERVER_ERROR);
        if (!httpResponse.getHeaders().containsKey("Content-Type")) {
          httpResponse.getHeaders().add("Content-Type", "application/json");
        }
        //此处无法触发全局异常处理，手动返回
        DataBuffer buffer = httpResponse.bufferFactory().wrap(("{\n"
            + "  \"code\": \"403\","
            + "  \"message\": \"服务器限流\","
            + "  \"data\": \"Server throttling\","
            + "  \"success\": false"
            + "}").getBytes(StandardCharsets.UTF_8));
        return httpResponse.writeWith(Mono.just(buffer));
      });
    });
  }

  private <T> T getOrDefault(T configValue, T defaultValue) {
    return (configValue != null) ? configValue : defaultValue;
  }
}