1 //  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
2 //  This source code is licensed under both the GPLv2 (found in the
3 //  COPYING file in the root directory) and Apache 2.0 License
4 //  (found in the LICENSE.Apache file in the root directory).
5 //
6 #include <ratio>
7 
8 #include "db/write_controller.h"
9 
10 #include "rocksdb/env.h"
11 #include "test_util/testharness.h"
12 
13 namespace ROCKSDB_NAMESPACE {
14 
15 class WriteControllerTest : public testing::Test {};
16 
17 class TimeSetEnv : public EnvWrapper {
18  public:
TimeSetEnv()19   explicit TimeSetEnv() : EnvWrapper(nullptr) {}
20   uint64_t now_micros_ = 6666;
NowNanos()21   uint64_t NowNanos() override { return now_micros_ * std::milli::den; }
22 };
23 
TEST_F(WriteControllerTest,ChangeDelayRateTest)24 TEST_F(WriteControllerTest, ChangeDelayRateTest) {
25   TimeSetEnv env;
26   WriteController controller(40000000u);  // also set max delayed rate
27   controller.set_delayed_write_rate(10000000u);
28   auto delay_token_0 =
29       controller.GetDelayToken(controller.delayed_write_rate());
30   ASSERT_EQ(static_cast<uint64_t>(2000000),
31             controller.GetDelay(&env, 20000000u));
32   auto delay_token_1 = controller.GetDelayToken(2000000u);
33   ASSERT_EQ(static_cast<uint64_t>(10000000),
34             controller.GetDelay(&env, 20000000u));
35   auto delay_token_2 = controller.GetDelayToken(1000000u);
36   ASSERT_EQ(static_cast<uint64_t>(20000000),
37             controller.GetDelay(&env, 20000000u));
38   auto delay_token_3 = controller.GetDelayToken(20000000u);
39   ASSERT_EQ(static_cast<uint64_t>(1000000),
40             controller.GetDelay(&env, 20000000u));
41   // This is more than max rate. Max delayed rate will be used.
42   auto delay_token_4 =
43       controller.GetDelayToken(controller.delayed_write_rate() * 3);
44   ASSERT_EQ(static_cast<uint64_t>(500000),
45             controller.GetDelay(&env, 20000000u));
46 }
47 
TEST_F(WriteControllerTest,SanityTest)48 TEST_F(WriteControllerTest, SanityTest) {
49   WriteController controller(10000000u);
50   auto stop_token_1 = controller.GetStopToken();
51   auto stop_token_2 = controller.GetStopToken();
52 
53   ASSERT_TRUE(controller.IsStopped());
54   stop_token_1.reset();
55   ASSERT_TRUE(controller.IsStopped());
56   stop_token_2.reset();
57   ASSERT_FALSE(controller.IsStopped());
58 
59   TimeSetEnv env;
60 
61   auto delay_token_1 = controller.GetDelayToken(10000000u);
62   ASSERT_EQ(static_cast<uint64_t>(2000000),
63             controller.GetDelay(&env, 20000000u));
64 
65   env.now_micros_ += 1999900u;  // sleep debt 1000
66 
67   auto delay_token_2 = controller.GetDelayToken(10000000u);
68   // Rate reset after changing the token.
69   ASSERT_EQ(static_cast<uint64_t>(2000000),
70             controller.GetDelay(&env, 20000000u));
71 
72   env.now_micros_ += 1999900u;  // sleep debt 1000
73 
74   // One refill: 10240 bytes allowed, 1000 used, 9240 left
75   ASSERT_EQ(static_cast<uint64_t>(1124), controller.GetDelay(&env, 1000u));
76   env.now_micros_ += 1124u;  // sleep debt 0
77 
78   delay_token_2.reset();
79   // 1000 used, 8240 left
80   ASSERT_EQ(static_cast<uint64_t>(0), controller.GetDelay(&env, 1000u));
81 
82   env.now_micros_ += 100u;  // sleep credit 100
83   // 1000 used, 7240 left
84   ASSERT_EQ(static_cast<uint64_t>(0), controller.GetDelay(&env, 1000u));
85 
86   env.now_micros_ += 100u;  // sleep credit 200
87   // One refill: 10240 fileed, sleep credit generates 2000. 8000 used
88   //             7240 + 10240 + 2000 - 8000 = 11480 left
89   ASSERT_EQ(static_cast<uint64_t>(1024u), controller.GetDelay(&env, 8000u));
90 
91   env.now_micros_ += 200u;  // sleep debt 824
92   // 1000 used, 10480 left.
93   ASSERT_EQ(static_cast<uint64_t>(0), controller.GetDelay(&env, 1000u));
94 
95   env.now_micros_ += 200u;  // sleep debt 624
96   // Out of bound sleep, still 10480 left
97   ASSERT_EQ(static_cast<uint64_t>(3000624u),
98             controller.GetDelay(&env, 30000000u));
99 
100   env.now_micros_ += 3000724u;  // sleep credit 100
101   // 6000 used, 4480 left.
102   ASSERT_EQ(static_cast<uint64_t>(0), controller.GetDelay(&env, 6000u));
103 
104   env.now_micros_ += 200u;  // sleep credit 300
105   // One refill, credit 4480 balance + 3000 credit + 10240 refill
106   // Use 8000, 9720 left
107   ASSERT_EQ(static_cast<uint64_t>(1024u), controller.GetDelay(&env, 8000u));
108 
109   env.now_micros_ += 3024u;  // sleep credit 2000
110 
111   // 1720 left
112   ASSERT_EQ(static_cast<uint64_t>(0u), controller.GetDelay(&env, 8000u));
113 
114   // 1720 balance + 20000 credit = 20170 left
115   // Use 8000, 12170 left
116   ASSERT_EQ(static_cast<uint64_t>(0u), controller.GetDelay(&env, 8000u));
117 
118   // 4170 left
119   ASSERT_EQ(static_cast<uint64_t>(0u), controller.GetDelay(&env, 8000u));
120 
121   // Need a refill
122   ASSERT_EQ(static_cast<uint64_t>(1024u), controller.GetDelay(&env, 9000u));
123 
124   delay_token_1.reset();
125   ASSERT_EQ(static_cast<uint64_t>(0), controller.GetDelay(&env, 30000000u));
126   delay_token_1.reset();
127   ASSERT_FALSE(controller.IsStopped());
128 }
129 
130 }  // namespace ROCKSDB_NAMESPACE
131 
main(int argc,char ** argv)132 int main(int argc, char** argv) {
133   ::testing::InitGoogleTest(&argc, argv);
134   return RUN_ALL_TESTS();
135 }
136