folly/test/TokenBucketTest.cpp

   1 /*
   2  * Copyright 2017 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(TokenBucket,
  73                         TokenBucketTest,
  74                         ::testing::ValuesIn(rateToConsumeSize));
  75
  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;
  89     }
  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);
  96   }
  97 }
  98
  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);
 105 }
 106
 107 TEST(TokenBucket, ReverseTime2) {
 108   const double rate = 1000;
 109   TokenBucket tokenBucket(rate, rate * 0.01 + 1e-6);
 110   size_t count = 0;
 111   while (tokenBucket.consume(1, 0.1)) {
 112     count += 1;
 113   }
 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());
 120 }
 121
 122 TEST(TokenBucket, drainOnFail) {
 123   DynamicTokenBucket tokenBucket;
 124
 125   // Almost empty the bucket
 126   EXPECT_TRUE(tokenBucket.consume(9, 10, 10, 1));
 127
 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));
 131
 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));
 135 }