[clangd] Narrow and document a loophole in blockUntilIdle

[sammccall]

Summary

blockUntilIdle of a parent can't always be correctly implemented as

return ChildA.blockUntilIdle() && ChildB.blockUntilIdle()

The problem is that B can schedule work on A while we're waiting on it.

I believe this is theoretically possible today between CDB and background index.
Modules open more possibilities and it's hard to reason about all of them.

I don't have a perfect fix, and the abstraction is too good to lose. this patch:

• calls out why we block on workscheduler first, and asserts correctness
• documents the issue
• reduces the practical possibility of spuriously returning true significantly

This function is ultimately only for testing, so we're driving down flake rate.

[kadircet]

LG for the blockUntilIdle update, but feels like two patches wind up together :(

[kadircet]

clang-tools-extra/clangd/ClangdServer.cpp

893

this is extending the Deadline in theory, e.g. if user requested idleness in 10 seconds, this can now wait for up to 20 seconds. but this was possible in the previous case too, e.g. CDB could block for 10 seconds, and then bgindex would block for another 10 seconds, and mentioned this is only for tests, so should be fine (but might be worth mentioning in the comments.)

[sammccall]

Revert unintended changes, add FIXME

[sammccall]

Doh, sorry about the garbled patch

clang-tools-extra/clangd/ClangdServer.cpp

893

Right - the idea with Deadline is it's something absolute you can pass around, and here we are creating a new one for each call.

This is wrong, but it's consistent with the code before this patch and I'd rather not fix both at once. Added a FIXME to the header.

(To fix this we'd want to change the signature of BackgroundIndex to take a Deadline too)

[kadircet]

thanks, lgtm!

clang-tools-extra/clangd/ClangdServer.cpp

ClangdServer::Options::operator TUScheduler::Options() const {
   Opts.AsyncThreadsCount = AsyncThreadsCount;
   Opts.RetentionPolicy = RetentionPolicy;
   Opts.StorePreamblesInMemory = StorePreamblesInMemory;
   Opts.UpdateDebounce = UpdateDebounce;
   Opts.ContextProvider = ContextProvider;
   return Opts;
 }
 
 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) {
-  return WorkScheduler.blockUntilIdle(timeoutSeconds(TimeoutSeconds)) &&
-         CDB.blockUntilIdle(timeoutSeconds(TimeoutSeconds)) &&
-         (!BackgroundIdx ||
-          BackgroundIdx->blockUntilIdleForTest(TimeoutSeconds));
+  Deadline D = timeoutSeconds(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(D))
+    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;
+  }
+
+  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