[clangd] Initial cancellation mechanism for LSP requests.

[ilya-biryukov]

I have left a few comments, but wanted to start with a higher-level design consideration.

I believe we should probably expose the cancellations in the ClangdServer API directly.
The reasons for that are:

1. We have an internal client that would also benefit from using the cancellation API and we should try to share as much code as possible between that and the rest of clangd. Creating the cancellation cookies, stashing them in the context are all details that could be handled by ClangdServer so two clients don't have to do the same things.
2. ClangdServer seems like a natural layer for adding this, since it's a user-facing async API. It spawns async tasks, so it is a perfect place to provide a way to cancel the spawned tasks as well.

The concrete API that I propose is:

• Relevant methods of ClangdServer (we could start we code completion) that spawn async reqeuests returns a cancellation object that allows to cancel the spawned job:

class TaskHandle {
  void cancel();
};

class ClangdServer {
  TaskHandle codeComplete(File, Position, Callback);
};

• Users of the API are free to ignore the TaskHandle if they don't need to cancel the corresponding requests.
• If the users are capable of cancelling the request, they are responsible for keeping the TaskHandle around before the corresponding callback is called.
• If the request was cancelled, the Callback might receive the TaskCancelledError in the callback and should handle it accordingly.

The rest of the design LG, i.e. the code that can actually be cancelled is responsible for checking CancellationHandler::HasCancelled and should return a corresponding error in the callback if it was cancelled.

clangd/Cancellation.h

1 ↗

(On Diff #159888)

This file could use some comments about the API and samples of how it can be used.
Could probably be done after we finalize the design, though.

22 ↗

(On Diff #159888)

I think we might want to expose the context key for the cancellation primitive here.
The reason is that checking if request is cancelled is a common operation (i.e. I can imagine it being called on every few parsed ast node, etc.), so we'd want to keep the overhead of checking for cancellation very low.

Exposing a key in the context would allow to avoid doing the lookups in the context every time we need to check if request was cancelled.

Note that exposing shared_ptr<atomic<bool>> is probably not a great idea, since it'll allow to write into that object too. I suggest we go with a class that wraps it and only allows to read the value of the variable, i.e. only has isCancelled method.

23 ↗

(On Diff #159888)

The LLVM_NODISCARD is misplaced. The current code applies the attribute to the type (i.e. Context), and we want to apply to the function.
There are two ways to do that:

1. LLVM_NODISCARD static Context SetCurrentCancellationToken()
2. static Context SetCurrentCancellationToken LLVM_NODISCARD ()

The first one is definitely less awkward and more widely used, so I suggest we stick to it

25 ↗

(On Diff #159888)

NIT: use lowerCamelCase for function names (per LLVM style guide)

36 ↗

(On Diff #159888)

It seems the contract for ErrorInfo actually requires this function to be implemented, so using llvm_unreachable doesn't seem right here.
We could get away without actually unique error codes by returning llvm::inconvertibleErrorCode(), though.

clangd/ClangdLSPServer.cpp

76 ↗

(On Diff #159888)

1. Maybe use getAsInteger() to convert directly to int64_t?
2. Maybe use itostr to avoid conversion from int64 to uint64, which is undefined for negative numbers?
3. What if input is neither string nor number? Maybe add an assertion and check at the call-sites? That's a form of user input and we don't want clangd to crash in case of incorrect ones.

Alternatively, could we simply pipe the json::Value to a raw_string_stream and get the pretty-printed string? Would avoid dealing with input validation issues.

clangd/ClangdLSPServer.h

170 ↗

(On Diff #159888)

Maybe clarify in the comment that the key of the map is a serialized request-id?

[kadircet]

• Polished API.
• Resolved discussions.

[ilya-biryukov]

Sorry for missing this one. The biggest concern that I have is about the thread-safety of the API, it's too easy to misuse from multiple threads at this point.

We should define thread-safety guarantees for TaskHandle and CancellationToken and make sure they're not easily broken. The suggested way is to make both move-only types.
For CancellationToken we might need a non-thread-safe copy operation, though. But maybe make it a separate method (clone?) with a clear statement that it's not thread-safe?

Other than that just NITs.

clangd/Cancellation.cpp

17 ↗

(On Diff #160531)

Having a shared_ptr key in the Context can cause data races (e.g. if we copy it concurrently from multiple threads).
I suggest we make CancellationToken move-only (i.e. disallow copies) and return const CancellationToken& when getting it from the context.

34 ↗

(On Diff #160531)

I'm not sure if this function buys us much in terms of clarity, maybe remove it and inline everywhere?

clangd/Cancellation.h

9 ↗

(On Diff #160531)

The comments seems to mention all the important bits, but maybe we could separate different concerns more clearly based on how the code should be used?
E.g. here's a rough draft to illustrate the idea:

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(function<void(llvm::Expected<ResultType>)> Callback);

The callers of async methods (2) can issue cancellations and should be prepared to handle `TaskCancelledError` result:
    TaskHandle H = someAsyncMethod([](llvm::Expected<ResultType> R) {
        if (/* code to check for TaskCancelledError */)
           llvm::errs() << "Task got cancelled";
         else
           llvm::errs() << "The result is: " << *R;        
        } 
    });

    sleep(5);
    H.cancel();
    
The worker code itself (3) should check for cancellations using `CancellationToken` that can be retrieved via `CancellationToken::isCancelled()`.
    void codeComplete() {
      const CancellationToken& CT = CancellationToken::isCancelled();
      if (CT.isCancelled()) 
        return llvm::makeError<TaskCancelledError>();

      runSema();
      if (CT.isCancelled()) 
        return llvm::makeError<TaskCancelledError>();

      queryIndex();  
      return results;
    }
If the operation was cancelled before it could run to completion, it should propagate the TaskCancelledError as a result.

60 ↗

(On Diff #160531)

We're missing docs on this and the other classes. I suggest a small doc describing what they're used for, e.g. used for checking if the operation was cancelled and used to signal the operation should be cancelled, etc.

Also noting the threading-safety guarantees is a probably a good idea.

61 ↗

(On Diff #160531)

NIT: move members to the end of the class to be consistent with the rest of clangd.

67 ↗

(On Diff #160531)

Can we make this ctor private?

71 ↗

(On Diff #160531)

I wonder if we should make TaskHandle move-only?

The reasoning is that is can be easily used from multiple threads (since it's used by code dealing with async tasks anyway) and copying it concurrently is a data race.
On the other hand, calling cancel() is perfectly thread-safe and shouldn't cause any problems.

74 ↗

(On Diff #160531)

NIT: maybe use a shorter name, e.g. TaskHandle::create?

82 ↗

(On Diff #160531)

Maybe remove CancellationHandler and put all its functions into the clangd namespace? It's somewhat easily confused with CancellationToken.
We could declare some function friends if we need to

89 ↗

(On Diff #160531)

NIT: maybe use a shorter name: CancelledError?

clangd/ClangdLSPServer.cpp

355 ↗

(On Diff #160531)

I expect the pattern if cancelled -> reply with ErrorCode::RequestCancelled else reply with ErrorCode::InvalidParams to be pretty common when we add cancellation to more operations.

Maybe add a small helper that does that? E.g.

/// Replies to a request with an error, passing proper error codes to the LSP client.
/// Properly handles cancellations.
void replyError(Error E);

627 ↗

(On Diff #160531)

NIT: use Lock or L for consistencty with the rest of clangd. it's somewhat easy to misread std::lock_guard<std::mutex> _(Mut) as std::lock_guard<std::mutex> (Mut)

[kadircet]

• Get rid of getCancellationError.
• Add replyError helper.

[kadircet]

clangd/Cancellation.cpp

17 ↗

(On Diff #160531)

As talked offline, copying std::shared_ptr is thread safe.

[kadircet]

• Resolve discussions.
• Get rid of CancellationHandler class.
• Change error class name.
• Improve documentation.

[kadircet]

clangd/Cancellation.h

71 ↗

(On Diff #160531)

As similar to above mentioned case copying doesn't introduce any data races.

[kadircet]

• Made TaskHandle move-only. Since it is costly and most likely unnecessary to copy it other than to move it into Context.
• Provided an explicit clone method for copying.

[ilya-biryukov]

Mostly LSP-specific comments and comments about making TaskHandle thread-safe.

I'm also starting to wonder if the scope of this patch is too big, we could potentially split it into three separate bits:

1. Generic cancellation API, i.e. Cancellation.h and unit-tests of the API,
2. ClangdServer using the cancellation API.
3. LSP-specific bits, i.e. cancelRequest etc.

Not sure if it's worth splitting the last two, but splitting out the Cancellation API in a separate change could make review simpler.
But also happy to review the whole thing and land after it's done, feel free to pick whatever suits you best.

clangd/Cancellation.cpp

17 ↗

(On Diff #160531)

My mistake, thanks for explaining this one to me.
Copying of CancellationTokens LG!

clangd/Cancellation.h

101 ↗

(On Diff #160652)

s/asyn/async

[ilya-biryukov]

clangd/Cancellation.h

86 ↗

(On Diff #160652)

NIT: Maybe replace 'Totally thread-safe' with just Thread-safe? The latter conveys the same information.

90 ↗

(On Diff #160652)

NIT: maybe use Token->load() instead of static_cast<bool>(Token)? Arguably, looks more concise.

92 ↗

(On Diff #160652)

It's a bit confusing that this name clashes with a member function.
We seem to optimize for making this work anywhere:

if (isCancelled())  // <-- get token from ctx, then use implicit conversion to check if it was cancelled
  return llvm::makeError<CancelledError>();

Which is probably a nice pattern. But we could have two functions that do the same thing with less magic (i.e. implicit conversions) involved at the call sites:

CancellationToken getCurrentCancellation();
void setCurrentCancellation(CancellationToken CT);

inline bool isCancelled() { return getCurrentCancellation().isCancelled(); }

Would allow to get rid of implicit conversion and friend function with the same signature as member. Would also make cases where we actually want to stash the token somewhere cleaner, e.g. instead of CancellationToken CT = isCancelled(), we'll have CancellationToken CT = getCurrentCancellation().
WDYT?

104 ↗

(On Diff #160652)

As discussed offline, maybe we should make TaskHandle move-only?
Conceptually, TaskHandle is an abstraction that represent some spawned async task and in that sense somewhat similar to futures.

108 ↗

(On Diff #160652)

I wonder if we should have a createCancellationToken() method and a method to stash it in the context instead?
Would help if we want to make TaskHandle move-only and in general seems like a simpler user API, because things you put and get back out of the context is the same, i.e. it's a CancellationToken.
WDYT?

clangd/ClangdLSPServer.cpp

74 ↗

(On Diff #160652)

NIT: return std::string, const does not buys us anything here.

354 ↗

(On Diff #160652)

NIT: remove braces around short, single-statement branches

631 ↗

(On Diff #160652)

Use a value type here, i.e. std::string instead of const std::string&.

642 ↗

(On Diff #160652)

Use a value type here, i.e. std::string instead of const std::string&.

643 ↗

(On Diff #160652)

NIT: remove these braces, they don't buy us anything here?

clangd/ClangdLSPServer.h

178 ↗

(On Diff #160652)

I'm not sure what associated with only their thread means.
The tokens are associated with the request-ids rather than threads, so having something like ..., they create and remove cancellation handles for the request-id stored in the current Context

179 ↗

(On Diff #160652)

These two methods look like a nice suit for better wrapped in a RAII-object. It would add the value to the associated map on construction and remove on destruction.
One problem is that all operations calls would become weird, because we'll have to move this RAII object around and moving into lambdas is not fun without C++14 (which allows [x=std::move(y)]() {}..

Maybe add a comment that calls to this function should always be paired and a FIXME that we should have an API that enforces this?

clangd/ClangdServer.cpp

167 ↗

(On Diff #160652)

NIT: remove braces around single-statement branches

clangd/JSONRPCDispatcher.cpp

135 ↗

(On Diff #160652)

Isn't there a form of handleErrors that can handle the second case in the last argument?
I.e. can we merge the following if into the handleErrors call?

clangd/JSONRPCDispatcher.h

67 ↗

(On Diff #160652)

Maybe adds docs explaining what this function does?
Specifically, what ErrorCode it chooses and how the message would look like

115 ↗

(On Diff #160652)

naming-nit: use lowerCamelCase for function names

115 ↗

(On Diff #160652)

I suggest we document this function and put to next to reply and replyError, since they are closely related (i.e. the latter put the current request-id into the replies).
We should probably discourage using this function in its documentation unless absolutely necessary, the users should probably directly call reply or replyError instead.

unittests/clangd/CancellationTests.cpp

64 ↗

(On Diff #160652)

Another interesting case to test:
create multiple threads, propagate cancellation into them and check they get properly cancelled.
Using Notification from Threading.h should allow to write it in a pretty simplistic manner.

[ilya-biryukov]

The last set of comments is for the previous version, might be missing some updates, sorry about that. Will make sure to review the one too!

[kadircet]

• Resolve discussions.

[kadircet]

clangd/Cancellation.h

92 ↗

(On Diff #160652)

Totally agreed, actually, didn't even notice they had the same signature :D, thanks!
But I still think implicit conversion of CancellationToken to bool seems like a good thing to have.

clangd/ClangdLSPServer.h

179 ↗

(On Diff #160652)

Yeah that would be really nice to have them in a RAII object, but the thing is cleanup can occur only after the task dies. Which we don't directly know since threads are being detached. So the only way to trigger destruction is at the callback, which is where I call cleanup currently. Maybe we can store the object within context to trigger destructor at the end of the thread, but still don't think it would look any better or ease the job of the caller.

[ilya-biryukov]

Mostly NITs, but please take a look at the CancelParams handling problem. I believe there might be a potential bug hiding there :-)

clangd/Cancellation.h

87 ↗

(On Diff #161018)

s/Can be created with clangd::isCancelled()/Tokens for the currently running task can be obtained via clangd::getCurrentCancellationToken/?
(to focus on which kind of tokens can be obtained from the tokens)

Another idea: maybe add a reference no why someone would want to use CancellationToken instead of isCancelled() helper, e.g.

/// If cancellation checks are rare, one could use the isCancelled() helper to simplify the code.
/// However, if cancellation checks are frequent, the guideline is first obtain the CancellationToken for the currently running task with getCurrentCancellationToken() and do cancel checks using it to avoid extra lookups in the Context

137 ↗

(On Diff #161018)

s/got/was?

clangd/ClangdLSPServer.cpp

621 ↗

(On Diff #161018)

Wouldn't it work incorrectly for string IDs?
When normalizing json::Value from getRequestID, we simply print out json. For strings, this should result in quoted output, e.g. "req". However, when parsing CancelParams, we parse this json. So we'll end up inserting with a key "req" and erasing with a key 'req' (without the quotes).

643 ↗

(On Diff #161018)

Could we elog if insert indicated that a value already exists in a map? This would probably mean malformed user input, but could also be immensely helpful for debugging.

clangd/ClangdLSPServer.h

179 ↗

(On Diff #160652)

Having these paired method calls is probably fine for the first iteration, but we should abstract away this pattern when supporting cancellations for more operations.
LG with a FIXME.

clangd/ClangdServer.cpp

152 ↗

(On Diff #161018)

Maybe use a shorter name? E.g. Task or TH :-)

153 ↗

(On Diff #161018)

Instead of moving cancellation token around explicitly, we could stash in a context right away and it would nicely propagate into the thread. This looks like less boilerplate.
E.g.

auto CancellableTaskHandle = TaskHandle::create();
WithContext ContextWithCancellation(setCurrentCancellationToken(TaskHandle.createCancellationToken()));

auto Task = [](File, Callback<CodeCompleteResult> CB, llvm::Expected<InputsAndPreamble> IP) {
  if (getCancellationToken().isCancelled())
    return CB(make_error<CancelledError>());
 };
 /// rest of the code is the same

178 ↗

(On Diff #161018)

Thinking out loud: the TaskHandle and CancellationToken handling will probably eventually move into TUScheduler::runWithPreamble and other thread-spawning code to avoid this kind of boilerplate here.

No need to change anything in this patch, though.

clangd/Protocol.h

871 ↗

(On Diff #161018)

Maybe add a comment from the LSP spec to be consistent with the rest of the file?
Also note that this can be either a number or a string, but we choose to print out a number in that case and just always use strings.

[kadircet]

• Resolve discussions & Delete enclosing quotes when normalizing.

[kadircet]

clangd/ClangdLSPServer.cpp

621 ↗

(On Diff #161018)

You are right, that's exactly the case, thanks!

[ilya-biryukov]

clangd/ClangdLSPServer.cpp

80 ↗

(On Diff #161261)

This still misses string escaping issues. E.g. " will get quoted as \", etc.
I suggest we actually switch on this json value kind.
Better yet, reuse the part of fromJSON(CancelParams) that does the request-id parsing.

clangd/ClangdServer.cpp

166 ↗

(On Diff #161261)

NIT: check for it even before looking at the preamble?

[kadircet]

• Check cancellation earlier && Fix normalization difference between string and integer request ids.

[ioeric]

clangd/Cancellation.h

96 ↗

(On Diff #161428)

As chatted offline, I have questions about the motivation of the CancellationToken class. Intuitively, it seems to me that this can be merged into TaskHandle, and we can simply stash the TaskHandle instead of a token into the context. There would be fewer states to maintain, and the design would be a bit more straightforward. I might be missing context/concerns here, so I'd like to hear what you and Ilya think. @ilya-biryukov

clangd/ClangdLSPServer.cpp

75 ↗

(On Diff #161428)

Use ID.getAsString()?

[kadircet]

clangd/ClangdLSPServer.cpp

75 ↗

(On Diff #161428)

Unfortunately id can be both a string or a number according to LSP specs therefore we normalize both cases into a string.

[ilya-biryukov]

clangd/Cancellation.h

96 ↗

(On Diff #161428)

I am for splitting cancellation checks from cancellation triggers.
The processing code only needs to check if it was cancelled and exposing the cancel() does not add any useful functionality, merely ways to misuse it.

[ioeric]

clangd/Cancellation.h

96 ↗

(On Diff #161428)

Couldn't we prevent this by passing only const handle references to users who are not expect to call cancel()?

[ilya-biryukov]

clangd/Cancellation.h

96 ↗

(On Diff #161428)

As long as TaskHandle does not have a copy constructor, that LG.
I.e. the scheme that seems to work is:

1. return shared_ptr<TaskHandle> from async computations like code complete.
2. stash shared_ptr<const TaskHandle> into the context.
3. return const TaskHandle* from the context.

This should give a simpler implementation.
The cons from my point of view are:

1. We return copyable task handles (i.e. the shared_ptr) from the API by default. This provokes usages that don't take ownership issues into account, i.e. it's easy to just copy the task everywhere without having clear ownership of it. As long as it's just cancellation there, we're probably fine, but if we add more stuff into it that might become a problem.
2. We expose implementation details (i.e. shared_ptr) by removing the layer of abstraction. This makes potential refactorings of the client code harder.

In general, I think having our own wrappers there opens up more room for experiments with the API later (the API is very constrained, so it's easy to refactor the implementation). OTOH, if we believe the design is right or changing the clients is not too much work (which is probably true), either way is fine with me.

But either solution has its pros and cons, I'm happy with trying out any of those on code completion, see how it goes and decide whether we want to change anything before finalizing our design and adding cancellation to all operations in clangd.

[kadircet]

• Get rid of CancellationToken && Resolve discussions.

[ilya-biryukov]

LG, just a few last nits

clangd/Cancellation.cpp

30 ↗

(On Diff #161716)

NIT: maybe consider inlining it? Since Task has a public default constructor, I don't think having this method buys us anything

clangd/Cancellation.h

22 ↗

(On Diff #161716)

s/Taskr/Task

23 ↗

(On Diff #161716)

remove mentions of CancellationToken

49 ↗

(On Diff #161716)

s/CancellationToken/getCurrentTask().isCancelled()

clangd/ClangdLSPServer.cpp

76 ↗

(On Diff #161716)

Maybe extract the relevant code into a helper to avoid creating the unrelated class (CancelParams)?

621 ↗

(On Diff #161716)

Invert condition to reduce nesting?
See LLVM Style Guide

[kadircet]

• Resolve discussions.

[ilya-biryukov]

Mostly NITs, but please take a look at the potential data race

clangd/Cancellation.cpp

23 ↗

(On Diff #161908)

Maybe add 'assert' that TaskKey is present?
To make sure we see the breakages as early as possible when asserts are enabled (otherwise it's undefined behavior and it can propagate to breakages in client code)

clangd/Cancellation.h

27 ↗

(On Diff #161908)

s/propagete/propagate

114 ↗

(On Diff #161908)

Maybe add a comment that the current Context must be have the task stashed in it?

clangd/ClangdLSPServer.cpp

351 ↗

(On Diff #161908)

CleanupTaskHandle() can run in a separate thread, so can potentially run before the StoreTaskHandle call.
To avoid memory leaks in that case, let's preallocate the entry before calling codeComplete

621 ↗

(On Diff #161908)

s/it/It

clangd/Protocol.cpp

629 ↗

(On Diff #161908)

Maybe switch on different kinds of json::Value instead of using the ObjectMapper?
The code should be simpler. Albeit the fromJson of CancelParams might get a little more complicated, but that looks like a small price to pay.

clangd/Protocol.h

883 ↗

(On Diff #161908)

Maybe simplify the signature to: Optional<string> parseNumberOrString(const llvm::json::Value&);
And call as parseNumberOrString(obj["id"]) instead of parseNumberOrString(obj, "id")?

Also, maybe call it normalizeRequestId?

[kadircet]

• Change a few comments.
• Add some assertions.
• Fix a data race.

[kadircet]

clangd/ClangdLSPServer.cpp

351 ↗

(On Diff #161908)

Thanks!

[kadircet]

• Change helper for normalizing Number | String.

clangd/Protocol.h

883 ↗

(On Diff #161908)

LSP has a few more places that can have paremeters of type number | string so I believe this one could be used with them as well. Therefore I think it is better to keep the name that way. The normalizeRequestID is a sepcial case and handled within ClangdLSPServer.

[ilya-biryukov]

LG with a few small comments

clangd/Protocol.cpp

635 ↗

(On Diff #162354)

Maybe use itostr?

641 ↗

(On Diff #162354)

What is Params is not an object?

[kadircet]

• Resolve comments.

clang-tools-extra/trunk/clangd/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   Support
   )
 
 set(CLANGD_ATOMIC_LIB "")
 if(NOT HAVE_CXX_ATOMICS64_WITHOUT_LIB)
   list(APPEND CLANGD_ATOMIC_LIB "atomic")
 endif()
 
 add_clang_library(clangDaemon
   AST.cpp
+  Cancellation.cpp
   ClangdLSPServer.cpp
   ClangdServer.cpp
   ClangdUnit.cpp
   CodeComplete.cpp
   CodeCompletionStrings.cpp
   Compiler.cpp
   Context.cpp
   Diagnostics.cpp

clang-tools-extra/trunk/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;
+// }
+//
+// 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/trunk/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;
+
+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/trunk/clangd/ClangdLSPServer.h

private:
   void onSwitchSourceHeader(TextDocumentIdentifier &Params) override;
   void onDocumentHighlight(TextDocumentPositionParams &Params) override;
   void onFileEvent(DidChangeWatchedFilesParams &Params) override;
   void onCommand(ExecuteCommandParams &Params) override;
   void onWorkspaceSymbol(WorkspaceSymbolParams &Params) override;
   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);
 
   /// 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);
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;
+  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
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDLSPSERVER_H

clang-tools-extra/trunk/clangd/ClangdLSPServer.cpp

 //===--- ClangdLSPServer.cpp - LSP server ------------------------*- C++-*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 
 using namespace clang::clangd;
 using namespace clang;
 using namespace llvm;
 

 SymbolKindBitset defaultSymbolKinds() {
   SymbolKindBitset Defaults;
   for (size_t I = SymbolKindMin; I <= static_cast<size_t>(SymbolKind::Array);
        ++I)
     Defaults.set(I);
   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) {
   if (Params.initializationOptions)
     applyConfiguration(*Params.initializationOptions);
 
   if (Params.rootUri && *Params.rootUri)
     Server.setRootPath(Params.rootUri->file());
for (auto &F : getFixes(Params.textDocument.uri.file(), D)) {
           {"arguments", {WE}},
       });
     }
   }
   reply(std::move(Commands));
 }
 
 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,
+        auto _ = llvm::make_scope_exit([this]() { CleanupTaskHandle(); });
+
+          return replyError(List.takeError());
+        return reply(std::move(LSPList));
+  StoreTaskHandle(std::move(TH));
 }
 
 void ClangdLSPServer::onSignatureHelp(TextDocumentPositionParams &Params) {
   Server.signatureHelp(Params.textDocument.uri.file(), Params.position,
                        [](llvm::Expected<SignatureHelp> SignatureHelp) {
                          if (!SignatureHelp)
                            return replyError(
                                ErrorCode::InvalidParams,
                                llvm::toString(SignatureHelp.takeError()));
                          reply(*SignatureHelp);
                        });
 }
 
 void ClangdLSPServer::onGoToDefinition(TextDocumentPositionParams &Params) {
-  Server.findDefinitions(
-      Params.textDocument.uri.file(), Params.position,
+  Server.findDefinitions(Params.textDocument.uri.file(), Params.position,
                          [](llvm::Expected<std::vector<Location>> Items) {
                            if (!Items)
-          return replyError(ErrorCode::InvalidParams,
+                             return replyError(
+                                 ErrorCode::InvalidParams,
                                  llvm::toString(Items.takeError()));
                            reply(json::Array(*Items));
                          });
 }
 
 void ClangdLSPServer::onSwitchSourceHeader(TextDocumentIdentifier &Params) {
   llvm::Optional<Path> Result = Server.switchSourceHeader(Params.uri.file());
   reply(Result ? URI::createFile(*Result).toString() : "");
 }
 
 void ClangdLSPServer::onDocumentHighlight(TextDocumentPositionParams &Params) {
void ClangdLSPServer::CompilationDB::setCompileCommandsDir(Path P) {
   CachingCDB->clear();
 }
 
 GlobalCompilationDatabase &ClangdLSPServer::CompilationDB::getCDB() {
   if (CachingCDB)
     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/trunk/clangd/ClangdServer.h

 //===--- ClangdServer.h - Main clangd server code ----------------*- 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_CLANGDSERVER_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDSERVER_H
 
+#include "Cancellation.h"
 #include "ClangdUnit.h"
 #include "CodeComplete.h"
 #include "FSProvider.h"
 #include "Function.h"
 #include "GlobalCompilationDatabase.h"
 #include "Protocol.h"
 #include "TUScheduler.h"
 #include "index/FileIndex.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.
-  void codeComplete(PathRef File, Position Pos,
+  TaskHandle 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.
   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);

clang-tools-extra/trunk/clangd/ClangdServer.cpp

 //===--- ClangdServer.cpp - Main clangd server code --------------*- C++-*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===-------------------------------------------------------------------===//
 
 #include "ClangdServer.h"
+#include "Cancellation.h"
 #include "CodeComplete.h"
 #include "FindSymbols.h"
 #include "Headers.h"
 #include "SourceCode.h"
 #include "Trace.h"
 #include "XRefs.h"
 #include "index/Merge.h"
 #include "clang/Format/Format.h"
WorkScheduler.update(File, std::move(Inputs), WantDiags,
                        });
 }
 
 void ClangdServer::removeDocument(PathRef File) {
   ++InternalVersion[File];
   WorkScheduler.remove(File);
 }
 
-void ClangdServer::codeComplete(PathRef File, Position Pos,
+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());
 
     auto PreambleData = IP->Preamble;
 
     llvm::Optional<SpeculativeFuzzyFind> SpecFuzzyFind;
     if (CodeCompleteOpts.Index && CodeCompleteOpts.SpeculativeIndexRequest) {
       SpecFuzzyFind.emplace();
auto Task = [PCHs, Pos, FS, CodeCompleteOpts,
     // SpecFuzzyFind is only destroyed after speculative fuzzy find finishes.
     // 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.
   };
 
   WorkScheduler.runWithPreamble("CodeComplete", File,
                                 Bind(Task, File.str(), std::move(CB)));
+  return TH;
 }
 
 void ClangdServer::signatureHelp(PathRef File, Position Pos,
                                  Callback<SignatureHelp> CB) {
 
   auto PCHs = this->PCHs;
   auto FS = FSProvider.getFileSystem();
   auto *Index = this->Index;

clang-tools-extra/trunk/clangd/JSONRPCDispatcher.h

 };
 
 /// 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 replyError(ErrorCode Code, const llvm::StringRef &Message);
+/// Implements ErrorCode and message extraction from a given llvm::Error. It
+/// fetches the related message from error's message method. If error doesn't
+/// match any known errors, uses ErrorCode::InvalidParams for the error.
+void replyError(llvm::Error E);
+/// Returns the request-id of the current request. Should not be used directly
+/// for replying to requests, use the above mentioned methods for that case.
+const llvm::json::Value *getRequestId();
 /// Sends a request to the client.
 /// 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.
 class JSONRPCDispatcher {
 public:

 /// 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

clang-tools-extra/trunk/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 "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"
void JSONOutput::mirrorInput(const Twine &Message) {
   if (!InputMirror)
     return;
 
   *InputMirror << Message;
   InputMirror->flush();
 }
 
 void clangd::reply(json::Value &&Result) {
-  auto ID = Context::current().get(RequestID);
+  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)},
       });
 }
 
 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 = Context::current().get(RequestID)) {
+  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}}},
         });
   }
 }
 
+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);
if (auto JSON = ReadMessage(In, Out)) {
       } else {
         // Parse error. Log the raw message.
         vlog("<<< {0}\n", *JSON);
         elog("JSON parse error: {0}", llvm::toString(Doc.takeError()));
       }
     }
   }
 }
+
+const json::Value *clangd::getRequestId() {
+  return Context::current().get(RequestID);
+}

clang-tools-extra/trunk/clangd/Protocol.h

struct DocumentHighlight {
   friend bool operator==(const DocumentHighlight &LHS,
                          const DocumentHighlight &RHS) {
     return LHS.kind == RHS.kind && LHS.range == RHS.range;
   }
 };
 llvm::json::Value toJSON(const DocumentHighlight &DH);
 llvm::raw_ostream &operator<<(llvm::raw_ostream &, const DocumentHighlight &);
 
+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
 
 namespace llvm {
 template <> struct format_provider<clang::clangd::Position> {
   static void format(const clang::clangd::Position &Pos, raw_ostream &OS,
                      StringRef Style) {
     assert(Style.empty() && "style modifiers for this type are not supported");
     OS << Pos;
   }
 };
 } // namespace llvm
 
 #endif

clang-tools-extra/trunk/clangd/Protocol.cpp

 
 bool fromJSON(const json::Value &Params,
               ClangdConfigurationParamsChange &CCPC) {
   json::ObjectMapper O(Params);
   return O && O.map("compilationDatabasePath", CCPC.compilationDatabasePath) &&
          O.map("compilationDatabaseChanges", CCPC.compilationDatabaseChanges);
 }
 
+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

clang-tools-extra/trunk/clangd/ProtocolHandlers.h

public:
   virtual void onSwitchSourceHeader(TextDocumentIdentifier &Params) = 0;
   virtual void onFileEvent(DidChangeWatchedFilesParams &Params) = 0;
   virtual void onCommand(ExecuteCommandParams &Params) = 0;
   virtual void onWorkspaceSymbol(WorkspaceSymbolParams &Params) = 0;
   virtual void onRename(RenameParams &Parames) = 0;
   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,
                               ProtocolCallbacks &Callbacks);
 
 } // namespace clangd
 } // namespace clang
 
 #endif

clang-tools-extra/trunk/clangd/ProtocolHandlers.cpp

void clangd::registerCallbackHandlers(JSONRPCDispatcher &Dispatcher,
   Register("textDocument/documentSymbol", &ProtocolCallbacks::onDocumentSymbol);
   Register("workspace/didChangeWatchedFiles", &ProtocolCallbacks::onFileEvent);
   Register("workspace/executeCommand", &ProtocolCallbacks::onCommand);
   Register("textDocument/documentHighlight",
            &ProtocolCallbacks::onDocumentHighlight);
   Register("workspace/didChangeConfiguration",
            &ProtocolCallbacks::onChangeConfiguration);
   Register("workspace/symbol", &ProtocolCallbacks::onWorkspaceSymbol);
+  Register("$/cancelRequest", &ProtocolCallbacks::onCancelRequest);
 }

clang-tools-extra/trunk/unittests/clangd/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   support
   )
 
 get_filename_component(CLANGD_SOURCE_DIR
   ${CMAKE_CURRENT_SOURCE_DIR}/../../clangd REALPATH)
 include_directories(
   ${CLANGD_SOURCE_DIR}
   )
 
 add_extra_unittest(ClangdTests
   Annotations.cpp
+  CancellationTests.cpp
   ClangdTests.cpp
   ClangdUnitTests.cpp
   CodeCompleteTests.cpp
   CodeCompletionStringsTests.cpp
   ContextTests.cpp
   DexIndexTests.cpp
   DraftStoreTests.cpp
   FileDistanceTests.cpp

clang-tools-extra/trunk/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));
+  EXPECT_FALSE(isCancelled());
+  TH->cancel();
+  EXPECT_TRUE(isCancelled());
+}
+
+TEST(CancellationTest, TaskTestHandleDiesContextLives) {
+  llvm::Optional<WithContext> ContextWithCancellation;
+  {
+    TaskHandle TH = Task::createHandle();
+    ContextWithCancellation.emplace(setCurrentTask(TH));
+    EXPECT_FALSE(isCancelled());
+    TH->cancel();
+    EXPECT_TRUE(isCancelled());
+  }
+  EXPECT_TRUE(isCancelled());
+}
+
+TEST(CancellationTest, TaskContextDiesHandleLives) {
+  TaskHandle TH = Task::createHandle();
+  {
+    WithContext ContextWithCancellation(setCurrentTask(TH));
+    EXPECT_FALSE(isCancelled());
+    TH->cancel();
+    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());
+}
+
+TEST(CancellationTest, AsynCancellationTest) {
+  std::atomic<bool> HasCancelled(false);
+  Notification Cancelled;
+  auto TaskToBeCancelled = [&](ConstTaskHandle CT) {
+    WithContext ContextGuard(setCurrentTask(std::move(CT)));
+    Cancelled.wait();
+    HasCancelled = isCancelled();
+  };
+  TaskHandle TH = Task::createHandle();
+  std::thread AsyncTask(TaskToBeCancelled, TH);
+  TH->cancel();
+  Cancelled.notify();
+  AsyncTask.join();
+
+  EXPECT_TRUE(HasCancelled);
+}
+} // namespace
+} // namespace clangd
+} // namespace clang