[clangd] Simplify cancellation public API

[sammccall]

Summary

Task is no longer exposed:

• task cancellation is hidden as a std::function
• task creation returns the new context directly
• checking is via free function only, with no way to avoid the context lookup

The implementation is essentially the same, but a bit terser as it's hidden.

isCancelled() is now safe to use outside any task (it returns false).
This will leave us free to sprinkle cancellation in e.g. TUScheduler without
needing elaborate test setup, and lets callers that don't cancel "just work".

Updated the docs to describe the new expected use pattern.
One thing I noticed: there's nothing async-specific about the cancellation.
Async tasks can be cancelled from any thread (typically the one that created
them), sync tasks can be cancelled from any *other* thread in the same way.
So the docs now refer to "long-running" tasks instead of async ones.

Updated usage in code complete, without any structural changes.
I didn't update all the names of the helpers in ClangdLSPServer (these will
likely be moved to JSONRPCDispatcher anyway).

[sammccall]

Fix includes, formatting tweak.

[kadircet]

Seems a lot cleaner now, thanks!

Do you plan to have other changes like moving control to JSONRPCDispatcher and recording timings for analysis on this patch? If not maybe we can add some fixme's so that we won't forget. Also the somewhat "caching" of cancellation token from previous implementation might still be useful in future if we really face "crowded" contexts and frequent cancellation checks, so maybe keep some notes about it?

clangd/Cancellation.h

21

s/TaskCancelledError/CancelledError

21

Maybe also mention propagating context into long-running work. runAsync does that implicitly and not sure if there will be other use cases that doesn't include it, but if there might be it would be nice to point it out as well.

[sammccall]

Address comments by adding comments!

[sammccall]

In D51996#1232459, @kadircet wrote:

Seems a lot cleaner now, thanks!

Do you plan to have other changes like moving control to JSONRPCDispatcher and recording timings for analysis on this patch?

D52004 is the JSONRPCDispatcher change.
Added a FIXME to the file comment about timings, hope to get to it soon.

If not maybe we can add some fixme's so that we won't forget. Also the somewhat "caching" of cancellation token from previous implementation might still be useful in future if we really face "crowded" contexts and frequent cancellation checks, so maybe keep some notes about it?

I think the best plan for frequent cancellation checks is "don't do that" ;-) The situation is that you're checking for cancellation to save CPU, but the thing you're spending the CPU on is checking for cancellation... I don't think the fix is to make cancellation checking cheaper! (Obviously cheaper is always better, but it would make the API worse here)
So i left a note on isCancelled(). We can revisit if we find cases where context traversal is really hurting us.

clangd/Cancellation.h

21

Yeah, added this to the caveats.
This shouldn't be a problem in practice (at least for pure open-source clangd): LLVM doesn't have a lot of multithreading utilities so we mostly define our own, and those are context aware.

[ilya-biryukov]

but the thing you're spending the CPU on is checking for cancellation...

Unless checking for cancellation is really cheap (which is doable, I think).
We should probably hit some of those in practice before doing something, though.

clangd/Cancellation.h

71

Allowing fast checking for cancellation seem important.
E.g. if we want to cancel in the middle of parsing, we'll probably be faced with sema callbacks, frequency of which we don't control, so we'd better design something that's fast to check.

I see two ways to deal with this:

• Keep the current design, but cache last access to context key.
• Add an API to get something that can quickly check for cancellation, i.e. something that would hold a reference to resolved context key.

Both could be added later, though, and I don't have any data, it's just my intuition.

[ilya-biryukov]

Will let Kadir finish the review, consider lgtm'ed from me ;-)

[kadircet]

LGTM

[sammccall]

In D51996#1232770, @ilya-biryukov wrote:

but the thing you're spending the CPU on is checking for cancellation...

Unless checking for cancellation is really cheap (which is doable, I think).
We should probably hit some of those in practice before doing something, though.

My point is that if a few pointer traversals/comparisons is significant, then the loop is so fast that checking for cancellations doesn't make sense anyway - we're aiming to save 10s of ms, not a few hundred nanos.

Like you say, maybe there are important cases where we can't tell that we're calling frequently, or can't easily avoid that. We don't have any yet though, so optimizing this seems premature.

clangd/Cancellation.h

71

Yeah, we have a few options if we get there here.
Maybe the simplest is

if (LLVM_UNLIKELY(++Counter % 1000 == 0))
 if (isCancelled())
  return;

(obviously this is limited, but it depends what problem we're facing)

clangd/Cancellation.h

 //===--- 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;
-// }
+// Cancellation mechanism for long-running tasks.
+//
+// This manages interactions between:
 //
-// The callers of async methods (2) can issue cancellations and should be
-// prepared to handle `TaskCancelledError` result:
+// 1. Client code that starts some long-running work, and maybe cancels later.
 //
-// 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();
-// }
+//   std::pair<Context, Canceler> Task = cancelableTask();
+//   {
+//     WithContext Cancelable(std::move(Task.first));
+//     Expected
+//     deepThoughtAsync([](int answer){ errs() << answer; });
+//   }
+//   // ...some time later...
+//   if (User.fellAsleep())
+//     Task.second();
+//
+//  (This example has an asynchronous computation, but synchronous examples
+//  work similarly - the Canceler should be invoked from another thread).
+//
+// 2. Library code that executes long-running work, and can exit early if the
+//   result is not needed.
+//
+//   void deepThoughtAsync(std::function<void(int)> Callback) {
+//     runAsync([Callback]{
+//       int A = ponder(6);
+//       if (isCancelled())
+//         return;
+//       int B = ponder(9);
+//       if (isCancelled())
+//         return;
+//       Callback(A * B);
+//     });
+//   }
 //
-// The worker code itself (3) should check for cancellations using
-// `Task::isCancelled` that can be retrieved via `getCurrentTask()`.
+//   (A real example may invoke the callback with an error on cancellation,
+//   the CancelledError is provided for this purpose).
 //
-// 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.
+// Cancellation has some caveats:
+//   - the work will only stop when/if the library code next checks for it.
+//     Code outside clangd such as Sema will not do this.
+//   - it's inherently racy: client code must be prepared to accept results
+//     even after requesting cancellation.
+//   - it's Context-based, so async work must be dispatched to threads in
+//     ways that preserve the context. (Like runAsync() or TUScheduler).
+//
+// FIXME: We could add timestamps to isCancelled() and CancelledError.
+//        Measuring the start -> cancel -> acknowledge -> finish timeline would
+//        help find where libraries' cancellation should be improved.
 
 #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 <functional>
 #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 async 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(); }
+/// A canceller requests cancellation of a task, when called.
+/// Calling it again has no effect.
+using Canceler = std::function<void()>;
+
+/// Defines a new task whose cancellation may be requested.
+/// The returned Context defines the scope of the task.
+/// When the context is active, isCancelled() is false until the Canceler is
+/// invoked, and true afterwards.
+std::pair<Context, Canceler> cancelableTask();
+
+/// True if the current context is within a cancelable task which was cancelled.
+/// Always false if there is no active cancelable task.
+/// This isn't free (context lookup) - don't call it in a tight loop.
+bool isCancelled();
 
+/// Conventional error when no result is returned due to cancellation.
 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

clangd/Cancellation.cpp

 //===--- 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;
+static Key<std::shared_ptr<std::atomic<bool>>> FlagKey;
 
-const Task &getCurrentTask() {
-  const auto TH = Context::current().getExisting(TaskKey);
-  assert(TH && "Fetched a nullptr for TaskHandle from context.");
-  return *TH;
+std::pair<Context, Canceler> cancelableTask() {
+  auto Flag = std::make_shared<std::atomic<bool>>();
+  return {
+      Context::current().derive(FlagKey, Flag),
+      [Flag] { *Flag = true; },
+  };
 }
 
-Context setCurrentTask(ConstTaskHandle TH) {
-  assert(TH && "Trying to stash a nullptr as TaskHandle into context.");
-  return Context::current().derive(TaskKey, std::move(TH));
+bool isCancelled() {
+  if (auto *Flag = Context::current().get(FlagKey))
+    return **Flag;
+  return false; // Not in scope of a task.
 }
 
 } // namespace clangd
 } // namespace clang

clangd/ClangdLSPServer.h

private:
   // 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.
   ClangdServer Server;
 
-  // Holds task handles for running requets. Key of the map is a serialized
-  // request id.
-  llvm::StringMap<TaskHandle> TaskHandles;
+  // 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(TaskHandle TH);
+  void StoreTaskHandle(Canceler TH);
 };
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDLSPSERVER_H

clangd/ClangdLSPServer.cpp

for (auto &F : getFixes(Params.textDocument.uri.file(), D)) {
       });
     }
   }
   reply(std::move(Commands));
 }
 
 void ClangdLSPServer::onCompletion(TextDocumentPositionParams &Params) {
   CreateSpaceForTaskHandle();
-  TaskHandle TH = Server.codeComplete(
+  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(TH));
+  StoreTaskHandle(std::move(Cancel));
 }
 
 void ClangdLSPServer::onSignatureHelp(TextDocumentPositionParams &Params) {
   Server.signatureHelp(Params.textDocument.uri.file(), Params.position,
                        [](llvm::Expected<SignatureHelp> SignatureHelp) {
                          if (!SignatureHelp)
                            return replyError(
                                ErrorCode::InvalidParams,
GlobalCompilationDatabase &ClangdLSPServer::CompilationDB::getCDB() {
   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();
+  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(TaskHandle TH) {
+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);
 }

clangd/ClangdServer.h

public:
   /// Run code completion for \p File at \p Pos.
   /// Request is processed asynchronously.
   ///
   /// This method should only be called for currently tracked files. However, it
   /// is safe to call removeDocument for \p File after this method returns, even
   /// while returned future is not yet ready.
   /// A version of `codeComplete` that runs \p Callback on the processing thread
   /// when codeComplete results become available.
-  TaskHandle codeComplete(PathRef File, Position Pos,
-                          const clangd::CodeCompleteOptions &Opts,
-                          Callback<CodeCompleteResult> CB);
+  Canceler codeComplete(PathRef File, Position Pos,
 
   /// Provide signature help for \p File at \p Pos.  This method should only be
   /// called for tracked files.
   void signatureHelp(PathRef File, Position Pos, Callback<SignatureHelp> CB);
 
   /// Get definition of symbol at a specified \p Line and \p Column in \p File.
   void findDefinitions(PathRef File, Position Pos,
                        Callback<std::vector<Location>> CB);

clangd/ClangdServer.cpp

WorkScheduler.update(File, std::move(Inputs), WantDiags,
                        });
 }
 
 void ClangdServer::removeDocument(PathRef File) {
   ++InternalVersion[File];
   WorkScheduler.remove(File);
 }
 
-TaskHandle ClangdServer::codeComplete(PathRef File, Position Pos,
-                                      const clangd::CodeCompleteOptions &Opts,
-                                      Callback<CodeCompleteResult> CB) {
+Canceler ClangdServer::codeComplete(PathRef File, Position Pos,
   // 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));
+  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();
 
   auto Task = [PCHs, Pos, FS, CodeCompleteOpts,
                this](Path File, Callback<CodeCompleteResult> CB,
                      llvm::Expected<InputsAndPreamble> IP) {
     if (!IP)
auto Task = [PCHs, Pos, FS, CodeCompleteOpts,
     // We don't want `codeComplete` to wait for the async call if it doesn't use
     // the result (e.g. non-index completion, speculation fails), so that `CB`
     // is called as soon as results are available.
   };
 
   // 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 TH;
+  return std::move(Cancelable.second);
 }
 
 void ClangdServer::signatureHelp(PathRef File, Position Pos,
                                  Callback<SignatureHelp> CB) {
 
   auto PCHs = this->PCHs;
   auto FS = FSProvider.getFileSystem();
   auto *Index = this->Index;

unittests/clangd/CancellationTests.cpp

 #include "Cancellation.h"
 #include "Context.h"
 #include "Threading.h"
 #include "llvm/Support/Error.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <atomic>
 #include <memory>
 #include <thread>
 
 namespace clang {
 namespace clangd {
 namespace {
 
 TEST(CancellationTest, CancellationTest) {
-  TaskHandle TH = Task::createHandle();
-  WithContext ContextWithCancellation(setCurrentTask(TH));
+  auto Task = cancelableTask();
+  WithContext ContextWithCancellation(std::move(Task.first));
   EXPECT_FALSE(isCancelled());
-  TH->cancel();
+  Task.second();
   EXPECT_TRUE(isCancelled());
 }
 
-TEST(CancellationTest, TaskTestHandleDiesContextLives) {
+TEST(CancellationTest, CancelerDiesContextLives) {
   llvm::Optional<WithContext> ContextWithCancellation;
   {
-    TaskHandle TH = Task::createHandle();
-    ContextWithCancellation.emplace(setCurrentTask(TH));
+    auto Task = cancelableTask();
+    ContextWithCancellation.emplace(std::move(Task.first));
     EXPECT_FALSE(isCancelled());
-    TH->cancel();
+    Task.second();
     EXPECT_TRUE(isCancelled());
   }
   EXPECT_TRUE(isCancelled());
 }
 
 TEST(CancellationTest, TaskContextDiesHandleLives) {
-  TaskHandle TH = Task::createHandle();
+  auto Task = cancelableTask();
   {
-    WithContext ContextWithCancellation(setCurrentTask(TH));
+    WithContext ContextWithCancellation(std::move(Task.first));
     EXPECT_FALSE(isCancelled());
-    TH->cancel();
+    Task.second();
     EXPECT_TRUE(isCancelled());
   }
   // Still should be able to cancel without any problems.
-  TH->cancel();
-}
-
-TEST(CancellationTest, CancellationToken) {
-  TaskHandle TH = Task::createHandle();
-  WithContext ContextWithCancellation(setCurrentTask(TH));
-  const auto &CT = getCurrentTask();
-  EXPECT_FALSE(CT.isCancelled());
-  TH->cancel();
-  EXPECT_TRUE(CT.isCancelled());
+  Task.second();
 }
 
 TEST(CancellationTest, AsynCancellationTest) {
   std::atomic<bool> HasCancelled(false);
   Notification Cancelled;
-  auto TaskToBeCancelled = [&](ConstTaskHandle CT) {
-    WithContext ContextGuard(setCurrentTask(std::move(CT)));
+  auto TaskToBeCancelled = [&](Context Ctx) {
+    WithContext ContextGuard(std::move(Ctx));
     Cancelled.wait();
     HasCancelled = isCancelled();
   };
-  TaskHandle TH = Task::createHandle();
-  std::thread AsyncTask(TaskToBeCancelled, TH);
-  TH->cancel();
+  auto Task = cancelableTask();
+  std::thread AsyncTask(TaskToBeCancelled, std::move(Task.first));
+  Task.second();
   Cancelled.notify();
   AsyncTask.join();
 
   EXPECT_TRUE(HasCancelled);
 }
 } // namespace
 } // namespace clangd
 } // namespace clang