1 /*
2  * Copyright (C) 2016 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.internal.util;
18 
19 import android.os.SystemClock;
20 
21 import static com.android.internal.util.Preconditions.checkArgumentNonnegative;
22 import static com.android.internal.util.Preconditions.checkArgumentPositive;
23 
24 /**
25  * A class useful for rate-limiting or throttling that stores and distributes tokens.
26  *
27  * A TokenBucket starts with a fixed capacity of tokens, an initial amount of tokens, and
28  * a fixed filling period (in milliseconds).
29  *
30  * For every filling period, the bucket gains one token, up to its maximum capacity from
31  * which point tokens simply overflow and are lost. Tokens can be obtained one by one or n by n.
32  *
33  * The available amount of tokens is computed lazily when the bucket state is inspected.
34  * Therefore it is purely synchronous and does not involve any asynchronous activity.
35  * It is not synchronized in any way and not a thread-safe object.
36  *
37  * {@hide}
38  */
39 public class TokenBucket {
40 
41     private final int mFillDelta; // Time in ms it takes to generate one token.
42     private final int mCapacity;  // Maximum number of tokens that can be stored.
43     private long mLastFill;       // Last time in ms the bucket generated tokens.
44     private int mAvailable;       // Current number of available tokens.
45 
46     /**
47      * Create a new TokenBucket.
48      * @param deltaMs the time in milliseconds it takes to generate a new token.
49      * Must be strictly positive.
50      * @param capacity the maximum token capacity. Must be strictly positive.
51      * @param tokens the starting amount of token. Must be positive or zero.
52      */
TokenBucket(int deltaMs, int capacity, int tokens)53     public TokenBucket(int deltaMs, int capacity, int tokens) {
54         mFillDelta = checkArgumentPositive(deltaMs, "deltaMs must be strictly positive");
55         mCapacity = checkArgumentPositive(capacity, "capacity must be strictly positive");
56         mAvailable = Math.min(checkArgumentNonnegative(tokens), mCapacity);
57         mLastFill = scaledTime();
58     }
59 
60     /**
61      * Create a new TokenBucket that starts completely filled.
62      * @param deltaMs the time in milliseconds it takes to generate a new token.
63      * Must be strictly positive.
64      * @param capacity the maximum token capacity. Must be strictly positive.
65      */
TokenBucket(int deltaMs, int capacity)66     public TokenBucket(int deltaMs, int capacity) {
67         this(deltaMs, capacity, capacity);
68     }
69 
70     /** Reset this TokenBucket and set its number of available tokens. */
reset(int tokens)71     public void reset(int tokens) {
72         checkArgumentNonnegative(tokens);
73         mAvailable = Math.min(tokens, mCapacity);
74         mLastFill = scaledTime();
75     }
76 
77     /** Returns this TokenBucket maximum token capacity. */
capacity()78     public int capacity() {
79         return mCapacity;
80     }
81 
82     /** Returns this TokenBucket currently number of available tokens. */
available()83     public int available() {
84         fill();
85         return mAvailable;
86     }
87 
88     /** Returns true if this TokenBucket as one or more tokens available. */
has()89     public boolean has() {
90         fill();
91         return mAvailable > 0;
92     }
93 
94     /** Consumes a token from this TokenBucket and returns true if a token is available. */
get()95     public boolean get() {
96         return (get(1) == 1);
97     }
98 
99     /**
100      * Try to consume many tokens from this TokenBucket.
101      * @param n the number of tokens to consume.
102      * @return the number of tokens that were actually consumed.
103      */
get(int n)104     public int get(int n) {
105         fill();
106         if (n <= 0) {
107             return 0;
108         }
109         if (n > mAvailable) {
110             int got = mAvailable;
111             mAvailable = 0;
112             return got;
113         }
114         mAvailable -= n;
115         return n;
116     }
117 
fill()118     private void fill() {
119         final long now = scaledTime();
120         final int diff = (int) (now - mLastFill);
121         mAvailable = Math.min(mCapacity, mAvailable + diff);
122         mLastFill = now;
123     }
124 
scaledTime()125     private long scaledTime() {
126         return SystemClock.elapsedRealtime() / mFillDelta;
127     }
128 }
129