[clangd] Initial cancellation mechanism for LSP requests.

Reviewers: ilya-biryukov, ioeric, hokein

Reviewed By: ilya-biryukov

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

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

llvm-svn: 340607

clang-tools-extra/clangd/CMakeLists.txt

@@ -9,6 +9,7 @@ endif()
 
 add_clang_library(clangDaemon
   AST.cpp
+  Cancellation.cpp
   ClangdLSPServer.cpp
   ClangdServer.cpp
   ClangdUnit.cpp

clang-tools-extra/clangd/Cancellation.cpp

@@ -0,0 +1,34 @@
+//===--- Cancellation.cpp -----------------------------------------*-C++-*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "Cancellation.h"
+#include <atomic>
+
+namespace clang {
+namespace clangd {
+
+namespace {
+static Key<ConstTaskHandle> TaskKey;
+} // namespace
+
+char CancelledError::ID = 0;
+
+const Task &getCurrentTask() {
+  const auto TH = Context::current().getExisting(TaskKey);
+  assert(TH && "Fetched a nullptr for TaskHandle from context.");
+  return *TH;
+}
+
+Context setCurrentTask(ConstTaskHandle TH) {
+  assert(TH && "Trying to stash a nullptr as TaskHandle into context.");
+  return Context::current().derive(TaskKey, std::move(TH));
+}
+
+} // namespace clangd
+} // namespace clang

clang-tools-extra/clangd/Cancellation.h

@@ -0,0 +1,142 @@
+//===--- Cancellation.h -------------------------------------------*-C++-*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+// Cancellation mechanism for async tasks. Roughly all the clients of this code
+// can be classified into three categories:
+// 1. The code that creates and schedules async tasks, e.g. TUScheduler.
+// 2. The callers of the async method that can cancel some of the running tasks,
+// e.g. `ClangdLSPServer`
+// 3. The code running inside the async task itself, i.e. code completion or
+// find definition implementation that run clang, etc.
+//
+// For (1), the guideline is to accept a callback for the result of async
+// operation and return a `TaskHandle` to allow cancelling the request.
+//
+//  TaskHandle someAsyncMethod(Runnable T,
+//  function<void(llvm::Expected<ResultType>)> Callback) {
+//   auto TH = Task::createHandle();
+//   WithContext ContextWithCancellationToken(TH);
+//   auto run = [](){
+//     Callback(T());
+//   }
+//   // Start run() in a new async thread, and make sure to propagate Context.
+//   return TH;
+// }
+//
+// The callers of async methods (2) can issue cancellations and should be
+// prepared to handle `TaskCancelledError` result:
+//
+// void Caller() {
+//   // You should store this handle if you wanna cancel the task later on.
+//   TaskHandle TH = someAsyncMethod(Task, [](llvm::Expected<ResultType> R) {
+//     if(/*check for task cancellation error*/)
+//       // Handle the error
+//     // Do other things on R.
+//   });
+//   // To cancel the task:
+//   sleep(5);
+//   TH->cancel();
+// }
+//
+// The worker code itself (3) should check for cancellations using
+// `Task::isCancelled` that can be retrieved via `getCurrentTask()`.
+//
+// llvm::Expected<ResultType> AsyncTask() {
+//    // You can either store the read only TaskHandle by calling getCurrentTask
+//    // once and just use the variable everytime you want to check for
+//    // cancellation, or call isCancelled everytime. The former is more
+//    // efficient if you are going to have multiple checks.
+//    const auto T = getCurrentTask();
+//    // DO SMTHNG...
+//    if(T.isCancelled()) {
+//      // Task has been cancelled, lets get out.
+//      return llvm::makeError<CancelledError>();
+//    }
+//    // DO SOME MORE THING...
+//    if(T.isCancelled()) {
+//      // Task has been cancelled, lets get out.
+//      return llvm::makeError<CancelledError>();
+//    }
+//    return ResultType(...);
+// }
+// If the operation was cancelled before task could run to completion, it should
+// propagate the TaskCancelledError as a result.
+
+#ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_CANCELLATION_H
+#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_CANCELLATION_H
+
+#include "Context.h"
+#include "llvm/Support/Error.h"
+#include <atomic>
+#include <memory>
+#include <system_error>
+
+namespace clang {
+namespace clangd {
+
+/// Enables signalling a cancellation on an async task or checking for
+/// cancellation. It is thread-safe to trigger cancellation from multiple
+/// threads or check for cancellation. Task object for the currently running
+/// task can be obtained via clangd::getCurrentTask().
+class Task {
+public:
+  void cancel() { CT = true; }
+  /// If cancellation checks are rare, one could use the isCancelled() helper in
+  /// the namespace to simplify the code. However, if cancellation checks are
+  /// frequent, the guideline is first obtain the Task object for the currently
+  /// running task with getCurrentTask() and do cancel checks using it to avoid
+  /// extra lookups in the Context.
+  bool isCancelled() const { return CT; }
+
+  /// Creates a task handle that can be used by an asyn task to check for
+  /// information that can change during it's runtime, like Cancellation.
+  static std::shared_ptr<Task> createHandle() {
+    return std::shared_ptr<Task>(new Task());
+  }
+
+  Task(const Task &) = delete;
+  Task &operator=(const Task &) = delete;
+  Task(Task &&) = delete;
+  Task &operator=(Task &&) = delete;
+
+private:
+  Task() : CT(false) {}
+  std::atomic<bool> CT;
+};
+using ConstTaskHandle = std::shared_ptr<const Task>;
+using TaskHandle = std::shared_ptr<Task>;
+
+/// Fetches current task information from Context. TaskHandle must have been
+/// stashed into context beforehand.
+const Task &getCurrentTask();
+
+/// Stashes current task information within the context.
+LLVM_NODISCARD Context setCurrentTask(ConstTaskHandle TH);
+
+/// Checks whether the current task has been cancelled or not.
+/// Consider storing the task handler returned by getCurrentTask and then
+/// calling isCancelled through it. getCurrentTask is expensive since it does a
+/// lookup in the context.
+inline bool isCancelled() { return getCurrentTask().isCancelled(); }
+
+class CancelledError : public llvm::ErrorInfo<CancelledError> {
+public:
+  static char ID;
+
+  void log(llvm::raw_ostream &OS) const override {
+    OS << "Task was cancelled.";
+  }
+  std::error_code convertToErrorCode() const override {
+    return std::make_error_code(std::errc::operation_canceled);
+  }
+};
+
+} // namespace clangd
+} // namespace clang
+
+#endif

clang-tools-extra/clangd/ClangdLSPServer.cpp

@@ -8,10 +8,12 @@
 //===----------------------------------------------------------------------===//
 
 #include "ClangdLSPServer.h"
+#include "Cancellation.h"
 #include "Diagnostics.h"
 #include "JSONRPCDispatcher.h"
 #include "SourceCode.h"
 #include "URI.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Path.h"
@@ -69,6 +71,11 @@ 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) {
@@ -339,17 +346,21 @@ void ClangdLSPServer::onCodeAction(CodeActionParams &Params) {
 }
 
 void ClangdLSPServer::onCompletion(TextDocumentPositionParams &Params) {
-  Server.codeComplete(Params.textDocument.uri.file(), Params.position, CCOpts,
-                      [this](llvm::Expected<CodeCompleteResult> List) {
-                        if (!List)
-                          return replyError(ErrorCode::InvalidParams,
-                                            llvm::toString(List.takeError()));
-                        CompletionList LSPList;
-                        LSPList.isIncomplete = List->HasMore;
-                        for (const auto &R : List->Completions)
-                          LSPList.items.push_back(R.render(CCOpts));
-                        reply(std::move(LSPList));
-                      });
+  CreateSpaceForTaskHandle();
+  TaskHandle TH = 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(TH));
 }
 
 void ClangdLSPServer::onSignatureHelp(TextDocumentPositionParams &Params) {
@@ -364,14 +375,14 @@ void ClangdLSPServer::onSignatureHelp(TextDocumentPositionParams &Params) {
 }
 
 void ClangdLSPServer::onGoToDefinition(TextDocumentPositionParams &Params) {
-  Server.findDefinitions(
-      Params.textDocument.uri.file(), Params.position,
-      [](llvm::Expected<std::vector<Location>> Items) {
-        if (!Items)
-          return replyError(ErrorCode::InvalidParams,
-                            llvm::toString(Items.takeError()));
-        reply(json::Array(*Items));
-      });
+  Server.findDefinitions(Params.textDocument.uri.file(), Params.position,
+                         [](llvm::Expected<std::vector<Location>> Items) {
+                           if (!Items)
+                             return replyError(
+                                 ErrorCode::InvalidParams,
+                                 llvm::toString(Items.takeError()));
+                           reply(json::Array(*Items));
+                         });
 }
 
 void ClangdLSPServer::onSwitchSourceHeader(TextDocumentIdentifier &Params) {
@@ -606,3 +617,49 @@ 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;
+  if (It->second)
+    It->second->cancel();
+  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(TaskHandle 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);
+}

clang-tools-extra/clangd/ClangdLSPServer.h

@@ -75,6 +75,7 @@ class ClangdLSPServer : private DiagnosticsConsumer, private ProtocolCallbacks {
   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);
 
@@ -167,8 +168,22 @@ class ClangdLSPServer : private DiagnosticsConsumer, private ProtocolCallbacks {
   // the worker thread that may otherwise run an async callback on partially
   // destructed instance of ClangdLSPServer.
   ClangdServer Server;
-};
 
+  // Holds task handles for running requets. Key of the map is a serialized
+  // request id.
+  llvm::StringMap<TaskHandle> 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(TaskHandle TH);
+};
 } // namespace clangd
 } // namespace clang
 

clang-tools-extra/clangd/ClangdServer.cpp

@@ -8,6 +8,7 @@
 //===-------------------------------------------------------------------===//
 
 #include "ClangdServer.h"
+#include "Cancellation.h"
 #include "CodeComplete.h"
 #include "FindSymbols.h"
 #include "Headers.h"
@@ -173,21 +174,26 @@ void ClangdServer::removeDocument(PathRef File) {
   WorkScheduler.remove(File);
 }
 
-void ClangdServer::codeComplete(PathRef File, Position Pos,
-                                const clangd::CodeCompleteOptions &Opts,
-                                Callback<CodeCompleteResult> CB) {
+TaskHandle 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;
 
+  TaskHandle TH = Task::createHandle();
+  WithContext ContextWithCancellation(setCurrentTask(TH));
   // Copy PCHs to avoid accessing this->PCHs concurrently
   std::shared_ptr<PCHContainerOperations> PCHs = this->PCHs;
   auto FS = FSProvider.getFileSystem();
 
   auto Task = [PCHs, Pos, FS, CodeCompleteOpts,
                this](Path File, Callback<CodeCompleteResult> CB,
                      llvm::Expected<InputsAndPreamble> IP) {
+    if (isCancelled())
+      return CB(llvm::make_error<CancelledError>());
+
     if (!IP)
       return CB(IP.takeError());
 
@@ -226,6 +232,7 @@ void ClangdServer::codeComplete(PathRef File, Position Pos,
 
   WorkScheduler.runWithPreamble("CodeComplete", File,
                                 Bind(Task, File.str(), std::move(CB)));
+  return TH;
 }
 
 void ClangdServer::signatureHelp(PathRef File, Position Pos,