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 // Exported to Mainline modules; cannot use annotations
40 // @android.ravenwood.annotation.RavenwoodKeepWholeClass
41 public class TokenBucket {
42 
43     private final int mFillDelta; // Time in ms it takes to generate one token.
44     private final int mCapacity;  // Maximum number of tokens that can be stored.
45     private long mLastFill;       // Last time in ms the bucket generated tokens.
46     private int mAvailable;       // Current number of available tokens.
47 
48     /**
49      * Create a new TokenBucket.
50      * @param deltaMs the time in milliseconds it takes to generate a new token.
51      * Must be strictly positive.
52      * @param capacity the maximum token capacity. Must be strictly positive.
53      * @param tokens the starting amount of token. Must be positive or zero.
54      */
TokenBucket(int deltaMs, int capacity, int tokens)55     public TokenBucket(int deltaMs, int capacity, int tokens) {
56         mFillDelta = checkArgumentPositive(deltaMs, "deltaMs must be strictly positive");
57         mCapacity = checkArgumentPositive(capacity, "capacity must be strictly positive");
58         mAvailable = Math.min(checkArgumentNonnegative(tokens), mCapacity);
59         mLastFill = scaledTime();
60     }
61 
62     /**
63      * Create a new TokenBucket that starts completely filled.
64      * @param deltaMs the time in milliseconds it takes to generate a new token.
65      * Must be strictly positive.
66      * @param capacity the maximum token capacity. Must be strictly positive.
67      */
TokenBucket(int deltaMs, int capacity)68     public TokenBucket(int deltaMs, int capacity) {
69         this(deltaMs, capacity, capacity);
70     }
71 
72     /** Reset this TokenBucket and set its number of available tokens. */
reset(int tokens)73     public void reset(int tokens) {
74         checkArgumentNonnegative(tokens);
75         mAvailable = Math.min(tokens, mCapacity);
76         mLastFill = scaledTime();
77     }
78 
79     /** Returns this TokenBucket maximum token capacity. */
capacity()80     public int capacity() {
81         return mCapacity;
82     }
83 
84     /** Returns this TokenBucket currently number of available tokens. */
available()85     public int available() {
86         fill();
87         return mAvailable;
88     }
89 
90     /** Returns true if this TokenBucket as one or more tokens available. */
has()91     public boolean has() {
92         fill();
93         return mAvailable > 0;
94     }
95 
96     /** Consumes a token from this TokenBucket and returns true if a token is available. */
get()97     public boolean get() {
98         return (get(1) == 1);
99     }
100 
101     /**
102      * Try to consume many tokens from this TokenBucket.
103      * @param n the number of tokens to consume.
104      * @return the number of tokens that were actually consumed.
105      */
get(int n)106     public int get(int n) {
107         fill();
108         if (n <= 0) {
109             return 0;
110         }
111         if (n > mAvailable) {
112             int got = mAvailable;
113             mAvailable = 0;
114             return got;
115         }
116         mAvailable -= n;
117         return n;
118     }
119 
fill()120     private void fill() {
121         final long now = scaledTime();
122         final int diff = (int) (now - mLastFill);
123         mAvailable = Math.min(mCapacity, mAvailable + diff);
124         mLastFill = now;
125     }
126 
scaledTime()127     private long scaledTime() {
128         return SystemClock.elapsedRealtime() / mFillDelta;
129     }
130 }
131