[WIP][clangd] Dont block for preamble builds

[kadircet]

Summary

Build preambles fully asynchronously.

[kadircet]

This requires patched-up ASTs and getting rid of consistent preamble needs first.

[kadircet]

• Rebase

clang-tools-extra/clangd/TUScheduler.cpp

PreambleThread(llvm::StringRef FileName, ParsingCallbacks &Callbacks,
         StoreInMemory(StorePreambleInMemory), RunSync(RunSync), Status(Status),
         ASTPeer(AW) {}
 
   /// It isn't guaranteed that each requested version will be built. If there
   /// are multiple update requests while building a preamble, only the last one
   /// will be built.
   void update(std::unique_ptr<CompilerInvocation> CI, ParseInputs PI,
               std::vector<Diag> CIDiags, WantDiagnostics WantDiags) {
+    dlog("Preamble worker queueing {0}", PI.Version);
     // Make possibly expensive copy while not holding the lock.
     Request Req = {std::move(CI), std::move(PI), std::move(CIDiags), WantDiags,
                    Context::current().clone()};
     if (RunSync) {
       build(std::move(Req));
       Status.update([&](TUStatus &Status) {
         Status.PreambleActivity = PreambleAction::Idle;
       });
       return;
     }
     {
-      std::lock_guard<std::mutex> Lock(Mutex);
+      std::unique_lock<std::mutex> Lock(Mutex);
       // If shutdown is issued, don't bother building.
       if (Done)
         return;
+      ReqCV.wait(Lock, [this] {
+        return !NextReq || NextReq->WantDiags != WantDiagnostics::Yes;
+      });
+      dlog("preamble worker queued {0}", Req.Inputs.Version);
       NextReq = std::move(Req);
     }
     // Let the worker thread know there's a request, notify_one is safe as there
     // should be a single worker thread waiting on it.
     ReqCV.notify_all();
   }
 
   void run() {
dlog("Preamble worker for {0} received stop", FileName);
       Done = true;
       NextReq.reset();
     }
     // Let the worker thread know that it should stop.
     ReqCV.notify_all();
   }
 
   bool blockUntilIdle(Deadline Timeout) const {
+    dlog("blocking for preamble worker");
     std::unique_lock<std::mutex> Lock(Mutex);
     return wait(Lock, ReqCV, Timeout, [&] { return !NextReq && !CurrentReq; });
   }
 
 private:
   /// Holds inputs required for building a preamble. CI is guaranteed to be
   /// non-null.
   struct Request {

 std::shared_ptr<const PreambleData>
 ASTWorker::getPossiblyStalePreamble() const {
   std::lock_guard<std::mutex> Lock(Mutex);
   return LatestPreamble;
 }
 
 void ASTWorker::getCurrentPreamble(
     llvm::unique_function<void(std::shared_ptr<const PreambleData>)> Callback) {
+  if (RunSync)
+    return Callback(getPossiblyStalePreamble());
   // We could just call startTask() to throw the read on the queue, knowing
   // it will run after any updates. But we know this task is cheap, so to
   // improve latency we cheat: insert it on the queue after the last update.
   std::unique_lock<std::mutex> Lock(Mutex);
   auto LastUpdate =
       std::find_if(Requests.rbegin(), Requests.rend(),
                    [](const Request &R) { return R.UpdateType.hasValue(); });
-  // If there were no writes in the queue, and CurrentRequest is not a write,
-  // the preamble is ready now.
-  if (LastUpdate == Requests.rend() &&
-      (!CurrentRequest || CurrentRequest->UpdateType.hasValue())) {
-    Lock.unlock();
-    return Callback(getPossiblyStalePreamble());
-  }
-  assert(!RunSync && "Running synchronously, but queue is non-empty!");
-  Requests.insert(LastUpdate.base(),
-                  Request{[Callback = std::move(Callback), this]() mutable {
-                            Callback(getPossiblyStalePreamble());
-                          },
-                          "GetPreamble", steady_clock::now(),
-                          Context::current().clone(),
-                          /*UpdateType=*/None,
-                          /*InvalidationPolicy=*/TUScheduler::NoInvalidation,
-                          /*Invalidate=*/nullptr});
+  // FIXME: We should only serve stale preambles or build it here instead.
+  Requests.insert(
+      LastUpdate.base(),
+      Request{
+          [Callback = std::move(Callback), this]() mutable {
+            PreamblePeer.blockUntilIdle(Deadline::infinity());
+            PreambleRequests.push(
+                {[Callback = std::move(Callback), this]() mutable {
+                 "GetPreamble", steady_clock::now(), Context::current().clone(),
+                 None, TUScheduler::NoInvalidation, nullptr});
+          },
+          "WaitForPreamble", steady_clock::now(), Context::current().clone(),
   Lock.unlock();
   RequestsCV.notify_all();
 }
 
 void ASTWorker::waitForFirstPreamble() const { BuiltFirstPreamble.wait(); }
 
 std::shared_ptr<const ParseInputs> ASTWorker::getCurrentFileInputs() const {
   std::unique_lock<std::mutex> Lock(Mutex);
void ASTWorker::stop() {
   Status.stop();
   RequestsCV.notify_all();
 }
 
 void ASTWorker::startTask(llvm::StringRef Name,
                           llvm::unique_function<void()> Task,
                           llvm::Optional<WantDiagnostics> UpdateType,
                           TUScheduler::ASTActionInvalidation Invalidation) {
+  dlog("ASTWorker queueing {0}", Name);
   if (RunSync) {
     assert(!Done && "running a task after stop()");
     trace::Span Tracer(Name + ":" + llvm::sys::path::filename(FileName));
     Task();
     return;
   }
 
   {
for (; Next != Requests.end(); ++Next)
         return true; // Prefer later diagnostics.
     return false;
   }
   llvm_unreachable("Unknown WantDiagnostics");
 }
 
 bool ASTWorker::blockUntilIdle(Deadline Timeout) const {
   std::unique_lock<std::mutex> Lock(Mutex);
-  return wait(Lock, RequestsCV, Timeout, [&] {
-    return PreambleRequests.empty() && Requests.empty() && !CurrentRequest;
-  });
+  // Make sure ASTWorker has processed all requests.
+  if (!wait(Lock, RequestsCV, Timeout,
+            [&] { return Requests.empty() && !CurrentRequest; }))
+    return false;
+  Lock.unlock();
+  if (!PreamblePeer.blockUntilIdle(Timeout))
+    return false;
+  Lock.lock();
+  assert(Requests.empty() && "Received new requests during blockUntilIdle");
+  // Make sure ASTWorker has processed all preambles.
+  return wait(Lock, RequestsCV, Timeout,
+              [&] { return PreambleRequests.empty() && !CurrentRequest; });
 }
 
 // Render a TUAction to a user-facing string representation.
 // TUAction represents clangd-internal states, we don't intend to expose them
 // to users (say C++ programmers) directly to avoid confusion, we use terms that
 // are familiar by C++ programmers.
 std::string renderTUAction(const PreambleAction PA, const ASTAction &AA) {
   llvm::SmallVector<std::string, 2> Result;

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

 )cpp";
 
   auto FooCpp = testPath("foo.cpp");
 
   FS.Files[testPath("foo.h")] = "int a;";
   FS.Files[FooCpp] = SourceContents;
 
   Server.addDocument(FooCpp, SourceContents);
-  auto DumpParse1 = dumpASTWithoutMemoryLocs(Server, FooCpp);
   ASSERT_TRUE(Server.blockUntilIdleForTest()) << "Waiting for diagnostics";
+  auto DumpParse1 = dumpASTWithoutMemoryLocs(Server, FooCpp);
   EXPECT_FALSE(DiagConsumer.hadErrorInLastDiags());
 
   Server.addDocument(FooCpp, "");
-  auto DumpParseEmpty = dumpASTWithoutMemoryLocs(Server, FooCpp);
   ASSERT_TRUE(Server.blockUntilIdleForTest()) << "Waiting for diagnostics";
+  auto DumpParseEmpty = dumpASTWithoutMemoryLocs(Server, FooCpp);
   EXPECT_FALSE(DiagConsumer.hadErrorInLastDiags());
 
   Server.addDocument(FooCpp, SourceContents);
-  auto DumpParse2 = dumpASTWithoutMemoryLocs(Server, FooCpp);
   ASSERT_TRUE(Server.blockUntilIdleForTest()) << "Waiting for diagnostics";
+  auto DumpParse2 = dumpASTWithoutMemoryLocs(Server, FooCpp);
   EXPECT_FALSE(DiagConsumer.hadErrorInLastDiags());
 
   EXPECT_EQ(DumpParse1, DumpParse2);
   EXPECT_NE(DumpParse1, DumpParseEmpty);
 }
 
 TEST_F(ClangdVFSTest, ReparseOnHeaderChange) {
   MockFSProvider FS;
   ErrorCheckingCallbacks DiagConsumer;
   MockCompilationDatabase CDB;
   ClangdServer Server(CDB, FS, ClangdServer::optsForTest(), &DiagConsumer);
 
   const auto SourceContents = R"cpp(
 #include "foo.h"
 int b = a;
 )cpp";
 
   auto FooCpp = testPath("foo.cpp");
   auto FooH = testPath("foo.h");
 
   FS.Files[FooH] = "int a;";
   FS.Files[FooCpp] = SourceContents;
 
   Server.addDocument(FooCpp, SourceContents);
-  auto DumpParse1 = dumpASTWithoutMemoryLocs(Server, FooCpp);
   ASSERT_TRUE(Server.blockUntilIdleForTest()) << "Waiting for diagnostics";
+  auto DumpParse1 = dumpASTWithoutMemoryLocs(Server, FooCpp);
   EXPECT_FALSE(DiagConsumer.hadErrorInLastDiags());
 
   FS.Files[FooH] = "";
   Server.addDocument(FooCpp, SourceContents);
-  auto DumpParseDifferent = dumpASTWithoutMemoryLocs(Server, FooCpp);
   ASSERT_TRUE(Server.blockUntilIdleForTest()) << "Waiting for diagnostics";
+  auto DumpParseDifferent = dumpASTWithoutMemoryLocs(Server, FooCpp);
   EXPECT_TRUE(DiagConsumer.hadErrorInLastDiags());
 
   FS.Files[FooH] = "int a;";
   Server.addDocument(FooCpp, SourceContents);
-  auto DumpParse2 = dumpASTWithoutMemoryLocs(Server, FooCpp);
   ASSERT_TRUE(Server.blockUntilIdleForTest()) << "Waiting for diagnostics";
+  auto DumpParse2 = dumpASTWithoutMemoryLocs(Server, FooCpp);
   EXPECT_FALSE(DiagConsumer.hadErrorInLastDiags());
 
   EXPECT_EQ(DumpParse1, DumpParse2);
   EXPECT_NE(DumpParse1, DumpParseDifferent);
 }
 
 TEST_F(ClangdVFSTest, PropagatesContexts) {
   static Key<int> Secret;

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

TEST_F(TUSchedulerTests, Cancellation) {
   std::vector<std::string> DiagsSeen, ReadsSeen, ReadsCanceled;
   {
     Notification Proceed; // Ensure we schedule everything.
     TUScheduler S(CDB, optsForTest(), captureDiags());
     auto Path = testPath("foo.cpp");
     // Helper to schedule a named update and return a function to cancel it.
     auto Update = [&](std::string ID) -> Canceler {
       auto T = cancelableTask();
+      auto Inp = getInputs(Path, "//" + ID);
+      Inp.Version = ID;
       WithContext C(std::move(T.first));
       updateWithDiags(
-          S, Path, "//" + ID, WantDiagnostics::Yes,
+          S, Path, std::move(Inp), WantDiagnostics::Yes,
           [&, ID](std::vector<Diag> Diags) { DiagsSeen.push_back(ID); });
       return std::move(T.second);
     };
     // Helper to schedule a named read and return a function to cancel it.
     auto Read = [&](std::string ID) -> Canceler {
       auto T = cancelableTask();
       WithContext C(std::move(T.first));
       S.runWithAST(ID, Path, [&, ID](llvm::Expected<InputsAndAST> E) {
for (int FileI = 0; FileI < FilesCount; ++FileI) {
         auto Contents = AllContents[(FileI + UpdateI) % AllContentsSize];
 
         auto File = Files[FileI];
         auto Inputs = getInputs(File, Contents.str());
         {
           WithContextValue WithNonce(NonceKey, ++Nonce);
           Inputs.Version = std::to_string(Nonce);
           updateWithDiags(
-              S, File, Inputs, WantDiagnostics::Auto,
+              S, File, Inputs, WantDiagnostics::Yes,
               [File, Nonce, &Mut, &TotalUpdates](std::vector<Diag>) {
                 EXPECT_THAT(Context::current().get(NonceKey), Pointee(Nonce));
 
                 std::lock_guard<std::mutex> Lock(Mut);
                 ++TotalUpdates;
                 EXPECT_EQ(File, *TUScheduler::getFileBeingProcessedInContext());
               });
         }
TEST_F(TUSchedulerTests, ForceRebuild) {
   auto Header = testPath("foo.h");
 
   auto SourceContents = R"cpp(
       #include "foo.h"
       int b = a;
     )cpp";
 
   ParseInputs Inputs = getInputs(Source, SourceContents);
-  std::atomic<size_t> DiagCount(0);
+  Inputs.Version = "MissingHeader";
+  std::atomic<int> DiagCount(0);
 
   // Update the source contents, which should trigger an initial build with
   // the header file missing.
   updateWithDiags(
       S, Source, Inputs, WantDiagnostics::Yes,
       [&DiagCount](std::vector<Diag> Diags) {
         ++DiagCount;
         EXPECT_THAT(Diags,
                     ElementsAre(Field(&Diag::Message, "'foo.h' file not found"),
                                 Field(&Diag::Message,
                                       "use of undeclared identifier 'a'")));
       });
 
   // Add the header file. We need to recreate the inputs since we changed a
   // file from underneath the test FS.
   Files[Header] = "int a;";
   Timestamps[Header] = time_t(1);
   Inputs = getInputs(Source, SourceContents);
+  Inputs.Version = "WithHeader";
 
   // The addition of the missing header file shouldn't trigger a rebuild since
   // we don't track missing files.
   updateWithDiags(
       S, Source, Inputs, WantDiagnostics::Yes,
       [&DiagCount](std::vector<Diag> Diags) {
         ++DiagCount;
         ADD_FAILURE() << "Did not expect diagnostics for missing header update";
       });
 
   // Forcing the reload should should cause a rebuild which no longer has any
   // errors.
   Inputs.ForceRebuild = true;
+  Inputs.Version = "ForceRebuild";
   updateWithDiags(S, Source, Inputs, WantDiagnostics::Yes,
                   [&DiagCount](std::vector<Diag> Diags) {
                     ++DiagCount;
                     EXPECT_THAT(Diags, IsEmpty());
                   });
 
   ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
   EXPECT_EQ(DiagCount, 2U);
Server.addDocument(testPath("foo.cpp"), Code.code(), "1",
                      WantDiagnostics::Yes);
   Server.locateSymbolAt(testPath("foo.cpp"), Code.point(),
                         [](Expected<std::vector<LocatedSymbol>> Result) {
                           ASSERT_TRUE((bool)Result);
                         });
 
   ASSERT_TRUE(Server.blockUntilIdleForTest());
 
-  EXPECT_THAT(CaptureTUStatus.allStatus(),
-              ElementsAre(
-                  // Everything starts with ASTWorker starting to execute an
-                  // update
-                  TUState(PreambleAction::Idle, ASTAction::RunningAction),
-                  // We build the preamble
-                  TUState(PreambleAction::Building, ASTAction::RunningAction),
-                  // We built the preamble, and issued ast built on ASTWorker
-                  // thread. Preambleworker goes idle afterwards.
-                  TUState(PreambleAction::Idle, ASTAction::RunningAction),
-                  // Start task for building the ast, as a result of building
-                  // preamble, on astworker thread.
-                  TUState(PreambleAction::Idle, ASTAction::RunningAction),
-                  // AST build starts.
-                  TUState(PreambleAction::Idle, ASTAction::Building),
-                  // AST built finished successfully
-                  TUState(PreambleAction::Idle, ASTAction::Building),
-                  // Running go to def
-                  TUState(PreambleAction::Idle, ASTAction::RunningAction),
-                  // ASTWorker goes idle.
-                  TUState(PreambleAction::Idle, ASTAction::Idle)));
+  auto Statuses = CaptureTUStatus.allStatus();
+  const std::vector<PreambleAction> PreambleStatuses = {
+      PreambleAction::Idle, PreambleAction::Building, PreambleAction::Idle};
+
+  llvm::ArrayRef<PreambleAction> RemainingStatuses =
+      llvm::makeArrayRef(PreambleStatuses);
+  ASSERT_FALSE(Statuses.empty());
+  for (size_t I = 0, E = Statuses.size(); I != E; ++I) {
+    if (RemainingStatuses.front() != Statuses[I].PreambleActivity)
+      RemainingStatuses = RemainingStatuses.drop_front();
+    ASSERT_FALSE(RemainingStatuses.empty());
+    EXPECT_EQ(Statuses[I].PreambleActivity, RemainingStatuses.front());
+  }
+  EXPECT_EQ(RemainingStatuses.size(), 1U);
 }
 
 TEST_F(TUSchedulerTests, CommandLineErrors) {
   // We should see errors from command-line parsing inside the main file.
   CDB.ExtraClangFlags = {"-fsome-unknown-flag"};
 
   // (!) 'Ready' must live longer than TUScheduler.
   Notification Ready;

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

TEST(LocateSymbol, WithPreamble) {
   // LocateSymbol goes to a #include file: the result comes from the preamble.
   EXPECT_THAT(
       cantFail(runLocateSymbolAt(Server, FooCpp, FooWithHeader.point())),
       ElementsAre(Sym("foo.h", FooHeader.range())));
 
   // Only preamble is built, and no AST is built in this request.
   Server.addDocument(FooCpp, FooWithoutHeader.code(), "null",
                      WantDiagnostics::No);
+  ASSERT_TRUE(Server.blockUntilIdleForTest());
   // We build AST here, and it should use the latest preamble rather than the
   // stale one.
   EXPECT_THAT(
       cantFail(runLocateSymbolAt(Server, FooCpp, FooWithoutHeader.point())),
       ElementsAre(Sym("foo", FooWithoutHeader.range())));
 
   // Reset test environment.
   runAddDocument(Server, FooCpp, FooWithHeader.code());
   // Both preamble and AST are built in this request.
   Server.addDocument(FooCpp, FooWithoutHeader.code(), "null",
                      WantDiagnostics::Yes);
+  ASSERT_TRUE(Server.blockUntilIdleForTest());
   // Use the AST being built in above request.
   EXPECT_THAT(
       cantFail(runLocateSymbolAt(Server, FooCpp, FooWithoutHeader.point())),
       ElementsAre(Sym("foo", FooWithoutHeader.range())));
 }
 
 TEST(FindReferences, WithinAST) {
   const char *Tests[] = {