[RFC][clangd] Move preamble index out of document open critical path

We would like to move the preamble index out of the critical path.
This patch is an RFC to get feedback on the correct implementation and potential pitfalls to keep into consideration.

I am not entirely sure if the lazy AST initialisation would create using Preamble AST in parallel. I tried with tsan enabled clangd but it seems to work OK (at least for the cases I tried)

Reviewed By: kadircet

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

clang-tools-extra/clangd/ClangdServer.cpp

@@ -62,22 +62,38 @@ struct UpdateIndexCallbacks : public ParsingCallbacks {
   UpdateIndexCallbacks(FileIndex *FIndex,
                        ClangdServer::Callbacks *ServerCallbacks,
                        const ThreadsafeFS &TFS, AsyncTaskRunner *Tasks,
-                       bool CollectInactiveRegions)
-      : FIndex(FIndex), ServerCallbacks(ServerCallbacks), TFS(TFS),
-        Stdlib{std::make_shared<StdLibSet>()}, Tasks(Tasks),
-        CollectInactiveRegions(CollectInactiveRegions) {}
-
-  void onPreambleAST(PathRef Path, llvm::StringRef Version,
-                     const CompilerInvocation &CI, ASTContext &Ctx,
-                     Preprocessor &PP,
-                     const CanonicalIncludes &CanonIncludes) override {
-    // If this preamble uses a standard library we haven't seen yet, index it.
-    if (FIndex)
-      if (auto Loc = Stdlib->add(*CI.getLangOpts(), PP.getHeaderSearchInfo()))
-        indexStdlib(CI, std::move(*Loc));
+                       bool CollectInactiveRegions,
+                       const ClangdServer::Options &Opts)
+      : FIndex(FIndex), ServerCallbacks(ServerCallbacks),
+        TFS(TFS), Stdlib{std::make_shared<StdLibSet>()}, Tasks(Tasks),
+        CollectInactiveRegions(CollectInactiveRegions), Opts(Opts) {}
+
+  void onPreambleAST(
+      PathRef Path, llvm::StringRef Version, CapturedASTCtx ASTCtx,
+      const std::shared_ptr<const CanonicalIncludes> CanonIncludes) override {
+
+    if (!FIndex)
+      return;
+
+    auto &PP = ASTCtx.getPreprocessor();
+    auto &CI = ASTCtx.getCompilerInvocation();
+    if (auto Loc = Stdlib->add(*CI.getLangOpts(), PP.getHeaderSearchInfo()))
+      indexStdlib(CI, std::move(*Loc));
+
+    // FIndex outlives the UpdateIndexCallbacks.
+    auto Task = [FIndex(FIndex), Path(Path.str()), Version(Version.str()),
+                 ASTCtx(std::move(ASTCtx)),
+                 CanonIncludes(CanonIncludes)]() mutable {
+      trace::Span Tracer("PreambleIndexing");
+      FIndex->updatePreamble(Path, Version, ASTCtx.getASTContext(),
+                             ASTCtx.getPreprocessor(), *CanonIncludes);
+    };
 
-    if (FIndex)
-      FIndex->updatePreamble(Path, Version, Ctx, PP, CanonIncludes);
+    if (Opts.AsyncPreambleIndexing && Tasks) {
+      Tasks->runAsync("Preamble indexing for:" + Path + Version,
+                      std::move(Task));
+    } else
+      Task();
   }
 
   void indexStdlib(const CompilerInvocation &CI, StdLibLocation Loc) {
@@ -146,6 +162,7 @@ struct UpdateIndexCallbacks : public ParsingCallbacks {
   std::shared_ptr<StdLibSet> Stdlib;
   AsyncTaskRunner *Tasks;
   bool CollectInactiveRegions;
+  const ClangdServer::Options &Opts;
 };
 
 class DraftStoreFS : public ThreadsafeFS {
@@ -210,7 +227,7 @@ ClangdServer::ClangdServer(const GlobalCompilationDatabase &CDB,
                         std::make_unique<UpdateIndexCallbacks>(
                             DynamicIdx.get(), Callbacks, TFS,
                             IndexTasks ? &*IndexTasks : nullptr,
-                            PublishInactiveRegions));
+                            PublishInactiveRegions, Opts));
   // Adds an index to the stack, at higher priority than existing indexes.
   auto AddIndex = [&](SymbolIndex *Idx) {
     if (this->Index != nullptr) {

clang-tools-extra/clangd/ClangdServer.h

@@ -185,6 +185,10 @@ class ClangdServer {
     /// regions in the document.
     bool PublishInactiveRegions = false;
 
+    /// Whether to run preamble indexing asynchronously in an independent
+    /// thread.
+    bool AsyncPreambleIndexing = false;
+
     explicit operator TUScheduler::Options() const;
   };
   // Sensible default options for use in tests.

clang-tools-extra/clangd/ParsedAST.cpp

@@ -666,7 +666,7 @@ ParsedAST::build(llvm::StringRef Filename, const ParseInputs &Inputs,
   // non-preamble includes below.
   CanonicalIncludes CanonIncludes;
   if (Preamble)
-    CanonIncludes = Preamble->CanonIncludes;
+    CanonIncludes = *Preamble->CanonIncludes;
   else
     CanonIncludes.addSystemHeadersMapping(Clang->getLangOpts());
   std::unique_ptr<CommentHandler> IWYUHandler =

clang-tools-extra/clangd/Preamble.cpp

@@ -26,7 +26,9 @@
 #include "clang/Basic/LangOptions.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
+#include "clang/Basic/TargetInfo.h"
 #include "clang/Basic/TokenKinds.h"
+#include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/FrontendActions.h"
 #include "clang/Frontend/PrecompiledPreamble.h"
@@ -35,6 +37,7 @@
 #include "clang/Lex/PPCallbacks.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Lex/PreprocessorOptions.h"
+#include "clang/Serialization/ASTReader.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
@@ -77,10 +80,9 @@ bool compileCommandsAreEqual(const tooling::CompileCommand &LHS,
 class CppFilePreambleCallbacks : public PreambleCallbacks {
 public:
   CppFilePreambleCallbacks(
-      PathRef File, PreambleParsedCallback ParsedCallback,
-      PreambleBuildStats *Stats, bool ParseForwardingFunctions,
+      PathRef File, PreambleBuildStats *Stats, bool ParseForwardingFunctions,
       std::function<void(CompilerInstance &)> BeforeExecuteCallback)
-      : File(File), ParsedCallback(ParsedCallback), Stats(Stats),
+      : File(File), Stats(Stats),
         ParseForwardingFunctions(ParseForwardingFunctions),
         BeforeExecuteCallback(std::move(BeforeExecuteCallback)) {}
 
@@ -95,13 +97,22 @@ class CppFilePreambleCallbacks : public PreambleCallbacks {
   }
   CanonicalIncludes takeCanonicalIncludes() { return std::move(CanonIncludes); }
 
+  std::optional<CapturedASTCtx> takeLife() { return std::move(CapturedCtx); }
+
   bool isMainFileIncludeGuarded() const { return IsMainFileIncludeGuarded; }
 
   void AfterExecute(CompilerInstance &CI) override {
-    if (ParsedCallback) {
-      trace::Span Tracer("Running PreambleCallback");
-      ParsedCallback(CI.getASTContext(), CI.getPreprocessor(), CanonIncludes);
+    // As part of the Preamble compilation, ASTConsumer
+    // PrecompilePreambleConsumer/PCHGenerator is setup. This would be called
+    // when Preamble consists of modules. Therefore while capturing AST context,
+    // we have to reset ast consumer and ASTMutationListener.
+    if (CI.getASTReader()) {
+      CI.getASTReader()->setDeserializationListener(nullptr);
+      // This just sets consumer to null when DeserializationListener is null.
+      CI.getASTReader()->StartTranslationUnit(nullptr);
     }
+    CI.getASTContext().setASTMutationListener(nullptr);
+    CapturedCtx.emplace(CI);
 
     const SourceManager &SM = CI.getSourceManager();
     const FileEntry *MainFE = SM.getFileEntryForID(SM.getMainFileID());
@@ -202,7 +213,6 @@ class CppFilePreambleCallbacks : public PreambleCallbacks {
 
 private:
   PathRef File;
-  PreambleParsedCallback ParsedCallback;
   IncludeStructure Includes;
   CanonicalIncludes CanonIncludes;
   include_cleaner::PragmaIncludes Pragmas;
@@ -216,6 +226,7 @@ class CppFilePreambleCallbacks : public PreambleCallbacks {
   PreambleBuildStats *Stats;
   bool ParseForwardingFunctions;
   std::function<void(CompilerInstance &)> BeforeExecuteCallback;
+  std::optional<CapturedASTCtx> CapturedCtx;
 };
 
 // Represents directives other than includes, where basic textual information is
@@ -635,16 +646,15 @@ buildPreamble(PathRef FileName, CompilerInvocation CI,
   CI.getPreprocessorOpts().WriteCommentListToPCH = false;
 
   CppFilePreambleCallbacks CapturedInfo(
-      FileName, PreambleCallback, Stats,
-      Inputs.Opts.PreambleParseForwardingFunctions,
+      FileName, Stats, Inputs.Opts.PreambleParseForwardingFunctions,
       [&ASTListeners](CompilerInstance &CI) {
         for (const auto &L : ASTListeners)
           L->beforeExecute(CI);
       });
   auto VFS = Inputs.TFS->view(Inputs.CompileCommand.Directory);
   llvm::SmallString<32> AbsFileName(FileName);
   VFS->makeAbsolute(AbsFileName);
-  auto StatCache = std::make_unique<PreambleFileStatusCache>(AbsFileName);
+  auto StatCache = std::make_shared<PreambleFileStatusCache>(AbsFileName);
   auto StatCacheFS = StatCache->getProducingFS(VFS);
   llvm::IntrusiveRefCntPtr<TimerFS> TimedFS(new TimerFS(StatCacheFS));
 
@@ -679,9 +689,22 @@ buildPreamble(PathRef FileName, CompilerInvocation CI,
     Result->Pragmas = CapturedInfo.takePragmaIncludes();
     Result->Macros = CapturedInfo.takeMacros();
     Result->Marks = CapturedInfo.takeMarks();
-    Result->CanonIncludes = CapturedInfo.takeCanonicalIncludes();
-    Result->StatCache = std::move(StatCache);
+    Result->CanonIncludes = std::make_shared<const CanonicalIncludes>(
+        (CapturedInfo.takeCanonicalIncludes()));
+    Result->StatCache = StatCache;
     Result->MainIsIncludeGuarded = CapturedInfo.isMainFileIncludeGuarded();
+    if (PreambleCallback) {
+      trace::Span Tracer("Running PreambleCallback");
+      auto Ctx = CapturedInfo.takeLife();
+      // Stat cache is thread safe only when there are no producers. Hence
+      // change the VFS underneath to a consuming fs.
+      Ctx->getFileManager().setVirtualFileSystem(
+          Result->StatCache->getConsumingFS(VFS));
+      // While extending the life of FileMgr and VFS, StatCache should also be
+      // extended.
+      Ctx->setStatCache(Result->StatCache);
+      PreambleCallback(std::move(*Ctx), Result->CanonIncludes);
+    }
     return Result;
   }
 

clang-tools-extra/clangd/Preamble.h

@@ -48,6 +48,46 @@
 namespace clang {
 namespace clangd {
 
+/// The captured AST conext.
+/// Keeps necessary structs for an ASTContext and Preprocessor alive.
+/// This enables consuming them after context that produced the AST is gone.
+/// (e.g. indexing a preamble ast on a separate thread). ASTContext stored
+/// inside is still not thread-safe.
+
+struct CapturedASTCtx {
+public:
+  CapturedASTCtx(CompilerInstance &Clang)
+      : Invocation(Clang.getInvocationPtr()),
+        Diagnostics(Clang.getDiagnosticsPtr()), Target(Clang.getTargetPtr()),
+        AuxTarget(Clang.getAuxTarget()), FileMgr(Clang.getFileManagerPtr()),
+        SourceMgr(Clang.getSourceManagerPtr()), PP(Clang.getPreprocessorPtr()),
+        Context(Clang.getASTContextPtr()) {}
+
+  CapturedASTCtx(const CapturedASTCtx &) = delete;
+  CapturedASTCtx &operator=(const CapturedASTCtx &) = delete;
+  CapturedASTCtx(CapturedASTCtx &&) = default;
+  CapturedASTCtx &operator=(CapturedASTCtx &&) = default;
+
+  ASTContext &getASTContext() { return *Context; }
+  Preprocessor &getPreprocessor() { return *PP; }
+  CompilerInvocation &getCompilerInvocation() { return *Invocation; }
+  FileManager &getFileManager() { return *FileMgr; }
+  void setStatCache(std::shared_ptr<PreambleFileStatusCache> StatCache) {
+    this->StatCache = StatCache;
+  }
+
+private:
+  std::shared_ptr<CompilerInvocation> Invocation;
+  IntrusiveRefCntPtr<DiagnosticsEngine> Diagnostics;
+  IntrusiveRefCntPtr<TargetInfo> Target;
+  IntrusiveRefCntPtr<TargetInfo> AuxTarget;
+  IntrusiveRefCntPtr<FileManager> FileMgr;
+  IntrusiveRefCntPtr<SourceManager> SourceMgr;
+  std::shared_ptr<Preprocessor> PP;
+  IntrusiveRefCntPtr<ASTContext> Context;
+  std::shared_ptr<PreambleFileStatusCache> StatCache;
+};
+
 /// The parsed preamble and associated data.
 ///
 /// As we must avoid re-parsing the preamble, any information that can only
@@ -73,15 +113,16 @@ struct PreambleData {
   std::vector<PragmaMark> Marks;
   // Cache of FS operations performed when building the preamble.
   // When reusing a preamble, this cache can be consumed to save IO.
-  std::unique_ptr<PreambleFileStatusCache> StatCache;
-  CanonicalIncludes CanonIncludes;
+  std::shared_ptr<PreambleFileStatusCache> StatCache;
+  std::shared_ptr<const CanonicalIncludes> CanonIncludes;
   // Whether there was a (possibly-incomplete) include-guard on the main file.
   // We need to propagate this information "by hand" to subsequent parses.
   bool MainIsIncludeGuarded = false;
 };
 
-using PreambleParsedCallback = std::function<void(ASTContext &, Preprocessor &,
-                                                  const CanonicalIncludes &)>;
+using PreambleParsedCallback =
+    std::function<void(CapturedASTCtx ASTCtx,
+                       std::shared_ptr<const CanonicalIncludes> CanonIncludes)>;
 
 /// Timings and statistics from the premble build. Unlike PreambleData, these
 /// do not need to be stored for later, but can be useful for logging, metrics,

clang-tools-extra/clangd/TUScheduler.cpp

@@ -1080,9 +1080,9 @@ void PreambleThread::build(Request Req) {
   bool IsFirstPreamble = !LatestBuild;
   LatestBuild = clang::clangd::buildPreamble(
       FileName, *Req.CI, Inputs, StoreInMemory,
-      [&](ASTContext &Ctx, Preprocessor &PP,
-          const CanonicalIncludes &CanonIncludes) {
-        Callbacks.onPreambleAST(FileName, Inputs.Version, *Req.CI, Ctx, PP,
+      [&](CapturedASTCtx ASTCtx,
+          std::shared_ptr<const CanonicalIncludes> CanonIncludes) {
+        Callbacks.onPreambleAST(FileName, Inputs.Version, std::move(ASTCtx),
                                 CanonIncludes);
       },
       &Stats);

clang-tools-extra/clangd/TUScheduler.h

@@ -162,8 +162,8 @@ class ParsingCallbacks {
   /// contains only AST nodes from the #include directives at the start of the
   /// file. AST node in the current file should be observed on onMainAST call.
   virtual void onPreambleAST(PathRef Path, llvm::StringRef Version,
-                             const CompilerInvocation &CI, ASTContext &Ctx,
-                             Preprocessor &PP, const CanonicalIncludes &) {}
+                             CapturedASTCtx Ctx,
+                             const std::shared_ptr<const CanonicalIncludes>) {}
 
   /// The argument function is run under the critical section guarding against
   /// races when closing the files.

clang-tools-extra/clangd/tool/Check.cpp

@@ -228,15 +228,16 @@ class Checker {
   // Build preamble and AST, and index them.
   bool buildAST() {
     log("Building preamble...");
-    Preamble = buildPreamble(File, *Invocation, Inputs, /*StoreInMemory=*/true,
-                             [&](ASTContext &Ctx, Preprocessor &PP,
-                                 const CanonicalIncludes &Includes) {
-                               if (!Opts.BuildDynamicSymbolIndex)
-                                 return;
-                               log("Indexing headers...");
-                               Index.updatePreamble(File, /*Version=*/"null",
-                                                    Ctx, PP, Includes);
-                             });
+    Preamble = buildPreamble(
+        File, *Invocation, Inputs, /*StoreInMemory=*/true,
+        [&](CapturedASTCtx Ctx,
+            const std::shared_ptr<const CanonicalIncludes> Includes) {
+          if (!Opts.BuildDynamicSymbolIndex)
+            return;
+          log("Indexing headers...");
+          Index.updatePreamble(File, /*Version=*/"null", Ctx.getASTContext(),
+                               Ctx.getPreprocessor(), *Includes);
+        });
     if (!Preamble) {
       elog("Failed to build preamble");
       return false;

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

@@ -305,16 +305,18 @@ TEST(FileIndexTest, RebuildWithPreamble) {
 
   FileIndex Index;
   bool IndexUpdated = false;
-  buildPreamble(FooCpp, *CI, PI,
-                /*StoreInMemory=*/true,
-                [&](ASTContext &Ctx, Preprocessor &PP,
-                    const CanonicalIncludes &CanonIncludes) {
-                  EXPECT_FALSE(IndexUpdated)
-                      << "Expected only a single index update";
-                  IndexUpdated = true;
-                  Index.updatePreamble(FooCpp, /*Version=*/"null", Ctx, PP,
-                                       CanonIncludes);
-                });
+  buildPreamble(
+      FooCpp, *CI, PI,
+      /*StoreInMemory=*/true,
+      [&](CapturedASTCtx ASTCtx,
+          const std::shared_ptr<const CanonicalIncludes> CanonIncludes) {
+        auto &Ctx = ASTCtx.getASTContext();
+        auto &PP = ASTCtx.getPreprocessor();
+        EXPECT_FALSE(IndexUpdated) << "Expected only a single index update";
+        IndexUpdated = true;
+        Index.updatePreamble(FooCpp, /*Version=*/"null", Ctx, PP,
+                             *CanonIncludes);
+      });
   ASSERT_TRUE(IndexUpdated);
 
   // Check the index contains symbols from the preamble, but not from the main

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

@@ -1129,9 +1129,9 @@ TEST_F(TUSchedulerTests, AsyncPreambleThread) {
   public:
     BlockPreambleThread(llvm::StringRef BlockVersion, Notification &N)
         : BlockVersion(BlockVersion), N(N) {}
-    void onPreambleAST(PathRef Path, llvm::StringRef Version,
-                       const CompilerInvocation &, ASTContext &Ctx,
-                       Preprocessor &, const CanonicalIncludes &) override {
+    void
+    onPreambleAST(PathRef Path, llvm::StringRef Version, CapturedASTCtx,
+                  const std::shared_ptr<const CanonicalIncludes>) override {
       if (Version == BlockVersion)
         N.wait();
     }
@@ -1208,9 +1208,9 @@ TEST_F(TUSchedulerTests, PublishWithStalePreamble) {
     BlockPreambleThread(Notification &UnblockPreamble, DiagsCB CB)
         : UnblockPreamble(UnblockPreamble), CB(std::move(CB)) {}
 
-    void onPreambleAST(PathRef Path, llvm::StringRef Version,
-                       const CompilerInvocation &, ASTContext &Ctx,
-                       Preprocessor &, const CanonicalIncludes &) override {
+    void
+    onPreambleAST(PathRef Path, llvm::StringRef Version, CapturedASTCtx,
+                  const std::shared_ptr<const CanonicalIncludes>) override {
       if (BuildBefore)
         ASSERT_TRUE(UnblockPreamble.wait(timeoutSeconds(5)))
             << "Expected notification";
@@ -1562,9 +1562,9 @@ TEST_F(TUSchedulerTests, PreambleThrottle) {
     std::vector<std::string> &Filenames;
     CaptureBuiltFilenames(std::vector<std::string> &Filenames)
         : Filenames(Filenames) {}
-    void onPreambleAST(PathRef Path, llvm::StringRef Version,
-                       const CompilerInvocation &CI, ASTContext &Ctx,
-                       Preprocessor &PP, const CanonicalIncludes &) override {
+    void
+    onPreambleAST(PathRef Path, llvm::StringRef Version, CapturedASTCtx,
+                  const std::shared_ptr<const CanonicalIncludes>) override {
       // Deliberately no synchronization.
       // The PreambleThrottler should serialize these calls, if not then tsan
       // will find a bug here.