1 /*
2 * Copyright (c) Facebook, Inc. and its affiliates.
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 <wangle/acceptor/ConnectionManager.h>
18
19 #include <folly/portability/GFlags.h>
20 #include <folly/portability/GMock.h>
21 #include <folly/portability/GTest.h>
22
23 using namespace testing;
24 using namespace wangle;
25
26 namespace {
27
28 class ConnectionManagerTest;
29
30 class MockConnection : public ManagedConnection {
31 public:
32 using Mock = StrictMock<MockConnection>;
33 using UniquePtr = folly::DelayedDestructionUniquePtr<Mock>;
34
makeUnique(ConnectionManagerTest * test)35 static UniquePtr makeUnique(ConnectionManagerTest* test) {
36 UniquePtr p(new StrictMock<MockConnection>(test));
37 return p;
38 }
39
MockConnection(ConnectionManagerTest * test)40 explicit MockConnection(ConnectionManagerTest* test) : test_(test) {
41 EXPECT_CALL(*this, isBusy()).WillRepeatedly(Return(false));
42 EXPECT_CALL(*this, dumpConnectionState(testing::_))
43 .WillRepeatedly(Return());
44 ON_CALL(*this, closeWhenIdle()).WillByDefault(Invoke([this] {
45 closeWhenIdle_ = true;
46 closeWhenIdleImpl();
47 }));
48 }
49
50 MOCK_METHOD0(timeoutExpired_, void());
timeoutExpired()51 void timeoutExpired() noexcept override {
52 timeoutExpired_();
53 }
54
55 MOCK_CONST_METHOD1(describe, void(std::ostream& os));
56
57 MOCK_CONST_METHOD0(isBusy, bool());
58
59 MOCK_CONST_METHOD0(getIdleTime, std::chrono::milliseconds());
60 MOCK_CONST_METHOD0(
61 getLastActivityElapsedTime,
62 folly::Optional<std::chrono::milliseconds>());
63
64 MOCK_METHOD0(notifyPendingShutdown, void());
65 MOCK_METHOD0(closeWhenIdle, void());
66 MOCK_METHOD1(dropConnection, void(const std::string&));
67 MOCK_METHOD1(dumpConnectionState, void(uint8_t));
68 MOCK_METHOD2(drainConnections, void(double, std::chrono::milliseconds));
69
setIdle(bool idle)70 void setIdle(bool idle) {
71 idle_ = idle;
72 closeWhenIdleImpl();
73 }
74
75 void closeWhenIdleImpl();
76
77 ConnectionManagerTest* test_{nullptr};
78 bool idle_{false};
79 bool closeWhenIdle_{false};
80 };
81
82 class ConnectionManagerTest : public testing::Test {
83 public:
ConnectionManagerTest()84 ConnectionManagerTest() {
85 cm_ = ConnectionManager::makeUnique(
86 &eventBase_, std::chrono::milliseconds(100), nullptr);
87 addConns(65);
88 }
89
SetUp()90 void SetUp() override {}
91
addConns(uint64_t n)92 void addConns(uint64_t n) {
93 for (size_t i = 0; i < n; i++) {
94 conns_.insert(conns_.begin(), MockConnection::makeUnique(this));
95 cm_->addConnection(conns_.front().get());
96 }
97 }
98
removeConn(MockConnection * connection)99 void removeConn(MockConnection* connection) {
100 for (auto& conn : conns_) {
101 if (conn.get() == connection) {
102 cm_->removeConnection(connection);
103 conn.reset();
104 }
105 }
106 }
107
108 protected:
109 void testAddDuringCloseWhenIdle(bool deactivate);
110
111 folly::EventBase eventBase_;
112 ConnectionManager::UniquePtr cm_;
113 std::vector<MockConnection::UniquePtr> conns_;
114 };
115
closeWhenIdleImpl()116 void MockConnection::closeWhenIdleImpl() {
117 if (idle_ && closeWhenIdle_) {
118 test_->removeConn(this);
119 }
120 }
121
TEST_F(ConnectionManagerTest,testShutdownSequence)122 TEST_F(ConnectionManagerTest, testShutdownSequence) {
123 InSequence enforceOrder;
124
125 // activate one connection, it should not be exempt from notifyPendingShutdown
126 cm_->onActivated(*conns_.front());
127 // make sure the idleIterator points to !end
128 cm_->onDeactivated(*conns_.back());
129 for (const auto& conn : conns_) {
130 EXPECT_CALL(*conn, notifyPendingShutdown());
131 }
132 cm_->initiateGracefulShutdown(std::chrono::milliseconds(50));
133 eventBase_.loopOnce();
134 for (const auto& conn : conns_) {
135 EXPECT_CALL(*conn, closeWhenIdle());
136 }
137
138 eventBase_.loop();
139 }
140
TEST_F(ConnectionManagerTest,testRemoveDrainIterator)141 TEST_F(ConnectionManagerTest, testRemoveDrainIterator) {
142 addConns(1);
143 InSequence enforceOrder;
144
145 // activate one connection, it should not be exempt from notifyPendingShutdown
146 cm_->onActivated(*conns_.front());
147 for (size_t i = 0; i < conns_.size() - 1; i++) {
148 EXPECT_CALL(*conns_[i], notifyPendingShutdown());
149 }
150 auto conn65 = conns_[conns_.size() - 2].get();
151 auto conn66 = conns_[conns_.size() - 1].get();
152 eventBase_.runInLoop([&] {
153 // deactivate the drain iterator
154 cm_->onDeactivated(*conn65);
155 // remove the drain iterator
156 cm_->removeConnection(conn66);
157 // deactivate the new drain iterator, now it's the end of the list
158 cm_->onDeactivated(*conn65);
159 });
160 cm_->initiateGracefulShutdown(std::chrono::milliseconds(50));
161 // Schedule a loop callback to remove the connection pointed to by the drain
162 // iterator
163 eventBase_.loopOnce();
164 for (size_t i = 0; i < conns_.size() - 1; i++) {
165 EXPECT_CALL(*conns_[i], closeWhenIdle());
166 }
167
168 eventBase_.loop();
169 }
170
TEST_F(ConnectionManagerTest,testIdleGraceTimeout)171 TEST_F(ConnectionManagerTest, testIdleGraceTimeout) {
172 InSequence enforceOrder;
173
174 // Slow down the notifyPendingShutdown calls enough so that the idle grace
175 // timeout fires before the end of the loop.
176 // I would prefer a non-sleep solution to this, but I can't think how to do it
177 // without changing the class to expose internal details
178 for (const auto& conn : conns_) {
179 EXPECT_CALL(*conn, notifyPendingShutdown()).WillOnce(Invoke([] {
180 /* sleep override */
181 usleep(1000);
182 }));
183 }
184 cm_->initiateGracefulShutdown(std::chrono::milliseconds(1));
185 eventBase_.loopOnce();
186 for (const auto& conn : conns_) {
187 EXPECT_CALL(*conn, closeWhenIdle());
188 }
189
190 eventBase_.loop();
191 }
192
TEST_F(ConnectionManagerTest,testDropAll)193 TEST_F(ConnectionManagerTest, testDropAll) {
194 InSequence enforceOrder;
195
196 for (const auto& conn : conns_) {
197 EXPECT_CALL(*conn, dropConnection(_))
198 .WillOnce(Invoke(
199 [&](const std::string&) { cm_->removeConnection(conn.get()); }));
200 }
201 cm_->dropAllConnections();
202 }
203
TEST_F(ConnectionManagerTest,testDropPercent)204 TEST_F(ConnectionManagerTest, testDropPercent) {
205 InSequence enforceOrder;
206
207 // Make sure we have exactly 100 connections.
208 const size_t numToAdd = 100 - conns_.size();
209 addConns(numToAdd);
210 const size_t numToRemove = conns_.size() - 100;
211 for (size_t i = 0; i < numToRemove; i++) {
212 removeConn(conns_.begin()->get());
213 }
214 EXPECT_EQ(100, cm_->getNumConnections());
215
216 // Drop 20% of connections.
217 double pct = 0.2;
218 int numToDrop = 100 * pct;
219 auto connIter = conns_.begin();
220 while (connIter != conns_.end() && numToDrop > 0) {
221 EXPECT_CALL(*(*connIter), dropConnection(_));
222 --numToDrop;
223 ++connIter;
224 }
225 cm_->dropConnections(pct);
226
227 // Make sure they are gone.
228 EXPECT_EQ(0, numToDrop);
229 EXPECT_EQ(80, cm_->getNumConnections());
230
231 // Then drop 50% of the remaining 80 connections.
232 pct = 0.5;
233 numToDrop = 80 * pct;
234 while (connIter != conns_.end() && numToDrop > 0) {
235 EXPECT_CALL(*(*connIter), dropConnection(_));
236 --numToDrop;
237 ++connIter;
238 }
239 cm_->dropConnections(pct);
240
241 // Make sure those are gone as well.
242 EXPECT_EQ(0, numToDrop);
243 EXPECT_EQ(40, cm_->getNumConnections());
244 }
245
TEST_F(ConnectionManagerTest,testDrainPercent)246 TEST_F(ConnectionManagerTest, testDrainPercent) {
247 InSequence enforceOrder;
248 double drain_percentage = .123;
249
250 for (size_t i = 58 /* tail .123 of all conns */; i < conns_.size(); ++i) {
251 EXPECT_CALL(*conns_[i], notifyPendingShutdown());
252 }
253
254 cm_->drainConnections(drain_percentage, std::chrono::milliseconds(50));
255
256 for (size_t i = 58; i < conns_.size(); ++i) {
257 EXPECT_CALL(*conns_[i], closeWhenIdle());
258 }
259
260 eventBase_.loop();
261 }
262
TEST_F(ConnectionManagerTest,testDrainPctAfterAll)263 TEST_F(ConnectionManagerTest, testDrainPctAfterAll) {
264 InSequence enforceOrder;
265 double drain_percentage = 0.1;
266
267 for (const auto& conn : conns_) {
268 EXPECT_CALL(*conn, notifyPendingShutdown());
269 }
270
271 cm_->initiateGracefulShutdown(std::chrono::milliseconds(50));
272 cm_->drainConnections(drain_percentage, std::chrono::milliseconds(50));
273 eventBase_.loopOnce();
274
275 for (const auto& conn : conns_) {
276 EXPECT_CALL(*conn, closeWhenIdle());
277 }
278
279 eventBase_.loop();
280 }
281
TEST_F(ConnectionManagerTest,testDrainAllAfterPct)282 TEST_F(ConnectionManagerTest, testDrainAllAfterPct) {
283 InSequence enforceOrder;
284 double drain_pct = 0.8;
285
286 for (auto i = conns_.size() - static_cast<int>(conns_.size() * drain_pct);
287 i < conns_.size();
288 ++i) {
289 EXPECT_CALL(*conns_[i], notifyPendingShutdown());
290 }
291
292 cm_->drainConnections(drain_pct, std::chrono::milliseconds(50));
293
294 for (size_t i = 0;
295 i < conns_.size() - static_cast<size_t>(conns_.size() * drain_pct);
296 ++i) {
297 EXPECT_CALL(*conns_[i], notifyPendingShutdown());
298 }
299
300 cm_->initiateGracefulShutdown(std::chrono::milliseconds(50));
301 eventBase_.loopOnce();
302
303 for (const auto& conn : conns_) {
304 EXPECT_CALL(*conn, closeWhenIdle());
305 }
306
307 eventBase_.loop();
308 }
309
TEST_F(ConnectionManagerTest,testDropIdle)310 TEST_F(ConnectionManagerTest, testDropIdle) {
311 for (const auto& conn : conns_) {
312 // Set everyone to be idle for 100ms
313 EXPECT_CALL(*conn, getIdleTime())
314 .WillRepeatedly(Return(std::chrono::milliseconds(100)));
315 }
316
317 // Mark the first half of the connections idle
318 for (size_t i = 0; i < conns_.size() / 2; i++) {
319 cm_->onDeactivated(*conns_[i]);
320 }
321 // reactivate conn 0
322 cm_->onActivated(*conns_[0]);
323 // remove the first idle conn
324 cm_->removeConnection(conns_[1].get());
325
326 InSequence enforceOrder;
327
328 // Expect the remaining idle conns to drop
329 for (size_t i = 2; i < conns_.size() / 2; i++) {
330 EXPECT_CALL(*conns_[i], dropConnection(_))
331 .WillOnce(Invoke([this, i](const std::string&) {
332 cm_->removeConnection(conns_[i].get());
333 }));
334 }
335
336 cm_->dropIdleConnections(conns_.size());
337 }
338
TEST_F(ConnectionManagerTest,testDropActive)339 TEST_F(ConnectionManagerTest, testDropActive) {
340 for (const auto& conn : conns_) {
341 // Set everyone to be idle for 100ms
342 EXPECT_CALL(*conn, getLastActivityElapsedTime())
343 .WillRepeatedly(Return(std::chrono::milliseconds(100)));
344 }
345
346 // Mark the first half of the connections idle
347 for (size_t i = 0; i < conns_.size(); i++) {
348 if (i < conns_.size() / 2) {
349 cm_->onDeactivated(*conns_[i]);
350 } else {
351 cm_->onActivated(*conns_[i]);
352 }
353 }
354 // remove the first idle conn
355 cm_->removeConnection(conns_[1].get());
356
357 InSequence enforceOrder;
358
359 // Expect all active connections to be dropped
360 for (size_t i = conns_.size() / 2; i < conns_.size(); i++) {
361 EXPECT_CALL(*conns_[i], dropConnection(_))
362 .WillOnce(Invoke([this, i](const std::string&) {
363 cm_->removeConnection(conns_[i].get());
364 }));
365 }
366 cm_->dropActiveConnections(conns_.size(), std::chrono::milliseconds(50));
367 }
368
TEST_F(ConnectionManagerTest,testAddDuringShutdown)369 TEST_F(ConnectionManagerTest, testAddDuringShutdown) {
370 auto extraConn = MockConnection::makeUnique(this);
371 InSequence enforceOrder;
372
373 // activate one connection, it should not be exempt from notifyPendingShutdown
374 cm_->onActivated(*conns_.front());
375 for (const auto& conn : conns_) {
376 EXPECT_CALL(*conn, notifyPendingShutdown());
377 }
378 cm_->initiateGracefulShutdown(std::chrono::milliseconds(50));
379 eventBase_.loopOnce();
380 conns_.insert(conns_.begin(), std::move(extraConn));
381 EXPECT_CALL(*conns_.front(), notifyPendingShutdown());
382 cm_->addConnection(conns_.front().get());
383 for (const auto& conn : conns_) {
384 EXPECT_CALL(*conn, closeWhenIdle());
385 }
386
387 eventBase_.loop();
388 }
389
TEST_F(ConnectionManagerTest,testAddDuringShutdownWithoutIdleGrace)390 TEST_F(ConnectionManagerTest, testAddDuringShutdownWithoutIdleGrace) {
391 auto extraConn = MockConnection::makeUnique(this);
392 InSequence enforceOrder;
393
394 cm_->onActivated(*conns_.front());
395 for (const auto& conn : conns_) {
396 EXPECT_CALL(*conn, closeWhenIdle());
397 }
398 cm_->initiateGracefulShutdown(std::chrono::milliseconds(0));
399 eventBase_.loopOnce();
400
401 conns_.insert(conns_.begin(), std::move(extraConn));
402 EXPECT_CALL(*conns_.front().get(), closeWhenIdle());
403 cm_->addConnection(conns_.front().get());
404 eventBase_.loop();
405 }
406
testAddDuringCloseWhenIdle(bool deactivate)407 void ConnectionManagerTest::testAddDuringCloseWhenIdle(bool deactivate) {
408 auto extraConn = MockConnection::makeUnique(this);
409 InSequence enforceOrder;
410
411 // All conns will get closeWhenIdle
412 for (const auto& conn : conns_) {
413 conn->setIdle(true);
414 EXPECT_CALL(*conn, closeWhenIdle());
415 }
416 cm_->initiateGracefulShutdown(std::chrono::milliseconds(0));
417 // Add the extra conn in this state
418 extraConn->setIdle(true);
419 conns_.insert(conns_.begin(), std::move(extraConn));
420 cm_->addConnection(conns_.front().get());
421 // Shouldn't be deleted yet, call is delayed
422 ASSERT_TRUE(conns_.front().get() != nullptr);
423
424 // Mark the connection as active
425 conns_.front()->setIdle(false);
426 if (deactivate) {
427 // mark it idle and move to the end of the list. The regular
428 // drainAllConnections code will find it and call closeWhenIdle. The
429 // second loop callback won't find the conn and be a no-op
430 cm_->onDeactivated(*conns_.front());
431 conns_.front()->setIdle(true);
432 }
433 EXPECT_CALL(*conns_.front(), closeWhenIdle());
434 eventBase_.loop();
435 if (!deactivate) {
436 // drainAllConnections didn't find it, closeWhenIdle was invoked by the
437 // second loop callback.
438 cm_->onDeactivated(*conns_.front());
439 conns_.front()->setIdle(true);
440 }
441 ASSERT_TRUE(conns_.front().get() == nullptr);
442 }
443
TEST_F(ConnectionManagerTest,testAddDuringCloseWhenIdleActive)444 TEST_F(ConnectionManagerTest, testAddDuringCloseWhenIdleActive) {
445 testAddDuringCloseWhenIdle(false);
446 }
447
TEST_F(ConnectionManagerTest,testAddDuringCloseWhenIdleInactive)448 TEST_F(ConnectionManagerTest, testAddDuringCloseWhenIdleInactive) {
449 testAddDuringCloseWhenIdle(true);
450 }
451 } // namespace
452