聊聊token bucket算法的实现

序

本文主要研究一下token bucket算法的实现

限流算法概述

主要有如下几种：

• 基于信号量Semaphore

只有数量维度，没有时间维度

• 基于fixed window

带上了时间维度，不过在两个窗口的临界点容易出现超出限流的情况，比如限制每分钟10个请求，在00:59请求了10次，在01:01又请求了10次，而从00:30-01:30这个时间窗口来看，这一分钟请求了20次，没有控制好

• 基于rolling window

就是要解决fixed window没解决的窗口临界问题，主要有基于token bucket的算法，以及基于leaky bucket的算法

token bucket算法

• token按指定速率添加到bucket中
• 一个bucket有其容量限制，超过其容量则多余的token会被丢弃
• 当请求到来时，先试图获取token，如果剩余token足够则放行，不够则不允许放行(可能等待token足够再继续)

简单实现

/**
 * The minimalistic token-bucket implementation
 */
public class MinimalisticTokenBucket {

    private final long capacity;
    private final double refillTokensPerOneMillis;

    private double availableTokens;
    private long lastRefillTimestamp;

    /**
     * Creates token-bucket with specified capacity and refill rate equals to refillTokens/refillPeriodMillis
     */
    public MinimalisticTokenBucket(long capacity, long refillTokens, long refillPeriodMillis) {
        this.capacity = capacity;
        this.refillTokensPerOneMillis = (double) refillTokens / (double) refillPeriodMillis;

        this.availableTokens = capacity;
        this.lastRefillTimestamp = System.currentTimeMillis();
    }

    synchronized public boolean tryConsume(int numberTokens) {
        refill();
        if (availableTokens < numberTokens) {
            return false;
        } else {
            availableTokens -= numberTokens;
            return true;
        }
    }

    private void refill() {
        long currentTimeMillis = System.currentTimeMillis();
        if (currentTimeMillis > lastRefillTimestamp) {
            long millisSinceLastRefill = currentTimeMillis - lastRefillTimestamp;
            double refill = millisSinceLastRefill * refillTokensPerOneMillis;
            this.availableTokens = Math.min(capacity, availableTokens + refill);
            this.lastRefillTimestamp = currentTimeMillis;
        }
    }

    private static final class Selftest {

        public static void main(String[] args) {
            // 100 tokens per 1 second
            MinimalisticTokenBucket limiter = new MinimalisticTokenBucket(100, 100, 1000);

            long startMillis = System.currentTimeMillis();
            long consumed = 0;
            while (System.currentTimeMillis() - startMillis < 10000) {
                if (limiter.tryConsume(1)) {
                    consumed++;
                }
            }
            System.out.println(consumed);
        }

    }

}

• 以上是bucket4j给出的一个简单实现，用于理解token bucket算法
• 这个算法没有采用线程去refill token，因为bucket太多的话，线程太多，耗cpu
• 这个算法没有存储每个period使用的token，设计了lastRefillTimestamp字段，用于计算需要填充的token
• 每次tryConsume的时候，方法内部首先调用refill，根据设定的速度以及时间差计算这个时间段需要补充的token，更新availableTokens以及lastRefillTimestamp
• 之后限流判断，就是判断availableTokens与请求的numberTokens

小结

token bucket算法，是基于qps来限流，其简单的实现，就是计算单位时间补充token的速率，然后每次tryConsume的时候根据速率修正availableTokens。

doc

• Brief overview of token-bucket algorithm