[clangd] Refactor JSON-over-stdin/stdout code into Transport abstraction.

[sammccall]

Summary

This paves the way for alternative transports (mac XPC, maybe messagepack?),
and also generally improves layering: testing ClangdLSPServer becomes less of
a pipe dream, we split up the JSONOutput monolith, etc.

This isn't a final state, much of what remains in JSONRPCDispatcher can go away,
handlers can call reply() on the transport directly, JSONOutput can be renamed
to StreamLogger and removed, etc. But this patch is sprawling already.

The main observable change (see tests) is that hitting EOF on input is now an
error: the client should send the 'exit' notification.
This is defensible: the protocol doesn't spell this case out. Reproducing the
current behavior for all combinations of shutdown/exit/EOF clutters interfaces.
We can iterate on this if desired.

[sammccall]

This patch is big and hard to navigate. I tried to keep it contained, but JSONRPCDispatcher is pretty tangled.

The main parts are:

• Transport.h is the key new abstraction that should be understood first
• JSONTransport.cpp is its standard implementation. The raw IO stuff in there is just lifted from the old JSONRPCDispatcher, and isn't very interesting. Everything else should be though - it shows how the new interface works.
• The ClangdLSPServer and ClangdMain changes show how the new interface is used
• JSONRPCDispatcher has had most of its guts ripped out. There's more that can be done there, that file may go away completely.
• The test changes just reflect the shutdown getting a little bit stricter (they also pass before this patch)

[sammccall]

For reference, @jkorous has a WIP in D53290.
It's scope is a superset, and I think everything in common is basically the same (they were both based on a common prototype).
Jan is out at the moment, so I think it makes sense to move ahead with this piece - it's progress, and the smaller scope makes landing the change easier.

I just realized though there are no JSONTransport unit tests in this patch, I'll add those.

[sammccall]

Add unit test for JSON transport.

[ioeric]

Nice! The code looks much clearer. Just some nits

clangd/ClangdLSPServer.cpp

156

?

clangd/ClangdLSPServer.h

45

doc is outdated now

clangd/JSONRPCDispatcher.cpp

218

The comment should also be moved?

clangd/JSONRPCDispatcher.h

29

nit: make the FIXME more concrete?

clangd/JSONTransport.cpp

24

Is E guaranteed to be an LSPError? Do we need a handler for other errors?

84

Maybe add a comment here that the return value of handleMessage indicates whether to shutdown. It' not clear without context.

130

nit: add 2.0 to the log?

clangd/Transport.h

51

nit: maybe handleNotify to make the two directions more distinguishable? This can get a bit counter-intuitive in the call sites.

unittests/clangd/JSONTransportTests.cpp

106 ↗

(On Diff #169807)

Would operator<< be used when the error is real (i.e. check passes)?

[sammccall]

Addressed review comments.
Also inverted the sense of the boolean return from onNotify etc, was "should
exit" and is now "should continue", which I think is slightly clearer.

[sammccall]

clangd/ClangdLSPServer.cpp

156

Fixed this comment.
I'm not totally happy with the shutdown mechanism, it's noisier and more magic than I'd like. May find something nicer to replace it with.

clangd/JSONTransport.cpp

24

This actually handles generic errors as well, Error just has a really confusing API.
Rewrote this to be more direct.

clangd/Transport.h

51

onNotify etc

unittests/clangd/JSONTransportTests.cpp

106 ↗

(On Diff #169807)

No, EXPECT_* are magic macros that only evaluate the right hand side of << if they fire.

EXPECT_FALSE(expr) expands to something like if (expr) CreateErrorStream("expr"), so the << blah becomes part of the if statement.
(And the temporary stream returned from CreateErrorStream does the actual logging/mark-the-test-as-failed in its destructor).

[ioeric]

clangd/JSONRPCDispatcher.h

70

s/true/false/

clangd/CMakeLists.txt

add_clang_library(clangDaemon
   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

clangd/ClangdLSPServer.h

 /// 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 {
 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(Transport &Transport,
+                  const clangd::CodeCompleteOptions &CCOpts,
                   llvm::Optional<Path> CompileCommandsDir,
                   bool ShouldUseInMemoryCDB, 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 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;
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();
   void applyConfiguration(const ClangdConfigurationParamsChange &Settings);
 
-  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;
-
   std::mutex FixItsMutex;
   typedef std::map<clangd::Diagnostic, std::vector<Fix>, LSPDiagnosticCompare>
       DiagnosticToReplacementMap;
   /// Caches FixIts per file and diagnostics
   llvm::StringMap<DiagnosticToReplacementMap> FixItsMap;
 
   /// Encapsulates the directory-based or the in-memory compilation database
   /// that's used by the LSP server.
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;
   };
 
+  clangd::Transport &Transport;
   // 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

 }
 
 void ClangdLSPServer::onShutdown(ShutdownParams &Params) {
   // Do essentially nothing, just say we're ready to exit.
   ShutdownRequestReceived = true;
   reply(nullptr);
 }
 
-void ClangdLSPServer::onExit(ExitParams &Params) { IsDone = true; }
+void ClangdLSPServer::onExit(ExitParams &Params) {
+  // XXX handler should return true.
+}
 
 void ClangdLSPServer::onDocumentDidOpen(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;
 
Server->findReferences(Params.textDocument.uri.file(), Params.position,
                            if (!Locations)
                              return replyError(
                                  ErrorCode::InternalError,
                                  llvm::toString(Locations.takeError()));
                            reply(llvm::json::Array(*Locations));
                          });
 }
 
-ClangdLSPServer::ClangdLSPServer(JSONOutput &Out,
+ClangdLSPServer::ClangdLSPServer(class Transport &Transport,
                                  const clangd::CodeCompleteOptions &CCOpts,
                                  llvm::Optional<Path> CompileCommandsDir,
                                  bool ShouldUseInMemoryCDB,
                                  const ClangdServer::Options &Opts)
-    : Out(Out), CDB(ShouldUseInMemoryCDB ? CompilationDB::makeInMemory()
-                                         : CompilationDB::makeDirectoryBased(
-                                               std::move(CompileCommandsDir))),
+    : Transport(Transport),
+      CDB(ShouldUseInMemoryCDB ? CompilationDB::makeInMemory()
       CCOpts(CCOpts), SupportedSymbolKinds(defaultSymbolKinds()),
       SupportedCompletionItemKinds(defaultCompletionItemKinds()),
       Server(new ClangdServer(CDB.getCDB(), FSProvider, /*DiagConsumer=*/*this,
                               Opts)) {}
 
-bool ClangdLSPServer::run(std::FILE *In, JSONStreamStyle InputStyle) {
-  assert(!IsDone && "Run was called before");
+bool ClangdLSPServer::run() {
   assert(Server);
 
   // Set up JSONRPCDispatcher.
   JSONRPCDispatcher Dispatcher([](const json::Value &Params) {
     replyError(ErrorCode::MethodNotFound, "method not found");
+    return false;
   });
   registerCallbackHandlers(Dispatcher, /*Callbacks=*/*this);
 
   // Run the Language Server loop.
-  runLanguageServerLoop(In, Out, InputStyle, Dispatcher, IsDone);
+  bool CleanExit = true;
+  if (auto Err = Dispatcher.runLanguageServerLoop(Transport)) {
+    elog("Transport error: {0}", std::move(Err));
+    CleanExit = false;
+  }
 
-  // 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;
   // Destroy ClangdServer to ensure all worker threads finish.
   Server.reset();
 
-  return ShutdownRequestReceived;
+  return CleanExit && 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 {};
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::Object{
-      {"jsonrpc", "2.0"},
-      {"method", "textDocument/publishDiagnostics"},
-      {"params",
+  Transport.notify("textDocument/publishDiagnostics",
                    json::Object{
                        {"uri", URIForFile{File}},
                        {"diagnostics", std::move(DiagnosticsJSON)},
-       }},
                    });
 }
 
 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 {
 
-/// Encapsulates output and logs streams and provides thread-safe access to
-/// them.
+// Logs to an output stream, such as stderr.
+// TODO: rename and move.
 class JSONOutput : public Logger {
-  // FIXME(ibiryukov): figure out if we can shrink the public interface of
-  // JSONOutput now that we pass Context everywhere.
 public:
-  JSONOutput(llvm::raw_ostream &Outs, llvm::raw_ostream &Logs,
-             Logger::Level MinLevel, llvm::raw_ostream *InputMirror = nullptr,
-             bool Pretty = false)
-      : Pretty(Pretty), MinLevel(MinLevel), Outs(Outs), Logs(Logs),
-        InputMirror(InputMirror) {}
-
-  /// Emit a JSONRPC message.
-  void writeMessage(const llvm::json::Value &Result);
+  JSONOutput(llvm::raw_ostream &Logs, Logger::Level MinLevel)
+      : MinLevel(MinLevel), Logs(Logs) {}
 
   /// Write a line to the logging stream.
   void log(Level, const llvm::formatv_object_base &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);
-
-  // Whether output should be pretty-printed.
-  const bool Pretty;
-
 private:
   Logger::Level MinLevel;
-  llvm::raw_ostream &Outs;
   llvm::raw_ostream &Logs;
-  llvm::raw_ostream *InputMirror;
 
   std::mutex StreamMutex;
 };
 
 /// 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 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 {
+class JSONRPCDispatcher : private Transport::MessageHandler {
 public:
   /// A handler responds to requests for a particular method name.
+  /// It returns true 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<void(const llvm::json::Value &)>;
+  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 a JSONRPC message and calls the Handler for it.
-  bool call(const llvm::json::Value &Message, JSONOutput &Out);
+  /// 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 reply(llvm::json::Value ID,
+             llvm::Expected<llvm::json::Value> Result) override;
+  bool notify(llvm::StringRef Method, llvm::json::Value Message) override;
+  bool call(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;
 };
 
-/// 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,
-                           JSONStreamStyle InputStyle,
-                           JSONRPCDispatcher &Dispatcher, bool &IsDone);
 } // 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/Chrono.h"
 #include "llvm/Support/Errno.h"
 #include "llvm/Support/FormatVariadic.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<JSONOutput *> RequestOut;
+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;
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::Value &Message) {
-  std::string S;
-  llvm::raw_string_ostream OS(S);
-  if (Pretty)
-    OS << llvm::formatv("{0:2}", Message);
-  else
-    OS << Message;
-  OS.flush();
-
-  {
-    std::lock_guard<std::mutex> Guard(StreamMutex);
-    Outs << "Content-Length: " << S.size() << "\r\n\r\n" << S;
-    Outs.flush();
-  }
-  vlog(">>> {0}\n", S);
-}
-
 void JSONOutput::log(Logger::Level Level,
                      const llvm::formatv_object_base &Message) {
   if (Level < MinLevel)
     return;
   llvm::sys::TimePoint<> Timestamp = std::chrono::system_clock::now();
   trace::log(Message);
   std::lock_guard<std::mutex> Guard(StreamMutex);
   Logs << llvm::formatv("{0}[{1:%H:%M:%S.%L}] {2}\n", indicator(Level),
                         Timestamp, Message);
   Logs.flush();
 }
 
-void JSONOutput::mirrorInput(const Twine &Message) {
-  if (!InputMirror)
-    return;
-
-  *InputMirror << Message;
-  InputMirror->flush();
-}
-
 void clangd::reply(json::Value &&Result) {
   auto ID = getRequestId();
   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(RequestOut)
-      ->writeMessage(json::Object{
-          {"jsonrpc", "2.0"},
-          {"id", *ID},
-          {"result", std::move(Result)},
-      });
+      .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 = getRequestId()) {
     log("--> reply({0}) error: {1}", *ID, Message);
     Context::current()
-        .getExisting(RequestOut)
-        ->writeMessage(json::Object{
-            {"jsonrpc", "2.0"},
-            {"id", *ID},
-            {"error", json::Object{{"code", static_cast<int>(Code)},
-                                   {"message", Message}}},
-        });
+        .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(RequestOut)
-      ->writeMessage(json::Object{
-          {"jsonrpc", "2.0"},
-          {"id", ID},
-          {"method", Method},
-          {"params", std::move(Params)},
-      });
+      .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"))
-        return cancelRequest(*ID);
+      if (auto *ID = O->get("id")) {
+        cancelRequest(*ID);
+        return false;
+      }
     log("Bad cancellation request: {0}", Params);
+    return false;
   });
 }
 
 void JSONRPCDispatcher::registerHandler(StringRef Method, Handler H) {
   assert(!Handlers.count(Method) && "Handler already registered!");
   Handlers[Method] = std::move(H);
 }
 
-static void logIncomingMessage(const llvm::Optional<json::Value> &ID,
-                               llvm::Optional<StringRef> Method,
-                               const json::Object *Error) {
-  if (Method) { // incoming request
-    if (ID)     // call
-      log("<-- {0}({1})", *Method, *ID);
-    else // notification
-      log("<-- {0}", *Method);
-  } else if (ID) { // response, ID must be provided
-    if (Error)
-      log("<-- reply({0}) error: {1}", *ID,
-          Error->getString("message").getValueOr("<no message>"));
-    else
-      log("<-- reply({0})", *ID);
-  }
-}
-
-bool JSONRPCDispatcher::call(const json::Value &Message, JSONOutput &Out) {
-  // Message must be an object with "jsonrpc":"2.0".
-  auto *Object = Message.getAsObject();
-  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::Value> ID;
-  if (auto *I = Object->get("id"))
-    ID = std::move(*I);
-  auto Method = Object->getString("method");
-  logIncomingMessage(ID, Method, Object->getObject("error"));
-  if (!Method) // We only handle incoming requests, and ignore responses.
-    return false;
-  // Params should be given, use null if not.
-  json::Value Params = nullptr;
-  if (auto *P = Object->get("params"))
-    Params = std::move(*P);
-
-  auto I = Handlers.find(*Method);
+bool JSONRPCDispatcher::call(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 WithRequestOut(RequestOut, &Out);
-  llvm::Optional<WithContextValue> WithID;
-  if (ID)
-    WithID.emplace(RequestID, *ID);
+  WithContextValue WithID(RequestID, ID);
 
   // Create a tracing Span covering the whole request lifetime.
-  trace::Span Tracer(*Method);
-  if (ID)
-    SPAN_ATTACH(Tracer, "ID", *ID);
+  trace::Span Tracer(Method);
+  SPAN_ATTACH(Tracer, "ID", ID);
   SPAN_ATTACH(Tracer, "Params", Params);
 
-  // Requests with IDs can be canceled by the client. Add cancellation context.
-  llvm::Optional<WithContext> WithCancel;
-  if (ID)
-    WithCancel.emplace(cancelableRequestContext(*ID));
+  // 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::notify(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));
-  Handler(std::move(Params));
-  return true;
+  return Handler(std::move(Params));
+}
+
+bool JSONRPCDispatcher::reply(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 false;
 }
 
 // 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();

 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.
 }
 
-// Tries to read a line up to and including \n.
-// If failing, feof() or ferror() will be set.
-static bool readLine(std::FILE *In, std::string &Out) {
-  static constexpr int BufSize = 1024;
-  size_t Size = 0;
-  Out.clear();
-  for (;;) {
-    Out.resize(Size + BufSize);
-    // Handle EINTR which is sent when a debugger attaches on some platforms.
-    if (!llvm::sys::RetryAfterSignal(nullptr, ::fgets, &Out[Size], BufSize, In))
-      return false;
-    clearerr(In);
-    // If the line contained null bytes, anything after it (including \n) will
-    // be ignored. Fortunately this is not a legal header or JSON.
-    size_t Read = std::strlen(&Out[Size]);
-    if (Read > 0 && Out[Size + Read - 1] == '\n') {
-      Out.resize(Size + Read);
-      return true;
-    }
-    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) {
-  // 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;
-
-    Out.mirrorInput(Line);
-    llvm::StringRef LineRef(Line);
-
-    // We allow comments in headers. Technically this isn't part
-    // of the LSP specification, but makes writing tests easier.
-    if (LineRef.startswith("#"))
-      continue;
-
-    // Content-Length is a mandatory header, and the only one we handle.
-    if (LineRef.consume_front("Content-Length: ")) {
-      if (ContentLength != 0) {
-        elog("Warning: Duplicate Content-Length header received. "
-             "The previous value for this message ({0}) was ignored.",
-             ContentLength);
-      }
-      llvm::getAsUnsignedInteger(LineRef.trim(), 0, ContentLength);
-      continue;
-    } else if (!LineRef.trim().empty()) {
-      // It's another header, ignore it.
-      continue;
-    } else {
-      // An empty line indicates the end of headers.
-      // Go ahead and read the JSON.
-      break;
-    }
-  }
-
-  // The fuzzer likes crashing us by sending "Content-Length: 9999999999999999"
-  if (ContentLength > 1 << 30) { // 1024M
-    elog("Refusing to read message with long Content-Length: {0}. "
-         "Expect protocol errors",
-         ContentLength);
-    return llvm::None;
-  }
-  if (ContentLength == 0) {
-    log("Warning: Missing Content-Length header, or zero-length message.");
-    return llvm::None;
-  }
-
-  std::string JSON(ContentLength, '\0');
-  for (size_t Pos = 0, Read; Pos < ContentLength; Pos += Read) {
-    // Handle EINTR which is sent when a debugger attaches on some platforms.
-    Read = llvm::sys::RetryAfterSignal(0u, ::fread, &JSON[Pos], 1,
-                                       ContentLength - Pos, In);
-    Out.mirrorInput(StringRef(&JSON[Pos], Read));
-    if (Read == 0) {
-      elog("Input was aborted. Read only {0} bytes of expected {1}.", Pos,
-           ContentLength);
-      return llvm::None;
-    }
-    clearerr(In); // If we're done, the error was transient. If we're not done,
-                  // either it was transient or we'll see it again on retry.
-    Pos += Read;
-  }
-  return std::move(JSON);
-}
-
-// 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) {
-  std::string JSON;
-  std::string Line;
-  while (readLine(In, Line)) {
-    auto LineRef = llvm::StringRef(Line).trim();
-    if (LineRef.startswith("#")) // comment
-      continue;
-
-    // found a delimiter
-    if (LineRef.rtrim() == "---")
-      break;
-
-    JSON += Line;
-  }
-
-  if (ferror(In)) {
-    elog("Input error while reading message!");
-    return llvm::None;
-  } else { // Including EOF
-    Out.mirrorInput(
-        llvm::formatv("Content-Length: {0}\r\n\r\n{1}", JSON.size(), JSON));
-    return std::move(JSON);
-  }
-}
-
 // 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,
-                                   JSONStreamStyle InputStyle,
-                                   JSONRPCDispatcher &Dispatcher,
-                                   bool &IsDone) {
-  auto &ReadMessage =
-      (InputStyle == Delimited) ? readDelimitedMessage : readStandardMessage;
-  while (!IsDone && !feof(In)) {
-    if (ferror(In)) {
-      elog("IO error: {0}", llvm::sys::StrError());
-      return;
-    }
-    if (auto JSON = ReadMessage(In, Out)) {
-      if (auto Doc = json::parse(*JSON)) {
-        // Log the formatted message.
-        vlog(Out.Pretty ? "<<< {0:2}\n" : "<<< {0}\n", *Doc);
-        // Finally, execute the action for this JSON message.
-        if (!Dispatcher.call(*Doc, Out))
-          elog("JSON dispatch failed!");
-      } else {
-        // Parse error. Log the raw message.
-        vlog("<<< {0}\n", *JSON);
-        elog("JSON parse error: {0}", llvm::toString(Doc.takeError()));
-      }
-    }
-  }
+llvm::Error JSONRPCDispatcher::runLanguageServerLoop(Transport &Transport) {
+  // Propagate transport to all handlers so they can reply.
+  WithContextValue WithTransport(CurrentTransport, &Transport);
+  return Transport.loop(*this);
 }
 
 const json::Value *clangd::getRequestId() {
   return Context::current().get(RequestID);
 }

clangd/JSONTransport.cpp

+//===--- JSONTransport.cpp - sending and receiving LSP messages over JSON -===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+#include "Logger.h"
+#include "Protocol.h" // For LSPError
+#include "Transport.h"
+#include "llvm/Support/Errno.h"
+
+using namespace llvm;
+namespace clang {
+namespace clangd {
+namespace {
+
+json::Object encodeError(Error E) {
+  ErrorCode Code = ErrorCode::UnknownErrorCode;
+  std::string Msg = toString(handleErrors(std::move(E), [&](const LSPError &L) {
+    Code = L.Code;
+    return make_error<StringError>(L.Message, inconvertibleErrorCode());
+  }));
+  return json::Object{
+      {"message", std::move(Msg)},
+      {"code", int64_t(Code)},
+  };
+}
+
+Error decodeError(const json::Object &O) {
+  std::string Msg = O.getString("message").getValueOr("Unspecified error");
+  if (auto Code = O.getInteger("code"))
+    return make_error<LSPError>(std::move(Msg), ErrorCode(*Code));
+  return make_error<StringError>(std::move(Msg), inconvertibleErrorCode());
+}
+
+class JSONTransport : public Transport {
+public:
+  JSONTransport(std::FILE *In, llvm::raw_ostream &Out,
+                llvm::raw_ostream *InMirror, bool Pretty, JSONStreamStyle Style)
+      : In(In), Out(Out), InMirror(InMirror ? *InMirror : nulls()),
+        Pretty(Pretty), Style(Style) {}
+
+  void notify(StringRef Method, json::Value Params) override {
+    sendMessage(json::Object{
+        {"jsonrpc", "2.0"},
+        {"method", Method},
+        {"params", std::move(Params)},
+    });
+  }
+  void call(StringRef Method, json::Value Params, json::Value ID) override {
+    sendMessage(json::Object{
+        {"jsonrpc", "2.0"},
+        {"id", std::move(ID)},
+        {"method", Method},
+        {"params", std::move(Params)},
+    });
+  }
+  void reply(json::Value ID, Expected<json::Value> Result) override {
+    if (Result) {
+      sendMessage(json::Object{
+          {"jsonrpc", "2.0"},
+          {"id", std::move(ID)},
+          {"result", std::move(*Result)},
+      });
+    } else {
+      sendMessage(json::Object{
+          {"jsonrpc", "2.0"},
+          {"id", std::move(ID)},
+          {"error", encodeError(Result.takeError())},
+      });
+    }
+  }
+
+  Error loop(MessageHandler &Handler) override {
+    while (!feof(In)) {
+      if (ferror(In))
+        return errorCodeToError(std::error_code(errno, std::system_category()));
+      if (auto JSON = readRawMessage()) {
+        if (auto Doc = json::parse(*JSON)) {
+          vlog(Pretty ? "<<< {0:2}\n" : "<<< {0}\n", *Doc);
+          if (handleMessage(std::move(*Doc), Handler))
+            return Error::success();
+        } else {
+          // Parse error. Log the raw message.
+          vlog("<<< {0}\n", *JSON);
+          elog("JSON parse error: {0}", llvm::toString(Doc.takeError()));
+        }
+      }
+    }
+    return errorCodeToError(std::make_error_code(std::errc::io_error));
+  }
+
+private:
+  // Dispatches incoming message to Handler notify/call/reply.
+  bool handleMessage(llvm::json::Value Message, MessageHandler &Handler);
+  // Writes outgoing message to Out stream.
+  void sendMessage(llvm::json::Value Message) {
+    std::string S;
+    llvm::raw_string_ostream OS(S);
+    OS << llvm::formatv(Pretty ? "{0:2}" : "{0}", Message);
+    OS.flush();
+    Out << "Content-Length: " << S.size() << "\r\n\r\n" << S;
+    Out.flush();
+    vlog(">>> {0}\n", S);
+  }
+
+  // Read raw string messages from input stream.
+  llvm::Optional<std::string> readRawMessage() {
+    return Style == JSONStreamStyle::Delimited ? readDelimitedMessage()
+                                               : readStandardMessage();
+  }
+  llvm::Optional<std::string> readDelimitedMessage();
+  llvm::Optional<std::string> readStandardMessage();
+
+  std::FILE *In;
+  llvm::raw_ostream &Out;
+  llvm::raw_ostream &InMirror;
+  bool Pretty;
+  JSONStreamStyle Style;
+};
+
+bool JSONTransport::handleMessage(llvm::json::Value Message,
+                                  MessageHandler &Handler) {
+  // Message must be an object with "jsonrpc":"2.0".
+  auto *Object = Message.getAsObject();
+  if (!Object || Object->getString("jsonrpc") != Optional<StringRef>("2.0")) {
+    elog("Not a JSON-RPC message: {0:2}", Message);
+    return false;
+  }
+  // ID may be any JSON value. If absent, this is a notification.
+  llvm::Optional<json::Value> ID;
+  if (auto *I = Object->get("id"))
+    ID = std::move(*I);
+  auto Method = Object->getString("method");
+  if (!Method) { // This is a response.
+    if (!ID) {
+      elog("No method and no response ID: {0:2}", Message);
+      return false;
+    }
+    if (auto *Err = Object->getObject("error"))
+      return Handler.reply(std::move(*ID), decodeError(*Err));
+    // Result should be given, use null if not.
+    json::Value Result = nullptr;
+    if (auto *R = Object->get("result"))
+      Result = std::move(*R);
+    return Handler.reply(std::move(*ID), std::move(Result));
+  }
+  // Params should be given, use null if not.
+  json::Value Params = nullptr;
+  if (auto *P = Object->get("params"))
+    Params = std::move(*P);
+
+  if (ID)
+    return Handler.call(*Method, std::move(Params), std::move(*ID));
+  else
+    return Handler.notify(*Method, std::move(Params));
+}
+
+// Tries to read a line up to and including \n.
+// If failing, feof() or ferror() will be set.
+bool readLine(std::FILE *In, std::string &Out) {
+  static constexpr int BufSize = 1024;
+  size_t Size = 0;
+  Out.clear();
+  for (;;) {
+    Out.resize(Size + BufSize);
+    // Handle EINTR which is sent when a debugger attaches on some platforms.
+    if (!llvm::sys::RetryAfterSignal(nullptr, ::fgets, &Out[Size], BufSize, In))
+      return false;
+    clearerr(In);
+    // If the line contained null bytes, anything after it (including \n) will
+    // be ignored. Fortunately this is not a legal header or JSON.
+    size_t Read = std::strlen(&Out[Size]);
+    if (Read > 0 && Out[Size + Read - 1] == '\n') {
+      Out.resize(Size + Read);
+      return true;
+    }
+    Size += Read;
+  }
+}
+
+// Returns None when:
+//  - ferror() or feof() are set.
+//  - Content-Length is missing or empty (protocol error)
+llvm::Optional<std::string> JSONTransport::readStandardMessage() {
+  // 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;
+    InMirror << Line;
+
+    llvm::StringRef LineRef(Line);
+
+    // We allow comments in headers. Technically this isn't part
+
+    // of the LSP specification, but makes writing tests easier.
+    if (LineRef.startswith("#"))
+      continue;
+
+    // Content-Length is a mandatory header, and the only one we handle.
+    if (LineRef.consume_front("Content-Length: ")) {
+      if (ContentLength != 0) {
+        elog("Warning: Duplicate Content-Length header received. "
+             "The previous value for this message ({0}) was ignored.",
+             ContentLength);
+      }
+      llvm::getAsUnsignedInteger(LineRef.trim(), 0, ContentLength);
+      continue;
+    } else if (!LineRef.trim().empty()) {
+      // It's another header, ignore it.
+      continue;
+    } else {
+      // An empty line indicates the end of headers.
+      // Go ahead and read the JSON.
+      break;
+    }
+  }
+
+  // The fuzzer likes crashing us by sending "Content-Length: 9999999999999999"
+  if (ContentLength > 1 << 30) { // 1024M
+    elog("Refusing to read message with long Content-Length: {0}. "
+         "Expect protocol errors",
+         ContentLength);
+    return llvm::None;
+  }
+  if (ContentLength == 0) {
+    log("Warning: Missing Content-Length header, or zero-length message.");
+    return llvm::None;
+  }
+
+  std::string JSON(ContentLength, '\0');
+  for (size_t Pos = 0, Read; Pos < ContentLength; Pos += Read) {
+    // Handle EINTR which is sent when a debugger attaches on some platforms.
+    Read = llvm::sys::RetryAfterSignal(0u, ::fread, &JSON[Pos], 1,
+                                       ContentLength - Pos, In);
+    if (Read == 0) {
+      elog("Input was aborted. Read only {0} bytes of expected {1}.", Pos,
+           ContentLength);
+      return llvm::None;
+    }
+    InMirror << StringRef(&JSON[Pos], Read);
+    clearerr(In); // If we're done, the error was transient. If we're not done,
+                  // either it was transient or we'll see it again on retry.
+    Pos += Read;
+  }
+  return std::move(JSON);
+}
+
+// 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.
+llvm::Optional<std::string> JSONTransport::readDelimitedMessage() {
+  std::string JSON;
+  std::string Line;
+  while (readLine(In, Line)) {
+    InMirror << Line;
+    auto LineRef = llvm::StringRef(Line).trim();
+    if (LineRef.startswith("#")) // comment
+      continue;
+
+    // found a delimiter
+    if (LineRef.rtrim() == "---")
+      break;
+
+    JSON += Line;
+  }
+
+  if (ferror(In)) {
+    elog("Input error while reading message!");
+    return llvm::None;
+  }
+  return std::move(JSON); // Including at EOF
+}
+
+} // namespace
+
+std::unique_ptr<Transport> newJSONTransport(std::FILE *In,
+                                            llvm::raw_ostream &Out,
+                                            llvm::raw_ostream *InMirror,
+                                            bool Pretty,
+                                            JSONStreamStyle Style) {
+  return llvm::make_unique<JSONTransport>(In, Out, InMirror, Pretty, Style);
+}
+
+} // namespace clangd
+} // namespace clang

clangd/Protocol.h

enum class ErrorCode {
   InternalError = -32603,
 
   ServerNotInitialized = -32002,
   UnknownErrorCode = -32001,
 
   // Defined by the protocol.
   RequestCancelled = -32800,
 };
+// Models an LSP error as an llvm::Error.
+class LSPError : public llvm::ErrorInfo<LSPError> {
+public:
+  std::string Message;
+  ErrorCode Code;
+  static char ID;
+
+  LSPError(std::string Message, ErrorCode Code)
+      : Message(std::move(Message)), Code(Code) {}
+
+  void log(llvm::raw_ostream &OS) const override {
+    OS << int(Code) << ": " << Message;
+  }
+  std::error_code convertToErrorCode() const override {
+    return llvm::inconvertibleErrorCode();
+  }
+};
 
 struct URIForFile {
   URIForFile() = default;
   explicit URIForFile(std::string AbsPath);
 
   /// Retrieves absolute path to the file.
   llvm::StringRef file() const { return File; }
 

clangd/Protocol.cpp

 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace clang {
 namespace clangd {
 using namespace llvm;
 
+char LSPError::ID;
+
 URIForFile::URIForFile(std::string AbsPath) {
   assert(llvm::sys::path::is_absolute(AbsPath) && "the path is relative");
   File = std::move(AbsPath);
 }
 
 bool fromJSON(const json::Value &E, URIForFile &R) {
   if (auto S = E.getAsString()) {
     auto U = URI::parse(*S);

clangd/ProtocolHandlers.cpp

void operator()(StringRef Method, void (ProtocolCallbacks::*Handler)(Param)) {
     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";
     });
   }
 
   JSONRPCDispatcher &Dispatcher;
   ProtocolCallbacks *Callbacks;
 };
 
 } // namespace

clangd/Transport.h

+//===--- Transport.h - sending and receiving LSP messages -------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// The language server protocol is usually implemented by writing messages as
+// JSON-RPC over the stdin/stdout of a subprocess. However other communications
+// mechanisms are possible, such as XPC on mac.
+//
+// The Transport interface allows the mechanism to be replaced, and the JSONRPC
+// Transport is the standard implementation.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_TRANSPORT_H_
+#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_TRANSPORT_H_
+
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/JSON.h"
+#include "llvm/Support/raw_ostream.h"
+
+namespace clang {
+namespace clangd {
+
+// A transport is responsible for maintaining the connection to a client
+// application, and reading/writing structured messages to it.
+//
+// Transports have limited thread safety requirements:
+//  - messages will not be sent concurrently
+//  - messages MAY be sent while loop() is reading, or its callback is active
+class Transport {
+public:
+  virtual ~Transport() = default;
+
+  // Called by Clangd to send messages to the client.
+  virtual void notify(llvm::StringRef Method, llvm::json::Value Params) = 0;
+  virtual void call(llvm::StringRef Method, llvm::json::Value Params,
+                    llvm::json::Value ID) = 0;
+  virtual void reply(llvm::json::Value ID,
+                     llvm::Expected<llvm::json::Value> Result) = 0;
+
+  // Implemented by Clangd to handle incoming messages. (See loop() below).
+  class MessageHandler {
+  public:
+    virtual ~MessageHandler() = default;
+    virtual bool notify(llvm::StringRef Method, llvm::json::Value) = 0;
+    virtual bool call(llvm::StringRef Method, llvm::json::Value Params,
+                      llvm::json::Value ID) = 0;
+    virtual bool reply(llvm::json::Value ID,
+                       llvm::Expected<llvm::json::Value> Result) = 0;
+  };
+  // Called by Clangd to receive messages from the client.
+  // The transport should in turn invoke the handler to process messages.
+  // If handler returns true, the transport should immedately return success.
+  // (This is used to implement the `exit` notification).
+  // Otherwise, it returns an error when the transport becomes unusable.
+  virtual llvm::Error loop(MessageHandler &) = 0;
+};
+
+// 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
+};
+
+// Returns a Transport that speaks JSON-RPC over a pair of streams.
+// The input stream must be opened in binary mode.
+// If InMirror is set, data read will be echoed to it.
+std::unique_ptr<Transport>
+newJSONTransport(std::FILE *In, llvm::raw_ostream &Out,
+                 llvm::raw_ostream *InMirror, bool Pretty,
+                 JSONStreamStyle = JSONStreamStyle::Standard);
+
+} // namespace clangd
+} // namespace clang
+
+#endif

clangd/tool/ClangdMain.cpp

int main(int argc, char *argv[]) {
 
   llvm::Optional<trace::Session> TracingSession;
   if (Tracer)
     TracingSession.emplace(*Tracer);
 
   // Use buffered stream to stderr (we still flush each log message). Unbuffered
   // stream can cause significant (non-deterministic) latency for the logger.
   llvm::errs().SetBuffered();
-  JSONOutput Out(llvm::outs(), llvm::errs(), LogLevel,
-                 InputMirrorStream ? InputMirrorStream.getPointer() : nullptr,
-                 PrettyPrint);
-
-  clangd::LoggingSession LoggingSession(Out);
+  JSONOutput Logger(llvm::errs(), LogLevel);
+  clangd::LoggingSession LoggingSession(Logger);
 
   // 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)) {
if (!HeaderInsertionDecorators) {
     CCOpts.IncludeIndicator.Insert.clear();
     CCOpts.IncludeIndicator.NoInsert.clear();
   }
   CCOpts.SpeculativeIndexRequest = Opts.StaticIndex;
   CCOpts.EnableFunctionArgSnippets = EnableFunctionArgSnippets;
   CCOpts.AllScopes = AllScopesCompletion;
 
   // Initialize and run ClangdLSPServer.
+  // Change stdin to binary to not lose \r\n on windows.
+  llvm::sys::ChangeStdinToBinary();
+  auto Transport = newJSONTransport(
+      stdin, llvm::outs(),
+      InputMirrorStream ? InputMirrorStream.getPointer() : nullptr, PrettyPrint,
+      InputStyle);
   ClangdLSPServer LSPServer(
-      Out, CCOpts, CompileCommandsDirPath,
+      *Transport, CCOpts, CompileCommandsDirPath,
       /*ShouldUseInMemoryCDB=*/CompileArgsFrom == LSPCompileArgs, Opts);
   constexpr int NoShutdownRequestErrorCode = 1;
   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;
+  return LSPServer.run() ? 0 : NoShutdownRequestErrorCode;
 }

test/clangd/compile-commands-path-in-initialize.test

 {"jsonrpc":"2.0","id":0,"method":"workspace/didChangeConfiguration","params":{"settings":{"compilationDatabasePath":"INPUT_DIR/build-2"}}}
 # CHECK:   "method": "textDocument/publishDiagnostics",
 # CHECK-NEXT:   "params": {
 # CHECK-NEXT:     "diagnostics": [
 # CHECK-NEXT:       {
 # CHECK-NEXT:         "message": "MACRO is two",
 ---
 {"jsonrpc":"2.0","id":10000,"method":"shutdown"}
+---
+{"jsonrpc":"2.0","method":"exit"}

test/clangd/compile-commands-path.test

 {"jsonrpc":"2.0","id":0,"method":"workspace/didChangeConfiguration","params":{"settings":{"compilationDatabasePath":"INPUT_DIR/build-2"}}}
 # CHECK:   "method": "textDocument/publishDiagnostics",
 # CHECK-NEXT:   "params": {
 # CHECK-NEXT:     "diagnostics": [
 # CHECK-NEXT:       {
 # CHECK-NEXT:         "message": "MACRO is two",
 ---
 {"jsonrpc":"2.0","id":10000,"method":"shutdown"}
+---
+{"jsonrpc":"2.0","method":"exit"}

test/clangd/completion-snippets.test

 # CHECK-NEXT:          }
 # CHECK-NEXT:        }
 # CHECK-NEXT:      }
 # CHECK-NEXT:    }
 # CHECK-NEXT:    ]
 # CHECK-NEXT:  }
 ---
 {"jsonrpc":"2.0","id":4,"method":"shutdown"}
+---
+{"jsonrpc":"2.0","method":"exit"}

test/clangd/completion.test

 # CHECK-NEXT:            "line": 2
 # CHECK-NEXT:          }
 # CHECK-NEXT:        }
 # CHECK-NEXT:      }
 # CHECK-NEXT:    }
 # CHECK-NEXT:  ]
 ---
 {"jsonrpc":"2.0","id":4,"method":"shutdown"}
+---
+{"jsonrpc":"2.0","method":"exit"}

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/execute-command.test

 ---
 # TextEdit not decoded
 {"jsonrpc":"2.0","id":14,"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":"","newText":")"}]}}]}}
 ---
 # Command name after arguments
 {"jsonrpc":"2.0","id":9,"method":"workspace/executeCommand","params":{"arguments":[{"custom":"foo", "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":")"}]}}],"command":"clangd.applyFix"}}
 ---
 {"jsonrpc":"2.0","id":3,"method":"shutdown"}
+---
+{"jsonrpc":"2.0","method":"exit"}

test/clangd/input-mirror.test

 # RUN: clangd -pretty -run-synchronously -input-mirror-file %t < %s
 # Note that we have to use '-b' as -input-mirror-file does not have a newline at the end of file.
 # RUN: diff -b %t %s
 # It is absolutely vital that this file has CRLF line endings.
 #
 Content-Length: 125
 
 {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"clangd","capabilities":{},"trace":"off"}}
 Content-Length: 172
 
 {"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{"uri":"test:///main.cpp","languageId":"cpp","version":1,"text":"int main() {\nint a;\na;\n}\n"}}}
 Content-Length: 44
 
 {"jsonrpc":"2.0","id":3,"method":"shutdown"}
+Content-Length: 33
+
+{"jsonrpc":"2.0","method":"exit"}

test/clangd/signature-help.test

 # CHECK-NEXT:           "label": "int"
 # CHECK-NEXT:         }
 # CHECK-NEXT:       ]
 # CHECK-NEXT:     }
 # CHECK-NEXT:   ]
 # CHECK-NEXT: }
 ---
 {"jsonrpc":"2.0","id":100000,"method":"shutdown"}
+---
+{"jsonrpc":"2.0","method":"exit"}

test/clangd/textdocument-didchange-fail.test

 # CHECK-NEXT:    "code": -32602,
 # CHECK-NEXT:    "message": "trying to get AST for non-added document"
 # CHECK-NEXT:  },
 # CHECK-NEXT:  "id": 1,
 # CHECK-NEXT:  "jsonrpc": "2.0"
 # CHECK-NEXT:}
 ---
 {"jsonrpc":"2.0","id":4,"method":"shutdown"}
+---
+{"jsonrpc":"2.0","method":"exit"}

test/clangd/trace.test

 # CHECK:     "File": "{{.*(/|\\)}}foo.c"
 # CHECK:   },
 # CHECK:   "name": "BuildAST",
 # CHECK:   "ph": "X",
 # CHECK: }
 # CHECK: },
 ---
 {"jsonrpc":"2.0","id":5,"method":"shutdown"}
+---
+{"jsonrpc":"2.0","method":"exit"}

test/clangd/xrefs.test

 # CHECK-NEXT:         "character": 8,
 # CHECK-NEXT:         "line": 1
 # CHECK-NEXT:       }
 # CHECK-NEXT:     }
 # CHECK-NEXT:   }
 # CHECK-NEXT: ]
 ---
 {"jsonrpc":"2.0","id":10000,"method":"shutdown"}
+---
+{"jsonrpc":"2.0","method":"exit"}