[clangd] Allow all LSP methods to signal cancellation via $/cancelRequest

Summary:
The cancelable scopes are managed by JSONRPCDispatcher so that all Handlers
run in cancelable contexts.
(Previously ClangdServer did this, for code completion only).

Cancellation request processing is therefore also in JSONRPCDispatcher.
(Previously it was in ClangdLSPServer).

This doesn't actually make any new commands *respect* cancellation - they'd
need to check isCancelled() and bail out. But it opens the door to doing
this incrementally, and putting such logic in common machinery like TUScheduler.

I also rewrote the ClangdServer class/threading comments because I wanted to
add to it and I got carried away.

Reviewers: ilya-biryukov, kadircet

Subscribers: ioeric, MaskRay, jkorous, arphaman, cfe-commits

Differential Revision: https://reviews.llvm.org/D52004

10 files changed, 103 insertions, 154 deletions

diff --git a/clang-tools-extra/clangd/ClangdLSPServer.cpp b/clang-tools-extra/clangd/ClangdLSPServer.cpp
index c01e5419c34..c68d422d319 100644
--- a/clang-tools-extra/clangd/ClangdLSPServer.cpp
+++ b/clang-tools-extra/clangd/ClangdLSPServer.cpp
@@ -8,7 +8,6 @@
 //===----------------------------------------------------------------------===//
 
 #include "ClangdLSPServer.h"
-#include "Cancellation.h"
 #include "Diagnostics.h"
 #include "JSONRPCDispatcher.h"
 #include "SourceCode.h"
@@ -71,11 +70,6 @@ SymbolKindBitset defaultSymbolKinds() {
   return Defaults;
 }
 
-std::string NormalizeRequestID(const json::Value &ID) {
-  auto NormalizedID = parseNumberOrString(&ID);
-  assert(NormalizedID && "Was not able to parse request id.");
-  return std::move(*NormalizedID);
-}
 } // namespace
 
 void ClangdLSPServer::onInitialize(InitializeParams &Params) {
@@ -347,21 +341,16 @@ void ClangdLSPServer::onCodeAction(CodeActionParams &Params) {
 }
 
 void ClangdLSPServer::onCompletion(TextDocumentPositionParams &Params) {
-  CreateSpaceForTaskHandle();
-  Canceler Cancel = Server.codeComplete(
-      Params.textDocument.uri.file(), Params.position, CCOpts,
-      [this](llvm::Expected<CodeCompleteResult> List) {
-        auto _ = llvm::make_scope_exit([this]() { CleanupTaskHandle(); });
-
-        if (!List)
-          return replyError(List.takeError());
-        CompletionList LSPList;
-        LSPList.isIncomplete = List->HasMore;
-        for (const auto &R : List->Completions)
-          LSPList.items.push_back(R.render(CCOpts));
-        return reply(std::move(LSPList));
-      });
-  StoreTaskHandle(std::move(Cancel));
+  Server.codeComplete(Params.textDocument.uri.file(), Params.position, CCOpts,
+                      [this](llvm::Expected<CodeCompleteResult> List) {
+                        if (!List)
+                          return replyError(List.takeError());
+                        CompletionList LSPList;
+                        LSPList.isIncomplete = List->HasMore;
+                        for (const auto &R : List->Completions)
+                          LSPList.items.push_back(R.render(CCOpts));
+                        return reply(std::move(LSPList));
+                      });
 }
 
 void ClangdLSPServer::onSignatureHelp(TextDocumentPositionParams &Params) {
@@ -629,48 +618,3 @@ GlobalCompilationDatabase &ClangdLSPServer::CompilationDB::getCDB() {
     return *CachingCDB;
   return *CDB;
 }
-
-void ClangdLSPServer::onCancelRequest(CancelParams &Params) {
-  std::lock_guard<std::mutex> Lock(TaskHandlesMutex);
-  const auto &It = TaskHandles.find(Params.ID);
-  if (It == TaskHandles.end())
-    return;
-  It->second();
-  TaskHandles.erase(It);
-}
-
-void ClangdLSPServer::CleanupTaskHandle() {
-  const json::Value *ID = getRequestId();
-  if (!ID)
-    return;
-  std::string NormalizedID = NormalizeRequestID(*ID);
-  std::lock_guard<std::mutex> Lock(TaskHandlesMutex);
-  TaskHandles.erase(NormalizedID);
-}
-
-void ClangdLSPServer::CreateSpaceForTaskHandle() {
-  const json::Value *ID = getRequestId();
-  if (!ID)
-    return;
-  std::string NormalizedID = NormalizeRequestID(*ID);
-  std::lock_guard<std::mutex> Lock(TaskHandlesMutex);
-  if (!TaskHandles.insert({NormalizedID, nullptr}).second)
-    elog("Creation of space for task handle: {0} failed.", NormalizedID);
-}
-
-void ClangdLSPServer::StoreTaskHandle(Canceler TH) {
-  const json::Value *ID = getRequestId();
-  if (!ID)
-    return;
-  std::string NormalizedID = NormalizeRequestID(*ID);
-  std::lock_guard<std::mutex> Lock(TaskHandlesMutex);
-  auto It = TaskHandles.find(NormalizedID);
-  if (It == TaskHandles.end()) {
-    elog("CleanupTaskHandle called before store can happen for request:{0}.",
-         NormalizedID);
-    return;
-  }
-  if (It->second != nullptr)
-    elog("TaskHandle didn't get cleared for: {0}.", NormalizedID);
-  It->second = std::move(TH);
-}
diff --git a/clang-tools-extra/clangd/ClangdLSPServer.h b/clang-tools-extra/clangd/ClangdLSPServer.h
index 4f444eae9a1..62260e2221b 100644
--- a/clang-tools-extra/clangd/ClangdLSPServer.h
+++ b/clang-tools-extra/clangd/ClangdLSPServer.h
@@ -26,8 +26,11 @@ namespace clangd {
 class JSONOutput;
 class SymbolIndex;
 
-/// This class provides implementation of an LSP server, glueing the JSON
-/// dispatch and ClangdServer together.
+/// This class exposes ClangdServer's capabilities via Language Server Protocol.
+///
+/// JSONRPCDispatcher binds the implemented ProtocolCallbacks methods
+/// (e.g. onInitialize) to corresponding JSON-RPC methods ("initialize").
+/// The server also supports $/cancelRequest (JSONRPCDispatcher provides this).
 class ClangdLSPServer : private DiagnosticsConsumer, private ProtocolCallbacks {
 public:
   /// If \p CompileCommandsDir has a value, compile_commands.json will be
@@ -76,7 +79,6 @@ private:
   void onRename(RenameParams &Parames) override;
   void onHover(TextDocumentPositionParams &Params) override;
   void onChangeConfiguration(DidChangeConfigurationParams &Params) override;
-  void onCancelRequest(CancelParams &Params) override;
 
   std::vector<Fix> getFixes(StringRef File, const clangd::Diagnostic &D);
 
@@ -169,21 +171,6 @@ private:
   // the worker thread that may otherwise run an async callback on partially
   // destructed instance of ClangdLSPServer.
   ClangdServer Server;
-
-  // Holds cancelers for running requets. Key of the map is a serialized
-  // request id. FIXME: handle cancellation generically in JSONRPCDispatcher.
-  llvm::StringMap<Canceler> TaskHandles;
-  std::mutex TaskHandlesMutex;
-
-  // Following three functions are for managing TaskHandles map. They store or
-  // remove a task handle for the request-id stored in current Context.
-  // FIXME(kadircet): Wrap the following three functions in a RAII object to
-  // make sure these do not get misused. The object might be stored in the
-  // Context of the thread or moved around until a reply is generated for the
-  // request.
-  void CleanupTaskHandle();
-  void CreateSpaceForTaskHandle();
-  void StoreTaskHandle(Canceler TH);
 };
 } // namespace clangd
 } // namespace clang
diff --git a/clang-tools-extra/clangd/ClangdServer.cpp b/clang-tools-extra/clangd/ClangdServer.cpp
index 20af72fb451..fb4f1be527e 100644
--- a/clang-tools-extra/clangd/ClangdServer.cpp
+++ b/clang-tools-extra/clangd/ClangdServer.cpp
@@ -8,7 +8,6 @@
 //===-------------------------------------------------------------------===//
 
 #include "ClangdServer.h"
-#include "Cancellation.h"
 #include "CodeComplete.h"
 #include "FindSymbols.h"
 #include "Headers.h"
@@ -201,16 +200,14 @@ void ClangdServer::removeDocument(PathRef File) {
   WorkScheduler.remove(File);
 }
 
-Canceler ClangdServer::codeComplete(PathRef File, Position Pos,
-                                    const clangd::CodeCompleteOptions &Opts,
-                                    Callback<CodeCompleteResult> CB) {
+void ClangdServer::codeComplete(PathRef File, Position Pos,
+                                const clangd::CodeCompleteOptions &Opts,
+                                Callback<CodeCompleteResult> CB) {
   // Copy completion options for passing them to async task handler.
   auto CodeCompleteOpts = Opts;
   if (!CodeCompleteOpts.Index) // Respect overridden index.
     CodeCompleteOpts.Index = Index;
 
-  auto Cancelable = cancelableTask();
-  WithContext ContextWithCancellation(std::move(Cancelable.first));
   // Copy PCHs to avoid accessing this->PCHs concurrently
   std::shared_ptr<PCHContainerOperations> PCHs = this->PCHs;
   auto FS = FSProvider.getFileSystem();
@@ -259,7 +256,6 @@ Canceler ClangdServer::codeComplete(PathRef File, Position Pos,
   // We use a potentially-stale preamble because latency is critical here.
   WorkScheduler.runWithPreamble("CodeComplete", File, TUScheduler::Stale,
                                 Bind(Task, File.str(), std::move(CB)));
-  return std::move(Cancelable.second);
 }
 
 void ClangdServer::signatureHelp(PathRef File, Position Pos,
diff --git a/clang-tools-extra/clangd/ClangdServer.h b/clang-tools-extra/clangd/ClangdServer.h
index 41b99b533a7..afbda4bd84e 100644
--- a/clang-tools-extra/clangd/ClangdServer.h
+++ b/clang-tools-extra/clangd/ClangdServer.h
@@ -45,15 +45,25 @@ public:
                                   std::vector<Diag> Diagnostics) = 0;
 };
 
-/// Provides API to manage ASTs for a collection of C++ files and request
-/// various language features.
-/// Currently supports async diagnostics, code completion, formatting and goto
-/// definition.
+/// Manages a collection of source files and derived data (ASTs, indexes),
+/// and provides language-aware features such as code completion.
+///
+/// The primary client is ClangdLSPServer which exposes these features via
+/// the Language Server protocol. ClangdServer may also be embedded directly,
+/// though its API is not stable over time.
+///
+/// ClangdServer should be used from a single thread. Many potentially-slow
+/// operations have asynchronous APIs and deliver their results on another
+/// thread.
+/// Such operations support cancellation: if the caller sets up a cancelable
+/// context, many operations will notice cancellation and fail early.
+/// (ClangdLSPServer uses this to implement $/cancelRequest).
 class ClangdServer {
 public:
   struct Options {
     /// To process requests asynchronously, ClangdServer spawns worker threads.
-    /// If 0, all requests are processed on the calling thread.
+    /// If this is zero, no threads are spawned. All work is done on the calling
+    /// thread, and callbacks are invoked before "async" functions return.
     unsigned AsyncThreadsCount = getDefaultAsyncThreadsCount();
 
     /// AST caching policy. The default is to keep up to 3 ASTs in memory.
@@ -125,9 +135,9 @@ public:
   /// while returned future is not yet ready.
   /// A version of `codeComplete` that runs \p Callback on the processing thread
   /// when codeComplete results become available.
-  Canceler codeComplete(PathRef File, Position Pos,
-                        const clangd::CodeCompleteOptions &Opts,
-                        Callback<CodeCompleteResult> CB);
+  void codeComplete(PathRef File, Position Pos,
+                    const clangd::CodeCompleteOptions &Opts,
+                    Callback<CodeCompleteResult> CB);
 
   /// Provide signature help for \p File at \p Pos.  This method should only be
   /// called for tracked files.
diff --git a/clang-tools-extra/clangd/JSONRPCDispatcher.cpp b/clang-tools-extra/clangd/JSONRPCDispatcher.cpp
index a3dd9ac4a49..4fe4f55968e 100644
--- a/clang-tools-extra/clangd/JSONRPCDispatcher.cpp
+++ b/clang-tools-extra/clangd/JSONRPCDispatcher.cpp
@@ -11,12 +11,14 @@
 #include "Cancellation.h"
 #include "ProtocolHandlers.h"
 #include "Trace.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>
 
@@ -158,6 +160,16 @@ void clangd::call(StringRef Method, json::Value &&Params) {
       });
 }
 
+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);
+    log("Bad cancellation request: {0}", Params);
+  });
+}
+
 void JSONRPCDispatcher::registerHandler(StringRef Method, Handler H) {
   assert(!Handlers.count(Method) && "Handler already registered!");
   Handlers[Method] = std::move(H);
@@ -180,8 +192,7 @@ static void logIncomingMessage(const llvm::Optional<json::Value> &ID,
   }
 }
 
-bool JSONRPCDispatcher::call(const json::Value &Message,
-                             JSONOutput &Out) const {
+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"))
@@ -214,12 +225,47 @@ bool JSONRPCDispatcher::call(const json::Value &Message,
     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));
+
   // Stash a reference to the span args, so later calls can add metadata.
   WithContext WithRequestSpan(RequestSpan::stash(Tracer));
   Handler(std::move(Params));
   return true;
 }
 
+// We run cancelable requests in a context that does two things:
+//  - allows cancellation using RequestCancelers[ID]
+//  - cleans up the entry in RequestCancelers when it's no longer needed
+// If a client reuses an ID, the last one wins and the first cannot be canceled.
+Context JSONRPCDispatcher::cancelableRequestContext(const json::Value &ID) {
+  auto Task = cancelableTask();
+  auto StrID = llvm::to_string(ID);  // JSON-serialize ID for map key.
+  auto Cookie = NextRequestCookie++; // No lock, only called on main thread.
+  {
+    std::lock_guard<std::mutex> Lock(RequestCancelersMutex);
+    RequestCancelers[StrID] = {std::move(Task.second), Cookie};
+  }
+  // When the request ends, we can clean up the entry we just added.
+  // The cookie lets us check that it hasn't been overwritten due to ID reuse.
+  return Task.first.derive(make_scope_exit([this, StrID, Cookie] {
+    std::lock_guard<std::mutex> Lock(RequestCancelersMutex);
+    auto It = RequestCancelers.find(StrID);
+    if (It != RequestCancelers.end() && It->second.second == Cookie)
+      RequestCancelers.erase(It);
+  }));
+}
+
+void JSONRPCDispatcher::cancelRequest(const json::Value &ID) {
+  auto StrID = llvm::to_string(ID);
+  std::lock_guard<std::mutex> Lock(RequestCancelersMutex);
+  auto It = RequestCancelers.find(StrID);
+  if (It != RequestCancelers.end())
+    It->second.first(); // Invoke the canceler.
+}
+
 // 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) {
diff --git a/clang-tools-extra/clangd/JSONRPCDispatcher.h b/clang-tools-extra/clangd/JSONRPCDispatcher.h
index fda37f86c66..b223ceb5f45 100644
--- a/clang-tools-extra/clangd/JSONRPCDispatcher.h
+++ b/clang-tools-extra/clangd/JSONRPCDispatcher.h
@@ -10,6 +10,7 @@
 #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"
@@ -77,23 +78,37 @@ 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 {
 public:
-  // A handler responds to requests for a particular method name.
+  /// A handler responds to requests for a particular method name.
+  ///
+  /// 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 &)>;
 
   /// Create a new JSONRPCDispatcher. UnknownHandler is called when an unknown
   /// method is received.
-  JSONRPCDispatcher(Handler UnknownHandler)
-      : UnknownHandler(std::move(UnknownHandler)) {}
+  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) const;
+  bool call(const llvm::json::Value &Message, JSONOutput &Out);
 
 private:
+  // 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;
 };
diff --git a/clang-tools-extra/clangd/Protocol.cpp b/clang-tools-extra/clangd/Protocol.cpp
index 9b232af176f..e6ffe66a129 100644
--- a/clang-tools-extra/clangd/Protocol.cpp
+++ b/clang-tools-extra/clangd/Protocol.cpp
@@ -623,38 +623,5 @@ bool fromJSON(const json::Value &Params, ReferenceParams &R) {
   return fromJSON(Params, Base);
 }
 
-json::Value toJSON(const CancelParams &CP) {
-  return json::Object{{"id", CP.ID}};
-}
-
-llvm::raw_ostream &operator<<(llvm::raw_ostream &O, const CancelParams &CP) {
-  O << toJSON(CP);
-  return O;
-}
-
-llvm::Optional<std::string> parseNumberOrString(const json::Value *Params) {
-  if (!Params)
-    return llvm::None;
-  // ID is either a number or a string, check for both.
-  if(const auto AsString = Params->getAsString())
-    return AsString->str();
-
-  if(const auto AsNumber = Params->getAsInteger())
-    return itostr(AsNumber.getValue());
-
-  return llvm::None;
-}
-
-bool fromJSON(const json::Value &Params, CancelParams &CP) {
-  const auto ParamsAsObject = Params.getAsObject();
-  if (!ParamsAsObject)
-    return false;
-  if (auto Parsed = parseNumberOrString(ParamsAsObject->get("id"))) {
-    CP.ID = std::move(*Parsed);
-    return true;
-  }
-  return false;
-}
-
 } // namespace clangd
 } // namespace clang
diff --git a/clang-tools-extra/clangd/Protocol.h b/clang-tools-extra/clangd/Protocol.h
index 9eec9c0f5c7..f09a63c95aa 100644
--- a/clang-tools-extra/clangd/Protocol.h
+++ b/clang-tools-extra/clangd/Protocol.h
@@ -886,20 +886,6 @@ struct ReferenceParams : public TextDocumentPositionParams {
 };
 bool fromJSON(const llvm::json::Value &, ReferenceParams &);
 
-struct CancelParams {
-  /// The request id to cancel.
-  /// This can be either a number or string, if it is a number simply print it
-  /// out and always use a string.
-  std::string ID;
-};
-llvm::json::Value toJSON(const CancelParams &);
-llvm::raw_ostream &operator<<(llvm::raw_ostream &, const CancelParams &);
-bool fromJSON(const llvm::json::Value &, CancelParams &);
-
-/// Param can be either of type string or number. Returns the result as a
-/// string.
-llvm::Optional<std::string> parseNumberOrString(const llvm::json::Value *Param);
-
 } // namespace clangd
 } // namespace clang
 
diff --git a/clang-tools-extra/clangd/ProtocolHandlers.cpp b/clang-tools-extra/clangd/ProtocolHandlers.cpp
index 9bb21e3e3d5..73da45eca67 100644
--- a/clang-tools-extra/clangd/ProtocolHandlers.cpp
+++ b/clang-tools-extra/clangd/ProtocolHandlers.cpp
@@ -76,5 +76,4 @@ void clangd::registerCallbackHandlers(JSONRPCDispatcher &Dispatcher,
   Register("workspace/didChangeConfiguration",
            &ProtocolCallbacks::onChangeConfiguration);
   Register("workspace/symbol", &ProtocolCallbacks::onWorkspaceSymbol);
-  Register("$/cancelRequest", &ProtocolCallbacks::onCancelRequest);
 }
diff --git a/clang-tools-extra/clangd/ProtocolHandlers.h b/clang-tools-extra/clangd/ProtocolHandlers.h
index 72ecd2b268d..983b7e2f423 100644
--- a/clang-tools-extra/clangd/ProtocolHandlers.h
+++ b/clang-tools-extra/clangd/ProtocolHandlers.h
@@ -56,7 +56,6 @@ public:
   virtual void onDocumentHighlight(TextDocumentPositionParams &Params) = 0;
   virtual void onHover(TextDocumentPositionParams &Params) = 0;
   virtual void onChangeConfiguration(DidChangeConfigurationParams &Params) = 0;
-  virtual void onCancelRequest(CancelParams &Params) = 0;
 };
 
 void registerCallbackHandlers(JSONRPCDispatcher &Dispatcher,