nixd: threaded RPC (#367)

nix-community · Mar 29, 2024 · f1d754a · f1d754a
2 parents 551ad9c + 0aa504f
commit f1d754a
Show file tree

Hide file tree

Showing 11 changed files with 253 additions and 136 deletions.
diff --git a/nixd/nix-node-eval/src/EvalProvider.cpp b/nixd/nix-node-eval/src/EvalProvider.cpp
@@ -35,6 +35,10 @@ EvalProvider::EvalProvider(std::unique_ptr<lspserver::InboundPort> In,
           new nix::EvalState{{}, nix::openStore("dummy://")})) {
   Registry.addMethod("exprValue", this, &EvalProvider::onExprValue);
   Registry.addNotification("registerBC", this, &EvalProvider::onRegisterBC);
+
+  Ready = mkOutNotifiction<int>("ready");
+
+  Ready(getpid());
 }
 
 void EvalProvider::onRegisterBC(const rpc::RegisterBCParams &Params) {

diff --git a/nixd/nix-node-eval/src/EvalProvider.h b/nixd/nix-node-eval/src/EvalProvider.h
@@ -19,7 +19,7 @@ class EvalProvider : public lspserver::LSPServer {
   nixt::EnvMap EMap;
   std::unique_ptr<nix::EvalState> State;
 
-  llvm::unique_function<void(int)> Exit;
+  llvm::unique_function<void(int)> Ready;
 
 public:
   EvalProvider(std::unique_ptr<lspserver::InboundPort> In,

diff --git a/nixd/tools/nixd/src/CodeAction.cpp b/nixd/tools/nixd/src/CodeAction.cpp
@@ -6,46 +6,55 @@
 #include "Controller.h"
 #include "Convert.h"
 
+#include <boost/asio/post.hpp>
+
 namespace nixd {
 
 using namespace llvm::json;
 using namespace lspserver;
 
 void Controller::onCodeAction(const lspserver::CodeActionParams &Params,
                               Callback<std::vector<CodeAction>> Reply) {
-  PathRef File = Params.textDocument.uri.file();
+  std::string File(Params.textDocument.uri.file());
   Range Range = Params.range;
-  const std::vector<nixf::Diagnostic> &Diagnostics = TUs[File].diagnostics();
-  std::vector<CodeAction> Actions;
-  Actions.reserve(Diagnostics.size());
-  for (const nixf::Diagnostic &D : Diagnostics) {
-    auto DRange = toLSPRange(D.range());
-    if (!Range.overlap(DRange))
-      continue;
+  auto Action = [Reply = std::move(Reply), File, Range, this]() mutable {
+    std::vector<nixf::Diagnostic> Diagnostics;
+    {
+      std::lock_guard TU(TUsLock);
+      Diagnostics = TUs[File].diagnostics();
+    }
+    std::vector<CodeAction> Actions;
+    Actions.reserve(Diagnostics.size());
+    for (const nixf::Diagnostic &D : Diagnostics) {
+      auto DRange = toLSPRange(D.range());
+      if (!Range.overlap(DRange))
+        continue;
 
-    // Add fixes.
-    for (const nixf::Fix &F : D.fixes()) {
-      std::vector<TextEdit> Edits;
-      Edits.reserve(F.edits().size());
-      for (const nixf::TextEdit &TE : F.edits()) {
-        Edits.emplace_back(TextEdit{
-            .range = toLSPRange(TE.oldRange()),
-            .newText = std::string(TE.newText()),
+      // Add fixes.
+      for (const nixf::Fix &F : D.fixes()) {
+        std::vector<TextEdit> Edits;
+        Edits.reserve(F.edits().size());
+        for (const nixf::TextEdit &TE : F.edits()) {
+          Edits.emplace_back(TextEdit{
+              .range = toLSPRange(TE.oldRange()),
+              .newText = std::string(TE.newText()),
+          });
+        }
+        using Changes = std::map<std::string, std::vector<TextEdit>>;
+        std::string FileURI = URIForFile::canonicalize(File, File).uri();
+        WorkspaceEdit WE{.changes = Changes{
+                             {std::move(FileURI), std::move(Edits)},
+                         }};
+        Actions.emplace_back(CodeAction{
+            .title = F.message(),
+            .kind = std::string(CodeAction::QUICKFIX_KIND),
+            .edit = std::move(WE),
         });
       }
-      using Changes = std::map<std::string, std::vector<TextEdit>>;
-      std::string FileURI = URIForFile::canonicalize(File, File).uri();
-      WorkspaceEdit WE{.changes = Changes{
-                           {std::move(FileURI), std::move(Edits)},
-                       }};
-      Actions.emplace_back(CodeAction{
-          .title = F.message(),
-          .kind = std::string(CodeAction::QUICKFIX_KIND),
-          .edit = std::move(WE),
-      });
     }
-  }
-  Reply(std::move(Actions));
+    Reply(std::move(Actions));
+  };
+  boost::asio::post(Pool, std::move(Action));
 }
 
 } // namespace nixd
diff --git a/nixd/tools/nixd/src/Controller.h b/nixd/tools/nixd/src/Controller.h
@@ -6,6 +6,8 @@
 #include "lspserver/DraftStore.h"
 #include "lspserver/LSPServer.h"
 
+#include <boost/asio/thread_pool.hpp>
+
 namespace nixd {
 
 class Controller : public lspserver::LSPServer {
@@ -15,8 +17,14 @@ class Controller : public lspserver::LSPServer {
   llvm::unique_function<void(const lspserver::PublishDiagnosticsParams &)>
       PublishDiagnostic;
 
+  std::mutex TUsLock;
   llvm::StringMap<NixTU> TUs;
 
+  boost::asio::thread_pool Pool;
+
+  /// In lit-test mode. Disable some concurrency for better text-testing.
+  bool LitTest;
+
   /// Action right after a document is added (including updates).
   void actOnDocumentAdd(lspserver::PathRef File,
                         std::optional<int64_t> Version);
@@ -51,6 +59,12 @@ class Controller : public lspserver::LSPServer {
 public:
   Controller(std::unique_ptr<lspserver::InboundPort> In,
              std::unique_ptr<lspserver::OutboundPort> Out);
+
+  ~Controller() { Pool.join(); }
+
+  void setLitTest(bool LitTest) { this->LitTest = LitTest; }
+
+  bool isReadyToEval() { return Eval && Eval->ready(); }
 };
 
 } // namespace nixd
diff --git a/nixd/tools/nixd/src/EvalClient.cpp b/nixd/tools/nixd/src/EvalClient.cpp
@@ -45,4 +45,13 @@ std::unique_ptr<OwnedEvalClient> OwnedEvalClient::create(int &Fail) {
                                            std::move(ProcFdStream));
 }
 
+EvalClient::EvalClient(std::unique_ptr<lspserver::InboundPort> In,
+                       std::unique_ptr<lspserver::OutboundPort> Out)
+    : lspserver::LSPServer(std::move(In), std::move(Out)), Ready(false) {
+  RegisterBC = mkOutNotifiction<rpc::RegisterBCParams>("registerBC");
+  ExprValue =
+      mkOutMethod<rpc::ExprValueParams, rpc::ExprValueResponse>("exprValue");
+  Registry.addNotification("ready", this, &EvalClient::onReady);
+}
+
 } // namespace nixd
diff --git a/nixd/tools/nixd/src/EvalClient.h b/nixd/tools/nixd/src/EvalClient.h
@@ -12,21 +12,27 @@
 namespace nixd {
 
 class EvalClient : public lspserver::LSPServer {
+  std::atomic<bool> Ready;
+
 public:
   llvm::unique_function<void(const rpc::RegisterBCParams &)> RegisterBC;
   llvm::unique_function<void(const rpc::ExprValueParams &,
                              lspserver::Callback<rpc::ExprValueResponse>)>
       ExprValue;
 
   EvalClient(std::unique_ptr<lspserver::InboundPort> In,
-             std::unique_ptr<lspserver::OutboundPort> Out)
-      : lspserver::LSPServer(std::move(In), std::move(Out)) {
-    RegisterBC = mkOutNotifiction<rpc::RegisterBCParams>("registerBC");
-    ExprValue =
-        mkOutMethod<rpc::ExprValueParams, rpc::ExprValueResponse>("exprValue");
-  }
+             std::unique_ptr<lspserver::OutboundPort> Out);
 
   virtual ~EvalClient() = default;
+
+  void onReady(const int &Flags) {
+    lspserver::log(
+        "nix-node-eval({0}) reported it's ready for processing requests",
+        Flags);
+    Ready = true;
+  }
+
+  bool ready() { return Ready; }
 };
 
 class OwnedEvalClient : public EvalClient {

diff --git a/nixd/tools/nixd/src/Hover.cpp b/nixd/tools/nixd/src/Hover.cpp
@@ -6,30 +6,42 @@
 #include "Controller.h"
 #include "Convert.h"
 
+#include <llvm/Support/Error.h>
+
+#include <boost/asio/post.hpp>
+
+#include <mutex>
+
 namespace nixd {
 
 using namespace llvm::json;
 using namespace lspserver;
+using namespace rpc;
 
 void Controller::onHover(const TextDocumentPositionParams &Params,
                          Callback<std::optional<Hover>> Reply) {
-  PathRef File = Params.textDocument.uri.file();
-  const nixf::Node *AST = TUs[File].ast().get();
-  nixf::Position Pos{Params.position.line, Params.position.character};
-  const nixf::Node *N = AST->descend({Pos, Pos});
-  if (!N) {
-    Reply(std::nullopt);
-    return;
-  }
-  std::string Name = N->name();
-  Reply(Hover{
-      .contents =
-          MarkupContent{
-              .kind = MarkupKind::Markdown,
-              .value = "`" + Name + "`",
-          },
-      .range = toLSPRange(N->range()),
-  });
+  auto Action = [Reply = std::move(Reply),
+                 File = std::string(Params.textDocument.uri.file()),
+                 RawPos = Params.position, this]() mutable {
+    std::lock_guard G(TUsLock);
+    const std::shared_ptr<nixf::Node> &AST = TUs[File].ast();
+    nixf::Position Pos{RawPos.line, RawPos.character};
+    const nixf::Node *N = AST->descend({Pos, Pos});
+    if (!N) {
+      Reply(std::nullopt);
+      return;
+    }
+    std::string Name = N->name();
+    Reply(Hover{
+        .contents =
+            MarkupContent{
+                .kind = MarkupKind::Markdown,
+                .value = "`" + Name + "`",
+            },
+        .range = toLSPRange(N->range()),
+    });
+  };
+  boost::asio::post(Pool, std::move(Action));
 }
 
 } // namespace nixd
diff --git a/nixd/tools/nixd/src/Main.cpp b/nixd/tools/nixd/src/Main.cpp
@@ -80,6 +80,9 @@ int main(int argc, char *argv[]) {
   auto Controller =
       std::make_unique<nixd::Controller>(std::move(In), std::move(Out));
 
+  if (LitTest)
+    Controller->setLitTest(LitTest);
+
   Controller->run();
 
   return 0;

diff --git a/nixd/tools/nixd/src/Support.cpp b/nixd/tools/nixd/src/Support.cpp
@@ -7,6 +7,10 @@
 #include "nixf/Bytecode/Write.h"
 #include "nixf/Parse/Parser.h"
 
+#include <boost/asio/post.hpp>
+
+#include <mutex>
+
 using namespace lspserver;
 using namespace nixd;
 
@@ -27,54 +31,70 @@ namespace nixd {
 
 void Controller::actOnDocumentAdd(PathRef File,
                                   std::optional<int64_t> Version) {
-  auto Draft = Store.getDraft(File);
-  assert(Draft && "Added document is not in the store?");
-
-  std::vector<nixf::Diagnostic> Diagnostics;
-  std::shared_ptr<nixf::Node> AST = nixf::parse(*Draft->Contents, Diagnostics);
-  publishDiagnostics(File, Version, Diagnostics);
-
-  if (!AST) {
-    TUs[File] = NixTU(std::move(Diagnostics), std::move(AST), std::nullopt);
-    return;
-  }
-
-  // Serialize the AST into shared memory. Prepare for evaluation.
-  std::stringstream OS;
-  nixf::writeBytecode(OS, *AST);
-  std::string Buf = OS.str();
-  if (Buf.empty()) {
-    lspserver::log("empty AST for {0}", File);
-    TUs[File] = NixTU(std::move(Diagnostics), std::move(AST), std::nullopt);
-    return;
-  }
-
-  // Create an mmap()-ed region, and write AST byte code there.
-  std::string ShmName = getShmName(File);
-
-  auto Shm = std::make_unique<util::AutoRemoveShm>(
-      ShmName, static_cast<bipc::offset_t>(Buf.size()));
-
-  auto Region =
-      std::make_unique<bipc::mapped_region>(Shm->get(), bipc::read_write);
-
-  std::memcpy(Region->get_address(), Buf.data(), Buf.size());
-
-  lspserver::log("serialized AST {0} to {1}, size: {2}", File, ShmName,
-                 Buf.size());
-
-  TUs[File] = NixTU(std::move(Diagnostics), std::move(AST),
-                    util::OwnedRegion{std::move(Shm), std::move(Region)});
-
-  if (Eval) {
-    Eval->RegisterBC(rpc::RegisterBCParams{
-        ShmName, ".", ".", static_cast<std::int64_t>(Buf.size())});
+  auto Action = [this, File = std::string(File), Version]() {
+    auto Draft = Store.getDraft(File);
+    assert(Draft && "Added document is not in the store?");
+
+    std::vector<nixf::Diagnostic> Diagnostics;
+    std::shared_ptr<nixf::Node> AST =
+        nixf::parse(*Draft->Contents, Diagnostics);
+    publishDiagnostics(File, Version, Diagnostics);
+
+    if (!AST) {
+      std::lock_guard G(TUsLock);
+      TUs[File] = NixTU(std::move(Diagnostics), std::move(AST), std::nullopt);
+      return;
+    }
+
+    // Serialize the AST into shared memory. Prepare for evaluation.
+    std::stringstream OS;
+    nixf::writeBytecode(OS, *AST);
+    std::string Buf = OS.str();
+    if (Buf.empty()) {
+      lspserver::log("empty AST for {0}", File);
+      std::lock_guard G(TUsLock);
+      TUs[File] = NixTU(std::move(Diagnostics), std::move(AST), std::nullopt);
+      return;
+    }
+
+    // Create an mmap()-ed region, and write AST byte code there.
+    std::string ShmName = getShmName(File);
+
+    auto Shm = std::make_unique<util::AutoRemoveShm>(
+        ShmName, static_cast<bipc::offset_t>(Buf.size()));
+
+    auto Region =
+        std::make_unique<bipc::mapped_region>(Shm->get(), bipc::read_write);
+
+    std::memcpy(Region->get_address(), Buf.data(), Buf.size());
+
+    lspserver::log("serialized AST {0} to {1}, size: {2}", File, ShmName,
+                   Buf.size());
+
+    std::lock_guard G(TUsLock);
+    TUs[File] = NixTU(std::move(Diagnostics), std::move(AST),
+                      util::OwnedRegion{std::move(Shm), std::move(Region)});
+
+    if (Eval) {
+      Eval->RegisterBC(rpc::RegisterBCParams{
+          ShmName, ".", ".", static_cast<std::int64_t>(Buf.size())});
+    }
+  };
+  if (LitTest) {
+    // In lit-testing mode we don't want to having requests processed before
+    // document updates. (e.g. Hover before parsing)
+    // So just invoke the action here.
+    Action();
+  } else {
+    // Otherwise we may want to concurently parse & serialize the file, so post
+    // it to the thread pool.
+    boost::asio::post(Pool, std::move(Action));
   }
 }
 
 Controller::Controller(std::unique_ptr<lspserver::InboundPort> In,
                        std::unique_ptr<lspserver::OutboundPort> Out)
-    : LSPServer(std::move(In), std::move(Out)) {
+    : LSPServer(std::move(In), std::move(Out)), LitTest(false) {
 
   // Life Cycle
   Registry.addMethod("initialize", this, &Controller::onInitialize);