[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.
+  /// However, Action will be scheduled to run after the first rebuild of the
+  /// preamble for the corresponding file finishes. Note that the rebuild can
+  /// still fail, in which case Action will get a null pointer instead of the
+  /// preamble.
   /// 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()});
   };