[clangd] Lay JSONRPCDispatcher to rest.

[sammccall]

Summary

Most of its functionality is moved into ClangdLSPServer.
The decoupling between JSONRPCDispatcher, ProtocolCallbacks, ClangdLSPServer
was never real, and only served to obfuscate.

Some previous implicit/magic stuff is now explicit:

• the return type of LSP method calls are now in the signature
• no more reply() that gets the ID using global context magic
• arg tracing no longer relies on RequestArgs::stash context magic either

This is mostly refactoring, but some deliberate fixes while here:

• LSP method params are now by const reference
• notifications and calls are now distinct namespaces. (some tests had protocol errors and needed updating)
• we now reply to calls we failed to decode
• outgoing calls use distinct IDs

A few error codes and message IDs changed in unimportant ways (see tests).

[ioeric]

lgtm

clangd/ClangdLSPServer.cpp

148

do we really want to stash all responses in trace? It sounds like it could get pretty spammy (e.g. with completion results).

184

nit: while we are here, could you add some documentation for these two fields?

409

We are not translating internal errors from ClangdServer to LSP errors anymore. It seems fine and makes code cleaner. Just want to ask if this is intentional.

clangd/ClangdLSPServer.h

34

Just out of curiosity, why does DiagnosticsConsumer implemented in ClangdLSPServer instead of ClangdServer?

[sammccall]

clangd/ClangdLSPServer.cpp

148

It's a fair point, I don't know the answer. I know I've looked at the replies, but not all that often.

It's not new behavior though (it used to be in JSONRPCDispatcher), and I'd rather not change it in this patch.

409

We used to explicitly mark some of these with the "internal error" code, and others with the "unknown error" code, without much consistency as to why.

I've tried to preserve the error codes that were semantically significant (e.g. invalid params).

clangd/ClangdLSPServer.h

34

This is clangd::DiagnosticsConsumer, it's exactly the callback that ClangdServer uses to send diagnostics to ClangdLSPServer.

So ClangdServer doesn't implement it but rather accepts it, because it doesn't know what to do with the diagnostics.

clangd/CMakeLists.txt

add_clang_library(clangDaemon
   Diagnostics.cpp
   DraftStore.cpp
   FindSymbols.cpp
   FileDistance.cpp
   FS.cpp
   FuzzyMatch.cpp
   GlobalCompilationDatabase.cpp
   Headers.cpp
-  JSONRPCDispatcher.cpp
   JSONTransport.cpp
   Logger.cpp
   Protocol.cpp
-  ProtocolHandlers.cpp
   Quality.cpp
   RIFF.cpp
   SourceCode.cpp
   Threading.cpp
   Trace.cpp
   TUScheduler.cpp
   URI.cpp
   XRefs.cpp

clangd/ClangdLSPServer.h

 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDLSPSERVER_H
 
 #include "ClangdServer.h"
 #include "DraftStore.h"
 #include "FindSymbols.h"
 #include "GlobalCompilationDatabase.h"
 #include "Path.h"
 #include "Protocol.h"
-#include "ProtocolHandlers.h"
+#include "Transport.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "llvm/ADT/Optional.h"
 #include <memory>
 
 namespace clang {
 namespace clangd {
 
 class SymbolIndex;
 
 /// This class exposes ClangdServer's capabilities via Language Server Protocol.
 ///
-/// JSONRPCDispatcher binds the implemented ProtocolCallbacks methods
-/// (e.g. onInitialize) to corresponding JSON-RPC methods ("initialize").
-/// The server also supports $/cancelRequest (JSONRPCDispatcher provides this).
-class ClangdLSPServer : private DiagnosticsConsumer, private ProtocolCallbacks {
+/// MessageHandler binds the implemented LSP methods (e.g. onInitialize) to
+/// corresponding JSON-RPC methods ("initialize").
+/// The server also supports $/cancelRequest (MessageHandler provides this).
+class ClangdLSPServer : private DiagnosticsConsumer {
 public:
   /// If \p CompileCommandsDir has a value, compile_commands.json will be
   /// loaded only from \p CompileCommandsDir. Otherwise, clangd will look
   /// for compile_commands.json in all parent directories of each file.
   ClangdLSPServer(Transport &Transp, const clangd::CodeCompleteOptions &CCOpts,
                   llvm::Optional<Path> CompileCommandsDir,
                   bool ShouldUseInMemoryCDB, const ClangdServer::Options &Opts);
+  ~ClangdLSPServer();
 
   /// Run LSP server loop, communicating with the Transport provided in the
   /// constructor. This method must not be executed more than once.
   ///
   /// \return Whether we shut down cleanly with a 'shutdown' -> 'exit' sequence.
   bool run();
 
 private:
   // Implement DiagnosticsConsumer.
   void onDiagnosticsReady(PathRef File, std::vector<Diag> Diagnostics) override;
 
-  // Implement ProtocolCallbacks.
-  void onInitialize(InitializeParams &Params) override;
-  void onShutdown(ShutdownParams &Params) override;
-  void onExit(ExitParams &Params) override;
-  void onDocumentDidOpen(DidOpenTextDocumentParams &Params) override;
-  void onDocumentDidChange(DidChangeTextDocumentParams &Params) override;
-  void onDocumentDidClose(DidCloseTextDocumentParams &Params) override;
-  void
-  onDocumentOnTypeFormatting(DocumentOnTypeFormattingParams &Params) override;
-  void
-  onDocumentRangeFormatting(DocumentRangeFormattingParams &Params) override;
-  void onDocumentFormatting(DocumentFormattingParams &Params) override;
-  void onDocumentSymbol(DocumentSymbolParams &Params) override;
-  void onCodeAction(CodeActionParams &Params) override;
-  void onCompletion(TextDocumentPositionParams &Params) override;
-  void onSignatureHelp(TextDocumentPositionParams &Params) override;
-  void onGoToDefinition(TextDocumentPositionParams &Params) override;
-  void onReference(ReferenceParams &Params) override;
-  void onSwitchSourceHeader(TextDocumentIdentifier &Params) override;
-  void onDocumentHighlight(TextDocumentPositionParams &Params) override;
-  void onFileEvent(DidChangeWatchedFilesParams &Params) override;
-  void onCommand(ExecuteCommandParams &Params) override;
-  void onWorkspaceSymbol(WorkspaceSymbolParams &Params) override;
-  void onRename(RenameParams &Parames) override;
-  void onHover(TextDocumentPositionParams &Params) override;
-  void onChangeConfiguration(DidChangeConfigurationParams &Params) override;
+  // LSP methods. Notifications have signature void(const Params&).
+  // Calls have signature void(const Params&, Callback<Response>).
+  void onInitialize(const InitializeParams &, Callback<llvm::json::Value>);
+  void onShutdown(const ShutdownParams &, Callback<std::nullptr_t>);
+  void onDocumentDidOpen(const DidOpenTextDocumentParams &);
+  void onDocumentDidChange(const DidChangeTextDocumentParams &);
+  void onDocumentDidClose(const DidCloseTextDocumentParams &);
+  void onDocumentOnTypeFormatting(const DocumentOnTypeFormattingParams &,
+                                  Callback<std::vector<TextEdit>>);
+  void onDocumentRangeFormatting(const DocumentRangeFormattingParams &,
+                                 Callback<std::vector<TextEdit>>);
+  void onDocumentFormatting(const DocumentFormattingParams &,
+                            Callback<std::vector<TextEdit>>);
+  void onDocumentSymbol(const DocumentSymbolParams &,
+                        Callback<std::vector<SymbolInformation>>);
+  void onCodeAction(const CodeActionParams &, Callback<llvm::json::Value>);
+  void onCompletion(const TextDocumentPositionParams &,
+                    Callback<CompletionList>);
+  void onSignatureHelp(const TextDocumentPositionParams &,
+                       Callback<SignatureHelp>);
+  void onGoToDefinition(const TextDocumentPositionParams &,
+                        Callback<std::vector<Location>>);
+  void onReference(const ReferenceParams &, Callback<std::vector<Location>>);
+  void onSwitchSourceHeader(const TextDocumentIdentifier &,
+                            Callback<std::string>);
+  void onDocumentHighlight(const TextDocumentPositionParams &,
+                           Callback<std::vector<DocumentHighlight>>);
+  void onFileEvent(const DidChangeWatchedFilesParams &);
+  void onCommand(const ExecuteCommandParams &, Callback<llvm::json::Value>);
+  void onWorkspaceSymbol(const WorkspaceSymbolParams &,
+                         Callback<std::vector<SymbolInformation>>);
+  void onRename(const RenameParams &, Callback<WorkspaceEdit>);
+  void onHover(const TextDocumentPositionParams &,
+               Callback<llvm::Optional<Hover>>);
+  void onChangeConfiguration(const DidChangeConfigurationParams &);
 
   std::vector<Fix> getFixes(StringRef File, const clangd::Diagnostic &D);
 
   /// Forces a reparse of all currently opened files.  As a result, this method
   /// may be very expensive.  This method is normally called when the
   /// compilation database is changed.
   void reparseOpenedFiles();
   void applyConfiguration(const ClangdConfigurationParamsChange &Settings);
private:
     // If IsDirectoryBased is false, an instance of DirectoryBasedCDB.
     // unique_ptr<GlobalCompilationDatabase> CDB;
     std::unique_ptr<GlobalCompilationDatabase> CDB;
     // Non-null only for directory-based CDB
     std::unique_ptr<CachingCompilationDb> CachingCDB;
     bool IsDirectoryBased;
   };
 
+  // Most code should not deal with Transport directly.
+  // MessageHandler deals with incoming messages, use call() etc for outgoing.
   clangd::Transport &Transp;
+  class MessageHandler;
+  std::unique_ptr<MessageHandler> MsgHandler;
+  std::atomic<int> NextCallID = {0};
+  std::mutex TranspWriter;
+  void call(StringRef Method, llvm::json::Value Params);
+  void notify(StringRef Method, llvm::json::Value Params);
+  void reply(llvm::json::Value ID, llvm::Expected<llvm::json::Value> Result);
+
   // Various ClangdServer parameters go here. It's important they're created
   // before ClangdServer.
   CompilationDB CDB;
 
   RealFileSystemProvider FSProvider;
   /// Options used for code completion
   clangd::CodeCompleteOptions CCOpts;
   /// Options used for diagnostics.

clangd/ClangdLSPServer.cpp

 //===--- ClangdLSPServer.cpp - LSP server ------------------------*- C++-*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #include "ClangdLSPServer.h"
 #include "Diagnostics.h"
-#include "JSONRPCDispatcher.h"
 #include "SourceCode.h"
+#include "Trace.h"
 #include "URI.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Path.h"
+#include "llvm/Support/ScopedPrinter.h"
 
 using namespace clang::clangd;
 using namespace clang;
 using namespace llvm;
 
 namespace {
 
 /// \brief Supports a test URI scheme with relaxed constraints for lit tests.
CompletionItemKindBitset defaultCompletionItemKinds() {
   for (size_t I = CompletionItemKindMin;
        I <= static_cast<size_t>(CompletionItemKind::Reference); ++I)
     Defaults.set(I);
   return Defaults;
 }
 
 } // namespace
 
-void ClangdLSPServer::onInitialize(InitializeParams &Params) {
+// MessageHandler dispatches incoming LSP messages.
+// It handles cross-cutting concerns:
+//  - serializes/deserializes protocol objects to JSON
+//  - logging of inbound messages
+//  - cancellation handling
+//  - basic call tracing
+class ClangdLSPServer::MessageHandler : public Transport::MessageHandler {
+public:
+  MessageHandler(ClangdLSPServer &Server) : Server(Server) {}
+
+  bool onNotify(StringRef Method, json::Value Params) override {
+    log("<-- {0}", Method);
+    if (Method == "exit")
+      return false;
+    if (Method == "$/cancelRequest")
+      onCancel(std::move(Params));
+    else if (auto Handler = Notifications.lookup(Method))
+      Handler(std::move(Params));
+    else
+      log("unhandled notification {0}", Method);
+    return true;
+  }
+
+  bool onCall(StringRef Method, json::Value Params, json::Value ID) override {
+    log("<-- {0}({1})", Method, ID);
+    if (auto Handler = Calls.lookup(Method))
+      Handler(std::move(Params), std::move(ID));
+    else
+      Server.reply(ID, llvm::make_error<LSPError>("method not found",
+                                                  ErrorCode::MethodNotFound));
+    return true;
+  }
+
+  bool onReply(json::Value ID, Expected<json::Value> Result) override {
+    // We ignore replies, just log them.
+    if (Result)
+      log("<-- reply({0})", ID);
+    else
+      log("<-- reply({0}) error: {1}", ID, llvm::toString(Result.takeError()));
+    return true;
+  }
+
+  // Bind an LSP method name to a call.
+  template <typename Param, typename Reply>
+  void bind(const char *Method,
+            void (ClangdLSPServer::*Handler)(const Param &, Callback<Reply>)) {
+    Calls[Method] = [Method, Handler, this](json::Value RawParams,
+                                            json::Value ID) {
+      Param P;
+      if (!fromJSON(RawParams, P)) {
+        elog("Failed to decode {0} request.", Method);
+        Server.reply(ID, make_error<LSPError>("failed to decode request",
+                                              ErrorCode::InvalidRequest));
+        return;
+      }
+      trace::Span Tracer(Method);
+      SPAN_ATTACH(Tracer, "Params", RawParams);
+      auto *Trace = Tracer.Args; // We attach reply from another thread.
+      // Calls can be canceled by the client. Add cancellation context.
+      WithContext WithCancel(cancelableRequestContext(ID));
+      // FIXME: this function should assert it's called exactly once.
+      (Server.*Handler)(P, [this, ID, Trace](llvm::Expected<Reply> Result) {
+        if (Result) {
+          if (Trace)
+            (*Trace)["Reply"] = *Result;
+          Server.reply(ID, json::Value(std::move(*Result)));
+        } else {
+          auto Err = Result.takeError();
+          if (Trace)
+            (*Trace)["Error"] = llvm::to_string(Err);
+          Server.reply(ID, std::move(Err));
+        }
+      });
+    };
+  }
+
+  // Bind an LSP method name to a notification.
+  template <typename Param>
+  void bind(const char *Method,
+            void (ClangdLSPServer::*Handler)(const Param &)) {
+    Notifications[Method] = [Method, Handler, this](json::Value RawParams) {
+      Param P;
+      if (!fromJSON(RawParams, P)) {
+        elog("Failed to decode {0} request.", Method);
+        return;
+      }
+      trace::Span Tracer(Method);
+      SPAN_ATTACH(Tracer, "Params", RawParams);
+      (Server.*Handler)(P);
+    };
+  }
+
+private:
+  llvm::StringMap<std::function<void(json::Value)>> Notifications;
+  llvm::StringMap<std::function<void(json::Value, json::Value)>> Calls;
+
+  // Method calls may be cancelled by ID, so keep track of their state.
+  // This needs a mutex: handlers may finish on a different thread, and that's
+  // when we clean up entries in the map.
+  mutable std::mutex RequestCancelersMutex;
+  llvm::StringMap<std::pair<Canceler, unsigned>> RequestCancelers;
+  unsigned NextRequestCookie = 0;
+  void onCancel(const llvm::json::Value &Params) {
+    const json::Value *ID = nullptr;
+    if (auto *O = Params.getAsObject())
+      ID = O->get("id");
+    if (!ID) {
+      elog("Bad cancellation request: {0}", Params);
+      return;
+    }
+    auto StrID = llvm::to_string(*ID);
+    std::lock_guard<std::mutex> Lock(RequestCancelersMutex);
+    auto It = RequestCancelers.find(StrID);
+    if (It != RequestCancelers.end())
+      It->second.first(); // Invoke the canceler.
+  }
+  // We run cancelable requests in a context that does two things:
+  //  - allows cancellation using RequestCancelers[ID]
+  //  - cleans up the entry in RequestCancelers when it's no longer needed
+  // If a client reuses an ID, the last wins and the first cannot be canceled.
+  Context cancelableRequestContext(const json::Value &ID) {
+    auto Task = cancelableTask();
+    auto StrID = llvm::to_string(ID);  // JSON-serialize ID for map key.
+    auto Cookie = NextRequestCookie++; // No lock, only called on main thread.
+    {
+      std::lock_guard<std::mutex> Lock(RequestCancelersMutex);
+      RequestCancelers[StrID] = {std::move(Task.second), Cookie};
+    }
+    // When the request ends, we can clean up the entry we just added.
+    // The cookie lets us check that it hasn't been overwritten due to ID
+    // reuse.
+    return Task.first.derive(make_scope_exit([this, StrID, Cookie] {
+      std::lock_guard<std::mutex> Lock(RequestCancelersMutex);
+      auto It = RequestCancelers.find(StrID);
+      if (It != RequestCancelers.end() && It->second.second == Cookie)
+        RequestCancelers.erase(It);
+    }));
+  }
+
+  ClangdLSPServer &Server;
+};
+
+// call(), notify(), and reply() wrap the Transport, adding logging and locking.
+void ClangdLSPServer::call(StringRef Method, json::Value Params) {
+  auto ID = NextCallID++;
+  log("--> {0}({1})", Method, ID);
+  // We currently don't handle responses, so no need to store ID anywhere.
+  std::lock_guard<std::mutex> Lock(TranspWriter);
+  Transp.call(Method, std::move(Params), ID);
+}
+
+void ClangdLSPServer::notify(StringRef Method, json::Value Params) {
+  log("--> {0}", Method);
+  std::lock_guard<std::mutex> Lock(TranspWriter);
+  Transp.notify(Method, std::move(Params));
+}
+
+void ClangdLSPServer::reply(llvm::json::Value ID,
+                            llvm::Expected<llvm::json::Value> Result) {
+  if (Result) {
+    log("--> reply({0})", ID);
+    std::lock_guard<std::mutex> Lock(TranspWriter);
+    Transp.reply(std::move(ID), std::move(Result));
+  } else {
+    Error Err = Result.takeError();
+    log("--> reply({0}) error: {1}", ID, Err);
+    std::lock_guard<std::mutex> Lock(TranspWriter);
+    Transp.reply(std::move(ID), std::move(Err));
+  }
+}
+
+void ClangdLSPServer::onInitialize(const InitializeParams &Params,
+                                   Callback<json::Value> Reply) {
   if (Params.initializationOptions) {
     const ClangdInitializationOptions &Opts = *Params.initializationOptions;
 
     // Explicit compilation database path.
     if (Opts.compilationDatabasePath.hasValue()) {
       CDB.setCompileCommandsDir(Opts.compilationDatabasePath.getValue());
     }
 
void ClangdLSPServer::onInitialize(const InitializeParams &Params,
   DiagOpts.EmbedFixesInDiagnostics = Params.capabilities.DiagnosticFixes;
   DiagOpts.SendDiagnosticCategory = Params.capabilities.DiagnosticCategory;
   if (Params.capabilities.WorkspaceSymbolKinds)
     SupportedSymbolKinds |= *Params.capabilities.WorkspaceSymbolKinds;
   if (Params.capabilities.CompletionItemKinds)
     SupportedCompletionItemKinds |= *Params.capabilities.CompletionItemKinds;
   SupportsCodeAction = Params.capabilities.CodeActionStructure;
 
-  reply(json::Object{
+  Reply(json::Object{
       {{"capabilities",
         json::Object{
             {"textDocumentSync", (int)TextDocumentSyncKind::Incremental},
             {"documentFormattingProvider", true},
             {"documentRangeFormattingProvider", true},
             {"documentOnTypeFormattingProvider",
              json::Object{
                  {"firstTriggerCharacter", "}"},
Reply(json::Object{
             {"referencesProvider", true},
             {"executeCommandProvider",
              json::Object{
                  {"commands", {ExecuteCommandParams::CLANGD_APPLY_FIX_COMMAND}},
              }},
         }}}});
 }
 
-void ClangdLSPServer::onShutdown(ShutdownParams &Params) {
+void ClangdLSPServer::onShutdown(const ShutdownParams &Params,
+                                 Callback<std::nullptr_t> Reply) {
   // Do essentially nothing, just say we're ready to exit.
   ShutdownRequestReceived = true;
-  reply(nullptr);
-}
-
-void ClangdLSPServer::onExit(ExitParams &Params) {
-  // No work to do.
-  // JSONRPCDispatcher shuts down the transport after this notification.
+  Reply(nullptr);
 }
 
-void ClangdLSPServer::onDocumentDidOpen(DidOpenTextDocumentParams &Params) {
+void ClangdLSPServer::onDocumentDidOpen(
+    const DidOpenTextDocumentParams &Params) {
   PathRef File = Params.textDocument.uri.file();
   if (Params.metadata && !Params.metadata->extraFlags.empty())
     CDB.setExtraFlagsForFile(File, std::move(Params.metadata->extraFlags));
 
-  std::string &Contents = Params.textDocument.text;
+  const std::string &Contents = Params.textDocument.text;
 
   DraftMgr.addDraft(File, Contents);
   Server->addDocument(File, Contents, WantDiagnostics::Yes);
 }
 
-void ClangdLSPServer::onDocumentDidChange(DidChangeTextDocumentParams &Params) {
+void ClangdLSPServer::onDocumentDidChange(
+    const DidChangeTextDocumentParams &Params) {
   auto WantDiags = WantDiagnostics::Auto;
   if (Params.wantDiagnostics.hasValue())
     WantDiags = Params.wantDiagnostics.getValue() ? WantDiagnostics::Yes
                                                   : WantDiagnostics::No;
 
   PathRef File = Params.textDocument.uri.file();
   llvm::Expected<std::string> Contents =
       DraftMgr.updateDraft(File, Params.contentChanges);
   if (!Contents) {
     // If this fails, we are most likely going to be not in sync anymore with
     // the client.  It is better to remove the draft and let further operations
     // fail rather than giving wrong results.
     DraftMgr.removeDraft(File);
     Server->removeDocument(File);
     CDB.invalidate(File);
     elog("Failed to update {0}: {1}", File, Contents.takeError());
     return;
   }
 
   Server->addDocument(File, *Contents, WantDiags);
 }
 
-void ClangdLSPServer::onFileEvent(DidChangeWatchedFilesParams &Params) {
+void ClangdLSPServer::onFileEvent(const DidChangeWatchedFilesParams &Params) {
   Server->onFileEvent(Params);
 }
 
-void ClangdLSPServer::onCommand(ExecuteCommandParams &Params) {
-  auto ApplyEdit = [](WorkspaceEdit WE) {
+void ClangdLSPServer::onCommand(const ExecuteCommandParams &Params,
+                                Callback<json::Value> Reply) {
+  auto ApplyEdit = [&](WorkspaceEdit WE) {
     ApplyWorkspaceEditParams Edit;
     Edit.edit = std::move(WE);
-    // We don't need the response so id == 1 is OK.
     // Ideally, we would wait for the response and if there is no error, we
     // would reply success/failure to the original RPC.
     call("workspace/applyEdit", Edit);
   };
   if (Params.command == ExecuteCommandParams::CLANGD_APPLY_FIX_COMMAND &&
       Params.workspaceEdit) {
     // The flow for "apply-fix" :
     // 1. We publish a diagnostic, including fixits
     // 2. The user clicks on the diagnostic, the editor asks us for code actions
     // 3. We send code actions, with the fixit embedded as context
     // 4. The user selects the fixit, the editor asks us to apply it
     // 5. We unwrap the changes and send them back to the editor
     // 6. The editor applies the changes (applyEdit), and sends us a reply (but
     // we ignore it)
 
-    reply("Fix applied.");
+    Reply("Fix applied.");
     ApplyEdit(*Params.workspaceEdit);
   } else {
     // We should not get here because ExecuteCommandParams would not have
     // parsed in the first place and this handler should not be called. But if
     // more commands are added, this will be here has a safe guard.
-    replyError(
-        ErrorCode::InvalidParams,
-        llvm::formatv("Unsupported command \"{0}\".", Params.command).str());
+    Reply(make_error<LSPError>(
+        llvm::formatv("Unsupported command \"{0}\".", Params.command).str(),
+        ErrorCode::InvalidParams));
   }
 }
 
-void ClangdLSPServer::onWorkspaceSymbol(WorkspaceSymbolParams &Params) {
+void ClangdLSPServer::onWorkspaceSymbol(
+    const WorkspaceSymbolParams &Params,
+    Callback<std::vector<SymbolInformation>> Reply) {
   Server->workspaceSymbols(
       Params.query, CCOpts.Limit,
-      [this](llvm::Expected<std::vector<SymbolInformation>> Items) {
+      Bind(
+          [this](decltype(Reply) Reply,
+                 llvm::Expected<std::vector<SymbolInformation>> Items) {
             if (!Items)
-          return replyError(ErrorCode::InternalError,
-                            llvm::toString(Items.takeError()));
+              return Reply(Items.takeError());
             for (auto &Sym : *Items)
               Sym.kind = adjustKindToCapability(Sym.kind, SupportedSymbolKinds);
 
-        reply(json::Array(*Items));
-      });
+            Reply(std::move(*Items));
+          },
+          std::move(Reply)));
 }
 
-void ClangdLSPServer::onRename(RenameParams &Params) {
+void ClangdLSPServer::onRename(const RenameParams &Params,
+                               Callback<WorkspaceEdit> Reply) {
   Path File = Params.textDocument.uri.file();
   llvm::Optional<std::string> Code = DraftMgr.getDraft(File);
   if (!Code)
-    return replyError(ErrorCode::InvalidParams,
-                      "onRename called for non-added file");
+    return Reply(make_error<LSPError>("onRename called for non-added file",
+                                      ErrorCode::InvalidParams));
 
   Server->rename(
       File, Params.position, Params.newName,
-      [File, Code,
-       Params](llvm::Expected<std::vector<tooling::Replacement>> Replacements) {
+      Bind(
+          [File, Code, Params](
+              decltype(Reply) Reply,
+              llvm::Expected<std::vector<tooling::Replacement>> Replacements) {
             if (!Replacements)
-          return replyError(ErrorCode::InternalError,
-                            llvm::toString(Replacements.takeError()));
+              return Reply(Replacements.takeError());
 
             // Turn the replacements into the format specified by the Language
             // Server Protocol. Fuse them into one big JSON array.
             std::vector<TextEdit> Edits;
             for (const auto &R : *Replacements)
               Edits.push_back(replacementToEdit(*Code, R));
             WorkspaceEdit WE;
             WE.changes = {{Params.textDocument.uri.uri(), Edits}};
-        reply(WE);
-      });
+            Reply(WE);
+          },
+          std::move(Reply)));
 }
 
-void ClangdLSPServer::onDocumentDidClose(DidCloseTextDocumentParams &Params) {
+void ClangdLSPServer::onDocumentDidClose(
+    const DidCloseTextDocumentParams &Params) {
   PathRef File = Params.textDocument.uri.file();
   DraftMgr.removeDraft(File);
   Server->removeDocument(File);
   CDB.invalidate(File);
 }
 
 void ClangdLSPServer::onDocumentOnTypeFormatting(
-    DocumentOnTypeFormattingParams &Params) {
+    const DocumentOnTypeFormattingParams &Params,
+    Callback<std::vector<TextEdit>> Reply) {
   auto File = Params.textDocument.uri.file();
   auto Code = DraftMgr.getDraft(File);
   if (!Code)
-    return replyError(ErrorCode::InvalidParams,
-                      "onDocumentOnTypeFormatting called for non-added file");
+    return Reply(make_error<LSPError>(
+        "onDocumentOnTypeFormatting called for non-added file",
+        ErrorCode::InvalidParams));
 
   auto ReplacementsOrError = Server->formatOnType(*Code, File, Params.position);
   if (ReplacementsOrError)
-    reply(json::Array(replacementsToEdits(*Code, ReplacementsOrError.get())));
+    Reply(replacementsToEdits(*Code, ReplacementsOrError.get()));
   else
-    replyError(ErrorCode::UnknownErrorCode,
-               llvm::toString(ReplacementsOrError.takeError()));
+    Reply(ReplacementsOrError.takeError());
 }
 
 void ClangdLSPServer::onDocumentRangeFormatting(
-    DocumentRangeFormattingParams &Params) {
+    const DocumentRangeFormattingParams &Params,
+    Callback<std::vector<TextEdit>> Reply) {
   auto File = Params.textDocument.uri.file();
   auto Code = DraftMgr.getDraft(File);
   if (!Code)
-    return replyError(ErrorCode::InvalidParams,
-                      "onDocumentRangeFormatting called for non-added file");
+    return Reply(make_error<LSPError>(
+        "onDocumentRangeFormatting called for non-added file",
+        ErrorCode::InvalidParams));
 
   auto ReplacementsOrError = Server->formatRange(*Code, File, Params.range);
   if (ReplacementsOrError)
-    reply(json::Array(replacementsToEdits(*Code, ReplacementsOrError.get())));
+    Reply(replacementsToEdits(*Code, ReplacementsOrError.get()));
   else
-    replyError(ErrorCode::UnknownErrorCode,
-               llvm::toString(ReplacementsOrError.takeError()));
+    Reply(ReplacementsOrError.takeError());
 }
 
-void ClangdLSPServer::onDocumentFormatting(DocumentFormattingParams &Params) {
+void ClangdLSPServer::onDocumentFormatting(
+    const DocumentFormattingParams &Params,
+    Callback<std::vector<TextEdit>> Reply) {
   auto File = Params.textDocument.uri.file();
   auto Code = DraftMgr.getDraft(File);
   if (!Code)
-    return replyError(ErrorCode::InvalidParams,
-                      "onDocumentFormatting called for non-added file");
+    return Reply(
+        make_error<LSPError>("onDocumentFormatting called for non-added file",
+                             ErrorCode::InvalidParams));
 
   auto ReplacementsOrError = Server->formatFile(*Code, File);
   if (ReplacementsOrError)
-    reply(json::Array(replacementsToEdits(*Code, ReplacementsOrError.get())));
+    Reply(replacementsToEdits(*Code, ReplacementsOrError.get()));
   else
-    replyError(ErrorCode::UnknownErrorCode,
-               llvm::toString(ReplacementsOrError.takeError()));
+    Reply(ReplacementsOrError.takeError());
 }
 
-void ClangdLSPServer::onDocumentSymbol(DocumentSymbolParams &Params) {
+void ClangdLSPServer::onDocumentSymbol(
+    const DocumentSymbolParams &Params,
+    Callback<std::vector<SymbolInformation>> Reply) {
   Server->documentSymbols(
       Params.textDocument.uri.file(),
-      [this](llvm::Expected<std::vector<SymbolInformation>> Items) {
+      Bind(
+          [this](decltype(Reply) Reply,
+                 llvm::Expected<std::vector<SymbolInformation>> Items) {
             if (!Items)
-          return replyError(ErrorCode::InvalidParams,
-                            llvm::toString(Items.takeError()));
+              return Reply(Items.takeError());
             for (auto &Sym : *Items)
               Sym.kind = adjustKindToCapability(Sym.kind, SupportedSymbolKinds);
-        reply(json::Array(*Items));
-      });
+            Reply(std::move(*Items));
+          },
+          std::move(Reply)));
 }
 
 static Optional<Command> asCommand(const CodeAction &Action) {
   Command Cmd;
   if (Action.command && Action.edit)
     return llvm::None; // Not representable. (We never emit these anyway).
   if (Action.command) {
     Cmd = *Action.command;
   } else if (Action.edit) {
     Cmd.command = Command::CLANGD_APPLY_FIX_COMMAND;
     Cmd.workspaceEdit = *Action.edit;
   } else {
     return llvm::None;
   }
   Cmd.title = Action.title;
   if (Action.kind && *Action.kind == CodeAction::QUICKFIX_KIND)
     Cmd.title = "Apply fix: " + Cmd.title;
   return Cmd;
 }
 
-void ClangdLSPServer::onCodeAction(CodeActionParams &Params) {
+void ClangdLSPServer::onCodeAction(const CodeActionParams &Params,
+                                   Callback<json::Value> Reply) {
+  auto Code = DraftMgr.getDraft(Params.textDocument.uri.file());
+  if (!Code)
+    return Reply(make_error<LSPError>("onCodeAction called for non-added file",
+                                      ErrorCode::InvalidParams));
   // We provide a code action for Fixes on the specified diagnostics.
-  if (!DraftMgr.getDraft(Params.textDocument.uri.file()))
-    return replyError(ErrorCode::InvalidParams,
-                      "onCodeAction called for non-added file");
-
   std::vector<CodeAction> Actions;
-  for (Diagnostic &D : Params.context.diagnostics) {
+  for (const Diagnostic &D : Params.context.diagnostics) {
     for (auto &F : getFixes(Params.textDocument.uri.file(), D)) {
       Actions.emplace_back();
       Actions.back().title = F.Message;
       Actions.back().kind = CodeAction::QUICKFIX_KIND;
       Actions.back().diagnostics = {D};
       Actions.back().edit.emplace();
       Actions.back().edit->changes.emplace();
       (*Actions.back().edit->changes)[Params.textDocument.uri.uri()] = {
           F.Edits.begin(), F.Edits.end()};
     }
   }
 
   if (SupportsCodeAction)
-    reply(json::Array(Actions));
+    Reply(json::Array(Actions));
   else {
     std::vector<Command> Commands;
     for (const auto &Action : Actions)
       if (auto Command = asCommand(Action))
         Commands.push_back(std::move(*Command));
-    reply(json::Array(Commands));
+    Reply(json::Array(Commands));
   }
 }
 
-void ClangdLSPServer::onCompletion(TextDocumentPositionParams &Params) {
+void ClangdLSPServer::onCompletion(const TextDocumentPositionParams &Params,
+                                   Callback<CompletionList> Reply) {
   Server->codeComplete(Params.textDocument.uri.file(), Params.position, CCOpts,
-                       [this](llvm::Expected<CodeCompleteResult> List) {
+                       Bind(
+                           [this](decltype(Reply) Reply,
+                                  llvm::Expected<CodeCompleteResult> List) {
                              if (!List)
-                           return replyError(List.takeError());
+                               return Reply(List.takeError());
                              CompletionList LSPList;
                              LSPList.isIncomplete = List->HasMore;
                              for (const auto &R : List->Completions) {
                                CompletionItem C = R.render(CCOpts);
                                C.kind = adjustKindToCapability(
                                    C.kind, SupportedCompletionItemKinds);
                                LSPList.items.push_back(std::move(C));
                              }
-                         return reply(std::move(LSPList));
-                       });
+                             return Reply(std::move(LSPList));
+                           },
+                           std::move(Reply)));
 }
 
-void ClangdLSPServer::onSignatureHelp(TextDocumentPositionParams &Params) {
+void ClangdLSPServer::onSignatureHelp(const TextDocumentPositionParams &Params,
+                                      Callback<SignatureHelp> Reply) {
   Server->signatureHelp(Params.textDocument.uri.file(), Params.position,
-                        [](llvm::Expected<SignatureHelp> SignatureHelp) {
-                          if (!SignatureHelp)
-                            return replyError(
-                                ErrorCode::InvalidParams,
-                                llvm::toString(SignatureHelp.takeError()));
-                          reply(*SignatureHelp);
-                        });
+                        std::move(Reply));
 }
 
-void ClangdLSPServer::onGoToDefinition(TextDocumentPositionParams &Params) {
+void ClangdLSPServer::onGoToDefinition(const TextDocumentPositionParams &Params,
+                                       Callback<std::vector<Location>> Reply) {
   Server->findDefinitions(Params.textDocument.uri.file(), Params.position,
-                          [](llvm::Expected<std::vector<Location>> Items) {
-                            if (!Items)
-                              return replyError(
-                                  ErrorCode::InvalidParams,
-                                  llvm::toString(Items.takeError()));
-                            reply(json::Array(*Items));
-                          });
+                          std::move(Reply));
 }
 
-void ClangdLSPServer::onSwitchSourceHeader(TextDocumentIdentifier &Params) {
+void ClangdLSPServer::onSwitchSourceHeader(const TextDocumentIdentifier &Params,
+                                           Callback<std::string> Reply) {
   llvm::Optional<Path> Result = Server->switchSourceHeader(Params.uri.file());
-  reply(Result ? URI::createFile(*Result).toString() : "");
+  Reply(Result ? URI::createFile(*Result).toString() : "");
 }
 
-void ClangdLSPServer::onDocumentHighlight(TextDocumentPositionParams &Params) {
-  Server->findDocumentHighlights(
-      Params.textDocument.uri.file(), Params.position,
-      [](llvm::Expected<std::vector<DocumentHighlight>> Highlights) {
-        if (!Highlights)
-          return replyError(ErrorCode::InternalError,
-                            llvm::toString(Highlights.takeError()));
-        reply(json::Array(*Highlights));
-      });
+void ClangdLSPServer::onDocumentHighlight(
+    const TextDocumentPositionParams &Params,
+    Callback<std::vector<DocumentHighlight>> Reply) {
+  Server->findDocumentHighlights(Params.textDocument.uri.file(),
+                                 Params.position, std::move(Reply));
 }
 
-void ClangdLSPServer::onHover(TextDocumentPositionParams &Params) {
+void ClangdLSPServer::onHover(const TextDocumentPositionParams &Params,
+                              Callback<llvm::Optional<Hover>> Reply) {
   Server->findHover(Params.textDocument.uri.file(), Params.position,
-                    [](llvm::Expected<llvm::Optional<Hover>> H) {
-                      if (!H) {
-                        replyError(ErrorCode::InternalError,
-                                   llvm::toString(H.takeError()));
-                        return;
-                      }
-
-                      reply(*H);
-                    });
+                    std::move(Reply));
 }
 
 void ClangdLSPServer::applyConfiguration(
     const ClangdConfigurationParamsChange &Params) {
   // Per-file update to the compilation database.
   if (Params.compilationDatabaseChanges) {
     const auto &CompileCommandUpdates = *Params.compilationDatabaseChanges;
     bool ShouldReparseOpenFiles = false;
if (Params.compilationDatabaseChanges) {
     }
     if (ShouldReparseOpenFiles)
       reparseOpenedFiles();
   }
 }
 
 // FIXME: This function needs to be properly tested.
 void ClangdLSPServer::onChangeConfiguration(
-    DidChangeConfigurationParams &Params) {
+    const DidChangeConfigurationParams &Params) {
   applyConfiguration(Params.settings);
 }
 
-void ClangdLSPServer::onReference(ReferenceParams &Params) {
+void ClangdLSPServer::onReference(const ReferenceParams &Params,
+                                  Callback<std::vector<Location>> Reply) {
   Server->findReferences(Params.textDocument.uri.file(), Params.position,
-                         [](llvm::Expected<std::vector<Location>> Locations) {
-                           if (!Locations)
-                             return replyError(
-                                 ErrorCode::InternalError,
-                                 llvm::toString(Locations.takeError()));
-                           reply(llvm::json::Array(*Locations));
-                         });
+                         std::move(Reply));
 }
 
 ClangdLSPServer::ClangdLSPServer(class Transport &Transp,
                                  const clangd::CodeCompleteOptions &CCOpts,
                                  llvm::Optional<Path> CompileCommandsDir,
                                  bool ShouldUseInMemoryCDB,
                                  const ClangdServer::Options &Opts)
-    : Transp(Transp),
+    : Transp(Transp), MsgHandler(new MessageHandler(*this)),
       CDB(ShouldUseInMemoryCDB ? CompilationDB::makeInMemory()
                                : CompilationDB::makeDirectoryBased(
                                      std::move(CompileCommandsDir))),
       CCOpts(CCOpts), SupportedSymbolKinds(defaultSymbolKinds()),
       SupportedCompletionItemKinds(defaultCompletionItemKinds()),
       Server(new ClangdServer(CDB.getCDB(), FSProvider, /*DiagConsumer=*/*this,
-                              Opts)) {}
+                              Opts)) {
+  // clang-format off
+  MsgHandler->bind("initialize", &ClangdLSPServer::onInitialize);
+  MsgHandler->bind("shutdown", &ClangdLSPServer::onShutdown);
+  MsgHandler->bind("textDocument/rangeFormatting", &ClangdLSPServer::onDocumentRangeFormatting);
+  MsgHandler->bind("textDocument/onTypeFormatting", &ClangdLSPServer::onDocumentOnTypeFormatting);
+  MsgHandler->bind("textDocument/formatting", &ClangdLSPServer::onDocumentFormatting);
+  MsgHandler->bind("textDocument/codeAction", &ClangdLSPServer::onCodeAction);
+  MsgHandler->bind("textDocument/completion", &ClangdLSPServer::onCompletion);
+  MsgHandler->bind("textDocument/signatureHelp", &ClangdLSPServer::onSignatureHelp);
+  MsgHandler->bind("textDocument/definition", &ClangdLSPServer::onGoToDefinition);
+  MsgHandler->bind("textDocument/references", &ClangdLSPServer::onReference);
+  MsgHandler->bind("textDocument/switchSourceHeader", &ClangdLSPServer::onSwitchSourceHeader);
+  MsgHandler->bind("textDocument/rename", &ClangdLSPServer::onRename);
+  MsgHandler->bind("textDocument/hover", &ClangdLSPServer::onHover);
+  MsgHandler->bind("textDocument/documentSymbol", &ClangdLSPServer::onDocumentSymbol);
+  MsgHandler->bind("workspace/executeCommand", &ClangdLSPServer::onCommand);
+  MsgHandler->bind("textDocument/documentHighlight", &ClangdLSPServer::onDocumentHighlight);
+  MsgHandler->bind("workspace/symbol", &ClangdLSPServer::onWorkspaceSymbol);
+  MsgHandler->bind("textDocument/didOpen", &ClangdLSPServer::onDocumentDidOpen);
+  MsgHandler->bind("textDocument/didClose", &ClangdLSPServer::onDocumentDidClose);
+  MsgHandler->bind("textDocument/didChange", &ClangdLSPServer::onDocumentDidChange);
+  MsgHandler->bind("workspace/didChangeWatchedFiles", &ClangdLSPServer::onFileEvent);
+  MsgHandler->bind("workspace/didChangeConfiguration", &ClangdLSPServer::onChangeConfiguration);
+  // clang-format on
+}
+
+ClangdLSPServer::~ClangdLSPServer() = default;
 
 bool ClangdLSPServer::run() {
   assert(Server);
 
-  // Set up JSONRPCDispatcher.
-  JSONRPCDispatcher Dispatcher([](const json::Value &Params) {
-    replyError(ErrorCode::MethodNotFound, "method not found");
-    return true;
-  });
-  registerCallbackHandlers(Dispatcher, /*Callbacks=*/*this);
-
   // Run the Language Server loop.
   bool CleanExit = true;
-  if (auto Err = Dispatcher.runLanguageServerLoop(Transp)) {
+  if (auto Err = Transp.loop(*MsgHandler)) {
     elog("Transport error: {0}", std::move(Err));
     CleanExit = false;
   }
 
   // Destroy ClangdServer to ensure all worker threads finish.
   Server.reset();
 
   return CleanExit && ShutdownRequestReceived;
void ClangdLSPServer::onDiagnosticsReady(PathRef File,
   // Cache FixIts
   {
     // FIXME(ibiryukov): should be deleted when documents are removed
     std::lock_guard<std::mutex> Lock(FixItsMutex);
     FixItsMap[File] = LocalFixIts;
   }
 
   // Publish diagnostics.
-  Transp.notify("textDocument/publishDiagnostics",
-                json::Object{
-                    {"uri", URIForFile{File}},
-                    {"diagnostics", std::move(DiagnosticsJSON)},
-                });
+  notify("textDocument/publishDiagnostics",
 }
 
 void ClangdLSPServer::reparseOpenedFiles() {
   for (const Path &FilePath : DraftMgr.getActiveFiles())
     Server->addDocument(FilePath, *DraftMgr.getDraft(FilePath),
                         WantDiagnostics::Auto);
 }
 

clangd/JSONRPCDispatcher.h

-//===--- JSONRPCDispatcher.h - Main JSON parser entry point -----*- C++ -*-===//
-//
-//                     The LLVM Compiler Infrastructure
-//
-// This file is distributed under the University of Illinois Open Source
-// License. See LICENSE.TXT for details.
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_JSONRPCDISPATCHER_H
-#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_JSONRPCDISPATCHER_H
-
-#include "Cancellation.h"
-#include "Logger.h"
-#include "Protocol.h"
-#include "Trace.h"
-#include "Transport.h"
-#include "clang/Basic/LLVM.h"
-#include "llvm/ADT/SmallString.h"
-#include "llvm/ADT/StringMap.h"
-#include "llvm/Support/JSON.h"
-#include <iosfwd>
-#include <mutex>
-
-namespace clang {
-namespace clangd {
-
-/// Sends a successful reply.
-/// Current context must derive from JSONRPCDispatcher::Handler.
-void reply(llvm::json::Value &&Result);
-/// Sends an error response to the client, and logs it.
-/// Current context must derive from JSONRPCDispatcher::Handler.
-void replyError(ErrorCode Code, const llvm::StringRef &Message);
-/// Implements ErrorCode and message extraction from a given llvm::Error. It
-/// fetches the related message from error's message method. If error doesn't
-/// match any known errors, uses ErrorCode::InvalidParams for the error.
-void replyError(llvm::Error E);
-/// Sends a request to the client.
-/// Current context must derive from JSONRPCDispatcher::Handler.
-void call(llvm::StringRef Method, llvm::json::Value &&Params);
-
-/// Main JSONRPC entry point. This parses the JSONRPC "header" and calls the
-/// registered Handler for the method received.
-///
-/// The `$/cancelRequest` notification is handled by the dispatcher itself.
-/// It marks the matching request as cancelled, if it's still running.
-class JSONRPCDispatcher : private Transport::MessageHandler {
-public:
-  /// A handler responds to requests for a particular method name.
-  /// It returns false if the server should now shut down.
-  ///
-  /// JSONRPCDispatcher will mark the handler's context as cancelled if a
-  /// matching cancellation request is received. Handlers are encouraged to
-  /// check for cancellation and fail quickly in this case.
-  using Handler = std::function<bool(const llvm::json::Value &)>;
-
-  /// Create a new JSONRPCDispatcher. UnknownHandler is called when an unknown
-  /// method is received.
-  JSONRPCDispatcher(Handler UnknownHandler);
-
-  /// Registers a Handler for the specified Method.
-  void registerHandler(StringRef Method, Handler H);
-
-  /// Parses input queries from LSP client (coming from \p In) and runs call
-  /// method for each query.
-  ///
-  /// Input stream(\p In) must be opened in binary mode to avoid
-  /// preliminary replacements of \r\n with \n. We use C-style FILE* for reading
-  /// as std::istream has unclear interaction with signals, which are sent by
-  /// debuggers on some OSs.
-  llvm::Error runLanguageServerLoop(Transport &);
-
-private:
-  bool onReply(llvm::json::Value ID,
-               llvm::Expected<llvm::json::Value> Result) override;
-  bool onNotify(llvm::StringRef Method, llvm::json::Value Message) override;
-  bool onCall(llvm::StringRef Method, llvm::json::Value Message,
-              llvm::json::Value ID) override;
-
-  // Tracking cancellations needs a mutex: handlers may finish on a different
-  // thread, and that's when we clean up entries in the map.
-  mutable std::mutex RequestCancelersMutex;
-  llvm::StringMap<std::pair<Canceler, unsigned>> RequestCancelers;
-  unsigned NextRequestCookie = 0;
-  Context cancelableRequestContext(const llvm::json::Value &ID);
-  void cancelRequest(const llvm::json::Value &ID);
-
-  llvm::StringMap<Handler> Handlers;
-  Handler UnknownHandler;
-};
-
-} // namespace clangd
-} // namespace clang
-
-#endif

clangd/JSONRPCDispatcher.cpp

-//===--- JSONRPCDispatcher.cpp - Main JSON parser entry point -------------===//
-//
-//                     The LLVM Compiler Infrastructure
-//
-// This file is distributed under the University of Illinois Open Source
-// License. See LICENSE.TXT for details.
-//
-//===----------------------------------------------------------------------===//
-
-#include "JSONRPCDispatcher.h"
-#include "Cancellation.h"
-#include "ProtocolHandlers.h"
-#include "Trace.h"
-#include "Transport.h"
-#include "llvm/ADT/ScopeExit.h"
-#include "llvm/ADT/SmallString.h"
-#include "llvm/ADT/StringExtras.h"
-#include "llvm/Support/Errno.h"
-#include "llvm/Support/JSON.h"
-#include "llvm/Support/ScopedPrinter.h"
-#include "llvm/Support/SourceMgr.h"
-#include <istream>
-
-using namespace llvm;
-using namespace clang;
-using namespace clangd;
-
-namespace {
-static Key<json::Value> RequestID;
-static Key<Transport *> CurrentTransport;
-
-// When tracing, we trace a request and attach the response in reply().
-// Because the Span isn't available, we find the current request using Context.
-class RequestSpan {
-  RequestSpan(llvm::json::Object *Args) : Args(Args) {}
-  std::mutex Mu;
-  llvm::json::Object *Args;
-  static Key<std::unique_ptr<RequestSpan>> RSKey;
-
-public:
-  // Return a context that's aware of the enclosing request, identified by Span.
-  static Context stash(const trace::Span &Span) {
-    return Context::current().derive(
-        RSKey, std::unique_ptr<RequestSpan>(new RequestSpan(Span.Args)));
-  }
-
-  // If there's an enclosing request and the tracer is interested, calls \p F
-  // with a json::Object where request info can be added.
-  template <typename Func> static void attach(Func &&F) {
-    auto *RequestArgs = Context::current().get(RSKey);
-    if (!RequestArgs || !*RequestArgs || !(*RequestArgs)->Args)
-      return;
-    std::lock_guard<std::mutex> Lock((*RequestArgs)->Mu);
-    F(*(*RequestArgs)->Args);
-  }
-};
-Key<std::unique_ptr<RequestSpan>> RequestSpan::RSKey;
-} // namespace
-
-void clangd::reply(json::Value &&Result) {
-  auto ID = Context::current().get(RequestID);
-  if (!ID) {
-    elog("Attempted to reply to a notification!");
-    return;
-  }
-  RequestSpan::attach([&](json::Object &Args) { Args["Reply"] = Result; });
-  log("--> reply({0})", *ID);
-  Context::current()
-      .getExisting(CurrentTransport)
-      ->reply(std::move(*ID), std::move(Result));
-}
-
-void clangd::replyError(ErrorCode Code, const llvm::StringRef &Message) {
-  elog("Error {0}: {1}", static_cast<int>(Code), Message);
-  RequestSpan::attach([&](json::Object &Args) {
-    Args["Error"] = json::Object{{"code", static_cast<int>(Code)},
-                                 {"message", Message.str()}};
-  });
-
-  if (auto ID = Context::current().get(RequestID)) {
-    log("--> reply({0}) error: {1}", *ID, Message);
-    Context::current()
-        .getExisting(CurrentTransport)
-        ->reply(std::move(*ID), make_error<LSPError>(Message, Code));
-  }
-}
-
-void clangd::replyError(Error E) {
-  handleAllErrors(std::move(E),
-                  [](const CancelledError &TCE) {
-                    replyError(ErrorCode::RequestCancelled, TCE.message());
-                  },
-                  [](const ErrorInfoBase &EIB) {
-                    replyError(ErrorCode::InvalidParams, EIB.message());
-                  });
-}
-
-void clangd::call(StringRef Method, json::Value &&Params) {
-  RequestSpan::attach([&](json::Object &Args) {
-    Args["Call"] = json::Object{{"method", Method.str()}, {"params", Params}};
-  });
-  // FIXME: Generate/Increment IDs for every request so that we can get proper
-  // replies once we need to.
-  auto ID = 1;
-  log("--> {0}({1})", Method, ID);
-  Context::current()
-      .getExisting(CurrentTransport)
-      ->call(Method, std::move(Params), ID);
-}
-
-JSONRPCDispatcher::JSONRPCDispatcher(Handler UnknownHandler)
-    : UnknownHandler(std::move(UnknownHandler)) {
-  registerHandler("$/cancelRequest", [this](const json::Value &Params) {
-    if (auto *O = Params.getAsObject())
-      if (auto *ID = O->get("id")) {
-        cancelRequest(*ID);
-        return true;
-      }
-    log("Bad cancellation request: {0}", Params);
-    return true;
-  });
-}
-
-void JSONRPCDispatcher::registerHandler(StringRef Method, Handler H) {
-  assert(!Handlers.count(Method) && "Handler already registered!");
-  Handlers[Method] = std::move(H);
-}
-
-bool JSONRPCDispatcher::onCall(StringRef Method, json::Value Params,
-                               json::Value ID) {
-  log("<-- {0}({1})", Method, ID);
-  auto I = Handlers.find(Method);
-  auto &Handler = I != Handlers.end() ? I->second : UnknownHandler;
-
-  // Create a Context that contains request information.
-  WithContextValue WithID(RequestID, ID);
-
-  // Create a tracing Span covering the whole request lifetime.
-  trace::Span Tracer(Method);
-  SPAN_ATTACH(Tracer, "ID", ID);
-  SPAN_ATTACH(Tracer, "Params", Params);
-
-  // Calls can be canceled by the client. Add cancellation context.
-  WithContext WithCancel(cancelableRequestContext(ID));
-
-  // Stash a reference to the span args, so later calls can add metadata.
-  WithContext WithRequestSpan(RequestSpan::stash(Tracer));
-  return Handler(std::move(Params));
-}
-
-bool JSONRPCDispatcher::onNotify(StringRef Method, json::Value Params) {
-  log("<-- {0}", Method);
-  auto I = Handlers.find(Method);
-  auto &Handler = I != Handlers.end() ? I->second : UnknownHandler;
-
-  // Create a tracing Span covering the whole request lifetime.
-  trace::Span Tracer(Method);
-  SPAN_ATTACH(Tracer, "Params", Params);
-
-  // Stash a reference to the span args, so later calls can add metadata.
-  WithContext WithRequestSpan(RequestSpan::stash(Tracer));
-  return Handler(std::move(Params));
-}
-
-bool JSONRPCDispatcher::onReply(json::Value ID, Expected<json::Value> Result) {
-  // We ignore replies, just log them.
-  if (Result)
-    log("<-- reply({0})", ID);
-  else
-    log("<-- reply({0}) error: {1}", ID, llvm::toString(Result.takeError()));
-  return true;
-}
-
-// We run cancelable requests in a context that does two things:
-//  - allows cancellation using RequestCancelers[ID]
-//  - cleans up the entry in RequestCancelers when it's no longer needed
-// If a client reuses an ID, the last one wins and the first cannot be canceled.
-Context JSONRPCDispatcher::cancelableRequestContext(const json::Value &ID) {
-  auto Task = cancelableTask();
-  auto StrID = llvm::to_string(ID);  // JSON-serialize ID for map key.
-  auto Cookie = NextRequestCookie++; // No lock, only called on main thread.
-  {
-    std::lock_guard<std::mutex> Lock(RequestCancelersMutex);
-    RequestCancelers[StrID] = {std::move(Task.second), Cookie};
-  }
-  // When the request ends, we can clean up the entry we just added.
-  // The cookie lets us check that it hasn't been overwritten due to ID reuse.
-  return Task.first.derive(make_scope_exit([this, StrID, Cookie] {
-    std::lock_guard<std::mutex> Lock(RequestCancelersMutex);
-    auto It = RequestCancelers.find(StrID);
-    if (It != RequestCancelers.end() && It->second.second == Cookie)
-      RequestCancelers.erase(It);
-  }));
-}
-
-void JSONRPCDispatcher::cancelRequest(const json::Value &ID) {
-  auto StrID = llvm::to_string(ID);
-  std::lock_guard<std::mutex> Lock(RequestCancelersMutex);
-  auto It = RequestCancelers.find(StrID);
-  if (It != RequestCancelers.end())
-    It->second.first(); // Invoke the canceler.
-}
-
-llvm::Error JSONRPCDispatcher::runLanguageServerLoop(Transport &Transport) {
-  // Propagate transport to all handlers so they can reply.
-  WithContextValue WithTransport(CurrentTransport, &Transport);
-  return Transport.loop(*this);
-}

clangd/ProtocolHandlers.h

-//===--- ProtocolHandlers.h - LSP callbacks ---------------------*- C++ -*-===//
-//
-//                     The LLVM Compiler Infrastructure
-//
-// This file is distributed under the University of Illinois Open Source
-// License. See LICENSE.TXT for details.
-//
-//===----------------------------------------------------------------------===//
-//
-// ProtocolHandlers translates incoming JSON requests from JSONRPCDispatcher
-// into method calls on ClangLSPServer.
-//
-// Currently it parses requests into objects, but the ClangLSPServer is
-// responsible for producing JSON responses. We should move that here, too.
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_PROTOCOLHANDLERS_H
-#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_PROTOCOLHANDLERS_H
-
-#include "JSONRPCDispatcher.h"
-#include "Protocol.h"
-#include "llvm/ADT/Twine.h"
-#include "llvm/Support/raw_ostream.h"
-
-namespace clang {
-namespace clangd {
-
-// The interface implemented by ClangLSPServer to handle incoming requests.
-class ProtocolCallbacks {
-public:
-  virtual ~ProtocolCallbacks() = default;
-
-  virtual void onInitialize(InitializeParams &Params) = 0;
-  virtual void onShutdown(ShutdownParams &Params) = 0;
-  virtual void onExit(ExitParams &Params) = 0;
-  virtual void onDocumentDidOpen(DidOpenTextDocumentParams &Params) = 0;
-  virtual void onDocumentDidChange(DidChangeTextDocumentParams &Params) = 0;
-  virtual void onDocumentDidClose(DidCloseTextDocumentParams &Params) = 0;
-  virtual void onDocumentFormatting(DocumentFormattingParams &Params) = 0;
-  virtual void onDocumentSymbol(DocumentSymbolParams &Params) = 0;
-  virtual void
-  onDocumentOnTypeFormatting(DocumentOnTypeFormattingParams &Params) = 0;
-  virtual void
-  onDocumentRangeFormatting(DocumentRangeFormattingParams &Params) = 0;
-  virtual void onCodeAction(CodeActionParams &Params) = 0;
-  virtual void onCompletion(TextDocumentPositionParams &Params) = 0;
-  virtual void onSignatureHelp(TextDocumentPositionParams &Params) = 0;
-  virtual void onGoToDefinition(TextDocumentPositionParams &Params) = 0;
-  virtual void onReference(ReferenceParams &Params) = 0;
-  virtual void onSwitchSourceHeader(TextDocumentIdentifier &Params) = 0;
-  virtual void onFileEvent(DidChangeWatchedFilesParams &Params) = 0;
-  virtual void onCommand(ExecuteCommandParams &Params) = 0;
-  virtual void onWorkspaceSymbol(WorkspaceSymbolParams &Params) = 0;
-  virtual void onRename(RenameParams &Parames) = 0;
-  virtual void onDocumentHighlight(TextDocumentPositionParams &Params) = 0;
-  virtual void onHover(TextDocumentPositionParams &Params) = 0;
-  virtual void onChangeConfiguration(DidChangeConfigurationParams &Params) = 0;
-};
-
-void registerCallbackHandlers(JSONRPCDispatcher &Dispatcher,
-                              ProtocolCallbacks &Callbacks);
-
-} // namespace clangd
-} // namespace clang
-
-#endif

clangd/ProtocolHandlers.cpp

-//===--- ProtocolHandlers.cpp - LSP callbacks -----------------------------===//
-//
-//                     The LLVM Compiler Infrastructure
-//
-// This file is distributed under the University of Illinois Open Source
-// License. See LICENSE.TXT for details.
-//
-//===----------------------------------------------------------------------===//
-
-#include "ProtocolHandlers.h"
-#include "ClangdLSPServer.h"
-#include "ClangdServer.h"
-#include "DraftStore.h"
-#include "Trace.h"
-
-using namespace clang;
-using namespace clang::clangd;
-using namespace llvm;
-
-namespace {
-
-// Helper for attaching ProtocolCallbacks methods to a JSONRPCDispatcher.
-// Invoke like: Registerer("foo", &ProtocolCallbacks::onFoo)
-// onFoo should be: void onFoo(Ctx &C, FooParams &Params)
-// FooParams should have a fromJSON function.
-struct HandlerRegisterer {
-  template <typename Param>
-  void operator()(StringRef Method, void (ProtocolCallbacks::*Handler)(Param)) {
-    // Capture pointers by value, as the lambda will outlive this object.
-    auto *Callbacks = this->Callbacks;
-    Dispatcher.registerHandler(Method, [=](const json::Value &RawParams) {
-      typename std::remove_reference<Param>::type P;
-      if (fromJSON(RawParams, P)) {
-        (Callbacks->*Handler)(P);
-      } else {
-        elog("Failed to decode {0} request.", Method);
-      }
-      return Method != "exit"; // Shut down after exit notification.
-    });
-  }
-
-  JSONRPCDispatcher &Dispatcher;
-  ProtocolCallbacks *Callbacks;
-};
-
-} // namespace
-
-void clangd::registerCallbackHandlers(JSONRPCDispatcher &Dispatcher,
-                                      ProtocolCallbacks &Callbacks) {
-  HandlerRegisterer Register{Dispatcher, &Callbacks};
-
-  Register("initialize", &ProtocolCallbacks::onInitialize);
-  Register("shutdown", &ProtocolCallbacks::onShutdown);
-  Register("exit", &ProtocolCallbacks::onExit);
-  Register("textDocument/didOpen", &ProtocolCallbacks::onDocumentDidOpen);
-  Register("textDocument/didClose", &ProtocolCallbacks::onDocumentDidClose);
-  Register("textDocument/didChange", &ProtocolCallbacks::onDocumentDidChange);
-  Register("textDocument/rangeFormatting",
-           &ProtocolCallbacks::onDocumentRangeFormatting);
-  Register("textDocument/onTypeFormatting",
-           &ProtocolCallbacks::onDocumentOnTypeFormatting);
-  Register("textDocument/formatting", &ProtocolCallbacks::onDocumentFormatting);
-  Register("textDocument/codeAction", &ProtocolCallbacks::onCodeAction);
-  Register("textDocument/completion", &ProtocolCallbacks::onCompletion);
-  Register("textDocument/signatureHelp", &ProtocolCallbacks::onSignatureHelp);
-  Register("textDocument/definition", &ProtocolCallbacks::onGoToDefinition);
-  Register("textDocument/references", &ProtocolCallbacks::onReference);
-  Register("textDocument/switchSourceHeader",
-           &ProtocolCallbacks::onSwitchSourceHeader);
-  Register("textDocument/rename", &ProtocolCallbacks::onRename);
-  Register("textDocument/hover", &ProtocolCallbacks::onHover);
-  Register("textDocument/documentSymbol", &ProtocolCallbacks::onDocumentSymbol);
-  Register("workspace/didChangeWatchedFiles", &ProtocolCallbacks::onFileEvent);
-  Register("workspace/executeCommand", &ProtocolCallbacks::onCommand);
-  Register("textDocument/documentHighlight",
-           &ProtocolCallbacks::onDocumentHighlight);
-  Register("workspace/didChangeConfiguration",
-           &ProtocolCallbacks::onChangeConfiguration);
-  Register("workspace/symbol", &ProtocolCallbacks::onWorkspaceSymbol);
-}

clangd/TUScheduler.cpp

elog("Trying to remove file from TUScheduler that is not tracked: {0}",
          File);
 }
 
 void TUScheduler::runWithAST(
     llvm::StringRef Name, PathRef File,
     llvm::unique_function<void(llvm::Expected<InputsAndAST>)> Action) {
   auto It = Files.find(File);
   if (It == Files.end()) {
-    Action(llvm::make_error<llvm::StringError>(
-        "trying to get AST for non-added document",
-        llvm::errc::invalid_argument));
+    Action(llvm::make_error<LSPError>(
+        "trying to get AST for non-added document", ErrorCode::InvalidParams));
     return;
   }
 
   It->second->Worker->runWithAST(Name, std::move(Action));
 }
 
 void TUScheduler::runWithPreamble(
     llvm::StringRef Name, PathRef File, PreambleConsistency Consistency,
     llvm::unique_function<void(llvm::Expected<InputsAndPreamble>)> Action) {
   auto It = Files.find(File);
   if (It == Files.end()) {
-    Action(llvm::make_error<llvm::StringError>(
+    Action(llvm::make_error<LSPError>(
         "trying to get preamble for non-added document",
-        llvm::errc::invalid_argument));
+        ErrorCode::InvalidParams));
     return;
   }
 
   if (!PreambleTasks) {
     trace::Span Tracer(Name);
     SPAN_ATTACH(Tracer, "file", File);
     std::shared_ptr<const PreambleData> Preamble =
         It->second->Worker->getPossiblyStalePreamble();

clangd/tool/ClangdMain.cpp

 //===--- ClangdMain.cpp - clangd server loop ------------------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #include "ClangdLSPServer.h"
-#include "JSONRPCDispatcher.h"
 #include "Path.h"
-#include "RIFF.h"
 #include "Trace.h"
+#include "Transport.h"
 #include "index/Serialization.h"
 #include "clang/Basic/Version.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Program.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/raw_ostream.h"

test/clangd/crash-non-added-files.test

 # CHECK-NEXT:    "code": -32602
 # CHECK-NEXT:    "message": "onDocumentOnTypeFormatting called for non-added file"
 # CHECK-NEXT:  },
 # CHECK-NEXT:  "id": 5,
 ---
 {"jsonrpc":"2.0","id":6,"method":"shutdown"}
 ---
 {"jsonrpc":"2.0","method":"exit"}
----
-{"jsonrpc":"2.0","method":"exit"}

test/clangd/delimited-input-comment-at-the-end.test

 # RUN: clangd -input-style=delimited -run-synchronously -input-mirror-file %t < %s
 # RUN: grep '{"jsonrpc":"2.0","id":3,"method":"exit"}' %t
 #
 # RUN: clangd -lit-test -input-mirror-file %t < %s
 # RUN: grep '{"jsonrpc":"2.0","id":3,"method":"exit"}' %t
 #
 {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"clangd","capabilities":{},"trace":"off"}}
 ---
 {"jsonrpc":"2.0","id":3,"method":"shutdown"}
 ---
-{"jsonrpc":"2.0","id":3,"method":"exit"}
-# comment at the end
+{"jsonrpc":"2.0","method":"exit"}

test/clangd/fixits-command.test

 # CHECK-NEXT:    }
 # CHECK-NEXT:  ]
 ---
 {"jsonrpc":"2.0","id":4,"method":"workspace/executeCommand","params":{"command":"clangd.applyFix","arguments":[{"changes":{"test:///foo.c":[{"range":{"start":{"line":0,"character":32},"end":{"line":0,"character":32}},"newText":"("},{"range":{"start":{"line":0,"character":37},"end":{"line":0,"character":37}},"newText":")"}]}}]}}
 #      CHECK:  "id": 4,
 # CHECK-NEXT:  "jsonrpc": "2.0",
 # CHECK-NEXT:  "result": "Fix applied."
 #
-#      CHECK:  "id": 1,
+#      CHECK:  "id": 0,
 # CHECK-NEXT:  "jsonrpc": "2.0",
 # CHECK-NEXT:  "method": "workspace/applyEdit",
 # CHECK-NEXT:  "params": {
 # CHECK-NEXT:    "edit": {
 # CHECK-NEXT:      "changes": {
 # CHECK-NEXT:        "{{.*}}/foo.c": [
 # CHECK-NEXT:          {
 # CHECK-NEXT:            "newText": "(",

test/clangd/rename.test

 # CHECK-NEXT:          }
 # CHECK-NEXT:        }
 # CHECK-NEXT:      ]
 # CHECK-NEXT:    }
 # CHECK-NEXT:  }
 ---
 {"jsonrpc":"2.0","id":2,"method":"textDocument/rename","params":{"textDocument":{"uri":"test:///foo.cpp"},"position":{"line":0,"character":2},"newName":"bar"}}
 #      CHECK:  "error": {
-# CHECK-NEXT:    "code": -32603,
+# CHECK-NEXT:    "code": -32001,
 # CHECK-NEXT:    "message": "clang diagnostic"
 # CHECK-NEXT:  },
 # CHECK-NEXT:  "id": 2,
 # CHECK-NEXT:  "jsonrpc": "2.0"
 ---
 {"jsonrpc":"2.0","id":3,"method":"shutdown"}
 ---
 {"jsonrpc":"2.0","method":"exit"}

test/clangd/spaces-in-delimited-input.test

 # RUN: clangd -input-style=delimited -run-synchronously < %s 2>&1 | FileCheck %s
 # RUN: clangd -lit-test -run-synchronously < %s 2>&1 | FileCheck %s
 #
 {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"clangd","capabilities":{},"trace":"off"}}
 
 ---
 
 {"jsonrpc":"2.0","id":3,"method":"shutdown"}
 
 ---
 
-{"jsonrpc":"2.0","id":3,"method":"exit"}
+{"jsonrpc":"2.0","method":"exit"}
 # CHECK-NOT: JSON parse error