[clangd] Fix raciness in code completion tests

Reviewers: sammccall, ilya-biryukov

Subscribers: MaskRay, jkorous, arphaman, usaxena95, cfe-commits

Tags: #clang

Differential Revision: https://reviews.llvm.org/D68273

llvm-svn: 373924

clang-tools-extra/clangd/unittests/CodeCompleteTests.cpp

@@ -18,6 +18,7 @@
 #include "TestFS.h"
 #include "TestIndex.h"
 #include "TestTU.h"
+#include "Threading.h"
 #include "index/Index.h"
 #include "index/MemIndex.h"
 #include "clang/Sema/CodeCompleteConsumer.h"
@@ -27,6 +28,8 @@
 #include "llvm/Testing/Support/Error.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include <condition_variable>
+#include <mutex>
 
 namespace clang {
 namespace clangd {
@@ -1112,8 +1115,9 @@ class IndexRequestCollector : public SymbolIndex {
   bool
   fuzzyFind(const FuzzyFindRequest &Req,
             llvm::function_ref<void(const Symbol &)> Callback) const override {
-    std::lock_guard<std::mutex> Lock(Mut);
+    std::unique_lock<std::mutex> Lock(Mut);
     Requests.push_back(Req);
+    ReceivedRequestCV.notify_one();
     return true;
   }
 
@@ -1131,25 +1135,32 @@ class IndexRequestCollector : public SymbolIndex {
   // isn't used in production code.
   size_t estimateMemoryUsage() const override { return 0; }
 
-  const std::vector<FuzzyFindRequest> consumeRequests() const {
-    std::lock_guard<std::mutex> Lock(Mut);
+  const std::vector<FuzzyFindRequest> consumeRequests(size_t Num) const {
+    std::unique_lock<std::mutex> Lock(Mut);
+    EXPECT_TRUE(wait(Lock, ReceivedRequestCV, timeoutSeconds(10),
+                     [this, Num] { return Requests.size() == Num; }));
     auto Reqs = std::move(Requests);
     Requests = {};
     return Reqs;
   }
 
 private:
   // We need a mutex to handle async fuzzy find requests.
+  mutable std::condition_variable ReceivedRequestCV;
   mutable std::mutex Mut;
   mutable std::vector<FuzzyFindRequest> Requests;
 };
 
-std::vector<FuzzyFindRequest> captureIndexRequests(llvm::StringRef Code) {
+// Clients have to consume exactly Num requests.
+std::vector<FuzzyFindRequest> captureIndexRequests(llvm::StringRef Code,
+                                                   size_t Num = 1) {
   clangd::CodeCompleteOptions Opts;
   IndexRequestCollector Requests;
   Opts.Index = &Requests;
   completions(Code, {}, Opts);
-  return Requests.consumeRequests();
+  const auto Reqs = Requests.consumeRequests(Num);
+  EXPECT_EQ(Reqs.size(), Num);
+  return Reqs;
 }
 
 TEST(CompletionTest, UnqualifiedIdQuery) {
@@ -2098,26 +2109,23 @@ TEST(CompletionTest, EnableSpeculativeIndexRequest) {
 
   auto CompleteAtPoint = [&](StringRef P) {
     cantFail(runCodeComplete(Server, File, Test.point(P), Opts));
-    // Sleep for a while to make sure asynchronous call (if applicable) is also
-    // triggered before callback is invoked.
-    std::this_thread::sleep_for(std::chrono::milliseconds(100));
   };
 
   CompleteAtPoint("1");
-  auto Reqs1 = Requests.consumeRequests();
+  auto Reqs1 = Requests.consumeRequests(1);
   ASSERT_EQ(Reqs1.size(), 1u);
   EXPECT_THAT(Reqs1[0].Scopes, UnorderedElementsAre("ns1::"));
 
   CompleteAtPoint("2");
-  auto Reqs2 = Requests.consumeRequests();
+  auto Reqs2 = Requests.consumeRequests(1);
   // Speculation succeeded. Used speculative index result.
   ASSERT_EQ(Reqs2.size(), 1u);
   EXPECT_EQ(Reqs2[0], Reqs1[0]);
 
   CompleteAtPoint("3");
   // Speculation failed. Sent speculative index request and the new index
   // request after sema.
-  auto Reqs3 = Requests.consumeRequests();
+  auto Reqs3 = Requests.consumeRequests(2);
   ASSERT_EQ(Reqs3.size(), 2u);
 }