聊聊token bucket算法的實現

時間 2019-12-13

標籤聊聊 token bucket 算法實現简体版

原文原文鏈接

序

本文主要研究一下token bucket算法的實現git

限流算法概述

主要有以下幾種：github

基於信號量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