[clang-tools-extra] f0e6927 - [clangd] Shutdown sequence for modules, and doc threading requirements

Sam McCall via cfe-commits cfe-commits at lists.llvm.org
Mon Feb 22 14:15:11 PST 2021


Author: Sam McCall
Date: 2021-02-22T23:14:47+01:00
New Revision: f0e69272c62f2b7a39687748db60d21445010ac5

URL: https://github.com/llvm/llvm-project/commit/f0e69272c62f2b7a39687748db60d21445010ac5
DIFF: https://github.com/llvm/llvm-project/commit/f0e69272c62f2b7a39687748db60d21445010ac5.diff

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

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

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

Added: 
    

Modified: 
    clang-tools-extra/clangd/ClangdServer.cpp
    clang-tools-extra/clangd/ClangdServer.h
    clang-tools-extra/clangd/Module.h
    clang-tools-extra/clangd/unittests/ClangdLSPServerTests.cpp

Removed: 
    


################################################################################
diff  --git a/clang-tools-extra/clangd/ClangdServer.cpp b/clang-tools-extra/clangd/ClangdServer.cpp
index 5ec0d816c95f..245a2d081f93 100644
--- a/clang-tools-extra/clangd/ClangdServer.cpp
+++ b/clang-tools-extra/clangd/ClangdServer.cpp
@@ -130,15 +130,13 @@ ClangdServer::ClangdServer(const GlobalCompilationDatabase &CDB,
     : 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) {
+  // 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.
+  WorkScheduler.emplace(
+      CDB, TUScheduler::Options(Opts),
+      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) {
@@ -170,7 +168,7 @@ ClangdServer::ClangdServer(const GlobalCompilationDatabase &CDB,
 
   if (Opts.Modules) {
     Module::Facilities F{
-        this->WorkScheduler,
+        *this->WorkScheduler,
         this->Index,
         this->TFS,
     };
@@ -179,6 +177,20 @@ ClangdServer::ClangdServer(const GlobalCompilationDatabase &CDB,
   }
 }
 
+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) {
@@ -193,7 +205,7 @@ void ClangdServer::addDocument(PathRef File, llvm::StringRef Contents,
   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);
@@ -276,7 +288,7 @@ ClangdServer::createConfiguredContextProvider(const config::Provider *Provider,
   };
 }
 
-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,
@@ -330,7 +342,7 @@ void ClangdServer::codeComplete(PathRef File, Position Pos,
   };
 
   // We use a potentially-stale preamble because latency is critical here.
-  WorkScheduler.runWithPreamble(
+  WorkScheduler->runWithPreamble(
       "CodeComplete", File,
       (Opts.RunParser == CodeCompleteOptions::AlwaysParse)
           ? TUScheduler::Stale
@@ -356,8 +368,8 @@ void ClangdServer::signatureHelp(PathRef File, Position Pos,
   };
 
   // Unlike code completion, we wait for a preamble here.
-  WorkScheduler.runWithPreamble("SignatureHelp", File, TUScheduler::Stale,
-                                std::move(Action));
+  WorkScheduler->runWithPreamble("SignatureHelp", File, TUScheduler::Stale,
+                                 std::move(Action));
 }
 
 void ClangdServer::formatRange(PathRef File, llvm::StringRef Code, Range Rng,
@@ -399,7 +411,7 @@ void ClangdServer::formatOnType(PathRef File, llvm::StringRef Code,
       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,
@@ -424,14 +436,14 @@ void ClangdServer::prepareRename(PathRef File, Position Pos,
     }
     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 {
@@ -466,7 +478,7 @@ void ClangdServer::rename(PathRef File, Position Pos, llvm::StringRef NewName,
     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.
@@ -522,8 +534,8 @@ void ClangdServer::enumerateTweaks(
     CB(std::move(Res));
   };
 
-  WorkScheduler.runWithAST("EnumerateTweaks", File, std::move(Action),
-                           TUScheduler::InvalidateOnUpdate);
+  WorkScheduler->runWithAST("EnumerateTweaks", File, std::move(Action),
+                            TUScheduler::InvalidateOnUpdate);
 }
 
 void ClangdServer::applyTweak(PathRef File, Range Sel, StringRef TweakID,
@@ -569,7 +581,7 @@ void ClangdServer::applyTweak(PathRef File, Range Sel, StringRef 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,
@@ -581,7 +593,7 @@ void ClangdServer::locateSymbolAt(PathRef File, Position Pos,
     CB(clangd::locateSymbolAt(InpAST->AST, Pos, Index));
   };
 
-  WorkScheduler.runWithAST("Definitions", File, std::move(Action));
+  WorkScheduler->runWithAST("Definitions", File, std::move(Action));
 }
 
 void ClangdServer::switchSourceHeader(
@@ -601,7 +613,7 @@ void ClangdServer::switchSourceHeader(
       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,
@@ -622,7 +634,7 @@ void ClangdServer::formatCode(PathRef File, llvm::StringRef Code,
         tooling::calculateRangesAfterReplacements(IncludeReplaces, Ranges),
         File)));
   };
-  WorkScheduler.runQuick("Format", File, std::move(Action));
+  WorkScheduler->runQuick("Format", File, std::move(Action));
 }
 
 void ClangdServer::findDocumentHighlights(
@@ -634,8 +646,8 @@ void ClangdServer::findDocumentHighlights(
         CB(clangd::findDocumentHighlights(InpAST->AST, Pos));
       };
 
-  WorkScheduler.runWithAST("Highlights", File, std::move(Action),
-                           TUScheduler::InvalidateOnUpdate);
+  WorkScheduler->runWithAST("Highlights", File, std::move(Action),
+                            TUScheduler::InvalidateOnUpdate);
 }
 
 void ClangdServer::findHover(PathRef File, Position Pos,
@@ -649,8 +661,8 @@ void ClangdServer::findHover(PathRef File, Position Pos,
     CB(clangd::getHover(InpAST->AST, Pos, std::move(Style), Index));
   };
 
-  WorkScheduler.runWithAST("Hover", File, std::move(Action),
-                           TUScheduler::InvalidateOnUpdate);
+  WorkScheduler->runWithAST("Hover", File, std::move(Action),
+                            TUScheduler::InvalidateOnUpdate);
 }
 
 void ClangdServer::typeHierarchy(PathRef File, Position Pos, int Resolve,
@@ -664,13 +676,13 @@ void ClangdServer::typeHierarchy(PathRef File, Position Pos, int Resolve,
                                 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);
@@ -685,16 +697,16 @@ void ClangdServer::prepareCallHierarchy(
       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", "",
-                    [CB = std::move(CB), Item, this]() mutable {
-                      CB(clangd::incomingCalls(Item, Index));
-                    });
+  WorkScheduler->run("Incoming Calls", "",
+                     [CB = std::move(CB), Item, this]() mutable {
+                       CB(clangd::incomingCalls(Item, Index));
+                     });
 }
 
 void ClangdServer::onFileEvent(const DidChangeWatchedFilesParams &Params) {
@@ -705,7 +717,7 @@ void ClangdServer::onFileEvent(const DidChangeWatchedFilesParams &Params) {
 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,
@@ -721,8 +733,8 @@ void ClangdServer::documentSymbols(llvm::StringRef File,
           return CB(InpAST.takeError());
         CB(clangd::getDocumentSymbols(InpAST->AST));
       };
-  WorkScheduler.runWithAST("DocumentSymbols", File, std::move(Action),
-                           TUScheduler::InvalidateOnUpdate);
+  WorkScheduler->runWithAST("DocumentSymbols", File, std::move(Action),
+                            TUScheduler::InvalidateOnUpdate);
 }
 
 void ClangdServer::foldingRanges(llvm::StringRef File,
@@ -733,8 +745,8 @@ void ClangdServer::foldingRanges(llvm::StringRef File,
           return CB(InpAST.takeError());
         CB(clangd::getFoldingRanges(InpAST->AST));
       };
-  WorkScheduler.runWithAST("FoldingRanges", File, std::move(Action),
-                           TUScheduler::InvalidateOnUpdate);
+  WorkScheduler->runWithAST("FoldingRanges", File, std::move(Action),
+                            TUScheduler::InvalidateOnUpdate);
 }
 
 void ClangdServer::findImplementations(
@@ -746,7 +758,7 @@ void ClangdServer::findImplementations(
     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,
@@ -758,7 +770,7 @@ void ClangdServer::findReferences(PathRef File, Position Pos, uint32_t Limit,
     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,
@@ -770,7 +782,7 @@ void ClangdServer::symbolInfo(PathRef File, Position Pos,
         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,
@@ -789,7 +801,7 @@ void ClangdServer::semanticRanges(PathRef File,
     }
     CB(std::move(Result));
   };
-  WorkScheduler.runWithAST("SemanticRanges", File, std::move(Action));
+  WorkScheduler->runWithAST("SemanticRanges", File, std::move(Action));
 }
 
 void ClangdServer::documentLinks(PathRef File,
@@ -800,8 +812,8 @@ void ClangdServer::documentLinks(PathRef File,
           return CB(InpAST.takeError());
         CB(clangd::getDocumentLinks(InpAST->AST));
       };
-  WorkScheduler.runWithAST("DocumentLinks", File, std::move(Action),
-                           TUScheduler::InvalidateOnUpdate);
+  WorkScheduler->runWithAST("DocumentLinks", File, std::move(Action),
+                            TUScheduler::InvalidateOnUpdate);
 }
 
 void ClangdServer::semanticHighlights(
@@ -812,8 +824,8 @@ void ClangdServer::semanticHighlights(
           return CB(InpAST.takeError());
         CB(clangd::getSemanticHighlightings(InpAST->AST));
       };
-  WorkScheduler.runWithAST("SemanticHighlights", File, std::move(Action),
-                           TUScheduler::InvalidateOnUpdate);
+  WorkScheduler->runWithAST("SemanticHighlights", File, std::move(Action),
+                            TUScheduler::InvalidateOnUpdate);
 }
 
 void ClangdServer::getAST(PathRef File, Range R,
@@ -845,16 +857,16 @@ void ClangdServer::getAST(PathRef File, Range R,
         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
@@ -864,7 +876,7 @@ ClangdServer::blockUntilIdleForTest(llvm::Optional<double> TimeoutSeconds) {
 
   // 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
@@ -882,9 +894,13 @@ ClangdServer::blockUntilIdleForTest(llvm::Optional<double> TimeoutSeconds) {
       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;
 }
@@ -894,7 +910,7 @@ void ClangdServer::profile(MemoryTree &MT) const {
     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

diff  --git a/clang-tools-extra/clangd/ClangdServer.h b/clang-tools-extra/clangd/ClangdServer.h
index 1c09c66e6406..198395c4a8f3 100644
--- a/clang-tools-extra/clangd/ClangdServer.h
+++ b/clang-tools-extra/clangd/ClangdServer.h
@@ -160,6 +160,7 @@ class ClangdServer {
   /// 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.
@@ -375,10 +376,7 @@ class ClangdServer {
   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

diff  --git a/clang-tools-extra/clangd/Module.h b/clang-tools-extra/clangd/Module.h
index 148da7789aef..f660eff85653 100644
--- a/clang-tools-extra/clangd/Module.h
+++ b/clang-tools-extra/clangd/Module.h
@@ -10,6 +10,7 @@
 #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"
@@ -36,13 +37,26 @@ class TUScheduler;
 ///    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,
@@ -63,6 +77,20 @@ class Module {
   /// 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();

diff  --git a/clang-tools-extra/clangd/unittests/ClangdLSPServerTests.cpp b/clang-tools-extra/clangd/unittests/ClangdLSPServerTests.cpp
index a079d75dd3af..4542f7423a48 100644
--- a/clang-tools-extra/clangd/unittests/ClangdLSPServerTests.cpp
+++ b/clang-tools-extra/clangd/unittests/ClangdLSPServerTests.cpp
@@ -16,6 +16,7 @@
 #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"
@@ -23,6 +24,7 @@
 namespace clang {
 namespace clangd {
 namespace {
+using llvm::Succeeded;
 using testing::ElementsAre;
 
 MATCHER_P(DiagMessage, M, "") {
@@ -40,6 +42,7 @@ class LSPTest : public ::testing::Test, private clangd::Logger {
     Base = ClangdServer::optsForTest();
     // This is needed to we can test index-based operations like call hierarchy.
     Base.BuildDynamicSymbolIndex = true;
+    Base.Modules = &Modules;
   }
 
   LSPClient &start() {
@@ -67,6 +70,7 @@ class LSPTest : public ::testing::Test, private clangd::Logger {
 
   MockFS FS;
   ClangdLSPServer::Options Opts;
+  ModuleSet Modules;
 
 private:
   // Color logs so we can distinguish them from test output.
@@ -244,9 +248,7 @@ TEST_F(LSPTest, ModulesTest) {
           [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);
@@ -256,6 +258,104 @@ TEST_F(LSPTest, ModulesTest) {
               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


        


More information about the cfe-commits mailing list