[clangd] Wait for first preamble before code completion

[ilya-biryukov]

Summary

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

[sammccall]

Nice fix!
Possible test: add a file with complicated preamble (billion laughs?) and immediately schedule 5 preamble actions. They should all get a non-null preamble and the pointers should all be the same.

clangd/TUScheduler.cpp

408

inline? should fit on one line...

clangd/TUScheduler.h

104

Not sure what "however" is contrasting with.

Could be just: "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."

[ilya-biryukov]

• Add a test

[ilya-biryukov]

In D48940#1152750, @sammccall wrote:

Possible test: add a file with complicated preamble (billion laughs?) and immediately schedule 5 preamble actions. They should all get a non-null preamble and the pointers should all be the same.

The most trivial test that updates a file and immediately schedules a few preamble actions failed without this change for me. I have added it to the patch.

clangd/TUScheduler.cpp

408

This is used in TUScheduler, where we don't have access to the fields of ASTWorker.

clangd/TUScheduler.h

public:
 
   /// Schedule an async read of the Preamble.
   /// The preamble may be stale, generated from an older version of the file.
   /// Reading from locations in the preamble may cause the files to be re-read.
   /// This gives callers two options:
   /// - 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,
                        Callback<InputsAndPreamble> Action);
 
   /// Wait until there are no scheduled or running tasks.
   /// Mostly useful for synchronizing tests.
   bool blockUntilIdle(Deadline D) const;

clangd/TUScheduler.cpp

public:
   void update(ParseInputs Inputs, WantDiagnostics,
               llvm::unique_function<void(std::vector<Diag>)> OnUpdated);
   void
   runWithAST(llvm::StringRef Name,
              llvm::unique_function<void(llvm::Expected<InputsAndAST>)> Action);
   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;
 
 private:
   // Must be called exactly once on processing thread. Will return after
   // stop() is called on a separate thread and all pending requests are
   // processed.
   void run();
private:
 
   Semaphore &Barrier;
   /// Inputs, corresponding to the current state of AST.
   ParseInputs FileInputs;
   /// Size of the last AST
   /// 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) */
   mutable std::condition_variable RequestsCV;
 };
 
 /// A smart-pointer-like class that points to an active ASTWorker.
 /// In destructor, signals to the underlying ASTWorker that no new requests will
auto Task = [=](decltype(OnUpdated) OnUpdated) mutable {
     log("Updating file " + FileName + " with command [" +
         Inputs.CompileCommand.Directory + "] " +
         llvm::join(Inputs.CompileCommand.CommandLine, " "));
     // Rebuild the preamble and the AST.
     std::unique_ptr<CompilerInvocation> Invocation =
         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;
     }
 
     std::shared_ptr<const PreambleData> NewPreamble = buildPreamble(
         FileName, *Invocation, getPossiblyStalePreamble(), OldCommand, Inputs,
         PCHs, StorePreambleInMemory, PreambleCallback);
     {
       std::lock_guard<std::mutex> Lock(Mutex);
       if (NewPreamble)
         LastBuiltPreamble = NewPreamble;
     }
+    PreambleWasBuilt.notify();
+
     // Build the AST for diagnostics.
     llvm::Optional<ParsedAST> AST =
         buildAST(FileName, std::move(Invocation), Inputs, NewPreamble, PCHs);
     // We want to report the diagnostics even if this update was cancelled.
     // It seems more useful than making the clients wait indefinitely if they
     // spam us with updates.
     if (WantDiags != WantDiagnostics::No && AST)
       OnUpdated(AST->getDiagnostics());

 }
 
 std::shared_ptr<const PreambleData>
 ASTWorker::getPossiblyStalePreamble() const {
   std::lock_guard<std::mutex> Lock(Mutex);
   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
   // only, so this should be fine.
   std::size_t Result = IdleASTs.getUsedBytes(this);
   if (auto Preamble = getPossiblyStalePreamble())
     Result += Preamble->Preamble.getSize();
   return Result;
if (!PreambleTasks) {
     return;
   }
 
   std::shared_ptr<const ASTWorker> Worker = It->second->Worker.lock();
   auto Task = [Worker, this](std::string Name, std::string File,
                              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);
     SPAN_ATTACH(Tracer, "file", File);
     std::shared_ptr<const PreambleData> Preamble =
         Worker->getPossiblyStalePreamble();
     Action(InputsAndPreamble{Contents, Command, Preamble.get()});
   };

unittests/clangd/TUSchedulerTests.cpp

 #include "gtest/gtest.h"
 #include <algorithm>
 #include <utility>
 
 namespace clang {
 namespace clangd {
 
 using ::testing::_;
+using ::testing::Each;
 using ::testing::AnyOf;
 using ::testing::Pair;
 using ::testing::Pointee;
 using ::testing::UnorderedElementsAre;
 
 void ignoreUpdate(llvm::Optional<std::vector<Diag>>) {}
 void ignoreError(llvm::Error Err) {
   handleAllErrors(std::move(Err), [](const llvm::ErrorInfoBase &) {});
TEST_F(TUSchedulerTests, EvictedAST) {
   ASSERT_EQ(BuiltASTCounter.load(), 4);
 
   // Check the AST for foo.cpp is retained now and one of the others got
   // evicted.
   EXPECT_THAT(S.getFilesWithCachedAST(),
               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