spring-gateway 自定义限流

某个服务需要对它的所有接口进行流量控制，已降低服务器资源的利用

pom: (jdk1.8)

java 复制代码

<spring-boot.version>2.7.13</spring-boot.version>
        <spring-cloud.version>2021.0.8</spring-cloud.version>
        <spring-cloud-alibaba.version>2021.0.5.0</spring-cloud-alibaba.version>

 
 <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-data-redis-reactive</artifactId>
        </dependency>

filter

java 复制代码

import com.sikaryofficial.common.core.utils.ServletUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.data.redis.core.ReactiveRedisTemplate;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.text.MessageFormat;
import java.time.Duration;

/**
 * @author : qinjinyuan
 * @desc : TODO  请填写你的功能描述
 * @date : 2024/08/14 15:26
 */
@Component
@Slf4j
public class RateLimitingFilterFactory extends AbstractGatewayFilterFactory<Object> {
    private static final int MAX_REQUESTS_PER_SECOND = 2;
    /**
     * 3秒内请求2次
     */
    private static final Duration EXPIRATION_TIME = Duration.ofSeconds(3);
    private static final String GATEWAY_RATE_LIMIT_KEY = "gateway:rate:limit:{0}";

    private final ReactiveRedisTemplate<Object, Object> reactiveRedisTemplate;

    public RateLimitingFilterFactory(ReactiveRedisTemplate<Object, Object> reactiveRedisTemplate) {
        super(Object.class);
        this.reactiveRedisTemplate = reactiveRedisTemplate;
    }

    @Override
    public GatewayFilter apply(Object config) {
        return (exchange, chain) -> {
            ServerHttpRequest request = exchange.getRequest();
            String path = request.getURI().getPath(); // 获取请求路径
            // 生成键
            String cacheKey = MessageFormat.format(GATEWAY_RATE_LIMIT_KEY, path);
            // 获取当前计数
            return reactiveRedisTemplate.opsForValue().increment(cacheKey, 1L)
                    .flatMap(count -> {
                        // 如果计数超过最大值，则拒绝请求
                        if (count > MAX_REQUESTS_PER_SECOND) {
                            return handleTooManyRequests(exchange);
                        } else {
                            // 设置过期时间
                            return reactiveRedisTemplate.expire(cacheKey, EXPIRATION_TIME)
                                    .then(chain.filter(exchange));
                        }
                    });
        };
    }

    private Mono<Void> handleTooManyRequests(ServerWebExchange exchange) {
        // 返回429 Too Many Requests HTTP状态码
        log.error("[Too Many Requests]请求路径:{}", exchange.getRequest().getPath());
          exchange.getResponse().setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
        exchange.getResponse().getHeaders().add(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON_VALUE);
        R<?> result = R.fail(HttpStatus.TOO_MANY_REQUESTS, "Too Many Requests");
        DataBuffer dataBuffer = exchange.getResponse().bufferFactory().wrap(JSON.toJSONString(result).getBytes());
        return exchange.getResponse().writeWith(Mono.just(dataBuffer));

    }


}

其中响应式 redis 配置如下

java 复制代码

@Configuration
@EnableCaching
@AutoConfigureBefore({RedisReactiveAutoConfiguration.class})
public class RedisConfig extends CachingConfigurerSupport
{
    

    /**
     * {@code @ConditionalOnMissingBean(name = "reactiveRedisTemplate")}
     * 表示组件可覆盖
     *
     * @see org.springframework.boot.autoconfigure.data.redis.RedisReactiveAutoConfiguration
     * #reactiveRedisTemplate(ReactiveRedisConnectionFactory, ResourceLoader)
     */
    @Bean
    @ConditionalOnMissingBean(name = "reactiveRedisTemplate")
    @SuppressWarnings(value = { "unchecked", "rawtypes" })
    public ReactiveRedisTemplate<Object, Object> reactiveRedisTemplate(
            ReactiveRedisConnectionFactory reactiveRedisConnectionFactory) {
        /**
         *   // RedisSerializer.string()
         *         // StringRedisSerializer.UTF_8
         *         // RedisSerializer.json()
         *         // GenericJackson2JsonRedisSerialize
         */
        FastJson2JsonRedisSerializer serializer = new FastJson2JsonRedisSerializer(Object.class);
        RedisSerializationContext<Object, Object> serializationContext = RedisSerializationContext
                .newSerializationContext(RedisSerializer.string())
                .value(serializer)
                .build();
        return new ReactiveRedisTemplate<>(reactiveRedisConnectionFactory, serializationContext);
    }

    
}

gateway 路由配置

python 复制代码

spring:
  redis:
    database: 1 
    host: localhost
    port: 6379
    password: 123456
    timeout: 6000ms  # 连接超时时长（毫秒）  
    lettuce:    
      pool:      
        max-active: 1000  # 连接池最大连接数（使用负值表示没有限制）      
        max-wait: -1ms      # 连接池最大阻塞等待时间（使用负值表示没有限制）      
        max-idle: 10      # 连接池中的最大空闲连接      
        min-idle: 5       # 连接池中的最小空闲连接
  cloud:
    gateway:
      discovery:
        locator:
          lowerCaseServiceId: true
          enabled: true
      routes:       
        # xxx 知识问答服务
        - id: xxx-qa-conversation
          uri: http://192.168.5.180:5000
          predicates:
            - Path=/qa/conversation/**
          filters:
            - StripPrefix=2
            # 服务限流控制过滤器
            - RateLimitingFilterFactory

这样就完成了一个完整的限流功能。

代码解释：

限流逻辑：

代码中的限流逻辑是基于每个路径的总请求数进行的，而不是在特定时间窗口内的请求数。这意味着如果一个路径在短时间内收到大量请求，只要总数不超过 MAX_REQUESTS_PER_SECOND，请求就会被允许通过

所以，如果需要根据当前用户来限流，在缓存key中，添加userID 来判定即可