folly/test/TokenBucketTest.cpp

   1 /*
   2  * Copyright 2016 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 <gtest/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));