[clangd] Wait for first preamble before code completion

Summary:
To avoid doing extra work of processing headers in the preamble
mutilple times in parallel.

Reviewers: sammccall

Reviewed By: sammccall

Subscribers: javed.absar, ioeric, MaskRay, jkorous, cfe-commits

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

llvm-svn: 336538

3 files changed, 61 insertions, 0 deletions

diff --git a/clang-tools-extra/clangd/TUScheduler.cpp b/clang-tools-extra/clangd/TUScheduler.cpp
index 6c2918cb871..ea543e5dcca 100644
--- a/clang-tools-extra/clangd/TUScheduler.cpp
+++ b/clang-tools-extra/clangd/TUScheduler.cpp
@@ -176,6 +176,12 @@ public:
   bool blockUntilIdle(Deadline Timeout) const;
 
   std::shared_ptr<const PreambleData> getPossiblyStalePreamble() const;
+  /// Wait for the first build of preamble to finish. Preamble itself can be
+  /// accessed via getPossibleStalePreamble(). Note that this function will
+  /// return after an unsuccessful build of the preamble too, i.e. result of
+  /// getPossiblyStalePreamble() can be null even after this function returns.
+  void waitForFirstPreamble() const;
+
   std::size_t getUsedBytes() const;
   bool isASTCached() const;
 
@@ -226,6 +232,8 @@ private:
   /// Guards members used by both TUScheduler and the worker thread.
   mutable std::mutex Mutex;
   std::shared_ptr<const PreambleData> LastBuiltPreamble; /* GUARDED_BY(Mutex) */
+  /// Becomes ready when the first preamble build finishes.
+  Notification PreambleWasBuilt;
   /// Set to true to signal run() to finish processing.
   bool Done;                    /* GUARDED_BY(Mutex) */
   std::deque<Request> Requests; /* GUARDED_BY(Mutex) */
@@ -329,6 +337,9 @@ void ASTWorker::update(
         buildCompilerInvocation(Inputs);
     if (!Invocation) {
       log("Could not build CompilerInvocation for file " + FileName);
+      // Make sure anyone waiting for the preamble gets notified it could not
+      // be built.
+      PreambleWasBuilt.notify();
       return;
     }
 
@@ -340,6 +351,8 @@ void ASTWorker::update(
       if (NewPreamble)
         LastBuiltPreamble = NewPreamble;
     }
+    PreambleWasBuilt.notify();
+
     // Build the AST for diagnostics.
     llvm::Optional<ParsedAST> AST =
         buildAST(FileName, std::move(Invocation), Inputs, NewPreamble, PCHs);
@@ -392,6 +405,10 @@ ASTWorker::getPossiblyStalePreamble() const {
   return LastBuiltPreamble;
 }
 
+void ASTWorker::waitForFirstPreamble() const {
+  PreambleWasBuilt.wait();
+}
+
 std::size_t ASTWorker::getUsedBytes() const {
   // Note that we don't report the size of ASTs currently used for processing
   // the in-flight requests. We used this information for debugging purposes
@@ -655,6 +672,11 @@ void TUScheduler::runWithPreamble(
                              std::string Contents,
                              tooling::CompileCommand Command, Context Ctx,
                              decltype(Action) Action) mutable {
+    // We don't want to be running preamble actions before the preamble was
+    // built for the first time. This avoids extra work of processing the
+    // preamble headers in parallel multiple times.
+    Worker->waitForFirstPreamble();
+
     std::lock_guard<Semaphore> BarrierLock(Barrier);
     WithContext Guard(std::move(Ctx));
     trace::Span Tracer(Name);
diff --git a/clang-tools-extra/clangd/TUScheduler.h b/clang-tools-extra/clangd/TUScheduler.h
index a0c67b56e11..af1ff3656e8 100644
--- a/clang-tools-extra/clangd/TUScheduler.h
+++ b/clang-tools-extra/clangd/TUScheduler.h
@@ -101,6 +101,9 @@ public:
   /// - validate that the preamble is still valid, and only use it in this case
   /// - accept that preamble contents may be outdated, and try to avoid reading
   ///   source code from headers.
+  /// If there's no preamble yet (because the file was just opened), we'll wait
+  /// for it to build. The preamble may still be null if it fails to build or is
+  /// empty.
   /// If an error occurs during processing, it is forwarded to the \p Action
   /// callback.
   void runWithPreamble(llvm::StringRef Name, PathRef File,
diff --git a/clang-tools-extra/unittests/clangd/TUSchedulerTests.cpp b/clang-tools-extra/unittests/clangd/TUSchedulerTests.cpp
index fbf623fb06e..0826e1fa6a6 100644
--- a/clang-tools-extra/unittests/clangd/TUSchedulerTests.cpp
+++ b/clang-tools-extra/unittests/clangd/TUSchedulerTests.cpp
@@ -19,6 +19,7 @@ namespace clang {
 namespace clangd {
 
 using ::testing::_;
+using ::testing::Each;
 using ::testing::AnyOf;
 using ::testing::Pair;
 using ::testing::Pointee;
@@ -299,5 +300,40 @@ TEST_F(TUSchedulerTests, EvictedAST) {
               UnorderedElementsAre(Foo, AnyOf(Bar, Baz)));
 }
 
+TEST_F(TUSchedulerTests, RunWaitsForPreamble) {
+  // Testing strategy: we update the file and schedule a few preamble reads at
+  // the same time. All reads should get the same non-null preamble.
+  TUScheduler S(
+      /*AsyncThreadsCount=*/4, /*StorePreambleInMemory=*/true,
+      PreambleParsedCallback(),
+      /*UpdateDebounce=*/std::chrono::steady_clock::duration::zero(),
+      ASTRetentionPolicy());
+  auto Foo = testPath("foo.cpp");
+  auto NonEmptyPreamble = R"cpp(
+    #define FOO 1
+    #define BAR 2
+
+    int main() {}
+  )cpp";
+  constexpr int ReadsToSchedule = 10;
+  std::mutex PreamblesMut;
+  std::vector<const void *> Preambles(ReadsToSchedule, nullptr);
+  S.update(Foo, getInputs(Foo, NonEmptyPreamble), WantDiagnostics::Auto,
+           [](std::vector<Diag>) {});
+  for (int I = 0; I < ReadsToSchedule; ++I) {
+    S.runWithPreamble(
+        "test", Foo,
+        [I, &PreamblesMut, &Preambles](llvm::Expected<InputsAndPreamble> IP) {
+          std::lock_guard<std::mutex> Lock(PreamblesMut);
+          Preambles[I] = cantFail(std::move(IP)).Preamble;
+        });
+  }
+  ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
+  // Check all actions got the same non-null preamble.
+  std::lock_guard<std::mutex> Lock(PreamblesMut);
+  ASSERT_NE(Preambles[0], nullptr);
+  ASSERT_THAT(Preambles, Each(Preambles[0]));
+}
+
 } // namespace clangd
 } // namespace clang