[clangd] Fix crashing race in ClangdServer shutdown with stdlib indexing

[sammccall]

Summary

In principle it's OK for stdlib-indexing tasks to run after the TUScheduler is
destroyed, as mostly they just update the dynamic index. We do drain the
stdlib-indexing queue before destroying the index.

However the task captures references to the PreambleCallbacks object, which is
owned by the TUScheduler. Once this is destroyed (explicitly, early in
~ClangdServer) an outstanding stdlib-indexing task may use-after-free.

The fix here is to avoid capturing references to the PreambleCallbacks.
Alternatives would be to have TUScheduler (exclusively) not own its callbacks
so they could live longer, or explicitly stopping the TUScheduler instead of
early-destroying it. These both seem more invasive.

See https://reviews.llvm.org/D115232 for some more context.

[kadircet]

thanks, lgtm!

clang-tools-extra/clangd/ClangdServer.cpp

86

formatting seems to be off here (looks like clang-format can't deal with this leading comment. do you mind dropping it? the comment above makes it clear that we shouldn't be keeping references already.

clang-tools-extra/clangd/ClangdServer.cpp

 namespace {
 
 // Update the FileIndex with new ASTs and plumb the diagnostics responses.
 struct UpdateIndexCallbacks : public ParsingCallbacks {
   UpdateIndexCallbacks(FileIndex *FIndex,
                        ClangdServer::Callbacks *ServerCallbacks,
                        const ThreadsafeFS &TFS, AsyncTaskRunner *Tasks)
       : FIndex(FIndex), ServerCallbacks(ServerCallbacks), TFS(TFS),
-        Tasks(Tasks) {}
+        Stdlib{std::make_shared<StdLibSet>()}, Tasks(Tasks) {}
 
   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()))
+      if (auto Loc = Stdlib->add(*CI.getLangOpts(), PP.getHeaderSearchInfo()))
         indexStdlib(CI, std::move(*Loc));
 
     if (FIndex)
       FIndex->updatePreamble(Path, Version, Ctx, PP, CanonIncludes);
   }
 
   void indexStdlib(const CompilerInvocation &CI, StdLibLocation Loc) {
-    auto Task = [this, LO(*CI.getLangOpts()), Loc(std::move(Loc)),
-                 CI(std::make_unique<CompilerInvocation>(CI))]() mutable {
+    // This task is owned by Tasks, which outlives the TUScheduler and
+    // therefore the UpdateIndexCallbacks.
+    // We must be careful that the references we capture outlive TUScheduler.
+    auto Task = [LO(*CI.getLangOpts()), Loc(std::move(Loc)),
+                 CI(std::make_unique<CompilerInvocation>(CI)),
+                 // External values that outlive ClangdServer
+                 TFS(&TFS),
+                 // Index outlives TUScheduler (declared first)
+                 FIndex(FIndex),
+                 // shared_ptr extends lifetime
+                 Stdlib(Stdlib)]() mutable {
       IndexFileIn IF;
-      IF.Symbols = indexStandardLibrary(std::move(CI), Loc, TFS);
-      if (Stdlib.isBest(LO))
+      IF.Symbols = indexStandardLibrary(std::move(CI), Loc, *TFS);
+      if (Stdlib->isBest(LO))
         FIndex->updatePreamble(std::move(IF));
     };
     if (Tasks)
       // This doesn't have a semaphore to enforce -j, but it's rare.
       Tasks->runAsync("IndexStdlib", std::move(Task));
     else
       Task();
   }
void onPreamblePublished(PathRef File) override {
     if (ServerCallbacks)
       ServerCallbacks->onSemanticsMaybeChanged(File);
   }
 
 private:
   FileIndex *FIndex;
   ClangdServer::Callbacks *ServerCallbacks;
   const ThreadsafeFS &TFS;
-  StdLibSet Stdlib;
+  std::shared_ptr<StdLibSet> Stdlib;
   AsyncTaskRunner *Tasks;
 };
 
 class DraftStoreFS : public ThreadsafeFS {
 public:
   DraftStoreFS(const ThreadsafeFS &Base, const DraftStore &Drafts)
       : Base(Base), DirtyFiles(Drafts) {}