2 * Copyright 2017 Facebook, Inc.
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
8 * http://www.apache.org/licenses/LICENSE-2.0
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.
17 #include <folly/test/TokenBucketTest.h>
19 #include <folly/portability/GTest.h>
21 using namespace folly;
23 TEST(TokenBucket, ReverseTime) {
24 const double rate = 1000;
25 TokenBucket tokenBucket(rate, rate * 0.01 + 1e-6, 0);
27 while (tokenBucket.consume(1, 0.1)) {
31 // Going backwards in time has no affect on the toke count (this protects
32 // against different threads providing out of order timestamps).
33 double tokensBefore = tokenBucket.available();
34 EXPECT_FALSE(tokenBucket.consume(1, 0.09999999));
35 EXPECT_EQ(tokensBefore, tokenBucket.available());
38 TEST_P(TokenBucketTest, sanity) {
39 std::pair<double, double> params = GetParam();
40 double rate = params.first;
41 double consumeSize = params.second;
43 const double tenMillisecondBurst = rate * 0.010;
44 // Select a burst size of 10 milliseconds at the max rate or the consume size
45 // if 10 ms at rate is too small.
46 const double burstSize = std::max(consumeSize, tenMillisecondBurst);
47 TokenBucket tokenBucket(rate, burstSize, 0);
48 double tokenCounter = 0;
49 double currentTime = 0;
50 // Simulate time advancing 10 seconds
51 for (; currentTime <= 10.0; currentTime += 0.001) {
52 EXPECT_FALSE(tokenBucket.consume(burstSize + 1, currentTime));
53 while (tokenBucket.consume(consumeSize, currentTime)) {
54 tokenCounter += consumeSize;
56 // Tokens consumed should exceed some lower bound based on rate.
57 // Note: The token bucket implementation is not precise, so the lower bound
58 // is somewhat fudged. The upper bound is accurate however.
59 EXPECT_LE(rate * currentTime * 0.9 - 1, tokenCounter);
60 // Tokens consumed should not exceed some upper bound based on rate.
61 EXPECT_GE(rate * currentTime + 1e-6, tokenCounter);
65 static std::vector<std::pair<double, double> > rateToConsumeSize = {
72 INSTANTIATE_TEST_CASE_P(TokenBucket,
74 ::testing::ValuesIn(rateToConsumeSize));
76 void doTokenBucketTest(double maxQps, double consumeSize) {
77 const double tenMillisecondBurst = maxQps * 0.010;
78 // Select a burst size of 10 milliseconds at the max rate or the consume size
79 // if 10 ms at maxQps is too small.
80 const double burstSize = std::max(consumeSize, tenMillisecondBurst);
81 TokenBucket tokenBucket(maxQps, burstSize, 0);
82 double tokenCounter = 0;
83 double currentTime = 0;
84 // Simulate time advancing 10 seconds
85 for (; currentTime <= 10.0; currentTime += 0.001) {
86 EXPECT_FALSE(tokenBucket.consume(burstSize + 1, currentTime));
87 while (tokenBucket.consume(consumeSize, currentTime)) {
88 tokenCounter += consumeSize;
90 // Tokens consumed should exceed some lower bound based on maxQps.
91 // Note: The token bucket implementation is not precise, so the lower bound
92 // is somewhat fudged. The upper bound is accurate however.
93 EXPECT_LE(maxQps * currentTime * 0.9 - 1, tokenCounter);
94 // Tokens consumed should not exceed some upper bound based on maxQps.
95 EXPECT_GE(maxQps * currentTime + 1e-6, tokenCounter);
99 TEST(TokenBucket, sanity) {
100 doTokenBucketTest(100, 1);
101 doTokenBucketTest(1000, 1);
102 doTokenBucketTest(10000, 1);
103 // Consume more than one at a time.
104 doTokenBucketTest(10000, 5);
107 TEST(TokenBucket, ReverseTime2) {
108 const double rate = 1000;
109 TokenBucket tokenBucket(rate, rate * 0.01 + 1e-6);
111 while (tokenBucket.consume(1, 0.1)) {
114 EXPECT_EQ(10, count);
115 // Going backwards in time has no affect on the toke count (this protects
116 // against different threads providing out of order timestamps).
117 double tokensBefore = tokenBucket.available();
118 EXPECT_FALSE(tokenBucket.consume(1, 0.09999999));
119 EXPECT_EQ(tokensBefore, tokenBucket.available());
122 TEST(TokenBucket, drainOnFail) {
123 DynamicTokenBucket tokenBucket;
125 // Almost empty the bucket
126 EXPECT_TRUE(tokenBucket.consume(9, 10, 10, 1));
128 // Request more tokens than available
129 EXPECT_FALSE(tokenBucket.consume(5, 10, 10, 1));
130 EXPECT_DOUBLE_EQ(1.0, tokenBucket.available(10, 10, 1));
132 // Again request more tokens than available, but ask to drain
133 EXPECT_DOUBLE_EQ(1.0, tokenBucket.consumeOrDrain(5, 10, 10, 1));
134 EXPECT_DOUBLE_EQ(0.0, tokenBucket.consumeOrDrain(1, 10, 10, 1));