手写一个简单的网关-限流过滤器

实现三种限流策略计数器限流，漏桶算法限流，令牌桶算法限流。网关中具体使用到的的限流策略同样可以通过配置文件进行配置。

计数器限流，该限流器实现简单，且经过压测性能也是最好的，tps达到3500左右。

package processors.limit;

import common.HttpStatueCode;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.http.FullHttpRequest;
import processors.Processor;

import java.util.Optional;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 计数器限流 通过后台定时线程每秒重置计数器
 * 放行速率较为平均 每秒放行一定量的请求
 */
public class CountLimit implements Processor {
    private Processor nextProcessor;
    private  final ScheduledExecutorService scheduler= Executors.newScheduledThreadPool(1);
    private  AtomicInteger currentRequests = new AtomicInteger(0);

    //TODO 通过配置文件读取
    private int maxRequests;
    public CountLimit(int maxRequests) {
        this.maxRequests = maxRequests;
        start();
    }

    public void start() {
        scheduler.scheduleAtFixedRate(() -> {
            currentRequests.set(0); // 重置计数器
        }, 1, 1, TimeUnit.SECONDS); // 每1秒重置一次
    }

    @Override
    public void process(ChannelHandlerContext ctx, FullHttpRequest request) {

        if (currentRequests.incrementAndGet() > maxRequests) {
            sendHttpResponse(ctx, request, HttpStatueCode.LIMIT_ERROR);
        }
        Optional.ofNullable(nextProcessor).ifPresent(p -> p.process(ctx, request));
    }

    @Override
    public void setNext(Processor processor) {
        this.nextProcessor = processor;
    }

    public CountLimit() {

    }
}

漏桶算法限流器

package processors.limit;

import common.HttpStatueCode;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.http.FullHttpRequest;
import processors.Processor;

import java.util.LinkedList;
import java.util.Optional;
import java.util.Queue;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 漏桶算法限流器
 * 算法思想 请求来到时判断漏桶是否有容量 有就存入漏桶中同时通过countDownlanch阻塞该线程继续向下执行 没有就拒绝请求
 * 后台线程每秒取出一定量的请求进行放行（通过countdownlangch放行执行请求的线程）
 *
 */
public class LeakyBucketLimit implements Processor {

    private Processor nextProcessor;
    private final int capacity; // 桶的容量
    private final int leakRate; // 漏水速率（每秒漏出的请求数）
    private final AtomicInteger currentWater = new AtomicInteger(0); // 当前桶中的水量（请求数）
    private final Queue<Task> requestQueue = new LinkedList<>(); // 存储请求的队列
    private  final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1);

    //TODO 通过配置文件读取


    public LeakyBucketLimit(int capacity, int leakRate) {
        this.capacity = capacity;
        this.leakRate = leakRate;
        start();
    }

    public void start() {
        scheduler.scheduleAtFixedRate(() -> {
            int allowedRequests = Math.min(leakRate, currentWater.get());
            currentWater.addAndGet(-allowedRequests);
            for (int i = 0; i < allowedRequests; i++) {
                Task task = requestQueue.poll();
                task.latch.countDown();
            }
        }, 1, 1, TimeUnit.SECONDS);
    }


    @Override
    public void process(ChannelHandlerContext ctx, FullHttpRequest request) {

        if (currentWater.incrementAndGet() > capacity) {
            sendHttpResponse(ctx, request, HttpStatueCode.LIMIT_ERROR);
            return;
        }
        CountDownLatch latch = new CountDownLatch(1);
        requestQueue.add(new Task(latch));
        try {
            latch.await();
        } catch (InterruptedException e) {
            System.out.println(e.getMessage());
        }
        Optional.ofNullable(nextProcessor).ifPresent(p -> p.process(ctx, request));
    }

    @Override
    public void setNext(Processor processor) {
        this.nextProcessor = processor;
    }

    /**
     * 用于控制线程阻塞
     */
    private class Task {
        CountDownLatch latch;

        public Task(CountDownLatch latch) {
            {
                this.latch = latch;
            }
        }
    }
}

令牌桶算法限流器

package processors.limit;

import common.HttpStatueCode;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.http.FullHttpRequest;
import processors.Processor;

import java.util.Optional;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 令牌桶算法限流器
 * 算法思想 后台线程每秒以一定速率补充令牌 获取到令牌的线程才能放行
 * 桶中可以存储一定量的令牌 允许一定的突发流量 但在令牌耗尽时也是以一定速率放行
 */
public class TokenBucketLimit implements Processor {
    private Processor nextProcessor;
    private final int capacity; // 桶的最大容量
    private final int refillRate; // 令牌生成速率（每秒生成的令牌数量）
    private AtomicInteger currentTokens = new AtomicInteger(0); // 当前桶中的令牌
    private  final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1);

    public TokenBucketLimit(int capacity, int refillRate) {
        this.capacity = capacity;
        this.refillRate = refillRate;
        start();
    }

    private void start() {
        scheduler.scheduleAtFixedRate(() -> {
            int filledTokens = Math.min(capacity, currentTokens.get()+refillRate);
            currentTokens.set(filledTokens);
        }, 1, 1, TimeUnit.SECONDS);
    }
    @Override
    public void process(ChannelHandlerContext ctx, FullHttpRequest request) {
        int current = currentTokens.get();
        if (currentTokens.get() < 0) {
            sendHttpResponse(ctx, request, HttpStatueCode.LIMIT_ERROR);
            return;
        }
        //尝试三次获取令牌
        for (int i = 0; i < 3; i++) {
            if (currentTokens.compareAndSet(current, current - 1)) {
                // 成功获取令牌，继续处理请求
                Optional.ofNullable(nextProcessor).ifPresent(p -> p.process(ctx, request));
                return;
            }
            try {
                Thread.sleep(200);
            } catch (InterruptedException e) {
                System.out.println(e.getMessage());
            }
        }
        sendHttpResponse(ctx, request, HttpStatueCode.LIMIT_ERROR);
    }

    @Override
    public void setNext(Processor processor) {
        this.nextProcessor = processor;
    }
}