[clangd] Enforce rules around "initialize" request, and create ClangdServer lazily.

[sammccall]

Summary

LSP is a slightly awkward map to C++ object lifetimes: the initialize request
is part of the protocol and provides information that doesn't change over the
lifetime of the server.

Until now, we handled this by initializing ClangdServer and ClangdLSPServer
right away, and making anything that can be set in the "initialize" request
mutable.
With this patch, we create ClangdLSPServer immediately, but defer creating
ClangdServer until "initialize". This opens the door to passing the relevant
initialize params in the constructor and storing them immutably.
(That change isn't actually done in this patch).

To make this safe, we have the MessageDispatcher enforce that the "initialize"
method is called before any other (as required by LSP). That way each method
handler can assume Server is initialized, as today.

As usual, while implementing this I found places where our test cases violated
the protocol.

[ioeric]

Looks good!

clangd/ClangdLSPServer.h

private:
   /// The supported completion item kinds of the client.
   CompletionItemKindBitset SupportedCompletionItemKinds;
   // Whether the client supports CodeAction response objects.
   bool SupportsCodeAction = false;
 
   // Store of the current versions of the open documents.
   DraftStore DraftMgr;
 
-  // Server must be the last member of the class to allow its destructor to exit
-  // the worker thread that may otherwise run an async callback on partially
-  // destructed instance of ClangdLSPServer.
-  // Set in construtor and destroyed when run() finishes. To ensure all worker
-  // threads exit before run() returns.
-  std::unique_ptr<ClangdServer> Server;
+  // The ClangdServer is created by the "initialize" LSP method.
+  // It is destroyed before run() returns, to ensure worker threads exit.
+  ClangdServer::Options ClangdServerOpts;
+  llvm::Optional<ClangdServer> Server;
 };
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDLSPSERVER_H

clangd/ClangdLSPServer.cpp

 } // namespace
 
 // MessageHandler dispatches incoming LSP messages.
 // It handles cross-cutting concerns:
 //  - serializes/deserializes protocol objects to JSON
 //  - logging of inbound messages
 //  - cancellation handling
 //  - basic call tracing
+// MessageHandler ensures that initialize() is called before any other handler.
 class ClangdLSPServer::MessageHandler : public Transport::MessageHandler {
 public:
   MessageHandler(ClangdLSPServer &Server) : Server(Server) {}
 
   bool onNotify(StringRef Method, json::Value Params) override {
     log("<-- {0}", Method);
     if (Method == "exit")
       return false;
-    if (Method == "$/cancelRequest")
+    if (!Server.Server)
+      elog("Notification {0} before initialization", Method);
+    else if (Method == "$/cancelRequest")
       onCancel(std::move(Params));
     else if (auto Handler = Notifications.lookup(Method))
       Handler(std::move(Params));
     else
       log("unhandled notification {0}", Method);
     return true;
   }
 
   bool onCall(StringRef Method, json::Value Params, json::Value ID) override {
     log("<-- {0}({1})", Method, ID);
-    if (auto Handler = Calls.lookup(Method))
+    if (!Server.Server && Method != "initialize") {
+      elog("Call {0} before initialization.", Method);
+      Server.reply(ID, make_error<LSPError>("server not initialized",
+                                            ErrorCode::ServerNotInitialized));
+    } else if (auto Handler = Calls.lookup(Method))
       Handler(std::move(Params), std::move(ID));
     else
       Server.reply(ID, llvm::make_error<LSPError>("method not found",
                                                   ErrorCode::MethodNotFound));
     return true;
   }
 
   bool onReply(json::Value ID, Expected<json::Value> Result) override {
if (Result) {
     log("--> reply({0}) error: {1}", ID, Err);
     std::lock_guard<std::mutex> Lock(TranspWriter);
     Transp.reply(std::move(ID), std::move(Err));
   }
 }
 
 void ClangdLSPServer::onInitialize(const InitializeParams &Params,
                                    Callback<json::Value> Reply) {
+  if (Server)
+    return Reply(make_error<LSPError>("server already initialized",
+                                      ErrorCode::InvalidRequest));
+  Server.emplace(CDB.getCDB(), FSProvider,
+                 static_cast<DiagnosticsConsumer &>(*this), ClangdServerOpts);
   if (Params.initializationOptions) {
     const ClangdInitializationOptions &Opts = *Params.initializationOptions;
 
     // Explicit compilation database path.
     if (Opts.compilationDatabasePath.hasValue()) {
       CDB.setCompileCommandsDir(Opts.compilationDatabasePath.getValue());
     }
 
ClangdLSPServer::ClangdLSPServer(class Transport &Transp,
                                  bool ShouldUseInMemoryCDB,
                                  const ClangdServer::Options &Opts)
     : Transp(Transp), MsgHandler(new MessageHandler(*this)),
       CDB(ShouldUseInMemoryCDB ? CompilationDB::makeInMemory()
                                : CompilationDB::makeDirectoryBased(
                                      std::move(CompileCommandsDir))),
       CCOpts(CCOpts), SupportedSymbolKinds(defaultSymbolKinds()),
       SupportedCompletionItemKinds(defaultCompletionItemKinds()),
-      Server(new ClangdServer(CDB.getCDB(), FSProvider, /*DiagConsumer=*/*this,
-                              Opts)) {
+      ClangdServerOpts(Opts) {
   // clang-format off
   MsgHandler->bind("initialize", &ClangdLSPServer::onInitialize);
   MsgHandler->bind("shutdown", &ClangdLSPServer::onShutdown);
   MsgHandler->bind("textDocument/rangeFormatting", &ClangdLSPServer::onDocumentRangeFormatting);
   MsgHandler->bind("textDocument/onTypeFormatting", &ClangdLSPServer::onDocumentOnTypeFormatting);
   MsgHandler->bind("textDocument/formatting", &ClangdLSPServer::onDocumentFormatting);
   MsgHandler->bind("textDocument/codeAction", &ClangdLSPServer::onCodeAction);
   MsgHandler->bind("textDocument/completion", &ClangdLSPServer::onCompletion);
ClangdLSPServer::ClangdLSPServer(class Transport &Transp,
   MsgHandler->bind("workspace/didChangeWatchedFiles", &ClangdLSPServer::onFileEvent);
   MsgHandler->bind("workspace/didChangeConfiguration", &ClangdLSPServer::onChangeConfiguration);
   // clang-format on
 }
 
 ClangdLSPServer::~ClangdLSPServer() = default;
 
 bool ClangdLSPServer::run() {
-  assert(Server);
-
   // Run the Language Server loop.
   bool CleanExit = true;
   if (auto Err = Transp.loop(*MsgHandler)) {
     elog("Transport error: {0}", std::move(Err));
     CleanExit = false;
   }
 
   // Destroy ClangdServer to ensure all worker threads finish.
   Server.reset();
-
   return CleanExit && ShutdownRequestReceived;
 }
 
 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())

clangd/ClangdServer.h

struct Options {
 
     /// If set, use this index to augment code completion results.
     SymbolIndex *StaticIndex = nullptr;
 
     /// The resource directory is used to find internal headers, overriding
     /// defaults and -resource-dir compiler flag).
     /// If None, ClangdServer calls CompilerInvocation::GetResourcePath() to
     /// obtain the standard resource directory.
-    llvm::Optional<StringRef> ResourceDir = llvm::None;
+    llvm::Optional<std::string> ResourceDir = llvm::None;
 
     /// Time to wait after a new file version before computing diagnostics.
     std::chrono::steady_clock::duration UpdateDebounce =
         std::chrono::milliseconds(500);
   };
   // Sensible default options for use in tests.
   // Features like indexing must be enabled if desired.
   static Options optsForTest();

clangd/ClangdServer.cpp

ClangdServer::Options ClangdServer::optsForTest() {
   return Opts;
 }
 
 ClangdServer::ClangdServer(const GlobalCompilationDatabase &CDB,
                            const FileSystemProvider &FSProvider,
                            DiagnosticsConsumer &DiagConsumer,
                            const Options &Opts)
     : CDB(CDB), DiagConsumer(DiagConsumer), FSProvider(FSProvider),
-      ResourceDir(Opts.ResourceDir ? Opts.ResourceDir->str()
+      ResourceDir(Opts.ResourceDir ? *Opts.ResourceDir
                                    : getStandardResourceDir()),
       DynamicIdx(Opts.BuildDynamicSymbolIndex
                      ? new FileIndex(Opts.URISchemes,
                                      Opts.HeavyweightDynamicSymbolIndex)
                      : nullptr),
       PCHs(std::make_shared<PCHContainerOperations>()),
       // Pass a callback into `WorkScheduler` to extract symbols from a newly
       // parsed file and rebuild the file index synchronously each time an AST

test/clangd/exit-with-shutdown.test

 # RUN: clangd -lit-test < %s
+{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"clangd","capabilities":{},"trace":"off"}}
+---
 {"jsonrpc":"2.0","id":3,"method":"shutdown"}
 ---
 {"jsonrpc":"2.0","method":"exit"}

test/clangd/exit-without-shutdown.test

 # RUN: not clangd -lit-test < %s
+{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"clangd","capabilities":{},"trace":"off"}}
+---
 {"jsonrpc":"2.0","method":"exit"}

test/clangd/initialize-sequence.test

+# RUN: clangd -lit-test < %s | FileCheck -strict-whitespace %s
+{"jsonrpc":"2.0","id":0,"method":"workspace/symbol","params":{"query":"std::basic_ostringstream"}}
+#      CHECK:  "error": {
+# CHECK-NEXT:    "code": -32002
+# CHECK-NEXT:    "message": "server not initialized"
+# CHECK-NEXT:  },
+# CHECK-NEXT:  "id": 0,
+---
+{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"processId":123,"rootPath":"clangd","capabilities":{},"trace":"off"}}
+---
+{"jsonrpc":"2.0","id":2,"method":"initialize","params":{"processId":123,"rootPath":"clangd","capabilities":{},"trace":"off"}}
+#      CHECK:  "error": {
+# CHECK-NEXT:    "code": -32600
+# CHECK-NEXT:    "message": "server already initialized"
+# CHECK-NEXT:  },
+# CHECK-NEXT:  "id": 2,
+---
+{"jsonrpc":"2.0","id":3,"method":"shutdown"}
+---
+{"jsonrpc":"2.0","method":"exit"}
+