简单单机版限流器

在高并发系统设计的时候，限流器是一个很重要的工具，现在的分布式框架大多集成了限流器
如何手撸一个简单限流器呢

限流种类

• 限制并发数：即限制同一时刻的请求数，例如限制1w的并发，就是说一个请求来了我们检查一下当前正在运行的请求，如果数量不超过1w我们就放行
• 限制qps：限制时间段内的请求数，例如限制qps为1w，就是说一秒钟最多有1w个请求放行

限制并发数

对于第一种我们需要一个信号量，请求开始的时候判断信号量，如果不大于0就直接拒绝请求，否则信号量减1并放行，然后请求结束的时候信号量加1，这种一般用在环绕切面上
demo如下：

public class ConcurrentLimiter {
    private final int STOCK;
    private final AtomicInteger TOKEN;
    public ConcurrentLimiter(int stock) {
        STOCK = stock;
        TOKEN = new AtomicInteger(stock);
    }
    public boolean sub() {
        int v;
        do {
            if ((v = TOKEN.get()) <= 0L) {
                return false;
            }
        } while (!TOKEN.compareAndSet(v, v - 1));
        return true;
    }
    public void add() {
        int v = TOKEN.incrementAndGet();
        if (v > STOCK) {
            throw new RuntimeException("concurrent used error");
        }
    }
}

测试代码如下：

final ConcurrentLimiter limiter = new ConcurrentLimiter(5);
final SimpleDateFormat sdf = new SimpleDateFormat("HH:mm:ss");
Runnable r = new Runnable() {
    @Override
    public void run() {
        System.out.println("thread[" + Thread.currentThread().getName() + "] start[" + sdf.format(new Date()) + "]");
        if (limiter.sub()) {
            try {
                Thread.sleep(5000L);
            } catch (InterruptedException e) {
                e.printStackTrace();
            } finally {
                limiter.add();
            }
        }
        System.out.println("thread[" + Thread.currentThread().getName() + "] end[" + sdf.format(new Date()) + "]");
    }
};
Thread[] ts = new Thread[10];
for (int i = 0; i < ts.length; i++) {
    ts[i] = new Thread(r, String.valueOf(i));
    ts[i].start();
}
for (int i = 0; i < ts.length; i++) {
    ts[i].join();
}

输出为：

thread[3] start[11:38:10]
thread[0] start[11:38:10]
thread[4] start[11:38:10]
thread[9] start[11:38:10]
thread[1] start[11:38:10]
thread[2] start[11:38:10]
thread[7] start[11:38:10]
thread[2] end[11:38:10]
thread[6] start[11:38:10]
thread[5] start[11:38:10]
thread[6] end[11:38:10]
thread[7] end[11:38:10]
thread[8] start[11:38:10]
thread[5] end[11:38:10]
thread[8] end[11:38:10]
thread[4] end[11:38:15]
thread[3] end[11:38:15]
thread[9] end[11:38:15]
thread[1] end[11:38:15]
thread[0] end[11:38:15]

限制qps

对于第二种，我们需要把时间划分成时间段，然后每个时间段设置一个限量值，请求的时候限量值减1，如果时间段内值小于0了就拒绝掉，因此需要：时间段，限流值

public class PeriodLimiter {
    private static final int EXPIRE = -1;
    private static final int OVERFLOW = 0;
    private static final int NORMAL = 1;

    private final long PERIOD;
    private final int STOCK;
    volatile Slot slot;
    public PeriodLimiter(long period, int stock) {
        PERIOD = period;
        STOCK = stock;
    }
    public boolean get() {
        int state;
        //此处还可以用unsafe乐观锁来实现
        if (slot == null || (state = slot.get()) == EXPIRE) {
            synchronized (this) {
                if (slot == null || (state = slot.get()) == EXPIRE) {
                    slot = new Slot(System.currentTimeMillis());
                    state = slot.get();
                }
            }
        }
        if (state == OVERFLOW) {
            return false;
        }
        return true;
    }

    private class Slot {
        private final AtomicInteger TOKEN;
        private final long START;
        Slot(long start) {
            TOKEN = new AtomicInteger(STOCK);
            START = start;
        }
        int get() {
            long time = System.currentTimeMillis();
            if (time > START + PERIOD) {
                return EXPIRE;
            }
            int v;
            do {
                if ((v = TOKEN.get()) <= 0L) {
                    return OVERFLOW;
                }
            } while (!TOKEN.compareAndSet(v, v - 1));
            return NORMAL;
        }
    }
}

测试代码如下

//限流每秒500w
final PeriodLimiter limiter = new PeriodLimiter(1000L, 5000000);
final AtomicInteger tCnt = new AtomicInteger();
final AtomicInteger fCnt = new AtomicInteger();
long start = System.currentTimeMillis();
Runnable r = new Runnable() {
    @Override
    public void run() {
        int cnt = 12000000;
        while (cnt-- > 0) {
            if (limiter.get()) {
                tCnt.incrementAndGet();
            } else {
                fCnt.incrementAndGet();
            }
        }
    }
};
Thread[] ts = new Thread[10];
for (int i = 0; i < ts.length; i++) {
    ts[i] = new Thread(r);
    ts[i].start();
}
for (int i = 0; i < ts.length; i++) {
    ts[i].join();
}
System.out.println("true:[" + tCnt.get() + "]false:[" + fCnt.get() + "]" + (System.currentTimeMillis() - start) + "ms");

运行结果如下：

true:[30000000]false:[90000000]5689ms

适用场景

rpc服务一般都是利用固定大小的线程池来运行服务的，当某个服务出现运行时间过长的时候，就会出现所有的服务线程都被其占用的情况(例如数据库连接池被慢SQL占满，dubbo服务线程被慢查询占满)，导致一些更重要的请求被拒绝服务，这个时候，我们可以限制单服务的并发数来做限制。当某个服务被多个第三方请求时，我们可以分别对不同来源的请求做qps的限流额度配置