diff --git a/mcrouter/test/cpp_unit_tests/McrouterClientUsage.cpp b/mcrouter/test/cpp_unit_tests/McrouterClientUsage.cpp index 1384b19e5..278cfe2b7 100644 --- a/mcrouter/test/cpp_unit_tests/McrouterClientUsage.cpp +++ b/mcrouter/test/cpp_unit_tests/McrouterClientUsage.cpp @@ -97,7 +97,9 @@ TEST(CarbonRouterClient, basicUsageSameThreadClient) { client->setProxyIndex(0); bool replyReceived = false; - eventBase.runInEventBaseThread([client = client.get(), &replyReceived]() { + folly::fibers::Baton baton; + + eventBase.runInEventBaseThread([client = client.get(), &replyReceived, &baton]() { // We must ensure that req will remain alive all the way through the reply // callback given to client->send(). This demonstrates one way of ensuring // this. @@ -105,18 +107,19 @@ TEST(CarbonRouterClient, basicUsageSameThreadClient) { auto reqRawPtr = req.get(); client->send( *reqRawPtr, - [req = std::move(req), &replyReceived]( + [req = std::move(req), &replyReceived, &baton]( const McGetRequest&, McGetReply&& reply) { EXPECT_EQ(carbon::Result::NOTFOUND, *reply.result_ref()); replyReceived = true; + baton.post(); }); }); // Wait for proxy threads to complete outstanding requests and exit // gracefully. This ensures graceful destruction of the static // CarbonRouterInstance instance. + baton.wait(); router->shutdown(); - ioThreadPool->join(); ioThreadPool.reset(); EXPECT_TRUE(replyReceived); }