[clangd] refactoring for XPC transport layer [NFCI]

[jkorous]

Summary

Hi all,

We finally finished a self-contained first version of our implementation of alternative transport layer for macOS based on XPC.

To enable this I did couple of changes to current clangd design. Generally I aimed for small amount of disruption and tried to keep it simple.

The notable changes are:

• enabled registration of request handlers to different dispatchers by templatizing registerCallbackHandlers()
• factored out parts of JsonDispatcher that could be reused (created DispatcherCommon.h/cpp)
• created abstraction over JsonOutput (class LSPOutput)
• removed of ClangdLSPServer::run() method so server can be run with different dispatcher
• published IsDone and ShutdownRequestReceived through methods in ClangdLSPServer class interface to support new dispatcher

We are also putting up the transport layer implementation itself for a review soon so it will be more obvious where are we going to and what motivated some of these changes.

Big thanks in advance to all reviewers!

Jan

[jkorous]

Followed by these two patches:
[clangd] JSON <-> XPC conversions
https://reviews.llvm.org/D48560
[clangd] XPC transport layer
https://reviews.llvm.org/D48562

[jkorous]

Ping. Added reviewers suggestion.

[sammccall]

Hi Jan,

Thanks for putting this together, and apologies - this is one of the places where we don't have nice abstractions/layering, so adding XPC was harder than it should be.

As I mentioned on the other review, I think maybe this patch isn't invasive enough, we could do a more thorough job of making the json transport a standalone, first-class thing. This makes it easier to swap out for XPC but also would improve the layering in the core clangd classes.

I put together a sketch of a Transport interface in D49389 (that patch is extremely unfinished and won't compile, but the idea might be interesting).
The idea is that we should be able to implement that class for XPC and then it should just drop-in as a replacement for the JSON one.
I've indicated there where XPC and JSON implementation can share code (assuming you wouldn't rather "tighten up" the protocol and eliminate some JSON-RPC noise).

If you like this direction I'm happy for you to pick the bits you like, or I can finish it as a refactoring and we can make sure it works for XPC.
If not, that's fine too - happy to look at other ways to reduce duplication between them.

DispatcherCommon.h

38

I think I wrote this bit... it's a hack that I'm not sure deserves to be visible in a header (it was bad enough in JSONRPCDispatcher.cpp!)

Rather than exposing it so we can use it twice, I'd hope we can manage to keep it hidden so that JSON and XPC can share it.

[jkorous]

Hi Sam, thanks for your feedback!

In general I agree that explicitly factoring out the transport layer and improving layering would be a good thing. Unfortunately it's highly probable that we'd like to drop JSON completely from XPC dispatch (XPC -> JSON -> ProtocolCallbacks -> JSON -> XPC) in not too distant future. (Please don't take this as a promise as our priorities might change but it's probable enough that I don't recommend to base any common abstraction on JSON.) I originally tried to create similar abstraction but ultimately gave up as it was way too complex. Not saying it's impossible or that I would not like to though!

I basically see these different parts of the problem:

1. dispatch implementation

I think that different dispatch implementations might be better kept separate without any common abstraction as it's not a lot of code and not much of opportunities for shared implementation - DispatcherCommon.h/cpp effectively vanishes in case we drop JSON from XPC dispatch.

2. LSP request deserialization

Deserialization of LSP requests for different transport layers is easy - it's a single line of code (currently lambda wrapper in HandlerRegisterer) so it would be easy to parametrize this.

3. LSP response serialization

This is where I failed. Basically I haven't found any nice way how to make it a part of transport layer interface. I could imagine having a set of response classes in Protocol.h and toXpc functions but still not sure what kind of interface to have for these in transport layer.

[sammccall]

In D48559#1165172, @jkorous wrote:

Hi Sam, thanks for your feedback!

In general I agree that explicitly factoring out the transport layer and improving layering would be a good thing. Unfortunately it's highly probable that we'd like to drop JSON completely from XPC dispatch (XPC -> JSON -> ProtocolCallbacks -> JSON -> XPC) in not too distant future. (Please don't take this as a promise as our priorities might change but it's probable enough that I don't recommend to base any common abstraction on JSON.) I originally tried to create similar abstraction but ultimately gave up as it was way too complex. Not saying it's impossible or that I would not like to though!

That makes sense and is good to know. I definitely don't want to add a transport abstraction that isn't a good fit for your long-term plans.

I think we need to clearly see what the goal is to get the design right though - it doesn't make sense to refactor in order to share code temporarily. Let's take code completion as an example, as it's one of the ones we've fleshed out most for our internal (non-LSP) clients.

• if the goal is to mirror JSON-RPC with a different encoding (current patches), and you're just concerned about the *efficiency* of CodeCompletion -> CompletionItem -> json::Value -> xpc_object_t, this seems like premature optimization to me, but let's discuss!
• if the goal is to basically follow the LSP (same textDocument/completion for), using the CompletionItem structs in Protocol.h but exposing more/fewer fields, renaming them etc, then we should probably have the feature code express replies as protocol structs, and make the transport define a way to transform specific objects (CompletionItem) into generic ones whose type (json::Value or xpc_object_t) are transport-specific
• if the goal is to maintain a stable protocol that can be evolved independently of LSP (different methods or semantics) then this isn't really a transport, it's a new protocol. I think the best option is probably to skip ClangdLSPServer and have the xpc service wrap ClangdServer instead, with a separate main entrypoint. The litmus test here: as people make changes to clangd that affect the LSP surface, do you want the XPC service's interface to change too? This is the hardest model to support as the C++ API (ClangdServer) is not a stable one, but it's what we do for the other non-LSP clients.

Does one of these seem close to what you want? Any thoughts on the conclusions?

[arphaman]

In D48559#1165511, @sammccall wrote:

In D48559#1165172, @jkorous wrote:

Hi Sam, thanks for your feedback!

In general I agree that explicitly factoring out the transport layer and improving layering would be a good thing. Unfortunately it's highly probable that we'd like to drop JSON completely from XPC dispatch (XPC -> JSON -> ProtocolCallbacks -> JSON -> XPC) in not too distant future. (Please don't take this as a promise as our priorities might change but it's probable enough that I don't recommend to base any common abstraction on JSON.) I originally tried to create similar abstraction but ultimately gave up as it was way too complex. Not saying it's impossible or that I would not like to though!

That makes sense and is good to know. I definitely don't want to add a transport abstraction that isn't a good fit for your long-term plans.

I think we need to clearly see what the goal is to get the design right though - it doesn't make sense to refactor in order to share code temporarily. Let's take code completion as an example, as it's one of the ones we've fleshed out most for our internal (non-LSP) clients.

• if the goal is to mirror JSON-RPC with a different encoding (current patches), and you're just concerned about the *efficiency* of CodeCompletion -> CompletionItem -> json::Value -> xpc_object_t, this seems like premature optimization to me, but let's discuss!
• if the goal is to basically follow the LSP (same textDocument/completion for), using the CompletionItem structs in Protocol.h but exposing more/fewer fields, renaming them etc, then we should probably have the feature code express replies as protocol structs, and make the transport define a way to transform specific objects (CompletionItem) into generic ones whose type (json::Value or xpc_object_t) are transport-specific
• if the goal is to maintain a stable protocol that can be evolved independently of LSP (different methods or semantics) then this isn't really a transport, it's a new protocol. I think the best option is probably to skip ClangdLSPServer and have the xpc service wrap ClangdServer instead, with a separate main entrypoint. The litmus test here: as people make changes to clangd that affect the LSP surface, do you want the XPC service's interface to change too? This is the hardest model to support as the C++ API (ClangdServer) is not a stable one, but it's what we do for the other non-LSP clients.

Does one of these seem close to what you want? Any thoughts on the conclusions?

Hi Sam,
Thanks for your feedback!

Generally right now will be mirroring JSON-RPC with a different encoding, but in the future we will send some optimized payloads for a subset of replies which are constructed out of things like CompletionItem directly (so bypassing the JSON step).

I believe that your patch (https://reviews.llvm.org/D49389) is implementing the JSON-RPC with a choice of encoding, which will work for us right now, and it doesn't look like we will have problems optimizing our payloads later when we choose to do so.
However, the second option that you proposed does sound quite compelling as well (i.e. then we should probably have the feature code express replies as protocol structs), as it would give us more flexibility. However, I think it's orthogonal to the first option (JSON-RPC with a choice of encoding), so maybe we can leave it as something that will be on the table if we need it later? I think for the transport layer idea you proposed in https://reviews.llvm.org/D49389 makes sense and will work for us, and in the future we always can implement the second option that your proposed on top of of your patch, or use a simpler way to optimize our payloads without much disruption.

Thanks,
Alex

[jkorous]

Sam, just out of curiosity - would it be possible for you to share any relevant experience gained by using porting clangd to protobuf-based transport layer?

[jkorous]

Hi Sam, we discussed a bit more with Alex offline. Ultimately we don't mind using JSON as part of the generic interface but we'd like to optimize specific cases in the future (thinking only about code completion so far) which might need to "work-around" the abstraction. What is your opinion about this - does it ruin the whole point of having the abstraction or do you consider pragmatic to have abstraction covering most of the functionality with maybe couple exceptions? I kind of assume (please correct me if I am wrong) that you might have done something similar with profobuf interface so we're definitely happy to discuss.

[jkorous]

BTW Just for the record - I put the rebased & updated patch here: [clangd] XPC WIP https://reviews.llvm.org/D49548

[sammccall]

Sorry for the delay here, and I'm sorry I haven't been into the patches in detail again yet - crazy week. After experiments, I am convinced the Transport abstraction patch can turn into something nice if we want XPC vs JSON to be a pure transport-level difference with the same semantics. But you need to decide the approach here based on your requirements.

In D48559#1168204, @jkorous wrote:

Sam, just out of curiosity - would it be possible for you to share any relevant experience gained by using porting clangd to protobuf-based transport layer?

Sure! (and sorry for the delay here).

This is done as part of an internal-only editor that's a hosted web application with tight integration to our internal source control, build system etc.
That's not my immediate team, but we talk a lot and collaborate on the integration, which is different than a typical LSP server/editor integration in a few ways.
(I expect some of these aspects will apply to XCode+clangd, and others won't). This is kind of a brain dump...

• like a regular LSP client, the editor wants to run clangd out-of-process for isolation reasons, and because of cross-language issues. Because of our hosted datacenter environment, a network RPC server was the most obvious choice, and protobufs are our lingua franca for such servers. This sounds analagous to the choice of XPC to me.
• the editor is multi-language, so the idea of using the existing LSP is appealing (reuse servers, we can't rewrite the world).
• adapting to protobuf is some work (write a wrapper for each LSP server). In principle a generic one is possible, but not for us due to VFS needs etc.
• you can wrap at the LSP level (CompletionList -> CompletionListProto) or at the transport level (json::Value --> JsonValueProto). Editor team team went for the former.
• consuming clangd as a C++ API rather than as an external process or opaque JSON stream, gives lots more flexibility, and we can add extension points relatively easily. Robust VFS support, performance tracing, logging, different threading model are examples of this outside the scope of LSP. The internal code uses unmodified upstream clangd code, but provides its own main().
• Inside the scope of LSP, changes are tempting too due to LSP limitations. Since the editor team controls the protocol and the frontend, they can add features. This is a source of tension: we (clangd folks) don't want to support two divergent APIs. So in limited cases we added to clangd a richer/more semantic C++ API that provides extra features over more-presentational LSP. Best examples are diagnostics (C++ API includes fixits and 'note' stack, before LSP supported these), and code completion (C++ API includes semantic things like "name of include to insert"). We've tried to keep this limited to what's both valuable and sensible.
• adding LSP extensions to the *protocol* causes skew across languages. So the divergences at that level tend to be more presentational (e.g. "completion item documentation subtitle"). The clangd-wrapper decides how to map the semantic Clangd C++ API fields onto the presentational mutant-LSP fields, in the same way that regular clangd decides how to map them onto regular LSP.
• it's annoying that the issues you run into with this setup need to be carefully dissected (editor vs wrapper vs clangd) to work out where the bug belongs. Compare to "plain" clangd where you can just check with a different editor and then know the bug is in clangd.

In summary:

• speaking a different wire protocol that's semantically identical to LSP is an easy win
• owning a separate main() so you can plug in cross-cutting facilities (logging, tracing, index) is very manageable and worthwhile
• maintaining a separate protocol is a bunch more work and decisions all the time, even if divergences from LSP are small. It *is* more flexible though.

One halfway possibility if you're on the fence: we could support some limited extensions to the protocol behind an option (e.g. extra fields serialized with the option, whether XPC or JSON is used). It's less flexibility than full ownership of the protocol semantics on Apple's side though.

[arphaman]

In D48559#1169975, @sammccall wrote:

Sorry for the delay here, and I'm sorry I haven't been into the patches in detail again yet - crazy week. After experiments, I am convinced the Transport abstraction patch can turn into something nice if we want XPC vs JSON to be a pure transport-level difference with the same semantics. But you need to decide the approach here based on your requirements.

In D48559#1168204, @jkorous wrote:

Sam, just out of curiosity - would it be possible for you to share any relevant experience gained by using porting clangd to protobuf-based transport layer?

Sure! (and sorry for the delay here).

This is done as part of an internal-only editor that's a hosted web application with tight integration to our internal source control, build system etc.
That's not my immediate team, but we talk a lot and collaborate on the integration, which is different than a typical LSP server/editor integration in a few ways.
(I expect some of these aspects will apply to XCode+clangd, and others won't). This is kind of a brain dump...

• like a regular LSP client, the editor wants to run clangd out-of-process for isolation reasons, and because of cross-language issues. Because of our hosted datacenter environment, a network RPC server was the most obvious choice, and protobufs are our lingua franca for such servers. This sounds analagous to the choice of XPC to me.
• the editor is multi-language, so the idea of using the existing LSP is appealing (reuse servers, we can't rewrite the world).
• adapting to protobuf is some work (write a wrapper for each LSP server). In principle a generic one is possible, but not for us due to VFS needs etc.
• you can wrap at the LSP level (CompletionList -> CompletionListProto) or at the transport level (json::Value --> JsonValueProto). Editor team team went for the former.
• consuming clangd as a C++ API rather than as an external process or opaque JSON stream, gives lots more flexibility, and we can add extension points relatively easily. Robust VFS support, performance tracing, logging, different threading model are examples of this outside the scope of LSP. The internal code uses unmodified upstream clangd code, but provides its own main().
• Inside the scope of LSP, changes are tempting too due to LSP limitations. Since the editor team controls the protocol and the frontend, they can add features. This is a source of tension: we (clangd folks) don't want to support two divergent APIs. So in limited cases we added to clangd a richer/more semantic C++ API that provides extra features over more-presentational LSP. Best examples are diagnostics (C++ API includes fixits and 'note' stack, before LSP supported these), and code completion (C++ API includes semantic things like "name of include to insert"). We've tried to keep this limited to what's both valuable and sensible.
• adding LSP extensions to the *protocol* causes skew across languages. So the divergences at that level tend to be more presentational (e.g. "completion item documentation subtitle"). The clangd-wrapper decides how to map the semantic Clangd C++ API fields onto the presentational mutant-LSP fields, in the same way that regular clangd decides how to map them onto regular LSP.
• it's annoying that the issues you run into with this setup need to be carefully dissected (editor vs wrapper vs clangd) to work out where the bug belongs. Compare to "plain" clangd where you can just check with a different editor and then know the bug is in clangd.

In summary:

• speaking a different wire protocol that's semantically identical to LSP is an easy win
• owning a separate main() so you can plug in cross-cutting facilities (logging, tracing, index) is very manageable and worthwhile
• maintaining a separate protocol is a bunch more work and decisions all the time, even if divergences from LSP are small. It *is* more flexible though.

One halfway possibility if you're on the fence: we could support some limited extensions to the protocol behind an option (e.g. extra fields serialized with the option, whether XPC or JSON is used). It's less flexibility than full ownership of the protocol semantics on Apple's side though.

A custom protocol does sound quite appealing. I previously was quite skeptical just because ClangdLSPServer seemed to have had a lot of important logic, but it seems that that's not the case right now. I will lead an internal discussion on this to figure out if that's something we'd like to do. I'll respond here sometime next week.

[jkorous]

Hi Sam, we are still discussing internally how to fit clangd and XPC together. Please bear with us and ignore our patches until we decide.

[jkorous]

Hi, we decided to go with a different solution:
https://reviews.llvm.org/D50452

CMakeLists.txt

add_clang_library(clangDaemon
   ClangdLSPServer.cpp
   ClangdServer.cpp
   ClangdUnit.cpp
   CodeComplete.cpp
   CodeCompletionStrings.cpp
   Compiler.cpp
   Context.cpp
   Diagnostics.cpp
+  DispatcherCommon.cpp
   DraftStore.cpp
   FindSymbols.cpp
   FuzzyMatch.cpp
   GlobalCompilationDatabase.cpp
   Headers.cpp
   JSONExpr.cpp
   JSONRPCDispatcher.cpp
   Logger.cpp
   Protocol.cpp
-  ProtocolHandlers.cpp
   Quality.cpp
   SourceCode.cpp
   Threading.cpp
   Trace.cpp
   TUScheduler.cpp
   URI.cpp
   XRefs.cpp
   index/CanonicalIncludes.cpp

ClangdLSPServer.h

 #include "Protocol.h"
 #include "ProtocolHandlers.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "llvm/ADT/Optional.h"
 
 namespace clang {
 namespace clangd {
 
-class JSONOutput;
+class LSPOutput;
 class SymbolIndex;
 
 /// This class provides implementation of an LSP server, glueing the JSON
 /// dispatch and ClangdServer together.
-class ClangdLSPServer : private DiagnosticsConsumer, private ProtocolCallbacks {
+class ClangdLSPServer : private DiagnosticsConsumer, public ProtocolCallbacks {
 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(JSONOutput &Out, const clangd::CodeCompleteOptions &CCOpts,
+  ClangdLSPServer(const clangd::CodeCompleteOptions &CCOpts,
                   llvm::Optional<Path> CompileCommandsDir,
                   const ClangdServer::Options &Opts);
 
-  /// Run LSP server loop, receiving input for it from \p In. \p In must be
-  /// opened in binary mode. Output will be written using Out variable passed to
-  /// class constructor. This method must not be executed more than once for
-  /// each instance of ClangdLSPServer.
-  ///
-  /// \return Whether we received a 'shutdown' request before an 'exit' request.
-  bool run(std::FILE *In,
-           JSONStreamStyle InputStyle = JSONStreamStyle::Standard);
-
+  bool& getIsDone();
+  bool getShutdownRequestReceived() const;
 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;
private:
 
   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();
 
-  JSONOutput &Out;
   /// Used to indicate that the 'shutdown' request was received from the
   /// Language Server client.
   bool ShutdownRequestReceived = false;
 
   /// Used to indicate that the 'exit' notification was received from the
   /// Language Server client.
   /// It's used to break out of the LSP parsing loop.
   bool IsDone = false;

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 "DispatcherCommon.h"
 #include "JSONRPCDispatcher.h"
 #include "SourceCode.h"
 #include "URI.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Path.h"
 
 using namespace clang::clangd;
if (Settings.compilationDatabasePath.hasValue()) {
     NonCachedCDB.setCompileCommandsDir(
         Settings.compilationDatabasePath.getValue());
     CDB.clear();
 
     reparseOpenedFiles();
   }
 }
 
-ClangdLSPServer::ClangdLSPServer(JSONOutput &Out,
-                                 const clangd::CodeCompleteOptions &CCOpts,
+ClangdLSPServer::ClangdLSPServer(const clangd::CodeCompleteOptions &CCOpts,
                                  llvm::Optional<Path> CompileCommandsDir,
                                  const ClangdServer::Options &Opts)
-    : Out(Out), NonCachedCDB(std::move(CompileCommandsDir)), CDB(NonCachedCDB),
+    : NonCachedCDB(std::move(CompileCommandsDir)), CDB(NonCachedCDB),
       CCOpts(CCOpts), SupportedSymbolKinds(defaultSymbolKinds()),
       Server(CDB, FSProvider, /*DiagConsumer=*/*this, Opts) {}
 
-bool ClangdLSPServer::run(std::FILE *In, JSONStreamStyle InputStyle) {
-  assert(!IsDone && "Run was called before");
-
-  // Set up JSONRPCDispatcher.
-  JSONRPCDispatcher Dispatcher([](const json::Expr &Params) {
-    replyError(ErrorCode::MethodNotFound, "method not found");
-  });
-  registerCallbackHandlers(Dispatcher, /*Callbacks=*/*this);
-
-  // Run the Language Server loop.
-  runLanguageServerLoop(In, Out, InputStyle, Dispatcher, IsDone);
-
-  // Make sure IsDone is set to true after this method exits to ensure assertion
-  // at the start of the method fires if it's ever executed again.
-  IsDone = true;
-
-  return ShutdownRequestReceived;
-}
-
 std::vector<Fix> ClangdLSPServer::getFixes(StringRef File,
                                            const clangd::Diagnostic &D) {
   std::lock_guard<std::mutex> Lock(FixItsMutex);
   auto DiagToFixItsIter = FixItsMap.find(File);
   if (DiagToFixItsIter == FixItsMap.end())
     return {};
 
   const auto &DiagToFixItsMap = DiagToFixItsIter->second;
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.
-  Out.writeMessage(json::obj{
+  sendMessage(json::obj{
       {"jsonrpc", "2.0"},
       {"method", "textDocument/publishDiagnostics"},
       {"params",
        json::obj{
            {"uri", URIForFile{File}},
            {"diagnostics", std::move(DiagnosticsJSON)},
        }},
   });
 }
 
 void ClangdLSPServer::reparseOpenedFiles() {
   for (const Path &FilePath : DraftMgr.getActiveFiles())
     Server.addDocument(FilePath, *DraftMgr.getDraft(FilePath),
                        WantDiagnostics::Auto);
 }
+
+bool& ClangdLSPServer::getIsDone() { return IsDone; }
+bool ClangdLSPServer::getShutdownRequestReceived() const { return ShutdownRequestReceived; }
 No newline at end of file

DispatcherCommon.h

+//===--- DispatcherCommon.h - Shared code for dispatchers -----------------===//
+//
+//                     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_DISPATCHERCOMMON_H
+#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_DISPATCHERCOMMON_H
+
+#include "JSONRPCDispatcher.h"
+#include "JSONExpr.h"
+#include "Trace.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/StringExtras.h"
+#include "llvm/Support/Chrono.h"
+#include "llvm/Support/SourceMgr.h"
+#include <istream>
+
+namespace clang {
+namespace clangd {
+
+extern Key<json::Expr> RequestID;
+extern Key<LSPOutput *> RequestOut;
+
+// When tracing, we trace a request and attach the repsonse in reply().
+// Because the Span isn't available, we find the current request using Context.
+class RequestSpan {
+  RequestSpan(json::obj *Args) : Args(Args) {}
+  std::mutex Mu;
+  json::obj *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);
+
+  // If there's an enclosing request and the tracer is interested, calls \p F
+  // with a json::obj 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);
+  }
+};
+
+void sendMessage(json::Expr &&Message);
+
+void reply(json::Expr &&Result);
+
+void replyError(ErrorCode code, const llvm::StringRef &Message);
+
+void call(StringRef Method, json::Expr &&Params);
+
+}
+}
+
+#endif
 No newline at end of file

DispatcherCommon.cpp

+//===--- DispatcherCommon.h - Shared code for dispatchers -----------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "DispatcherCommon.h"
+
+namespace clang {
+namespace clangd {
+
+Key<json::Expr> RequestID;
+Key<LSPOutput *> RequestOut;
+
+Context RequestSpan::stash(const trace::Span &Span) {
+  return Context::current().derive(
+      RSKey, std::unique_ptr<RequestSpan>(new RequestSpan(Span.Args)));
+}
+
+Key<std::unique_ptr<RequestSpan>> RequestSpan::RSKey;
+
+void sendMessage(json::Expr &&Message) {
+  Context::current()
+      .getExisting(RequestOut)
+      ->writeMessage(Message);
+}
+
+void reply(json::Expr &&Result) {
+  auto ID = Context::current().get(RequestID);
+  if (!ID) {
+    log("Attempted to reply to a notification!");
+    return;
+  }
+  RequestSpan::attach([&](json::obj &Args) { Args["Reply"] = Result; });
+  Context::current()
+      .getExisting(RequestOut)
+      ->writeMessage(json::obj{
+          {"jsonrpc", "2.0"},
+          {"id", *ID},
+          {"result", std::move(Result)},
+      });
+}
+
+void replyError(ErrorCode code, const llvm::StringRef &Message) {
+  log("Error " + Twine(static_cast<int>(code)) + ": " + Message);
+  RequestSpan::attach([&](json::obj &Args) {
+    Args["Error"] =
+        json::obj{{"code", static_cast<int>(code)}, {"message", Message.str()}};
+  });
+
+  if (auto ID = Context::current().get(RequestID)) {
+    Context::current()
+        .getExisting(RequestOut)
+        ->writeMessage(json::obj{
+            {"jsonrpc", "2.0"},
+            {"id", *ID},
+            {"error",
+             json::obj{{"code", static_cast<int>(code)}, {"message", Message}}},
+        });
+  }
+}
+
+void call(StringRef Method, json::Expr &&Params) {
+  // FIXME: Generate/Increment IDs for every request so that we can get proper
+  // replies once we need to.
+  RequestSpan::attach([&](json::obj &Args) {
+    Args["Call"] = json::obj{{"method", Method.str()}, {"params", Params}};
+  });
+  Context::current()
+      .getExisting(RequestOut)
+      ->writeMessage(json::obj{
+          {"jsonrpc", "2.0"},
+          {"id", 1},
+          {"method", Method},
+          {"params", std::move(Params)},
+      });
+}
+
+}
+}

JSONRPCDispatcher.h

-//===--- JSONRPCDispatcher.h - Main JSON parser entry point -----*- C++ -*-===//
+//===--- JSONRPCDispatcher.h - Main JSON RPC 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 "JSONExpr.h"
 #include "Logger.h"
+#include "LSPOutput.h"
 #include "Protocol.h"
 #include "Trace.h"
 #include "clang/Basic/LLVM.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringMap.h"
 #include <iosfwd>
 #include <mutex>
 
 namespace clang {
 namespace clangd {
 
 /// Encapsulates output and logs streams and provides thread-safe access to
 /// them.
-class JSONOutput : public Logger {
-  // FIXME(ibiryukov): figure out if we can shrink the public interface of
-  // JSONOutput now that we pass Context everywhere.
+class JSONOutput : public LSPOutput {
 public:
   JSONOutput(llvm::raw_ostream &Outs, llvm::raw_ostream &Logs,
              llvm::raw_ostream *InputMirror = nullptr, bool Pretty = false)
       : Pretty(Pretty), Outs(Outs), Logs(Logs), InputMirror(InputMirror) {}
 
   /// Emit a JSONRPC message.
-  void writeMessage(const json::Expr &Result);
+  void writeMessage(const json::Expr &Result) override;
 
   /// Write a line to the logging stream.
   void log(const Twine &Message) override;
 
   /// Mirror \p Message into InputMirror stream. Does nothing if InputMirror is
   /// null.
   /// Unlike other methods of JSONOutput, mirrorInput is not thread-safe.
-  void mirrorInput(const Twine &Message);
+  void mirrorInput(const Twine &Message) override;
 
   // Whether output should be pretty-printed.
   const bool Pretty;
 
 private:
   llvm::raw_ostream &Outs;
   llvm::raw_ostream &Logs;
   llvm::raw_ostream *InputMirror;
public:
   /// method is received.
   JSONRPCDispatcher(Handler UnknownHandler)
       : UnknownHandler(std::move(UnknownHandler)) {}
 
   /// Registers a Handler for the specified Method.
   void registerHandler(StringRef Method, Handler H);
 
   /// Parses a JSONRPC message and calls the Handler for it.
-  bool call(const json::Expr &Message, JSONOutput &Out) const;
+  bool call(const json::Expr &Message, LSPOutput &Out) const;
 
 private:
   llvm::StringMap<Handler> Handlers;
   Handler UnknownHandler;
 };
 
 /// Controls the way JSON-RPC messages are encoded (both input and output).
 enum JSONStreamStyle {
   /// Encoding per the LSP specification, with mandatory Content-Length header.
   Standard,
   /// Messages are delimited by a '---' line. Comment lines start with #.
   Delimited
 };
 
 /// Parses input queries from LSP client (coming from \p In) and runs call
 /// method of \p Dispatcher for each query.
 /// After handling each query checks if \p IsDone is set true and exits the loop
 /// if it is.
 /// 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.
-void runLanguageServerLoop(std::FILE *In, JSONOutput &Out,
+void runJSONRPCServerLoop(std::FILE *In, JSONOutput &Out,
                            JSONStreamStyle InputStyle,
                            JSONRPCDispatcher &Dispatcher, bool &IsDone);
 
 } // namespace clangd
 } // namespace clang
 
 #endif

JSONRPCDispatcher.cpp

-//===--- JSONRPCDispatcher.cpp - Main JSON parser entry point -------------===//
+//===--- JSONRPCDispatcher.h - Main JSON RPC entry point --------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
+#include "DispatcherCommon.h"
 #include "JSONRPCDispatcher.h"
 #include "JSONExpr.h"
-#include "ProtocolHandlers.h"
 #include "Trace.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/Chrono.h"
 #include "llvm/Support/Errno.h"
 #include "llvm/Support/SourceMgr.h"
 #include <istream>
 
 using namespace clang;
 using namespace clangd;
 
-namespace {
-static Key<json::Expr> RequestID;
-static Key<JSONOutput *> RequestOut;
-
-// When tracing, we trace a request and attach the repsonse in reply().
-// Because the Span isn't available, we find the current request using Context.
-class RequestSpan {
-  RequestSpan(json::obj *Args) : Args(Args) {}
-  std::mutex Mu;
-  json::obj *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::obj 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 JSONOutput::writeMessage(const json::Expr &Message) {
   std::string S;
   llvm::raw_string_ostream OS(S);
   if (Pretty)
     OS << llvm::formatv("{0:2}", Message);
   else
     OS << Message;
   OS.flush();

 void JSONOutput::mirrorInput(const Twine &Message) {
   if (!InputMirror)
     return;
 
   *InputMirror << Message;
   InputMirror->flush();
 }
 
-void clangd::reply(json::Expr &&Result) {
-  auto ID = Context::current().get(RequestID);
-  if (!ID) {
-    log("Attempted to reply to a notification!");
-    return;
-  }
-  RequestSpan::attach([&](json::obj &Args) { Args["Reply"] = Result; });
-  Context::current()
-      .getExisting(RequestOut)
-      ->writeMessage(json::obj{
-          {"jsonrpc", "2.0"},
-          {"id", *ID},
-          {"result", std::move(Result)},
-      });
-}
-
-void clangd::replyError(ErrorCode code, const llvm::StringRef &Message) {
-  log("Error " + Twine(static_cast<int>(code)) + ": " + Message);
-  RequestSpan::attach([&](json::obj &Args) {
-    Args["Error"] =
-        json::obj{{"code", static_cast<int>(code)}, {"message", Message.str()}};
-  });
-
-  if (auto ID = Context::current().get(RequestID)) {
-    Context::current()
-        .getExisting(RequestOut)
-        ->writeMessage(json::obj{
-            {"jsonrpc", "2.0"},
-            {"id", *ID},
-            {"error",
-             json::obj{{"code", static_cast<int>(code)}, {"message", Message}}},
-        });
-  }
-}
-
-void clangd::call(StringRef Method, json::Expr &&Params) {
-  // FIXME: Generate/Increment IDs for every request so that we can get proper
-  // replies once we need to.
-  RequestSpan::attach([&](json::obj &Args) {
-    Args["Call"] = json::obj{{"method", Method.str()}, {"params", Params}};
-  });
-  Context::current()
-      .getExisting(RequestOut)
-      ->writeMessage(json::obj{
-          {"jsonrpc", "2.0"},
-          {"id", 1},
-          {"method", Method},
-          {"params", std::move(Params)},
-      });
-}
-
 void JSONRPCDispatcher::registerHandler(StringRef Method, Handler H) {
   assert(!Handlers.count(Method) && "Handler already registered!");
   Handlers[Method] = std::move(H);
 }
 
-bool JSONRPCDispatcher::call(const json::Expr &Message, JSONOutput &Out) const {
+bool JSONRPCDispatcher::call(const json::Expr &Message, LSPOutput &Out) const {
   // Message must be an object with "jsonrpc":"2.0".
   auto *Object = Message.asObject();
   if (!Object || Object->getString("jsonrpc") != Optional<StringRef>("2.0"))
     return false;
   // ID may be any JSON value. If absent, this is a notification.
   llvm::Optional<json::Expr> ID;
   if (auto *I = Object->get("id"))
     ID = std::move(*I);
for (;;) {
     Size += Read;
   }
 }
 
 // Returns None when:
 //  - ferror() or feof() are set.
 //  - Content-Length is missing or empty (protocol error)
 static llvm::Optional<std::string> readStandardMessage(std::FILE *In,
-                                                       JSONOutput &Out) {
+                                                       LSPOutput &Out) {
   // A Language Server Protocol message starts with a set of HTTP headers,
   // delimited  by \r\n, and terminated by an empty line (\r\n).
   unsigned long long ContentLength = 0;
   std::string Line;
   while (true) {
     if (feof(In) || ferror(In) || !readLine(In, Line))
       return llvm::None;
 

 }
 
 // For lit tests we support a simplified syntax:
 // - messages are delimited by '---' on a line by itself
 // - lines starting with # are ignored.
 // This is a testing path, so favor simplicity over performance here.
 // When returning None, feof() or ferror() will be set.
 static llvm::Optional<std::string> readDelimitedMessage(std::FILE *In,
-                                                        JSONOutput &Out) {
+                                                        LSPOutput &Out) {
   std::string JSON;
   std::string Line;
   while (readLine(In, Line)) {
     auto LineRef = llvm::StringRef(Line).trim();
     if (LineRef.startswith("#")) // comment
       continue;
 
     // found a delimiter

 
 // The use of C-style std::FILE* IO deserves some explanation.
 // Previously, std::istream was used. When a debugger attached on MacOS, the
 // process received EINTR, the stream went bad, and clangd exited.
 // A retry-on-EINTR loop around reads solved this problem, but caused clangd to
 // sometimes hang rather than exit on other OSes. The interaction between
 // istreams and signals isn't well-specified, so it's hard to get this right.
 // The C APIs seem to be clearer in this respect.
-void clangd::runLanguageServerLoop(std::FILE *In, JSONOutput &Out,
+void clangd::runJSONRPCServerLoop(std::FILE *In, JSONOutput &Out,
                                    JSONStreamStyle InputStyle,
                                    JSONRPCDispatcher &Dispatcher,
                                    bool &IsDone) {
   auto &ReadMessage =
       (InputStyle == Delimited) ? readDelimitedMessage : readStandardMessage;
   while (!IsDone && !feof(In)) {
     if (ferror(In)) {
       log("IO error: " + llvm::sys::StrError());

LSPOutput.h

+//===--- LSPOutput.h - LSP output interface ---------------------*- 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_LSPOUTPUT_H
+#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_LSPOUTPUT_H
+
+#include "JSONExpr.h"
+#include "Logger.h"
+#include "clang/Basic/LLVM.h"
+
+namespace clang {
+namespace clangd {
+
+/// Encapsulates output and logs streams and provides thread-safe access to
+/// them.
+class LSPOutput : public Logger {
+  // FIXME(ibiryukov): figure out if we can shrink the public interface of
+  // LSPOutput now that we pass Context everywhere.
+public:
+  /// Emit LSP message.
+  virtual void writeMessage(const json::Expr &Result) = 0;
+
+  /// Write a line to log.
+  virtual void log(const Twine &Message) = 0;
+
+  /// FIXME: Does it make sense for all implementations?
+  /// Mirror \p Message into InputMirror stream. Does nothing if InputMirror is
+  /// null.
+  /// Unlike other methods of LSPOutput, mirrorInput is not thread-safe.
+  virtual void mirrorInput(const Twine &Message) = 0;
+};
+
+} // namespace clangd
+} // namespace clang
+
+#endif

ProtocolHandlers.h

public:
   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);
+// 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.
+template<typename DispatcherType>
+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::Expr &RawParams) {
+      typename std::remove_reference<Param>::type P;
+      if (fromJSON(RawParams, P)) {
+        (Callbacks->*Handler)(P);
+      } else {
+        log("Failed to decode " + Method + " request.");
+      }
+    });
+  }
+
+  DispatcherType &Dispatcher;
+  ProtocolCallbacks *Callbacks;
+};
+
+template<typename DispatcherType>
+void registerCallbackHandlers(DispatcherType &Dispatcher,
+                              ProtocolCallbacks &Callbacks) {
+  HandlerRegisterer<DispatcherType> 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/switchSourceHeader",
+           &ProtocolCallbacks::onSwitchSourceHeader);
+  Register("textDocument/rename", &ProtocolCallbacks::onRename);
+  Register("textDocument/hover", &ProtocolCallbacks::onHover);
+  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);
+}
+
 
 } // namespace clangd
 } // namespace clang
 
 #endif

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;
-
-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::Expr &RawParams) {
-      typename std::remove_reference<Param>::type P;
-      if (fromJSON(RawParams, P)) {
-        (Callbacks->*Handler)(P);
-      } else {
-        log("Failed to decode " + Method + " request.");
-      }
-    });
-  }
-
-  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/switchSourceHeader",
-           &ProtocolCallbacks::onSwitchSourceHeader);
-  Register("textDocument/rename", &ProtocolCallbacks::onRename);
-  Register("textDocument/hover", &ProtocolCallbacks::onHover);
-  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);
-}

tool/ClangdMain.cpp

 #include "index/SymbolYAML.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"
 #include <cstdlib>
-#include <iostream>
 #include <memory>
 #include <string>
 #include <thread>
 
 using namespace clang;
 using namespace clang::clangd;
 
 namespace {
if (EC) {
       Tracer = trace::createJSONTracer(*TraceStream, PrettyPrint);
     }
   }
 
   llvm::Optional<trace::Session> TracingSession;
   if (Tracer)
     TracingSession.emplace(*Tracer);
 
-  JSONOutput Out(llvm::outs(), llvm::errs(),
-                 InputMirrorStream ? InputMirrorStream.getPointer() : nullptr,
-                 PrettyPrint);
-
-  clangd::LoggingSession LoggingSession(Out);
-
   // If --compile-commands-dir arg was invoked, check value and override default
   // path.
   llvm::Optional<Path> CompileCommandsDirPath;
   if (CompileCommandsDir.empty()) {
     CompileCommandsDirPath = llvm::None;
   } else if (!llvm::sys::path::is_absolute(CompileCommandsDir) ||
              !llvm::sys::fs::exists(CompileCommandsDir)) {
     llvm::errs() << "Path specified by --compile-commands-dir either does not "
int main(int argc, char *argv[]) {
   }
   Opts.AsyncThreadsCount = WorkerThreadsCount;
 
   clangd::CodeCompleteOptions CCOpts;
   CCOpts.IncludeIneligibleResults = IncludeIneligibleResults;
   CCOpts.Limit = LimitResults;
   CCOpts.BundleOverloads = CompletionStyle != Detailed;
 
-  // Initialize and run ClangdLSPServer.
-  ClangdLSPServer LSPServer(Out, CCOpts, CompileCommandsDirPath, Opts);
-  constexpr int NoShutdownRequestErrorCode = 1;
+  ClangdLSPServer LSPServer(CCOpts, CompileCommandsDirPath, Opts);
+
   llvm::set_thread_name("clangd.main");
   // Change stdin to binary to not lose \r\n on windows.
   llvm::sys::ChangeStdinToBinary();
-  return LSPServer.run(stdin, InputStyle) ? 0 : NoShutdownRequestErrorCode;
+  {
+    JSONRPCDispatcher Dispatcher([](const json::Expr &Params) {
+      replyError(ErrorCode::MethodNotFound, "method not found");
+    });
+    registerCallbackHandlers(Dispatcher, /*Callbacks=*/LSPServer);
+
+
+    JSONOutput Out(llvm::outs(), llvm::errs(),
+                 InputMirrorStream ? InputMirrorStream.getPointer() : nullptr,
+                 PrettyPrint);
+
+    clangd::LoggingSession LoggingSession(Out);
+
+    runJSONRPCServerLoop(stdin, Out, InputStyle, Dispatcher, LSPServer.getIsDone());
+  }
+
+  constexpr int NoShutdownRequestErrorCode = 1;
+  return LSPServer.getShutdownRequestReceived() ? 0 : NoShutdownRequestErrorCode;
 }