1998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi/* 2998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * Copyright (C) 2016 The Android Open Source Project 3998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * 4998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * Licensed under the Apache License, Version 2.0 (the "License"); 5998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * you may not use this file except in compliance with the License. 6998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * You may obtain a copy of the License at 7998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * 8998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * http://www.apache.org/licenses/LICENSE-2.0 9998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * 10998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * Unless required by applicable law or agreed to in writing, software 11998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * distributed under the License is distributed on an "AS IS" BASIS, 12998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * See the License for the specific language governing permissions and 14998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * limitations under the License. 15998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi */ 16998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi 17998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichipackage com.android.internal.util; 18998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi 19998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichiimport android.os.SystemClock; 20998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi 21998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichiimport static com.android.internal.util.Preconditions.checkArgumentNonnegative; 22998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichiimport static com.android.internal.util.Preconditions.checkArgumentPositive; 23998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi 24998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi/** 25998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * A class useful for rate-limiting or throttling that stores and distributes tokens. 26998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * 27998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * A TokenBucket starts with a fixed capacity of tokens, an initial amount of tokens, and 28998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * a fixed filling period (in milliseconds). 29998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * 30998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * For every filling period, the bucket gains one token, up to its maximum capacity from 31998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * which point tokens simply overflow and are lost. Tokens can be obtained one by one or n by n. 32998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * 33998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * The available amount of tokens is computed lazily when the bucket state is inspected. 34998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * Therefore it is purely synchronous and does not involve any asynchronous activity. 35998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * It is not synchronized in any way and not a thread-safe object. 360d4a398b7846f26e4452180915ecb5a9d4566148Hugo Benichi * 370d4a398b7846f26e4452180915ecb5a9d4566148Hugo Benichi * {@hide} 38998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi */ 39998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichipublic class TokenBucket { 40998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi 41998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi private final int mFillDelta; // Time in ms it takes to generate one token. 42998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi private final int mCapacity; // Maximum number of tokens that can be stored. 43998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi private long mLastFill; // Last time in ms the bucket generated tokens. 44998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi private int mAvailable; // Current number of available tokens. 45998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi 46998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi /** 47998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * Create a new TokenBucket. 48998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * @param deltaMs the time in milliseconds it takes to generate a new token. 49998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * Must be strictly positive. 50998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * @param capacity the maximum token capacity. Must be strictly positive. 51998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * @param tokens the starting amount of token. Must be positive or zero. 52998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi */ 53998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi public TokenBucket(int deltaMs, int capacity, int tokens) { 54998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi mFillDelta = checkArgumentPositive(deltaMs, "deltaMs must be strictly positive"); 55998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi mCapacity = checkArgumentPositive(capacity, "capacity must be strictly positive"); 56998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi mAvailable = Math.min(checkArgumentNonnegative(tokens), mCapacity); 57998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi mLastFill = scaledTime(); 58998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi } 59998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi 60998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi /** 61998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * Create a new TokenBucket that starts completely filled. 62998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * @param deltaMs the time in milliseconds it takes to generate a new token. 63998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * Must be strictly positive. 64998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * @param capacity the maximum token capacity. Must be strictly positive. 65998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi */ 66998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi public TokenBucket(int deltaMs, int capacity) { 67998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi this(deltaMs, capacity, capacity); 68998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi } 69998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi 70998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi /** Reset this TokenBucket and set its number of available tokens. */ 71998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi public void reset(int tokens) { 72998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi checkArgumentNonnegative(tokens); 73998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi mAvailable = Math.min(tokens, mCapacity); 74998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi mLastFill = scaledTime(); 75998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi } 76998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi 77998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi /** Returns this TokenBucket maximum token capacity. */ 78998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi public int capacity() { 79998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi return mCapacity; 80998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi } 81998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi 82998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi /** Returns this TokenBucket currently number of available tokens. */ 83998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi public int available() { 84998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi fill(); 85998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi return mAvailable; 86998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi } 87998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi 88998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi /** Returns true if this TokenBucket as one or more tokens available. */ 89998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi public boolean has() { 90998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi fill(); 91998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi return mAvailable > 0; 92998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi } 93998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi 94998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi /** Consumes a token from this TokenBucket and returns true if a token is available. */ 95998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi public boolean get() { 96998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi return (get(1) == 1); 97998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi } 98998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi 99998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi /** 100998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * Try to consume many tokens from this TokenBucket. 101998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * @param n the number of tokens to consume. 102998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi * @return the number of tokens that were actually consumed. 103998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi */ 104998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi public int get(int n) { 105998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi fill(); 106998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi if (n <= 0) { 107998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi return 0; 108998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi } 109998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi if (n > mAvailable) { 110998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi int got = mAvailable; 111998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi mAvailable = 0; 112998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi return got; 113998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi } 114998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi mAvailable -= n; 115998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi return n; 116998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi } 117998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi 118998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi private void fill() { 119998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi final long now = scaledTime(); 120998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi final int diff = (int) (now - mLastFill); 121998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi mAvailable = Math.min(mCapacity, mAvailable + diff); 122998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi mLastFill = now; 123998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi } 124998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi 125998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi private long scaledTime() { 126998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi return SystemClock.elapsedRealtime() / mFillDelta; 127998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi } 128998493f0ee39ae0e9ffdea27f48f1b11b0807fcbHugo Benichi} 129