demo系列-限流算法-简单实现

demo系列-限流算法-简单实现

策略

1. 计数器
2. 漏桶算法
3. 令牌桶算法

使用(guava)

        RateLimiter limiter = RateLimiter.create(10);

        ThreadFactory factory = ThreadFactoryBuilder.create()
            .setNamePrefix("limit-")
            .build();
        ThreadPoolExecutor executor = new ThreadPoolExecutor(30, 30, 0L, TimeUnit.SECONDS, new LinkedBlockingQueue<>(10000));
        for (int i = 0; i < 30; ++i) {
            executor.execute(() -> {
                limiter.acquire();
                System.out.println(Thread.currentThread().getName() + " 请求成功");
            });
        }
        Thread.sleep(2000);
        executor.shutdown();

实现：ReentrantLock

Limiter

@NoArgsConstructor
@AllArgsConstructor
public abstract class Limiter {
    /**
     * 每秒请求次数
     */
    protected int rate = 1;

    /**
     * 非阻塞
     *
     * @return 成功
     */
    public abstract boolean tryAcquire();

    /**
     * 阻塞
     */
    public void acquire() {
        do {
            if (tryAcquire()) {
                return;
            }

            try {
                TimeUnit.MILLISECONDS.sleep(100);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        } while (true);
    }
}

计数器

public class CountLimiter extends Limiter {
    private volatile int counter = 0;

    private volatile long lastTime = 0L;

    private Lock lock = new ReentrantLock();

    public CountLimiter(int rate) {
        super(rate);
    }

    @Override
    public boolean tryAcquire() {
        try {
            lock.lock();
            if (System.currentTimeMillis() - lastTime > 1000L) {
                counter = 0;
                lastTime = System.currentTimeMillis();
            }

            counter++;
            if (counter <= rate) {
                return true;
            }
        } finally {
            lock.unlock();
        }
        return false;
    }
}

LeakyBucketLimiter

public class LeakyBucketLimiter extends Limiter {
    /*
    最大容量
    当前剩余：匀速出去
    上次时间
     */
    private int capacity;

    private volatile int water = 0;

    private volatile long lastTime = 0L;

    private Lock lock = new ReentrantLock();

    public LeakyBucketLimiter(int rate) {
        super(rate);
        this.capacity = rate;
    }

    @Override
    public boolean tryAcquire() {
        try {
            lock.lock();
            long now = System.currentTimeMillis();
            // 匀速漏出
            int outWater = Math.round((now - lastTime) / 1000L * rate);
            if (outWater > 0) {
                lastTime = now;
            }

            water = Math.max(0, water - outWater);
            if (water < capacity) {
                water++;
                return true;
            }
        } finally {
            lock.unlock();
        }

        return false;
    }
}

TokenBucketLimiter

public class TokenBucketLimiter extends Limiter {
    /*
    最大容量
    当前剩余：匀速进来
    上次时间
     */
    private int capacity;

    private volatile int token = 0;

    private volatile long lastTime = 0L;

    private Lock lock = new ReentrantLock();

    public TokenBucketLimiter(int rate) {
        super(rate);
        this.capacity = rate;
    }

    @Override
    public boolean tryAcquire() {
        try {
            lock.lock();
            long now = System.currentTimeMillis();
            // 匀速增加令牌
            int inToken = Math.round((now - lastTime) / 1000L * rate);
            if (inToken > 0) {
                lastTime = now;
            }

            token = Math.min(capacity, token + inToken);
            if (token > 0) {
                token--;
                return true;
            }
        } finally {
            lock.unlock();
        }

        return false;
    }
}