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
17package com.android.internal.util;
18
19import android.os.SystemClock;
20
21import static com.android.internal.util.Preconditions.checkArgumentNonnegative;
22import 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 */
39public 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     */
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     */
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. */
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. */
78    public int capacity() {
79        return mCapacity;
80    }
81
82    /** Returns this TokenBucket currently number of available tokens. */
83    public int available() {
84        fill();
85        return mAvailable;
86    }
87
88    /** Returns true if this TokenBucket as one or more tokens available. */
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. */
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     */
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
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
125    private long scaledTime() {
126        return SystemClock.elapsedRealtime() / mFillDelta;
127    }
128}
129