[clangd] Shutdown sequence for modules, and doc threading requirements

[sammccall]

Summary

This allows modules to do work on non-TUScheduler background threads.

[kadircet]

clang-tools-extra/clangd/ClangdServer.cpp

190

why is our contract saying that just calling stop is not enough?
i think clangdserver should just signal shutdown to modules, and our contract should say that server facilities will be undefined from this point forward.
that way modules accessing the facilities, could block stop until they are done, and never make use of it afterwards? it'll make modules a little more complicated, at the very least they would need some stricter control whenever they are accessing facilities, but I think it is worth for keeping clangdserver shutdown cleaner. wdyt?

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

273

= false;

361

not sure what this is testing in addition to final callback receiving 3 as a value.

[sammccall]

trim tests

clang-tools-extra/clangd/ClangdServer.cpp

190

We need Module::blockUntilIdle anyway for tests and stuff. And in practice requesting shutdown of background work doesn't naturally block until that work is complete IME.

So we could ask every module to call blockUntilIdle() in stop(), or make stop() non-virtual and have it wrap requestStop()+blockUntilIdle().

Neither really seems simpler/cleaner to me overall, and it means modules shut down in serial instead of parallel.

[kadircet]

thanks, lgtm!

clang-tools-extra/clangd/ClangdServer.cpp

190

We need Module::blockUntilIdle anyway for tests and stuff

i was trying to ensure it is *only* used for tests, but i suppose you are right, in practice making stop block until in-flight tasks finishes vs stop+blockuntilidle is likely to have same effect, with the latter having the benefit of triggering shut down in parallel. so SGTM.

[sammccall]

Fix racy test - forgot to blockUntilIdle in sync!

Found a bug in blockUntilIdle, i'll send it separately.

clang-tools-extra/clangd/ClangdServer.h

public:
   ///
   /// ClangdServer uses \p CDB to obtain compilation arguments for parsing. Note
   /// that ClangdServer only obtains compilation arguments once for each newly
   /// added file (i.e., when processing a first call to addDocument) and reuses
   /// those arguments for subsequent reparses. However, ClangdServer will check
   /// if compilation arguments changed on calls to forceReparse().
   ClangdServer(const GlobalCompilationDatabase &CDB, const ThreadsafeFS &TFS,
                const Options &Opts, Callbacks *Callbacks = nullptr);
+  ~ClangdServer();
 
   /// Gets the installed module of a given type, if any.
   /// This exposes access the public interface of modules that have one.
   template <typename Mod> Mod *getModule() {
     return Modules ? Modules->get<Mod>() : nullptr;
   }
   template <typename Mod> const Mod *getModule() const {
     return Modules ? Modules->get<Mod>() : nullptr;
private:
   TidyProviderRef ClangTidyProvider;
 
   // GUARDED_BY(CachedCompletionFuzzyFindRequestMutex)
   llvm::StringMap<llvm::Optional<FuzzyFindRequest>>
       CachedCompletionFuzzyFindRequestByFile;
   mutable std::mutex CachedCompletionFuzzyFindRequestMutex;
 
   llvm::Optional<std::string> WorkspaceRoot;
-  // WorkScheduler has to be the last member, because its destructor has to be
-  // called before all other members to stop the worker thread that references
-  // ClangdServer.
-  TUScheduler WorkScheduler;
+  llvm::Optional<TUScheduler> WorkScheduler;
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif

clang-tools-extra/clangd/ClangdServer.cpp

 }
 
 ClangdServer::ClangdServer(const GlobalCompilationDatabase &CDB,
                            const ThreadsafeFS &TFS, const Options &Opts,
                            Callbacks *Callbacks)
     : Modules(Opts.Modules), CDB(CDB), TFS(TFS),
       DynamicIdx(Opts.BuildDynamicSymbolIndex ? new FileIndex() : nullptr),
       ClangTidyProvider(Opts.ClangTidyProvider),
-      WorkspaceRoot(Opts.WorkspaceRoot),
-      // Pass a callback into `WorkScheduler` to extract symbols from a newly
-      // parsed file and rebuild the file index synchronously each time an AST
-      // is parsed.
-      // FIXME(ioeric): this can be slow and we may be able to index on less
-      // critical paths.
-      WorkScheduler(
-          CDB, TUScheduler::Options(Opts),
-          std::make_unique<UpdateIndexCallbacks>(DynamicIdx.get(), Callbacks)) {
+      WorkspaceRoot(Opts.WorkspaceRoot) {
+  WorkScheduler.emplace(
+      std::make_unique<UpdateIndexCallbacks>(DynamicIdx.get(), Callbacks));
   // Adds an index to the stack, at higher priority than existing indexes.
   auto AddIndex = [&](SymbolIndex *Idx) {
     if (this->Index != nullptr) {
       MergedIdx.push_back(std::make_unique<MergedIndex>(Idx, this->Index));
       this->Index = MergedIdx.back().get();
     } else {
       this->Index = Idx;
     }
BackgroundIdx = std::make_unique<BackgroundIndex>(
         std::move(BGOpts));
     AddIndex(BackgroundIdx.get());
   }
   if (DynamicIdx)
     AddIndex(DynamicIdx.get());
 
   if (Opts.Modules) {
     Module::Facilities F{
-        this->WorkScheduler,
+        *this->WorkScheduler,
         this->Index,
         this->TFS,
     };
     for (auto &Mod : *Opts.Modules)
       Mod.initialize(F);
   }
 }
 
+ClangdServer::~ClangdServer() {
+  // Destroying TUScheduler first shuts down request threads that might
+  // otherwise access members concurrently.
+  // (Nobody can be using TUScheduler because we're on the main thread).
+  WorkScheduler.reset();
+  // Now requests have stopped, we can shut down modules.
+  if (Modules) {
+    for (auto &Mod : *Modules)
+      Mod.stop();
+    for (auto &Mod : *Modules)
+      Mod.blockUntilIdle(Deadline::infinity());
+  }
+}
+
 void ClangdServer::addDocument(PathRef File, llvm::StringRef Contents,
                                llvm::StringRef Version,
                                WantDiagnostics WantDiags, bool ForceRebuild) {
   ParseOptions Opts;
 
   // Compile command is set asynchronously during update, as it can be slow.
   ParseInputs Inputs;
   Inputs.TFS = &TFS;
   Inputs.Contents = std::string(Contents);
   Inputs.Version = Version.str();
   Inputs.ForceRebuild = ForceRebuild;
   Inputs.Opts = std::move(Opts);
   Inputs.Index = Index;
   Inputs.ClangTidyProvider = ClangTidyProvider;
-  bool NewFile = WorkScheduler.update(File, Inputs, WantDiags);
+  bool NewFile = WorkScheduler->update(File, Inputs, WantDiags);
   // If we loaded Foo.h, we want to make sure Foo.cpp is indexed.
   if (NewFile && BackgroundIdx)
     BackgroundIdx->boostRelated(File);
 }
 
 std::function<Context(PathRef)>
 ClangdServer::createConfiguredContextProvider(const config::Provider *Provider,
                                               Callbacks *Publish) {
ClangdServer::createConfiguredContextProvider(const config::Provider *Provider,
   };
 
   // Copyable wrapper.
   return [I(std::make_shared<Impl>(Provider, Publish))](llvm::StringRef Path) {
     return (*I)(Path);
   };
 }
 
-void ClangdServer::removeDocument(PathRef File) { WorkScheduler.remove(File); }
+void ClangdServer::removeDocument(PathRef File) { WorkScheduler->remove(File); }
 
 void ClangdServer::codeComplete(PathRef File, Position Pos,
                                 const clangd::CodeCompleteOptions &Opts,
                                 Callback<CodeCompleteResult> CB) {
   // Copy completion options for passing them to async task handler.
   auto CodeCompleteOpts = Opts;
   if (!CodeCompleteOpts.Index) // Respect overridden index.
     CodeCompleteOpts.Index = Index;
auto Task = [Pos, CodeCompleteOpts, File = File.str(), CB = std::move(CB),
     }
     // SpecFuzzyFind is only destroyed after speculative fuzzy find finishes.
     // We don't want `codeComplete` to wait for the async call if it doesn't use
     // the result (e.g. non-index completion, speculation fails), so that `CB`
     // is called as soon as results are available.
   };
 
   // We use a potentially-stale preamble because latency is critical here.
-  WorkScheduler.runWithPreamble(
+  WorkScheduler->runWithPreamble(
       "CodeComplete", File,
       (Opts.RunParser == CodeCompleteOptions::AlwaysParse)
           ? TUScheduler::Stale
           : TUScheduler::StaleOrAbsent,
       std::move(Task));
 }
 
 void ClangdServer::signatureHelp(PathRef File, Position Pos,
if (!PreambleData)
       return CB(error("Failed to parse includes"));
 
     ParseInputs ParseInput{IP->Command, &TFS, IP->Contents.str()};
     ParseInput.Index = Index;
     CB(clangd::signatureHelp(File, Pos, *PreambleData, ParseInput));
   };
 
   // Unlike code completion, we wait for a preamble here.
-  WorkScheduler.runWithPreamble("SignatureHelp", File, TUScheduler::Stale,
+  WorkScheduler->runWithPreamble("SignatureHelp", File, TUScheduler::Stale,
                                  std::move(Action));
 }
 
 void ClangdServer::formatRange(PathRef File, llvm::StringRef Code, Range Rng,
                                Callback<tooling::Replacements> CB) {
   llvm::Expected<size_t> Begin = positionToOffset(Code, Rng.start);
   if (!Begin)
     return CB(Begin.takeError());
   llvm::Expected<size_t> End = positionToOffset(Code, Rng.end);
if (!Style)
       return CB(Style.takeError());
 
     std::vector<TextEdit> Result;
     for (const tooling::Replacement &R :
          formatIncremental(Code, CursorPos, TriggerText, *Style))
       Result.push_back(replacementToEdit(Code, R));
     return CB(Result);
   };
-  WorkScheduler.runQuick("FormatOnType", File, std::move(Action));
+  WorkScheduler->runQuick("FormatOnType", File, std::move(Action));
 }
 
 void ClangdServer::prepareRename(PathRef File, Position Pos,
                                  llvm::Optional<std::string> NewName,
                                  const RenameOptions &RenameOpts,
                                  Callback<RenameResult> CB) {
   auto Action = [Pos, File = File.str(), CB = std::move(CB),
                  NewName = std::move(NewName),
                  RenameOpts](llvm::Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)
       return CB(InpAST.takeError());
     // prepareRename is latency-sensitive: we don't query the index, as we
     // only need main-file references
     auto Results = clangd::rename(
         {Pos, NewName.getValueOr("__clangd_rename_dummy"), InpAST->AST, File,
          /*Index=*/nullptr, RenameOpts});
     if (!Results) {
       // LSP says to return null on failure, but that will result in a generic
       // failure message. If we send an LSP error response, clients can surface
       // the message to users (VSCode does).
       return CB(Results.takeError());
     }
     return CB(*Results);
   };
-  WorkScheduler.runWithAST("PrepareRename", File, std::move(Action));
+  WorkScheduler->runWithAST("PrepareRename", File, std::move(Action));
 }
 
 void ClangdServer::rename(PathRef File, Position Pos, llvm::StringRef NewName,
                           const RenameOptions &Opts,
                           Callback<RenameResult> CB) {
   // A snapshot of all file dirty buffers.
-  llvm::StringMap<std::string> Snapshot = WorkScheduler.getAllFileContents();
+  llvm::StringMap<std::string> Snapshot = WorkScheduler->getAllFileContents();
   auto Action = [File = File.str(), NewName = NewName.str(), Pos, Opts,
                  CB = std::move(CB), Snapshot = std::move(Snapshot),
                  this](llvm::Expected<InputsAndAST> InpAST) mutable {
     // Tracks number of files edited per invocation.
     static constexpr trace::Metric RenameFiles("rename_files",
                                                trace::Metric::Distribution);
     if (!InpAST)
       return CB(InpAST.takeError());
if (Opts.WantFormat) {
             llvm::joinErrors(reformatEdit(E.getValue(), Style), std::move(Err));
 
       if (Err)
         return CB(std::move(Err));
     }
     RenameFiles.record(R->GlobalChanges.size());
     return CB(*R);
   };
-  WorkScheduler.runWithAST("Rename", File, std::move(Action));
+  WorkScheduler->runWithAST("Rename", File, std::move(Action));
 }
 
 // May generate several candidate selections, due to SelectionTree ambiguity.
 // vector of pointers because GCC doesn't like non-copyable Selection.
 static llvm::Expected<std::vector<std::unique_ptr<Tweak::Selection>>>
 tweakSelection(const Range &Sel, const InputsAndAST &AST) {
   auto Begin = positionToOffset(AST.Inputs.Contents, Sel.start);
   if (!Begin)
for (const auto &Sel : *Selections) {
         PreparedTweaks.insert(T->id());
         TweakAvailable.record(1, T->id());
       }
     }
 
     CB(std::move(Res));
   };
 
-  WorkScheduler.runWithAST("EnumerateTweaks", File, std::move(Action),
+  WorkScheduler->runWithAST("EnumerateTweaks", File, std::move(Action),
                             TUScheduler::InvalidateOnUpdate);
 }
 
 void ClangdServer::applyTweak(PathRef File, Range Sel, StringRef TweakID,
                               Callback<Tweak::Effect> CB) {
   // Tracks number of times a tweak has been attempted.
   static constexpr trace::Metric TweakAttempt(
       "tweak_attempt", trace::Metric::Counter, "tweak_id");
   // Tracks number of times a tweak has failed to produce edits.
if (*Effect) {
         if (llvm::Error Err = reformatEdit(E, Style))
           elog("Failed to format {0}: {1}", It.first(), std::move(Err));
       }
     } else {
       TweakFailed.record(1, TweakID);
     }
     return CB(std::move(*Effect));
   };
-  WorkScheduler.runWithAST("ApplyTweak", File, std::move(Action));
+  WorkScheduler->runWithAST("ApplyTweak", File, std::move(Action));
 }
 
 void ClangdServer::locateSymbolAt(PathRef File, Position Pos,
                                   Callback<std::vector<LocatedSymbol>> CB) {
   auto Action = [Pos, CB = std::move(CB),
                  this](llvm::Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)
       return CB(InpAST.takeError());
     CB(clangd::locateSymbolAt(InpAST->AST, Pos, Index));
   };
 
-  WorkScheduler.runWithAST("Definitions", File, std::move(Action));
+  WorkScheduler->runWithAST("Definitions", File, std::move(Action));
 }
 
 void ClangdServer::switchSourceHeader(
     PathRef Path, Callback<llvm::Optional<clangd::Path>> CB) {
   // We want to return the result as fast as possible, strategy is:
   //  1) use the file-only heuristic, it requires some IO but it is much
   //     faster than building AST, but it only works when .h/.cc files are in
   //     the same directory.
   //  2) if 1) fails, we use the AST&Index approach, it is slower but supports
   //     different code layout.
   if (auto CorrespondingFile =
           getCorrespondingHeaderOrSource(Path, TFS.view(llvm::None)))
     return CB(std::move(CorrespondingFile));
   auto Action = [Path = Path.str(), CB = std::move(CB),
                  this](llvm::Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)
       return CB(InpAST.takeError());
     CB(getCorrespondingHeaderOrSource(Path, InpAST->AST, Index));
   };
-  WorkScheduler.runWithAST("SwitchHeaderSource", Path, std::move(Action));
+  WorkScheduler->runWithAST("SwitchHeaderSource", Path, std::move(Action));
 }
 
 void ClangdServer::formatCode(PathRef File, llvm::StringRef Code,
                               llvm::ArrayRef<tooling::Range> Ranges,
                               Callback<tooling::Replacements> CB) {
   // Call clang-format.
   auto Action = [File = File.str(), Code = Code.str(), Ranges = Ranges.vec(),
                  CB = std::move(CB), this]() mutable {
     format::FormatStyle Style = getFormatStyleForFile(File, Code, TFS);
     tooling::Replacements IncludeReplaces =
         format::sortIncludes(Style, Code, Ranges, File);
     auto Changed = tooling::applyAllReplacements(Code, IncludeReplaces);
     if (!Changed)
       return CB(Changed.takeError());
 
     CB(IncludeReplaces.merge(format::reformat(
         Style, *Changed,
         tooling::calculateRangesAfterReplacements(IncludeReplaces, Ranges),
         File)));
   };
-  WorkScheduler.runQuick("Format", File, std::move(Action));
+  WorkScheduler->runQuick("Format", File, std::move(Action));
 }
 
 void ClangdServer::findDocumentHighlights(
     PathRef File, Position Pos, Callback<std::vector<DocumentHighlight>> CB) {
   auto Action =
       [Pos, CB = std::move(CB)](llvm::Expected<InputsAndAST> InpAST) mutable {
         if (!InpAST)
           return CB(InpAST.takeError());
         CB(clangd::findDocumentHighlights(InpAST->AST, Pos));
       };
 
-  WorkScheduler.runWithAST("Highlights", File, std::move(Action),
+  WorkScheduler->runWithAST("Highlights", File, std::move(Action),
                             TUScheduler::InvalidateOnUpdate);
 }
 
 void ClangdServer::findHover(PathRef File, Position Pos,
                              Callback<llvm::Optional<HoverInfo>> CB) {
   auto Action = [File = File.str(), Pos, CB = std::move(CB),
                  this](llvm::Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)
       return CB(InpAST.takeError());
     format::FormatStyle Style = getFormatStyleForFile(
         File, InpAST->Inputs.Contents, *InpAST->Inputs.TFS);
     CB(clangd::getHover(InpAST->AST, Pos, std::move(Style), Index));
   };
 
-  WorkScheduler.runWithAST("Hover", File, std::move(Action),
+  WorkScheduler->runWithAST("Hover", File, std::move(Action),
                             TUScheduler::InvalidateOnUpdate);
 }
 
 void ClangdServer::typeHierarchy(PathRef File, Position Pos, int Resolve,
                                  TypeHierarchyDirection Direction,
                                  Callback<Optional<TypeHierarchyItem>> CB) {
   auto Action = [File = File.str(), Pos, Resolve, Direction, CB = std::move(CB),
                  this](Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)
       return CB(InpAST.takeError());
     CB(clangd::getTypeHierarchy(InpAST->AST, Pos, Resolve, Direction, Index,
                                 File));
   };
 
-  WorkScheduler.runWithAST("TypeHierarchy", File, std::move(Action));
+  WorkScheduler->runWithAST("TypeHierarchy", File, std::move(Action));
 }
 
 void ClangdServer::resolveTypeHierarchy(
     TypeHierarchyItem Item, int Resolve, TypeHierarchyDirection Direction,
     Callback<llvm::Optional<TypeHierarchyItem>> CB) {
-  WorkScheduler.run(
+  WorkScheduler->run(
       "Resolve Type Hierarchy", "", [=, CB = std::move(CB)]() mutable {
         clangd::resolveTypeHierarchy(Item, Resolve, Direction, Index);
         CB(Item);
       });
 }
 
 void ClangdServer::prepareCallHierarchy(
     PathRef File, Position Pos, Callback<std::vector<CallHierarchyItem>> CB) {
   auto Action = [File = File.str(), Pos,
                  CB = std::move(CB)](Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)
       return CB(InpAST.takeError());
     CB(clangd::prepareCallHierarchy(InpAST->AST, Pos, File));
   };
-  WorkScheduler.runWithAST("CallHierarchy", File, std::move(Action));
+  WorkScheduler->runWithAST("CallHierarchy", File, std::move(Action));
 }
 
 void ClangdServer::incomingCalls(
     const CallHierarchyItem &Item,
     Callback<std::vector<CallHierarchyIncomingCall>> CB) {
-  WorkScheduler.run("Incoming Calls", "",
+  WorkScheduler->run("Incoming Calls", "",
                      [CB = std::move(CB), Item, this]() mutable {
                        CB(clangd::incomingCalls(Item, Index));
                      });
 }
 
 void ClangdServer::onFileEvent(const DidChangeWatchedFilesParams &Params) {
   // FIXME: Do nothing for now. This will be used for indexing and potentially
   // invalidating other caches.
 }
 
 void ClangdServer::workspaceSymbols(
     llvm::StringRef Query, int Limit,
     Callback<std::vector<SymbolInformation>> CB) {
-  WorkScheduler.run(
+  WorkScheduler->run(
       "getWorkspaceSymbols", /*Path=*/"",
       [Query = Query.str(), Limit, CB = std::move(CB), this]() mutable {
         CB(clangd::getWorkspaceSymbols(Query, Limit, Index,
                                        WorkspaceRoot.getValueOr("")));
       });
 }
 
 void ClangdServer::documentSymbols(llvm::StringRef File,
                                    Callback<std::vector<DocumentSymbol>> CB) {
   auto Action =
       [CB = std::move(CB)](llvm::Expected<InputsAndAST> InpAST) mutable {
         if (!InpAST)
           return CB(InpAST.takeError());
         CB(clangd::getDocumentSymbols(InpAST->AST));
       };
-  WorkScheduler.runWithAST("DocumentSymbols", File, std::move(Action),
+  WorkScheduler->runWithAST("DocumentSymbols", File, std::move(Action),
                             TUScheduler::InvalidateOnUpdate);
 }
 
 void ClangdServer::foldingRanges(llvm::StringRef File,
                                  Callback<std::vector<FoldingRange>> CB) {
   auto Action =
       [CB = std::move(CB)](llvm::Expected<InputsAndAST> InpAST) mutable {
         if (!InpAST)
           return CB(InpAST.takeError());
         CB(clangd::getFoldingRanges(InpAST->AST));
       };
-  WorkScheduler.runWithAST("FoldingRanges", File, std::move(Action),
+  WorkScheduler->runWithAST("FoldingRanges", File, std::move(Action),
                             TUScheduler::InvalidateOnUpdate);
 }
 
 void ClangdServer::findImplementations(
     PathRef File, Position Pos, Callback<std::vector<LocatedSymbol>> CB) {
   auto Action = [Pos, CB = std::move(CB),
                  this](llvm::Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)
       return CB(InpAST.takeError());
     CB(clangd::findImplementations(InpAST->AST, Pos, Index));
   };
 
-  WorkScheduler.runWithAST("Implementations", File, std::move(Action));
+  WorkScheduler->runWithAST("Implementations", File, std::move(Action));
 }
 
 void ClangdServer::findReferences(PathRef File, Position Pos, uint32_t Limit,
                                   Callback<ReferencesResult> CB) {
   auto Action = [Pos, Limit, CB = std::move(CB),
                  this](llvm::Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)
       return CB(InpAST.takeError());
     CB(clangd::findReferences(InpAST->AST, Pos, Limit, Index));
   };
 
-  WorkScheduler.runWithAST("References", File, std::move(Action));
+  WorkScheduler->runWithAST("References", File, std::move(Action));
 }
 
 void ClangdServer::symbolInfo(PathRef File, Position Pos,
                               Callback<std::vector<SymbolDetails>> CB) {
   auto Action =
       [Pos, CB = std::move(CB)](llvm::Expected<InputsAndAST> InpAST) mutable {
         if (!InpAST)
           return CB(InpAST.takeError());
         CB(clangd::getSymbolInfo(InpAST->AST, Pos));
       };
 
-  WorkScheduler.runWithAST("SymbolInfo", File, std::move(Action));
+  WorkScheduler->runWithAST("SymbolInfo", File, std::move(Action));
 }
 
 void ClangdServer::semanticRanges(PathRef File,
                                   const std::vector<Position> &Positions,
                                   Callback<std::vector<SelectionRange>> CB) {
   auto Action = [Positions, CB = std::move(CB)](
                     llvm::Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)
       return CB(InpAST.takeError());
     std::vector<SelectionRange> Result;
     for (const auto &Pos : Positions) {
       if (auto Range = clangd::getSemanticRanges(InpAST->AST, Pos))
         Result.push_back(std::move(*Range));
       else
         return CB(Range.takeError());
     }
     CB(std::move(Result));
   };
-  WorkScheduler.runWithAST("SemanticRanges", File, std::move(Action));
+  WorkScheduler->runWithAST("SemanticRanges", File, std::move(Action));
 }
 
 void ClangdServer::documentLinks(PathRef File,
                                  Callback<std::vector<DocumentLink>> CB) {
   auto Action =
       [CB = std::move(CB)](llvm::Expected<InputsAndAST> InpAST) mutable {
         if (!InpAST)
           return CB(InpAST.takeError());
         CB(clangd::getDocumentLinks(InpAST->AST));
       };
-  WorkScheduler.runWithAST("DocumentLinks", File, std::move(Action),
+  WorkScheduler->runWithAST("DocumentLinks", File, std::move(Action),
                             TUScheduler::InvalidateOnUpdate);
 }
 
 void ClangdServer::semanticHighlights(
     PathRef File, Callback<std::vector<HighlightingToken>> CB) {
   auto Action =
       [CB = std::move(CB)](llvm::Expected<InputsAndAST> InpAST) mutable {
         if (!InpAST)
           return CB(InpAST.takeError());
         CB(clangd::getSemanticHighlightings(InpAST->AST));
       };
-  WorkScheduler.runWithAST("SemanticHighlights", File, std::move(Action),
+  WorkScheduler->runWithAST("SemanticHighlights", File, std::move(Action),
                             TUScheduler::InvalidateOnUpdate);
 }
 
 void ClangdServer::getAST(PathRef File, Range R,
                           Callback<llvm::Optional<ASTNode>> CB) {
   auto Action =
       [R, CB(std::move(CB))](llvm::Expected<InputsAndAST> Inputs) mutable {
         if (!Inputs)
           return CB(Inputs.takeError());
auto Action =
                            Inputs->AST.getASTContext()));
                 return true;
               }
               return false;
             });
         if (!Success)
           CB(llvm::None);
       };
-  WorkScheduler.runWithAST("GetAST", File, std::move(Action));
+  WorkScheduler->runWithAST("GetAST", File, std::move(Action));
 }
 
 void ClangdServer::customAction(PathRef File, llvm::StringRef Name,
                                 Callback<InputsAndAST> Action) {
-  WorkScheduler.runWithAST(Name, File, std::move(Action));
+  WorkScheduler->runWithAST(Name, File, std::move(Action));
 }
 
 llvm::StringMap<TUScheduler::FileStats> ClangdServer::fileStats() const {
-  return WorkScheduler.fileStats();
+  return WorkScheduler->fileStats();
 }
 
 LLVM_NODISCARD bool
 ClangdServer::blockUntilIdleForTest(llvm::Optional<double> TimeoutSeconds) {
   // Order is important here: we don't want to block on A and then B,
   // if B might schedule work on A.
 
   // Nothing else can schedule work on TUScheduler, because it's not threadsafe
   // and we're blocking the main thread.
-  if (!WorkScheduler.blockUntilIdle(timeoutSeconds(TimeoutSeconds)))
+  if (!WorkScheduler->blockUntilIdle(timeoutSeconds(TimeoutSeconds)))
     return false;
 
   // Unfortunately we don't have strict topological order between the rest of
   // the components. E.g. CDB broadcast triggers backrgound indexing.
   // This queries the CDB which may discover new work if disk has changed.
   //
   // So try each one a few times in a loop.
   // If there are no tricky interactions then all after the first are no-ops.
   // Then on the last iteration, verify they're idle without waiting.
   //
   // There's a small chance they're juggling work and we didn't catch them :-(
   for (llvm::Optional<double> Timeout :
        {TimeoutSeconds, TimeoutSeconds, llvm::Optional<double>(0)}) {
     if (!CDB.blockUntilIdle(timeoutSeconds(Timeout)))
       return false;
     if (BackgroundIdx && !BackgroundIdx->blockUntilIdleForTest(Timeout))
       return false;
+    if (Modules && llvm::any_of(*Modules, [&](Module &M) {
+          return !M.blockUntilIdle(timeoutSeconds(Timeout));
+        }))
+      return false;
   }
 
-  assert(WorkScheduler.blockUntilIdle(Deadline::zero()) &&
+  assert(WorkScheduler->blockUntilIdle(Deadline::zero()) &&
          "Something scheduled work while we're blocking the main thread!");
   return true;
 }
 
 void ClangdServer::profile(MemoryTree &MT) const {
   if (DynamicIdx)
     DynamicIdx->profile(MT.child("dynamic_index"));
   if (BackgroundIdx)
     BackgroundIdx->profile(MT.child("background_index"));
-  WorkScheduler.profile(MT.child("tuscheduler"));
+  WorkScheduler->profile(MT.child("tuscheduler"));
 }
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/Module.h

 //===--- Module.h - Plugging features into clangd -----------------*-C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_MODULE_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_MODULE_H
 
 #include "support/Function.h"
+#include "support/Threading.h"
 #include "llvm/ADT/FunctionExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/JSON.h"
 #include <memory>
 #include <type_traits>
 #include <vector>
 

 ///  - modules are created before the LSP server, in ClangdMain.cpp
 ///  - these modules are then passed to ClangdLSPServer in a ModuleSet
 ///  - initializeLSP() is called when the editor calls initialize.
 //   - initialize() is then called by ClangdServer as it is constructed.
 ///  - module hooks can be called by the server at this point.
 ///    Server facilities (scheduler etc) are available.
 ///  - ClangdServer will not be destroyed until all the requests are done.
 ///    FIXME: Block server shutdown until all the modules are idle.
+///  - When shutting down, ClangdServer will wait for all requests to
+///    finish, call stop(), and then blockUntilIdle().
 ///  - modules will be destroyed after ClangdLSPServer is destroyed.
 ///
+/// Modules are not threadsafe in general. A module's entrypoints are:
+///   - method handlers registered in initializeLSP()
+///   - public methods called directly via ClangdServer.getModule<T>()->...
+///   - specific overridable "hook" methods inherited from Module
+/// Unless otherwise specified, these are only called on the main thread.
+///
 /// Conventionally, standard modules live in the `clangd` namespace, and other
 /// exposed details live in a sub-namespace.
 class Module {
 public:
-  virtual ~Module() = default;
+  virtual ~Module() {
+    /// Perform shutdown sequence on destruction in case the ClangdServer was
+    /// never initialized. Usually redundant, but shutdown is idempotent.
+    stop();
+    blockUntilIdle(Deadline::infinity());
+  }
 
   /// Called by the server to connect this module to LSP.
   /// The module should register the methods/notifications/commands it handles,
   /// and update the server capabilities to advertise them.
   ///
   /// This is only called if the module is running in ClangdLSPServer!
   /// Modules with a public interface should satisfy it without LSP bindings.
   virtual void initializeLSP(LSPBinder &Bind,
                              const llvm::json::Object &ClientCaps,
                              llvm::json::Object &ServerCaps) {}
 
   /// Shared server facilities needed by the module to get its work done.
   struct Facilities {
     TUScheduler &Scheduler;
     const SymbolIndex *Index;
     const ThreadsafeFS &FS;
   };
   /// Called by the server to prepare this module for use.
   void initialize(const Facilities &F);
 
+  /// Requests that the module cancel background work and go idle soon.
+  /// Does not block, the caller will call blockUntilIdle() instead.
+  /// After a module is stop()ed, it should not receive any more requests.
+  /// Called by the server when shutting down.
+  /// May be called multiple times, should be idempotent.
+  virtual void stop() {}
+
+  /// Waits until the module is idle (no background work) or a deadline expires.
+  /// In general all modules should eventually go idle, though it may take a
+  /// long time (e.g. background indexing).
+  /// Modules should go idle quickly if stop() has been called.
+  /// Called by the server when shutting down, and also by tests.
+  virtual bool blockUntilIdle(Deadline) { return true; }
+
 protected:
   /// Accessors for modules to access shared server facilities they depend on.
   Facilities &facilities();
   /// The scheduler is used to run tasks on worker threads and access ASTs.
   TUScheduler &scheduler() { return facilities().Scheduler; }
   /// The index is used to get information about the whole codebase.
   const SymbolIndex *index() { return facilities().Index; }
   /// The filesystem is used to read source files on disk.

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

 #include "LSPClient.h"
 #include "Protocol.h"
 #include "TestFS.h"
 #include "support/Logger.h"
 #include "support/TestTracer.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/JSON.h"
+#include "llvm/Testing/Support/Error.h"
 #include "llvm/Testing/Support/SupportHelpers.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 namespace clang {
 namespace clangd {
 namespace {
+using llvm::Succeeded;
 using testing::ElementsAre;
 
 MATCHER_P(DiagMessage, M, "") {
   if (const auto *O = arg.getAsObject()) {
     if (const auto Msg = O->getString("message"))
       return *Msg == M;
   }
   return false;
 }
 
 class LSPTest : public ::testing::Test, private clangd::Logger {
 protected:
   LSPTest() : LogSession(*this) {
     ClangdServer::Options &Base = Opts;
     Base = ClangdServer::optsForTest();
     // This is needed to we can test index-based operations like call hierarchy.
     Base.BuildDynamicSymbolIndex = true;
+    Base.Modules = &Modules;
   }
 
   LSPClient &start() {
     EXPECT_FALSE(Server.hasValue()) << "Already initialized";
     Server.emplace(Client.transport(), FS, Opts);
     ServerThread.emplace([&] { EXPECT_TRUE(Server->run()); });
     Client.call("initialize", llvm::json::Object{});
     return Client;
protected:
 
   ~LSPTest() {
     if (Server)
       stop();
   }
 
   MockFS FS;
   ClangdLSPServer::Options Opts;
+  ModuleSet Modules;
 
 private:
   // Color logs so we can distinguish them from test output.
   void log(Level L, const char *Fmt,
            const llvm::formatv_object_base &Message) override {
     raw_ostream::Colors Color;
     switch (L) {
     case Level::Verbose:
void add(const int &X) {
       Changed(Value);
     }
     void get(const std::nullptr_t &, Callback<int> Reply) {
       scheduler().runQuick(
           "get", "",
           [Reply(std::move(Reply)), Value(Value)]() mutable { Reply(Value); });
     }
   };
-  ModuleSet Mods;
-  Mods.add(std::make_unique<MathModule>());
-  Opts.Modules = &Mods;
+  Modules.add(std::make_unique<MathModule>());
 
   auto &Client = start();
   Client.notify("add", 2);
   Client.notify("add", 8);
   EXPECT_EQ(10, Client.call("get", nullptr).takeValue());
   EXPECT_THAT(Client.takeNotifications("changed"),
               ElementsAre(llvm::json::Value(2), llvm::json::Value(10)));
 }
 
+// Creates a Callback that writes its received value into an Optional<Expected>.
+template <typename T>
+llvm::unique_function<void(llvm::Expected<T>)>
+capture(llvm::Optional<llvm::Expected<T>> &Out) {
+  Out.reset();
+  return [&Out](llvm::Expected<T> V) { Out.emplace(std::move(V)); };
+}
+
+TEST_F(LSPTest, ModulesThreadingTest) {
+  // A module that does its work on a background thread, and so exercises the
+  // block/shutdown protocol.
+  class AsyncCounter final : public Module {
+    bool ShouldStop = false;
+    int State = 0;
+    std::deque<Callback<int>> Queue; // null = increment, non-null = read.
+    std::condition_variable CV;
+    std::mutex Mu;
+    std::thread Thread;
+
+    void run() {
+      std::unique_lock<std::mutex> Lock(Mu);
+      while (true) {
+        CV.wait(Lock, [&] { return ShouldStop || !Queue.empty(); });
+        if (ShouldStop) {
+          Queue.clear();
+          CV.notify_all();
+          return;
+        }
+        Callback<int> &Task = Queue.front();
+        if (Task)
+          Task(State);
+        else
+          ++State;
+        Queue.pop_front();
+        CV.notify_all();
+      }
+    }
+
+    bool blockUntilIdle(Deadline D) override {
+      std::unique_lock<std::mutex> Lock(Mu);
+      return clangd::wait(Lock, CV, D, [this] { return Queue.empty(); });
+    }
+
+    void stop() override {
+      {
+        std::lock_guard<std::mutex> Lock(Mu);
+        ShouldStop = true;
+      }
+      CV.notify_all();
+    }
+
+  public:
+    AsyncCounter() : Thread([this] { run(); }) {}
+    ~AsyncCounter() {
+      // Verify shutdown sequence was performed.
+      // Real modules would not do this, to be robust to no ClangdServer.
+      EXPECT_TRUE(ShouldStop) << "ClangdServer should request shutdown";
+      EXPECT_EQ(Queue.size(), 0u) << "ClangdServer should block until idle";
+      Thread.join();
+    }
+
+    void initializeLSP(LSPBinder &Bind, const llvm::json::Object &ClientCaps,
+                       llvm::json::Object &ServerCaps) override {
+      Bind.notification("increment", this, &AsyncCounter::increment);
+    }
+
+    // Get the current value, bypassing the queue.
+    // Used to verify that sync->blockUntilIdle avoids races in tests.
+    int getSync() {
+      std::lock_guard<std::mutex> Lock(Mu);
+      return State;
+    }
+
+    // Increment the current value asynchronously.
+    void increment(const std::nullptr_t &) {
+      {
+        std::lock_guard<std::mutex> Lock(Mu);
+        Queue.push_back(nullptr);
+      }
+      CV.notify_all();
+    }
+  };
+
+  Modules.add(std::make_unique<AsyncCounter>());
+  auto &Client = start();
+
+  Client.notify("increment", nullptr);
+  Client.notify("increment", nullptr);
+  Client.notify("increment", nullptr);
+  EXPECT_THAT_EXPECTED(Client.call("sync", nullptr).take(), Succeeded());
+  EXPECT_EQ(3, Modules.get<AsyncCounter>()->getSync());
+  // Throw some work on the queue to make sure shutdown blocks on it.
+  Client.notify("increment", nullptr);
+  Client.notify("increment", nullptr);
+  Client.notify("increment", nullptr);
+  // And immediately shut down. Module destructor verifies that we blocked.
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang