1 | /* |
2 | * Copyright 2015-present Facebook, Inc. |
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 | #include <folly/test/TokenBucketTest.h> |
18 | |
19 | #include <folly/portability/GTest.h> |
20 | |
21 | using namespace folly; |
22 | |
23 | TEST(TokenBucket, ReverseTime) { |
24 | const double rate = 1000; |
25 | TokenBucket tokenBucket(rate, rate * 0.01 + 1e-6, 0); |
26 | size_t count = 0; |
27 | while (tokenBucket.consume(1, 0.1)) { |
28 | count += 1; |
29 | } |
30 | EXPECT_EQ(10, count); |
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()); |
36 | } |
37 | |
38 | TEST_P(TokenBucketTest, sanity) { |
39 | std::pair<double, double> params = GetParam(); |
40 | double rate = params.first; |
41 | double consumeSize = params.second; |
42 | |
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; |
55 | } |
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); |
62 | } |
63 | } |
64 | |
65 | static std::vector<std::pair<double, double>> rateToConsumeSize = { |
66 | {100, 1}, |
67 | {1000, 1}, |
68 | {10000, 1}, |
69 | {10000, 5}, |
70 | }; |
71 | |
72 | INSTANTIATE_TEST_CASE_P( |
73 | TokenBucket, |
74 | TokenBucketTest, |
75 | ::testing::ValuesIn(rateToConsumeSize)); |
76 | |
77 | void doTokenBucketTest(double maxQps, double consumeSize) { |
78 | const double tenMillisecondBurst = maxQps * 0.010; |
79 | // Select a burst size of 10 milliseconds at the max rate or the consume size |
80 | // if 10 ms at maxQps is too small. |
81 | const double burstSize = std::max(consumeSize, tenMillisecondBurst); |
82 | TokenBucket tokenBucket(maxQps, burstSize, 0); |
83 | double tokenCounter = 0; |
84 | double currentTime = 0; |
85 | // Simulate time advancing 10 seconds |
86 | for (; currentTime <= 10.0; currentTime += 0.001) { |
87 | EXPECT_FALSE(tokenBucket.consume(burstSize + 1, currentTime)); |
88 | while (tokenBucket.consume(consumeSize, currentTime)) { |
89 | tokenCounter += consumeSize; |
90 | } |
91 | // Tokens consumed should exceed some lower bound based on maxQps. |
92 | // Note: The token bucket implementation is not precise, so the lower bound |
93 | // is somewhat fudged. The upper bound is accurate however. |
94 | EXPECT_LE(maxQps * currentTime * 0.9 - 1, tokenCounter); |
95 | // Tokens consumed should not exceed some upper bound based on maxQps. |
96 | EXPECT_GE(maxQps * currentTime + 1e-6, tokenCounter); |
97 | } |
98 | } |
99 | |
100 | TEST(TokenBucket, sanity) { |
101 | doTokenBucketTest(100, 1); |
102 | doTokenBucketTest(1000, 1); |
103 | doTokenBucketTest(10000, 1); |
104 | // Consume more than one at a time. |
105 | doTokenBucketTest(10000, 5); |
106 | } |
107 | |
108 | TEST(TokenBucket, ReverseTime2) { |
109 | const double rate = 1000; |
110 | TokenBucket tokenBucket(rate, rate * 0.01 + 1e-6); |
111 | size_t count = 0; |
112 | while (tokenBucket.consume(1, 0.1)) { |
113 | count += 1; |
114 | } |
115 | EXPECT_EQ(10, count); |
116 | // Going backwards in time has no affect on the toke count (this protects |
117 | // against different threads providing out of order timestamps). |
118 | double tokensBefore = tokenBucket.available(); |
119 | EXPECT_FALSE(tokenBucket.consume(1, 0.09999999)); |
120 | EXPECT_EQ(tokensBefore, tokenBucket.available()); |
121 | } |
122 | |
123 | TEST(TokenBucket, drainOnFail) { |
124 | DynamicTokenBucket tokenBucket; |
125 | |
126 | // Almost empty the bucket |
127 | EXPECT_TRUE(tokenBucket.consume(9, 10, 10, 1)); |
128 | |
129 | // Request more tokens than available |
130 | EXPECT_FALSE(tokenBucket.consume(5, 10, 10, 1)); |
131 | EXPECT_DOUBLE_EQ(1.0, tokenBucket.available(10, 10, 1)); |
132 | |
133 | // Again request more tokens than available, but ask to drain |
134 | EXPECT_DOUBLE_EQ(1.0, tokenBucket.consumeOrDrain(5, 10, 10, 1)); |
135 | EXPECT_DOUBLE_EQ(0.0, tokenBucket.consumeOrDrain(1, 10, 10, 1)); |
136 | } |
137 | |
138 | TEST(TokenBucket, returnTokensTest) { |
139 | DynamicTokenBucket tokenBucket; |
140 | |
141 | // Empty the bucket. |
142 | EXPECT_TRUE(tokenBucket.consume(10, 10, 10, 5)); |
143 | // consume should fail now. |
144 | EXPECT_FALSE(tokenBucket.consume(1, 10, 10, 5)); |
145 | EXPECT_DOUBLE_EQ(0.0, tokenBucket.consumeOrDrain(1, 10, 10, 5)); |
146 | |
147 | // Return tokens. Return 40 'excess' tokens but they wont be available to |
148 | // later callers. |
149 | tokenBucket.returnTokens(50, 10); |
150 | // Should be able to allocate 10 tokens again but the extra 40 returned in |
151 | // previous call are gone. |
152 | EXPECT_TRUE(tokenBucket.consume(10, 10, 10, 5)); |
153 | EXPECT_FALSE(tokenBucket.consume(1, 10, 10, 5)); |
154 | } |
155 | |
156 | TEST(TokenBucket, consumeOrBorrowTest) { |
157 | DynamicTokenBucket tokenBucket; |
158 | |
159 | // Empty the bucket. |
160 | EXPECT_TRUE(tokenBucket.consume(10, 10, 10, 1)); |
161 | // consume should fail now. |
162 | EXPECT_FALSE(tokenBucket.consume(1, 10, 10, 1)); |
163 | // Now borrow from future allocations. Each call is asking for 1s worth of |
164 | // allocations so it should return (i+1)*1s in the ith iteration as the time |
165 | // caller needs to wait. |
166 | for (int i = 0; i < 10; ++i) { |
167 | auto waitTime = tokenBucket.consumeWithBorrowNonBlocking(10, 10, 10, 1); |
168 | EXPECT_TRUE(waitTime.has_value()); |
169 | EXPECT_DOUBLE_EQ((i + 1) * 1.0, *waitTime); |
170 | } |
171 | |
172 | // No allocation will succeed until nowInSeconds goes higher than 11s. |
173 | EXPECT_FALSE(tokenBucket.consume(1, 10, 10, 11)); |
174 | } |
175 | |