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 "db/write_controller.h"
7 
8 #include <array>
9 #include <ratio>
10 
11 #include "rocksdb/system_clock.h"
12 #include "test_util/testharness.h"
13 
14 namespace ROCKSDB_NAMESPACE {
15 namespace {
16 class TimeSetClock : public SystemClockWrapper {
17  public:
TimeSetClock()18   explicit TimeSetClock() : SystemClockWrapper(nullptr) {}
Name() const19   const char* Name() const override { return "TimeSetClock"; }
20   uint64_t now_micros_ = 6666;
NowNanos()21   uint64_t NowNanos() override { return now_micros_ * std::milli::den; }
22 };
23 }  // namespace
24 class WriteControllerTest : public testing::Test {
25  public:
WriteControllerTest()26   WriteControllerTest() { clock_ = std::make_shared<TimeSetClock>(); }
27   std::shared_ptr<TimeSetClock> clock_;
28 };
29 
30 // Make tests easier to read
31 #define MILLION *1000000u
32 #define MB MILLION
33 #define MBPS MILLION
34 #define SECS MILLION  // in microseconds
35 
TEST_F(WriteControllerTest,BasicAPI)36 TEST_F(WriteControllerTest, BasicAPI) {
37   WriteController controller(40 MBPS);  // also set max delayed rate
38   EXPECT_EQ(controller.delayed_write_rate(), 40 MBPS);
39   EXPECT_FALSE(controller.IsStopped());
40   EXPECT_FALSE(controller.NeedsDelay());
41   EXPECT_EQ(0, controller.GetDelay(clock_.get(), 100 MB));
42 
43   // set, get
44   controller.set_delayed_write_rate(20 MBPS);
45   EXPECT_EQ(controller.delayed_write_rate(), 20 MBPS);
46   EXPECT_FALSE(controller.IsStopped());
47   EXPECT_FALSE(controller.NeedsDelay());
48   EXPECT_EQ(0, controller.GetDelay(clock_.get(), 100 MB));
49 
50   {
51     // set with token, get
52     auto delay_token_0 = controller.GetDelayToken(10 MBPS);
53     EXPECT_EQ(controller.delayed_write_rate(), 10 MBPS);
54     EXPECT_FALSE(controller.IsStopped());
55     EXPECT_TRUE(controller.NeedsDelay());
56     // test with delay
57     EXPECT_EQ(2 SECS, controller.GetDelay(clock_.get(), 20 MB));
58     clock_->now_micros_ += 2 SECS;  // pay the "debt"
59 
60     auto delay_token_1 = controller.GetDelayToken(2 MBPS);
61     EXPECT_EQ(10 SECS, controller.GetDelay(clock_.get(), 20 MB));
62     clock_->now_micros_ += 10 SECS;  // pay the "debt"
63 
64     auto delay_token_2 = controller.GetDelayToken(1 MBPS);
65     EXPECT_EQ(20 SECS, controller.GetDelay(clock_.get(), 20 MB));
66     clock_->now_micros_ += 20 SECS;  // pay the "debt"
67 
68     auto delay_token_3 = controller.GetDelayToken(20 MBPS);
69     EXPECT_EQ(1 SECS, controller.GetDelay(clock_.get(), 20 MB));
70     clock_->now_micros_ += 1 SECS;  // pay the "debt"
71 
72     // 60M is more than the max rate of 40M. Max rate will be used.
73     EXPECT_EQ(controller.delayed_write_rate(), 20 MBPS);
74     auto delay_token_4 =
75         controller.GetDelayToken(controller.delayed_write_rate() * 3);
76     EXPECT_EQ(controller.delayed_write_rate(), 40 MBPS);
77     EXPECT_EQ(static_cast<uint64_t>(0.5 SECS),
78               controller.GetDelay(clock_.get(), 20 MB));
79 
80     EXPECT_FALSE(controller.IsStopped());
81     EXPECT_TRUE(controller.NeedsDelay());
82 
83     // Test stop tokens
84     {
85       auto stop_token_1 = controller.GetStopToken();
86       EXPECT_TRUE(controller.IsStopped());
87       EXPECT_EQ(0, controller.GetDelay(clock_.get(), 100 MB));
88       {
89         auto stop_token_2 = controller.GetStopToken();
90         EXPECT_TRUE(controller.IsStopped());
91         EXPECT_EQ(0, controller.GetDelay(clock_.get(), 100 MB));
92       }
93       EXPECT_TRUE(controller.IsStopped());
94       EXPECT_EQ(0, controller.GetDelay(clock_.get(), 100 MB));
95     }
96     // Stop tokens released
97     EXPECT_FALSE(controller.IsStopped());
98     EXPECT_TRUE(controller.NeedsDelay());
99     EXPECT_EQ(controller.delayed_write_rate(), 40 MBPS);
100     // pay the previous "debt"
101     clock_->now_micros_ += static_cast<uint64_t>(0.5 SECS);
102     EXPECT_EQ(1 SECS, controller.GetDelay(clock_.get(), 40 MB));
103   }
104 
105   // Delay tokens released
106   EXPECT_FALSE(controller.NeedsDelay());
107 }
108 
TEST_F(WriteControllerTest,StartFilled)109 TEST_F(WriteControllerTest, StartFilled) {
110   WriteController controller(10 MBPS);
111 
112   // Attempt to write two things that combined would be allowed within
113   // a single refill interval
114   auto delay_token_0 =
115       controller.GetDelayToken(controller.delayed_write_rate());
116 
117   // Verify no delay because write rate has not been exceeded within
118   // refill interval.
119   EXPECT_EQ(0U, controller.GetDelay(clock_.get(), 2000u /*bytes*/));
120   EXPECT_EQ(0U, controller.GetDelay(clock_.get(), 2000u /*bytes*/));
121 
122   // Allow refill (kMicrosPerRefill)
123   clock_->now_micros_ += 1000;
124 
125   // Again
126   EXPECT_EQ(0U, controller.GetDelay(clock_.get(), 2000u /*bytes*/));
127   EXPECT_EQ(0U, controller.GetDelay(clock_.get(), 2000u /*bytes*/));
128 
129   // Control: something bigger that would exceed write rate within interval
130   uint64_t delay = controller.GetDelay(clock_.get(), 10 MB);
131   EXPECT_GT(1.0 * delay, 0.999 SECS);
132   EXPECT_LT(1.0 * delay, 1.001 SECS);
133 }
134 
TEST_F(WriteControllerTest,DebtAccumulation)135 TEST_F(WriteControllerTest, DebtAccumulation) {
136   WriteController controller(10 MBPS);
137 
138   std::array<std::unique_ptr<WriteControllerToken>, 10> tokens;
139 
140   // Accumulate a time delay debt with no passage of time, like many column
141   // families delaying writes simultaneously. (Old versions of WriteController
142   // would reset the debt on every GetDelayToken.)
143   uint64_t debt = 0;
144   for (unsigned i = 0; i < tokens.size(); ++i) {
145     tokens[i] = controller.GetDelayToken((i + 1u) MBPS);
146     uint64_t delay = controller.GetDelay(clock_.get(), 63 MB);
147     ASSERT_GT(delay, debt);
148     uint64_t incremental = delay - debt;
149     ASSERT_EQ(incremental, (63 SECS) / (i + 1u));
150     debt += incremental;
151   }
152 
153   // Pay down the debt
154   clock_->now_micros_ += debt;
155   debt = 0;
156 
157   // Now accumulate debt with some passage of time.
158   for (unsigned i = 0; i < tokens.size(); ++i) {
159     // Debt is accumulated in time, not in bytes, so this new write
160     // limit is not applied to prior requested delays, even it they are
161     // in progress.
162     tokens[i] = controller.GetDelayToken((i + 1u) MBPS);
163     uint64_t delay = controller.GetDelay(clock_.get(), 63 MB);
164     ASSERT_GT(delay, debt);
165     uint64_t incremental = delay - debt;
166     ASSERT_EQ(incremental, (63 SECS) / (i + 1u));
167     debt += incremental;
168     uint64_t credit = debt / 2;
169     clock_->now_micros_ += credit;
170     debt -= credit;
171   }
172 
173   // Pay down the debt
174   clock_->now_micros_ += debt;
175   debt = 0;    // consistent state
176   (void)debt;  // appease clang-analyze
177 
178   // Verify paid down
179   EXPECT_EQ(0U, controller.GetDelay(clock_.get(), 100u /*small bytes*/));
180 
181   // Accumulate another debt, without accounting, and releasing tokens
182   for (unsigned i = 0; i < tokens.size(); ++i) {
183     // Big and small are delayed
184     ASSERT_LT(0U, controller.GetDelay(clock_.get(), 63 MB));
185     ASSERT_LT(0U, controller.GetDelay(clock_.get(), 100u /*small bytes*/));
186     tokens[i].reset();
187   }
188   // All tokens released.
189   // Verify that releasing all tokens pays down debt, even with no time passage.
190   tokens[0] = controller.GetDelayToken(1 MBPS);
191   ASSERT_EQ(0U, controller.GetDelay(clock_.get(), 100u /*small bytes*/));
192 }
193 
194 // This may or may not be a "good" feature, but it's an old feature
TEST_F(WriteControllerTest,CreditAccumulation)195 TEST_F(WriteControllerTest, CreditAccumulation) {
196   WriteController controller(10 MBPS);
197 
198   std::array<std::unique_ptr<WriteControllerToken>, 10> tokens;
199 
200   // Ensure started
201   tokens[0] = controller.GetDelayToken(1 MBPS);
202   ASSERT_EQ(10 SECS, controller.GetDelay(clock_.get(), 10 MB));
203   clock_->now_micros_ += 10 SECS;
204 
205   // Accumulate a credit
206   uint64_t credit = 1000 SECS /* see below: * 1 MB / 1 SEC */;
207   clock_->now_micros_ += credit;
208 
209   // Spend some credit (burst of I/O)
210   for (unsigned i = 0; i < tokens.size(); ++i) {
211     tokens[i] = controller.GetDelayToken((i + 1u) MBPS);
212     ASSERT_EQ(0U, controller.GetDelay(clock_.get(), 63 MB));
213     // In WriteController, credit is accumulated in bytes, not in time.
214     // After an "unnecessary" delay, all of our time credit will be
215     // translated to bytes on the next operation, in this case with
216     // setting 1 MBPS. So regardless of the rate at delay time, we just
217     // account for the bytes.
218     credit -= 63 MB;
219   }
220   // Spend remaining credit
221   tokens[0] = controller.GetDelayToken(1 MBPS);
222   ASSERT_EQ(0U, controller.GetDelay(clock_.get(), credit));
223   // Verify
224   ASSERT_EQ(10 SECS, controller.GetDelay(clock_.get(), 10 MB));
225   clock_->now_micros_ += 10 SECS;
226 
227   // Accumulate a credit, no accounting
228   clock_->now_micros_ += 1000 SECS;
229 
230   // Spend a small amount, releasing tokens
231   for (unsigned i = 0; i < tokens.size(); ++i) {
232     ASSERT_EQ(0U, controller.GetDelay(clock_.get(), 3 MB));
233     tokens[i].reset();
234   }
235 
236   // All tokens released.
237   // Verify credit is wiped away on new delay.
238   tokens[0] = controller.GetDelayToken(1 MBPS);
239   ASSERT_EQ(10 SECS, controller.GetDelay(clock_.get(), 10 MB));
240 }
241 
242 }  // namespace ROCKSDB_NAMESPACE
243 
main(int argc,char ** argv)244 int main(int argc, char** argv) {
245   ::testing::InitGoogleTest(&argc, argv);
246   return RUN_ALL_TESTS();
247 }
248