Spring Cloud Gateway接口限流处理

添加依赖#

       <!-- 令牌桶限流 -->
        <dependency>
            <groupId>com.github.vladimir-bukhtoyarov</groupId>
            <artifactId>bucket4j-core</artifactId>
            <version>7.5.0</version>
        </dependency>

过滤器#

import com.alibaba.fastjson.JSON;
import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.Refill;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.annotation.Order;
import org.springframework.core.io.buffer.DataBuffer;
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.Flux;
import reactor.core.publisher.Mono;

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

/**
 * 限流处理
 */
@Component
@Order(-200)
public class RateLimitFilter implements GlobalFilter {

    /**
     * key=url
     * value=Map<String, Bucket>
     * key=ip
     * value=Bucket
     */
    private static final ConcurrentHashMap<String, ConcurrentHashMap<String, Bucket>> BUCKET_CACHE_MAP = new ConcurrentHashMap<>();

    private static final String HEADER_IP = "X-Forwarded-For";

    // 每秒钟允许的请求次数
    private static final int REQUEST_PER_SECOND = 10;

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        String ip = getIpAddress(request);
        //模拟不同ip请求
//        String ip = request.getHeaders().getFirst("ip");
        String url = request.getURI().getPath();
        Bucket bucket;
        ConcurrentHashMap<String, Bucket> bucketMap = BUCKET_CACHE_MAP.get(url);
        if (bucketMap == null) {
            ConcurrentHashMap<String, Bucket> valueMap = new ConcurrentHashMap<>();
            bucket = valueMap.computeIfAbsent(ip, k -> {
                // 每秒钟可以处理的请求数量
                Bandwidth limit = Bandwidth.classic(REQUEST_PER_SECOND, Refill.intervally(REQUEST_PER_SECOND, Duration.ofSeconds(1)));
                // 创建Bucket4j桶
                return Bucket.builder().addLimit(limit).build();
            });
            BUCKET_CACHE_MAP.put(url, valueMap);
        } else {
            bucket = bucketMap.computeIfAbsent(ip, k -> {
                // 1、每秒钟可以处理的请求数量
                // 2、桶每秒补充多少个
                Bandwidth limit = Bandwidth.classic(REQUEST_PER_SECOND, Refill.intervally(REQUEST_PER_SECOND, Duration.ofSeconds(1)));
                // 创建Bucket4j桶
                return Bucket.builder().addLimit(limit).build();
            });
        }
        if (bucket.tryConsume(1)) {
            //放行
            return chain.filter(exchange);
        }
        exchange.getResponse().setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
        Map<String, Object> resultMap = new HashMap<>();
        resultMap.put("code", 110);
        resultMap.put("msg", "请求过于频繁，请稍后重试");
        DataBuffer buffer = exchange.getResponse()
                .bufferFactory().wrap(JSON.toJSONString(resultMap).getBytes());
        return exchange.getResponse().writeWith(Flux.just(buffer));
    }

    private String getIpAddress(ServerHttpRequest request) {
        String ipAddress = request.getHeaders().getFirst(HEADER_IP);
        if (ipAddress == null || ipAddress.isEmpty() || "unknown".equalsIgnoreCase(ipAddress)) {
            InetSocketAddress remoteAddress = request.getRemoteAddress();
            if (remoteAddress != null) {
                ipAddress = remoteAddress.getAddress().getHostAddress();
            }
        }
        return ipAddress;
    }
}